[client] introduce client-side event aggregation (#6627)

* added an implementation of aggregating memory store

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* initial support for aggregation of events

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* added tcp-aggregation test

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* added manager integration test

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* added tracking of the number of start-, drop, and end-events in an aggregation window

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* fixes based on sonarcube checks

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* regenerated proto files

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* removed inadvertenly added google proto files

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* pacifying linter

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* update test to validate event aggregation over tcp, udp, icmp, and icmpv6

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* updated event aggregation test

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* regenerate protobufs with expected versions of protoc and protoc-gen-go

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* remove protoc/protoc-gen headers from flow_grpc.pb.go

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* updated openapi spec

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* updated openapi NetworkTrafficEvent spec, regenerated types

Signed-off-by: Dmitri <dmitri.external@netbird.io>

* respond to feedback

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* fixed an issue with how we track events that shouldn't be aggregated

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* fixed mapping of events to protobuf

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* icmp code values in aggregated events do not matter

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* regenerate openapi types

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* added a comment re: unbounded unacked events

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* reset aggregated event type to unknown

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* fix event aggregation test

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* used the source port of the earliest event

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* add tracking of window starts and ends

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* updated openapi spec

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* reverted changes to generate.sh

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* cleanup handling of not-aggregated events + test

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* responded to feedback + small fixes

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* small fix in a test

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* another test

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* force setting non-empty rule id on aggregated events

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* fixed a couple of issues flagged by coderabbit

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* fix spelling

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* handle exhausted retry backoffs

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

---------

Signed-off-by: Dmitri <dmitri.external@netbird.io>
Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>
This commit is contained in:
dmitri-netbird
2026-07-06 16:04:26 +02:00
committed by GitHub
parent 91acb8147c
commit 47352e6e45
12 changed files with 1081 additions and 149 deletions

View File

@@ -27,7 +27,7 @@ type Logger struct {
wgIfaceNetV6 netip.Prefix
dnsCollection atomic.Bool
exitNodeCollection atomic.Bool
Store types.Store
Store types.AggregatingStore
}
func New(statusRecorder *peer.Status, wgIfaceIPNet, wgIfaceIPNetV6 netip.Prefix) *Logger {
@@ -35,7 +35,7 @@ func New(statusRecorder *peer.Status, wgIfaceIPNet, wgIfaceIPNetV6 netip.Prefix)
statusRecorder: statusRecorder,
wgIfaceNet: wgIfaceIPNet,
wgIfaceNetV6: wgIfaceIPNetV6,
Store: store.NewMemoryStore(),
Store: store.NewAggregatingMemoryStore(),
}
}
@@ -125,6 +125,10 @@ func (l *Logger) stop() {
l.mux.Unlock()
}
func (l *Logger) ResetAggregationWindow() types.FlowEventAggregator {
return l.Store.ResetAggregationWindow()
}
func (l *Logger) GetEvents() []*types.Event {
return l.Store.GetEvents()
}

View File

@@ -9,12 +9,14 @@ import (
"sync"
"time"
"github.com/cenkalti/backoff/v4"
"github.com/google/uuid"
log "github.com/sirupsen/logrus"
"google.golang.org/protobuf/types/known/timestamppb"
"github.com/netbirdio/netbird/client/internal/netflow/conntrack"
"github.com/netbirdio/netbird/client/internal/netflow/logger"
"github.com/netbirdio/netbird/client/internal/netflow/store"
nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
"github.com/netbirdio/netbird/client/internal/peer"
"github.com/netbirdio/netbird/flow/client"
@@ -23,14 +25,16 @@ import (
// Manager handles netflow tracking and logging
type Manager struct {
mux sync.Mutex
shutdownWg sync.WaitGroup
logger nftypes.FlowLogger
flowConfig *nftypes.FlowConfig
conntrack nftypes.ConnTracker
receiverClient *client.GRPCClient
publicKey []byte
cancel context.CancelFunc
mux sync.Mutex
shutdownWg sync.WaitGroup
logger nftypes.FlowLogger
flowConfig *nftypes.FlowConfig
conntrack nftypes.ConnTracker
receiverClient *client.GRPCClient
eventsWithoutAcks nftypes.Store
publicKey []byte
cancel context.CancelFunc
retryInterval time.Duration
}
// NewManager creates a new netflow manager
@@ -48,9 +52,11 @@ func NewManager(iface nftypes.IFaceMapper, publicKey []byte, statusRecorder *pee
}
return &Manager{
logger: flowLogger,
conntrack: ct,
publicKey: publicKey,
logger: flowLogger,
conntrack: ct,
publicKey: publicKey,
retryInterval: time.Second,
eventsWithoutAcks: store.NewMemoryStore(),
}
}
@@ -66,6 +72,7 @@ func (m *Manager) needsNewClient(previous *nftypes.FlowConfig) bool {
}
// enableFlow starts components for flow tracking
// must be called under m.mux lock
func (m *Manager) enableFlow(previous *nftypes.FlowConfig) error {
// first make sender ready so events don't pile up
if m.needsNewClient(previous) {
@@ -85,6 +92,7 @@ func (m *Manager) enableFlow(previous *nftypes.FlowConfig) error {
return nil
}
// must be called under m.mux lock
func (m *Manager) resetClient() error {
if m.receiverClient != nil {
if err := m.receiverClient.Close(); err != nil {
@@ -107,14 +115,19 @@ func (m *Manager) resetClient() error {
ctx, cancel := context.WithCancel(context.Background())
m.cancel = cancel
m.shutdownWg.Add(2)
m.shutdownWg.Add(3)
flowConfigInterval := m.flowConfig.Interval
go func() {
defer m.shutdownWg.Done()
m.receiveACKs(ctx, flowClient)
m.receiveACKs(ctx, flowClient, flowConfigInterval)
}()
go func() {
defer m.shutdownWg.Done()
m.startSender(ctx)
m.startSender(ctx, flowConfigInterval)
}()
go func() {
defer m.shutdownWg.Done()
m.startRetries(ctx, flowConfigInterval)
}()
return nil
@@ -198,8 +211,8 @@ func (m *Manager) GetLogger() nftypes.FlowLogger {
return m.logger
}
func (m *Manager) startSender(ctx context.Context) {
ticker := time.NewTicker(m.flowConfig.Interval)
func (m *Manager) startSender(ctx context.Context, flowConfigInterval time.Duration) {
ticker := time.NewTicker(flowConfigInterval)
defer ticker.Stop()
for {
@@ -207,27 +220,29 @@ func (m *Manager) startSender(ctx context.Context) {
case <-ctx.Done():
return
case <-ticker.C:
events := m.logger.GetEvents()
collectedEvents := m.logger.ResetAggregationWindow()
events := collectedEvents.GetAggregatedEvents()
for _, event := range events {
m.eventsWithoutAcks.StoreEvent(event)
if err := m.send(event); err != nil {
log.Errorf("failed to send flow event to server: %v", err)
continue
} else {
log.Tracef("sent flow event: %s", event.ID)
}
log.Tracef("sent flow event: %s", event.ID)
}
}
}
}
func (m *Manager) receiveACKs(ctx context.Context, client *client.GRPCClient) {
err := client.Receive(ctx, m.flowConfig.Interval, func(ack *proto.FlowEventAck) error {
func (m *Manager) receiveACKs(ctx context.Context, client *client.GRPCClient, flowConfigInterval time.Duration) {
err := client.Receive(ctx, flowConfigInterval, func(ack *proto.FlowEventAck) error {
id, err := uuid.FromBytes(ack.EventId)
if err != nil {
log.Warnf("failed to convert ack event id to uuid: %v", err)
return nil
}
log.Tracef("received flow event ack: %s", id)
m.logger.DeleteEvents([]uuid.UUID{id})
m.eventsWithoutAcks.DeleteEvents([]uuid.UUID{id})
return nil
})
@@ -236,6 +251,51 @@ func (m *Manager) receiveACKs(ctx context.Context, client *client.GRPCClient) {
}
}
// We effectively never drop events (see MaxInterval), which makes eventsWithoutAcks unbounded.
// We may want to limit the max size of the store, and start dropping oldest events when the threshold is reached.
func (m *Manager) startRetries(ctx context.Context, flowConfigInterval time.Duration) {
timer := time.NewTimer(m.retryInterval)
retryBackoff := backoff.WithContext(&backoff.ExponentialBackOff{
InitialInterval: 1 * time.Second,
RandomizationFactor: 0.5,
Multiplier: 1.7,
MaxInterval: flowConfigInterval / 2,
MaxElapsedTime: 3 * 30 * 24 * time.Hour, // 3 months
Stop: backoff.Stop,
Clock: backoff.SystemClock,
}, ctx)
defer timer.Stop()
for {
select {
case <-ctx.Done():
return
case <-timer.C:
resetBackoff := true
for _, e := range m.eventsWithoutAcks.GetEvents() {
if e.Timestamp.Add(time.Second).After(time.Now()) {
// grace period on retries to avoid early retries
// do not retry if the event is less than 1 sec old
continue
}
if err := m.send(e); err != nil {
if nextBackoff := retryBackoff.NextBackOff(); nextBackoff != backoff.Stop {
timer = time.NewTimer(nextBackoff)
resetBackoff = false
} else {
resetBackoff = true // we exhausted retries, reset retry loop
}
break
}
}
if resetBackoff { // use regular retry interval in absence of network errors
retryBackoff.Reset()
timer = time.NewTimer(m.retryInterval)
}
}
}
}
func (m *Manager) send(event *nftypes.Event) error {
m.mux.Lock()
client := m.receiverClient
@@ -250,9 +310,11 @@ func (m *Manager) send(event *nftypes.Event) error {
func toProtoEvent(publicKey []byte, event *nftypes.Event) *proto.FlowEvent {
protoEvent := &proto.FlowEvent{
EventId: event.ID[:],
Timestamp: timestamppb.New(event.Timestamp),
PublicKey: publicKey,
EventId: event.ID[:],
Timestamp: timestamppb.New(event.Timestamp),
PublicKey: publicKey,
WindowStart: timestamppb.New(event.WindowStart),
WindowEnd: timestamppb.New(event.WindowEnd),
FlowFields: &proto.FlowFields{
FlowId: event.FlowID[:],
RuleId: event.RuleID,
@@ -267,6 +329,9 @@ func toProtoEvent(publicKey []byte, event *nftypes.Event) *proto.FlowEvent {
TxBytes: event.TxBytes,
SourceResourceId: event.SourceResourceID,
DestResourceId: event.DestResourceID,
NumOfStarts: event.NumOfStarts,
NumOfEnds: event.NumOfEnds,
NumOfDrops: event.NumOfDrops,
},
}

View File

@@ -0,0 +1,291 @@
package netflow
import (
"context"
"errors"
"fmt"
"net"
"net/netip"
"slices"
"testing"
"time"
"github.com/google/uuid"
"github.com/netbirdio/netbird/client/iface/wgaddr"
"github.com/netbirdio/netbird/client/internal/netflow/types"
"github.com/netbirdio/netbird/flow/proto"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
)
type testServer struct {
proto.UnimplementedFlowServiceServer
events chan *proto.FlowEvent
acks chan *proto.FlowEventAck
grpcSrv *grpc.Server
addr string
handlerDone chan struct{} // signaled each time Events() exits
handlerStarted chan struct{} // signaled each time Events() begins
}
func newTestServer(t *testing.T) *testServer {
listener, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
s := &testServer{
events: make(chan *proto.FlowEvent, 100),
acks: make(chan *proto.FlowEventAck, 100),
grpcSrv: grpc.NewServer(),
addr: listener.Addr().String(),
handlerDone: make(chan struct{}, 10),
handlerStarted: make(chan struct{}, 10),
}
proto.RegisterFlowServiceServer(s.grpcSrv, s)
go func() {
if err := s.grpcSrv.Serve(listener); err != nil && !errors.Is(err, grpc.ErrServerStopped) {
t.Logf("server error: %v", err)
}
}()
t.Cleanup(func() {
s.grpcSrv.Stop()
})
return s
}
func (s *testServer) Events(stream proto.FlowService_EventsServer) error {
defer func() {
select {
case s.handlerDone <- struct{}{}:
default:
}
}()
err := stream.Send(&proto.FlowEventAck{IsInitiator: true})
if err != nil {
return err
}
select {
case s.handlerStarted <- struct{}{}:
default:
}
ctx, cancel := context.WithCancel(stream.Context())
defer cancel()
go func() {
defer cancel()
for {
event, err := stream.Recv()
if err != nil {
return
}
if !event.IsInitiator {
select {
case s.events <- event:
case <-ctx.Done():
return
}
}
}
}()
for {
select {
case ack := <-s.acks:
if err := stream.Send(ack); err != nil {
return err
}
case <-ctx.Done():
return ctx.Err()
}
}
}
func TestSendEventReceiveAck(t *testing.T) {
_, cancel := context.WithTimeout(context.Background(), 10*time.Second)
t.Cleanup(cancel)
server := newTestServer(t)
manager := createManager(t, server.addr, 60*time.Second) // set high to prevent retries in this test
defer manager.Close()
assert.Eventually(t, func() bool {
select {
case <-server.handlerStarted:
return true
default:
return false
}
}, 3*time.Second, 100*time.Millisecond)
event1 := types.EventFields{
FlowID: uuid.New(),
Type: types.TypeStart,
Direction: types.Ingress,
DestIP: ipAddr("172.16.1.2"),
DestPort: 2345,
Protocol: 6,
}
manager.logger.StoreEvent(event1)
event2 := types.EventFields{
FlowID: uuid.New(),
Type: types.TypeStart,
Direction: types.Ingress,
DestIP: ipAddr("172.16.1.1"),
DestPort: 1234,
Protocol: 6,
}
manager.logger.StoreEvent(event2)
// verify the server received logged events
serverSideEvents := make([]*proto.FlowEvent, 0)
assert.Eventually(t, func() bool {
select {
case event := <-server.events:
serverSideEvents = append(serverSideEvents, event)
if len(serverSideEvents) == 2 {
return true
}
default:
if len(serverSideEvents) == 2 {
return true
}
}
return false
}, 5*time.Second, 100*time.Millisecond)
serverSideFlowIds := make([]uuid.UUID, 0, 2)
slices.Values(serverSideEvents)(func(e *proto.FlowEvent) bool {
id, err := uuid.FromBytes(e.FlowFields.FlowId)
assert.NoError(t, err)
serverSideFlowIds = append(serverSideFlowIds, id)
return true
})
assert.ElementsMatch(t, []uuid.UUID{event1.FlowID, event2.FlowID}, serverSideFlowIds)
// verify the manager tracks un-acked events
unackedEvents := manager.eventsWithoutAcks.GetEvents()
assert.Len(t, unackedEvents, 2)
flowIds := make([]uuid.UUID, 0)
slices.Values(unackedEvents)(func(e *types.Event) bool {
flowIds = append(flowIds, e.FlowID)
return true
})
assert.ElementsMatch(t, flowIds, []uuid.UUID{event1.FlowID, event2.FlowID})
}
// verify handling of retries:
// - unacked events are retried
// - when acks arrive, events are removed from the un-acked event tracker
func TestRetryEvents(t *testing.T) {
_, cancel := context.WithTimeout(context.Background(), 10*time.Second)
t.Cleanup(cancel)
server := newTestServer(t)
manager := createManager(t, server.addr, time.Second) // set low to start retries sooner
defer manager.Close()
assert.Eventually(t, func() bool {
select {
case <-server.handlerStarted:
return true
default:
return false
}
}, 3*time.Second, 100*time.Millisecond)
event1 := types.EventFields{
FlowID: uuid.New(),
Type: types.TypeStart,
Direction: types.Ingress,
DestIP: ipAddr("172.16.1.2"),
DestPort: 2345,
Protocol: 6,
}
manager.logger.StoreEvent(event1)
event2 := types.EventFields{
FlowID: uuid.New(),
Type: types.TypeStart,
Direction: types.Ingress,
DestIP: ipAddr("172.16.1.1"),
DestPort: 1234,
Protocol: 6,
}
manager.logger.StoreEvent(event2)
// verify the server received retries of logged events
serverSideEvents := make([]*proto.FlowEvent, 0)
func() {
c := time.After(2500 * time.Millisecond)
for {
select {
case event := <-server.events:
serverSideEvents = append(serverSideEvents, event)
case <-c:
return
}
}
}()
assert.True(t, len(serverSideEvents) > 2) // must see retries
uniqueServerSideEvents := make(map[uuid.UUID]*proto.FlowEvent)
slices.Values(serverSideEvents)(func(e *proto.FlowEvent) bool {
id, err := uuid.FromBytes(e.FlowFields.FlowId)
assert.NoError(t, err)
uniqueServerSideEvents[id] = e
return true
})
assert.Contains(t, uniqueServerSideEvents, event1.FlowID)
assert.Contains(t, uniqueServerSideEvents, event2.FlowID)
// ack events
server.acks <- &proto.FlowEventAck{EventId: uniqueServerSideEvents[event1.FlowID].EventId}
server.acks <- &proto.FlowEventAck{EventId: uniqueServerSideEvents[event2.FlowID].EventId}
assert.EventuallyWithT(t, func(c *assert.CollectT) {
unackedEvents := manager.eventsWithoutAcks.GetEvents()
assert.Empty(c, unackedEvents)
}, 3*time.Second, 100*time.Millisecond)
}
func createManager(t *testing.T, serverAddr string, retryInterval time.Duration) *Manager {
t.Helper()
mockIFace := &mockIFaceMapper{
address: wgaddr.Address{
Network: netip.MustParsePrefix("192.168.1.1/32"),
},
isUserspaceBind: true,
}
publicKey := []byte("test-public-key")
manager := NewManager(mockIFace, publicKey, nil)
manager.retryInterval = retryInterval
initialConfig := &types.FlowConfig{
Enabled: true,
URL: fmt.Sprintf("http://%s", serverAddr),
TokenPayload: "initial-payload",
TokenSignature: "initial-signature",
Interval: 500 * time.Millisecond,
}
err := manager.Update(initialConfig)
require.NoError(t, err)
return manager
}
func ipAddr(a string) netip.Addr {
addr, _ := netip.ParseAddr(a)
return addr
}

View File

@@ -0,0 +1,362 @@
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
addresses [][]netip.Addr
dstPort uint16
eventTypes []types.Type
}{
{
description: "start and stop",
addresses: [][]netip.Addr{{netip.MustParseAddr("1.1.1.1"), netip.MustParseAddr("2.2.2.2")}, {netip.MustParseAddr("3.3.3.3"), netip.MustParseAddr("2.2.2.2")}},
dstPort: uint16(random.Uint32() >> 16),
eventTypes: []types.Type{types.TypeStart, types.TypeEnd},
},
{
description: "start and drop",
addresses: [][]netip.Addr{{netip.MustParseAddr("1.1.1.1"), netip.MustParseAddr("2.2.2.2")}, {netip.MustParseAddr("3.3.3.3"), netip.MustParseAddr("2.2.2.2")}},
dstPort: uint16(random.Uint32() >> 16),
eventTypes: []types.Type{types.TypeStart, types.TypeDrop},
},
{
description: "start only",
addresses: [][]netip.Addr{{netip.MustParseAddr("1.1.1.1"), netip.MustParseAddr("2.2.2.2")}, {netip.MustParseAddr("3.3.3.3"), netip.MustParseAddr("2.2.2.2")}},
dstPort: uint16(random.Uint32() >> 16),
eventTypes: []types.Type{types.TypeStart},
},
{
description: "drop only",
addresses: [][]netip.Addr{{netip.MustParseAddr("1.1.1.1"), netip.MustParseAddr("2.2.2.2")}, {netip.MustParseAddr("3.3.3.3"), netip.MustParseAddr("2.2.2.2")}},
dstPort: uint16(random.Uint32() >> 16),
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()
store.WindowEnd = time.Now().Add(5 * time.Second)
allExpected := make([]*types.Event, 0)
for _, srcAndDst := range tt.addresses {
inEvents, expected := generateEvents(srcAndDst[0], srcAndDst[1], tt.dstPort, tt.eventTypes, protocol, types.Ingress, 0, store.WindowStart, store.WindowEnd)
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
addresses [][]netip.Addr
eventTypes []types.Type
}{
{
description: "start and stop",
addresses: [][]netip.Addr{{netip.MustParseAddr("1.1.1.1"), netip.MustParseAddr("2.2.2.2")}},
eventTypes: []types.Type{types.TypeStart, types.TypeEnd},
},
{
description: "start and drop",
addresses: [][]netip.Addr{{netip.MustParseAddr("1.1.1.1"), netip.MustParseAddr("2.2.2.2")}},
eventTypes: []types.Type{types.TypeStart, types.TypeDrop},
},
{
description: "start only",
addresses: [][]netip.Addr{{netip.MustParseAddr("1.1.1.1"), netip.MustParseAddr("2.2.2.2")}},
eventTypes: []types.Type{types.TypeStart},
},
{
description: "drop only",
addresses: [][]netip.Addr{{netip.MustParseAddr("1.1.1.1"), netip.MustParseAddr("2.2.2.2")}},
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()
store.WindowEnd = time.Now().Add(5 * time.Second)
allExpected := make([]*types.Event, 0)
for _, icmpType := range icmpTypes {
events, expected := generateEvents(tt.addresses[0][0], tt.addresses[0][1], 0, tt.eventTypes, protocol, types.Ingress, icmpType, store.WindowStart, store.WindowEnd)
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 TestFlowAggregationOfUnknownProtocols(t *testing.T) {
var tests = []struct {
description string
addresses [][]netip.Addr
dstPort uint16
eventTypes []types.Type
}{
{
description: "start and stop",
addresses: [][]netip.Addr{{netip.MustParseAddr("1.1.1.1"), netip.MustParseAddr("2.2.2.2")}, {netip.MustParseAddr("3.3.3.3"), netip.MustParseAddr("2.2.2.2")}},
dstPort: uint16(random.Uint32() >> 16),
eventTypes: []types.Type{types.TypeStart, types.TypeEnd},
},
{
description: "start and drop",
addresses: [][]netip.Addr{{netip.MustParseAddr("1.1.1.1"), netip.MustParseAddr("2.2.2.2")}, {netip.MustParseAddr("3.3.3.3"), netip.MustParseAddr("2.2.2.2")}},
dstPort: uint16(random.Uint32() >> 16),
eventTypes: []types.Type{types.TypeStart, types.TypeDrop},
},
{
description: "start only",
addresses: [][]netip.Addr{{netip.MustParseAddr("1.1.1.1"), netip.MustParseAddr("2.2.2.2")}, {netip.MustParseAddr("3.3.3.3"), netip.MustParseAddr("2.2.2.2")}},
dstPort: uint16(random.Uint32() >> 16),
eventTypes: []types.Type{types.TypeStart},
},
{
description: "drop only",
addresses: [][]netip.Addr{{netip.MustParseAddr("1.1.1.1"), netip.MustParseAddr("2.2.2.2")}, {netip.MustParseAddr("3.3.3.3"), netip.MustParseAddr("2.2.2.2")}},
dstPort: uint16(random.Uint32() >> 16),
eventTypes: []types.Type{types.TypeDrop},
}}
for _, tt := range tests {
t.Run(tt.description+" "+types.ProtocolUnknown.String(), 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 TestResetAggregationWindow(t *testing.T) {
store := NewAggregatingMemoryStore()
store.StoreEvent(&types.Event{
ID: uuid.New(),
Timestamp: time.Now(),
EventFields: types.EventFields{
FlowID: uuid.New(),
Type: types.TypeStart,
Protocol: types.TCP,
RuleID: []byte("rule-id-1"),
Direction: types.Ingress,
SourceIP: netip.MustParseAddr("1.1.1.1"),
SourcePort: 1234,
DestIP: netip.MustParseAddr("2.2.2.2"),
DestPort: 5678,
SourceResourceID: []byte("source-resource-id"),
DestResourceID: []byte("dest-resource-id"),
RxPackets: random.Uint64(),
TxPackets: random.Uint64(),
RxBytes: random.Uint64(),
TxBytes: random.Uint64(),
},
})
reset := store.ResetAggregationWindow()
previousEvents, ok := reset.(*AggregatingMemory)
assert.True(t, ok)
assert.NotEqual(t, previousEvents.WindowStart, store.WindowStart)
assert.Equal(t, previousEvents.WindowEnd, store.WindowStart)
assert.NotEmpty(t, previousEvents.events)
assert.Empty(t, store.events)
}
func generateEvents(srcIp, dstIp netip.Addr, dstPort uint16, eventTypes []types.Type, protocol types.Protocol,
direction types.Direction, icmpType uint8, windowStart, windowEnd time.Time) ([]*types.Event, *types.Event) {
var rxPackets, txPackets, rxBytes, txBytes uint64
inEvents := 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(),
}}
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,
WindowStart: windowStart,
WindowEnd: windowEnd,
EventFields: types.EventFields{
FlowID: flowId,
Type: types.TypeUnknown,
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
}
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

@@ -1,10 +1,15 @@
package store
import (
"maps"
"math/rand"
v2 "math/rand/v2"
"net/netip"
"slices"
"sync"
"time"
"github.com/google/uuid"
"github.com/netbirdio/netbird/client/internal/netflow/types"
)
@@ -19,6 +24,13 @@ type Memory struct {
events map[uuid.UUID]*types.Event
}
type AggregatingMemory struct {
Memory
WindowStart time.Time
WindowEnd time.Time
rnd *v2.PCG
}
func (m *Memory) StoreEvent(event *types.Event) {
m.mux.Lock()
defer m.mux.Unlock()
@@ -48,3 +60,95 @@ func (m *Memory) DeleteEvents(ids []uuid.UUID) {
delete(m.events, id)
}
}
func NewAggregatingMemoryStore() *AggregatingMemory {
return &AggregatingMemory{WindowStart: time.Now(), Memory: Memory{events: make(map[uuid.UUID]*types.Event)}, rnd: v2.NewPCG(rand.Uint64(), rand.Uint64())}
}
func (am *AggregatingMemory) ResetAggregationWindow() types.FlowEventAggregator {
am.mux.Lock()
defer am.mux.Unlock()
now := time.Now()
toret := AggregatingMemory{WindowStart: am.WindowStart, WindowEnd: now, Memory: Memory{events: am.events}, rnd: v2.NewPCG(rand.Uint64(), rand.Uint64())}
am.events = make(map[uuid.UUID]*types.Event)
am.WindowStart = now
return &toret
}
type aggregationKey struct {
srcAddr netip.Addr
destAddr netip.Addr
destPort uint16
direction int
protocol uint8
icmpType uint8
unique uint64 // used to prevent aggregation on non icmp/udp/tcp events
}
func (am *AggregatingMemory) GetAggregatedEvents() []*types.Event {
am.mux.Lock()
defer am.mux.Unlock()
aggregated := make(map[aggregationKey]*types.Event)
for _, v := range am.events {
lookupKey := aggregationKey{srcAddr: v.SourceIP, destAddr: v.DestIP, destPort: v.DestPort, direction: int(v.Direction), protocol: uint8(v.Protocol), icmpType: v.ICMPType}
if _, ok := aggregated[lookupKey]; !ok {
event := v.Clone()
switch event.Type {
case types.TypeStart:
event.NumOfStarts += 1
case types.TypeDrop:
event.NumOfDrops += 1
case types.TypeEnd:
event.NumOfEnds += 1
}
event.Type = types.TypeUnknown
// Please note that ICMPCode field isn't propagated by the manager (see flow/proto/flow.pb.go, FlowFields struct)
// so the field value in an icmp event in the "aggregated" doesn't matter
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 = am.rnd.Uint64() // to make the lookup key unique so we don't aggregate on it
}
aggregated[lookupKey] = event
continue
}
aggregatedEvent := aggregated[lookupKey]
if aggregatedEvent.Protocol != types.ICMP && aggregatedEvent.Protocol != types.ICMPv6 && aggregatedEvent.Protocol != types.UDP && aggregatedEvent.Protocol != types.TCP {
continue // we don't aggregate this type of events; shouldn't ever get here
}
// track the number of connections, duration?, open and close events?
aggregatedEvent.RxBytes += v.RxBytes
aggregatedEvent.RxPackets += v.RxPackets
aggregatedEvent.TxBytes += v.TxBytes
aggregatedEvent.TxPackets += v.TxPackets
switch v.Type {
case types.TypeStart:
aggregatedEvent.NumOfStarts += 1
case types.TypeDrop:
aggregatedEvent.NumOfDrops += 1
case types.TypeEnd:
aggregatedEvent.NumOfEnds += 1
}
if aggregatedEvent.Timestamp.Compare(v.Timestamp) > 0 {
aggregatedEvent.Timestamp = v.Timestamp
aggregatedEvent.ID = v.ID
aggregatedEvent.SourcePort = v.SourcePort
}
if len(aggregatedEvent.RuleID) == 0 && len(v.RuleID) != 0 {
aggregatedEvent.RuleID = slices.Clone(v.RuleID)
}
}
return slices.Collect(maps.Values(aggregated)) // could return an iterator instead here
}

View File

@@ -2,6 +2,7 @@ package types
import (
"net/netip"
"slices"
"strconv"
"time"
@@ -69,8 +70,10 @@ const (
)
type Event struct {
ID uuid.UUID
Timestamp time.Time
ID uuid.UUID
Timestamp time.Time
WindowStart time.Time
WindowEnd time.Time
EventFields
}
@@ -92,6 +95,17 @@ type EventFields struct {
TxPackets uint64
RxBytes uint64
TxBytes uint64
NumOfStarts uint64
NumOfEnds uint64
NumOfDrops uint64
}
func (e *Event) Clone() *Event {
toret := *e
toret.RuleID = slices.Clone(e.RuleID)
toret.SourceResourceID = slices.Clone(e.SourceResourceID)
toret.DestResourceID = slices.Clone(e.DestResourceID)
return &toret
}
type FlowConfig struct {
@@ -114,13 +128,15 @@ type FlowManager interface {
GetLogger() FlowLogger
}
type FlowEventAggregator interface {
ResetAggregationWindow() FlowEventAggregator
GetAggregatedEvents() []*Event
}
type FlowLogger interface {
ResetAggregationWindow() FlowEventAggregator
// StoreEvent stores a flow event
StoreEvent(flowEvent EventFields)
// GetEvents returns all stored events
GetEvents() []*Event
// DeleteEvents deletes events from the store
DeleteEvents([]uuid.UUID)
// Close closes the logger
Close()
// Enable enables the flow logger receiver
@@ -140,6 +156,11 @@ type Store interface {
Close()
}
type AggregatingStore interface {
FlowEventAggregator
Store
}
// ConnTracker defines the interface for connection tracking functionality
type ConnTracker interface {
// Start begins tracking connections by listening for conntrack events.