cleanup handling of not-aggregated events + test

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>
This commit is contained in:
Dmitri Dolguikh
2026-06-22 13:16:09 +02:00
parent 4a1341883b
commit 41a15f6221
2 changed files with 131 additions and 7 deletions

View File

@@ -121,6 +121,59 @@ func TestIcmpEventAggregation(t *testing.T) {
}
}
func TestFlowAggregationOfUnknownProtocols(t *testing.T) {
var tests = []struct {
description string
addresses [][]netip.Addr
dstPort uint16
eventTypes []types.Type
}{
{
description: "start and stop",
addresses: [][]netip.Addr{{ipAddr("1.1.1.1"), ipAddr("2.2.2.2")}, {ipAddr("3.3.3.3"), ipAddr("2.2.2.2")}},
dstPort: uint16(random.Uint32() >> 16),
eventTypes: []types.Type{types.TypeStart, types.TypeEnd},
},
{
description: "start and drop",
addresses: [][]netip.Addr{{ipAddr("1.1.1.1"), ipAddr("2.2.2.2")}, {ipAddr("3.3.3.3"), ipAddr("2.2.2.2")}},
dstPort: uint16(random.Uint32() >> 16),
eventTypes: []types.Type{types.TypeStart, types.TypeDrop},
},
{
description: "start only",
addresses: [][]netip.Addr{{ipAddr("1.1.1.1"), ipAddr("2.2.2.2")}, {ipAddr("3.3.3.3"), ipAddr("2.2.2.2")}},
dstPort: uint16(random.Uint32() >> 16),
eventTypes: []types.Type{types.TypeStart},
},
{
description: "drop only",
addresses: [][]netip.Addr{{ipAddr("1.1.1.1"), ipAddr("2.2.2.2")}, {ipAddr("3.3.3.3"), ipAddr("2.2.2.2")}},
dstPort: uint16(random.Uint32() >> 16),
eventTypes: []types.Type{types.TypeDrop},
}}
for _, tt := range tests {
t.Run(tt.description+" "+string(types.ProtocolUnknown), func(t *testing.T) {
store := NewAggregatingMemoryStore()
store.WindowEnd = time.Now().Add(5 * time.Second)
allExpected := make([]*types.Event, 0)
for _, srcAndDst := range tt.addresses {
inEvents, expected := generateEventsForUnknownProtocol(srcAndDst[0], srcAndDst[1], tt.dstPort, tt.eventTypes, types.ProtocolUnknown, types.Ingress, store.WindowStart, store.WindowEnd)
for _, e := range inEvents {
store.StoreEvent(e)
}
allExpected = append(allExpected, expected...)
}
events := store.GetAggregatedEvents()
assert.ElementsMatch(t, events, allExpected)
})
}
}
func ipAddr(a string) netip.Addr {
addr, _ := netip.ParseAddr(a)
return addr
@@ -207,3 +260,75 @@ func generateEvents(srcIp, dstIp netip.Addr, dstPort uint16, eventTypes []types.
return inEvents, aggregatedEvent
}
func generateEventsForUnknownProtocol(srcIp, dstIp netip.Addr, dstPort uint16, eventTypes []types.Type, protocol types.Protocol,
direction types.Direction, windowStart, windowEnd time.Time) ([]*types.Event, []*types.Event) {
inEvents := make([]*types.Event, 0)
expectedEvents := make([]*types.Event, 0)
ts := time.Now()
flowId := uuid.New()
srcPort := uint16(random.Uint32() >> 16)
for idx, eventType := range eventTypes {
e := &types.Event{
ID: uuid.New(),
Timestamp: ts.Add(time.Duration(idx) * time.Second),
EventFields: types.EventFields{
FlowID: flowId,
Type: eventType,
Protocol: protocol,
RuleID: []byte("rule-id-1"),
Direction: direction,
SourceIP: srcIp,
SourcePort: srcPort,
DestIP: dstIp,
DestPort: dstPort,
SourceResourceID: []byte("source-resource-id"),
DestResourceID: []byte("dest-resource-id"),
RxPackets: random.Uint64(),
TxPackets: random.Uint64(),
RxBytes: random.Uint64(),
TxBytes: random.Uint64(),
}}
inEvents = append(inEvents, e)
var start, end, drop uint64
switch eventType {
case types.TypeStart:
start = 1
case types.TypeDrop:
drop = 1
case types.TypeEnd:
end = 1
}
expectedEvents = append(expectedEvents, &types.Event{
ID: e.ID,
Timestamp: e.Timestamp,
WindowStart: windowStart,
WindowEnd: windowEnd,
EventFields: types.EventFields{
FlowID: flowId,
Type: types.TypeUnknown,
Protocol: e.Protocol,
RuleID: []byte("rule-id-1"),
Direction: e.Direction,
SourceIP: srcIp,
SourcePort: srcPort,
DestIP: dstIp,
DestPort: dstPort,
SourceResourceID: []byte("source-resource-id"),
DestResourceID: []byte("dest-resource-id"),
RxPackets: e.RxPackets,
TxPackets: e.TxPackets,
RxBytes: e.RxBytes,
TxBytes: e.TxBytes,
NumOfStarts: start,
NumOfEnds: end,
NumOfDrops: drop,
}})
}
return inEvents, expectedEvents
}

View File

@@ -94,13 +94,6 @@ func (am *AggregatingMemory) GetAggregatedEvents() []*types.Event {
if _, ok := aggregated[lookupKey]; !ok {
event := v.Clone()
if event.Protocol != types.ICMP && event.Protocol != types.ICMPv6 && event.Protocol != types.UDP && event.Protocol != types.TCP {
lookupKey.unique = time.Now().UnixNano() // to make the lookup key unique so we don't aggregate on it
aggregated[lookupKey] = event
continue
}
aggregated[lookupKey] = event
switch event.Type {
case types.TypeStart:
event.NumOfStarts += 1
@@ -116,6 +109,12 @@ func (am *AggregatingMemory) GetAggregatedEvents() []*types.Event {
event.WindowStart = am.WindowStart
event.WindowEnd = am.WindowEnd
if event.Protocol != types.ICMP && event.Protocol != types.ICMPv6 && event.Protocol != types.UDP && event.Protocol != types.TCP {
lookupKey.unique = time.Now().UnixNano() // to make the lookup key unique so we don't aggregate on it
}
aggregated[lookupKey] = event
continue
}