diff --git a/client/internal/netflow/manager.go b/client/internal/netflow/manager.go index f6397952d..0e65c880e 100644 --- a/client/internal/netflow/manager.go +++ b/client/internal/netflow/manager.go @@ -218,7 +218,6 @@ func (m *Manager) startSender(ctx context.Context) { collectedEvents := m.logger.ResetAggregationWindow() events := collectedEvents.GetAggregatedEvents() for _, event := range events { - // handle retries, grace period? if err := m.send(event); err != nil { log.Errorf("failed to send flow event to server: %v", err) } else { @@ -265,6 +264,7 @@ func (m *Manager) startRetries(ctx context.Context) { case <-ctx.Done(): return case <-ticker.C: + // TODO: grace period on retries to avoid early retries? for _, e := range m.eventsWithoutAcks.GetEvents() { if err := m.send(e); err != nil { ticker = time.NewTimer(retryBackoff.NextBackOff()) diff --git a/client/internal/netflow/store/memory.go b/client/internal/netflow/store/memory.go index 92e4578fc..10da74d8b 100644 --- a/client/internal/netflow/store/memory.go +++ b/client/internal/netflow/store/memory.go @@ -91,8 +91,10 @@ func (am *AggregatingMemory) GetAggregatedEvents() []*types.Event { aggregatedEvent.RxPackets += v.RxPackets aggregatedEvent.TxBytes += v.TxBytes aggregatedEvent.TxPackets += v.TxPackets - if aggregatedEvent.Timestamp.Compare(v.Timestamp) < 0 { + if aggregatedEvent.Timestamp.Compare(v.Timestamp) > 0 { aggregatedEvent.Timestamp = v.Timestamp + aggregatedEvent.ID = v.ID + aggregatedEvent.Type = v.Type } // do we aggregate icmp by code? default: diff --git a/client/internal/netflow/store/tcp_aggregation_test.go b/client/internal/netflow/store/tcp_aggregation_test.go new file mode 100644 index 000000000..8b023ce69 --- /dev/null +++ b/client/internal/netflow/store/tcp_aggregation_test.go @@ -0,0 +1,272 @@ +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 TestTcpAggregation(t *testing.T) { + 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, + Protocol: types.TCP, + 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, + Protocol: types.TCP, + 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, + Protocol: types.TCP, + 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, + }}, + }, + }, + { + 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, + Protocol: types.TCP, + 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, + Protocol: types.TCP, + 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, + Protocol: types.TCP, + 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, + }}, + }, + }, + { + 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, + Protocol: types.TCP, + 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, + Protocol: types.TCP, + 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, + }}, + }, + }, + { + 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.TypeEnd, + RuleID: []byte("rule-id-1"), + Direction: types.Egress, + Protocol: types.TCP, + 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.TypeEnd, + RuleID: []byte("rule-id-1"), + Direction: types.Egress, + Protocol: types.TCP, + 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, + }}, + }, + }} + + for _, tt := range tests { + t.Run(tt.description, func(t *testing.T) { + store := NewAggregatingMemoryStore() + for _, e := range tt.events { + store.StoreEvent(e) + } + 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 +}