From 60bcf7dfc35ec498175b2913185032cb015eb1a4 Mon Sep 17 00:00:00 2001
From: Dmitri <dmitri.external@netbird.io>
Date: Tue, 9 Jun 2026 11:38:02 +0200
Subject: [PATCH] added an implementation of aggregating memory store

Signed-off-by: Dmitri <dmitri.external@netbird.io>
---
 client/internal/netflow/store/memory.go | 58 +++++++++++++++++++++++++
 1 file changed, 58 insertions(+)

diff --git a/client/internal/netflow/store/memory.go b/client/internal/netflow/store/memory.go
index a44505e96..f9151726e 100644
--- a/client/internal/netflow/store/memory.go
+++ b/client/internal/netflow/store/memory.go
@@ -1,7 +1,11 @@
 package store
 
 import (
+	"maps"
+	"net/netip"
+	"slices"
 	"sync"
+	"time"
 
 	"github.com/google/uuid"
 
@@ -19,6 +23,10 @@ type Memory struct {
 	events map[uuid.UUID]*types.Event
 }
 
+type AggregatingMemory struct {
+	Memory
+}
+
 func (m *Memory) StoreEvent(event *types.Event) {
 	m.mux.Lock()
 	defer m.mux.Unlock()
@@ -48,3 +56,53 @@ func (m *Memory) DeleteEvents(ids []uuid.UUID) {
 		delete(m.events, id)
 	}
 }
+
+func (am *AggregatingMemory) StartAggregationWindow() *AggregatingMemory {
+	am.mux.Lock()
+	defer am.mux.Unlock()
+
+	toret := AggregatingMemory{Memory: Memory{events: am.Memory.events}}
+	am.events = make(map[uuid.UUID]*types.Event)
+
+	return &toret
+}
+
+type aggregationKey struct {
+	destAddr netip.Addr
+	destPort uint16
+	protocol uint8
+	icmpType uint8
+	ts       int64 // used to prevent aggregation on non icmp/udp/tcp events
+}
+
+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}
+		if aggregatedEvent, ok := aggregated[lookupKey]; ok {
+			switch aggregatedEvent.Protocol {
+			case types.ICMP, types.ICMPv6, types.UDP, types.TCP:
+				aggregatedEvent.RxBytes += v.RxBytes
+				aggregatedEvent.RxPackets += v.RxPackets
+				aggregatedEvent.TxBytes += v.TxBytes
+				aggregatedEvent.TxPackets += v.TxPackets
+				if aggregatedEvent.Timestamp.Compare(v.Timestamp) < 0 {
+					aggregatedEvent.Timestamp = v.Timestamp
+				}
+				// do we aggregate icmp by code?
+			default:
+				// shouldn't get here
+			}
+		} else {
+			switch v.Protocol {
+			case types.ICMP, types.ICMPv6, types.TCP, types.UDP:
+				aggregated[lookupKey] = v
+			default:
+				lookupKey.ts = time.Now().UnixNano()
+				aggregated[lookupKey] = v
+			}
+		}
+	}
+
+	return slices.Collect(maps.Values(aggregated)) // could return an iterator instead here
+}