engine: more performance improvements

2026-05-18 09:17:24 +05:30
parent 77dba0c4fa
commit 581041b1a7
12 changed files with 1005 additions and 103 deletions
@@ -22,6 +22,89 @@ func (r fixedRuleset) Match(ruleset.StreamInfo) ruleset.MatchResult {
 	return ruleset.MatchResult{Action: r.action}
 }

+type analyzerRuleset struct {
+	action ruleset.Action
+	ans    []analyzer.Analyzer
+}
+
+func (r analyzerRuleset) Analyzers(ruleset.StreamInfo) []analyzer.Analyzer {
+	return r.ans
+}
+
+func (r analyzerRuleset) Match(ruleset.StreamInfo) ruleset.MatchResult {
+	return ruleset.MatchResult{Action: r.action}
+}
+
+type countingTCPAnalyzer struct {
+	newCalls  *int
+	feedCalls *int
+}
+
+func (a countingTCPAnalyzer) Name() string { return "tls" }
+func (a countingTCPAnalyzer) Limit() int   { return 0 }
+func (a countingTCPAnalyzer) NewTCP(analyzer.TCPInfo, analyzer.Logger) analyzer.TCPStream {
+	(*a.newCalls)++
+	return countingTCPStream{feedCalls: a.feedCalls}
+}
+
+type countingTCPStream struct {
+	feedCalls *int
+}
+
+func (s countingTCPStream) Feed(bool, bool, bool, int, []byte) (*analyzer.PropUpdate, bool) {
+	(*s.feedCalls)++
+	return nil, false
+}
+
+func (s countingTCPStream) Close(bool) *analyzer.PropUpdate {
+	return nil
+}
+
+type logFinalizingRuleset struct {
+	ans []analyzer.Analyzer
+}
+
+func (r logFinalizingRuleset) Analyzers(ruleset.StreamInfo) []analyzer.Analyzer {
+	return r.ans
+}
+
+func (r logFinalizingRuleset) Match(info ruleset.StreamInfo) ruleset.MatchResult {
+	if _, ok := info.Props["tls"]; ok {
+		return ruleset.MatchResult{Action: ruleset.ActionMaybe, Logged: true}
+	}
+	return ruleset.MatchResult{Action: ruleset.ActionMaybe}
+}
+
+func (r logFinalizingRuleset) CanFinalizeAfterLog(ruleset.StreamInfo, []string) bool {
+	return true
+}
+
+type requestPropTCPAnalyzer struct {
+	closeCalls *int
+}
+
+func (a requestPropTCPAnalyzer) Name() string { return "tls" }
+func (a requestPropTCPAnalyzer) Limit() int   { return 0 }
+func (a requestPropTCPAnalyzer) NewTCP(analyzer.TCPInfo, analyzer.Logger) analyzer.TCPStream {
+	return requestPropTCPStream{closeCalls: a.closeCalls}
+}
+
+type requestPropTCPStream struct {
+	closeCalls *int
+}
+
+func (s requestPropTCPStream) Feed(bool, bool, bool, int, []byte) (*analyzer.PropUpdate, bool) {
+	return &analyzer.PropUpdate{
+		Type: analyzer.PropUpdateMerge,
+		M:    analyzer.PropMap{"req": analyzer.PropMap{"sni": "good.example"}},
+	}, false
+}
+
+func (s requestPropTCPStream) Close(bool) *analyzer.PropUpdate {
+	(*s.closeCalls)++
+	return nil
+}
+
 type noopTestLogger struct{}

 func (noopTestLogger) WorkerStart(int) {}
@@ -207,3 +290,90 @@ func TestTCPFlowReevaluatesAfterRulesetVersionChange(t *testing.T) {
 		t.Fatalf("cached verdict after update=%v want=%v", v, io.VerdictDropStream)
 	}
 }
+
+func TestTCPFlowDelaysAnalyzerCreationUntilPayload(t *testing.T) {
+	node, err := snowflake.NewNode(0)
+	if err != nil {
+		t.Fatalf("create node: %v", err)
+	}
+	newCalls := 0
+	feedCalls := 0
+	mgr := newTCPFlowManager(0, noopTestLogger{}, nil, node, newAnalyzerSelector(AnalyzerSelectionModeSignature, &statsCounters{}))
+	mgr.updateRuleset(analyzerRuleset{
+		action: ruleset.ActionMaybe,
+		ans: []analyzer.Analyzer{countingTCPAnalyzer{
+			newCalls:  &newCalls,
+			feedCalls: &feedCalls,
+		}},
+	}, 0)
+
+	l3 := L3Info{
+		Version:  4,
+		Protocol: 6,
+		SrcIP:    [4]byte{10, 0, 0, 1},
+		DstIP:    [4]byte{10, 0, 0, 2},
+	}
+	tcp := TCPInfo{
+		SrcPort: 12345,
+		DstPort: 443,
+		Seq:     100,
+	}
+
+	v := mgr.handle(1, l3, tcp, nil, nil, nil)
+	if v != io.VerdictAccept {
+		t.Fatalf("empty packet verdict=%v want=%v", v, io.VerdictAccept)
+	}
+	if newCalls != 0 || feedCalls != 0 {
+		t.Fatalf("empty packet created/feed analyzer: new=%d feed=%d", newCalls, feedCalls)
+	}
+
+	tcp.Seq = 101
+	v = mgr.handle(1, l3, tcp, []byte{0x16, 0x03, 0x01}, nil, nil)
+	if v != io.VerdictAccept {
+		t.Fatalf("payload verdict=%v want=%v", v, io.VerdictAccept)
+	}
+	if newCalls != 1 || feedCalls != 1 {
+		t.Fatalf("payload should create/feed analyzer once: new=%d feed=%d", newCalls, feedCalls)
+	}
+}
+
+func TestTCPFlowFinalizesAfterLogClassification(t *testing.T) {
+	node, err := snowflake.NewNode(0)
+	if err != nil {
+		t.Fatalf("create node: %v", err)
+	}
+	closeCalls := 0
+	mgr := newTCPFlowManager(0, noopTestLogger{}, nil, node, newAnalyzerSelector(AnalyzerSelectionModeSignature, &statsCounters{}))
+	mgr.updateRuleset(logFinalizingRuleset{
+		ans: []analyzer.Analyzer{requestPropTCPAnalyzer{closeCalls: &closeCalls}},
+	}, 0)
+
+	l3 := L3Info{
+		Version:  4,
+		Protocol: 6,
+		SrcIP:    [4]byte{10, 0, 0, 1},
+		DstIP:    [4]byte{10, 0, 0, 2},
+	}
+	tcp := TCPInfo{
+		SrcPort: 12345,
+		DstPort: 443,
+		Seq:     100,
+	}
+
+	v := mgr.handle(1, l3, tcp, nil, nil, nil)
+	if v != io.VerdictAccept {
+		t.Fatalf("empty packet verdict=%v want=%v", v, io.VerdictAccept)
+	}
+
+	tcp.Seq = 101
+	v = mgr.handle(1, l3, tcp, []byte{0x16, 0x03, 0x01}, nil, nil)
+	if v != io.VerdictAcceptStream {
+		t.Fatalf("payload verdict=%v want=%v", v, io.VerdictAcceptStream)
+	}
+	if closeCalls != 1 {
+		t.Fatalf("expected analyzer to be closed once after finalization, got %d", closeCalls)
+	}
+	if _, ok := mgr.flows[1]; ok {
+		t.Fatal("expected finalized TCP flow to be removed from manager")
+	}
+}