engine: more performance improvements

engine/tcp_flow.go

@@ -27,6 +27,8 @@ type tcpFlow struct {
 	streamID uint32
 	srcPort  uint16
 	dstPort  uint16
+	srcIP    net.IP
+	dstIP    net.IP
 
 	dirSeq [2]uint32
 	dirBuf [2][]byte
@@ -41,6 +43,9 @@ type tcpFlow struct {
 	lastVerdict    io.Verdict
 	feedCalled     [2]bool
 	lastSeen       time.Time
+
+	pendingAnalyzers []analyzer.Analyzer
+	selector         *analyzerSelector
 }
 
 type tcpFlowEntry struct {
@@ -54,7 +59,7 @@ func (f *tcpFlow) feed(l3 L3Info, tcp TCPInfo, payload []byte) io.Verdict {
 	rs, version := f.currentRuleset()
 	rulesetChanged := version != f.rulesetVersion
 
-	if !f.virgin && !rulesetChanged && len(f.activeEntries) == 0 {
+	if !f.virgin && !rulesetChanged && len(f.activeEntries) == 0 && !f.hasPendingAnalyzers() {
 		return f.lastVerdict
 	}
 
@@ -68,6 +73,9 @@ func (f *tcpFlow) feed(l3 L3Info, tcp TCPInfo, payload []byte) io.Verdict {
 	propUpdated := false
 	if len(payload) > 0 {
 		dir, rev := f.resolveDirection(tcp)
+		if len(f.pendingAnalyzers) > 0 {
+			f.initPendingAnalyzers(payload)
+		}
 		expected := f.dirSeq[dir]
 		if !f.feedCalled[dir] || expected == 0 || tcp.Seq == expected {
 			f.feedCalled[dir] = true
@@ -108,6 +116,52 @@ func (f *tcpFlow) feedAnalyzers(rev bool) bool {
 	return updated
 }
 
+func (f *tcpFlow) initPendingAnalyzers(payload []byte) {
+	baseAns := f.pendingAnalyzers
+	f.pendingAnalyzers = nil
+	if f.selector != nil {
+		baseAns = f.selector.SelectTCP(baseAns, payload)
+	}
+	ans := analyzersToTCPAnalyzers(baseAns)
+	if len(ans) == 0 {
+		return
+	}
+	entries := make([]*tcpFlowEntry, 0, len(ans))
+	for _, a := range ans {
+		entries = append(entries, &tcpFlowEntry{
+			Name: a.Name(),
+			Stream: a.NewTCP(analyzer.TCPInfo{
+				SrcIP:   f.srcIP,
+				DstIP:   f.dstIP,
+				SrcPort: f.srcPort,
+				DstPort: f.dstPort,
+			}, &analyzerLogger{
+				StreamID: f.info.ID,
+				Name:     a.Name(),
+				Logger:   f.logger,
+			}),
+			HasLimit: a.Limit() > 0,
+			Quota:    a.Limit(),
+		})
+	}
+	f.activeEntries = append(f.activeEntries, entries...)
+}
+
+func (f *tcpFlow) hasPendingAnalyzers() bool {
+	return len(f.pendingAnalyzers) > 0
+}
+
+func (f *tcpFlow) analyzerNames() []string {
+	names := make([]string, 0, len(f.activeEntries)+len(f.pendingAnalyzers))
+	for _, entry := range f.activeEntries {
+		names = append(names, entry.Name)
+	}
+	for _, a := range f.pendingAnalyzers {
+		names = append(names, a.Name())
+	}
+	return names
+}
+
 func (f *tcpFlow) runMatch(rs ruleset.Ruleset, version uint64, rulesetChanged bool, propUpdated bool) {
 	if !propUpdated && !f.virgin && !rulesetChanged {
 		return
@@ -125,11 +179,15 @@ func (f *tcpFlow) runMatch(rs ruleset.Ruleset, version uint64, rulesetChanged bo
 		f.lastVerdict = verdict
 		f.closeActiveEntries()
 		f.logger.TCPStreamAction(f.info, action, false)
+	} else if result.Logged && canFinalizeAfterLog(rs, f.info, f.analyzerNames()) {
+		f.lastVerdict = io.VerdictAcceptStream
+		f.closeActiveEntries()
+		f.logger.TCPStreamAction(f.info, ruleset.ActionAllow, true)
 	}
 }
 
 func (f *tcpFlow) maybeFinalizeVerdict() {
-	if len(f.activeEntries) == 0 && f.lastVerdict == io.VerdictAccept {
+	if len(f.activeEntries) == 0 && !f.hasPendingAnalyzers() && f.lastVerdict == io.VerdictAccept {
 		f.lastVerdict = io.VerdictAcceptStream
 		f.logger.TCPStreamAction(f.info, ruleset.ActionAllow, true)
 	}
@@ -231,8 +289,10 @@ func (m *tcpFlowManager) createFlow(streamID uint32, l3 L3Info, tcp TCPInfo, pay
 	var ans []analyzer.TCPAnalyzer
 	if rs != nil {
 		baseAns := rs.Analyzers(info)
-		baseAns = m.selector.SelectTCP(baseAns, payload)
-		ans = analyzersToTCPAnalyzers(baseAns)
+		if len(payload) > 0 {
+			baseAns = m.selector.SelectTCP(baseAns, payload)
+			ans = analyzersToTCPAnalyzers(baseAns)
+		}
 	}
 	entries := make([]*tcpFlowEntry, 0, len(ans))
 	for _, a := range ans {
@@ -257,6 +317,8 @@ func (m *tcpFlowManager) createFlow(streamID uint32, l3 L3Info, tcp TCPInfo, pay
 		streamID:       streamID,
 		srcPort:        tcp.SrcPort,
 		dstPort:        tcp.DstPort,
+		srcIP:          ipSrc,
+		dstIP:          ipDst,
 		info:           info,
 		virgin:         true,
 		logger:         m.logger,
@@ -265,6 +327,10 @@ func (m *tcpFlowManager) createFlow(streamID uint32, l3 L3Info, tcp TCPInfo, pay
 		activeEntries:  entries,
 		lastVerdict:    io.VerdictAccept,
 		lastSeen:       time.Now(),
+		selector:       m.selector,
+	}
+	if len(payload) == 0 && rs != nil {
+		flow.pendingAnalyzers = rs.Analyzers(info)
 	}
 	flow.dirSeq[tcpDirC2S] = tcp.Seq + 1
 	return flow
@@ -325,3 +391,8 @@ func actionToTCPVerdict(a ruleset.Action) io.Verdict {
 		return io.VerdictAcceptStream
 	}
 }
+
+func canFinalizeAfterLog(rs ruleset.Ruleset, info ruleset.StreamInfo, activeAnalyzers []string) bool {
+	finalizer, ok := rs.(ruleset.LogFinalizer)
+	return ok && finalizer.CanFinalizeAfterLog(info, activeAnalyzers)
+}