[client] Add DNS and Exit Node collection configuration to Netflow (#3522)

[client] Add DNS and Exit Node collection configuration to Netflow (#3522)

client/internal/engine.go

@@ -734,12 +734,14 @@ func toFlowLoggerConfig(config *mgmProto.FlowConfig) (*nftypes.FlowConfig, error
 		return nil, errors.New("flow interval is nil")
 	}
 	return &nftypes.FlowConfig{
-		Enabled:        config.GetEnabled(),
-		Counters:       config.GetCounters(),
-		URL:            config.GetUrl(),
-		TokenPayload:   config.GetTokenPayload(),
-		TokenSignature: config.GetTokenSignature(),
-		Interval:       config.GetInterval().AsDuration(),
+		Enabled:            config.GetEnabled(),
+		Counters:           config.GetCounters(),
+		URL:                config.GetUrl(),
+		TokenPayload:       config.GetTokenPayload(),
+		TokenSignature:     config.GetTokenSignature(),
+		Interval:           config.GetInterval().AsDuration(),
+		DNSCollection:      config.GetDnsCollection(),
+		ExitNodeCollection: config.GetExitNodeCollection(),
 	}, nil
 }
 

client/internal/netflow/logger/logger.go

@@ -10,6 +10,7 @@ import (
 	"github.com/google/uuid"
 	log "github.com/sirupsen/logrus"
 
+	"github.com/netbirdio/netbird/client/internal/dnsfwd"
 	"github.com/netbirdio/netbird/client/internal/netflow/store"
 	"github.com/netbirdio/netbird/client/internal/netflow/types"
 	"github.com/netbirdio/netbird/client/internal/peer"
@@ -17,15 +18,17 @@ import (
 
 type rcvChan chan *types.EventFields
 type Logger struct {
-	mux            sync.Mutex
-	ctx            context.Context
-	cancel         context.CancelFunc
-	enabled        atomic.Bool
-	rcvChan        atomic.Pointer[rcvChan]
-	cancelReceiver context.CancelFunc
-	statusRecorder *peer.Status
-	wgIfaceIPNet   net.IPNet
-	Store          types.Store
+	mux                sync.Mutex
+	ctx                context.Context
+	cancel             context.CancelFunc
+	enabled            atomic.Bool
+	rcvChan            atomic.Pointer[rcvChan]
+	cancelReceiver     context.CancelFunc
+	statusRecorder     *peer.Status
+	wgIfaceIPNet       net.IPNet
+	dnsCollection      atomic.Bool
+	exitNodeCollection atomic.Bool
+	Store              types.Store
 }
 
 func New(ctx context.Context, statusRecorder *peer.Status, wgIfaceIPNet net.IPNet) *Logger {
@@ -88,17 +91,20 @@ func (l *Logger) startReceiver() {
 				Timestamp:   time.Now(),
 			}
 
+			var isExitNode bool
 			if event.Direction == types.Ingress {
 				if !l.wgIfaceIPNet.Contains(net.IP(event.SourceIP.AsSlice())) {
-					event.SourceResourceID = []byte(l.statusRecorder.CheckRoutes(event.SourceIP))
+					event.SourceResourceID, isExitNode = l.statusRecorder.CheckRoutes(event.SourceIP)
 				}
 			} else if event.Direction == types.Egress {
 				if !l.wgIfaceIPNet.Contains(net.IP(event.DestIP.AsSlice())) {
-					event.DestResourceID = []byte(l.statusRecorder.CheckRoutes(event.DestIP))
+					event.DestResourceID, isExitNode = l.statusRecorder.CheckRoutes(event.DestIP)
 				}
 			}
 
-			l.Store.StoreEvent(&event)
+			if l.shouldStore(eventFields, isExitNode) {
+				l.Store.StoreEvent(&event)
+			}
 		}
 	}
 }
@@ -131,7 +137,26 @@ func (l *Logger) DeleteEvents(ids []uuid.UUID) {
 	l.Store.DeleteEvents(ids)
 }
 
+func (l *Logger) UpdateConfig(dnsCollection, exitNodeCollection bool) {
+	l.dnsCollection.Store(dnsCollection)
+	l.exitNodeCollection.Store(exitNodeCollection)
+}
+
 func (l *Logger) Close() {
 	l.stop()
 	l.cancel()
 }
+
+func (l *Logger) shouldStore(event *types.EventFields, isExitNode bool) bool {
+	// check dns collection
+	if !l.dnsCollection.Load() && event.Protocol == types.UDP && (event.DestPort == 53 || event.DestPort == dnsfwd.ListenPort) {
+		return false
+	}
+
+	// check exit node collection
+	if !l.exitNodeCollection.Load() && isExitNode {
+		return false
+	}
+
+	return true
+}

client/internal/netflow/manager.go

@@ -122,6 +122,8 @@ func (m *Manager) Update(update *nftypes.FlowConfig) error {
 	previous := m.flowConfig
 	m.flowConfig = update
 
+	m.logger.UpdateConfig(update.DNSCollection, update.ExitNodeCollection)
+
 	if update.Enabled {
 		return m.enableFlow(previous)
 	}

client/internal/netflow/types/types.go

@@ -90,12 +90,14 @@ type EventFields struct {
 }
 
 type FlowConfig struct {
-	URL            string
-	Interval       time.Duration
-	Enabled        bool
-	Counters       bool
-	TokenPayload   string
-	TokenSignature string
+	URL                string
+	Interval           time.Duration
+	Enabled            bool
+	Counters           bool
+	TokenPayload       string
+	TokenSignature     string
+	DNSCollection      bool
+	ExitNodeCollection bool
 }
 
 type FlowManager interface {
@@ -120,6 +122,9 @@ type FlowLogger interface {
 	Enable()
 	// Disable disables the flow logger receiver
 	Disable()
+
+	// UpdateConfig updates the flow manager configuration
+	UpdateConfig(dnsCollection, exitNodeCollection bool)
 }
 
 type Store interface {

client/internal/peer/route.go

@@ -43,16 +43,22 @@ func (r *routeIDLookup) RemoveResolvedIP(route netip.Prefix) {
 	r.resolvedIPs.Delete(route.Addr())
 }
 
-func (r *routeIDLookup) Lookup(ip netip.Addr) string {
+// Lookup returns the resource ID for the given IP address
+// and a bool indicating if the IP is an exit node
+func (r *routeIDLookup) Lookup(ip netip.Addr) (string, bool) {
+	var isExitNode bool
+
 	resId, ok := r.resolvedIPs.Load(ip)
 	if ok {
-		return resId.(string)
+		return resId.(string), false
 	}
 
 	var resourceID string
 	r.localMap.Range(func(key, value interface{}) bool {
-		if key.(netip.Prefix).Contains(ip) {
+		pref := key.(netip.Prefix)
+		if pref.Contains(ip) {
 			resourceID = value.(string)
+			isExitNode = pref.Bits() == 0
 			return false
 
 		}
@@ -61,13 +67,15 @@ func (r *routeIDLookup) Lookup(ip netip.Addr) string {
 
 	if resourceID == "" {
 		r.remoteMap.Range(func(key, value interface{}) bool {
-			if key.(netip.Prefix).Contains(ip) {
+			pref := key.(netip.Prefix)
+			if pref.Contains(ip) {
 				resourceID = value.(string)
+				isExitNode = pref.Bits() == 0
 				return false
 			}
 			return true
 		})
 	}
 
-	return resourceID
+	return resourceID, isExitNode
 }

client/internal/peer/status.go

@@ -364,11 +364,12 @@ func (d *Status) RemovePeerStateRoute(peer string, route string) error {
 
 // CheckRoutes checks if the source and destination addresses are within the same route
 // and returns the resource ID of the route that contains the addresses
-func (d *Status) CheckRoutes(ip netip.Addr) (resId string) {
+func (d *Status) CheckRoutes(ip netip.Addr) ([]byte, bool) {
 	if d == nil {
-		return
+		return nil, false
 	}
-	return d.routeIDLookup.Lookup(ip)
+	resId, isExitNode := d.routeIDLookup.Lookup(ip)
+	return []byte(resId), isExitNode
 }
 
 func (d *Status) UpdatePeerICEState(receivedState State) error {