Difuse/Mellaris
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: eliminate stale verdict poisoning, memory leaks, data races, and per-packet allocations in engine

Browse Source
This commit is contained in:
hayzam
2026-05-15 02:08:22 +00:00
parent bc25169f41
commit 301c252c43
15 changed files with 222 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/udp.go
+50 -54
View File
@@ -12,8 +12,6 @@ import (
"git.difuse.io/Difuse/Mellaris/ruleset"
"github.com/bwmarrin/snowflake"
"github.com/google/gopacket"
"github.com/google/gopacket/layers"
lru "github.com/hashicorp/golang-lru/v2"
)
@@ -49,9 +47,10 @@ type udpStreamFactory struct {
RulesetVersion uint64
}
func (f *udpStreamFactory) New(ipFlow, udpFlow gopacket.Flow, udp *layers.UDP, uc *udpContext) *udpStream {
func (f *udpStreamFactory) New(k udpTupleKey, payload []byte, uc *udpContext) *udpStream {
id := f.Node.Generate()
ipSrc, ipDst := net.IP(ipFlow.Src().Raw()), net.IP(ipFlow.Dst().Raw())
ipSrc := net.IP(k.AIP[:k.ALen])
ipDst := net.IP(k.BIP[:k.BLen])
info := ruleset.StreamInfo{
ID: id.Int64(),
Protocol: ruleset.ProtocolUDP,
@@ -59,8 +58,8 @@ func (f *udpStreamFactory) New(ipFlow, udpFlow gopacket.Flow, udp *layers.UDP, u
DstMAC: append(net.HardwareAddr(nil), uc.DstMAC...),
SrcIP: ipSrc,
DstIP: ipDst,
SrcPort: uint16(udp.SrcPort),
DstPort: uint16(udp.DstPort),
SrcPort: k.APort,
DstPort: k.BPort,
Props: make(analyzer.CombinedPropMap),
}
f.Logger.UDPStreamNew(f.WorkerID, info)
@@ -69,11 +68,10 @@ func (f *udpStreamFactory) New(ipFlow, udpFlow gopacket.Flow, udp *layers.UDP, u
if rs != nil {
baseAns := rs.Analyzers(info)
if f.Selector != nil {
baseAns = f.Selector.SelectUDP(baseAns, udp.Payload)
baseAns = f.Selector.SelectUDP(baseAns, payload)
}
ans = analyzersToUDPAnalyzers(baseAns)
}
// Create entries for each analyzer
entries := make([]*udpStreamEntry, 0, len(ans))
for _, a := range ans {
entries = append(entries, &udpStreamEntry{
@@ -81,8 +79,8 @@ func (f *udpStreamFactory) New(ipFlow, udpFlow gopacket.Flow, udp *layers.UDP, u
Stream: a.NewUDP(analyzer.UDPInfo{
SrcIP: ipSrc,
DstIP: ipDst,
SrcPort: uint16(udp.SrcPort),
DstPort: uint16(udp.DstPort),
SrcPort: k.APort,
DstPort: k.BPort,
}, &analyzerLogger{
StreamID: id.Int64(),
Name: a.Name(),
@@ -125,9 +123,14 @@ type udpStreamManager struct {
}
type udpStreamValue struct {
Stream *udpStream
IPFlow gopacket.Flow
UDPFlow gopacket.Flow
Stream *udpStream
Tuple udpTupleKey
}
func (v *udpStreamValue) Match(k udpTupleKey) (ok, rev bool) {
fwd := v.Tuple == k
rev = v.Tuple == reverseTuple(k)
return fwd || rev, rev
}
type udpTupleKey struct {
@@ -139,12 +142,6 @@ type udpTupleKey struct {
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, stats *statsCounters) (*udpStreamManager, error) {
m := &udpStreamManager{
factory: factory,
@@ -153,6 +150,9 @@ func newUDPStreamManager(factory *udpStreamFactory, maxStreams int, stats *stats
stats: stats,
}
ss, err := lru.NewWithEvict[uint32, *udpStreamValue](maxStreams, func(k uint32, v *udpStreamValue) {
if v != nil && v.Stream != nil {
v.Stream.Close()
}
m.removeTupleMappingLocked(k)
})
if err != nil {
@@ -162,16 +162,12 @@ func newUDPStreamManager(factory *udpStreamFactory, maxStreams int, stats *stats
return m, nil
}
func (m *udpStreamManager) MatchWithContext(streamID uint32, ipFlow gopacket.Flow, udp *layers.UDP, uc *udpContext) {
rev := false
func (m *udpStreamManager) MatchWithContext(streamID uint32, tuple udpTupleKey, rev bool, payload []byte, uc *udpContext) {
value, ok := m.streams.Get(streamID)
tuple := canonicalUDPTupleKey(ipFlow, udp)
if !ok {
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
@@ -188,7 +184,7 @@ func (m *udpStreamManager) MatchWithContext(streamID uint32, ipFlow gopacket.Flo
}
}
if found {
_, matchedRev = matchedValue.Match(ipFlow, udp.TransportFlow())
_, matchedRev = matchedValue.Match(tuple)
value = matchedValue
rev = matchedRev
if matchedKey != streamID {
@@ -197,32 +193,27 @@ func (m *udpStreamManager) MatchWithContext(streamID uint32, ipFlow gopacket.Flo
m.bindTupleLocked(streamID, tuple)
}
} else {
// New stream
value = &udpStreamValue{
Stream: m.factory.New(ipFlow, udp.TransportFlow(), udp, uc),
IPFlow: ipFlow,
UDPFlow: udp.TransportFlow(),
Stream: m.factory.New(tuple, payload, uc),
Tuple: tuple,
}
m.streams.Add(streamID, value)
m.bindTupleLocked(streamID, tuple)
}
} else {
// Stream ID exists, but is it really the same stream?
ok, rev = value.Match(ipFlow, udp.TransportFlow())
ok, rev = value.Match(tuple)
if !ok {
// It's not - close the old stream & replace it with a new one
value.Stream.Close()
value = &udpStreamValue{
Stream: m.factory.New(ipFlow, udp.TransportFlow(), udp, uc),
IPFlow: ipFlow,
UDPFlow: udp.TransportFlow(),
Stream: m.factory.New(tuple, payload, uc),
Tuple: tuple,
}
m.streams.Add(streamID, value)
m.bindTupleLocked(streamID, tuple)
}
}
if value.Stream.Accept(udp, rev, uc) {
value.Stream.Feed(udp, rev, uc)
if value.Stream.Accept(rev, uc) {
value.Stream.Feed(rev, payload, uc)
}
}
@@ -242,25 +233,34 @@ func (m *udpStreamManager) removeTupleMappingLocked(streamID uint32) {
}
}
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)
func canonicalUDPTupleKey(srcIP, dstIP net.IP, srcPort, dstPort uint16) udpTupleKey {
srcRaw := []byte(srcIP)
dstRaw := []byte(dstIP)
if compareIPEndpoint(srcIP, srcPort, dstIP, dstPort) > 0 {
srcIP, dstIP = dstIP, srcIP
if compareIPEndpoint(srcRaw, srcPort, dstRaw, dstPort) > 0 {
srcRaw, dstRaw = dstRaw, srcRaw
srcPort, dstPort = dstPort, srcPort
}
var key udpTupleKey
key.ALen = uint8(copy(key.AIP[:], srcIP))
key.BLen = uint8(copy(key.BIP[:], dstIP))
key.ALen = uint8(copy(key.AIP[:], srcRaw))
key.BLen = uint8(copy(key.BIP[:], dstRaw))
key.APort = srcPort
key.BPort = dstPort
return key
}
func reverseTuple(k udpTupleKey) udpTupleKey {
var r udpTupleKey
r.ALen = k.BLen
r.BLen = k.ALen
r.AIP = k.BIP
r.BIP = k.AIP
r.APort = k.BPort
r.BPort = k.APort
return r
}
func compareIPEndpoint(aIP []byte, aPort uint16, bIP []byte, bPort uint16) int {
if len(aIP) != len(bIP) {
if len(aIP) < len(bIP) {
@@ -298,11 +298,8 @@ type udpStreamEntry struct {
Quota int
}
func (s *udpStream) Accept(udp *layers.UDP, rev bool, uc *udpContext) bool {
func (s *udpStream) Accept(rev bool, uc *udpContext) bool {
if len(s.activeEntries) > 0 || s.virgin || s.rulesetChanged() {
// Make sure every stream matches against the ruleset at least once,
// even if there are no activeEntries, as the ruleset may have built-in
// properties that need to be matched.
return true
} else {
uc.Verdict = s.lastVerdict
@@ -310,12 +307,11 @@ func (s *udpStream) Accept(udp *layers.UDP, rev bool, uc *udpContext) bool {
}
}
func (s *udpStream) Feed(udp *layers.UDP, rev bool, uc *udpContext) {
func (s *udpStream) Feed(rev bool, payload []byte, uc *udpContext) {
updated := false
for i := len(s.activeEntries) - 1; i >= 0; i-- {
// Important: reverse order so we can remove entries
entry := s.activeEntries[i]
update, closeUpdate, done := s.feedEntry(entry, rev, udp.Payload)
update, closeUpdate, done := s.feedEntry(entry, rev, payload)
up1 := processPropUpdate(s.info.Props, entry.Name, update)
up2 := processPropUpdate(s.info.Props, entry.Name, closeUpdate)
updated = updated || up1 || up2
@@ -345,7 +341,7 @@ func (s *udpStream) Feed(udp *layers.UDP, rev bool, uc *udpContext) {
action = ruleset.ActionMaybe
} else {
var err error
uc.Packet, err = udpMI.Process(udp.Payload)
uc.Packet, err = udpMI.Process(payload)
if err != nil {
// Modifier error, fallback to maybe
s.logger.ModifyError(s.info, err)