From 41a15f622150eb75d157904b30199bf7951ffe95 Mon Sep 17 00:00:00 2001 From: Dmitri Dolguikh Date: Mon, 22 Jun 2026 13:16:09 +0200 Subject: [PATCH] cleanup handling of not-aggregated events + test Signed-off-by: Dmitri Dolguikh --- .../netflow/store/event_aggregation_test.go | 125 ++++++++++++++++++ client/internal/netflow/store/memory.go | 13 +- 2 files changed, 131 insertions(+), 7 deletions(-) diff --git a/client/internal/netflow/store/event_aggregation_test.go b/client/internal/netflow/store/event_aggregation_test.go index 3bd1b21ea..dbadef485 100644 --- a/client/internal/netflow/store/event_aggregation_test.go +++ b/client/internal/netflow/store/event_aggregation_test.go @@ -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 +} diff --git a/client/internal/netflow/store/memory.go b/client/internal/netflow/store/memory.go index 02fe88903..87a2f250b 100644 --- a/client/internal/netflow/store/memory.go +++ b/client/internal/netflow/store/memory.go @@ -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 }