ktid/kanhole
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: clean up code (#5308)

Browse Source
This commit is contained in:
fatedier
2026-05-12 11:13:50 +08:00
committed by GitHub
Unverified
parent ad07d27914
commit a88e0e9a49
49 changed files with 2082 additions and 931 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/nathole/analysis.go
+41 -23
View File
@@ -21,6 +21,7 @@ import (
"time"
"github.com/samber/lo"
"k8s.io/utils/clock"
)
var (
@@ -144,19 +145,19 @@ func getBehaviorByModeAndIndex(mode int, index int) (RecommandBehavior, Recomman
return behaviors[index].A, behaviors[index].B
}
func getBehaviorScoresByMode(mode int, defaultScore int) []*BehaviorScore {
func getBehaviorScoresByMode(mode int, defaultScore int) []*behaviorScore {
return getBehaviorScoresByMode2(mode, defaultScore, defaultScore)
}
func getBehaviorScoresByMode2(mode int, senderScore, receiverScore int) []*BehaviorScore {
func getBehaviorScoresByMode2(mode int, senderScore, receiverScore int) []*behaviorScore {
behaviors := getBehaviorByMode(mode)
scores := make([]*BehaviorScore, 0, len(behaviors))
scores := make([]*behaviorScore, 0, len(behaviors))
for i := range behaviors {
score := receiverScore
if behaviors[i].A.Role == DetectRoleSender {
score = senderScore
}
scores = append(scores, &BehaviorScore{Mode: mode, Index: i, Score: score})
scores = append(scores, &behaviorScore{Mode: mode, Index: i, Score: score})
}
return scores
}
@@ -170,14 +171,18 @@ type RecommandBehavior struct {
ListenRandomPorts int
}
type MakeHoleRecords struct {
type makeHoleRecords struct {
mu sync.Mutex
scores []*BehaviorScore
LastUpdateTime time.Time
scores []*behaviorScore
clock clock.PassiveClock
lastUpdateTime time.Time
}
func NewMakeHoleRecords(c, v *NatFeature) *MakeHoleRecords {
scores := []*BehaviorScore{}
func newMakeHoleRecordsWithClock(c, v *NatFeature, clk clock.PassiveClock) *makeHoleRecords {
if clk == nil {
clk = clock.RealClock{}
}
scores := []*behaviorScore{}
easyCount, hardCount, portsChangedRegularCount := ClassifyFeatureCount([]*NatFeature{c, v})
appendMode0 := func() {
switch {
@@ -212,13 +217,17 @@ func NewMakeHoleRecords(c, v *NatFeature) *MakeHoleRecords {
scores = append(scores, getBehaviorScoresByMode(DetectMode1, 1)...)
scores = append(scores, getBehaviorScoresByMode(DetectMode3, 1)...)
}
return &MakeHoleRecords{scores: scores, LastUpdateTime: time.Now()}
return &makeHoleRecords{
scores: scores,
clock: clk,
lastUpdateTime: clk.Now(),
}
}
func (mhr *MakeHoleRecords) ReportSuccess(mode int, index int) {
func (mhr *makeHoleRecords) reportSuccess(mode int, index int) {
mhr.mu.Lock()
defer mhr.mu.Unlock()
mhr.LastUpdateTime = time.Now()
mhr.lastUpdateTime = mhr.clock.Now()
for i := range mhr.scores {
score := mhr.scores[i]
if score.Mode != mode || score.Index != index {
@@ -231,22 +240,22 @@ func (mhr *MakeHoleRecords) ReportSuccess(mode int, index int) {
}
}
func (mhr *MakeHoleRecords) Recommand() (mode, index int) {
func (mhr *makeHoleRecords) recommand() (mode, index int) {
mhr.mu.Lock()
defer mhr.mu.Unlock()
if len(mhr.scores) == 0 {
return 0, 0
}
maxScore := slices.MaxFunc(mhr.scores, func(a, b *BehaviorScore) int {
maxScore := slices.MaxFunc(mhr.scores, func(a, b *behaviorScore) int {
return cmp.Compare(a.Score, b.Score)
})
maxScore.Score--
mhr.LastUpdateTime = time.Now()
mhr.lastUpdateTime = mhr.clock.Now()
return maxScore.Mode, maxScore.Index
}
type BehaviorScore struct {
type behaviorScore struct {
Mode int
Index int
// between -10 and 10
@@ -255,16 +264,25 @@ type BehaviorScore struct {
type Analyzer struct {
// key is client ip + visitor ip
records map[string]*MakeHoleRecords
records map[string]*makeHoleRecords
dataReserveDuration time.Duration
clock clock.PassiveClock
mu sync.Mutex
}
func NewAnalyzer(dataReserveDuration time.Duration) *Analyzer {
return newAnalyzerWithClock(dataReserveDuration, clock.RealClock{})
}
func newAnalyzerWithClock(dataReserveDuration time.Duration, clk clock.PassiveClock) *Analyzer {
if clk == nil {
clk = clock.RealClock{}
}
return &Analyzer{
records: make(map[string]*MakeHoleRecords),
records: make(map[string]*makeHoleRecords),
dataReserveDuration: dataReserveDuration,
clock: clk,
}
}
@@ -272,12 +290,12 @@ func (a *Analyzer) GetRecommandBehaviors(key string, c, v *NatFeature) (mode, in
a.mu.Lock()
records, ok := a.records[key]
if !ok {
records = NewMakeHoleRecords(c, v)
records = newMakeHoleRecordsWithClock(c, v, a.clock)
a.records[key] = records
}
a.mu.Unlock()
mode, index = records.Recommand()
mode, index = records.recommand()
cBehavior, vBehavior := getBehaviorByModeAndIndex(mode, index)
switch mode {
@@ -307,11 +325,11 @@ func (a *Analyzer) ReportSuccess(key string, mode, index int) {
if !ok {
return
}
records.ReportSuccess(mode, index)
records.reportSuccess(mode, index)
}
func (a *Analyzer) Clean() (int, int) {
now := time.Now()
now := a.clock.Now()
total := 0
count := 0
@@ -321,7 +339,7 @@ func (a *Analyzer) Clean() (int, int) {
total = len(a.records)
// clean up records that have not been used for a period of time.
for key, records := range a.records {
if now.Sub(records.LastUpdateTime) > a.dataReserveDuration {
if now.Sub(records.lastUpdateTime) > a.dataReserveDuration {
delete(a.records, key)
count++
}
pkg/nathole/analysis_test.go
+33
View File
@@ -0,0 +1,33 @@
package nathole
import (
"testing"
"time"
"github.com/stretchr/testify/require"
clocktesting "k8s.io/utils/clock/testing"
)
func TestAnalyzerUsesClockForRecordTimestamps(t *testing.T) {
require := require.New(t)
start := time.Date(2026, time.May, 8, 12, 30, 0, 0, time.UTC)
clk := clocktesting.NewFakeClock(start)
analyzer := newAnalyzerWithClock(time.Hour, clk)
clientFeature := &NatFeature{NatType: EasyNAT, Behavior: BehaviorNoChange}
visitorFeature := &NatFeature{NatType: EasyNAT, Behavior: BehaviorNoChange}
mode, index, _, _ := analyzer.GetRecommandBehaviors("key", clientFeature, visitorFeature)
require.Equal(start, analyzer.records["key"].lastUpdateTime)
updatedAt := start.Add(time.Minute)
clk.SetTime(updatedAt)
analyzer.ReportSuccess("key", mode, index)
require.Equal(updatedAt, analyzer.records["key"].lastUpdateTime)
clk.SetTime(start.Add(2 * time.Hour))
count, total := analyzer.Clean()
require.Equal(1, count)
require.Equal(1, total)
require.Empty(analyzer.records)
}
pkg/nathole/controller.go
+42 -34
View File
@@ -326,40 +326,16 @@ func (c *Controller) analysis(session *Session) (*msg.NatHoleResp, *msg.NatHoleR
}
protocol := vm.Protocol
vResp := &msg.NatHoleResp{
TransactionID: vm.TransactionID,
Sid: session.sid,
Protocol: protocol,
CandidateAddrs: slices.Compact(cm.MappedAddrs),
AssistedAddrs: slices.Compact(cm.AssistedAddrs),
DetectBehavior: msg.NatHoleDetectBehavior{
Mode: mode,
Role: vBehavior.Role,
TTL: vBehavior.TTL,
SendDelayMs: vBehavior.SendDelayMs,
ReadTimeoutMs: timeoutMs - vBehavior.SendDelayMs,
SendRandomPorts: vBehavior.PortsRandomNumber,
ListenRandomPorts: vBehavior.ListenRandomPorts,
CandidatePorts: getRangePorts(cm.MappedAddrs, cNatFeature.PortsDifference, vBehavior.PortsRangeNumber),
},
}
cResp := &msg.NatHoleResp{
TransactionID: cm.TransactionID,
Sid: session.sid,
Protocol: protocol,
CandidateAddrs: slices.Compact(vm.MappedAddrs),
AssistedAddrs: slices.Compact(vm.AssistedAddrs),
DetectBehavior: msg.NatHoleDetectBehavior{
Mode: mode,
Role: cBehavior.Role,
TTL: cBehavior.TTL,
SendDelayMs: cBehavior.SendDelayMs,
ReadTimeoutMs: timeoutMs - cBehavior.SendDelayMs,
SendRandomPorts: cBehavior.PortsRandomNumber,
ListenRandomPorts: cBehavior.ListenRandomPorts,
CandidatePorts: getRangePorts(vm.MappedAddrs, vNatFeature.PortsDifference, cBehavior.PortsRangeNumber),
},
}
vResp := newNatHoleResponse(
vm.TransactionID, session.sid, protocol, mode,
cm.MappedAddrs, cm.AssistedAddrs, vBehavior,
timeoutMs-vBehavior.SendDelayMs, cNatFeature.PortsDifference,
)
cResp := newNatHoleResponse(
cm.TransactionID, session.sid, protocol, mode,
vm.MappedAddrs, vm.AssistedAddrs, cBehavior,
timeoutMs-cBehavior.SendDelayMs, vNatFeature.PortsDifference,
)
log.Debugf("sid [%s] visitor nat: %+v, candidateAddrs: %v; client nat: %+v, candidateAddrs: %v, protocol: %s",
session.sid, *vNatFeature, vm.MappedAddrs, *cNatFeature, cm.MappedAddrs, protocol)
@@ -368,6 +344,38 @@ func (c *Controller) analysis(session *Session) (*msg.NatHoleResp, *msg.NatHoleR
return vResp, cResp, nil
}
func newNatHoleResponse(
transactionID string,
sid string,
protocol string,
mode int,
candidateAddrs []string,
assistedAddrs []string,
behavior RecommandBehavior,
readTimeoutMs int,
portsDifference int,
) *msg.NatHoleResp {
compactCandidateAddrs := slices.Compact(candidateAddrs)
compactAssistedAddrs := slices.Compact(assistedAddrs)
return &msg.NatHoleResp{
TransactionID: transactionID,
Sid: sid,
Protocol: protocol,
CandidateAddrs: compactCandidateAddrs,
AssistedAddrs: compactAssistedAddrs,
DetectBehavior: msg.NatHoleDetectBehavior{
Mode: mode,
Role: behavior.Role,
TTL: behavior.TTL,
SendDelayMs: behavior.SendDelayMs,
ReadTimeoutMs: readTimeoutMs,
SendRandomPorts: behavior.PortsRandomNumber,
ListenRandomPorts: behavior.ListenRandomPorts,
CandidatePorts: getRangePorts(candidateAddrs, portsDifference, behavior.PortsRangeNumber),
},
}
}
func getRangePorts(addrs []string, difference, maxNumber int) []msg.PortsRange {
if maxNumber <= 0 {
return nil