diff --git a/client/internal/netflow/store/event_aggregation_test.go b/client/internal/netflow/store/event_aggregation_test.go new file mode 100644 index 000000000..a0fe8ef8e --- /dev/null +++ b/client/internal/netflow/store/event_aggregation_test.go @@ -0,0 +1,190 @@ +package store + +import ( + "math/rand" + "net/netip" + "testing" + "time" + + "github.com/google/uuid" + "github.com/netbirdio/netbird/client/internal/netflow/types" + "github.com/stretchr/testify/assert" +) + +var random = rand.New(rand.NewSource(time.Now().UnixNano())) + +func TestFlowAggregation(t *testing.T) { + var protocols = []types.Protocol{types.ICMP, types.ICMPv6, types.TCP, types.UDP} + var tests = []struct { + description string + eventTypes []types.Type + }{ + { + description: "start and stop", + eventTypes: []types.Type{types.TypeStart, types.TypeEnd}, + }, + { + description: "start and drop", + eventTypes: []types.Type{types.TypeStart, types.TypeDrop}, + }, + { + description: "start only", + eventTypes: []types.Type{types.TypeStart}, + }, + { + description: "drop only", + eventTypes: []types.Type{types.TypeDrop}, + }} + + for _, protocol := range protocols { + for _, tt := range tests { + t.Run(tt.description+" "+protocol.String(), func(t *testing.T) { + store := NewAggregatingMemoryStore() + allExpected := make([]*types.Event, 0) + + for i := 0; i < 2; i++ { + inEvents, expected := generateEvents(tt.eventTypes, protocol, types.Ingress, 0) + for _, e := range inEvents { + store.StoreEvent(e) + } + allExpected = append(allExpected, expected) + } + + events := store.GetAggregatedEvents() + assert.ElementsMatch(t, events, allExpected) + }) + } + } +} + +func TestIcmpEventAggregation(t *testing.T) { + var protocols = []types.Protocol{types.ICMP, types.ICMPv6} + var icmpTypes = []uint8{1, 2, 3} + + var tests = []struct { + description string + eventTypes []types.Type + }{ + { + description: "start and stop", + eventTypes: []types.Type{types.TypeStart, types.TypeEnd}, + }, + { + description: "start and drop", + eventTypes: []types.Type{types.TypeStart, types.TypeDrop}, + }, + { + description: "start only", + eventTypes: []types.Type{types.TypeStart}, + }, + { + description: "drop only", + eventTypes: []types.Type{types.TypeDrop}, + }} + + for _, protocol := range protocols { + for _, tt := range tests { + t.Run(tt.description+" "+protocol.String(), func(t *testing.T) { + store := NewAggregatingMemoryStore() + allExpected := make([]*types.Event, 0) + for _, icmpType := range icmpTypes { + events, expected := generateEvents(tt.eventTypes, protocol, types.Ingress, icmpType) + for _, e := range events { + store.StoreEvent(e) + } + allExpected = append(allExpected, expected) + } + aggregatedEvents := store.GetAggregatedEvents() + assert.Len(t, aggregatedEvents, len(allExpected)) + assert.ElementsMatch(t, aggregatedEvents, allExpected) + }) + } + } +} + +func ipAddr(a string) netip.Addr { + addr, _ := netip.ParseAddr(a) + return addr +} + +func generateEvents(eventTypes []types.Type, protocol types.Protocol, direction types.Direction, icmpType uint8) ([]*types.Event, *types.Event) { + var rxPackets, txPackets, rxBytes, txBytes uint64 + inEvents := make([]*types.Event, 0) + ts := time.Now() + flowId := uuid.New() + srcIp := ipAddr("1.1.1.1") + srcPort := uint16(random.Uint32() >> 16) + dstIp := ipAddr("2.2.2.2") + dstPort := 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(), + }} + rxBytes += e.RxBytes + txBytes += e.TxBytes + rxPackets += e.RxPackets + txPackets += e.TxPackets + inEvents = append(inEvents, e) + if protocol == types.ICMP || protocol == types.ICMPv6 { + e.ICMPType = icmpType + } + } + + var start, end, drop uint64 + for _, eventType := range eventTypes { + switch eventType { + case types.TypeStart: + start += 1 + case types.TypeDrop: + drop += 1 + case types.TypeEnd: + end += 1 + } + } + aggregatedEvent := &types.Event{ + ID: inEvents[0].ID, + Timestamp: inEvents[0].Timestamp, + EventFields: types.EventFields{ + FlowID: flowId, + Type: inEvents[0].Type, + Protocol: inEvents[0].Protocol, + RuleID: []byte("rule-id-1"), + Direction: inEvents[0].Direction, + SourceIP: srcIp, + SourcePort: srcPort, + DestIP: dstIp, + DestPort: dstPort, + SourceResourceID: []byte("source-resource-id"), + DestResourceID: []byte("dest-resource-id"), + RxPackets: rxPackets, + TxPackets: txPackets, + RxBytes: rxBytes, + TxBytes: txBytes, + NumOfStarts: start, + NumOfEnds: end, + NumOfDrops: drop, + }} + if protocol == types.ICMP || protocol == types.ICMPv6 { + aggregatedEvent.ICMPType = icmpType + } + + return inEvents, aggregatedEvent +} diff --git a/client/internal/netflow/store/memory.go b/client/internal/netflow/store/memory.go index 321a9b81f..1a3b9227a 100644 --- a/client/internal/netflow/store/memory.go +++ b/client/internal/netflow/store/memory.go @@ -81,7 +81,7 @@ type aggregationKey struct { func (am *AggregatingMemory) GetAggregatedEvents() []*types.Event { aggregated := make(map[aggregationKey]*types.Event) for _, v := range am.events { - lookupKey := aggregationKey{destAddr: v.DestIP, destPort: v.DestPort, protocol: uint8(v.Protocol), icmpType: v.ICMPCode} + lookupKey := aggregationKey{destAddr: v.DestIP, destPort: v.DestPort, protocol: uint8(v.Protocol), icmpType: v.ICMPType} if _, ok := aggregated[lookupKey]; !ok { aggregated[lookupKey] = v.Clone() event := aggregated[lookupKey] diff --git a/client/internal/netflow/store/tcp_aggregation_test.go b/client/internal/netflow/store/tcp_aggregation_test.go deleted file mode 100644 index a4c9eab95..000000000 --- a/client/internal/netflow/store/tcp_aggregation_test.go +++ /dev/null @@ -1,281 +0,0 @@ -package store - -import ( - "net/netip" - "testing" - "time" - - "github.com/google/uuid" - "github.com/netbirdio/netbird/client/internal/netflow/types" - "github.com/stretchr/testify/assert" -) - -var pregeneratedUUIDs = func() []uuid.UUID { - toret := make([]uuid.UUID, 0) - for range make([]int, 10) { - toret = append(toret, uuid.New()) - } - return toret -}() - -func TestFlowAggregation(t *testing.T) { - var protocols = []types.Protocol{types.ICMP, types.ICMPv6, types.TCP, types.UDP} - var tests = []struct { - description string - events []*types.Event - expected []*types.Event - }{ - { - description: "start and stop", - events: []*types.Event{ - { - ID: pregeneratedUUIDs[0], - Timestamp: time.Unix(100, 100), - EventFields: types.EventFields{ - FlowID: pregeneratedUUIDs[1], - Type: types.TypeStart, - RuleID: []byte("rule-id-1"), - Direction: types.Egress, - SourceIP: ipAddr("1.1.1.1"), - SourcePort: 1234, - DestIP: ipAddr("2.2.2.2"), - DestPort: 443, - SourceResourceID: []byte("source-resource-id"), - DestResourceID: []byte("dest-resource-id"), - RxPackets: 10, - TxPackets: 20, - RxBytes: 10000, - TxBytes: 20000, - }}, - { - ID: pregeneratedUUIDs[2], - Timestamp: time.Unix(100, 100).Add(time.Second), - EventFields: types.EventFields{ - FlowID: pregeneratedUUIDs[1], - Type: types.TypeEnd, - RuleID: []byte("rule-id-1"), - Direction: types.Egress, - SourceIP: ipAddr("1.1.1.1"), - SourcePort: 1234, - DestIP: ipAddr("2.2.2.2"), - DestPort: 443, - SourceResourceID: []byte("source-resource-id"), - DestResourceID: []byte("dest-resource-id"), - RxPackets: 30, - TxPackets: 40, - RxBytes: 30000, - TxBytes: 40000, - }}, - }, - expected: []*types.Event{ - { - ID: pregeneratedUUIDs[0], - Timestamp: time.Unix(100, 100), - EventFields: types.EventFields{ - FlowID: pregeneratedUUIDs[1], - Type: types.TypeStart, - RuleID: []byte("rule-id-1"), - Direction: types.Egress, - SourceIP: ipAddr("1.1.1.1"), - SourcePort: 1234, - DestIP: ipAddr("2.2.2.2"), - DestPort: 443, - SourceResourceID: []byte("source-resource-id"), - DestResourceID: []byte("dest-resource-id"), - RxPackets: 40, - TxPackets: 60, - RxBytes: 40000, - TxBytes: 60000, - NumOfStarts: 1, - NumOfEnds: 1, - NumOfDrops: 0, - }}, - }, - }, - { - description: "start and drop", - events: []*types.Event{ - { - ID: pregeneratedUUIDs[0], - Timestamp: time.Unix(100, 100), - EventFields: types.EventFields{ - FlowID: pregeneratedUUIDs[1], - Type: types.TypeStart, - RuleID: []byte("rule-id-1"), - Direction: types.Egress, - SourceIP: ipAddr("1.1.1.1"), - SourcePort: 1234, - DestIP: ipAddr("2.2.2.2"), - DestPort: 443, - SourceResourceID: []byte("source-resource-id"), - DestResourceID: []byte("dest-resource-id"), - RxPackets: 10, - TxPackets: 20, - RxBytes: 10000, - TxBytes: 20000, - }}, - { - ID: pregeneratedUUIDs[2], - Timestamp: time.Unix(100, 100).Add(time.Second), - EventFields: types.EventFields{ - FlowID: pregeneratedUUIDs[1], - Type: types.TypeDrop, - RuleID: []byte("rule-id-1"), - Direction: types.Egress, - SourceIP: ipAddr("1.1.1.1"), - SourcePort: 1234, - DestIP: ipAddr("2.2.2.2"), - DestPort: 443, - SourceResourceID: []byte("source-resource-id"), - DestResourceID: []byte("dest-resource-id"), - RxPackets: 30, - TxPackets: 40, - RxBytes: 30000, - TxBytes: 40000, - }}, - }, - expected: []*types.Event{ - { - ID: pregeneratedUUIDs[0], - Timestamp: time.Unix(100, 100), - EventFields: types.EventFields{ - FlowID: pregeneratedUUIDs[1], - Type: types.TypeStart, - RuleID: []byte("rule-id-1"), - Direction: types.Egress, - SourceIP: ipAddr("1.1.1.1"), - SourcePort: 1234, - DestIP: ipAddr("2.2.2.2"), - DestPort: 443, - SourceResourceID: []byte("source-resource-id"), - DestResourceID: []byte("dest-resource-id"), - RxPackets: 40, - TxPackets: 60, - RxBytes: 40000, - TxBytes: 60000, - NumOfStarts: 1, - NumOfEnds: 0, - NumOfDrops: 1, - }}, - }, - }, - { - description: "start only", - events: []*types.Event{ - { - ID: pregeneratedUUIDs[0], - Timestamp: time.Unix(100, 100), - EventFields: types.EventFields{ - FlowID: pregeneratedUUIDs[1], - Type: types.TypeStart, - RuleID: []byte("rule-id-1"), - Direction: types.Egress, - SourceIP: ipAddr("1.1.1.1"), - SourcePort: 1234, - DestIP: ipAddr("2.2.2.2"), - DestPort: 443, - SourceResourceID: []byte("source-resource-id"), - DestResourceID: []byte("dest-resource-id"), - RxPackets: 10, - TxPackets: 20, - RxBytes: 10000, - TxBytes: 20000, - }}, - }, - expected: []*types.Event{ - { - ID: pregeneratedUUIDs[0], - Timestamp: time.Unix(100, 100), - EventFields: types.EventFields{ - FlowID: pregeneratedUUIDs[1], - Type: types.TypeStart, - RuleID: []byte("rule-id-1"), - Direction: types.Egress, - SourceIP: ipAddr("1.1.1.1"), - SourcePort: 1234, - DestIP: ipAddr("2.2.2.2"), - DestPort: 443, - SourceResourceID: []byte("source-resource-id"), - DestResourceID: []byte("dest-resource-id"), - RxPackets: 10, - TxPackets: 20, - RxBytes: 10000, - TxBytes: 20000, - NumOfStarts: 1, - NumOfEnds: 0, - NumOfDrops: 0, - }}, - }, - }, - { - description: "drop only", - events: []*types.Event{ - { - ID: pregeneratedUUIDs[2], - Timestamp: time.Unix(100, 100).Add(time.Second), - EventFields: types.EventFields{ - FlowID: pregeneratedUUIDs[1], - Type: types.TypeDrop, - RuleID: []byte("rule-id-1"), - Direction: types.Egress, - SourceIP: ipAddr("1.1.1.1"), - SourcePort: 1234, - DestIP: ipAddr("2.2.2.2"), - DestPort: 443, - SourceResourceID: []byte("source-resource-id"), - DestResourceID: []byte("dest-resource-id"), - RxPackets: 30, - TxPackets: 40, - RxBytes: 30000, - TxBytes: 40000, - }}, - }, - expected: []*types.Event{ - { - ID: pregeneratedUUIDs[2], - Timestamp: time.Unix(100, 100).Add(time.Second), - EventFields: types.EventFields{ - FlowID: pregeneratedUUIDs[1], - Type: types.TypeDrop, - RuleID: []byte("rule-id-1"), - Direction: types.Egress, - SourceIP: ipAddr("1.1.1.1"), - SourcePort: 1234, - DestIP: ipAddr("2.2.2.2"), - DestPort: 443, - SourceResourceID: []byte("source-resource-id"), - DestResourceID: []byte("dest-resource-id"), - RxPackets: 30, - TxPackets: 40, - RxBytes: 30000, - TxBytes: 40000, - NumOfStarts: 0, - NumOfEnds: 0, - NumOfDrops: 1, - }}, - }, - }} - - for _, protocol := range protocols { - for _, tt := range tests { - t.Run(tt.description+" "+protocol.String(), func(t *testing.T) { - store := NewAggregatingMemoryStore() - for _, e := range tt.events { - e.Protocol = protocol - store.StoreEvent(e) - } - for _, e := range tt.expected { - e.Protocol = protocol - } - events := store.GetAggregatedEvents() - assert.Len(t, events, len(tt.expected)) - assert.ElementsMatch(t, events, tt.expected) - }) - } - } -} - -func ipAddr(a string) netip.Addr { - addr, _ := netip.ParseAddr(a) - return addr -}