Improves flow handling and adds runtime stats APIs
Refactors TCP and UDP flow managers to enhance analyzer selection and flow binding accuracy, including O(1) UDP stream rebinding by 5-tuple. Introduces runtime stats tracking for engine and ruleset operations, exposing new APIs for granular performance and error metrics. Optimizes GeoMatcher with result caching and supports efficient geosite set matching, reducing redundant computation in ruleset expressions.
This commit is contained in:
+109
-21
@@ -1,6 +1,7 @@
|
||||
package engine
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"net"
|
||||
"sync"
|
||||
@@ -40,6 +41,8 @@ type udpStreamFactory struct {
|
||||
WorkerID int
|
||||
Logger Logger
|
||||
Node *snowflake.Node
|
||||
Selector *analyzerSelector
|
||||
Stats *statsCounters
|
||||
|
||||
RulesetMutex sync.RWMutex
|
||||
Ruleset ruleset.Ruleset
|
||||
@@ -64,7 +67,11 @@ func (f *udpStreamFactory) New(ipFlow, udpFlow gopacket.Flow, udp *layers.UDP, u
|
||||
rs, version := f.currentRuleset()
|
||||
var ans []analyzer.UDPAnalyzer
|
||||
if rs != nil {
|
||||
ans = analyzersToUDPAnalyzers(rs.Analyzers(info))
|
||||
baseAns := rs.Analyzers(info)
|
||||
if f.Selector != nil {
|
||||
baseAns = f.Selector.SelectUDP(baseAns, udp.Payload)
|
||||
}
|
||||
ans = analyzersToUDPAnalyzers(baseAns)
|
||||
}
|
||||
// Create entries for each analyzer
|
||||
entries := make([]*udpStreamEntry, 0, len(ans))
|
||||
@@ -110,8 +117,11 @@ func (f *udpStreamFactory) currentRuleset() (ruleset.Ruleset, uint64) {
|
||||
}
|
||||
|
||||
type udpStreamManager struct {
|
||||
factory *udpStreamFactory
|
||||
streams *lru.Cache[uint32, *udpStreamValue]
|
||||
factory *udpStreamFactory
|
||||
streams *lru.Cache[uint32, *udpStreamValue]
|
||||
tupleIndex map[udpTupleKey]uint32
|
||||
streamTuples map[uint32]udpTupleKey
|
||||
stats *statsCounters
|
||||
}
|
||||
|
||||
type udpStreamValue struct {
|
||||
@@ -120,36 +130,71 @@ type udpStreamValue struct {
|
||||
UDPFlow gopacket.Flow
|
||||
}
|
||||
|
||||
type udpTupleKey struct {
|
||||
AIP [16]byte
|
||||
BIP [16]byte
|
||||
ALen uint8
|
||||
BLen uint8
|
||||
APort uint16
|
||||
BPort uint16
|
||||
}
|
||||
|
||||
func (v *udpStreamValue) Match(ipFlow, udpFlow gopacket.Flow) (ok, rev bool) {
|
||||
fwd := v.IPFlow == ipFlow && v.UDPFlow == udpFlow
|
||||
rev = v.IPFlow == ipFlow.Reverse() && v.UDPFlow == udpFlow.Reverse()
|
||||
return fwd || rev, rev
|
||||
}
|
||||
|
||||
func newUDPStreamManager(factory *udpStreamFactory, maxStreams int) (*udpStreamManager, error) {
|
||||
ss, err := lru.New[uint32, *udpStreamValue](maxStreams)
|
||||
func newUDPStreamManager(factory *udpStreamFactory, maxStreams int, stats *statsCounters) (*udpStreamManager, error) {
|
||||
m := &udpStreamManager{
|
||||
factory: factory,
|
||||
tupleIndex: make(map[udpTupleKey]uint32, maxStreams),
|
||||
streamTuples: make(map[uint32]udpTupleKey, maxStreams),
|
||||
stats: stats,
|
||||
}
|
||||
ss, err := lru.NewWithEvict[uint32, *udpStreamValue](maxStreams, func(k uint32, v *udpStreamValue) {
|
||||
m.removeTupleMappingLocked(k)
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &udpStreamManager{
|
||||
factory: factory,
|
||||
streams: ss,
|
||||
}, nil
|
||||
m.streams = ss
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func (m *udpStreamManager) MatchWithContext(streamID uint32, ipFlow gopacket.Flow, udp *layers.UDP, uc *udpContext) {
|
||||
rev := false
|
||||
value, ok := m.streams.Get(streamID)
|
||||
tuple := canonicalUDPTupleKey(ipFlow, udp)
|
||||
if !ok {
|
||||
// Fallback: conntrack IDs can change during early flow lifetime on some systems.
|
||||
// Try to find an existing stream by 5-tuple before creating a new stream.
|
||||
matchedKey, matchedValue, matchedRev, found := m.findByFlow(ipFlow, udp.TransportFlow())
|
||||
if m.stats != nil {
|
||||
m.stats.UDPTupleLookups.Add(1)
|
||||
}
|
||||
// Conntrack IDs can change during early flow lifetime on some systems.
|
||||
// Rebind by canonical 5-tuple in O(1).
|
||||
matchedKey, found := m.tupleIndex[tuple]
|
||||
var matchedValue *udpStreamValue
|
||||
var matchedRev bool
|
||||
if found {
|
||||
if m.stats != nil {
|
||||
m.stats.UDPTupleHits.Add(1)
|
||||
}
|
||||
var hasValue bool
|
||||
matchedValue, hasValue = m.streams.Get(matchedKey)
|
||||
if !hasValue || matchedValue == nil {
|
||||
delete(m.tupleIndex, tuple)
|
||||
delete(m.streamTuples, matchedKey)
|
||||
found = false
|
||||
}
|
||||
}
|
||||
if found {
|
||||
_, matchedRev = matchedValue.Match(ipFlow, udp.TransportFlow())
|
||||
value = matchedValue
|
||||
rev = matchedRev
|
||||
if matchedKey != streamID {
|
||||
m.streams.Remove(matchedKey)
|
||||
m.streams.Add(streamID, matchedValue)
|
||||
m.bindTupleLocked(streamID, tuple)
|
||||
}
|
||||
} else {
|
||||
// New stream
|
||||
@@ -159,6 +204,7 @@ func (m *udpStreamManager) MatchWithContext(streamID uint32, ipFlow gopacket.Flo
|
||||
UDPFlow: udp.TransportFlow(),
|
||||
}
|
||||
m.streams.Add(streamID, value)
|
||||
m.bindTupleLocked(streamID, tuple)
|
||||
}
|
||||
} else {
|
||||
// Stream ID exists, but is it really the same stream?
|
||||
@@ -172,6 +218,7 @@ func (m *udpStreamManager) MatchWithContext(streamID uint32, ipFlow gopacket.Flo
|
||||
UDPFlow: udp.TransportFlow(),
|
||||
}
|
||||
m.streams.Add(streamID, value)
|
||||
m.bindTupleLocked(streamID, tuple)
|
||||
}
|
||||
}
|
||||
if value.Stream.Accept(udp, rev, uc) {
|
||||
@@ -179,17 +226,58 @@ func (m *udpStreamManager) MatchWithContext(streamID uint32, ipFlow gopacket.Flo
|
||||
}
|
||||
}
|
||||
|
||||
func (m *udpStreamManager) findByFlow(ipFlow, udpFlow gopacket.Flow) (key uint32, value *udpStreamValue, rev bool, found bool) {
|
||||
for _, k := range m.streams.Keys() {
|
||||
v, ok := m.streams.Peek(k)
|
||||
if !ok || v == nil {
|
||||
continue
|
||||
}
|
||||
if ok2, rev2 := v.Match(ipFlow, udpFlow); ok2 {
|
||||
return k, v, rev2, true
|
||||
func (m *udpStreamManager) bindTupleLocked(streamID uint32, key udpTupleKey) {
|
||||
m.removeTupleMappingLocked(streamID)
|
||||
m.tupleIndex[key] = streamID
|
||||
m.streamTuples[streamID] = key
|
||||
}
|
||||
|
||||
func (m *udpStreamManager) removeTupleMappingLocked(streamID uint32) {
|
||||
if key, ok := m.streamTuples[streamID]; ok {
|
||||
delete(m.streamTuples, streamID)
|
||||
current, exists := m.tupleIndex[key]
|
||||
if exists && current == streamID {
|
||||
delete(m.tupleIndex, key)
|
||||
}
|
||||
}
|
||||
return 0, nil, false, false
|
||||
}
|
||||
|
||||
func canonicalUDPTupleKey(ipFlow gopacket.Flow, udp *layers.UDP) udpTupleKey {
|
||||
srcIP := ipFlow.Src().Raw()
|
||||
dstIP := ipFlow.Dst().Raw()
|
||||
srcPort := uint16(udp.SrcPort)
|
||||
dstPort := uint16(udp.DstPort)
|
||||
|
||||
if compareIPEndpoint(srcIP, srcPort, dstIP, dstPort) > 0 {
|
||||
srcIP, dstIP = dstIP, srcIP
|
||||
srcPort, dstPort = dstPort, srcPort
|
||||
}
|
||||
|
||||
var key udpTupleKey
|
||||
key.ALen = uint8(copy(key.AIP[:], srcIP))
|
||||
key.BLen = uint8(copy(key.BIP[:], dstIP))
|
||||
key.APort = srcPort
|
||||
key.BPort = dstPort
|
||||
return key
|
||||
}
|
||||
|
||||
func compareIPEndpoint(aIP []byte, aPort uint16, bIP []byte, bPort uint16) int {
|
||||
if len(aIP) != len(bIP) {
|
||||
if len(aIP) < len(bIP) {
|
||||
return -1
|
||||
}
|
||||
return 1
|
||||
}
|
||||
if c := bytes.Compare(aIP, bIP); c != 0 {
|
||||
return c
|
||||
}
|
||||
if aPort < bPort {
|
||||
return -1
|
||||
}
|
||||
if aPort > bPort {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
type udpStream struct {
|
||||
|
||||
Reference in New Issue
Block a user