mirror of
https://github.com/netbirdio/netbird.git
synced 2026-07-16 11:39:57 +00:00
133 lines
3.3 KiB
Go
133 lines
3.3 KiB
Go
package store
|
|
|
|
import (
|
|
"maps"
|
|
"net/netip"
|
|
"slices"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/netbirdio/netbird/client/internal/netflow/types"
|
|
)
|
|
|
|
func NewMemoryStore() *Memory {
|
|
return &Memory{
|
|
events: make(map[uuid.UUID]*types.Event),
|
|
}
|
|
}
|
|
|
|
type Memory struct {
|
|
mux sync.Mutex
|
|
events map[uuid.UUID]*types.Event
|
|
}
|
|
|
|
type AggregatingMemory struct {
|
|
Memory
|
|
}
|
|
|
|
func (m *Memory) StoreEvent(event *types.Event) {
|
|
m.mux.Lock()
|
|
defer m.mux.Unlock()
|
|
m.events[event.ID] = event
|
|
}
|
|
|
|
func (m *Memory) Close() {
|
|
m.mux.Lock()
|
|
defer m.mux.Unlock()
|
|
clear(m.events)
|
|
}
|
|
|
|
func (m *Memory) GetEvents() []*types.Event {
|
|
m.mux.Lock()
|
|
defer m.mux.Unlock()
|
|
events := make([]*types.Event, 0, len(m.events))
|
|
for _, event := range m.events {
|
|
events = append(events, event)
|
|
}
|
|
return events
|
|
}
|
|
|
|
func (m *Memory) DeleteEvents(ids []uuid.UUID) {
|
|
m.mux.Lock()
|
|
defer m.mux.Unlock()
|
|
for _, id := range ids {
|
|
delete(m.events, id)
|
|
}
|
|
}
|
|
|
|
func NewAggregatingMemoryStore() *AggregatingMemory {
|
|
return &AggregatingMemory{Memory{events: make(map[uuid.UUID]*types.Event)}}
|
|
}
|
|
|
|
func (am *AggregatingMemory) ResetAggregationWindow() types.FlowEventAggregator {
|
|
am.mux.Lock()
|
|
defer am.mux.Unlock()
|
|
|
|
toret := AggregatingMemory{Memory: Memory{events: am.events}}
|
|
am.events = make(map[uuid.UUID]*types.Event)
|
|
|
|
return &toret
|
|
}
|
|
|
|
type aggregationKey struct {
|
|
destAddr netip.Addr
|
|
destPort uint16
|
|
protocol uint8
|
|
icmpType uint8
|
|
unique 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.ICMPType}
|
|
if _, ok := aggregated[lookupKey]; !ok {
|
|
aggregated[lookupKey] = v.Clone()
|
|
event := aggregated[lookupKey]
|
|
|
|
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
|
|
continue
|
|
}
|
|
|
|
switch event.Type {
|
|
case types.TypeStart:
|
|
event.NumOfStarts += 1
|
|
case types.TypeDrop:
|
|
event.NumOfDrops += 1
|
|
case types.TypeEnd:
|
|
event.NumOfEnds += 1
|
|
}
|
|
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.Type = v.Type
|
|
}
|
|
// do we aggregate icmp by code?
|
|
}
|
|
|
|
return slices.Collect(maps.Values(aggregated)) // could return an iterator instead here
|
|
}
|