Add go.uber.org/mock dependency to go.mod

Refactor IdP user data caching by introducing a Redis cache implementation alongside an in-memory fallback, adding a Marshaler interface for flexible serialization, and updating related tests and account management code.

- Added a new cache store implementation with support for Redis and in-memory backends.
- Introduced Marshaler and wrapper types for handling serialization with msgpack and JSON.
- Updated account and user management modules to integrate and test the new caching strategy.

client/cmd/testutil_test.go

@@ -90,13 +90,13 @@ func startManagement(t *testing.T, config *mgmt.Config, testFile string) (*grpc.
 	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
 	require.NoError(t, err)
 
-	accountManager, err := mgmt.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, iv, metrics, port_forwarding.NewControllerMock())
+	accountManager, err := mgmt.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, iv, metrics, port_forwarding.NewControllerMock(), settings.NewManagerMock())
 	if err != nil {
 		t.Fatal(err)
 	}
 
 	secretsManager := mgmt.NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig, config.Relay)
-	mgmtServer, err := mgmt.NewServer(context.Background(), config, accountManager, settings.NewManager(store), peersUpdateManager, secretsManager, nil, nil, nil)
+	mgmtServer, err := mgmt.NewServer(context.Background(), config, accountManager, settings.NewManagerMock(), peersUpdateManager, secretsManager, nil, nil, nil)
 	if err != nil {
 		t.Fatal(err)
 	}

client/firewall/create.go

@@ -10,17 +10,18 @@ import (
 
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/firewall/uspfilter"
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 )
 
 // NewFirewall creates a firewall manager instance
-func NewFirewall(iface IFaceMapper, _ *statemanager.Manager, disableServerRoutes bool) (firewall.Manager, error) {
+func NewFirewall(iface IFaceMapper, _ statemanager.Manager, flowLogger nftypes.FlowLogger, disableServerRoutes bool) (firewall.Manager, error) {
 	if !iface.IsUserspaceBind() {
 		return nil, fmt.Errorf("not implemented for this OS: %s", runtime.GOOS)
 	}
 
 	// use userspace packet filtering firewall
-	fm, err := uspfilter.Create(iface, disableServerRoutes)
+	fm, err := uspfilter.Create(iface, disableServerRoutes, flowLogger)
 	if err != nil {
 		return nil, err
 	}

client/firewall/create_linux.go

@@ -15,6 +15,7 @@ import (
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	nbnftables "github.com/netbirdio/netbird/client/firewall/nftables"
 	"github.com/netbirdio/netbird/client/firewall/uspfilter"
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 )
 
@@ -33,7 +34,7 @@ const SKIP_NFTABLES_ENV = "NB_SKIP_NFTABLES_CHECK"
 // FWType is the type for the firewall type
 type FWType int
 
-func NewFirewall(iface IFaceMapper, stateManager *statemanager.Manager, disableServerRoutes bool) (firewall.Manager, error) {
+func NewFirewall(iface IFaceMapper, stateManager statemanager.Manager, flowLogger nftypes.FlowLogger, disableServerRoutes bool) (firewall.Manager, error) {
 	// on the linux system we try to user nftables or iptables
 	// in any case, because we need to allow netbird interface traffic
 	// so we use AllowNetbird traffic from these firewall managers
@@ -47,10 +48,10 @@ func NewFirewall(iface IFaceMapper, stateManager *statemanager.Manager, disableS
 	if err != nil {
 		log.Warnf("failed to create native firewall: %v. Proceeding with userspace", err)
 	}
-	return createUserspaceFirewall(iface, fm, disableServerRoutes)
+	return createUserspaceFirewall(iface, fm, disableServerRoutes, flowLogger)
 }
 
-func createNativeFirewall(iface IFaceMapper, stateManager *statemanager.Manager, routes bool) (firewall.Manager, error) {
+func createNativeFirewall(iface IFaceMapper, stateManager statemanager.Manager, routes bool) (firewall.Manager, error) {
 	fm, err := createFW(iface)
 	if err != nil {
 		return nil, fmt.Errorf("create firewall: %s", err)
@@ -77,12 +78,12 @@ func createFW(iface IFaceMapper) (firewall.Manager, error) {
 	}
 }
 
-func createUserspaceFirewall(iface IFaceMapper, fm firewall.Manager, disableServerRoutes bool) (firewall.Manager, error) {
+func createUserspaceFirewall(iface IFaceMapper, fm firewall.Manager, disableServerRoutes bool, flowLogger nftypes.FlowLogger) (firewall.Manager, error) {
 	var errUsp error
 	if fm != nil {
-		fm, errUsp = uspfilter.CreateWithNativeFirewall(iface, fm, disableServerRoutes)
+		fm, errUsp = uspfilter.CreateWithNativeFirewall(iface, fm, disableServerRoutes, flowLogger)
 	} else {
-		fm, errUsp = uspfilter.Create(iface, disableServerRoutes)
+		fm, errUsp = uspfilter.Create(iface, disableServerRoutes, flowLogger)
 	}
 
 	if errUsp != nil {

client/firewall/iptables/acl_linux.go

@@ -36,7 +36,7 @@ type aclManager struct {
 	optionalEntries map[string][]entry
 	ipsetStore      *ipsetStore
 
-	stateManager *statemanager.Manager
+	stateManager statemanager.Manager
 }
 
 func newAclManager(iptablesClient *iptables.IPTables, wgIface iFaceMapper) (*aclManager, error) {
@@ -55,7 +55,7 @@ func newAclManager(iptablesClient *iptables.IPTables, wgIface iFaceMapper) (*acl
 	return m, nil
 }
 
-func (m *aclManager) init(stateManager *statemanager.Manager) error {
+func (m *aclManager) init(stateManager statemanager.Manager) error {
 	m.stateManager = stateManager
 
 	m.seedInitialEntries()
@@ -75,6 +75,7 @@ func (m *aclManager) init(stateManager *statemanager.Manager) error {
 }
 
 func (m *aclManager) AddPeerFiltering(
+	id []byte,
 	ip net.IP,
 	protocol firewall.Protocol,
 	sPort *firewall.Port,

client/firewall/iptables/manager_linux.go

@@ -60,7 +60,7 @@ func Create(wgIface iFaceMapper) (*Manager, error) {
 	return m, nil
 }
 
-func (m *Manager) Init(stateManager *statemanager.Manager) error {
+func (m *Manager) Init(stateManager statemanager.Manager) error {
 	state := &ShutdownState{
 		InterfaceState: &InterfaceState{
 			NameStr:       m.wgIface.Name(),
@@ -96,21 +96,22 @@ func (m *Manager) Init(stateManager *statemanager.Manager) error {
 //
 // Comment will be ignored because some system this feature is not supported
 func (m *Manager) AddPeerFiltering(
+	id []byte,
 	ip net.IP,
-	protocol firewall.Protocol,
+	proto firewall.Protocol,
 	sPort *firewall.Port,
 	dPort *firewall.Port,
 	action firewall.Action,
 	ipsetName string,
-	_ string,
 ) ([]firewall.Rule, error) {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	return m.aclMgr.AddPeerFiltering(ip, protocol, sPort, dPort, action, ipsetName)
+	return m.aclMgr.AddPeerFiltering(id, ip, proto, sPort, dPort, action, ipsetName)
 }
 
 func (m *Manager) AddRouteFiltering(
+	id []byte,
 	sources []netip.Prefix,
 	destination netip.Prefix,
 	proto firewall.Protocol,
@@ -125,7 +126,7 @@ func (m *Manager) AddRouteFiltering(
 		return nil, fmt.Errorf("unsupported IP version: %s", destination.Addr().String())
 	}
 
-	return m.router.AddRouteFiltering(sources, destination, proto, sPort, dPort, action)
+	return m.router.AddRouteFiltering(id, sources, destination, proto, sPort, dPort, action)
 }
 
 // DeletePeerRule from the firewall by rule definition
@@ -166,7 +167,7 @@ func (m *Manager) SetLegacyManagement(isLegacy bool) error {
 }
 
 // Reset firewall to the default state
-func (m *Manager) Close(stateManager *statemanager.Manager) error {
+func (m *Manager) Close(stateManager statemanager.Manager) error {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
@@ -196,13 +197,13 @@ func (m *Manager) AllowNetbird() error {
 	}
 
 	_, err := m.AddPeerFiltering(
+		nil,
 		net.IP{0, 0, 0, 0},
 		"all",
 		nil,
 		nil,
 		firewall.ActionAccept,
 		"",
-		"",
 	)
 	if err != nil {
 		return fmt.Errorf("allow netbird interface traffic: %w", err)

client/firewall/iptables/manager_linux_test.go

@@ -75,7 +75,7 @@ func TestIptablesManager(t *testing.T) {
 			IsRange: true,
 			Values:  []uint16{8043, 8046},
 		}
-		rule2, err = manager.AddPeerFiltering(ip, "tcp", port, nil, fw.ActionAccept, "", "accept HTTPS traffic from ports range")
+		rule2, err = manager.AddPeerFiltering(nil, ip, "tcp", port, nil, fw.ActionAccept, "")
 		require.NoError(t, err, "failed to add rule")
 
 		for _, r := range rule2 {
@@ -97,7 +97,7 @@ func TestIptablesManager(t *testing.T) {
 		// add second rule
 		ip := net.ParseIP("10.20.0.3")
 		port := &fw.Port{Values: []uint16{5353}}
-		_, err = manager.AddPeerFiltering(ip, "udp", nil, port, fw.ActionAccept, "", "accept Fake DNS traffic")
+		_, err = manager.AddPeerFiltering(nil, ip, "udp", nil, port, fw.ActionAccept, "")
 		require.NoError(t, err, "failed to add rule")
 
 		err = manager.Close(nil)
@@ -148,7 +148,7 @@ func TestIptablesManagerIPSet(t *testing.T) {
 		port := &fw.Port{
 			Values: []uint16{443},
 		}
-		rule2, err = manager.AddPeerFiltering(ip, "tcp", port, nil, fw.ActionAccept, "default", "accept HTTPS traffic from ports range")
+		rule2, err = manager.AddPeerFiltering(nil, ip, "tcp", port, nil, fw.ActionAccept, "default")
 		for _, r := range rule2 {
 			require.NoError(t, err, "failed to add rule")
 			require.Equal(t, r.(*Rule).ipsetName, "default-sport", "ipset name must be set")
@@ -216,7 +216,7 @@ func TestIptablesCreatePerformance(t *testing.T) {
 			start := time.Now()
 			for i := 0; i < testMax; i++ {
 				port := &fw.Port{Values: []uint16{uint16(1000 + i)}}
-				_, err = manager.AddPeerFiltering(ip, "tcp", nil, port, fw.ActionAccept, "", "accept HTTP traffic")
+				_, err = manager.AddPeerFiltering(nil, ip, "tcp", nil, port, fw.ActionAccept, "")
 
 				require.NoError(t, err, "failed to add rule")
 			}

client/firewall/iptables/router_linux.go

@@ -15,7 +15,7 @@ import (
 
 	nberrors "github.com/netbirdio/netbird/client/errors"
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
-	"github.com/netbirdio/netbird/client/internal/acl/id"
+	nbid "github.com/netbirdio/netbird/client/internal/acl/id"
 	"github.com/netbirdio/netbird/client/internal/routemanager/ipfwdstate"
 	"github.com/netbirdio/netbird/client/internal/routemanager/refcounter"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
@@ -76,7 +76,7 @@ type router struct {
 	wgIface          iFaceMapper
 	legacyManagement bool
 
-	stateManager *statemanager.Manager
+	stateManager statemanager.Manager
 	ipFwdState   *ipfwdstate.IPForwardingState
 }
 
@@ -104,7 +104,7 @@ func newRouter(iptablesClient *iptables.IPTables, wgIface iFaceMapper) (*router,
 	return r, nil
 }
 
-func (r *router) init(stateManager *statemanager.Manager) error {
+func (r *router) init(stateManager statemanager.Manager) error {
 	r.stateManager = stateManager
 
 	if err := r.cleanUpDefaultForwardRules(); err != nil {
@@ -121,6 +121,7 @@ func (r *router) init(stateManager *statemanager.Manager) error {
 }
 
 func (r *router) AddRouteFiltering(
+	id []byte,
 	sources []netip.Prefix,
 	destination netip.Prefix,
 	proto firewall.Protocol,
@@ -128,7 +129,7 @@ func (r *router) AddRouteFiltering(
 	dPort *firewall.Port,
 	action firewall.Action,
 ) (firewall.Rule, error) {
-	ruleKey := id.GenerateRouteRuleKey(sources, destination, proto, sPort, dPort, action)
+	ruleKey := nbid.GenerateRouteRuleKey(sources, destination, proto, sPort, dPort, action)
 	if _, ok := r.rules[string(ruleKey)]; ok {
 		return ruleKey, nil
 	}

client/firewall/iptables/router_linux_test.go

@@ -330,7 +330,7 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			ruleKey, err := r.AddRouteFiltering(tt.sources, tt.destination, tt.proto, tt.sPort, tt.dPort, tt.action)
+			ruleKey, err := r.AddRouteFiltering(nil, tt.sources, tt.destination, tt.proto, tt.sPort, tt.dPort, tt.action)
 			require.NoError(t, err, "AddRouteFiltering failed")
 
 			// Check if the rule is in the internal map

client/firewall/manager/firewall.go

@@ -55,7 +55,7 @@ const (
 // It declares methods which handle actions required by the
 // Netbird client for ACL and routing functionality
 type Manager interface {
-	Init(stateManager *statemanager.Manager) error
+	Init(stateManager statemanager.Manager) error
 
 	// AllowNetbird allows netbird interface traffic
 	AllowNetbird() error
@@ -65,13 +65,13 @@ type Manager interface {
 	// If comment argument is empty firewall manager should set
 	// rule ID as comment for the rule
 	AddPeerFiltering(
+		id []byte,
 		ip net.IP,
 		proto Protocol,
 		sPort *Port,
 		dPort *Port,
 		action Action,
 		ipsetName string,
-		comment string,
 	) ([]Rule, error)
 
 	// DeletePeerRule from the firewall by rule definition
@@ -80,7 +80,15 @@ type Manager interface {
 	// IsServerRouteSupported returns true if the firewall supports server side routing operations
 	IsServerRouteSupported() bool
 
-	AddRouteFiltering(source []netip.Prefix, destination netip.Prefix, proto Protocol, sPort *Port, dPort *Port, action Action) (Rule, error)
+	AddRouteFiltering(
+		id []byte,
+		sources []netip.Prefix,
+		destination netip.Prefix,
+		proto Protocol,
+		sPort *Port,
+		dPort *Port,
+		action Action,
+	) (Rule, error)
 
 	// DeleteRouteRule deletes a routing rule
 	DeleteRouteRule(rule Rule) error
@@ -95,7 +103,7 @@ type Manager interface {
 	SetLegacyManagement(legacy bool) error
 
 	// Close closes the firewall manager
-	Close(stateManager *statemanager.Manager) error
+	Close(stateManager statemanager.Manager) error
 
 	// Flush the changes to firewall controller
 	Flush() error

client/firewall/nftables/acl_linux.go

@@ -84,13 +84,13 @@ func (m *AclManager) init(workTable *nftables.Table) error {
 // If comment argument is empty firewall manager should set
 // rule ID as comment for the rule
 func (m *AclManager) AddPeerFiltering(
+	id []byte,
 	ip net.IP,
 	proto firewall.Protocol,
 	sPort *firewall.Port,
 	dPort *firewall.Port,
 	action firewall.Action,
 	ipsetName string,
-	comment string,
 ) ([]firewall.Rule, error) {
 	var ipset *nftables.Set
 	if ipsetName != "" {
@@ -102,7 +102,7 @@ func (m *AclManager) AddPeerFiltering(
 	}
 
 	newRules := make([]firewall.Rule, 0, 2)
-	ioRule, err := m.addIOFiltering(ip, proto, sPort, dPort, action, ipset, comment)
+	ioRule, err := m.addIOFiltering(ip, proto, sPort, dPort, action, ipset)
 	if err != nil {
 		return nil, err
 	}
@@ -256,7 +256,6 @@ func (m *AclManager) addIOFiltering(
 	dPort *firewall.Port,
 	action firewall.Action,
 	ipset *nftables.Set,
-	comment string,
 ) (*Rule, error) {
 	ruleId := generatePeerRuleId(ip, sPort, dPort, action, ipset)
 	if r, ok := m.rules[ruleId]; ok {
@@ -338,7 +337,7 @@ func (m *AclManager) addIOFiltering(
 		mainExpressions = append(mainExpressions, &expr.Verdict{Kind: expr.VerdictDrop})
 	}
 
-	userData := []byte(strings.Join([]string{ruleId, comment}, " "))
+	userData := []byte(ruleId)
 
 	chain := m.chainInputRules
 	nftRule := m.rConn.AddRule(&nftables.Rule{

client/firewall/nftables/manager_linux.go

@@ -67,7 +67,7 @@ func Create(wgIface iFaceMapper) (*Manager, error) {
 }
 
 // Init nftables firewall manager
-func (m *Manager) Init(stateManager *statemanager.Manager) error {
+func (m *Manager) Init(stateManager statemanager.Manager) error {
 	workTable, err := m.createWorkTable()
 	if err != nil {
 		return fmt.Errorf("create work table: %w", err)
@@ -113,13 +113,13 @@ func (m *Manager) Init(stateManager *statemanager.Manager) error {
 // If comment argument is empty firewall manager should set
 // rule ID as comment for the rule
 func (m *Manager) AddPeerFiltering(
+	id []byte,
 	ip net.IP,
 	proto firewall.Protocol,
 	sPort *firewall.Port,
 	dPort *firewall.Port,
 	action firewall.Action,
 	ipsetName string,
-	comment string,
 ) ([]firewall.Rule, error) {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
@@ -129,10 +129,11 @@ func (m *Manager) AddPeerFiltering(
 		return nil, fmt.Errorf("unsupported IP version: %s", ip.String())
 	}
 
-	return m.aclManager.AddPeerFiltering(ip, proto, sPort, dPort, action, ipsetName, comment)
+	return m.aclManager.AddPeerFiltering(id, ip, proto, sPort, dPort, action, ipsetName)
 }
 
 func (m *Manager) AddRouteFiltering(
+	id []byte,
 	sources []netip.Prefix,
 	destination netip.Prefix,
 	proto firewall.Protocol,
@@ -147,7 +148,7 @@ func (m *Manager) AddRouteFiltering(
 		return nil, fmt.Errorf("unsupported IP version: %s", destination.Addr().String())
 	}
 
-	return m.router.AddRouteFiltering(sources, destination, proto, sPort, dPort, action)
+	return m.router.AddRouteFiltering(id, sources, destination, proto, sPort, dPort, action)
 }
 
 // DeletePeerRule from the firewall by rule definition
@@ -242,7 +243,7 @@ func (m *Manager) SetLegacyManagement(isLegacy bool) error {
 }
 
 // Reset firewall to the default state
-func (m *Manager) Close(stateManager *statemanager.Manager) error {
+func (m *Manager) Close(stateManager statemanager.Manager) error {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 

client/firewall/nftables/manager_linux_test.go

@@ -74,7 +74,7 @@ func TestNftablesManager(t *testing.T) {
 
 	testClient := &nftables.Conn{}
 
-	rule, err := manager.AddPeerFiltering(ip, fw.ProtocolTCP, nil, &fw.Port{Values: []uint16{53}}, fw.ActionDrop, "", "")
+	rule, err := manager.AddPeerFiltering(nil, ip, fw.ProtocolTCP, nil, &fw.Port{Values: []uint16{53}}, fw.ActionDrop, "")
 	require.NoError(t, err, "failed to add rule")
 
 	err = manager.Flush()
@@ -201,7 +201,7 @@ func TestNFtablesCreatePerformance(t *testing.T) {
 			start := time.Now()
 			for i := 0; i < testMax; i++ {
 				port := &fw.Port{Values: []uint16{uint16(1000 + i)}}
-				_, err = manager.AddPeerFiltering(ip, "tcp", nil, port, fw.ActionAccept, "", "accept HTTP traffic")
+				_, err = manager.AddPeerFiltering(nil, ip, "tcp", nil, port, fw.ActionAccept, "")
 				require.NoError(t, err, "failed to add rule")
 
 				if i%100 == 0 {
@@ -283,10 +283,11 @@ func TestNftablesManagerCompatibilityWithIptables(t *testing.T) {
 	})
 
 	ip := net.ParseIP("100.96.0.1")
-	_, err = manager.AddPeerFiltering(ip, fw.ProtocolTCP, nil, &fw.Port{Values: []uint16{80}}, fw.ActionAccept, "", "test rule")
+	_, err = manager.AddPeerFiltering(nil, ip, fw.ProtocolTCP, nil, &fw.Port{Values: []uint16{80}}, fw.ActionAccept, "")
 	require.NoError(t, err, "failed to add peer filtering rule")
 
 	_, err = manager.AddRouteFiltering(
+		nil,
 		[]netip.Prefix{netip.MustParsePrefix("192.168.2.0/24")},
 		netip.MustParsePrefix("10.1.0.0/24"),
 		fw.ProtocolTCP,

client/firewall/nftables/router_linux.go

@@ -20,7 +20,7 @@ import (
 
 	nberrors "github.com/netbirdio/netbird/client/errors"
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
-	"github.com/netbirdio/netbird/client/internal/acl/id"
+	nbid "github.com/netbirdio/netbird/client/internal/acl/id"
 	"github.com/netbirdio/netbird/client/internal/routemanager/ipfwdstate"
 	"github.com/netbirdio/netbird/client/internal/routemanager/refcounter"
 	nbnet "github.com/netbirdio/netbird/util/net"
@@ -228,6 +228,7 @@ func (r *router) createContainers() error {
 
 // AddRouteFiltering appends a nftables rule to the routing chain
 func (r *router) AddRouteFiltering(
+	id []byte,
 	sources []netip.Prefix,
 	destination netip.Prefix,
 	proto firewall.Protocol,
@@ -236,7 +237,7 @@ func (r *router) AddRouteFiltering(
 	action firewall.Action,
 ) (firewall.Rule, error) {
 
-	ruleKey := id.GenerateRouteRuleKey(sources, destination, proto, sPort, dPort, action)
+	ruleKey := nbid.GenerateRouteRuleKey(sources, destination, proto, sPort, dPort, action)
 	if _, ok := r.rules[string(ruleKey)]; ok {
 		return ruleKey, nil
 	}

client/firewall/nftables/router_linux_test.go

@@ -311,7 +311,7 @@ func TestRouter_AddRouteFiltering(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			ruleKey, err := r.AddRouteFiltering(tt.sources, tt.destination, tt.proto, tt.sPort, tt.dPort, tt.action)
+			ruleKey, err := r.AddRouteFiltering(nil, tt.sources, tt.destination, tt.proto, tt.sPort, tt.dPort, tt.action)
 			require.NoError(t, err, "AddRouteFiltering failed")
 
 			t.Cleanup(func() {

client/firewall/uspfilter/allow_netbird.go

@@ -4,6 +4,7 @@ package uspfilter
 
 import (
 	"context"
+	"net/netip"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -12,12 +13,12 @@ import (
 )
 
 // Reset firewall to the default state
-func (m *Manager) Close(stateManager *statemanager.Manager) error {
+func (m *Manager) Close(stateManager statemanager.Manager) error {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	m.outgoingRules = make(map[string]RuleSet)
-	m.incomingRules = make(map[string]RuleSet)
+	m.outgoingRules = make(map[netip.Addr]RuleSet)
+	m.incomingRules = make(map[netip.Addr]RuleSet)
 
 	if m.udpTracker != nil {
 		m.udpTracker.Close()
@@ -31,8 +32,8 @@ func (m *Manager) Close(stateManager *statemanager.Manager) error {
 		m.tcpTracker.Close()
 	}
 
-	if m.forwarder != nil {
-		m.forwarder.Stop()
+	if fwder := m.forwarder.Load(); fwder != nil {
+		fwder.Stop()
 	}
 
 	if m.logger != nil {

client/firewall/uspfilter/allow_netbird_windows.go

@@ -3,12 +3,14 @@ package uspfilter
 import (
 	"context"
 	"fmt"
+	"net/netip"
 	"os/exec"
 	"syscall"
 	"time"
 
 	log "github.com/sirupsen/logrus"
 
+	"github.com/netbirdio/netbird/client/firewall/uspfilter/conntrack"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 )
 
@@ -20,28 +22,31 @@ const (
 	firewallRuleName        = "Netbird"
 )
 
-// Close closes the firewall manager
-func (m *Manager) Close(*statemanager.Manager) error {
+// Reset firewall to the default state
+func (m *Manager) Close(statemanager.Manager) error {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	m.outgoingRules = make(map[string]RuleSet)
-	m.incomingRules = make(map[string]RuleSet)
+	m.outgoingRules = make(map[netip.Addr]RuleSet)
+	m.incomingRules = make(map[netip.Addr]RuleSet)
 
 	if m.udpTracker != nil {
 		m.udpTracker.Close()
+		m.udpTracker = conntrack.NewUDPTracker(conntrack.DefaultUDPTimeout, m.logger, m.flowLogger)
 	}
 
 	if m.icmpTracker != nil {
 		m.icmpTracker.Close()
+		m.icmpTracker = conntrack.NewICMPTracker(conntrack.DefaultICMPTimeout, m.logger, m.flowLogger)
 	}
 
 	if m.tcpTracker != nil {
 		m.tcpTracker.Close()
+		m.tcpTracker = conntrack.NewTCPTracker(conntrack.DefaultTCPTimeout, m.logger, m.flowLogger)
 	}
 
-	if m.forwarder != nil {
-		m.forwarder.Stop()
+	if fwder := m.forwarder.Load(); fwder != nil {
+		fwder.Stop()
 	}
 
 	if m.logger != nil {

client/firewall/uspfilter/conntrack/common.go

@@ -1,20 +1,27 @@
-// common.go
 package conntrack
 
 import (
-	"net"
-	"sync"
+	"fmt"
+	"net/netip"
 	"sync/atomic"
 	"time"
+
+	"github.com/google/uuid"
+
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 )
 
 // BaseConnTrack provides common fields and locking for all connection types
 type BaseConnTrack struct {
-	SourceIP   net.IP
-	DestIP     net.IP
-	SourcePort uint16
-	DestPort   uint16
-	lastSeen   atomic.Int64 // Unix nano for atomic access
+	FlowId    uuid.UUID
+	Direction nftypes.Direction
+	SourceIP  netip.Addr
+	DestIP    netip.Addr
+	lastSeen  atomic.Int64
+	PacketsTx atomic.Uint64
+	PacketsRx atomic.Uint64
+	BytesTx   atomic.Uint64
+	BytesRx   atomic.Uint64
 }
 
 // these small methods will be inlined by the compiler
@@ -24,6 +31,17 @@ func (b *BaseConnTrack) UpdateLastSeen() {
 	b.lastSeen.Store(time.Now().UnixNano())
 }
 
+// UpdateCounters safely updates the packet and byte counters
+func (b *BaseConnTrack) UpdateCounters(direction nftypes.Direction, bytes int) {
+	if direction == nftypes.Egress {
+		b.PacketsTx.Add(1)
+		b.BytesTx.Add(uint64(bytes))
+	} else {
+		b.PacketsRx.Add(1)
+		b.BytesRx.Add(uint64(bytes))
+	}
+}
+
 // GetLastSeen safely gets the last seen timestamp
 func (b *BaseConnTrack) GetLastSeen() time.Time {
 	return time.Unix(0, b.lastSeen.Load())
@@ -35,92 +53,14 @@ func (b *BaseConnTrack) timeoutExceeded(timeout time.Duration) bool {
 	return time.Since(lastSeen) > timeout
 }
 
-// IPAddr is a fixed-size IP address to avoid allocations
-type IPAddr [16]byte
-
-// MakeIPAddr creates an IPAddr from net.IP
-func MakeIPAddr(ip net.IP) (addr IPAddr) {
-	// Optimization: check for v4 first as it's more common
-	if ip4 := ip.To4(); ip4 != nil {
-		copy(addr[12:], ip4)
-	} else {
-		copy(addr[:], ip.To16())
-	}
-	return addr
-}
-
 // ConnKey uniquely identifies a connection
 type ConnKey struct {
-	SrcIP   IPAddr
-	DstIP   IPAddr
+	SrcIP   netip.Addr
+	DstIP   netip.Addr
 	SrcPort uint16
 	DstPort uint16
 }
 
-// makeConnKey creates a connection key
-func makeConnKey(srcIP net.IP, dstIP net.IP, srcPort uint16, dstPort uint16) ConnKey {
-	return ConnKey{
-		SrcIP:   MakeIPAddr(srcIP),
-		DstIP:   MakeIPAddr(dstIP),
-		SrcPort: srcPort,
-		DstPort: dstPort,
-	}
-}
-
-// ValidateIPs checks if IPs match without allocation
-func ValidateIPs(connIP IPAddr, pktIP net.IP) bool {
-	if ip4 := pktIP.To4(); ip4 != nil {
-		// Compare IPv4 addresses (last 4 bytes)
-		for i := 0; i < 4; i++ {
-			if connIP[12+i] != ip4[i] {
-				return false
-			}
-		}
-		return true
-	}
-	// Compare full IPv6 addresses
-	ip6 := pktIP.To16()
-	for i := 0; i < 16; i++ {
-		if connIP[i] != ip6[i] {
-			return false
-		}
-	}
-	return true
-}
-
-// PreallocatedIPs is a pool of IP byte slices to reduce allocations
-type PreallocatedIPs struct {
-	sync.Pool
-}
-
-// NewPreallocatedIPs creates a new IP pool
-func NewPreallocatedIPs() *PreallocatedIPs {
-	return &PreallocatedIPs{
-		Pool: sync.Pool{
-			New: func() interface{} {
-				ip := make(net.IP, 16)
-				return &ip
-			},
-		},
-	}
-}
-
-// Get retrieves an IP from the pool
-func (p *PreallocatedIPs) Get() net.IP {
-	return *p.Pool.Get().(*net.IP)
-}
-
-// Put returns an IP to the pool
-func (p *PreallocatedIPs) Put(ip net.IP) {
-	p.Pool.Put(&ip)
-}
-
-// copyIP copies an IP address efficiently
-func copyIP(dst, src net.IP) {
-	if len(src) == 16 {
-		copy(dst, src)
-	} else {
-		// Handle IPv4
-		copy(dst[12:], src.To4())
-	}
+func (c ConnKey) String() string {
+	return fmt.Sprintf("%s:%d -> %s:%d", c.SrcIP.Unmap(), c.SrcPort, c.DstIP.Unmap(), c.DstPort)
 }

client/firewall/uspfilter/conntrack/common_test.go

@@ -1,94 +1,67 @@
 package conntrack
 
 import (
-	"net"
+	"context"
+	"net/netip"
 	"testing"
 
 	"github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/client/firewall/uspfilter/log"
+	"github.com/netbirdio/netbird/client/internal/netflow"
 )
 
 var logger = log.NewFromLogrus(logrus.StandardLogger())
-
-func BenchmarkIPOperations(b *testing.B) {
-	b.Run("MakeIPAddr", func(b *testing.B) {
-		ip := net.ParseIP("192.168.1.1")
-		b.ResetTimer()
-		for i := 0; i < b.N; i++ {
-			_ = MakeIPAddr(ip)
-		}
-	})
-
-	b.Run("ValidateIPs", func(b *testing.B) {
-		ip1 := net.ParseIP("192.168.1.1")
-		ip2 := net.ParseIP("192.168.1.1")
-		addr := MakeIPAddr(ip1)
-		b.ResetTimer()
-		for i := 0; i < b.N; i++ {
-			_ = ValidateIPs(addr, ip2)
-		}
-	})
-
-	b.Run("IPPool", func(b *testing.B) {
-		pool := NewPreallocatedIPs()
-		b.ResetTimer()
-		for i := 0; i < b.N; i++ {
-			ip := pool.Get()
-			pool.Put(ip)
-		}
-	})
-
-}
+var flowLogger = netflow.NewManager(context.Background(), nil, []byte{}, nil).GetLogger()
 
 // Memory pressure tests
 func BenchmarkMemoryPressure(b *testing.B) {
 	b.Run("TCPHighLoad", func(b *testing.B) {
-		tracker := NewTCPTracker(DefaultTCPTimeout, logger)
+		tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
 		defer tracker.Close()
 
 		// Generate different IPs
-		srcIPs := make([]net.IP, 100)
-		dstIPs := make([]net.IP, 100)
+		srcIPs := make([]netip.Addr, 100)
+		dstIPs := make([]netip.Addr, 100)
 		for i := 0; i < 100; i++ {
-			srcIPs[i] = net.IPv4(192, 168, byte(i/256), byte(i%256))
-			dstIPs[i] = net.IPv4(10, 0, byte(i/256), byte(i%256))
+			srcIPs[i] = netip.AddrFrom4([4]byte{192, 168, byte(i / 256), byte(i % 256)})
+			dstIPs[i] = netip.AddrFrom4([4]byte{10, 0, byte(i / 256), byte(i % 256)})
 		}
 
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
 			srcIdx := i % len(srcIPs)
 			dstIdx := (i + 1) % len(dstIPs)
-			tracker.TrackOutbound(srcIPs[srcIdx], dstIPs[dstIdx], uint16(i%65535), 80, TCPSyn)
+			tracker.TrackOutbound(srcIPs[srcIdx], dstIPs[dstIdx], uint16(i%65535), 80, TCPSyn, 0)
 
 			// Simulate some valid inbound packets
 			if i%3 == 0 {
-				tracker.IsValidInbound(dstIPs[dstIdx], srcIPs[srcIdx], 80, uint16(i%65535), TCPAck)
+				tracker.IsValidInbound(dstIPs[dstIdx], srcIPs[srcIdx], 80, uint16(i%65535), TCPAck, 0)
 			}
 		}
 	})
 
 	b.Run("UDPHighLoad", func(b *testing.B) {
-		tracker := NewUDPTracker(DefaultUDPTimeout, logger)
+		tracker := NewUDPTracker(DefaultUDPTimeout, logger, flowLogger)
 		defer tracker.Close()
 
 		// Generate different IPs
-		srcIPs := make([]net.IP, 100)
-		dstIPs := make([]net.IP, 100)
+		srcIPs := make([]netip.Addr, 100)
+		dstIPs := make([]netip.Addr, 100)
 		for i := 0; i < 100; i++ {
-			srcIPs[i] = net.IPv4(192, 168, byte(i/256), byte(i%256))
-			dstIPs[i] = net.IPv4(10, 0, byte(i/256), byte(i%256))
+			srcIPs[i] = netip.AddrFrom4([4]byte{192, 168, byte(i / 256), byte(i % 256)})
+			dstIPs[i] = netip.AddrFrom4([4]byte{10, 0, byte(i / 256), byte(i % 256)})
 		}
 
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
 			srcIdx := i % len(srcIPs)
 			dstIdx := (i + 1) % len(dstIPs)
-			tracker.TrackOutbound(srcIPs[srcIdx], dstIPs[dstIdx], uint16(i%65535), 80)
+			tracker.TrackOutbound(srcIPs[srcIdx], dstIPs[dstIdx], uint16(i%65535), 80, 0)
 
 			// Simulate some valid inbound packets
 			if i%3 == 0 {
-				tracker.IsValidInbound(dstIPs[dstIdx], srcIPs[srcIdx], 80, uint16(i%65535))
+				tracker.IsValidInbound(dstIPs[dstIdx], srcIPs[srcIdx], 80, uint16(i%65535), 0)
 			}
 		}
 	})

client/firewall/uspfilter/conntrack/icmp.go

@@ -2,13 +2,16 @@ package conntrack
 
 import (
 	"context"
-	"net"
+	"fmt"
+	"net/netip"
 	"sync"
 	"time"
 
 	"github.com/google/gopacket/layers"
+	"github.com/google/uuid"
 
 	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 )
 
 const (
@@ -20,18 +23,20 @@ const (
 
 // ICMPConnKey uniquely identifies an ICMP connection
 type ICMPConnKey struct {
-	// Supports both IPv4 and IPv6
-	SrcIP    [16]byte
-	DstIP    [16]byte
-	Sequence uint16 // ICMP sequence number
-	ID       uint16 // ICMP identifier
+	SrcIP netip.Addr
+	DstIP netip.Addr
+	ID    uint16
+}
+
+func (i ICMPConnKey) String() string {
+	return fmt.Sprintf("%s -> %s (id %d)", i.SrcIP, i.DstIP, i.ID)
 }
 
 // ICMPConnTrack represents an ICMP connection state
 type ICMPConnTrack struct {
 	BaseConnTrack
-	Sequence uint16
-	ID       uint16
+	ICMPType uint8
+	ICMPCode uint8
 }
 
 // ICMPTracker manages ICMP connection states
@@ -42,11 +47,11 @@ type ICMPTracker struct {
 	cleanupTicker *time.Ticker
 	tickerCancel  context.CancelFunc
 	mutex         sync.RWMutex
-	ipPool        *PreallocatedIPs
+	flowLogger    nftypes.FlowLogger
 }
 
 // NewICMPTracker creates a new ICMP connection tracker
-func NewICMPTracker(timeout time.Duration, logger *nblog.Logger) *ICMPTracker {
+func NewICMPTracker(timeout time.Duration, logger *nblog.Logger, flowLogger nftypes.FlowLogger) *ICMPTracker {
 	if timeout == 0 {
 		timeout = DefaultICMPTimeout
 	}
@@ -59,67 +64,107 @@ func NewICMPTracker(timeout time.Duration, logger *nblog.Logger) *ICMPTracker {
 		timeout:       timeout,
 		cleanupTicker: time.NewTicker(ICMPCleanupInterval),
 		tickerCancel:  cancel,
-		ipPool:        NewPreallocatedIPs(),
+		flowLogger:    flowLogger,
 	}
 
 	go tracker.cleanupRoutine(ctx)
 	return tracker
 }
 
-// TrackOutbound records an outbound ICMP Echo Request
-func (t *ICMPTracker) TrackOutbound(srcIP net.IP, dstIP net.IP, id uint16, seq uint16) {
-	key := makeICMPKey(srcIP, dstIP, id, seq)
-
-	t.mutex.Lock()
-	conn, exists := t.connections[key]
-	if !exists {
-		srcIPCopy := t.ipPool.Get()
-		dstIPCopy := t.ipPool.Get()
-		copyIP(srcIPCopy, srcIP)
-		copyIP(dstIPCopy, dstIP)
-
-		conn = &ICMPConnTrack{
-			BaseConnTrack: BaseConnTrack{
-				SourceIP: srcIPCopy,
-				DestIP:   dstIPCopy,
-			},
-			ID:       id,
-			Sequence: seq,
-		}
-		conn.UpdateLastSeen()
-		t.connections[key] = conn
-
-		t.logger.Trace("New ICMP connection %v", key)
+func (t *ICMPTracker) updateIfExists(srcIP netip.Addr, dstIP netip.Addr, id uint16, direction nftypes.Direction, size int) (ICMPConnKey, bool) {
+	key := ICMPConnKey{
+		SrcIP: srcIP,
+		DstIP: dstIP,
+		ID:    id,
 	}
-	t.mutex.Unlock()
-
-	conn.UpdateLastSeen()
-}
-
-// IsValidInbound checks if an inbound ICMP Echo Reply matches a tracked request
-func (t *ICMPTracker) IsValidInbound(srcIP net.IP, dstIP net.IP, id uint16, seq uint16, icmpType uint8) bool {
-	if icmpType != uint8(layers.ICMPv4TypeEchoReply) {
-		return false
-	}
-
-	key := makeICMPKey(dstIP, srcIP, id, seq)
 
 	t.mutex.RLock()
 	conn, exists := t.connections[key]
 	t.mutex.RUnlock()
 
-	if !exists {
+	if exists {
+		conn.UpdateLastSeen()
+		conn.UpdateCounters(direction, size)
+
+		return key, true
+	}
+
+	return key, false
+}
+
+// TrackOutbound records an outbound ICMP connection
+func (t *ICMPTracker) TrackOutbound(srcIP netip.Addr, dstIP netip.Addr, id uint16, typecode layers.ICMPv4TypeCode, size int) {
+	if _, exists := t.updateIfExists(dstIP, srcIP, id, nftypes.Egress, size); !exists {
+		// if (inverted direction) conn is not tracked, track this direction
+		t.track(srcIP, dstIP, id, typecode, nftypes.Egress, nil, size)
+	}
+}
+
+// TrackInbound records an inbound ICMP Echo Request
+func (t *ICMPTracker) TrackInbound(srcIP netip.Addr, dstIP netip.Addr, id uint16, typecode layers.ICMPv4TypeCode, ruleId []byte, size int) {
+	t.track(srcIP, dstIP, id, typecode, nftypes.Ingress, ruleId, size)
+}
+
+// track is the common implementation for tracking both inbound and outbound ICMP connections
+func (t *ICMPTracker) track(srcIP netip.Addr, dstIP netip.Addr, id uint16, typecode layers.ICMPv4TypeCode, direction nftypes.Direction, ruleId []byte, size int) {
+	key, exists := t.updateIfExists(srcIP, dstIP, id, direction, size)
+	if exists {
+		return
+	}
+
+	typ, code := typecode.Type(), typecode.Code()
+
+	// non echo requests don't need tracking
+	if typ != uint8(layers.ICMPv4TypeEchoRequest) {
+		t.logger.Trace("New %s ICMP connection %s type %d code %d", direction, key, typ, code)
+		t.sendStartEvent(direction, srcIP, dstIP, typ, code, ruleId, size)
+		return
+	}
+
+	conn := &ICMPConnTrack{
+		BaseConnTrack: BaseConnTrack{
+			FlowId:    uuid.New(),
+			Direction: direction,
+			SourceIP:  srcIP,
+			DestIP:    dstIP,
+		},
+		ICMPType: typ,
+		ICMPCode: code,
+	}
+	conn.UpdateLastSeen()
+
+	t.mutex.Lock()
+	t.connections[key] = conn
+	t.mutex.Unlock()
+
+	t.logger.Trace("New %s ICMP connection %s type %d code %d", direction, key, typ, code)
+	t.sendEvent(nftypes.TypeStart, conn, ruleId)
+}
+
+// IsValidInbound checks if an inbound ICMP Echo Reply matches a tracked request
+func (t *ICMPTracker) IsValidInbound(srcIP netip.Addr, dstIP netip.Addr, id uint16, icmpType uint8, size int) bool {
+	if icmpType != uint8(layers.ICMPv4TypeEchoReply) {
 		return false
 	}
 
-	if conn.timeoutExceeded(t.timeout) {
+	key := ICMPConnKey{
+		SrcIP: dstIP,
+		DstIP: srcIP,
+		ID:    id,
+	}
+
+	t.mutex.RLock()
+	conn, exists := t.connections[key]
+	t.mutex.RUnlock()
+
+	if !exists || conn.timeoutExceeded(t.timeout) {
 		return false
 	}
 
-	return ValidateIPs(MakeIPAddr(srcIP), conn.DestIP) &&
-		ValidateIPs(MakeIPAddr(dstIP), conn.SourceIP) &&
-		conn.ID == id &&
-		conn.Sequence == seq
+	conn.UpdateLastSeen()
+	conn.UpdateCounters(nftypes.Ingress, size)
+
+	return true
 }
 
 func (t *ICMPTracker) cleanupRoutine(ctx context.Context) {
@@ -134,17 +179,18 @@ func (t *ICMPTracker) cleanupRoutine(ctx context.Context) {
 		}
 	}
 }
+
 func (t *ICMPTracker) cleanup() {
 	t.mutex.Lock()
 	defer t.mutex.Unlock()
 
 	for key, conn := range t.connections {
 		if conn.timeoutExceeded(t.timeout) {
-			t.ipPool.Put(conn.SourceIP)
-			t.ipPool.Put(conn.DestIP)
 			delete(t.connections, key)
 
-			t.logger.Debug("Removed ICMP connection %v (timeout)", key)
+			t.logger.Debug("Removed ICMP connection %s (timeout) [in: %d Pkts/%d B out: %d Pkts/%d B]",
+				key, conn.PacketsRx.Load(), conn.BytesRx.Load(), conn.PacketsTx.Load(), conn.BytesTx.Load())
+			t.sendEvent(nftypes.TypeEnd, conn, nil)
 		}
 	}
 }
@@ -154,20 +200,46 @@ func (t *ICMPTracker) Close() {
 	t.tickerCancel()
 
 	t.mutex.Lock()
-	for _, conn := range t.connections {
-		t.ipPool.Put(conn.SourceIP)
-		t.ipPool.Put(conn.DestIP)
-	}
 	t.connections = nil
 	t.mutex.Unlock()
 }
 
-// makeICMPKey creates an ICMP connection key
-func makeICMPKey(srcIP net.IP, dstIP net.IP, id uint16, seq uint16) ICMPConnKey {
-	return ICMPConnKey{
-		SrcIP:    MakeIPAddr(srcIP),
-		DstIP:    MakeIPAddr(dstIP),
-		ID:       id,
-		Sequence: seq,
-	}
+func (t *ICMPTracker) sendEvent(typ nftypes.Type, conn *ICMPConnTrack, ruleID []byte) {
+	t.flowLogger.StoreEvent(nftypes.EventFields{
+		FlowID:    conn.FlowId,
+		Type:      typ,
+		RuleID:    ruleID,
+		Direction: conn.Direction,
+		Protocol:  nftypes.ICMP, // TODO: adjust for IPv6/icmpv6
+		SourceIP:  conn.SourceIP,
+		DestIP:    conn.DestIP,
+		ICMPType:  conn.ICMPType,
+		ICMPCode:  conn.ICMPCode,
+		RxPackets: conn.PacketsRx.Load(),
+		TxPackets: conn.PacketsTx.Load(),
+		RxBytes:   conn.BytesRx.Load(),
+		TxBytes:   conn.BytesTx.Load(),
+	})
+}
+
+func (t *ICMPTracker) sendStartEvent(direction nftypes.Direction, srcIP netip.Addr, dstIP netip.Addr, typ uint8, code uint8, ruleID []byte, size int) {
+	fields := nftypes.EventFields{
+		FlowID:    uuid.New(),
+		Type:      nftypes.TypeStart,
+		RuleID:    ruleID,
+		Direction: direction,
+		Protocol:  nftypes.ICMP,
+		SourceIP:  srcIP,
+		DestIP:    dstIP,
+		ICMPType:  typ,
+		ICMPCode:  code,
+	}
+	if direction == nftypes.Ingress {
+		fields.RxPackets = 1
+		fields.RxBytes = uint64(size)
+	} else {
+		fields.TxPackets = 1
+		fields.TxBytes = uint64(size)
+	}
+	t.flowLogger.StoreEvent(fields)
 }

client/firewall/uspfilter/conntrack/icmp_test.go

@@ -1,39 +1,211 @@
 package conntrack
 
 import (
-	"net"
+	"context"
+	"net/netip"
 	"testing"
+	"time"
+
+	"github.com/google/gopacket/layers"
+	"github.com/stretchr/testify/require"
+	// Assume these come from your internal packages.
 )
 
+func TestICMPTracker_TrackOutbound_NonEcho(t *testing.T) {
+	// Use a non-echo type (for example, 3 = Destination Unreachable)
+	nonEchoTypeCode := layers.CreateICMPv4TypeCode(3, 0)
+
+	flowLogger.Enable()
+	defer flowLogger.Disable()
+
+	// Use a reasonable timeout.
+	tracker := NewICMPTracker(30*time.Second, logger, flowLogger)
+	defer tracker.Close()
+
+	localIP := netip.MustParseAddr("192.0.2.1")
+	remoteIP := netip.MustParseAddr("192.0.2.2")
+	id := uint16(1000)
+	size := 120
+
+	// For outbound, the function first checks for an inverse connection.
+	// Since none exists, TrackOutbound will call track().
+	tracker.TrackOutbound(localIP, remoteIP, id, nonEchoTypeCode, size)
+
+	// Since type != EchoRequest the connection should not be stored.
+	require.Equal(t, 0, len(tracker.connections), "Non-echo request should not be tracked")
+}
+
+func TestICMPTracker_TrackOutbound_Echo(t *testing.T) {
+	// Use EchoRequest type.
+	echoTypeCode := layers.CreateICMPv4TypeCode(8, 0)
+
+	flowLogger.Enable()
+	defer flowLogger.Disable()
+
+	tracker := NewICMPTracker(30*time.Second, logger, flowLogger)
+	defer tracker.Close()
+
+	localIP := netip.MustParseAddr("192.0.2.10")
+	remoteIP := netip.MustParseAddr("192.0.2.20")
+	id := uint16(2000)
+	size := 150
+
+	// This call should track the connection since it is an echo request.
+	tracker.TrackOutbound(localIP, remoteIP, id, echoTypeCode, size)
+
+	// The connection key is formed with (srcIP, dstIP, id).
+	key := ICMPConnKey{SrcIP: localIP, DstIP: remoteIP, ID: id}
+	tracker.mutex.RLock()
+	_, exists := tracker.connections[key]
+	tracker.mutex.RUnlock()
+	require.True(t, exists, "Echo request should be tracked as a connection")
+
+}
+
+func TestICMPTracker_TrackInbound(t *testing.T) {
+	// For inbound, we pass a rule ID.
+	echoTypeCode := layers.CreateICMPv4TypeCode(8, 0)
+	ruleID := []byte("rule-123")
+
+	tracker := NewICMPTracker(30*time.Second, logger, flowLogger)
+	defer tracker.Close()
+
+	// Here srcIP is the remote host and dstIP is local.
+	remoteIP := netip.MustParseAddr("203.0.113.5")
+	localIP := netip.MustParseAddr("203.0.113.10")
+	id := uint16(3000)
+	size := 180
+
+	tracker.TrackInbound(remoteIP, localIP, id, echoTypeCode, ruleID, size)
+
+	// The connection key for inbound echo request is (srcIP, dstIP, id).
+	key := ICMPConnKey{SrcIP: remoteIP, DstIP: localIP, ID: id}
+	tracker.mutex.RLock()
+	_, exists := tracker.connections[key]
+	tracker.mutex.RUnlock()
+	require.True(t, exists, "Inbound echo request should be tracked")
+}
+
+func TestICMPTracker_IsValidInbound(t *testing.T) {
+	// For a valid echo reply, the tracker must have stored the echo request.
+	echoRequest := layers.CreateICMPv4TypeCode(8, 0)
+
+	// Use a slightly short timeout for testing expiry.
+	tracker := NewICMPTracker(1*time.Second, logger, flowLogger)
+	defer tracker.Close()
+
+	localIP := netip.MustParseAddr("10.0.0.1")
+	remoteIP := netip.MustParseAddr("10.0.0.2")
+	id := uint16(4000)
+	size := 100
+
+	// Initiate the echo request.
+	tracker.TrackOutbound(localIP, remoteIP, id, echoRequest, size)
+
+	// For an echo reply, the src and dst are swapped relative to the request.
+	valid := tracker.IsValidInbound(remoteIP, localIP, id, uint8(layers.ICMPv4TypeEchoReply), size)
+	require.True(t, valid, "Valid echo reply should return true")
+
+	// Test with a wrong ICMP type (not echo reply).
+	invalid := tracker.IsValidInbound(remoteIP, localIP, id, 99, size)
+	require.False(t, invalid, "Invalid echo type should return false")
+
+	// Let the connection expire.
+	time.Sleep(1100 * time.Millisecond)
+	expired := tracker.IsValidInbound(remoteIP, localIP, id, uint8(layers.ICMPv4TypeEchoReply), size)
+	require.False(t, expired, "Expired connection should return false")
+}
+
+func TestICMPTracker_cleanup(t *testing.T) {
+	// Use a very short timeout to force cleanup.
+	echoRequest := layers.CreateICMPv4TypeCode(8, 0)
+
+	tracker := NewICMPTracker(50*time.Millisecond, logger, flowLogger)
+	defer tracker.Close()
+
+	localIP := netip.MustParseAddr("172.16.0.1")
+	remoteIP := netip.MustParseAddr("172.16.0.2")
+	id := uint16(5000)
+	size := 100
+
+	tracker.TrackOutbound(localIP, remoteIP, id, echoRequest, size)
+	key := ICMPConnKey{SrcIP: localIP, DstIP: remoteIP, ID: id}
+
+	// Confirm the connection is present.
+	tracker.mutex.RLock()
+	_, exists := tracker.connections[key]
+	tracker.mutex.RUnlock()
+	require.True(t, exists, "Connection should exist before cleanup")
+
+	// Wait for the timeout to expire.
+	time.Sleep(100 * time.Millisecond)
+
+	// Manually trigger cleanup.
+	tracker.cleanup()
+
+	tracker.mutex.RLock()
+	_, exists = tracker.connections[key]
+	tracker.mutex.RUnlock()
+	require.False(t, exists, "Expired connection should have been cleaned up")
+}
+
+func TestICMPTracker_Close(t *testing.T) {
+	echoRequest := layers.CreateICMPv4TypeCode(8, 0)
+
+	tracker := NewICMPTracker(30*time.Second, logger, flowLogger)
+
+	// Add a connection.
+	localIP := netip.MustParseAddr("198.51.100.1")
+	remoteIP := netip.MustParseAddr("198.51.100.2")
+	id := uint16(6000)
+	size := 100
+	tracker.TrackOutbound(localIP, remoteIP, id, echoRequest, size)
+
+	// Close the tracker.
+	tracker.Close()
+
+	// After Close the connections map should be nil.
+	tracker.mutex.RLock()
+	require.Nil(t, tracker.connections, "Connections map should be nil after Close")
+	tracker.mutex.RUnlock()
+
+	// The cleanup goroutine should also be stopped. Canceling the ticker context should end cleanupRoutine.
+	select {
+	case <-time.After(50 * time.Millisecond):
+		// no panic or deadlock indicates Close worked correctly.
+	case <-context.Background().Done():
+	}
+}
+
 func BenchmarkICMPTracker(b *testing.B) {
 	b.Run("TrackOutbound", func(b *testing.B) {
-		tracker := NewICMPTracker(DefaultICMPTimeout, logger)
+		tracker := NewICMPTracker(DefaultICMPTimeout, logger, flowLogger)
 		defer tracker.Close()
 
-		srcIP := net.ParseIP("192.168.1.1")
-		dstIP := net.ParseIP("192.168.1.2")
+		srcIP := netip.MustParseAddr("192.168.1.1")
+		dstIP := netip.MustParseAddr("192.168.1.2")
 
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
-			tracker.TrackOutbound(srcIP, dstIP, uint16(i%65535), uint16(i%65535))
+			tracker.TrackOutbound(srcIP, dstIP, uint16(i%65535), 0, 0)
 		}
 	})
 
 	b.Run("IsValidInbound", func(b *testing.B) {
-		tracker := NewICMPTracker(DefaultICMPTimeout, logger)
+		tracker := NewICMPTracker(DefaultICMPTimeout, logger, flowLogger)
 		defer tracker.Close()
 
-		srcIP := net.ParseIP("192.168.1.1")
-		dstIP := net.ParseIP("192.168.1.2")
+		srcIP := netip.MustParseAddr("192.168.1.1")
+		dstIP := netip.MustParseAddr("192.168.1.2")
 
 		// Pre-populate some connections
 		for i := 0; i < 1000; i++ {
-			tracker.TrackOutbound(srcIP, dstIP, uint16(i), uint16(i))
+			tracker.TrackOutbound(srcIP, dstIP, uint16(i), 0, 0)
 		}
 
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
-			tracker.IsValidInbound(dstIP, srcIP, uint16(i%1000), uint16(i%1000), 0)
+			tracker.IsValidInbound(dstIP, srcIP, uint16(i%1000), 0, 0)
 		}
 	})
 }

client/firewall/uspfilter/conntrack/tcp.go

@@ -4,12 +4,15 @@ package conntrack
 
 import (
 	"context"
-	"net"
+	"net/netip"
 	"sync"
 	"sync/atomic"
 	"time"
 
+	"github.com/google/uuid"
+
 	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 )
 
 const (
@@ -40,6 +43,35 @@ const (
 // TCPState represents the state of a TCP connection
 type TCPState int
 
+func (s TCPState) String() string {
+	switch s {
+	case TCPStateNew:
+		return "New"
+	case TCPStateSynSent:
+		return "SYN Sent"
+	case TCPStateSynReceived:
+		return "SYN Received"
+	case TCPStateEstablished:
+		return "Established"
+	case TCPStateFinWait1:
+		return "FIN Wait 1"
+	case TCPStateFinWait2:
+		return "FIN Wait 2"
+	case TCPStateClosing:
+		return "Closing"
+	case TCPStateTimeWait:
+		return "Time Wait"
+	case TCPStateCloseWait:
+		return "Close Wait"
+	case TCPStateLastAck:
+		return "Last ACK"
+	case TCPStateClosed:
+		return "Closed"
+	default:
+		return "Unknown"
+	}
+}
+
 const (
 	TCPStateNew TCPState = iota
 	TCPStateSynSent
@@ -54,19 +86,14 @@ const (
 	TCPStateClosed
 )
 
-// TCPConnKey uniquely identifies a TCP connection
-type TCPConnKey struct {
-	SrcIP   [16]byte
-	DstIP   [16]byte
-	SrcPort uint16
-	DstPort uint16
-}
-
 // TCPConnTrack represents a TCP connection state
 type TCPConnTrack struct {
 	BaseConnTrack
+	SourcePort  uint16
+	DestPort    uint16
 	State       TCPState
 	established atomic.Bool
+	tombstone   atomic.Bool
 	sync.RWMutex
 }
 
@@ -80,6 +107,16 @@ func (t *TCPConnTrack) SetEstablished(state bool) {
 	t.established.Store(state)
 }
 
+// IsTombstone safely checks if the connection is marked for deletion
+func (t *TCPConnTrack) IsTombstone() bool {
+	return t.tombstone.Load()
+}
+
+// SetTombstone safely marks the connection for deletion
+func (t *TCPConnTrack) SetTombstone() {
+	t.tombstone.Store(true)
+}
+
 // TCPTracker manages TCP connection states
 type TCPTracker struct {
 	logger        *nblog.Logger
@@ -88,11 +125,14 @@ type TCPTracker struct {
 	cleanupTicker *time.Ticker
 	tickerCancel  context.CancelFunc
 	timeout       time.Duration
-	ipPool        *PreallocatedIPs
+	flowLogger    nftypes.FlowLogger
 }
 
 // NewTCPTracker creates a new TCP connection tracker
-func NewTCPTracker(timeout time.Duration, logger *nblog.Logger) *TCPTracker {
+func NewTCPTracker(timeout time.Duration, logger *nblog.Logger, flowLogger nftypes.FlowLogger) *TCPTracker {
+	if timeout == 0 {
+		timeout = DefaultTCPTimeout
+	}
 
 	ctx, cancel := context.WithCancel(context.Background())
 
@@ -102,59 +142,91 @@ func NewTCPTracker(timeout time.Duration, logger *nblog.Logger) *TCPTracker {
 		cleanupTicker: time.NewTicker(TCPCleanupInterval),
 		tickerCancel:  cancel,
 		timeout:       timeout,
-		ipPool:        NewPreallocatedIPs(),
+		flowLogger:    flowLogger,
 	}
 
 	go tracker.cleanupRoutine(ctx)
 	return tracker
 }
 
-// TrackOutbound processes an outbound TCP packet and updates connection state
-func (t *TCPTracker) TrackOutbound(srcIP net.IP, dstIP net.IP, srcPort uint16, dstPort uint16, flags uint8) {
-	// Create key before lock
-	key := makeConnKey(srcIP, dstIP, srcPort, dstPort)
+func (t *TCPTracker) updateIfExists(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, flags uint8, direction nftypes.Direction, size int) (ConnKey, bool) {
+	key := ConnKey{
+		SrcIP:   srcIP,
+		DstIP:   dstIP,
+		SrcPort: srcPort,
+		DstPort: dstPort,
+	}
+
+	t.mutex.RLock()
+	conn, exists := t.connections[key]
+	t.mutex.RUnlock()
+
+	if exists {
+		conn.Lock()
+		t.updateState(key, conn, flags, conn.Direction == nftypes.Egress)
+		conn.Unlock()
+
+		conn.UpdateCounters(direction, size)
+
+		return key, true
+	}
+
+	return key, false
+}
+
+// TrackOutbound records an outbound TCP connection
+func (t *TCPTracker) TrackOutbound(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, flags uint8, size int) {
+	if _, exists := t.updateIfExists(dstIP, srcIP, dstPort, srcPort, flags, 0, 0); !exists {
+		// if (inverted direction) conn is not tracked, track this direction
+		t.track(srcIP, dstIP, srcPort, dstPort, flags, nftypes.Egress, nil, size)
+	}
+}
+
+// TrackInbound processes an inbound TCP packet and updates connection state
+func (t *TCPTracker) TrackInbound(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, flags uint8, ruleID []byte, size int) {
+	t.track(srcIP, dstIP, srcPort, dstPort, flags, nftypes.Ingress, ruleID, size)
+}
+
+// track is the common implementation for tracking both inbound and outbound connections
+func (t *TCPTracker) track(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, flags uint8, direction nftypes.Direction, ruleID []byte, size int) {
+	key, exists := t.updateIfExists(srcIP, dstIP, srcPort, dstPort, flags, direction, size)
+	if exists {
+		return
+	}
+
+	conn := &TCPConnTrack{
+		BaseConnTrack: BaseConnTrack{
+			FlowId:    uuid.New(),
+			Direction: direction,
+			SourceIP:  srcIP,
+			DestIP:    dstIP,
+		},
+		SourcePort: srcPort,
+		DestPort:   dstPort,
+	}
+
+	conn.established.Store(false)
+	conn.tombstone.Store(false)
+
+	t.logger.Trace("New %s TCP connection: %s", direction, key)
+	t.updateState(key, conn, flags, direction == nftypes.Egress)
 
 	t.mutex.Lock()
-	conn, exists := t.connections[key]
-	if !exists {
-		// Use preallocated IPs
-		srcIPCopy := t.ipPool.Get()
-		dstIPCopy := t.ipPool.Get()
-		copyIP(srcIPCopy, srcIP)
-		copyIP(dstIPCopy, dstIP)
-
-		conn = &TCPConnTrack{
-			BaseConnTrack: BaseConnTrack{
-				SourceIP:   srcIPCopy,
-				DestIP:     dstIPCopy,
-				SourcePort: srcPort,
-				DestPort:   dstPort,
-			},
-			State: TCPStateNew,
-		}
-		conn.UpdateLastSeen()
-		conn.established.Store(false)
-		t.connections[key] = conn
-
-		t.logger.Trace("New TCP connection: %s:%d -> %s:%d", srcIP, srcPort, dstIP, dstPort)
-	}
+	t.connections[key] = conn
 	t.mutex.Unlock()
 
-	// Lock individual connection for state update
-	conn.Lock()
-	t.updateState(conn, flags, true)
-	conn.Unlock()
-	conn.UpdateLastSeen()
+	t.sendEvent(nftypes.TypeStart, conn, ruleID)
 }
 
 // IsValidInbound checks if an inbound TCP packet matches a tracked connection
-func (t *TCPTracker) IsValidInbound(srcIP net.IP, dstIP net.IP, srcPort uint16, dstPort uint16, flags uint8) bool {
-	if !isValidFlagCombination(flags) {
-		return false
+func (t *TCPTracker) IsValidInbound(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, flags uint8, size int) bool {
+	key := ConnKey{
+		SrcIP:   dstIP,
+		DstIP:   srcIP,
+		SrcPort: dstPort,
+		DstPort: srcPort,
 	}
 
-	key := makeConnKey(dstIP, srcIP, dstPort, srcPort)
-
 	t.mutex.RLock()
 	conn, exists := t.connections[key]
 	t.mutex.RUnlock()
@@ -163,22 +235,26 @@ func (t *TCPTracker) IsValidInbound(srcIP net.IP, dstIP net.IP, srcPort uint16,
 		return false
 	}
 
-	// Handle RST packets
+	// Handle RST flag specially - it always causes transition to closed
 	if flags&TCPRst != 0 {
-		conn.Lock()
-		if conn.IsEstablished() || conn.State == TCPStateSynSent || conn.State == TCPStateSynReceived {
-			conn.State = TCPStateClosed
-			conn.SetEstablished(false)
-			conn.Unlock()
+		if conn.IsTombstone() {
 			return true
 		}
+
+		conn.Lock()
+		conn.SetTombstone()
+		conn.State = TCPStateClosed
+		conn.SetEstablished(false)
 		conn.Unlock()
-		return false
+		conn.UpdateCounters(nftypes.Ingress, size)
+
+		t.logger.Trace("TCP connection reset: %s", key)
+		t.sendEvent(nftypes.TypeEnd, conn, nil)
+		return true
 	}
 
 	conn.Lock()
-	t.updateState(conn, flags, false)
-	conn.UpdateLastSeen()
+	t.updateState(key, conn, flags, false)
 	isEstablished := conn.IsEstablished()
 	isValidState := t.isValidStateForFlags(conn.State, flags)
 	conn.Unlock()
@@ -187,18 +263,17 @@ func (t *TCPTracker) IsValidInbound(srcIP net.IP, dstIP net.IP, srcPort uint16,
 }
 
 // updateState updates the TCP connection state based on flags
-func (t *TCPTracker) updateState(conn *TCPConnTrack, flags uint8, isOutbound bool) {
-	// Handle RST flag specially - it always causes transition to closed
-	if flags&TCPRst != 0 {
-		conn.State = TCPStateClosed
-		conn.SetEstablished(false)
+func (t *TCPTracker) updateState(key ConnKey, conn *TCPConnTrack, flags uint8, isOutbound bool) {
+	conn.UpdateLastSeen()
 
-		t.logger.Trace("TCP connection reset: %s:%d -> %s:%d",
-			conn.SourceIP, conn.SourcePort, conn.DestIP, conn.DestPort)
-		return
-	}
+	state := conn.State
+	defer func() {
+		if state != conn.State {
+			t.logger.Trace("TCP connection %s transitioned from %s to %s", key, state, conn.State)
+		}
+	}()
 
-	switch conn.State {
+	switch state {
 	case TCPStateNew:
 		if flags&TCPSyn != 0 && flags&TCPAck == 0 {
 			conn.State = TCPStateSynSent
@@ -207,11 +282,11 @@ func (t *TCPTracker) updateState(conn *TCPConnTrack, flags uint8, isOutbound boo
 	case TCPStateSynSent:
 		if flags&TCPSyn != 0 && flags&TCPAck != 0 {
 			if isOutbound {
-				conn.State = TCPStateSynReceived
-			} else {
-				// Simultaneous open
 				conn.State = TCPStateEstablished
 				conn.SetEstablished(true)
+			} else {
+				// Simultaneous open
+				conn.State = TCPStateSynReceived
 			}
 		}
 
@@ -229,22 +304,32 @@ func (t *TCPTracker) updateState(conn *TCPConnTrack, flags uint8, isOutbound boo
 				conn.State = TCPStateCloseWait
 			}
 			conn.SetEstablished(false)
+		} else if flags&TCPRst != 0 {
+			conn.State = TCPStateClosed
+			conn.SetTombstone()
+			t.sendEvent(nftypes.TypeEnd, conn, nil)
 		}
 
 	case TCPStateFinWait1:
 		switch {
 		case flags&TCPFin != 0 && flags&TCPAck != 0:
-			// Simultaneous close - both sides sent FIN
 			conn.State = TCPStateClosing
 		case flags&TCPFin != 0:
 			conn.State = TCPStateFinWait2
 		case flags&TCPAck != 0:
 			conn.State = TCPStateFinWait2
+		case flags&TCPRst != 0:
+			conn.State = TCPStateClosed
+			conn.SetTombstone()
+			t.sendEvent(nftypes.TypeEnd, conn, nil)
 		}
 
 	case TCPStateFinWait2:
 		if flags&TCPFin != 0 {
 			conn.State = TCPStateTimeWait
+
+			t.logger.Trace("TCP connection %s completed", key)
+			t.sendEvent(nftypes.TypeEnd, conn, nil)
 		}
 
 	case TCPStateClosing:
@@ -252,8 +337,8 @@ func (t *TCPTracker) updateState(conn *TCPConnTrack, flags uint8, isOutbound boo
 			conn.State = TCPStateTimeWait
 			// Keep established = false from previous state
 
-			t.logger.Trace("TCP connection closed (simultaneous) - %s:%d -> %s:%d",
-				conn.SourceIP, conn.SourcePort, conn.DestIP, conn.DestPort)
+			t.logger.Trace("TCP connection %s closed (simultaneous)", key)
+			t.sendEvent(nftypes.TypeEnd, conn, nil)
 		}
 
 	case TCPStateCloseWait:
@@ -264,17 +349,12 @@ func (t *TCPTracker) updateState(conn *TCPConnTrack, flags uint8, isOutbound boo
 	case TCPStateLastAck:
 		if flags&TCPAck != 0 {
 			conn.State = TCPStateClosed
+			conn.SetTombstone()
 
-			t.logger.Trace("TCP connection gracefully closed: %s:%d -> %s:%d",
-				conn.SourceIP, conn.SourcePort, conn.DestIP, conn.DestPort)
+			// Send close event for gracefully closed connections
+			t.sendEvent(nftypes.TypeEnd, conn, nil)
+			t.logger.Trace("TCP connection %s closed gracefully", key)
 		}
-
-	case TCPStateTimeWait:
-		// Stay in TIME-WAIT for 2MSL before transitioning to closed
-		// This is handled by the cleanup routine
-
-		t.logger.Trace("TCP connection completed - %s:%d -> %s:%d",
-			conn.SourceIP, conn.SourcePort, conn.DestIP, conn.DestPort)
 	}
 }
 
@@ -337,6 +417,12 @@ func (t *TCPTracker) cleanup() {
 	defer t.mutex.Unlock()
 
 	for key, conn := range t.connections {
+		if conn.IsTombstone() {
+			// Clean up tombstoned connections without sending an event
+			delete(t.connections, key)
+			continue
+		}
+
 		var timeout time.Duration
 		switch {
 		case conn.State == TCPStateTimeWait:
@@ -347,14 +433,16 @@ func (t *TCPTracker) cleanup() {
 			timeout = TCPHandshakeTimeout
 		}
 
-		lastSeen := conn.GetLastSeen()
-		if time.Since(lastSeen) > timeout {
+		if conn.timeoutExceeded(timeout) {
 			// Return IPs to pool
-			t.ipPool.Put(conn.SourceIP)
-			t.ipPool.Put(conn.DestIP)
 			delete(t.connections, key)
 
-			t.logger.Trace("Cleaned up TCP connection: %s:%d -> %s:%d", conn.SourceIP, conn.SourcePort, conn.DestIP, conn.DestPort)
+			t.logger.Trace("Cleaned up timed-out TCP connection %s", key)
+
+			// event already handled by state change
+			if conn.State != TCPStateTimeWait {
+				t.sendEvent(nftypes.TypeEnd, conn, nil)
+			}
 		}
 	}
 }
@@ -365,10 +453,6 @@ func (t *TCPTracker) Close() {
 
 	// Clean up all remaining IPs
 	t.mutex.Lock()
-	for _, conn := range t.connections {
-		t.ipPool.Put(conn.SourceIP)
-		t.ipPool.Put(conn.DestIP)
-	}
 	t.connections = nil
 	t.mutex.Unlock()
 }
@@ -386,3 +470,21 @@ func isValidFlagCombination(flags uint8) bool {
 
 	return true
 }
+
+func (t *TCPTracker) sendEvent(typ nftypes.Type, conn *TCPConnTrack, ruleID []byte) {
+	t.flowLogger.StoreEvent(nftypes.EventFields{
+		FlowID:     conn.FlowId,
+		Type:       typ,
+		RuleID:     ruleID,
+		Direction:  conn.Direction,
+		Protocol:   nftypes.TCP,
+		SourceIP:   conn.SourceIP,
+		DestIP:     conn.DestIP,
+		SourcePort: conn.SourcePort,
+		DestPort:   conn.DestPort,
+		RxPackets:  conn.PacketsRx.Load(),
+		TxPackets:  conn.PacketsTx.Load(),
+		RxBytes:    conn.BytesRx.Load(),
+		TxBytes:    conn.BytesTx.Load(),
+	})
+}

client/firewall/uspfilter/conntrack/tcp_test.go

@@ -1,7 +1,7 @@
 package conntrack
 
 import (
-	"net"
+	"net/netip"
 	"testing"
 	"time"
 
@@ -9,11 +9,11 @@ import (
 )
 
 func TestTCPStateMachine(t *testing.T) {
-	tracker := NewTCPTracker(DefaultTCPTimeout, logger)
+	tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
 	defer tracker.Close()
 
-	srcIP := net.ParseIP("100.64.0.1")
-	dstIP := net.ParseIP("100.64.0.2")
+	srcIP := netip.MustParseAddr("100.64.0.1")
+	dstIP := netip.MustParseAddr("100.64.0.2")
 	srcPort := uint16(12345)
 	dstPort := uint16(80)
 
@@ -58,7 +58,7 @@ func TestTCPStateMachine(t *testing.T) {
 
 		for _, tt := range tests {
 			t.Run(tt.name, func(t *testing.T) {
-				isValid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, tt.flags)
+				isValid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, tt.flags, 0)
 				require.Equal(t, !tt.wantDrop, isValid, tt.desc)
 			})
 		}
@@ -76,17 +76,17 @@ func TestTCPStateMachine(t *testing.T) {
 					t.Helper()
 
 					// Send initial SYN
-					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPSyn)
+					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPSyn, 0)
 
 					// Receive SYN-ACK
-					valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPSyn|TCPAck)
+					valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPSyn|TCPAck, 0)
 					require.True(t, valid, "SYN-ACK should be allowed")
 
 					// Send ACK
-					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck)
+					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, 0)
 
 					// Test data transfer
-					valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPPush|TCPAck)
+					valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPPush|TCPAck, 0)
 					require.True(t, valid, "Data should be allowed after handshake")
 				},
 			},
@@ -99,18 +99,18 @@ func TestTCPStateMachine(t *testing.T) {
 					establishConnection(t, tracker, srcIP, dstIP, srcPort, dstPort)
 
 					// Send FIN
-					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPFin|TCPAck)
+					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPFin|TCPAck, 0)
 
 					// Receive ACK for FIN
-					valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPAck)
+					valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPAck, 0)
 					require.True(t, valid, "ACK for FIN should be allowed")
 
 					// Receive FIN from other side
-					valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPFin|TCPAck)
+					valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPFin|TCPAck, 0)
 					require.True(t, valid, "FIN should be allowed")
 
 					// Send final ACK
-					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck)
+					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, 0)
 				},
 			},
 			{
@@ -122,7 +122,7 @@ func TestTCPStateMachine(t *testing.T) {
 					establishConnection(t, tracker, srcIP, dstIP, srcPort, dstPort)
 
 					// Receive RST
-					valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPRst)
+					valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPRst, 0)
 					require.True(t, valid, "RST should be allowed for established connection")
 
 					// Connection is logically dead but we don't enforce blocking subsequent packets
@@ -138,13 +138,13 @@ func TestTCPStateMachine(t *testing.T) {
 					establishConnection(t, tracker, srcIP, dstIP, srcPort, dstPort)
 
 					// Both sides send FIN+ACK
-					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPFin|TCPAck)
-					valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPFin|TCPAck)
+					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPFin|TCPAck, 0)
+					valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPFin|TCPAck, 0)
 					require.True(t, valid, "Simultaneous FIN should be allowed")
 
 					// Both sides send final ACK
-					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck)
-					valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPAck)
+					tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, 0)
+					valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPAck, 0)
 					require.True(t, valid, "Final ACKs should be allowed")
 				},
 			},
@@ -154,7 +154,7 @@ func TestTCPStateMachine(t *testing.T) {
 			t.Run(tt.name, func(t *testing.T) {
 				t.Helper()
 
-				tracker = NewTCPTracker(DefaultTCPTimeout, logger)
+				tracker = NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
 				tt.test(t)
 			})
 		}
@@ -162,11 +162,11 @@ func TestTCPStateMachine(t *testing.T) {
 }
 
 func TestRSTHandling(t *testing.T) {
-	tracker := NewTCPTracker(DefaultTCPTimeout, logger)
+	tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
 	defer tracker.Close()
 
-	srcIP := net.ParseIP("100.64.0.1")
-	dstIP := net.ParseIP("100.64.0.2")
+	srcIP := netip.MustParseAddr("100.64.0.1")
+	dstIP := netip.MustParseAddr("100.64.0.2")
 	srcPort := uint16(12345)
 	dstPort := uint16(80)
 
@@ -181,12 +181,12 @@ func TestRSTHandling(t *testing.T) {
 			name: "RST in established",
 			setupState: func() {
 				// Establish connection first
-				tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPSyn)
-				tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPSyn|TCPAck)
-				tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck)
+				tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPSyn, 0)
+				tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPSyn|TCPAck, 0)
+				tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, 0)
 			},
 			sendRST: func() {
-				tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPRst)
+				tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPRst, 0)
 			},
 			wantValid: true,
 			desc:      "Should accept RST for established connection",
@@ -195,7 +195,7 @@ func TestRSTHandling(t *testing.T) {
 			name:       "RST without connection",
 			setupState: func() {},
 			sendRST: func() {
-				tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPRst)
+				tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPRst, 0)
 			},
 			wantValid: false,
 			desc:      "Should reject RST without connection",
@@ -208,7 +208,12 @@ func TestRSTHandling(t *testing.T) {
 			tt.sendRST()
 
 			// Verify connection state is as expected
-			key := makeConnKey(srcIP, dstIP, srcPort, dstPort)
+			key := ConnKey{
+				SrcIP:   srcIP,
+				DstIP:   dstIP,
+				SrcPort: srcPort,
+				DstPort: dstPort,
+			}
 			conn := tracker.connections[key]
 			if tt.wantValid {
 				require.NotNil(t, conn)
@@ -220,63 +225,63 @@ func TestRSTHandling(t *testing.T) {
 }
 
 // Helper to establish a TCP connection
-func establishConnection(t *testing.T, tracker *TCPTracker, srcIP, dstIP net.IP, srcPort, dstPort uint16) {
+func establishConnection(t *testing.T, tracker *TCPTracker, srcIP, dstIP netip.Addr, srcPort, dstPort uint16) {
 	t.Helper()
 
-	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPSyn)
+	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPSyn, 0)
 
-	valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPSyn|TCPAck)
+	valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPSyn|TCPAck, 0)
 	require.True(t, valid, "SYN-ACK should be allowed")
 
-	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck)
+	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, 0)
 }
 
 func BenchmarkTCPTracker(b *testing.B) {
 	b.Run("TrackOutbound", func(b *testing.B) {
-		tracker := NewTCPTracker(DefaultTCPTimeout, logger)
+		tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
 		defer tracker.Close()
 
-		srcIP := net.ParseIP("192.168.1.1")
-		dstIP := net.ParseIP("192.168.1.2")
+		srcIP := netip.MustParseAddr("192.168.1.1")
+		dstIP := netip.MustParseAddr("192.168.1.2")
 
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
-			tracker.TrackOutbound(srcIP, dstIP, uint16(i%65535), 80, TCPSyn)
+			tracker.TrackOutbound(srcIP, dstIP, uint16(i%65535), 80, TCPSyn, 0)
 		}
 	})
 
 	b.Run("IsValidInbound", func(b *testing.B) {
-		tracker := NewTCPTracker(DefaultTCPTimeout, logger)
+		tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
 		defer tracker.Close()
 
-		srcIP := net.ParseIP("192.168.1.1")
-		dstIP := net.ParseIP("192.168.1.2")
+		srcIP := netip.MustParseAddr("192.168.1.1")
+		dstIP := netip.MustParseAddr("192.168.1.2")
 
 		// Pre-populate some connections
 		for i := 0; i < 1000; i++ {
-			tracker.TrackOutbound(srcIP, dstIP, uint16(i), 80, TCPSyn)
+			tracker.TrackOutbound(srcIP, dstIP, uint16(i), 80, TCPSyn, 0)
 		}
 
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
-			tracker.IsValidInbound(dstIP, srcIP, 80, uint16(i%1000), TCPAck)
+			tracker.IsValidInbound(dstIP, srcIP, 80, uint16(i%1000), TCPAck, 0)
 		}
 	})
 
 	b.Run("ConcurrentAccess", func(b *testing.B) {
-		tracker := NewTCPTracker(DefaultTCPTimeout, logger)
+		tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
 		defer tracker.Close()
 
-		srcIP := net.ParseIP("192.168.1.1")
-		dstIP := net.ParseIP("192.168.1.2")
+		srcIP := netip.MustParseAddr("192.168.1.1")
+		dstIP := netip.MustParseAddr("192.168.1.2")
 
 		b.RunParallel(func(pb *testing.PB) {
 			i := 0
 			for pb.Next() {
 				if i%2 == 0 {
-					tracker.TrackOutbound(srcIP, dstIP, uint16(i%65535), 80, TCPSyn)
+					tracker.TrackOutbound(srcIP, dstIP, uint16(i%65535), 80, TCPSyn, 0)
 				} else {
-					tracker.IsValidInbound(dstIP, srcIP, 80, uint16(i%65535), TCPAck)
+					tracker.IsValidInbound(dstIP, srcIP, 80, uint16(i%65535), TCPAck, 0)
 				}
 				i++
 			}
@@ -287,14 +292,14 @@ func BenchmarkTCPTracker(b *testing.B) {
 // Benchmark connection cleanup
 func BenchmarkCleanup(b *testing.B) {
 	b.Run("TCPCleanup", func(b *testing.B) {
-		tracker := NewTCPTracker(100*time.Millisecond, logger) // Short timeout for testing
+		tracker := NewTCPTracker(100*time.Millisecond, logger, flowLogger) // Short timeout for testing
 		defer tracker.Close()
 
 		// Pre-populate with expired connections
-		srcIP := net.ParseIP("192.168.1.1")
-		dstIP := net.ParseIP("192.168.1.2")
+		srcIP := netip.MustParseAddr("192.168.1.1")
+		dstIP := netip.MustParseAddr("192.168.1.2")
 		for i := 0; i < 10000; i++ {
-			tracker.TrackOutbound(srcIP, dstIP, uint16(i), 80, TCPSyn)
+			tracker.TrackOutbound(srcIP, dstIP, uint16(i), 80, TCPSyn, 0)
 		}
 
 		// Wait for connections to expire

client/firewall/uspfilter/conntrack/udp.go

@@ -2,11 +2,14 @@ package conntrack
 
 import (
 	"context"
-	"net"
+	"net/netip"
 	"sync"
 	"time"
 
+	"github.com/google/uuid"
+
 	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 )
 
 const (
@@ -19,6 +22,8 @@ const (
 // UDPConnTrack represents a UDP connection state
 type UDPConnTrack struct {
 	BaseConnTrack
+	SourcePort uint16
+	DestPort   uint16
 }
 
 // UDPTracker manages UDP connection states
@@ -29,11 +34,11 @@ type UDPTracker struct {
 	cleanupTicker *time.Ticker
 	tickerCancel  context.CancelFunc
 	mutex         sync.RWMutex
-	ipPool        *PreallocatedIPs
+	flowLogger    nftypes.FlowLogger
 }
 
 // NewUDPTracker creates a new UDP connection tracker
-func NewUDPTracker(timeout time.Duration, logger *nblog.Logger) *UDPTracker {
+func NewUDPTracker(timeout time.Duration, logger *nblog.Logger, flowLogger nftypes.FlowLogger) *UDPTracker {
 	if timeout == 0 {
 		timeout = DefaultUDPTimeout
 	}
@@ -46,7 +51,7 @@ func NewUDPTracker(timeout time.Duration, logger *nblog.Logger) *UDPTracker {
 		timeout:       timeout,
 		cleanupTicker: time.NewTicker(UDPCleanupInterval),
 		tickerCancel:  cancel,
-		ipPool:        NewPreallocatedIPs(),
+		flowLogger:    flowLogger,
 	}
 
 	go tracker.cleanupRoutine(ctx)
@@ -54,55 +59,87 @@ func NewUDPTracker(timeout time.Duration, logger *nblog.Logger) *UDPTracker {
 }
 
 // TrackOutbound records an outbound UDP connection
-func (t *UDPTracker) TrackOutbound(srcIP net.IP, dstIP net.IP, srcPort uint16, dstPort uint16) {
-	key := makeConnKey(srcIP, dstIP, srcPort, dstPort)
-
-	t.mutex.Lock()
-	conn, exists := t.connections[key]
-	if !exists {
-		srcIPCopy := t.ipPool.Get()
-		dstIPCopy := t.ipPool.Get()
-		copyIP(srcIPCopy, srcIP)
-		copyIP(dstIPCopy, dstIP)
-
-		conn = &UDPConnTrack{
-			BaseConnTrack: BaseConnTrack{
-				SourceIP:   srcIPCopy,
-				DestIP:     dstIPCopy,
-				SourcePort: srcPort,
-				DestPort:   dstPort,
-			},
-		}
-		conn.UpdateLastSeen()
-		t.connections[key] = conn
-
-		t.logger.Trace("New UDP connection: %v", conn)
+func (t *UDPTracker) TrackOutbound(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, size int) {
+	if _, exists := t.updateIfExists(dstIP, srcIP, dstPort, srcPort, nftypes.Egress, size); !exists {
+		// if (inverted direction) conn is not tracked, track this direction
+		t.track(srcIP, dstIP, srcPort, dstPort, nftypes.Egress, nil, size)
 	}
-	t.mutex.Unlock()
-
-	conn.UpdateLastSeen()
 }
 
-// IsValidInbound checks if an inbound packet matches a tracked connection
-func (t *UDPTracker) IsValidInbound(srcIP net.IP, dstIP net.IP, srcPort uint16, dstPort uint16) bool {
-	key := makeConnKey(dstIP, srcIP, dstPort, srcPort)
+// TrackInbound records an inbound UDP connection
+func (t *UDPTracker) TrackInbound(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, ruleID []byte, size int) {
+	t.track(srcIP, dstIP, srcPort, dstPort, nftypes.Ingress, ruleID, size)
+}
+
+func (t *UDPTracker) updateIfExists(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, direction nftypes.Direction, size int) (ConnKey, bool) {
+	key := ConnKey{
+		SrcIP:   srcIP,
+		DstIP:   dstIP,
+		SrcPort: srcPort,
+		DstPort: dstPort,
+	}
 
 	t.mutex.RLock()
 	conn, exists := t.connections[key]
 	t.mutex.RUnlock()
 
-	if !exists {
+	if exists {
+		conn.UpdateLastSeen()
+		conn.UpdateCounters(direction, size)
+		return key, true
+	}
+
+	return key, false
+}
+
+// track is the common implementation for tracking both inbound and outbound connections
+func (t *UDPTracker) track(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, direction nftypes.Direction, ruleID []byte, size int) {
+	key, exists := t.updateIfExists(srcIP, dstIP, srcPort, dstPort, direction, size)
+	if exists {
+		return
+	}
+
+	conn := &UDPConnTrack{
+		BaseConnTrack: BaseConnTrack{
+			FlowId:    uuid.New(),
+			Direction: direction,
+			SourceIP:  srcIP,
+			DestIP:    dstIP,
+		},
+		SourcePort: srcPort,
+		DestPort:   dstPort,
+	}
+	conn.UpdateLastSeen()
+
+	t.mutex.Lock()
+	t.connections[key] = conn
+	t.mutex.Unlock()
+
+	t.logger.Trace("New %s UDP connection: %s", direction, key)
+	t.sendEvent(nftypes.TypeStart, conn, ruleID)
+}
+
+// IsValidInbound checks if an inbound packet matches a tracked connection
+func (t *UDPTracker) IsValidInbound(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, size int) bool {
+	key := ConnKey{
+		SrcIP:   dstIP,
+		DstIP:   srcIP,
+		SrcPort: dstPort,
+		DstPort: srcPort,
+	}
+
+	t.mutex.RLock()
+	conn, exists := t.connections[key]
+	t.mutex.RUnlock()
+
+	if !exists || conn.timeoutExceeded(t.timeout) {
 		return false
 	}
 
-	if conn.timeoutExceeded(t.timeout) {
-		return false
-	}
+	conn.UpdateLastSeen()
+	conn.UpdateCounters(nftypes.Ingress, size)
 
-	return ValidateIPs(MakeIPAddr(srcIP), conn.DestIP) &&
-		ValidateIPs(MakeIPAddr(dstIP), conn.SourceIP) &&
-		conn.DestPort == srcPort &&
-		conn.SourcePort == dstPort
+	return true
 }
 
 // cleanupRoutine periodically removes stale connections
@@ -125,11 +162,11 @@ func (t *UDPTracker) cleanup() {
 
 	for key, conn := range t.connections {
 		if conn.timeoutExceeded(t.timeout) {
-			t.ipPool.Put(conn.SourceIP)
-			t.ipPool.Put(conn.DestIP)
 			delete(t.connections, key)
 
-			t.logger.Trace("Removed UDP connection %v (timeout)", conn)
+			t.logger.Trace("Removed UDP connection %s (timeout) [in: %d Pkts/%d B out: %d Pkts/%d B]",
+				key, conn.PacketsRx.Load(), conn.BytesRx.Load(), conn.PacketsTx.Load(), conn.BytesTx.Load())
+			t.sendEvent(nftypes.TypeEnd, conn, nil)
 		}
 	}
 }
@@ -139,29 +176,44 @@ func (t *UDPTracker) Close() {
 	t.tickerCancel()
 
 	t.mutex.Lock()
-	for _, conn := range t.connections {
-		t.ipPool.Put(conn.SourceIP)
-		t.ipPool.Put(conn.DestIP)
-	}
 	t.connections = nil
 	t.mutex.Unlock()
 }
 
 // GetConnection safely retrieves a connection state
-func (t *UDPTracker) GetConnection(srcIP net.IP, srcPort uint16, dstIP net.IP, dstPort uint16) (*UDPConnTrack, bool) {
+func (t *UDPTracker) GetConnection(srcIP netip.Addr, srcPort uint16, dstIP netip.Addr, dstPort uint16) (*UDPConnTrack, bool) {
 	t.mutex.RLock()
 	defer t.mutex.RUnlock()
 
-	key := makeConnKey(srcIP, dstIP, srcPort, dstPort)
-	conn, exists := t.connections[key]
-	if !exists {
-		return nil, false
+	key := ConnKey{
+		SrcIP:   srcIP,
+		DstIP:   dstIP,
+		SrcPort: srcPort,
+		DstPort: dstPort,
 	}
-
-	return conn, true
+	conn, exists := t.connections[key]
+	return conn, exists
 }
 
 // Timeout returns the configured timeout duration for the tracker
 func (t *UDPTracker) Timeout() time.Duration {
 	return t.timeout
 }
+
+func (t *UDPTracker) sendEvent(typ nftypes.Type, conn *UDPConnTrack, ruleID []byte) {
+	t.flowLogger.StoreEvent(nftypes.EventFields{
+		FlowID:     conn.FlowId,
+		Type:       typ,
+		RuleID:     ruleID,
+		Direction:  conn.Direction,
+		Protocol:   nftypes.UDP,
+		SourceIP:   conn.SourceIP,
+		DestIP:     conn.DestIP,
+		SourcePort: conn.SourcePort,
+		DestPort:   conn.DestPort,
+		RxPackets:  conn.PacketsRx.Load(),
+		TxPackets:  conn.PacketsTx.Load(),
+		RxBytes:    conn.BytesRx.Load(),
+		TxBytes:    conn.BytesTx.Load(),
+	})
+}

client/firewall/uspfilter/conntrack/udp_test.go

@@ -2,7 +2,7 @@ package conntrack
 
 import (
 	"context"
-	"net"
+	"net/netip"
 	"testing"
 	"time"
 
@@ -30,7 +30,7 @@ func TestNewUDPTracker(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			tracker := NewUDPTracker(tt.timeout, logger)
+			tracker := NewUDPTracker(tt.timeout, logger, flowLogger)
 			assert.NotNil(t, tracker)
 			assert.Equal(t, tt.wantTimeout, tracker.timeout)
 			assert.NotNil(t, tracker.connections)
@@ -41,43 +41,48 @@ func TestNewUDPTracker(t *testing.T) {
 }
 
 func TestUDPTracker_TrackOutbound(t *testing.T) {
-	tracker := NewUDPTracker(DefaultUDPTimeout, logger)
+	tracker := NewUDPTracker(DefaultUDPTimeout, logger, flowLogger)
 	defer tracker.Close()
 
-	srcIP := net.ParseIP("192.168.1.2")
-	dstIP := net.ParseIP("192.168.1.3")
+	srcIP := netip.MustParseAddr("192.168.1.2")
+	dstIP := netip.MustParseAddr("192.168.1.3")
 	srcPort := uint16(12345)
 	dstPort := uint16(53)
 
-	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort)
+	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, 0)
 
 	// Verify connection was tracked
-	key := makeConnKey(srcIP, dstIP, srcPort, dstPort)
+	key := ConnKey{
+		SrcIP:   srcIP,
+		DstIP:   dstIP,
+		SrcPort: srcPort,
+		DstPort: dstPort,
+	}
 	conn, exists := tracker.connections[key]
 	require.True(t, exists)
-	assert.True(t, conn.SourceIP.Equal(srcIP))
-	assert.True(t, conn.DestIP.Equal(dstIP))
+	assert.True(t, conn.SourceIP.Compare(srcIP) == 0)
+	assert.True(t, conn.DestIP.Compare(dstIP) == 0)
 	assert.Equal(t, srcPort, conn.SourcePort)
 	assert.Equal(t, dstPort, conn.DestPort)
 	assert.WithinDuration(t, time.Now(), conn.GetLastSeen(), 1*time.Second)
 }
 
 func TestUDPTracker_IsValidInbound(t *testing.T) {
-	tracker := NewUDPTracker(1*time.Second, logger)
+	tracker := NewUDPTracker(1*time.Second, logger, flowLogger)
 	defer tracker.Close()
 
-	srcIP := net.ParseIP("192.168.1.2")
-	dstIP := net.ParseIP("192.168.1.3")
+	srcIP := netip.MustParseAddr("192.168.1.2")
+	dstIP := netip.MustParseAddr("192.168.1.3")
 	srcPort := uint16(12345)
 	dstPort := uint16(53)
 
 	// Track outbound connection
-	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort)
+	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, 0)
 
 	tests := []struct {
 		name    string
-		srcIP   net.IP
-		dstIP   net.IP
+		srcIP   netip.Addr
+		dstIP   netip.Addr
 		srcPort uint16
 		dstPort uint16
 		sleep   time.Duration
@@ -94,7 +99,7 @@ func TestUDPTracker_IsValidInbound(t *testing.T) {
 		},
 		{
 			name:    "invalid source IP",
-			srcIP:   net.ParseIP("192.168.1.4"),
+			srcIP:   netip.MustParseAddr("192.168.1.4"),
 			dstIP:   srcIP,
 			srcPort: dstPort,
 			dstPort: srcPort,
@@ -104,7 +109,7 @@ func TestUDPTracker_IsValidInbound(t *testing.T) {
 		{
 			name:    "invalid destination IP",
 			srcIP:   dstIP,
-			dstIP:   net.ParseIP("192.168.1.4"),
+			dstIP:   netip.MustParseAddr("192.168.1.4"),
 			srcPort: dstPort,
 			dstPort: srcPort,
 			sleep:   0,
@@ -144,7 +149,7 @@ func TestUDPTracker_IsValidInbound(t *testing.T) {
 			if tt.sleep > 0 {
 				time.Sleep(tt.sleep)
 			}
-			got := tracker.IsValidInbound(tt.srcIP, tt.dstIP, tt.srcPort, tt.dstPort)
+			got := tracker.IsValidInbound(tt.srcIP, tt.dstIP, tt.srcPort, tt.dstPort, 0)
 			assert.Equal(t, tt.want, got)
 		})
 	}
@@ -164,8 +169,8 @@ func TestUDPTracker_Cleanup(t *testing.T) {
 		timeout:       timeout,
 		cleanupTicker: time.NewTicker(cleanupInterval),
 		tickerCancel:  tickerCancel,
-		ipPool:        NewPreallocatedIPs(),
 		logger:        logger,
+		flowLogger:    flowLogger,
 	}
 
 	// Start cleanup routine
@@ -173,27 +178,27 @@ func TestUDPTracker_Cleanup(t *testing.T) {
 
 	// Add some connections
 	connections := []struct {
-		srcIP   net.IP
-		dstIP   net.IP
+		srcIP   netip.Addr
+		dstIP   netip.Addr
 		srcPort uint16
 		dstPort uint16
 	}{
 		{
-			srcIP:   net.ParseIP("192.168.1.2"),
-			dstIP:   net.ParseIP("192.168.1.3"),
+			srcIP:   netip.MustParseAddr("192.168.1.2"),
+			dstIP:   netip.MustParseAddr("192.168.1.3"),
 			srcPort: 12345,
 			dstPort: 53,
 		},
 		{
-			srcIP:   net.ParseIP("192.168.1.4"),
-			dstIP:   net.ParseIP("192.168.1.5"),
+			srcIP:   netip.MustParseAddr("192.168.1.4"),
+			dstIP:   netip.MustParseAddr("192.168.1.5"),
 			srcPort: 12346,
 			dstPort: 53,
 		},
 	}
 
 	for _, conn := range connections {
-		tracker.TrackOutbound(conn.srcIP, conn.dstIP, conn.srcPort, conn.dstPort)
+		tracker.TrackOutbound(conn.srcIP, conn.dstIP, conn.srcPort, conn.dstPort, 0)
 	}
 
 	// Verify initial connections
@@ -215,33 +220,33 @@ func TestUDPTracker_Cleanup(t *testing.T) {
 
 func BenchmarkUDPTracker(b *testing.B) {
 	b.Run("TrackOutbound", func(b *testing.B) {
-		tracker := NewUDPTracker(DefaultUDPTimeout, logger)
+		tracker := NewUDPTracker(DefaultUDPTimeout, logger, flowLogger)
 		defer tracker.Close()
 
-		srcIP := net.ParseIP("192.168.1.1")
-		dstIP := net.ParseIP("192.168.1.2")
+		srcIP := netip.MustParseAddr("192.168.1.1")
+		dstIP := netip.MustParseAddr("192.168.1.2")
 
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
-			tracker.TrackOutbound(srcIP, dstIP, uint16(i%65535), 80)
+			tracker.TrackOutbound(srcIP, dstIP, uint16(i%65535), 80, 0)
 		}
 	})
 
 	b.Run("IsValidInbound", func(b *testing.B) {
-		tracker := NewUDPTracker(DefaultUDPTimeout, logger)
+		tracker := NewUDPTracker(DefaultUDPTimeout, logger, flowLogger)
 		defer tracker.Close()
 
-		srcIP := net.ParseIP("192.168.1.1")
-		dstIP := net.ParseIP("192.168.1.2")
+		srcIP := netip.MustParseAddr("192.168.1.1")
+		dstIP := netip.MustParseAddr("192.168.1.2")
 
 		// Pre-populate some connections
 		for i := 0; i < 1000; i++ {
-			tracker.TrackOutbound(srcIP, dstIP, uint16(i), 80)
+			tracker.TrackOutbound(srcIP, dstIP, uint16(i), 80, 0)
 		}
 
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
-			tracker.IsValidInbound(dstIP, srcIP, 80, uint16(i%1000))
+			tracker.IsValidInbound(dstIP, srcIP, 80, uint16(i%1000), 0)
 		}
 	})
 }

client/firewall/uspfilter/forwarder/endpoint.go

@@ -1,6 +1,8 @@
 package forwarder
 
 import (
+	"fmt"
+
 	wgdevice "golang.zx2c4.com/wireguard/device"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
@@ -79,3 +81,10 @@ func (e *endpoint) AddHeader(*stack.PacketBuffer) {
 func (e *endpoint) ParseHeader(*stack.PacketBuffer) bool {
 	return true
 }
+
+type epID stack.TransportEndpointID
+
+func (i epID) String() string {
+	// src and remote is swapped
+	return fmt.Sprintf("%s:%d -> %s:%d", i.RemoteAddress, i.RemotePort, i.LocalAddress, i.LocalPort)
+}

client/firewall/uspfilter/forwarder/forwarder.go

@@ -18,6 +18,7 @@ import (
 
 	"github.com/netbirdio/netbird/client/firewall/uspfilter/common"
 	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 )
 
 const (
@@ -29,6 +30,7 @@ const (
 
 type Forwarder struct {
 	logger       *nblog.Logger
+	flowLogger   nftypes.FlowLogger
 	stack        *stack.Stack
 	endpoint     *endpoint
 	udpForwarder *udpForwarder
@@ -38,7 +40,7 @@ type Forwarder struct {
 	netstack     bool
 }
 
-func New(iface common.IFaceMapper, logger *nblog.Logger, netstack bool) (*Forwarder, error) {
+func New(iface common.IFaceMapper, logger *nblog.Logger, flowLogger nftypes.FlowLogger, netstack bool) (*Forwarder, error) {
 	s := stack.New(stack.Options{
 		NetworkProtocols: []stack.NetworkProtocolFactory{ipv4.NewProtocol},
 		TransportProtocols: []stack.TransportProtocolFactory{
@@ -102,9 +104,10 @@ func New(iface common.IFaceMapper, logger *nblog.Logger, netstack bool) (*Forwar
 	ctx, cancel := context.WithCancel(context.Background())
 	f := &Forwarder{
 		logger:       logger,
+		flowLogger:   flowLogger,
 		stack:        s,
 		endpoint:     endpoint,
-		udpForwarder: newUDPForwarder(mtu, logger),
+		udpForwarder: newUDPForwarder(mtu, logger, flowLogger),
 		ctx:          ctx,
 		cancel:       cancel,
 		netstack:     netstack,

client/firewall/uspfilter/forwarder/icmp.go

@@ -3,14 +3,30 @@ package forwarder
 import (
 	"context"
 	"net"
+	"net/netip"
 	"time"
 
+	"github.com/google/uuid"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
+
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 )
 
 // handleICMP handles ICMP packets from the network stack
 func (f *Forwarder) handleICMP(id stack.TransportEndpointID, pkt stack.PacketBufferPtr) bool {
+	icmpHdr := header.ICMPv4(pkt.TransportHeader().View().AsSlice())
+	icmpType := uint8(icmpHdr.Type())
+	icmpCode := uint8(icmpHdr.Code())
+
+	if header.ICMPv4Type(icmpType) == header.ICMPv4EchoReply {
+		// dont process our own replies
+		return true
+	}
+
+	flowID := uuid.New()
+	f.sendICMPEvent(nftypes.TypeStart, flowID, id, icmpType, icmpCode)
+
 	ctx, cancel := context.WithTimeout(f.ctx, 5*time.Second)
 	defer cancel()
 
@@ -18,7 +34,7 @@ func (f *Forwarder) handleICMP(id stack.TransportEndpointID, pkt stack.PacketBuf
 	// TODO: support non-root
 	conn, err := lc.ListenPacket(ctx, "ip4:icmp", "0.0.0.0")
 	if err != nil {
-		f.logger.Error("Failed to create ICMP socket for %v: %v", id, err)
+		f.logger.Error("Failed to create ICMP socket for %v: %v", epID(id), err)
 
 		// This will make netstack reply on behalf of the original destination, that's ok for now
 		return false
@@ -32,47 +48,31 @@ func (f *Forwarder) handleICMP(id stack.TransportEndpointID, pkt stack.PacketBuf
 	dstIP := f.determineDialAddr(id.LocalAddress)
 	dst := &net.IPAddr{IP: dstIP}
 
-	// Get the complete ICMP message (header + data)
 	fullPacket := stack.PayloadSince(pkt.TransportHeader())
 	payload := fullPacket.AsSlice()
 
-	icmpHdr := header.ICMPv4(pkt.TransportHeader().View().AsSlice())
+	if _, err = conn.WriteTo(payload, dst); err != nil {
+		f.logger.Error("Failed to write ICMP packet for %v: %v", epID(id), err)
+		return true
+	}
+
+	f.logger.Trace("Forwarded ICMP packet %v type %v code %v",
+		epID(id), icmpHdr.Type(), icmpHdr.Code())
 
 	// For Echo Requests, send and handle response
-	switch icmpHdr.Type() {
-	case header.ICMPv4Echo:
-		return f.handleEchoResponse(icmpHdr, payload, dst, conn, id)
-	case header.ICMPv4EchoReply:
-		// dont process our own replies
-		return true
-	default:
+	if header.ICMPv4Type(icmpType) == header.ICMPv4Echo {
+		f.handleEchoResponse(icmpHdr, conn, id)
+		f.sendICMPEvent(nftypes.TypeEnd, flowID, id, icmpType, icmpCode)
 	}
 
-	// For other ICMP types (Time Exceeded, Destination Unreachable, etc)
-	_, err = conn.WriteTo(payload, dst)
-	if err != nil {
-		f.logger.Error("Failed to write ICMP packet for %v: %v", id, err)
-		return true
-	}
-
-	f.logger.Trace("Forwarded ICMP packet %v type=%v code=%v",
-		id, icmpHdr.Type(), icmpHdr.Code())
-
+	// For other ICMP types (Time Exceeded, Destination Unreachable, etc) do nothing
 	return true
 }
 
-func (f *Forwarder) handleEchoResponse(icmpHdr header.ICMPv4, payload []byte, dst *net.IPAddr, conn net.PacketConn, id stack.TransportEndpointID) bool {
-	if _, err := conn.WriteTo(payload, dst); err != nil {
-		f.logger.Error("Failed to write ICMP packet for %v: %v", id, err)
-		return true
-	}
-
-	f.logger.Trace("Forwarded ICMP packet %v type=%v code=%v",
-		id, icmpHdr.Type(), icmpHdr.Code())
-
+func (f *Forwarder) handleEchoResponse(icmpHdr header.ICMPv4, conn net.PacketConn, id stack.TransportEndpointID) {
 	if err := conn.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil {
 		f.logger.Error("Failed to set read deadline for ICMP response: %v", err)
-		return true
+		return
 	}
 
 	response := make([]byte, f.endpoint.mtu)
@@ -81,7 +81,7 @@ func (f *Forwarder) handleEchoResponse(icmpHdr header.ICMPv4, payload []byte, ds
 		if !isTimeout(err) {
 			f.logger.Error("Failed to read ICMP response: %v", err)
 		}
-		return true
+		return
 	}
 
 	ipHdr := make([]byte, header.IPv4MinimumSize)
@@ -101,9 +101,27 @@ func (f *Forwarder) handleEchoResponse(icmpHdr header.ICMPv4, payload []byte, ds
 
 	if err := f.InjectIncomingPacket(fullPacket); err != nil {
 		f.logger.Error("Failed to inject ICMP response: %v", err)
-		return true
+
+		return
 	}
 
-	f.logger.Trace("Forwarded ICMP echo reply for %v", id)
-	return true
+	f.logger.Trace("Forwarded ICMP echo reply for %v type %v code %v",
+		epID(id), icmpHdr.Type(), icmpHdr.Code())
+}
+
+// sendICMPEvent stores flow events for ICMP packets
+func (f *Forwarder) sendICMPEvent(typ nftypes.Type, flowID uuid.UUID, id stack.TransportEndpointID, icmpType, icmpCode uint8) {
+	f.flowLogger.StoreEvent(nftypes.EventFields{
+		FlowID:    flowID,
+		Type:      typ,
+		Direction: nftypes.Ingress,
+		Protocol:  nftypes.ICMP,
+		// TODO: handle ipv6
+		SourceIP: netip.AddrFrom4(id.RemoteAddress.As4()),
+		DestIP:   netip.AddrFrom4(id.LocalAddress.As4()),
+		ICMPType: icmpType,
+		ICMPCode: icmpCode,
+
+		// TODO: get packets/bytes
+	})
 }

client/firewall/uspfilter/forwarder/tcp.go

@@ -5,24 +5,38 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"net/netip"
 
+	"github.com/google/uuid"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 	"gvisor.dev/gvisor/pkg/tcpip/transport/tcp"
 	"gvisor.dev/gvisor/pkg/waiter"
+
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 )
 
 // handleTCP is called by the TCP forwarder for new connections.
 func (f *Forwarder) handleTCP(r *tcp.ForwarderRequest) {
 	id := r.ID()
 
+	flowID := uuid.New()
+
+	f.sendTCPEvent(nftypes.TypeStart, flowID, id, nil)
+	var success bool
+	defer func() {
+		if !success {
+			f.sendTCPEvent(nftypes.TypeEnd, flowID, id, nil)
+		}
+	}()
+
 	dialAddr := fmt.Sprintf("%s:%d", f.determineDialAddr(id.LocalAddress), id.LocalPort)
 
 	outConn, err := (&net.Dialer{}).DialContext(f.ctx, "tcp", dialAddr)
 	if err != nil {
 		r.Complete(true)
-		f.logger.Trace("forwarder: dial error for %v: %v", id, err)
+		f.logger.Trace("forwarder: dial error for %v: %v", epID(id), err)
 		return
 	}
 
@@ -44,12 +58,13 @@ func (f *Forwarder) handleTCP(r *tcp.ForwarderRequest) {
 
 	inConn := gonet.NewTCPConn(&wq, ep)
 
-	f.logger.Trace("forwarder: established TCP connection %v", id)
+	success = true
+	f.logger.Trace("forwarder: established TCP connection %v", epID(id))
 
-	go f.proxyTCP(id, inConn, outConn, ep)
+	go f.proxyTCP(id, inConn, outConn, ep, flowID)
 }
 
-func (f *Forwarder) proxyTCP(id stack.TransportEndpointID, inConn *gonet.TCPConn, outConn net.Conn, ep tcpip.Endpoint) {
+func (f *Forwarder) proxyTCP(id stack.TransportEndpointID, inConn *gonet.TCPConn, outConn net.Conn, ep tcpip.Endpoint, flowID uuid.UUID) {
 	defer func() {
 		if err := inConn.Close(); err != nil {
 			f.logger.Debug("forwarder: inConn close error: %v", err)
@@ -58,6 +73,8 @@ func (f *Forwarder) proxyTCP(id stack.TransportEndpointID, inConn *gonet.TCPConn
 			f.logger.Debug("forwarder: outConn close error: %v", err)
 		}
 		ep.Close()
+
+		f.sendTCPEvent(nftypes.TypeEnd, flowID, id, ep)
 	}()
 
 	// Create context for managing the proxy goroutines
@@ -78,13 +95,38 @@ func (f *Forwarder) proxyTCP(id stack.TransportEndpointID, inConn *gonet.TCPConn
 
 	select {
 	case <-ctx.Done():
-		f.logger.Trace("forwarder: tearing down TCP connection %v due to context done", id)
+		f.logger.Trace("forwarder: tearing down TCP connection %v due to context done", epID(id))
 		return
 	case err := <-errChan:
 		if err != nil && !isClosedError(err) {
 			f.logger.Error("proxyTCP: copy error: %v", err)
 		}
-		f.logger.Trace("forwarder: tearing down TCP connection %v", id)
+		f.logger.Trace("forwarder: tearing down TCP connection %v", epID(id))
 		return
 	}
 }
+
+func (f *Forwarder) sendTCPEvent(typ nftypes.Type, flowID uuid.UUID, id stack.TransportEndpointID, ep tcpip.Endpoint) {
+	fields := nftypes.EventFields{
+		FlowID:    flowID,
+		Type:      typ,
+		Direction: nftypes.Ingress,
+		Protocol:  nftypes.TCP,
+		// TODO: handle ipv6
+		SourceIP:   netip.AddrFrom4(id.RemoteAddress.As4()),
+		DestIP:     netip.AddrFrom4(id.LocalAddress.As4()),
+		SourcePort: id.RemotePort,
+		DestPort:   id.LocalPort,
+	}
+
+	if ep != nil {
+		if tcpStats, ok := ep.Stats().(*tcp.Stats); ok {
+			// fields are flipped since this is the in conn
+			// TODO: get bytes
+			fields.RxPackets = tcpStats.SegmentsSent.Value()
+			fields.TxPackets = tcpStats.SegmentsReceived.Value()
+		}
+	}
+
+	f.flowLogger.StoreEvent(fields)
+}

client/firewall/uspfilter/forwarder/udp.go

@@ -5,10 +5,12 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"net/netip"
 	"sync"
 	"sync/atomic"
 	"time"
 
+	"github.com/google/uuid"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
@@ -16,6 +18,7 @@ import (
 	"gvisor.dev/gvisor/pkg/waiter"
 
 	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 )
 
 const (
@@ -28,15 +31,17 @@ type udpPacketConn struct {
 	lastSeen atomic.Int64
 	cancel   context.CancelFunc
 	ep       tcpip.Endpoint
+	flowID   uuid.UUID
 }
 
 type udpForwarder struct {
 	sync.RWMutex
-	logger  *nblog.Logger
-	conns   map[stack.TransportEndpointID]*udpPacketConn
-	bufPool sync.Pool
-	ctx     context.Context
-	cancel  context.CancelFunc
+	logger     *nblog.Logger
+	flowLogger nftypes.FlowLogger
+	conns      map[stack.TransportEndpointID]*udpPacketConn
+	bufPool    sync.Pool
+	ctx        context.Context
+	cancel     context.CancelFunc
 }
 
 type idleConn struct {
@@ -44,13 +49,14 @@ type idleConn struct {
 	conn *udpPacketConn
 }
 
-func newUDPForwarder(mtu int, logger *nblog.Logger) *udpForwarder {
+func newUDPForwarder(mtu int, logger *nblog.Logger, flowLogger nftypes.FlowLogger) *udpForwarder {
 	ctx, cancel := context.WithCancel(context.Background())
 	f := &udpForwarder{
-		logger: logger,
-		conns:  make(map[stack.TransportEndpointID]*udpPacketConn),
-		ctx:    ctx,
-		cancel: cancel,
+		logger:     logger,
+		flowLogger: flowLogger,
+		conns:      make(map[stack.TransportEndpointID]*udpPacketConn),
+		ctx:        ctx,
+		cancel:     cancel,
 		bufPool: sync.Pool{
 			New: func() any {
 				b := make([]byte, mtu)
@@ -72,10 +78,10 @@ func (f *udpForwarder) Stop() {
 	for id, conn := range f.conns {
 		conn.cancel()
 		if err := conn.conn.Close(); err != nil {
-			f.logger.Debug("forwarder: UDP conn close error for %v: %v", id, err)
+			f.logger.Debug("forwarder: UDP conn close error for %v: %v", epID(id), err)
 		}
 		if err := conn.outConn.Close(); err != nil {
-			f.logger.Debug("forwarder: UDP outConn close error for %v: %v", id, err)
+			f.logger.Debug("forwarder: UDP outConn close error for %v: %v", epID(id), err)
 		}
 
 		conn.ep.Close()
@@ -106,10 +112,10 @@ func (f *udpForwarder) cleanup() {
 			for _, idle := range idleConns {
 				idle.conn.cancel()
 				if err := idle.conn.conn.Close(); err != nil {
-					f.logger.Debug("forwarder: UDP conn close error for %v: %v", idle.id, err)
+					f.logger.Debug("forwarder: UDP conn close error for %v: %v", epID(idle.id), err)
 				}
 				if err := idle.conn.outConn.Close(); err != nil {
-					f.logger.Debug("forwarder: UDP outConn close error for %v: %v", idle.id, err)
+					f.logger.Debug("forwarder: UDP outConn close error for %v: %v", epID(idle.id), err)
 				}
 
 				idle.conn.ep.Close()
@@ -118,7 +124,7 @@ func (f *udpForwarder) cleanup() {
 				delete(f.conns, idle.id)
 				f.Unlock()
 
-				f.logger.Trace("forwarder: cleaned up idle UDP connection %v", idle.id)
+				f.logger.Trace("forwarder: cleaned up idle UDP connection %v", epID(idle.id))
 			}
 		}
 	}
@@ -137,14 +143,24 @@ func (f *Forwarder) handleUDP(r *udp.ForwarderRequest) {
 	_, exists := f.udpForwarder.conns[id]
 	f.udpForwarder.RUnlock()
 	if exists {
-		f.logger.Trace("forwarder: existing UDP connection for %v", id)
+		f.logger.Trace("forwarder: existing UDP connection for %v", epID(id))
 		return
 	}
 
+	flowID := uuid.New()
+
+	f.sendUDPEvent(nftypes.TypeStart, flowID, id, nil)
+	var success bool
+	defer func() {
+		if !success {
+			f.sendUDPEvent(nftypes.TypeEnd, flowID, id, nil)
+		}
+	}()
+
 	dstAddr := fmt.Sprintf("%s:%d", f.determineDialAddr(id.LocalAddress), id.LocalPort)
 	outConn, err := (&net.Dialer{}).DialContext(f.ctx, "udp", dstAddr)
 	if err != nil {
-		f.logger.Debug("forwarder: UDP dial error for %v: %v", id, err)
+		f.logger.Debug("forwarder: UDP dial error for %v: %v", epID(id), err)
 		// TODO: Send ICMP error message
 		return
 	}
@@ -155,7 +171,7 @@ func (f *Forwarder) handleUDP(r *udp.ForwarderRequest) {
 	if epErr != nil {
 		f.logger.Debug("forwarder: failed to create UDP endpoint: %v", epErr)
 		if err := outConn.Close(); err != nil {
-			f.logger.Debug("forwarder: UDP outConn close error for %v: %v", id, err)
+			f.logger.Debug("forwarder: UDP outConn close error for %v: %v", epID(id), err)
 		}
 		return
 	}
@@ -168,6 +184,7 @@ func (f *Forwarder) handleUDP(r *udp.ForwarderRequest) {
 		outConn: outConn,
 		cancel:  connCancel,
 		ep:      ep,
+		flowID:  flowID,
 	}
 	pConn.updateLastSeen()
 
@@ -177,17 +194,20 @@ func (f *Forwarder) handleUDP(r *udp.ForwarderRequest) {
 		f.udpForwarder.Unlock()
 		pConn.cancel()
 		if err := inConn.Close(); err != nil {
-			f.logger.Debug("forwarder: UDP inConn close error for %v: %v", id, err)
+			f.logger.Debug("forwarder: UDP inConn close error for %v: %v", epID(id), err)
 		}
 		if err := outConn.Close(); err != nil {
-			f.logger.Debug("forwarder: UDP outConn close error for %v: %v", id, err)
+			f.logger.Debug("forwarder: UDP outConn close error for %v: %v", epID(id), err)
 		}
+
 		return
 	}
 	f.udpForwarder.conns[id] = pConn
 	f.udpForwarder.Unlock()
 
-	f.logger.Trace("forwarder: established UDP connection to %v", id)
+	success = true
+	f.logger.Trace("forwarder: established UDP connection %v", epID(id))
+
 	go f.proxyUDP(connCtx, pConn, id, ep)
 }
 
@@ -195,10 +215,10 @@ func (f *Forwarder) proxyUDP(ctx context.Context, pConn *udpPacketConn, id stack
 	defer func() {
 		pConn.cancel()
 		if err := pConn.conn.Close(); err != nil {
-			f.logger.Debug("forwarder: UDP inConn close error for %v: %v", id, err)
+			f.logger.Debug("forwarder: UDP inConn close error for %v: %v", epID(id), err)
 		}
 		if err := pConn.outConn.Close(); err != nil {
-			f.logger.Debug("forwarder: UDP outConn close error for %v: %v", id, err)
+			f.logger.Debug("forwarder: UDP outConn close error for %v: %v", epID(id), err)
 		}
 
 		ep.Close()
@@ -206,6 +226,8 @@ func (f *Forwarder) proxyUDP(ctx context.Context, pConn *udpPacketConn, id stack
 		f.udpForwarder.Lock()
 		delete(f.udpForwarder.conns, id)
 		f.udpForwarder.Unlock()
+
+		f.sendUDPEvent(nftypes.TypeEnd, pConn.flowID, id, ep)
 	}()
 
 	errChan := make(chan error, 2)
@@ -220,17 +242,43 @@ func (f *Forwarder) proxyUDP(ctx context.Context, pConn *udpPacketConn, id stack
 
 	select {
 	case <-ctx.Done():
-		f.logger.Trace("forwarder: tearing down UDP connection %v due to context done", id)
+		f.logger.Trace("forwarder: tearing down UDP connection %v due to context done", epID(id))
 		return
 	case err := <-errChan:
 		if err != nil && !isClosedError(err) {
 			f.logger.Error("proxyUDP: copy error: %v", err)
 		}
-		f.logger.Trace("forwarder: tearing down UDP connection %v", id)
+		f.logger.Trace("forwarder: tearing down UDP connection %v", epID(id))
 		return
 	}
 }
 
+// sendUDPEvent stores flow events for UDP connections
+func (f *Forwarder) sendUDPEvent(typ nftypes.Type, flowID uuid.UUID, id stack.TransportEndpointID, ep tcpip.Endpoint) {
+	fields := nftypes.EventFields{
+		FlowID:    flowID,
+		Type:      typ,
+		Direction: nftypes.Ingress,
+		Protocol:  nftypes.UDP,
+		// TODO: handle ipv6
+		SourceIP:   netip.AddrFrom4(id.RemoteAddress.As4()),
+		DestIP:     netip.AddrFrom4(id.LocalAddress.As4()),
+		SourcePort: id.RemotePort,
+		DestPort:   id.LocalPort,
+	}
+
+	if ep != nil {
+		if tcpStats, ok := ep.Stats().(*tcpip.TransportEndpointStats); ok {
+			// fields are flipped since this is the in conn
+			// TODO: get bytes
+			fields.RxPackets = tcpStats.PacketsSent.Value()
+			fields.TxPackets = tcpStats.PacketsReceived.Value()
+		}
+	}
+
+	f.flowLogger.StoreEvent(fields)
+}
+
 func (c *udpPacketConn) updateLastSeen() {
 	c.lastSeen.Store(time.Now().UnixNano())
 }

client/firewall/uspfilter/localip.go

@@ -3,6 +3,7 @@ package uspfilter
 import (
 	"fmt"
 	"net"
+	"net/netip"
 	"sync"
 
 	log "github.com/sirupsen/logrus"
@@ -31,13 +32,9 @@ func (m *localIPManager) setBitmapBit(ip net.IP) {
 	m.ipv4Bitmap[high] |= 1 << (low % 32)
 }
 
-func (m *localIPManager) checkBitmapBit(ip net.IP) bool {
-	ipv4 := ip.To4()
-	if ipv4 == nil {
-		return false
-	}
-	high := (uint16(ipv4[0]) << 8) | uint16(ipv4[1])
-	low := (uint16(ipv4[2]) << 8) | uint16(ipv4[3])
+func (m *localIPManager) checkBitmapBit(ip []byte) bool {
+	high := (uint16(ip[0]) << 8) | uint16(ip[1])
+	low := (uint16(ip[2]) << 8) | uint16(ip[3])
 	return (m.ipv4Bitmap[high] & (1 << (low % 32))) != 0
 }
 
@@ -122,12 +119,12 @@ func (m *localIPManager) UpdateLocalIPs(iface common.IFaceMapper) (err error) {
 	return nil
 }
 
-func (m *localIPManager) IsLocalIP(ip net.IP) bool {
+func (m *localIPManager) IsLocalIP(ip netip.Addr) bool {
 	m.mu.RLock()
 	defer m.mu.RUnlock()
 
-	if ipv4 := ip.To4(); ipv4 != nil {
-		return m.checkBitmapBit(ipv4)
+	if ip.Is4() {
+		return m.checkBitmapBit(ip.AsSlice())
 	}
 
 	return false

client/firewall/uspfilter/localip_test.go

@@ -2,6 +2,7 @@ package uspfilter
 
 import (
 	"net"
+	"net/netip"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -13,7 +14,7 @@ func TestLocalIPManager(t *testing.T) {
 	tests := []struct {
 		name      string
 		setupAddr wgaddr.Address
-		testIP    net.IP
+		testIP    netip.Addr
 		expected  bool
 	}{
 		{
@@ -25,7 +26,7 @@ func TestLocalIPManager(t *testing.T) {
 					Mask: net.CIDRMask(24, 32),
 				},
 			},
-			testIP:   net.ParseIP("127.0.0.2"),
+			testIP:   netip.MustParseAddr("127.0.0.2"),
 			expected: true,
 		},
 		{
@@ -37,7 +38,7 @@ func TestLocalIPManager(t *testing.T) {
 					Mask: net.CIDRMask(24, 32),
 				},
 			},
-			testIP:   net.ParseIP("127.0.0.1"),
+			testIP:   netip.MustParseAddr("127.0.0.1"),
 			expected: true,
 		},
 		{
@@ -49,7 +50,7 @@ func TestLocalIPManager(t *testing.T) {
 					Mask: net.CIDRMask(24, 32),
 				},
 			},
-			testIP:   net.ParseIP("127.255.255.255"),
+			testIP:   netip.MustParseAddr("127.255.255.255"),
 			expected: true,
 		},
 		{
@@ -61,7 +62,7 @@ func TestLocalIPManager(t *testing.T) {
 					Mask: net.CIDRMask(24, 32),
 				},
 			},
-			testIP:   net.ParseIP("192.168.1.1"),
+			testIP:   netip.MustParseAddr("192.168.1.1"),
 			expected: true,
 		},
 		{
@@ -73,7 +74,7 @@ func TestLocalIPManager(t *testing.T) {
 					Mask: net.CIDRMask(24, 32),
 				},
 			},
-			testIP:   net.ParseIP("192.168.1.2"),
+			testIP:   netip.MustParseAddr("192.168.1.2"),
 			expected: false,
 		},
 		{
@@ -85,7 +86,7 @@ func TestLocalIPManager(t *testing.T) {
 					Mask: net.CIDRMask(64, 128),
 				},
 			},
-			testIP:   net.ParseIP("fe80::1"),
+			testIP:   netip.MustParseAddr("fe80::1"),
 			expected: false,
 		},
 	}
@@ -174,7 +175,7 @@ func TestLocalIPManager_AllInterfaces(t *testing.T) {
 	t.Logf("Testing %d IPs", len(tests))
 	for _, tt := range tests {
 		t.Run(tt.ip, func(t *testing.T) {
-			result := manager.IsLocalIP(net.ParseIP(tt.ip))
+			result := manager.IsLocalIP(netip.MustParseAddr(tt.ip))
 			require.Equal(t, tt.expected, result, "IP: %s", tt.ip)
 		})
 	}

client/firewall/uspfilter/log/log.go

@@ -1,4 +1,4 @@
-// Package logger provides a high-performance, non-blocking logger for userspace networking
+// Package log provides a high-performance, non-blocking logger for userspace networking
 package log
 
 import (
@@ -13,13 +13,12 @@ import (
 )
 
 const (
-	maxBatchSize         = 1024 * 16  // 16KB max batch size
-	maxMessageSize       = 1024 * 2   // 2KB per message
-	bufferSize           = 1024 * 256 // 256KB ring buffer
+	maxBatchSize         = 1024 * 16
+	maxMessageSize       = 1024 * 2
 	defaultFlushInterval = 2 * time.Second
+	logChannelSize       = 1000
 )
 
-// Level represents log severity
 type Level uint32
 
 const (
@@ -42,32 +41,37 @@ var levelStrings = map[Level]string{
 	LevelTrace: "TRAC",
 }
 
-// Logger is a high-performance, non-blocking logger
-type Logger struct {
-	output    io.Writer
-	level     atomic.Uint32
-	buffer    *ringBuffer
-	shutdown  chan struct{}
-	closeOnce sync.Once
-	wg        sync.WaitGroup
-
-	// Reusable buffer pool for formatting messages
-	bufPool sync.Pool
+type logMessage struct {
+	level  Level
+	format string
+	args   []any
 }
 
+// Logger is a high-performance, non-blocking logger
+type Logger struct {
+	output     io.Writer
+	level      atomic.Uint32
+	msgChannel chan logMessage
+	shutdown   chan struct{}
+	closeOnce  sync.Once
+	wg         sync.WaitGroup
+	bufPool    sync.Pool
+}
+
+// NewFromLogrus creates a new Logger that writes to the same output as the given logrus logger
 func NewFromLogrus(logrusLogger *log.Logger) *Logger {
 	l := &Logger{
-		output:   logrusLogger.Out,
-		buffer:   newRingBuffer(bufferSize),
-		shutdown: make(chan struct{}),
+		output:     logrusLogger.Out,
+		msgChannel: make(chan logMessage, logChannelSize),
+		shutdown:   make(chan struct{}),
 		bufPool: sync.Pool{
-			New: func() interface{} {
-				// Pre-allocate buffer for message formatting
+			New: func() any {
 				b := make([]byte, 0, maxMessageSize)
 				return &b
 			},
 		},
 	}
+
 	logrusLevel := logrusLogger.GetLevel()
 	l.level.Store(uint32(logrusLevel))
 	level := levelStrings[Level(logrusLevel)]
@@ -79,97 +83,149 @@ func NewFromLogrus(logrusLogger *log.Logger) *Logger {
 	return l
 }
 
+// SetLevel sets the logging level
 func (l *Logger) SetLevel(level Level) {
 	l.level.Store(uint32(level))
-
 	log.Debugf("Set uspfilter logger loglevel to %v", levelStrings[level])
 }
 
-func (l *Logger) formatMessage(buf *[]byte, level Level, format string, args ...interface{}) {
-	*buf = (*buf)[:0]
-
-	// Timestamp
-	*buf = time.Now().AppendFormat(*buf, "2006-01-02T15:04:05-07:00")
-	*buf = append(*buf, ' ')
-
-	// Level
-	*buf = append(*buf, levelStrings[level]...)
-	*buf = append(*buf, ' ')
-
-	// Message
-	if len(args) > 0 {
-		*buf = append(*buf, fmt.Sprintf(format, args...)...)
-	} else {
-		*buf = append(*buf, format...)
+func (l *Logger) log(level Level, format string, args ...any) {
+	select {
+	case l.msgChannel <- logMessage{level: level, format: format, args: args}:
+	default:
 	}
-
-	*buf = append(*buf, '\n')
 }
 
-func (l *Logger) log(level Level, format string, args ...interface{}) {
-	bufp := l.bufPool.Get().(*[]byte)
-	l.formatMessage(bufp, level, format, args...)
-
-	if len(*bufp) > maxMessageSize {
-		*bufp = (*bufp)[:maxMessageSize]
-	}
-	_, _ = l.buffer.Write(*bufp)
-
-	l.bufPool.Put(bufp)
-}
-
-func (l *Logger) Error(format string, args ...interface{}) {
+// Error logs a message at error level
+func (l *Logger) Error(format string, args ...any) {
 	if l.level.Load() >= uint32(LevelError) {
 		l.log(LevelError, format, args...)
 	}
 }
 
-func (l *Logger) Warn(format string, args ...interface{}) {
+// Warn logs a message at warning level
+func (l *Logger) Warn(format string, args ...any) {
 	if l.level.Load() >= uint32(LevelWarn) {
 		l.log(LevelWarn, format, args...)
 	}
 }
 
-func (l *Logger) Info(format string, args ...interface{}) {
+// Info logs a message at info level
+func (l *Logger) Info(format string, args ...any) {
 	if l.level.Load() >= uint32(LevelInfo) {
 		l.log(LevelInfo, format, args...)
 	}
 }
 
-func (l *Logger) Debug(format string, args ...interface{}) {
+// Debug logs a message at debug level
+func (l *Logger) Debug(format string, args ...any) {
 	if l.level.Load() >= uint32(LevelDebug) {
 		l.log(LevelDebug, format, args...)
 	}
 }
 
-func (l *Logger) Trace(format string, args ...interface{}) {
+// Trace logs a message at trace level
+func (l *Logger) Trace(format string, args ...any) {
 	if l.level.Load() >= uint32(LevelTrace) {
 		l.log(LevelTrace, format, args...)
 	}
 }
 
-// worker periodically flushes the buffer
+func (l *Logger) formatMessage(buf *[]byte, level Level, format string, args ...any) {
+	*buf = (*buf)[:0]
+	*buf = time.Now().AppendFormat(*buf, "2006-01-02T15:04:05-07:00")
+	*buf = append(*buf, ' ')
+	*buf = append(*buf, levelStrings[level]...)
+	*buf = append(*buf, ' ')
+
+	var msg string
+	if len(args) > 0 {
+		msg = fmt.Sprintf(format, args...)
+	} else {
+		msg = format
+	}
+	*buf = append(*buf, msg...)
+	*buf = append(*buf, '\n')
+
+	if len(*buf) > maxMessageSize {
+		*buf = (*buf)[:maxMessageSize]
+	}
+}
+
+// processMessage handles a single log message and adds it to the buffer
+func (l *Logger) processMessage(msg logMessage, buffer *[]byte) {
+	bufp := l.bufPool.Get().(*[]byte)
+	defer l.bufPool.Put(bufp)
+
+	l.formatMessage(bufp, msg.level, msg.format, msg.args...)
+
+	if len(*buffer)+len(*bufp) > maxBatchSize {
+		_, _ = l.output.Write(*buffer)
+		*buffer = (*buffer)[:0]
+	}
+	*buffer = append(*buffer, *bufp...)
+}
+
+// flushBuffer writes the accumulated buffer to output
+func (l *Logger) flushBuffer(buffer *[]byte) {
+	if len(*buffer) > 0 {
+		_, _ = l.output.Write(*buffer)
+		*buffer = (*buffer)[:0]
+	}
+}
+
+// processBatch processes as many messages as possible without blocking
+func (l *Logger) processBatch(buffer *[]byte) {
+	for len(*buffer) < maxBatchSize {
+		select {
+		case msg := <-l.msgChannel:
+			l.processMessage(msg, buffer)
+		default:
+			return
+		}
+	}
+}
+
+// handleShutdown manages the graceful shutdown sequence with timeout
+func (l *Logger) handleShutdown(buffer *[]byte) {
+	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cancel()
+
+	for {
+		select {
+		case msg := <-l.msgChannel:
+			l.processMessage(msg, buffer)
+		case <-ctx.Done():
+			l.flushBuffer(buffer)
+			return
+		}
+
+		if len(l.msgChannel) == 0 {
+			l.flushBuffer(buffer)
+			return
+		}
+	}
+}
+
+// worker is the main goroutine that processes log messages
 func (l *Logger) worker() {
 	defer l.wg.Done()
 
 	ticker := time.NewTicker(defaultFlushInterval)
 	defer ticker.Stop()
 
-	buf := make([]byte, 0, maxBatchSize)
+	buffer := make([]byte, 0, maxBatchSize)
 
 	for {
 		select {
 		case <-l.shutdown:
+			l.handleShutdown(&buffer)
 			return
 		case <-ticker.C:
-			// Read accumulated messages
-			n, _ := l.buffer.Read(buf[:cap(buf)])
-			if n == 0 {
-				continue
-			}
-
-			// Write batch
-			_, _ = l.output.Write(buf[:n])
+			l.flushBuffer(&buffer)
+		case msg := <-l.msgChannel:
+			l.processMessage(msg, &buffer)
+			l.processBatch(&buffer)
 		}
 	}
 }

client/firewall/uspfilter/log/log_test.go

@@ -0,0 +1,121 @@
+package log_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/client/firewall/uspfilter/log"
+)
+
+type discard struct{}
+
+func (d *discard) Write(p []byte) (n int, err error) {
+	return len(p), nil
+}
+
+func BenchmarkLogger(b *testing.B) {
+	simpleMessage := "Connection established"
+
+	conntrackMessage := "TCP connection %s:%d -> %s:%d state changed to %d"
+	srcIP := "192.168.1.1"
+	srcPort := uint16(12345)
+	dstIP := "10.0.0.1"
+	dstPort := uint16(443)
+	state := 4 // TCPStateEstablished
+
+	complexMessage := "Packet inspection result: protocol=%s, direction=%s, flags=0x%x, sequence=%d, acknowledged=%d, payload_size=%d, fragmented=%v, connection_id=%s"
+	protocol := "TCP"
+	direction := "outbound"
+	flags := uint16(0x18) // ACK + PSH
+	sequence := uint32(123456789)
+	acknowledged := uint32(987654321)
+	payloadSize := 1460
+	fragmented := false
+	connID := "f7a12b3e-c456-7890-d123-456789abcdef"
+
+	b.Run("SimpleMessage", func(b *testing.B) {
+		logger := createTestLogger()
+		defer cleanupLogger(logger)
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			logger.Trace(simpleMessage)
+		}
+	})
+
+	b.Run("ConntrackMessage", func(b *testing.B) {
+		logger := createTestLogger()
+		defer cleanupLogger(logger)
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			logger.Trace(conntrackMessage, srcIP, srcPort, dstIP, dstPort, state)
+		}
+	})
+
+	b.Run("ComplexMessage", func(b *testing.B) {
+		logger := createTestLogger()
+		defer cleanupLogger(logger)
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			logger.Trace(complexMessage, protocol, direction, flags, sequence, acknowledged, payloadSize, fragmented, connID)
+		}
+	})
+}
+
+// BenchmarkLoggerParallel tests the logger under concurrent load
+func BenchmarkLoggerParallel(b *testing.B) {
+	logger := createTestLogger()
+	defer cleanupLogger(logger)
+
+	conntrackMessage := "TCP connection %s:%d -> %s:%d state changed to %d"
+	srcIP := "192.168.1.1"
+	srcPort := uint16(12345)
+	dstIP := "10.0.0.1"
+	dstPort := uint16(443)
+	state := 4
+
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			logger.Trace(conntrackMessage, srcIP, srcPort, dstIP, dstPort, state)
+		}
+	})
+}
+
+// BenchmarkLoggerBurst tests how the logger handles bursts of messages
+func BenchmarkLoggerBurst(b *testing.B) {
+	logger := createTestLogger()
+	defer cleanupLogger(logger)
+
+	conntrackMessage := "TCP connection %s:%d -> %s:%d state changed to %d"
+	srcIP := "192.168.1.1"
+	srcPort := uint16(12345)
+	dstIP := "10.0.0.1"
+	dstPort := uint16(443)
+	state := 4
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		for j := 0; j < 100; j++ {
+			logger.Trace(conntrackMessage, srcIP, srcPort, dstIP, dstPort, state)
+		}
+	}
+}
+
+func createTestLogger() *log.Logger {
+	logrusLogger := logrus.New()
+	logrusLogger.SetOutput(&discard{})
+	logrusLogger.SetLevel(logrus.TraceLevel)
+	return log.NewFromLogrus(logrusLogger)
+}
+
+func cleanupLogger(logger *log.Logger) {
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+	_ = logger.Stop(ctx)
+}

client/firewall/uspfilter/log/ringbuffer.go

@@ -1,85 +0,0 @@
-package log
-
-import "sync"
-
-// ringBuffer is a simple ring buffer implementation
-type ringBuffer struct {
-	buf  []byte
-	size int
-	r, w int64 // Read and write positions
-	mu   sync.Mutex
-}
-
-func newRingBuffer(size int) *ringBuffer {
-	return &ringBuffer{
-		buf:  make([]byte, size),
-		size: size,
-	}
-}
-
-func (r *ringBuffer) Write(p []byte) (n int, err error) {
-	if len(p) == 0 {
-		return 0, nil
-	}
-
-	r.mu.Lock()
-	defer r.mu.Unlock()
-
-	if len(p) > r.size {
-		p = p[:r.size]
-	}
-
-	n = len(p)
-
-	// Write data, handling wrap-around
-	pos := int(r.w % int64(r.size))
-	writeLen := min(len(p), r.size-pos)
-	copy(r.buf[pos:], p[:writeLen])
-
-	// If we have more data and need to wrap around
-	if writeLen < len(p) {
-		copy(r.buf, p[writeLen:])
-	}
-
-	// Update write position
-	r.w += int64(n)
-
-	return n, nil
-}
-
-func (r *ringBuffer) Read(p []byte) (n int, err error) {
-	r.mu.Lock()
-	defer r.mu.Unlock()
-
-	if r.w == r.r {
-		return 0, nil
-	}
-
-	// Calculate available data accounting for wraparound
-	available := int(r.w - r.r)
-	if available < 0 {
-		available += r.size
-	}
-	available = min(available, r.size)
-
-	// Limit read to buffer size
-	toRead := min(available, len(p))
-	if toRead == 0 {
-		return 0, nil
-	}
-
-	// Read data, handling wrap-around
-	pos := int(r.r % int64(r.size))
-	readLen := min(toRead, r.size-pos)
-	n = copy(p, r.buf[pos:pos+readLen])
-
-	// If we need more data and need to wrap around
-	if readLen < toRead {
-		n += copy(p[readLen:toRead], r.buf[:toRead-readLen])
-	}
-
-	// Update read position
-	r.r += int64(n)
-
-	return n, nil
-}

client/firewall/uspfilter/rule.go

@@ -1,7 +1,6 @@
 package uspfilter
 
 import (
-	"net"
 	"net/netip"
 
 	"github.com/google/gopacket"
@@ -12,14 +11,14 @@ import (
 // PeerRule to handle management of rules
 type PeerRule struct {
 	id         string
-	ip         net.IP
+	mgmtId     []byte
+	ip         netip.Addr
 	ipLayer    gopacket.LayerType
 	matchByIP  bool
 	protoLayer gopacket.LayerType
 	sPort      *firewall.Port
 	dPort      *firewall.Port
 	drop       bool
-	comment    string
 
 	udpHook func([]byte) bool
 }
@@ -31,6 +30,7 @@ func (r *PeerRule) ID() string {
 
 type RouteRule struct {
 	id          string
+	mgmtId      []byte
 	sources     []netip.Prefix
 	destination netip.Prefix
 	proto       firewall.Protocol

client/firewall/uspfilter/tracer.go

@@ -2,7 +2,7 @@ package uspfilter
 
 import (
 	"fmt"
-	"net"
+	"net/netip"
 	"time"
 
 	"github.com/google/gopacket"
@@ -53,8 +53,8 @@ type TraceResult struct {
 }
 
 type PacketTrace struct {
-	SourceIP        net.IP
-	DestinationIP   net.IP
+	SourceIP        netip.Addr
+	DestinationIP   netip.Addr
 	Protocol        string
 	SourcePort      uint16
 	DestinationPort uint16
@@ -72,8 +72,8 @@ type TCPState struct {
 }
 
 type PacketBuilder struct {
-	SrcIP       net.IP
-	DstIP       net.IP
+	SrcIP       netip.Addr
+	DstIP       netip.Addr
 	Protocol    fw.Protocol
 	SrcPort     uint16
 	DstPort     uint16
@@ -126,8 +126,8 @@ func (p *PacketBuilder) buildIPLayer() *layers.IPv4 {
 		Version:  4,
 		TTL:      64,
 		Protocol: layers.IPProtocol(getIPProtocolNumber(p.Protocol)),
-		SrcIP:    p.SrcIP,
-		DstIP:    p.DstIP,
+		SrcIP:    p.SrcIP.AsSlice(),
+		DstIP:    p.DstIP.AsSlice(),
 	}
 }
 
@@ -260,28 +260,30 @@ func (m *Manager) TracePacket(packetData []byte, direction fw.RuleDirection) *Pa
 	return m.traceInbound(packetData, trace, d, srcIP, dstIP)
 }
 
-func (m *Manager) traceInbound(packetData []byte, trace *PacketTrace, d *decoder, srcIP net.IP, dstIP net.IP) *PacketTrace {
+func (m *Manager) traceInbound(packetData []byte, trace *PacketTrace, d *decoder, srcIP netip.Addr, dstIP netip.Addr) *PacketTrace {
 	if m.stateful && m.handleConntrackState(trace, d, srcIP, dstIP) {
 		return trace
 	}
 
-	if m.handleLocalDelivery(trace, packetData, d, srcIP, dstIP) {
-		return trace
+	if m.localipmanager.IsLocalIP(dstIP) {
+		if m.handleLocalDelivery(trace, packetData, d, srcIP, dstIP) {
+			return trace
+		}
 	}
 
 	if !m.handleRouting(trace) {
 		return trace
 	}
 
-	if m.nativeRouter {
+	if m.nativeRouter.Load() {
 		return m.handleNativeRouter(trace)
 	}
 
 	return m.handleRouteACLs(trace, d, srcIP, dstIP)
 }
 
-func (m *Manager) handleConntrackState(trace *PacketTrace, d *decoder, srcIP, dstIP net.IP) bool {
-	allowed := m.isValidTrackedConnection(d, srcIP, dstIP)
+func (m *Manager) handleConntrackState(trace *PacketTrace, d *decoder, srcIP, dstIP netip.Addr) bool {
+	allowed := m.isValidTrackedConnection(d, srcIP, dstIP, 0)
 	msg := "No existing connection found"
 	if allowed {
 		msg = m.buildConntrackStateMessage(d)
@@ -309,32 +311,46 @@ func (m *Manager) buildConntrackStateMessage(d *decoder) string {
 	return msg
 }
 
-func (m *Manager) handleLocalDelivery(trace *PacketTrace, packetData []byte, d *decoder, srcIP, dstIP net.IP) bool {
-	if !m.localForwarding {
-		trace.AddResult(StageRouting, "Local forwarding disabled", false)
-		trace.AddResult(StageCompleted, "Packet dropped - local forwarding disabled", false)
+func (m *Manager) handleLocalDelivery(trace *PacketTrace, packetData []byte, d *decoder, srcIP, dstIP netip.Addr) bool {
+	trace.AddResult(StageRouting, "Packet destined for local delivery", true)
+
+	ruleId, blocked := m.peerACLsBlock(srcIP, packetData, m.incomingRules, d)
+
+	strRuleId := "<no id>"
+	if ruleId != nil {
+		strRuleId = string(ruleId)
+	}
+	msg := fmt.Sprintf("Allowed by peer ACL rules (%s)", strRuleId)
+	if blocked {
+		msg = fmt.Sprintf("Blocked by peer ACL rules (%s)", strRuleId)
+		trace.AddResult(StagePeerACL, msg, false)
+		trace.AddResult(StageCompleted, "Packet dropped - ACL denied", false)
 		return true
 	}
 
-	trace.AddResult(StageRouting, "Packet destined for local delivery", true)
-	blocked := m.peerACLsBlock(srcIP, packetData, m.incomingRules, d)
-
-	msg := "Allowed by peer ACL rules"
-	if blocked {
-		msg = "Blocked by peer ACL rules"
-	}
-	trace.AddResult(StagePeerACL, msg, !blocked)
+	trace.AddResult(StagePeerACL, msg, true)
 
+	// Handle netstack mode
 	if m.netstack {
-		m.addForwardingResult(trace, "proxy-local", "127.0.0.1", !blocked)
+		switch {
+		case !m.localForwarding:
+			trace.AddResult(StageCompleted, "Packet sent to virtual stack", true)
+		case m.forwarder.Load() != nil:
+			m.addForwardingResult(trace, "proxy-local", "127.0.0.1", true)
+			trace.AddResult(StageCompleted, msgProcessingCompleted, true)
+		default:
+			trace.AddResult(StageCompleted, "Packet dropped - forwarder not initialized", false)
+		}
+		return true
 	}
 
-	trace.AddResult(StageCompleted, msgProcessingCompleted, !blocked)
+	// In normal mode, packets are allowed through for local delivery
+	trace.AddResult(StageCompleted, msgProcessingCompleted, true)
 	return true
 }
 
 func (m *Manager) handleRouting(trace *PacketTrace) bool {
-	if !m.routingEnabled {
+	if !m.routingEnabled.Load() {
 		trace.AddResult(StageRouting, "Routing disabled", false)
 		trace.AddResult(StageCompleted, "Packet dropped - routing disabled", false)
 		return false
@@ -350,18 +366,23 @@ func (m *Manager) handleNativeRouter(trace *PacketTrace) *PacketTrace {
 	return trace
 }
 
-func (m *Manager) handleRouteACLs(trace *PacketTrace, d *decoder, srcIP, dstIP net.IP) *PacketTrace {
-	proto := getProtocolFromPacket(d)
+func (m *Manager) handleRouteACLs(trace *PacketTrace, d *decoder, srcIP, dstIP netip.Addr) *PacketTrace {
+	proto, _ := getProtocolFromPacket(d)
 	srcPort, dstPort := getPortsFromPacket(d)
-	allowed := m.routeACLsPass(srcIP, dstIP, proto, srcPort, dstPort)
+	id, allowed := m.routeACLsPass(srcIP, dstIP, proto, srcPort, dstPort)
 
-	msg := "Allowed by route ACLs"
+	strId := string(id)
+	if id == nil {
+		strId = "<no id>"
+	}
+
+	msg := fmt.Sprintf("Allowed by route ACLs (%s)", strId)
 	if !allowed {
-		msg = "Blocked by route ACLs"
+		msg = fmt.Sprintf("Blocked by route ACLs (%s)", strId)
 	}
 	trace.AddResult(StageRouteACL, msg, allowed)
 
-	if allowed && m.forwarder != nil {
+	if allowed && m.forwarder.Load() != nil {
 		m.addForwardingResult(trace, "proxy-remote", fmt.Sprintf("%s:%d", dstIP, dstPort), true)
 	}
 
@@ -380,7 +401,7 @@ func (m *Manager) addForwardingResult(trace *PacketTrace, action, remoteAddr str
 
 func (m *Manager) traceOutbound(packetData []byte, trace *PacketTrace) *PacketTrace {
 	// will create or update the connection state
-	dropped := m.processOutgoingHooks(packetData)
+	dropped := m.processOutgoingHooks(packetData, 0)
 	if dropped {
 		trace.AddResult(StageCompleted, "Packet dropped by outgoing hook", false)
 	} else {

client/firewall/uspfilter/tracer_test.go

@@ -0,0 +1,440 @@
+package uspfilter
+
+import (
+	"net"
+	"net/netip"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	fw "github.com/netbirdio/netbird/client/firewall/manager"
+	"github.com/netbirdio/netbird/client/firewall/uspfilter/conntrack"
+	"github.com/netbirdio/netbird/client/firewall/uspfilter/forwarder"
+	"github.com/netbirdio/netbird/client/iface/device"
+	"github.com/netbirdio/netbird/client/iface/wgaddr"
+)
+
+func verifyTraceStages(t *testing.T, trace *PacketTrace, expectedStages []PacketStage) {
+	t.Logf("Trace results: %v", trace.Results)
+	actualStages := make([]PacketStage, 0, len(trace.Results))
+	for _, result := range trace.Results {
+		actualStages = append(actualStages, result.Stage)
+		t.Logf("Stage: %s, Message: %s, Allowed: %v", result.Stage, result.Message, result.Allowed)
+	}
+
+	require.ElementsMatch(t, expectedStages, actualStages, "Trace stages don't match expected stages")
+}
+
+func verifyFinalDisposition(t *testing.T, trace *PacketTrace, expectedAllowed bool) {
+	require.NotEmpty(t, trace.Results, "Trace should have results")
+	lastResult := trace.Results[len(trace.Results)-1]
+	require.Equal(t, StageCompleted, lastResult.Stage, "Last stage should be 'Completed'")
+	require.Equal(t, expectedAllowed, lastResult.Allowed, "Final disposition incorrect")
+}
+
+func TestTracePacket(t *testing.T) {
+	setupTracerTest := func(statefulMode bool) *Manager {
+		ifaceMock := &IFaceMock{
+			SetFilterFunc: func(device.PacketFilter) error { return nil },
+			AddressFunc: func() wgaddr.Address {
+				return wgaddr.Address{
+					IP: net.ParseIP("100.10.0.100"),
+					Network: &net.IPNet{
+						IP:   net.ParseIP("100.10.0.0"),
+						Mask: net.CIDRMask(16, 32),
+					},
+				}
+			},
+		}
+
+		m, err := Create(ifaceMock, false, flowLogger)
+		require.NoError(t, err)
+
+		if !statefulMode {
+			m.stateful = false
+		}
+
+		return m
+	}
+
+	createPacketBuilder := func(srcIP, dstIP string, protocol fw.Protocol, srcPort, dstPort uint16, direction fw.RuleDirection) *PacketBuilder {
+		builder := &PacketBuilder{
+			SrcIP:     netip.MustParseAddr(srcIP),
+			DstIP:     netip.MustParseAddr(dstIP),
+			Protocol:  protocol,
+			SrcPort:   srcPort,
+			DstPort:   dstPort,
+			Direction: direction,
+		}
+
+		if protocol == "tcp" {
+			builder.TCPState = &TCPState{SYN: true}
+		}
+
+		return builder
+	}
+
+	createICMPPacketBuilder := func(srcIP, dstIP string, icmpType, icmpCode uint8, direction fw.RuleDirection) *PacketBuilder {
+		return &PacketBuilder{
+			SrcIP:     netip.MustParseAddr(srcIP),
+			DstIP:     netip.MustParseAddr(dstIP),
+			Protocol:  "icmp",
+			ICMPType:  icmpType,
+			ICMPCode:  icmpCode,
+			Direction: direction,
+		}
+	}
+
+	testCases := []struct {
+		name           string
+		setup          func(*Manager)
+		packetBuilder  func() *PacketBuilder
+		expectedStages []PacketStage
+		expectedAllow  bool
+	}{
+		{
+			name: "LocalTraffic_ACLAllowed",
+			setup: func(m *Manager) {
+				ip := net.ParseIP("1.1.1.1")
+				proto := fw.ProtocolTCP
+				port := &fw.Port{Values: []uint16{80}}
+				action := fw.ActionAccept
+				_, err := m.AddPeerFiltering(nil, ip, proto, nil, port, action, "")
+				require.NoError(t, err)
+			},
+			packetBuilder: func() *PacketBuilder {
+				return createPacketBuilder("1.1.1.1", "100.10.0.100", "tcp", 12345, 80, fw.RuleDirectionIN)
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageConntrack,
+				StageRouting,
+				StagePeerACL,
+				StageCompleted,
+			},
+			expectedAllow: true,
+		},
+		{
+			name: "LocalTraffic_ACLDenied",
+			setup: func(m *Manager) {
+				ip := net.ParseIP("1.1.1.1")
+				proto := fw.ProtocolTCP
+				port := &fw.Port{Values: []uint16{80}}
+				action := fw.ActionDrop
+				_, err := m.AddPeerFiltering(nil, ip, proto, nil, port, action, "")
+				require.NoError(t, err)
+			},
+			packetBuilder: func() *PacketBuilder {
+				return createPacketBuilder("1.1.1.1", "100.10.0.100", "tcp", 12345, 80, fw.RuleDirectionIN)
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageConntrack,
+				StageRouting,
+				StagePeerACL,
+				StageCompleted,
+			},
+			expectedAllow: false,
+		},
+		{
+			name: "LocalTraffic_WithForwarder",
+			setup: func(m *Manager) {
+				m.netstack = true
+				m.localForwarding = true
+
+				m.forwarder.Store(&forwarder.Forwarder{})
+
+				ip := net.ParseIP("1.1.1.1")
+				proto := fw.ProtocolTCP
+				port := &fw.Port{Values: []uint16{80}}
+				action := fw.ActionAccept
+				_, err := m.AddPeerFiltering(nil, ip, proto, nil, port, action, "")
+				require.NoError(t, err)
+			},
+			packetBuilder: func() *PacketBuilder {
+				return createPacketBuilder("1.1.1.1", "100.10.0.100", "tcp", 12345, 80, fw.RuleDirectionIN)
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageConntrack,
+				StageRouting,
+				StagePeerACL,
+				StageForwarding,
+				StageCompleted,
+			},
+			expectedAllow: true,
+		},
+		{
+			name: "LocalTraffic_WithoutForwarder",
+			setup: func(m *Manager) {
+				m.netstack = true
+				m.localForwarding = false
+
+				ip := net.ParseIP("1.1.1.1")
+				proto := fw.ProtocolTCP
+				port := &fw.Port{Values: []uint16{80}}
+				action := fw.ActionAccept
+				_, err := m.AddPeerFiltering(nil, ip, proto, nil, port, action, "")
+				require.NoError(t, err)
+			},
+			packetBuilder: func() *PacketBuilder {
+				return createPacketBuilder("1.1.1.1", "100.10.0.100", "tcp", 12345, 80, fw.RuleDirectionIN)
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageConntrack,
+				StageRouting,
+				StagePeerACL,
+				StageCompleted,
+			},
+			expectedAllow: true,
+		},
+		{
+			name: "RoutedTraffic_ACLAllowed",
+			setup: func(m *Manager) {
+				m.routingEnabled.Store(true)
+				m.nativeRouter.Store(false)
+
+				m.forwarder.Store(&forwarder.Forwarder{})
+
+				src := netip.PrefixFrom(netip.AddrFrom4([4]byte{1, 1, 1, 1}), 32)
+				dst := netip.PrefixFrom(netip.AddrFrom4([4]byte{172, 17, 0, 2}), 32)
+				_, err := m.AddRouteFiltering(nil, []netip.Prefix{src}, dst, fw.ProtocolTCP, nil, &fw.Port{Values: []uint16{80}}, fw.ActionAccept)
+				require.NoError(t, err)
+			},
+			packetBuilder: func() *PacketBuilder {
+				return createPacketBuilder("1.1.1.1", "172.17.0.2", "tcp", 12345, 80, fw.RuleDirectionIN)
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageConntrack,
+				StageRouting,
+				StageRouteACL,
+				StageForwarding,
+				StageCompleted,
+			},
+			expectedAllow: true,
+		},
+		{
+			name: "RoutedTraffic_ACLDenied",
+			setup: func(m *Manager) {
+				m.routingEnabled.Store(true)
+				m.nativeRouter.Store(false)
+
+				src := netip.PrefixFrom(netip.AddrFrom4([4]byte{1, 1, 1, 1}), 32)
+				dst := netip.PrefixFrom(netip.AddrFrom4([4]byte{172, 17, 0, 2}), 32)
+				_, err := m.AddRouteFiltering(nil, []netip.Prefix{src}, dst, fw.ProtocolTCP, nil, &fw.Port{Values: []uint16{80}}, fw.ActionDrop)
+				require.NoError(t, err)
+			},
+			packetBuilder: func() *PacketBuilder {
+				return createPacketBuilder("1.1.1.1", "172.17.0.2", "tcp", 12345, 80, fw.RuleDirectionIN)
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageConntrack,
+				StageRouting,
+				StageRouteACL,
+				StageCompleted,
+			},
+			expectedAllow: false,
+		},
+		{
+			name: "RoutedTraffic_NativeRouter",
+			setup: func(m *Manager) {
+				m.routingEnabled.Store(true)
+				m.nativeRouter.Store(true)
+			},
+			packetBuilder: func() *PacketBuilder {
+				return createPacketBuilder("1.1.1.1", "172.17.0.2", "tcp", 12345, 80, fw.RuleDirectionIN)
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageConntrack,
+				StageRouting,
+				StageRouteACL,
+				StageForwarding,
+				StageCompleted,
+			},
+			expectedAllow: true,
+		},
+		{
+			name: "RoutedTraffic_RoutingDisabled",
+			setup: func(m *Manager) {
+				m.routingEnabled.Store(false)
+			},
+			packetBuilder: func() *PacketBuilder {
+				return createPacketBuilder("1.1.1.1", "172.17.0.2", "tcp", 12345, 80, fw.RuleDirectionIN)
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageConntrack,
+				StageRouting,
+				StageCompleted,
+			},
+			expectedAllow: false,
+		},
+		{
+			name: "ConnectionTracking_Hit",
+			setup: func(m *Manager) {
+				srcIP := netip.MustParseAddr("100.10.0.100")
+				dstIP := netip.MustParseAddr("1.1.1.1")
+				srcPort := uint16(12345)
+				dstPort := uint16(80)
+
+				m.tcpTracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, conntrack.TCPSyn, 0)
+			},
+			packetBuilder: func() *PacketBuilder {
+				pb := createPacketBuilder("1.1.1.1", "100.10.0.100", "tcp", 80, 12345, fw.RuleDirectionIN)
+				pb.TCPState = &TCPState{SYN: true, ACK: true}
+				return pb
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageConntrack,
+				StageCompleted,
+			},
+			expectedAllow: true,
+		},
+		{
+			name: "OutboundTraffic",
+			setup: func(m *Manager) {
+			},
+			packetBuilder: func() *PacketBuilder {
+				return createPacketBuilder("100.10.0.100", "1.1.1.1", "tcp", 12345, 80, fw.RuleDirectionOUT)
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageCompleted,
+			},
+			expectedAllow: true,
+		},
+		{
+			name: "ICMPEchoRequest",
+			setup: func(m *Manager) {
+				ip := net.ParseIP("1.1.1.1")
+				proto := fw.ProtocolICMP
+				action := fw.ActionAccept
+				_, err := m.AddPeerFiltering(nil, ip, proto, nil, nil, action, "")
+				require.NoError(t, err)
+			},
+			packetBuilder: func() *PacketBuilder {
+				return createICMPPacketBuilder("1.1.1.1", "100.10.0.100", 8, 0, fw.RuleDirectionIN)
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageConntrack,
+				StageRouting,
+				StagePeerACL,
+				StageCompleted,
+			},
+			expectedAllow: true,
+		},
+		{
+			name: "ICMPDestinationUnreachable",
+			setup: func(m *Manager) {
+				ip := net.ParseIP("1.1.1.1")
+				proto := fw.ProtocolICMP
+				action := fw.ActionDrop
+				_, err := m.AddPeerFiltering(nil, ip, proto, nil, nil, action, "")
+				require.NoError(t, err)
+			},
+			packetBuilder: func() *PacketBuilder {
+				return createICMPPacketBuilder("1.1.1.1", "100.10.0.100", 3, 0, fw.RuleDirectionIN)
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageConntrack,
+				StageRouting,
+				StagePeerACL,
+				StageCompleted,
+			},
+			expectedAllow: true,
+		},
+		{
+			name: "UDPTraffic_WithoutHook",
+			setup: func(m *Manager) {
+				ip := net.ParseIP("1.1.1.1")
+				proto := fw.ProtocolUDP
+				port := &fw.Port{Values: []uint16{53}}
+				action := fw.ActionAccept
+				_, err := m.AddPeerFiltering(nil, ip, proto, nil, port, action, "")
+				require.NoError(t, err)
+			},
+			packetBuilder: func() *PacketBuilder {
+				return createPacketBuilder("1.1.1.1", "100.10.0.100", "udp", 12345, 53, fw.RuleDirectionIN)
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageConntrack,
+				StageRouting,
+				StagePeerACL,
+				StageCompleted,
+			},
+			expectedAllow: true,
+		},
+		{
+			name: "UDPTraffic_WithHook",
+			setup: func(m *Manager) {
+				hookFunc := func([]byte) bool {
+					return true
+				}
+				m.AddUDPPacketHook(true, netip.MustParseAddr("1.1.1.1"), 53, hookFunc)
+			},
+			packetBuilder: func() *PacketBuilder {
+				return createPacketBuilder("1.1.1.1", "100.10.0.100", "udp", 12345, 53, fw.RuleDirectionIN)
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageConntrack,
+				StageRouting,
+				StagePeerACL,
+				StageCompleted,
+			},
+			expectedAllow: false,
+		},
+		{
+			name: "StatefulDisabled_NoTracking",
+			setup: func(m *Manager) {
+				m.stateful = false
+
+				ip := net.ParseIP("1.1.1.1")
+				proto := fw.ProtocolTCP
+				port := &fw.Port{Values: []uint16{80}}
+				action := fw.ActionDrop
+				_, err := m.AddPeerFiltering(nil, ip, proto, nil, port, action, "")
+				require.NoError(t, err)
+			},
+			packetBuilder: func() *PacketBuilder {
+				return createPacketBuilder("1.1.1.1", "100.10.0.100", "tcp", 12345, 80, fw.RuleDirectionIN)
+			},
+			expectedStages: []PacketStage{
+				StageReceived,
+				StageRouting,
+				StagePeerACL,
+				StageCompleted,
+			},
+			expectedAllow: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			m := setupTracerTest(true)
+
+			tc.setup(m)
+
+			require.True(t, m.localipmanager.IsLocalIP(netip.MustParseAddr("100.10.0.100")),
+				"100.10.0.100 should be recognized as a local IP")
+			require.False(t, m.localipmanager.IsLocalIP(netip.MustParseAddr("172.17.0.2")),
+				"172.17.0.2 should not be recognized as a local IP")
+
+			pb := tc.packetBuilder()
+
+			trace, err := m.TracePacketFromBuilder(pb)
+			require.NoError(t, err)
+
+			verifyTraceStages(t, trace, tc.expectedStages)
+			verifyFinalDisposition(t, trace, tc.expectedAllow)
+		})
+	}
+}

client/firewall/uspfilter/uspfilter.go

@@ -10,6 +10,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	"github.com/google/gopacket"
 	"github.com/google/gopacket/layers"
@@ -22,6 +23,7 @@ import (
 	"github.com/netbirdio/netbird/client/firewall/uspfilter/forwarder"
 	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
 	"github.com/netbirdio/netbird/client/iface/netstack"
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 )
 
@@ -65,9 +67,9 @@ func (r RouteRules) Sort() {
 // Manager userspace firewall manager
 type Manager struct {
 	// outgoingRules is used for hooks only
-	outgoingRules map[string]RuleSet
+	outgoingRules map[netip.Addr]RuleSet
 	// incomingRules is used for filtering and hooks
-	incomingRules  map[string]RuleSet
+	incomingRules  map[netip.Addr]RuleSet
 	routeRules     RouteRules
 	wgNetwork      *net.IPNet
 	decoders       sync.Pool
@@ -79,9 +81,9 @@ type Manager struct {
 	// indicates whether server routes are disabled
 	disableServerRoutes bool
 	// indicates whether we forward packets not destined for ourselves
-	routingEnabled bool
+	routingEnabled atomic.Bool
 	// indicates whether we leave forwarding and filtering to the native firewall
-	nativeRouter bool
+	nativeRouter atomic.Bool
 	// indicates whether we track outbound connections
 	stateful bool
 	// indicates whether wireguards runs in netstack mode
@@ -94,8 +96,9 @@ type Manager struct {
 	udpTracker  *conntrack.UDPTracker
 	icmpTracker *conntrack.ICMPTracker
 	tcpTracker  *conntrack.TCPTracker
-	forwarder   *forwarder.Forwarder
+	forwarder   atomic.Pointer[forwarder.Forwarder]
 	logger      *nblog.Logger
+	flowLogger  nftypes.FlowLogger
 }
 
 // decoder for packages
@@ -112,16 +115,16 @@ type decoder struct {
 }
 
 // Create userspace firewall manager constructor
-func Create(iface common.IFaceMapper, disableServerRoutes bool) (*Manager, error) {
-	return create(iface, nil, disableServerRoutes)
+func Create(iface common.IFaceMapper, disableServerRoutes bool, flowLogger nftypes.FlowLogger) (*Manager, error) {
+	return create(iface, nil, disableServerRoutes, flowLogger)
 }
 
-func CreateWithNativeFirewall(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableServerRoutes bool) (*Manager, error) {
+func CreateWithNativeFirewall(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableServerRoutes bool, flowLogger nftypes.FlowLogger) (*Manager, error) {
 	if nativeFirewall == nil {
 		return nil, errors.New("native firewall is nil")
 	}
 
-	mgr, err := create(iface, nativeFirewall, disableServerRoutes)
+	mgr, err := create(iface, nativeFirewall, disableServerRoutes, flowLogger)
 	if err != nil {
 		return nil, err
 	}
@@ -148,7 +151,7 @@ func parseCreateEnv() (bool, bool) {
 	return disableConntrack, enableLocalForwarding
 }
 
-func create(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableServerRoutes bool) (*Manager, error) {
+func create(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableServerRoutes bool, flowLogger nftypes.FlowLogger) (*Manager, error) {
 	disableConntrack, enableLocalForwarding := parseCreateEnv()
 
 	m := &Manager{
@@ -166,17 +169,18 @@ func create(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableSe
 			},
 		},
 		nativeFirewall:      nativeFirewall,
-		outgoingRules:       make(map[string]RuleSet),
-		incomingRules:       make(map[string]RuleSet),
+		outgoingRules:       make(map[netip.Addr]RuleSet),
+		incomingRules:       make(map[netip.Addr]RuleSet),
 		wgIface:             iface,
 		localipmanager:      newLocalIPManager(),
 		disableServerRoutes: disableServerRoutes,
-		routingEnabled:      false,
 		stateful:            !disableConntrack,
 		logger:              nblog.NewFromLogrus(log.StandardLogger()),
+		flowLogger:          flowLogger,
 		netstack:            netstack.IsEnabled(),
 		localForwarding:     enableLocalForwarding,
 	}
+	m.routingEnabled.Store(false)
 
 	if err := m.localipmanager.UpdateLocalIPs(iface); err != nil {
 		return nil, fmt.Errorf("update local IPs: %w", err)
@@ -185,9 +189,9 @@ func create(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableSe
 	if disableConntrack {
 		log.Info("conntrack is disabled")
 	} else {
-		m.udpTracker = conntrack.NewUDPTracker(conntrack.DefaultUDPTimeout, m.logger)
-		m.icmpTracker = conntrack.NewICMPTracker(conntrack.DefaultICMPTimeout, m.logger)
-		m.tcpTracker = conntrack.NewTCPTracker(conntrack.DefaultTCPTimeout, m.logger)
+		m.udpTracker = conntrack.NewUDPTracker(conntrack.DefaultUDPTimeout, m.logger, flowLogger)
+		m.icmpTracker = conntrack.NewICMPTracker(conntrack.DefaultICMPTimeout, m.logger, flowLogger)
+		m.tcpTracker = conntrack.NewTCPTracker(conntrack.DefaultTCPTimeout, m.logger, flowLogger)
 	}
 
 	// netstack needs the forwarder for local traffic
@@ -208,7 +212,7 @@ func create(iface common.IFaceMapper, nativeFirewall firewall.Manager, disableSe
 }
 
 func (m *Manager) blockInvalidRouted(iface common.IFaceMapper) error {
-	if m.forwarder == nil {
+	if m.forwarder.Load() == nil {
 		return nil
 	}
 	wgPrefix, err := netip.ParsePrefix(iface.Address().Network.String())
@@ -218,6 +222,7 @@ func (m *Manager) blockInvalidRouted(iface common.IFaceMapper) error {
 	log.Debugf("blocking invalid routed traffic for %s", wgPrefix)
 
 	if _, err := m.AddRouteFiltering(
+		nil,
 		[]netip.Prefix{netip.PrefixFrom(netip.IPv4Unspecified(), 0)},
 		wgPrefix,
 		firewall.ProtocolALL,
@@ -251,20 +256,20 @@ func (m *Manager) determineRouting() error {
 
 	switch {
 	case disableUspRouting:
-		m.routingEnabled = false
-		m.nativeRouter = false
+		m.routingEnabled.Store(false)
+		m.nativeRouter.Store(false)
 		log.Info("userspace routing is disabled")
 
 	case m.disableServerRoutes:
 		//  if server routes are disabled we will let packets pass to the native stack
-		m.routingEnabled = true
-		m.nativeRouter = true
+		m.routingEnabled.Store(true)
+		m.nativeRouter.Store(true)
 
 		log.Info("server routes are disabled")
 
 	case forceUserspaceRouter:
-		m.routingEnabled = true
-		m.nativeRouter = false
+		m.routingEnabled.Store(true)
+		m.nativeRouter.Store(false)
 
 		log.Info("userspace routing is forced")
 
@@ -272,19 +277,19 @@ func (m *Manager) determineRouting() error {
 		// if the OS supports routing natively, then we don't need to filter/route ourselves
 		// netstack mode won't support native routing as there is no interface
 
-		m.routingEnabled = true
-		m.nativeRouter = true
+		m.routingEnabled.Store(true)
+		m.nativeRouter.Store(true)
 
 		log.Info("native routing is enabled")
 
 	default:
-		m.routingEnabled = true
-		m.nativeRouter = false
+		m.routingEnabled.Store(true)
+		m.nativeRouter.Store(false)
 
 		log.Info("userspace routing enabled by default")
 	}
 
-	if m.routingEnabled && !m.nativeRouter {
+	if m.routingEnabled.Load() && !m.nativeRouter.Load() {
 		return m.initForwarder()
 	}
 
@@ -293,31 +298,31 @@ func (m *Manager) determineRouting() error {
 
 // initForwarder initializes the forwarder, it disables routing on errors
 func (m *Manager) initForwarder() error {
-	if m.forwarder != nil {
+	if m.forwarder.Load() != nil {
 		return nil
 	}
 
 	// Only supported in userspace mode as we need to inject packets back into wireguard directly
 	intf := m.wgIface.GetWGDevice()
 	if intf == nil {
-		m.routingEnabled = false
+		m.routingEnabled.Store(false)
 		return errors.New("forwarding not supported")
 	}
 
-	forwarder, err := forwarder.New(m.wgIface, m.logger, m.netstack)
+	forwarder, err := forwarder.New(m.wgIface, m.logger, m.flowLogger, m.netstack)
 	if err != nil {
-		m.routingEnabled = false
+		m.routingEnabled.Store(false)
 		return fmt.Errorf("create forwarder: %w", err)
 	}
 
-	m.forwarder = forwarder
+	m.forwarder.Store(forwarder)
 
 	log.Debug("forwarder initialized")
 
 	return nil
 }
 
-func (m *Manager) Init(*statemanager.Manager) error {
+func (m *Manager) Init(statemanager.Manager) error {
 	return nil
 }
 
@@ -326,7 +331,7 @@ func (m *Manager) IsServerRouteSupported() bool {
 }
 
 func (m *Manager) AddNatRule(pair firewall.RouterPair) error {
-	if m.nativeRouter && m.nativeFirewall != nil {
+	if m.nativeRouter.Load() && m.nativeFirewall != nil {
 		return m.nativeFirewall.AddNatRule(pair)
 	}
 
@@ -337,7 +342,7 @@ func (m *Manager) AddNatRule(pair firewall.RouterPair) error {
 
 // RemoveNatRule removes a routing firewall rule
 func (m *Manager) RemoveNatRule(pair firewall.RouterPair) error {
-	if m.nativeRouter && m.nativeFirewall != nil {
+	if m.nativeRouter.Load() && m.nativeFirewall != nil {
 		return m.nativeFirewall.RemoveNatRule(pair)
 	}
 	return nil
@@ -348,25 +353,31 @@ func (m *Manager) RemoveNatRule(pair firewall.RouterPair) error {
 // If comment argument is empty firewall manager should set
 // rule ID as comment for the rule
 func (m *Manager) AddPeerFiltering(
+	id []byte,
 	ip net.IP,
 	proto firewall.Protocol,
 	sPort *firewall.Port,
 	dPort *firewall.Port,
 	action firewall.Action,
 	_ string,
-	comment string,
 ) ([]firewall.Rule, error) {
+	// TODO: fix in upper layers
+	i, ok := netip.AddrFromSlice(ip)
+	if !ok {
+		return nil, fmt.Errorf("invalid IP: %s", ip)
+	}
+
+	i = i.Unmap()
 	r := PeerRule{
 		id:        uuid.New().String(),
-		ip:        ip,
+		mgmtId:    id,
+		ip:        i,
 		ipLayer:   layers.LayerTypeIPv6,
 		matchByIP: true,
 		drop:      action == firewall.ActionDrop,
-		comment:   comment,
 	}
-	if ipNormalized := ip.To4(); ipNormalized != nil {
+	if i.Is4() {
 		r.ipLayer = layers.LayerTypeIPv4
-		r.ip = ipNormalized
 	}
 
 	if s := r.ip.String(); s == "0.0.0.0" || s == "::" {
@@ -391,15 +402,16 @@ func (m *Manager) AddPeerFiltering(
 	}
 
 	m.mutex.Lock()
-	if _, ok := m.incomingRules[r.ip.String()]; !ok {
-		m.incomingRules[r.ip.String()] = make(RuleSet)
+	if _, ok := m.incomingRules[r.ip]; !ok {
+		m.incomingRules[r.ip] = make(RuleSet)
 	}
-	m.incomingRules[r.ip.String()][r.id] = r
+	m.incomingRules[r.ip][r.id] = r
 	m.mutex.Unlock()
 	return []firewall.Rule{&r}, nil
 }
 
 func (m *Manager) AddRouteFiltering(
+	id []byte,
 	sources []netip.Prefix,
 	destination netip.Prefix,
 	proto firewall.Protocol,
@@ -407,16 +419,15 @@ func (m *Manager) AddRouteFiltering(
 	dPort *firewall.Port,
 	action firewall.Action,
 ) (firewall.Rule, error) {
-	if m.nativeRouter && m.nativeFirewall != nil {
-		return m.nativeFirewall.AddRouteFiltering(sources, destination, proto, sPort, dPort, action)
+	if m.nativeRouter.Load() && m.nativeFirewall != nil {
+		return m.nativeFirewall.AddRouteFiltering(id, sources, destination, proto, sPort, dPort, action)
 	}
 
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
 	ruleID := uuid.New().String()
 	rule := RouteRule{
+		// TODO: consolidate these IDs
 		id:          ruleID,
+		mgmtId:      id,
 		sources:     sources,
 		destination: destination,
 		proto:       proto,
@@ -425,14 +436,16 @@ func (m *Manager) AddRouteFiltering(
 		action:      action,
 	}
 
+	m.mutex.Lock()
 	m.routeRules = append(m.routeRules, rule)
 	m.routeRules.Sort()
+	m.mutex.Unlock()
 
 	return &rule, nil
 }
 
 func (m *Manager) DeleteRouteRule(rule firewall.Rule) error {
-	if m.nativeRouter && m.nativeFirewall != nil {
+	if m.nativeRouter.Load() && m.nativeFirewall != nil {
 		return m.nativeFirewall.DeleteRouteRule(rule)
 	}
 
@@ -461,10 +474,10 @@ func (m *Manager) DeletePeerRule(rule firewall.Rule) error {
 		return fmt.Errorf("delete rule: invalid rule type: %T", rule)
 	}
 
-	if _, ok := m.incomingRules[r.ip.String()][r.id]; !ok {
+	if _, ok := m.incomingRules[r.ip][r.id]; !ok {
 		return fmt.Errorf("delete rule: no rule with such id: %v", r.id)
 	}
-	delete(m.incomingRules[r.ip.String()], r.id)
+	delete(m.incomingRules[r.ip], r.id)
 
 	return nil
 }
@@ -497,13 +510,13 @@ func (m *Manager) DeleteDNATRule(rule firewall.Rule) error {
 }
 
 // DropOutgoing filter outgoing packets
-func (m *Manager) DropOutgoing(packetData []byte) bool {
-	return m.processOutgoingHooks(packetData)
+func (m *Manager) DropOutgoing(packetData []byte, size int) bool {
+	return m.processOutgoingHooks(packetData, size)
 }
 
 // DropIncoming filter incoming packets
-func (m *Manager) DropIncoming(packetData []byte) bool {
-	return m.dropFilter(packetData)
+func (m *Manager) DropIncoming(packetData []byte, size int) bool {
+	return m.dropFilter(packetData, size)
 }
 
 // UpdateLocalIPs updates the list of local IPs
@@ -511,10 +524,7 @@ func (m *Manager) UpdateLocalIPs() error {
 	return m.localipmanager.UpdateLocalIPs(m.wgIface)
 }
 
-func (m *Manager) processOutgoingHooks(packetData []byte) bool {
-	m.mutex.RLock()
-	defer m.mutex.RUnlock()
-
+func (m *Manager) processOutgoingHooks(packetData []byte, size int) bool {
 	d := m.decoders.Get().(*decoder)
 	defer m.decoders.Put(d)
 
@@ -527,52 +537,37 @@ func (m *Manager) processOutgoingHooks(packetData []byte) bool {
 	}
 
 	srcIP, dstIP := m.extractIPs(d)
-	if srcIP == nil {
+	if !srcIP.IsValid() {
+		m.logger.Error("Unknown network layer: %v", d.decoded[0])
 		return false
 	}
 
-	// Track all protocols if stateful mode is enabled
-	if m.stateful {
-		switch d.decoded[1] {
-		case layers.LayerTypeUDP:
-			m.trackUDPOutbound(d, srcIP, dstIP)
-		case layers.LayerTypeTCP:
-			m.trackTCPOutbound(d, srcIP, dstIP)
-		case layers.LayerTypeICMPv4:
-			m.trackICMPOutbound(d, srcIP, dstIP)
-		}
+	if d.decoded[1] == layers.LayerTypeUDP && m.udpHooksDrop(uint16(d.udp.DstPort), dstIP, packetData) {
+		return true
 	}
 
-	// Process UDP hooks even if stateful mode is disabled
-	if d.decoded[1] == layers.LayerTypeUDP {
-		return m.checkUDPHooks(d, dstIP, packetData)
+	if m.stateful {
+		m.trackOutbound(d, srcIP, dstIP, size)
 	}
 
 	return false
 }
 
-func (m *Manager) extractIPs(d *decoder) (srcIP, dstIP net.IP) {
+func (m *Manager) extractIPs(d *decoder) (srcIP, dstIP netip.Addr) {
 	switch d.decoded[0] {
 	case layers.LayerTypeIPv4:
-		return d.ip4.SrcIP, d.ip4.DstIP
+		src, _ := netip.AddrFromSlice(d.ip4.SrcIP)
+		dst, _ := netip.AddrFromSlice(d.ip4.DstIP)
+		return src, dst
 	case layers.LayerTypeIPv6:
-		return d.ip6.SrcIP, d.ip6.DstIP
+		src, _ := netip.AddrFromSlice(d.ip6.SrcIP)
+		dst, _ := netip.AddrFromSlice(d.ip6.DstIP)
+		return src, dst
 	default:
-		return nil, nil
+		return netip.Addr{}, netip.Addr{}
 	}
 }
 
-func (m *Manager) trackTCPOutbound(d *decoder, srcIP, dstIP net.IP) {
-	flags := getTCPFlags(&d.tcp)
-	m.tcpTracker.TrackOutbound(
-		srcIP,
-		dstIP,
-		uint16(d.tcp.SrcPort),
-		uint16(d.tcp.DstPort),
-		flags,
-	)
-}
-
 func getTCPFlags(tcp *layers.TCP) uint8 {
 	var flags uint8
 	if tcp.SYN {
@@ -596,45 +591,70 @@ func getTCPFlags(tcp *layers.TCP) uint8 {
 	return flags
 }
 
-func (m *Manager) trackUDPOutbound(d *decoder, srcIP, dstIP net.IP) {
-	m.udpTracker.TrackOutbound(
-		srcIP,
-		dstIP,
-		uint16(d.udp.SrcPort),
-		uint16(d.udp.DstPort),
-	)
+func (m *Manager) trackOutbound(d *decoder, srcIP, dstIP netip.Addr, size int) {
+	transport := d.decoded[1]
+	switch transport {
+	case layers.LayerTypeUDP:
+		m.udpTracker.TrackOutbound(srcIP, dstIP, uint16(d.udp.SrcPort), uint16(d.udp.DstPort), size)
+	case layers.LayerTypeTCP:
+		flags := getTCPFlags(&d.tcp)
+		m.tcpTracker.TrackOutbound(srcIP, dstIP, uint16(d.tcp.SrcPort), uint16(d.tcp.DstPort), flags, size)
+	case layers.LayerTypeICMPv4:
+		m.icmpTracker.TrackOutbound(srcIP, dstIP, d.icmp4.Id, d.icmp4.TypeCode, size)
+	}
 }
 
-func (m *Manager) checkUDPHooks(d *decoder, dstIP net.IP, packetData []byte) bool {
-	for _, ipKey := range []string{dstIP.String(), "0.0.0.0", "::"} {
-		if rules, exists := m.outgoingRules[ipKey]; exists {
-			for _, rule := range rules {
-				if rule.udpHook != nil && portsMatch(rule.dPort, uint16(d.udp.DstPort)) {
-					return rule.udpHook(packetData)
-				}
+func (m *Manager) trackInbound(d *decoder, srcIP, dstIP netip.Addr, ruleID []byte, size int) {
+	transport := d.decoded[1]
+	switch transport {
+	case layers.LayerTypeUDP:
+		m.udpTracker.TrackInbound(srcIP, dstIP, uint16(d.udp.SrcPort), uint16(d.udp.DstPort), ruleID, size)
+	case layers.LayerTypeTCP:
+		flags := getTCPFlags(&d.tcp)
+		m.tcpTracker.TrackInbound(srcIP, dstIP, uint16(d.tcp.SrcPort), uint16(d.tcp.DstPort), flags, ruleID, size)
+	case layers.LayerTypeICMPv4:
+		m.icmpTracker.TrackInbound(srcIP, dstIP, d.icmp4.Id, d.icmp4.TypeCode, ruleID, size)
+	}
+}
+
+// udpHooksDrop checks if any UDP hooks should drop the packet
+func (m *Manager) udpHooksDrop(dport uint16, dstIP netip.Addr, packetData []byte) bool {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	// Check specific destination IP first
+	if rules, exists := m.outgoingRules[dstIP]; exists {
+		for _, rule := range rules {
+			if rule.udpHook != nil && portsMatch(rule.dPort, dport) {
+				return rule.udpHook(packetData)
 			}
 		}
 	}
-	return false
-}
 
-func (m *Manager) trackICMPOutbound(d *decoder, srcIP, dstIP net.IP) {
-	if d.icmp4.TypeCode.Type() == layers.ICMPv4TypeEchoRequest {
-		m.icmpTracker.TrackOutbound(
-			srcIP,
-			dstIP,
-			d.icmp4.Id,
-			d.icmp4.Seq,
-		)
+	// Check IPv4 unspecified address
+	if rules, exists := m.outgoingRules[netip.IPv4Unspecified()]; exists {
+		for _, rule := range rules {
+			if rule.udpHook != nil && portsMatch(rule.dPort, dport) {
+				return rule.udpHook(packetData)
+			}
+		}
 	}
+
+	// Check IPv6 unspecified address
+	if rules, exists := m.outgoingRules[netip.IPv6Unspecified()]; exists {
+		for _, rule := range rules {
+			if rule.udpHook != nil && portsMatch(rule.dPort, dport) {
+				return rule.udpHook(packetData)
+			}
+		}
+	}
+
+	return false
 }
 
 // dropFilter implements filtering logic for incoming packets.
 // If it returns true, the packet should be dropped.
-func (m *Manager) dropFilter(packetData []byte) bool {
-	m.mutex.RLock()
-	defer m.mutex.RUnlock()
-
+func (m *Manager) dropFilter(packetData []byte, size int) bool {
 	d := m.decoders.Get().(*decoder)
 	defer m.decoders.Put(d)
 
@@ -643,19 +663,19 @@ func (m *Manager) dropFilter(packetData []byte) bool {
 	}
 
 	srcIP, dstIP := m.extractIPs(d)
-	if srcIP == nil {
+	if !srcIP.IsValid() {
 		m.logger.Error("Unknown network layer: %v", d.decoded[0])
 		return true
 	}
 
 	// For all inbound traffic, first check if it matches a tracked connection.
 	// This must happen before any other filtering because the packets are statefully tracked.
-	if m.stateful && m.isValidTrackedConnection(d, srcIP, dstIP) {
+	if m.stateful && m.isValidTrackedConnection(d, srcIP, dstIP, size) {
 		return false
 	}
 
 	if m.localipmanager.IsLocalIP(dstIP) {
-		return m.handleLocalTraffic(d, srcIP, dstIP, packetData)
+		return m.handleLocalTraffic(d, srcIP, dstIP, packetData, size)
 	}
 
 	return m.handleRoutedTraffic(d, srcIP, dstIP, packetData)
@@ -663,10 +683,29 @@ func (m *Manager) dropFilter(packetData []byte) bool {
 
 // handleLocalTraffic handles local traffic.
 // If it returns true, the packet should be dropped.
-func (m *Manager) handleLocalTraffic(d *decoder, srcIP, dstIP net.IP, packetData []byte) bool {
-	if m.peerACLsBlock(srcIP, packetData, m.incomingRules, d) {
-		m.logger.Trace("Dropping local packet (ACL denied): src=%s dst=%s",
-			srcIP, dstIP)
+func (m *Manager) handleLocalTraffic(d *decoder, srcIP, dstIP netip.Addr, packetData []byte, size int) bool {
+	ruleID, blocked := m.peerACLsBlock(srcIP, packetData, m.incomingRules, d)
+	if blocked {
+		_, pnum := getProtocolFromPacket(d)
+		srcPort, dstPort := getPortsFromPacket(d)
+
+		m.logger.Trace("Dropping local packet (ACL denied): rule_id=%s proto=%v src=%s:%d dst=%s:%d",
+			ruleID, pnum, srcIP, srcPort, dstIP, dstPort)
+
+		m.flowLogger.StoreEvent(nftypes.EventFields{
+			FlowID:     uuid.New(),
+			Type:       nftypes.TypeDrop,
+			RuleID:     ruleID,
+			Direction:  nftypes.Ingress,
+			Protocol:   pnum,
+			SourceIP:   srcIP,
+			DestIP:     dstIP,
+			SourcePort: srcPort,
+			DestPort:   dstPort,
+			// TODO: icmp type/code
+			RxPackets: 1,
+			RxBytes:   uint64(size),
+		})
 		return true
 	}
 
@@ -675,6 +714,9 @@ func (m *Manager) handleLocalTraffic(d *decoder, srcIP, dstIP net.IP, packetData
 		return m.handleNetstackLocalTraffic(packetData)
 	}
 
+	// track inbound packets to get the correct direction and session id for flows
+	m.trackInbound(d, srcIP, dstIP, ruleID, size)
+
 	return false
 }
 
@@ -684,12 +726,13 @@ func (m *Manager) handleNetstackLocalTraffic(packetData []byte) bool {
 		return false
 	}
 
-	if m.forwarder == nil {
+	fwd := m.forwarder.Load()
+	if fwd == nil {
 		m.logger.Trace("Dropping local packet (forwarder not initialized)")
 		return true
 	}
 
-	if err := m.forwarder.InjectIncomingPacket(packetData); err != nil {
+	if err := fwd.InjectIncomingPacket(packetData); err != nil {
 		m.logger.Error("Failed to inject local packet: %v", err)
 	}
 
@@ -699,47 +742,65 @@ func (m *Manager) handleNetstackLocalTraffic(packetData []byte) bool {
 
 // handleRoutedTraffic handles routed traffic.
 // If it returns true, the packet should be dropped.
-func (m *Manager) handleRoutedTraffic(d *decoder, srcIP, dstIP net.IP, packetData []byte) bool {
+func (m *Manager) handleRoutedTraffic(d *decoder, srcIP, dstIP netip.Addr, packetData []byte) bool {
 	// Drop if routing is disabled
-	if !m.routingEnabled {
+	if !m.routingEnabled.Load() {
 		m.logger.Trace("Dropping routed packet (routing disabled): src=%s dst=%s",
 			srcIP, dstIP)
 		return true
 	}
 
 	// Pass to native stack if native router is enabled or forced
-	if m.nativeRouter {
+	if m.nativeRouter.Load() {
 		return false
 	}
 
-	proto := getProtocolFromPacket(d)
+	proto, pnum := getProtocolFromPacket(d)
 	srcPort, dstPort := getPortsFromPacket(d)
 
-	if !m.routeACLsPass(srcIP, dstIP, proto, srcPort, dstPort) {
-		m.logger.Trace("Dropping routed packet (ACL denied): src=%s:%d dst=%s:%d proto=%v",
-			srcIP, srcPort, dstIP, dstPort, proto)
+	if ruleID, pass := m.routeACLsPass(srcIP, dstIP, proto, srcPort, dstPort); !pass {
+		m.logger.Trace("Dropping routed packet (ACL denied): rule_id=%s proto=%v src=%s:%d dst=%s:%d",
+			ruleID, pnum, srcIP, srcPort, dstIP, dstPort)
+
+		m.flowLogger.StoreEvent(nftypes.EventFields{
+			FlowID:     uuid.New(),
+			Type:       nftypes.TypeDrop,
+			RuleID:     ruleID,
+			Direction:  nftypes.Ingress,
+			Protocol:   pnum,
+			SourceIP:   srcIP,
+			DestIP:     dstIP,
+			SourcePort: srcPort,
+			DestPort:   dstPort,
+			// TODO: icmp type/code
+		})
 		return true
 	}
 
 	// Let forwarder handle the packet if it passed route ACLs
-	if err := m.forwarder.InjectIncomingPacket(packetData); err != nil {
-		m.logger.Error("Failed to inject incoming packet: %v", err)
+	fwd := m.forwarder.Load()
+	if fwd == nil {
+		m.logger.Trace("failed to forward routed packet (forwarder not initialized)")
+	} else {
+		if err := fwd.InjectIncomingPacket(packetData); err != nil {
+			m.logger.Error("Failed to inject routed packet: %v", err)
+		}
 	}
 
 	// Forwarded packets shouldn't reach the native stack, hence they won't be visible in a packet capture
 	return true
 }
 
-func getProtocolFromPacket(d *decoder) firewall.Protocol {
+func getProtocolFromPacket(d *decoder) (firewall.Protocol, nftypes.Protocol) {
 	switch d.decoded[1] {
 	case layers.LayerTypeTCP:
-		return firewall.ProtocolTCP
+		return firewall.ProtocolTCP, nftypes.TCP
 	case layers.LayerTypeUDP:
-		return firewall.ProtocolUDP
+		return firewall.ProtocolUDP, nftypes.UDP
 	case layers.LayerTypeICMPv4, layers.LayerTypeICMPv6:
-		return firewall.ProtocolICMP
+		return firewall.ProtocolICMP, nftypes.ICMP
 	default:
-		return firewall.ProtocolALL
+		return firewall.ProtocolALL, nftypes.ProtocolUnknown
 	}
 }
 
@@ -767,7 +828,7 @@ func (m *Manager) isValidPacket(d *decoder, packetData []byte) bool {
 	return true
 }
 
-func (m *Manager) isValidTrackedConnection(d *decoder, srcIP, dstIP net.IP) bool {
+func (m *Manager) isValidTrackedConnection(d *decoder, srcIP, dstIP netip.Addr, size int) bool {
 	switch d.decoded[1] {
 	case layers.LayerTypeTCP:
 		return m.tcpTracker.IsValidInbound(
@@ -776,6 +837,7 @@ func (m *Manager) isValidTrackedConnection(d *decoder, srcIP, dstIP net.IP) bool
 			uint16(d.tcp.SrcPort),
 			uint16(d.tcp.DstPort),
 			getTCPFlags(&d.tcp),
+			size,
 		)
 
 	case layers.LayerTypeUDP:
@@ -784,6 +846,7 @@ func (m *Manager) isValidTrackedConnection(d *decoder, srcIP, dstIP net.IP) bool
 			dstIP,
 			uint16(d.udp.SrcPort),
 			uint16(d.udp.DstPort),
+			size,
 		)
 
 	case layers.LayerTypeICMPv4:
@@ -791,8 +854,8 @@ func (m *Manager) isValidTrackedConnection(d *decoder, srcIP, dstIP net.IP) bool
 			srcIP,
 			dstIP,
 			d.icmp4.Id,
-			d.icmp4.Seq,
 			d.icmp4.TypeCode.Type(),
+			size,
 		)
 
 		// TODO: ICMPv6
@@ -812,25 +875,27 @@ func (m *Manager) isSpecialICMP(d *decoder) bool {
 		icmpType == layers.ICMPv4TypeTimeExceeded
 }
 
-func (m *Manager) peerACLsBlock(srcIP net.IP, packetData []byte, rules map[string]RuleSet, d *decoder) bool {
+func (m *Manager) peerACLsBlock(srcIP netip.Addr, packetData []byte, rules map[netip.Addr]RuleSet, d *decoder) ([]byte, bool) {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
 	if m.isSpecialICMP(d) {
-		return false
+		return nil, false
 	}
 
-	if filter, ok := validateRule(srcIP, packetData, rules[srcIP.String()], d); ok {
-		return filter
+	if mgmtId, filter, ok := validateRule(srcIP, packetData, rules[srcIP], d); ok {
+		return mgmtId, filter
 	}
 
-	if filter, ok := validateRule(srcIP, packetData, rules["0.0.0.0"], d); ok {
-		return filter
+	if mgmtId, filter, ok := validateRule(srcIP, packetData, rules[netip.IPv4Unspecified()], d); ok {
+		return mgmtId, filter
 	}
 
-	if filter, ok := validateRule(srcIP, packetData, rules["::"], d); ok {
-		return filter
+	if mgmtId, filter, ok := validateRule(srcIP, packetData, rules[netip.IPv6Unspecified()], d); ok {
+		return mgmtId, filter
 	}
 
 	// Default policy: DROP ALL
-	return true
+	return nil, true
 }
 
 func portsMatch(rulePort *firewall.Port, packetPort uint16) bool {
@@ -850,15 +915,15 @@ func portsMatch(rulePort *firewall.Port, packetPort uint16) bool {
 	return false
 }
 
-func validateRule(ip net.IP, packetData []byte, rules map[string]PeerRule, d *decoder) (bool, bool) {
+func validateRule(ip netip.Addr, packetData []byte, rules map[string]PeerRule, d *decoder) ([]byte, bool, bool) {
 	payloadLayer := d.decoded[1]
 	for _, rule := range rules {
-		if rule.matchByIP && !ip.Equal(rule.ip) {
+		if rule.matchByIP && ip.Compare(rule.ip) != 0 {
 			continue
 		}
 
 		if rule.protoLayer == layerTypeAll {
-			return rule.drop, true
+			return rule.mgmtId, rule.drop, true
 		}
 
 		if payloadLayer != rule.protoLayer {
@@ -868,39 +933,36 @@ func validateRule(ip net.IP, packetData []byte, rules map[string]PeerRule, d *de
 		switch payloadLayer {
 		case layers.LayerTypeTCP:
 			if portsMatch(rule.sPort, uint16(d.tcp.SrcPort)) && portsMatch(rule.dPort, uint16(d.tcp.DstPort)) {
-				return rule.drop, true
+				return rule.mgmtId, rule.drop, true
 			}
 		case layers.LayerTypeUDP:
 			// if rule has UDP hook (and if we are here we match this rule)
 			// we ignore rule.drop and call this hook
 			if rule.udpHook != nil {
-				return rule.udpHook(packetData), true
+				return rule.mgmtId, rule.udpHook(packetData), true
 			}
 
 			if portsMatch(rule.sPort, uint16(d.udp.SrcPort)) && portsMatch(rule.dPort, uint16(d.udp.DstPort)) {
-				return rule.drop, true
+				return rule.mgmtId, rule.drop, true
 			}
 		case layers.LayerTypeICMPv4, layers.LayerTypeICMPv6:
-			return rule.drop, true
+			return rule.mgmtId, rule.drop, true
 		}
 	}
-	return false, false
+	return nil, false, false
 }
 
-// routeACLsPass returns treu if the packet is allowed by the route ACLs
-func (m *Manager) routeACLsPass(srcIP, dstIP net.IP, proto firewall.Protocol, srcPort, dstPort uint16) bool {
+// routeACLsPass returns true if the packet is allowed by the route ACLs
+func (m *Manager) routeACLsPass(srcIP, dstIP netip.Addr, proto firewall.Protocol, srcPort, dstPort uint16) ([]byte, bool) {
 	m.mutex.RLock()
 	defer m.mutex.RUnlock()
 
-	srcAddr := netip.AddrFrom4([4]byte(srcIP.To4()))
-	dstAddr := netip.AddrFrom4([4]byte(dstIP.To4()))
-
 	for _, rule := range m.routeRules {
-		if m.ruleMatches(rule, srcAddr, dstAddr, proto, srcPort, dstPort) {
-			return rule.action == firewall.ActionAccept
+		if matches := m.ruleMatches(rule, srcIP, dstIP, proto, srcPort, dstPort); matches {
+			return rule.mgmtId, rule.action == firewall.ActionAccept
 		}
 	}
-	return false
+	return nil, false
 }
 
 func (m *Manager) ruleMatches(rule RouteRule, srcAddr, dstAddr netip.Addr, proto firewall.Protocol, srcPort, dstPort uint16) bool {
@@ -940,36 +1002,32 @@ func (m *Manager) SetNetwork(network *net.IPNet) {
 // AddUDPPacketHook calls hook when UDP packet from given direction matched
 //
 // Hook function returns flag which indicates should be the matched package dropped or not
-func (m *Manager) AddUDPPacketHook(
-	in bool, ip net.IP, dPort uint16, hook func([]byte) bool,
-) string {
+func (m *Manager) AddUDPPacketHook(in bool, ip netip.Addr, dPort uint16, hook func(packet []byte) bool) string {
 	r := PeerRule{
 		id:         uuid.New().String(),
 		ip:         ip,
 		protoLayer: layers.LayerTypeUDP,
 		dPort:      &firewall.Port{Values: []uint16{dPort}},
 		ipLayer:    layers.LayerTypeIPv6,
-		comment:    fmt.Sprintf("UDP Hook direction: %v, ip:%v, dport:%d", in, ip, dPort),
 		udpHook:    hook,
 	}
 
-	if ip.To4() != nil {
+	if ip.Is4() {
 		r.ipLayer = layers.LayerTypeIPv4
 	}
 
 	m.mutex.Lock()
 	if in {
-		if _, ok := m.incomingRules[r.ip.String()]; !ok {
-			m.incomingRules[r.ip.String()] = make(map[string]PeerRule)
+		if _, ok := m.incomingRules[r.ip]; !ok {
+			m.incomingRules[r.ip] = make(map[string]PeerRule)
 		}
-		m.incomingRules[r.ip.String()][r.id] = r
+		m.incomingRules[r.ip][r.id] = r
 	} else {
-		if _, ok := m.outgoingRules[r.ip.String()]; !ok {
-			m.outgoingRules[r.ip.String()] = make(map[string]PeerRule)
+		if _, ok := m.outgoingRules[r.ip]; !ok {
+			m.outgoingRules[r.ip] = make(map[string]PeerRule)
 		}
-		m.outgoingRules[r.ip.String()][r.id] = r
+		m.outgoingRules[r.ip][r.id] = r
 	}
-
 	m.mutex.Unlock()
 
 	return r.id
@@ -1017,20 +1075,21 @@ func (m *Manager) DisableRouting() error {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	if m.forwarder == nil {
+	fwder := m.forwarder.Load()
+	if fwder == nil {
 		return nil
 	}
 
-	m.routingEnabled = false
-	m.nativeRouter = false
+	m.routingEnabled.Store(false)
+	m.nativeRouter.Store(false)
 
 	// don't stop forwarder if in use by netstack
 	if m.netstack && m.localForwarding {
 		return nil
 	}
 
-	m.forwarder.Stop()
-	m.forwarder = nil
+	fwder.Stop()
+	m.forwarder.Store(nil)
 
 	log.Debug("forwarder stopped")
 

client/firewall/uspfilter/uspfilter_bench_test.go

@@ -93,8 +93,7 @@ func BenchmarkCoreFiltering(b *testing.B) {
 			stateful: false,
 			setupFunc: func(m *Manager) {
 				// Single rule allowing all traffic
-				_, err := m.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolALL, nil, nil,
-					fw.ActionAccept, "", "allow all")
+				_, err := m.AddPeerFiltering(nil, net.ParseIP("0.0.0.0"), fw.ProtocolALL, nil, nil, fw.ActionAccept, "")
 				require.NoError(b, err)
 			},
 			desc: "Baseline: Single 'allow all' rule without connection tracking",
@@ -114,10 +113,15 @@ func BenchmarkCoreFiltering(b *testing.B) {
 				// Add explicit rules matching return traffic pattern
 				for i := 0; i < 1000; i++ { // Simulate realistic ruleset size
 					ip := generateRandomIPs(1)[0]
-					_, err := m.AddPeerFiltering(ip, fw.ProtocolTCP,
+					_, err := m.AddPeerFiltering(
+						nil,
+						ip,
+						fw.ProtocolTCP,
 						&fw.Port{Values: []uint16{uint16(1024 + i)}},
 						&fw.Port{Values: []uint16{80}},
-						fw.ActionAccept, "", "explicit return")
+						fw.ActionAccept,
+						"",
+					)
 					require.NoError(b, err)
 				}
 			},
@@ -128,8 +132,15 @@ func BenchmarkCoreFiltering(b *testing.B) {
 			stateful: true,
 			setupFunc: func(m *Manager) {
 				// Add some basic rules but rely on state for established connections
-				_, err := m.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolTCP, nil, nil,
-					fw.ActionDrop, "", "default drop")
+				_, err := m.AddPeerFiltering(
+					nil,
+					net.ParseIP("0.0.0.0"),
+					fw.ProtocolTCP,
+					nil,
+					nil,
+					fw.ActionDrop,
+					"",
+				)
 				require.NoError(b, err)
 			},
 			desc: "Connection tracking with established connections",
@@ -158,7 +169,7 @@ func BenchmarkCoreFiltering(b *testing.B) {
 				// Create manager and basic setup
 				manager, _ := Create(&IFaceMock{
 					SetFilterFunc: func(device.PacketFilter) error { return nil },
-				}, false)
+				}, false, flowLogger)
 				defer b.Cleanup(func() {
 					require.NoError(b, manager.Close(nil))
 				})
@@ -182,13 +193,13 @@ func BenchmarkCoreFiltering(b *testing.B) {
 
 				// For stateful scenarios, establish the connection
 				if sc.stateful {
-					manager.processOutgoingHooks(outbound)
+					manager.processOutgoingHooks(outbound, 0)
 				}
 
 				// Measure inbound packet processing
 				b.ResetTimer()
 				for i := 0; i < b.N; i++ {
-					manager.dropFilter(inbound)
+					manager.dropFilter(inbound, 0)
 				}
 			})
 		}
@@ -203,7 +214,7 @@ func BenchmarkStateScaling(b *testing.B) {
 		b.Run(fmt.Sprintf("conns_%d", count), func(b *testing.B) {
 			manager, _ := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			}, false, flowLogger)
 			b.Cleanup(func() {
 				require.NoError(b, manager.Close(nil))
 			})
@@ -219,7 +230,7 @@ func BenchmarkStateScaling(b *testing.B) {
 			for i := 0; i < count; i++ {
 				outbound := generatePacket(b, srcIPs[i], dstIPs[i],
 					uint16(1024+i), 80, layers.IPProtocolTCP)
-				manager.processOutgoingHooks(outbound)
+				manager.processOutgoingHooks(outbound, 0)
 			}
 
 			// Test packet
@@ -227,11 +238,11 @@ func BenchmarkStateScaling(b *testing.B) {
 			testIn := generatePacket(b, dstIPs[0], srcIPs[0], 80, 1024, layers.IPProtocolTCP)
 
 			// First establish our test connection
-			manager.processOutgoingHooks(testOut)
+			manager.processOutgoingHooks(testOut, 0)
 
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				manager.dropFilter(testIn)
+				manager.dropFilter(testIn, 0)
 			}
 		})
 	}
@@ -251,7 +262,7 @@ func BenchmarkEstablishmentOverhead(b *testing.B) {
 		b.Run(sc.name, func(b *testing.B) {
 			manager, _ := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			}, false, flowLogger)
 			b.Cleanup(func() {
 				require.NoError(b, manager.Close(nil))
 			})
@@ -267,12 +278,12 @@ func BenchmarkEstablishmentOverhead(b *testing.B) {
 			inbound := generatePacket(b, dstIP, srcIP, 80, 1024, layers.IPProtocolTCP)
 
 			if sc.established {
-				manager.processOutgoingHooks(outbound)
+				manager.processOutgoingHooks(outbound, 0)
 			}
 
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				manager.dropFilter(inbound)
+				manager.dropFilter(inbound, 0)
 			}
 		})
 	}
@@ -450,7 +461,7 @@ func BenchmarkRoutedNetworkReturn(b *testing.B) {
 		b.Run(sc.name, func(b *testing.B) {
 			manager, _ := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			}, false, flowLogger)
 			b.Cleanup(func() {
 				require.NoError(b, manager.Close(nil))
 			})
@@ -466,25 +477,25 @@ func BenchmarkRoutedNetworkReturn(b *testing.B) {
 			// For stateful cases and established connections
 			if !strings.Contains(sc.name, "allow_non_wg") ||
 				(strings.Contains(sc.state, "established") || sc.state == "post_handshake") {
-				manager.processOutgoingHooks(outbound)
+				manager.processOutgoingHooks(outbound, 0)
 
 				// For TCP post-handshake, simulate full handshake
 				if sc.state == "post_handshake" {
 					// SYN
 					syn := generateTCPPacketWithFlags(b, srcIP, dstIP, 1024, 80, uint16(conntrack.TCPSyn))
-					manager.processOutgoingHooks(syn)
+					manager.processOutgoingHooks(syn, 0)
 					// SYN-ACK
 					synack := generateTCPPacketWithFlags(b, dstIP, srcIP, 80, 1024, uint16(conntrack.TCPSyn|conntrack.TCPAck))
-					manager.dropFilter(synack)
+					manager.dropFilter(synack, 0)
 					// ACK
 					ack := generateTCPPacketWithFlags(b, srcIP, dstIP, 1024, 80, uint16(conntrack.TCPAck))
-					manager.processOutgoingHooks(ack)
+					manager.processOutgoingHooks(ack, 0)
 				}
 			}
 
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				manager.dropFilter(inbound)
+				manager.dropFilter(inbound, 0)
 			}
 		})
 	}
@@ -577,7 +588,7 @@ func BenchmarkLongLivedConnections(b *testing.B) {
 
 			manager, _ := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			}, false, flowLogger)
 			defer b.Cleanup(func() {
 				require.NoError(b, manager.Close(nil))
 			})
@@ -590,10 +601,7 @@ func BenchmarkLongLivedConnections(b *testing.B) {
 			// Setup initial state based on scenario
 			if sc.rules {
 				// Single rule to allow all return traffic from port 80
-				_, err := manager.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolTCP,
-					&fw.Port{Values: []uint16{80}},
-					nil,
-					fw.ActionAccept, "", "return traffic")
+				_, err := manager.AddPeerFiltering(nil, net.ParseIP("0.0.0.0"), fw.ProtocolTCP, &fw.Port{Values: []uint16{80}}, nil, fw.ActionAccept, "")
 				require.NoError(b, err)
 			}
 
@@ -616,17 +624,17 @@ func BenchmarkLongLivedConnections(b *testing.B) {
 				// Initial SYN
 				syn := generateTCPPacketWithFlags(b, srcIPs[i], dstIPs[i],
 					uint16(1024+i), 80, uint16(conntrack.TCPSyn))
-				manager.processOutgoingHooks(syn)
+				manager.processOutgoingHooks(syn, 0)
 
 				// SYN-ACK
 				synack := generateTCPPacketWithFlags(b, dstIPs[i], srcIPs[i],
 					80, uint16(1024+i), uint16(conntrack.TCPSyn|conntrack.TCPAck))
-				manager.dropFilter(synack)
+				manager.dropFilter(synack, 0)
 
 				// ACK
 				ack := generateTCPPacketWithFlags(b, srcIPs[i], dstIPs[i],
 					uint16(1024+i), 80, uint16(conntrack.TCPAck))
-				manager.processOutgoingHooks(ack)
+				manager.processOutgoingHooks(ack, 0)
 			}
 
 			// Prepare test packets simulating bidirectional traffic
@@ -647,9 +655,9 @@ func BenchmarkLongLivedConnections(b *testing.B) {
 
 				// Simulate bidirectional traffic
 				// First outbound data
-				manager.processOutgoingHooks(outPackets[connIdx])
+				manager.processOutgoingHooks(outPackets[connIdx], 0)
 				// Then inbound response - this is what we're actually measuring
-				manager.dropFilter(inPackets[connIdx])
+				manager.dropFilter(inPackets[connIdx], 0)
 			}
 		})
 	}
@@ -668,7 +676,7 @@ func BenchmarkShortLivedConnections(b *testing.B) {
 
 			manager, _ := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			}, false, flowLogger)
 			defer b.Cleanup(func() {
 				require.NoError(b, manager.Close(nil))
 			})
@@ -681,10 +689,7 @@ func BenchmarkShortLivedConnections(b *testing.B) {
 			// Setup initial state based on scenario
 			if sc.rules {
 				// Single rule to allow all return traffic from port 80
-				_, err := manager.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolTCP,
-					&fw.Port{Values: []uint16{80}},
-					nil,
-					fw.ActionAccept, "", "return traffic")
+				_, err := manager.AddPeerFiltering(nil, net.ParseIP("0.0.0.0"), fw.ProtocolTCP, &fw.Port{Values: []uint16{80}}, nil, fw.ActionAccept, "")
 				require.NoError(b, err)
 			}
 
@@ -756,19 +761,19 @@ func BenchmarkShortLivedConnections(b *testing.B) {
 				p := patterns[connIdx]
 
 				// Connection establishment
-				manager.processOutgoingHooks(p.syn)
-				manager.dropFilter(p.synAck)
-				manager.processOutgoingHooks(p.ack)
+				manager.processOutgoingHooks(p.syn, 0)
+				manager.dropFilter(p.synAck, 0)
+				manager.processOutgoingHooks(p.ack, 0)
 
 				// Data transfer
-				manager.processOutgoingHooks(p.request)
-				manager.dropFilter(p.response)
+				manager.processOutgoingHooks(p.request, 0)
+				manager.dropFilter(p.response, 0)
 
 				// Connection teardown
-				manager.processOutgoingHooks(p.finClient)
-				manager.dropFilter(p.ackServer)
-				manager.dropFilter(p.finServer)
-				manager.processOutgoingHooks(p.ackClient)
+				manager.processOutgoingHooks(p.finClient, 0)
+				manager.dropFilter(p.ackServer, 0)
+				manager.dropFilter(p.finServer, 0)
+				manager.processOutgoingHooks(p.ackClient, 0)
 			}
 		})
 	}
@@ -787,7 +792,7 @@ func BenchmarkParallelLongLivedConnections(b *testing.B) {
 
 			manager, _ := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			}, false, flowLogger)
 			defer b.Cleanup(func() {
 				require.NoError(b, manager.Close(nil))
 			})
@@ -799,10 +804,7 @@ func BenchmarkParallelLongLivedConnections(b *testing.B) {
 
 			// Setup initial state based on scenario
 			if sc.rules {
-				_, err := manager.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolTCP,
-					&fw.Port{Values: []uint16{80}},
-					nil,
-					fw.ActionAccept, "", "return traffic")
+				_, err := manager.AddPeerFiltering(nil, net.ParseIP("0.0.0.0"), fw.ProtocolTCP, &fw.Port{Values: []uint16{80}}, nil, fw.ActionAccept, "")
 				require.NoError(b, err)
 			}
 
@@ -824,15 +826,15 @@ func BenchmarkParallelLongLivedConnections(b *testing.B) {
 			for i := 0; i < sc.connCount; i++ {
 				syn := generateTCPPacketWithFlags(b, srcIPs[i], dstIPs[i],
 					uint16(1024+i), 80, uint16(conntrack.TCPSyn))
-				manager.processOutgoingHooks(syn)
+				manager.processOutgoingHooks(syn, 0)
 
 				synack := generateTCPPacketWithFlags(b, dstIPs[i], srcIPs[i],
 					80, uint16(1024+i), uint16(conntrack.TCPSyn|conntrack.TCPAck))
-				manager.dropFilter(synack)
+				manager.dropFilter(synack, 0)
 
 				ack := generateTCPPacketWithFlags(b, srcIPs[i], dstIPs[i],
 					uint16(1024+i), 80, uint16(conntrack.TCPAck))
-				manager.processOutgoingHooks(ack)
+				manager.processOutgoingHooks(ack, 0)
 			}
 
 			// Pre-generate test packets
@@ -854,8 +856,8 @@ func BenchmarkParallelLongLivedConnections(b *testing.B) {
 					counter++
 
 					// Simulate bidirectional traffic
-					manager.processOutgoingHooks(outPackets[connIdx])
-					manager.dropFilter(inPackets[connIdx])
+					manager.processOutgoingHooks(outPackets[connIdx], 0)
+					manager.dropFilter(inPackets[connIdx], 0)
 				}
 			})
 		})
@@ -875,7 +877,7 @@ func BenchmarkParallelShortLivedConnections(b *testing.B) {
 
 			manager, _ := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			}, false, flowLogger)
 			defer b.Cleanup(func() {
 				require.NoError(b, manager.Close(nil))
 			})
@@ -886,10 +888,7 @@ func BenchmarkParallelShortLivedConnections(b *testing.B) {
 			})
 
 			if sc.rules {
-				_, err := manager.AddPeerFiltering(net.ParseIP("0.0.0.0"), fw.ProtocolTCP,
-					&fw.Port{Values: []uint16{80}},
-					nil,
-					fw.ActionAccept, "", "return traffic")
+				_, err := manager.AddPeerFiltering(nil, net.ParseIP("0.0.0.0"), fw.ProtocolTCP, &fw.Port{Values: []uint16{80}}, nil, fw.ActionAccept, "")
 				require.NoError(b, err)
 			}
 
@@ -951,17 +950,17 @@ func BenchmarkParallelShortLivedConnections(b *testing.B) {
 					p := patterns[connIdx]
 
 					// Full connection lifecycle
-					manager.processOutgoingHooks(p.syn)
-					manager.dropFilter(p.synAck)
-					manager.processOutgoingHooks(p.ack)
+					manager.processOutgoingHooks(p.syn, 0)
+					manager.dropFilter(p.synAck, 0)
+					manager.processOutgoingHooks(p.ack, 0)
 
-					manager.processOutgoingHooks(p.request)
-					manager.dropFilter(p.response)
+					manager.processOutgoingHooks(p.request, 0)
+					manager.dropFilter(p.response, 0)
 
-					manager.processOutgoingHooks(p.finClient)
-					manager.dropFilter(p.ackServer)
-					manager.dropFilter(p.finServer)
-					manager.processOutgoingHooks(p.ackClient)
+					manager.processOutgoingHooks(p.finClient, 0)
+					manager.dropFilter(p.ackServer, 0)
+					manager.dropFilter(p.finServer, 0)
+					manager.processOutgoingHooks(p.ackClient, 0)
 				}
 			})
 		})
@@ -1033,14 +1032,7 @@ func BenchmarkRouteACLs(b *testing.B) {
 	}
 
 	for _, r := range rules {
-		_, err := manager.AddRouteFiltering(
-			r.sources,
-			r.dest,
-			r.proto,
-			nil,
-			r.port,
-			fw.ActionAccept,
-		)
+		_, err := manager.AddRouteFiltering(nil, r.sources, r.dest, r.proto, nil, r.port, fw.ActionAccept)
 		if err != nil {
 			b.Fatal(err)
 		}
@@ -1062,8 +1054,8 @@ func BenchmarkRouteACLs(b *testing.B) {
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		for _, tc := range cases {
-			srcIP := net.ParseIP(tc.srcIP)
-			dstIP := net.ParseIP(tc.dstIP)
+			srcIP := netip.MustParseAddr(tc.srcIP)
+			dstIP := netip.MustParseAddr(tc.dstIP)
 			manager.routeACLsPass(srcIP, dstIP, tc.proto, 0, tc.dstPort)
 		}
 	}

client/firewall/uspfilter/uspfilter_filter_test.go

@@ -34,7 +34,7 @@ func TestPeerACLFiltering(t *testing.T) {
 		},
 	}
 
-	manager, err := Create(ifaceMock, false)
+	manager, err := Create(ifaceMock, false, flowLogger)
 	require.NoError(t, err)
 	require.NotNil(t, manager)
 
@@ -192,20 +192,20 @@ func TestPeerACLFiltering(t *testing.T) {
 
 	t.Run("Implicit DROP (no rules)", func(t *testing.T) {
 		packet := createTestPacket(t, "100.10.0.1", "100.10.0.100", fw.ProtocolTCP, 12345, 443)
-		isDropped := manager.DropIncoming(packet)
+		isDropped := manager.DropIncoming(packet, 0)
 		require.True(t, isDropped, "Packet should be dropped when no rules exist")
 	})
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			rules, err := manager.AddPeerFiltering(
+				nil,
 				net.ParseIP(tc.ruleIP),
 				tc.ruleProto,
 				tc.ruleSrcPort,
 				tc.ruleDstPort,
 				tc.ruleAction,
 				"",
-				tc.name,
 			)
 			require.NoError(t, err)
 			require.NotEmpty(t, rules)
@@ -217,7 +217,7 @@ func TestPeerACLFiltering(t *testing.T) {
 			})
 
 			packet := createTestPacket(t, tc.srcIP, tc.dstIP, tc.proto, tc.srcPort, tc.dstPort)
-			isDropped := manager.DropIncoming(packet)
+			isDropped := manager.DropIncoming(packet, 0)
 			require.Equal(t, tc.shouldBeBlocked, isDropped)
 		})
 	}
@@ -302,12 +302,12 @@ func setupRoutedManager(tb testing.TB, network string) *Manager {
 		},
 	}
 
-	manager, err := Create(ifaceMock, false)
+	manager, err := Create(ifaceMock, false, flowLogger)
 	require.NoError(tb, manager.EnableRouting())
 	require.NoError(tb, err)
 	require.NotNil(tb, manager)
-	require.True(tb, manager.routingEnabled)
-	require.False(tb, manager.nativeRouter)
+	require.True(tb, manager.routingEnabled.Load())
+	require.False(tb, manager.nativeRouter.Load())
 
 	tb.Cleanup(func() {
 		require.NoError(tb, manager.Close(nil))
@@ -803,6 +803,7 @@ func TestRouteACLFiltering(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			rule, err := manager.AddRouteFiltering(
+				nil,
 				tc.rule.sources,
 				tc.rule.dest,
 				tc.rule.proto,
@@ -817,12 +818,12 @@ func TestRouteACLFiltering(t *testing.T) {
 				require.NoError(t, manager.DeleteRouteRule(rule))
 			})
 
-			srcIP := net.ParseIP(tc.srcIP)
-			dstIP := net.ParseIP(tc.dstIP)
+			srcIP := netip.MustParseAddr(tc.srcIP)
+			dstIP := netip.MustParseAddr(tc.dstIP)
 
 			// testing routeACLsPass only and not DropIncoming, as routed packets are dropped after being passed
 			// to the forwarder
-			isAllowed := manager.routeACLsPass(srcIP, dstIP, tc.proto, tc.srcPort, tc.dstPort)
+			_, isAllowed := manager.routeACLsPass(srcIP, dstIP, tc.proto, tc.srcPort, tc.dstPort)
 			require.Equal(t, tc.shouldPass, isAllowed)
 		})
 	}
@@ -985,6 +986,7 @@ func TestRouteACLOrder(t *testing.T) {
 			var rules []fw.Rule
 			for _, r := range tc.rules {
 				rule, err := manager.AddRouteFiltering(
+					nil,
 					r.sources,
 					r.dest,
 					r.proto,
@@ -1004,10 +1006,10 @@ func TestRouteACLOrder(t *testing.T) {
 			})
 
 			for i, p := range tc.packets {
-				srcIP := net.ParseIP(p.srcIP)
-				dstIP := net.ParseIP(p.dstIP)
+				srcIP := netip.MustParseAddr(p.srcIP)
+				dstIP := netip.MustParseAddr(p.dstIP)
 
-				isAllowed := manager.routeACLsPass(srcIP, dstIP, p.proto, p.srcPort, p.dstPort)
+				_, isAllowed := manager.routeACLsPass(srcIP, dstIP, p.proto, p.srcPort, p.dstPort)
 				require.Equal(t, p.shouldPass, isAllowed, "packet %d failed", i)
 			}
 		})

client/firewall/uspfilter/uspfilter_test.go

@@ -1,8 +1,10 @@
 package uspfilter
 
 import (
+	"context"
 	"fmt"
 	"net"
+	"net/netip"
 	"sync"
 	"testing"
 	"time"
@@ -18,9 +20,11 @@ import (
 	"github.com/netbirdio/netbird/client/firewall/uspfilter/log"
 	"github.com/netbirdio/netbird/client/iface/device"
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
+	"github.com/netbirdio/netbird/client/internal/netflow"
 )
 
 var logger = log.NewFromLogrus(logrus.StandardLogger())
+var flowLogger = netflow.NewManager(context.Background(), nil, []byte{}, nil).GetLogger()
 
 type IFaceMock struct {
 	SetFilterFunc   func(device.PacketFilter) error
@@ -62,7 +66,7 @@ func TestManagerCreate(t *testing.T) {
 		SetFilterFunc: func(device.PacketFilter) error { return nil },
 	}
 
-	m, err := Create(ifaceMock, false)
+	m, err := Create(ifaceMock, false, flowLogger)
 	if err != nil {
 		t.Errorf("failed to create Manager: %v", err)
 		return
@@ -82,7 +86,7 @@ func TestManagerAddPeerFiltering(t *testing.T) {
 		},
 	}
 
-	m, err := Create(ifaceMock, false)
+	m, err := Create(ifaceMock, false, flowLogger)
 	if err != nil {
 		t.Errorf("failed to create Manager: %v", err)
 		return
@@ -92,9 +96,8 @@ func TestManagerAddPeerFiltering(t *testing.T) {
 	proto := fw.ProtocolTCP
 	port := &fw.Port{Values: []uint16{80}}
 	action := fw.ActionDrop
-	comment := "Test rule"
 
-	rule, err := m.AddPeerFiltering(ip, proto, nil, port, action, "", comment)
+	rule, err := m.AddPeerFiltering(nil, ip, proto, nil, port, action, "")
 	if err != nil {
 		t.Errorf("failed to add filtering: %v", err)
 		return
@@ -116,26 +119,25 @@ func TestManagerDeleteRule(t *testing.T) {
 		SetFilterFunc: func(device.PacketFilter) error { return nil },
 	}
 
-	m, err := Create(ifaceMock, false)
+	m, err := Create(ifaceMock, false, flowLogger)
 	if err != nil {
 		t.Errorf("failed to create Manager: %v", err)
 		return
 	}
 
-	ip := net.ParseIP("192.168.1.1")
+	ip := netip.MustParseAddr("192.168.1.1")
 	proto := fw.ProtocolTCP
 	port := &fw.Port{Values: []uint16{80}}
 	action := fw.ActionDrop
-	comment := "Test rule 2"
 
-	rule2, err := m.AddPeerFiltering(ip, proto, nil, port, action, "", comment)
+	rule2, err := m.AddPeerFiltering(nil, ip.AsSlice(), proto, nil, port, action, "")
 	if err != nil {
 		t.Errorf("failed to add filtering: %v", err)
 		return
 	}
 
 	for _, r := range rule2 {
-		if _, ok := m.incomingRules[ip.String()][r.ID()]; !ok {
+		if _, ok := m.incomingRules[ip][r.ID()]; !ok {
 			t.Errorf("rule2 is not in the incomingRules")
 		}
 	}
@@ -149,7 +151,7 @@ func TestManagerDeleteRule(t *testing.T) {
 	}
 
 	for _, r := range rule2 {
-		if _, ok := m.incomingRules[ip.String()][r.ID()]; ok {
+		if _, ok := m.incomingRules[ip][r.ID()]; ok {
 			t.Errorf("rule2 is not in the incomingRules")
 		}
 	}
@@ -160,7 +162,7 @@ func TestAddUDPPacketHook(t *testing.T) {
 		name       string
 		in         bool
 		expDir     fw.RuleDirection
-		ip         net.IP
+		ip         netip.Addr
 		dPort      uint16
 		hook       func([]byte) bool
 		expectedID string
@@ -169,7 +171,7 @@ func TestAddUDPPacketHook(t *testing.T) {
 			name:   "Test Outgoing UDP Packet Hook",
 			in:     false,
 			expDir: fw.RuleDirectionOUT,
-			ip:     net.IPv4(10, 168, 0, 1),
+			ip:     netip.MustParseAddr("10.168.0.1"),
 			dPort:  8000,
 			hook:   func([]byte) bool { return true },
 		},
@@ -177,7 +179,7 @@ func TestAddUDPPacketHook(t *testing.T) {
 			name:   "Test Incoming UDP Packet Hook",
 			in:     true,
 			expDir: fw.RuleDirectionIN,
-			ip:     net.IPv6loopback,
+			ip:     netip.MustParseAddr("::1"),
 			dPort:  9000,
 			hook:   func([]byte) bool { return false },
 		},
@@ -187,18 +189,18 @@ func TestAddUDPPacketHook(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			manager, err := Create(&IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
-			}, false)
+			}, false, flowLogger)
 			require.NoError(t, err)
 
 			manager.AddUDPPacketHook(tt.in, tt.ip, tt.dPort, tt.hook)
 
 			var addedRule PeerRule
 			if tt.in {
-				if len(manager.incomingRules[tt.ip.String()]) != 1 {
+				if len(manager.incomingRules[tt.ip]) != 1 {
 					t.Errorf("expected 1 incoming rule, got %d", len(manager.incomingRules))
 					return
 				}
-				for _, rule := range manager.incomingRules[tt.ip.String()] {
+				for _, rule := range manager.incomingRules[tt.ip] {
 					addedRule = rule
 				}
 			} else {
@@ -206,12 +208,12 @@ func TestAddUDPPacketHook(t *testing.T) {
 					t.Errorf("expected 1 outgoing rule, got %d", len(manager.outgoingRules))
 					return
 				}
-				for _, rule := range manager.outgoingRules[tt.ip.String()] {
+				for _, rule := range manager.outgoingRules[tt.ip] {
 					addedRule = rule
 				}
 			}
 
-			if !tt.ip.Equal(addedRule.ip) {
+			if tt.ip.Compare(addedRule.ip) != 0 {
 				t.Errorf("expected ip %s, got %s", tt.ip, addedRule.ip)
 				return
 			}
@@ -236,7 +238,7 @@ func TestManagerReset(t *testing.T) {
 		SetFilterFunc: func(device.PacketFilter) error { return nil },
 	}
 
-	m, err := Create(ifaceMock, false)
+	m, err := Create(ifaceMock, false, flowLogger)
 	if err != nil {
 		t.Errorf("failed to create Manager: %v", err)
 		return
@@ -246,9 +248,8 @@ func TestManagerReset(t *testing.T) {
 	proto := fw.ProtocolTCP
 	port := &fw.Port{Values: []uint16{80}}
 	action := fw.ActionDrop
-	comment := "Test rule"
 
-	_, err = m.AddPeerFiltering(ip, proto, nil, port, action, "", comment)
+	_, err = m.AddPeerFiltering(nil, ip, proto, nil, port, action, "")
 	if err != nil {
 		t.Errorf("failed to add filtering: %v", err)
 		return
@@ -279,7 +280,7 @@ func TestNotMatchByIP(t *testing.T) {
 		},
 	}
 
-	m, err := Create(ifaceMock, false)
+	m, err := Create(ifaceMock, false, flowLogger)
 	if err != nil {
 		t.Errorf("failed to create Manager: %v", err)
 		return
@@ -292,9 +293,8 @@ func TestNotMatchByIP(t *testing.T) {
 	ip := net.ParseIP("0.0.0.0")
 	proto := fw.ProtocolUDP
 	action := fw.ActionAccept
-	comment := "Test rule"
 
-	_, err = m.AddPeerFiltering(ip, proto, nil, nil, action, "", comment)
+	_, err = m.AddPeerFiltering(nil, ip, proto, nil, nil, action, "")
 	if err != nil {
 		t.Errorf("failed to add filtering: %v", err)
 		return
@@ -328,7 +328,7 @@ func TestNotMatchByIP(t *testing.T) {
 		return
 	}
 
-	if m.dropFilter(buf.Bytes()) {
+	if m.dropFilter(buf.Bytes(), 0) {
 		t.Errorf("expected packet to be accepted")
 		return
 	}
@@ -347,7 +347,7 @@ func TestRemovePacketHook(t *testing.T) {
 	}
 
 	// creating manager instance
-	manager, err := Create(iface, false)
+	manager, err := Create(iface, false, flowLogger)
 	if err != nil {
 		t.Fatalf("Failed to create Manager: %s", err)
 	}
@@ -357,7 +357,7 @@ func TestRemovePacketHook(t *testing.T) {
 
 	// Add a UDP packet hook
 	hookFunc := func(data []byte) bool { return true }
-	hookID := manager.AddUDPPacketHook(false, net.IPv4(192, 168, 0, 1), 8080, hookFunc)
+	hookID := manager.AddUDPPacketHook(false, netip.MustParseAddr("192.168.0.1"), 8080, hookFunc)
 
 	// Assert the hook is added by finding it in the manager's outgoing rules
 	found := false
@@ -393,7 +393,7 @@ func TestRemovePacketHook(t *testing.T) {
 func TestProcessOutgoingHooks(t *testing.T) {
 	manager, err := Create(&IFaceMock{
 		SetFilterFunc: func(device.PacketFilter) error { return nil },
-	}, false)
+	}, false, flowLogger)
 	require.NoError(t, err)
 
 	manager.wgNetwork = &net.IPNet{
@@ -401,7 +401,7 @@ func TestProcessOutgoingHooks(t *testing.T) {
 		Mask: net.CIDRMask(16, 32),
 	}
 	manager.udpTracker.Close()
-	manager.udpTracker = conntrack.NewUDPTracker(100*time.Millisecond, logger)
+	manager.udpTracker = conntrack.NewUDPTracker(100*time.Millisecond, logger, flowLogger)
 	defer func() {
 		require.NoError(t, manager.Close(nil))
 	}()
@@ -423,7 +423,7 @@ func TestProcessOutgoingHooks(t *testing.T) {
 	hookCalled := false
 	hookID := manager.AddUDPPacketHook(
 		false,
-		net.ParseIP("100.10.0.100"),
+		netip.MustParseAddr("100.10.0.100"),
 		53,
 		func([]byte) bool {
 			hookCalled = true
@@ -458,7 +458,7 @@ func TestProcessOutgoingHooks(t *testing.T) {
 	require.NoError(t, err)
 
 	// Test hook gets called
-	result := manager.processOutgoingHooks(buf.Bytes())
+	result := manager.processOutgoingHooks(buf.Bytes(), 0)
 	require.True(t, result)
 	require.True(t, hookCalled)
 
@@ -468,7 +468,7 @@ func TestProcessOutgoingHooks(t *testing.T) {
 	err = gopacket.SerializeLayers(buf, opts, ipv4)
 	require.NoError(t, err)
 
-	result = manager.processOutgoingHooks(buf.Bytes())
+	result = manager.processOutgoingHooks(buf.Bytes(), 0)
 	require.False(t, result)
 }
 
@@ -479,7 +479,7 @@ func TestUSPFilterCreatePerformance(t *testing.T) {
 			ifaceMock := &IFaceMock{
 				SetFilterFunc: func(device.PacketFilter) error { return nil },
 			}
-			manager, err := Create(ifaceMock, false)
+			manager, err := Create(ifaceMock, false, flowLogger)
 			require.NoError(t, err)
 			time.Sleep(time.Second)
 
@@ -494,7 +494,7 @@ func TestUSPFilterCreatePerformance(t *testing.T) {
 			start := time.Now()
 			for i := 0; i < testMax; i++ {
 				port := &fw.Port{Values: []uint16{uint16(1000 + i)}}
-				_, err = manager.AddPeerFiltering(ip, "tcp", nil, port, fw.ActionAccept, "", "accept HTTP traffic")
+				_, err = manager.AddPeerFiltering(nil, ip, "tcp", nil, port, fw.ActionAccept, "")
 
 				require.NoError(t, err, "failed to add rule")
 			}
@@ -506,7 +506,7 @@ func TestUSPFilterCreatePerformance(t *testing.T) {
 func TestStatefulFirewall_UDPTracking(t *testing.T) {
 	manager, err := Create(&IFaceMock{
 		SetFilterFunc: func(device.PacketFilter) error { return nil },
-	}, false)
+	}, false, flowLogger)
 	require.NoError(t, err)
 
 	manager.wgNetwork = &net.IPNet{
@@ -515,7 +515,7 @@ func TestStatefulFirewall_UDPTracking(t *testing.T) {
 	}
 
 	manager.udpTracker.Close() // Close the existing tracker
-	manager.udpTracker = conntrack.NewUDPTracker(200*time.Millisecond, logger)
+	manager.udpTracker = conntrack.NewUDPTracker(200*time.Millisecond, logger, flowLogger)
 	manager.decoders = sync.Pool{
 		New: func() any {
 			d := &decoder{
@@ -534,8 +534,8 @@ func TestStatefulFirewall_UDPTracking(t *testing.T) {
 	}()
 
 	// Set up packet parameters
-	srcIP := net.ParseIP("100.10.0.1")
-	dstIP := net.ParseIP("100.10.0.100")
+	srcIP := netip.MustParseAddr("100.10.0.1")
+	dstIP := netip.MustParseAddr("100.10.0.100")
 	srcPort := uint16(51334)
 	dstPort := uint16(53)
 
@@ -543,8 +543,8 @@ func TestStatefulFirewall_UDPTracking(t *testing.T) {
 	outboundIPv4 := &layers.IPv4{
 		TTL:      64,
 		Version:  4,
-		SrcIP:    srcIP,
-		DstIP:    dstIP,
+		SrcIP:    srcIP.AsSlice(),
+		DstIP:    dstIP.AsSlice(),
 		Protocol: layers.IPProtocolUDP,
 	}
 	outboundUDP := &layers.UDP{
@@ -569,15 +569,15 @@ func TestStatefulFirewall_UDPTracking(t *testing.T) {
 	require.NoError(t, err)
 
 	// Process outbound packet and verify connection tracking
-	drop := manager.DropOutgoing(outboundBuf.Bytes())
+	drop := manager.DropOutgoing(outboundBuf.Bytes(), 0)
 	require.False(t, drop, "Initial outbound packet should not be dropped")
 
 	// Verify connection was tracked
 	conn, exists := manager.udpTracker.GetConnection(srcIP, srcPort, dstIP, dstPort)
 
 	require.True(t, exists, "Connection should be tracked after outbound packet")
-	require.True(t, conntrack.ValidateIPs(conntrack.MakeIPAddr(srcIP), conn.SourceIP), "Source IP should match")
-	require.True(t, conntrack.ValidateIPs(conntrack.MakeIPAddr(dstIP), conn.DestIP), "Destination IP should match")
+	require.True(t, srcIP.Compare(conn.SourceIP) == 0, "Source IP should match")
+	require.True(t, dstIP.Compare(conn.DestIP) == 0, "Destination IP should match")
 	require.Equal(t, srcPort, conn.SourcePort, "Source port should match")
 	require.Equal(t, dstPort, conn.DestPort, "Destination port should match")
 
@@ -585,8 +585,8 @@ func TestStatefulFirewall_UDPTracking(t *testing.T) {
 	inboundIPv4 := &layers.IPv4{
 		TTL:      64,
 		Version:  4,
-		SrcIP:    dstIP, // Original destination is now source
-		DstIP:    srcIP, // Original source is now destination
+		SrcIP:    dstIP.AsSlice(), // Original destination is now source
+		DstIP:    srcIP.AsSlice(), // Original source is now destination
 		Protocol: layers.IPProtocolUDP,
 	}
 	inboundUDP := &layers.UDP{
@@ -636,7 +636,7 @@ func TestStatefulFirewall_UDPTracking(t *testing.T) {
 	for _, cp := range checkPoints {
 		time.Sleep(cp.sleep)
 
-		drop = manager.dropFilter(inboundBuf.Bytes())
+		drop = manager.dropFilter(inboundBuf.Bytes(), 0)
 		require.Equal(t, cp.shouldAllow, !drop, cp.description)
 
 		// If the connection should still be valid, verify it exists
@@ -685,7 +685,7 @@ func TestStatefulFirewall_UDPTracking(t *testing.T) {
 	}
 
 	// Create a new outbound connection for invalid tests
-	drop = manager.processOutgoingHooks(outboundBuf.Bytes())
+	drop = manager.processOutgoingHooks(outboundBuf.Bytes(), 0)
 	require.False(t, drop, "Second outbound packet should not be dropped")
 
 	for _, tc := range invalidCases {
@@ -707,7 +707,7 @@ func TestStatefulFirewall_UDPTracking(t *testing.T) {
 			require.NoError(t, err)
 
 			// Verify the invalid packet is dropped
-			drop = manager.dropFilter(testBuf.Bytes())
+			drop = manager.dropFilter(testBuf.Bytes(), 0)
 			require.True(t, drop, tc.description)
 		})
 	}

client/iface/device/device_filter.go

@@ -2,6 +2,7 @@ package device
 
 import (
 	"net"
+	"net/netip"
 	"sync"
 
 	"golang.zx2c4.com/wireguard/tun"
@@ -10,16 +11,16 @@ import (
 // PacketFilter interface for firewall abilities
 type PacketFilter interface {
 	// DropOutgoing filter outgoing packets from host to external destinations
-	DropOutgoing(packetData []byte) bool
+	DropOutgoing(packetData []byte, size int) bool
 
 	// DropIncoming filter incoming packets from external sources to host
-	DropIncoming(packetData []byte) bool
+	DropIncoming(packetData []byte, size int) bool
 
 	// AddUDPPacketHook calls hook when UDP packet from given direction matched
 	//
 	// Hook function returns flag which indicates should be the matched package dropped or not.
 	// Hook function receives raw network packet data as argument.
-	AddUDPPacketHook(in bool, ip net.IP, dPort uint16, hook func(packet []byte) bool) string
+	AddUDPPacketHook(in bool, ip netip.Addr, dPort uint16, hook func(packet []byte) bool) string
 
 	// RemovePacketHook removes hook by ID
 	RemovePacketHook(hookID string) error
@@ -57,7 +58,7 @@ func (d *FilteredDevice) Read(bufs [][]byte, sizes []int, offset int) (n int, er
 	}
 
 	for i := 0; i < n; i++ {
-		if filter.DropOutgoing(bufs[i][offset : offset+sizes[i]]) {
+		if filter.DropOutgoing(bufs[i][offset:offset+sizes[i]], sizes[i]) {
 			bufs = append(bufs[:i], bufs[i+1:]...)
 			sizes = append(sizes[:i], sizes[i+1:]...)
 			n--
@@ -81,7 +82,7 @@ func (d *FilteredDevice) Write(bufs [][]byte, offset int) (int, error) {
 	filteredBufs := make([][]byte, 0, len(bufs))
 	dropped := 0
 	for _, buf := range bufs {
-		if !filter.DropIncoming(buf[offset:]) {
+		if !filter.DropIncoming(buf[offset:], len(buf)) {
 			filteredBufs = append(filteredBufs, buf)
 			dropped++
 		}

client/iface/device/device_filter_test.go

@@ -146,7 +146,7 @@ func TestDeviceWrapperRead(t *testing.T) {
 		tun.EXPECT().Write(mockBufs, 0).Return(0, nil)
 
 		filter := mocks.NewMockPacketFilter(ctrl)
-		filter.EXPECT().DropIncoming(gomock.Any()).Return(true)
+		filter.EXPECT().DropIncoming(gomock.Any(), gomock.Any()).Return(true)
 
 		wrapped := newDeviceFilter(tun)
 		wrapped.filter = filter
@@ -201,7 +201,7 @@ func TestDeviceWrapperRead(t *testing.T) {
 				return 1, nil
 			})
 		filter := mocks.NewMockPacketFilter(ctrl)
-		filter.EXPECT().DropOutgoing(gomock.Any()).Return(true)
+		filter.EXPECT().DropOutgoing(gomock.Any(), gomock.Any()).Return(true)
 
 		wrapped := newDeviceFilter(tun)
 		wrapped.filter = filter

client/iface/mocks/filter.go

@@ -6,6 +6,7 @@ package mocks
 
 import (
 	net "net"
+	"net/netip"
 	reflect "reflect"
 
 	gomock "github.com/golang/mock/gomock"
@@ -35,7 +36,7 @@ func (m *MockPacketFilter) EXPECT() *MockPacketFilterMockRecorder {
 }
 
 // AddUDPPacketHook mocks base method.
-func (m *MockPacketFilter) AddUDPPacketHook(arg0 bool, arg1 net.IP, arg2 uint16, arg3 func([]byte) bool) string {
+func (m *MockPacketFilter) AddUDPPacketHook(arg0 bool, arg1 netip.Addr, arg2 uint16, arg3 func([]byte) bool) string {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "AddUDPPacketHook", arg0, arg1, arg2, arg3)
 	ret0, _ := ret[0].(string)
@@ -49,31 +50,31 @@ func (mr *MockPacketFilterMockRecorder) AddUDPPacketHook(arg0, arg1, arg2, arg3
 }
 
 // DropIncoming mocks base method.
-func (m *MockPacketFilter) DropIncoming(arg0 []byte) bool {
+func (m *MockPacketFilter) DropIncoming(arg0 []byte, arg1 int) bool {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "DropIncoming", arg0)
+	ret := m.ctrl.Call(m, "DropIncoming", arg0, arg1)
 	ret0, _ := ret[0].(bool)
 	return ret0
 }
 
 // DropIncoming indicates an expected call of DropIncoming.
-func (mr *MockPacketFilterMockRecorder) DropIncoming(arg0 interface{}) *gomock.Call {
+func (mr *MockPacketFilterMockRecorder) DropIncoming(arg0 interface{}, arg1 any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DropIncoming", reflect.TypeOf((*MockPacketFilter)(nil).DropIncoming), arg0)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DropIncoming", reflect.TypeOf((*MockPacketFilter)(nil).DropIncoming), arg0, arg1)
 }
 
 // DropOutgoing mocks base method.
-func (m *MockPacketFilter) DropOutgoing(arg0 []byte) bool {
+func (m *MockPacketFilter) DropOutgoing(arg0 []byte, arg1 int) bool {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "DropOutgoing", arg0)
+	ret := m.ctrl.Call(m, "DropOutgoing", arg0, arg1)
 	ret0, _ := ret[0].(bool)
 	return ret0
 }
 
 // DropOutgoing indicates an expected call of DropOutgoing.
-func (mr *MockPacketFilterMockRecorder) DropOutgoing(arg0 interface{}) *gomock.Call {
+func (mr *MockPacketFilterMockRecorder) DropOutgoing(arg0 interface{}, arg1 any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DropOutgoing", reflect.TypeOf((*MockPacketFilter)(nil).DropOutgoing), arg0)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DropOutgoing", reflect.TypeOf((*MockPacketFilter)(nil).DropOutgoing), arg0, arg1)
 }
 
 // RemovePacketHook mocks base method.

client/internal/acl/manager.go

@@ -28,6 +28,11 @@ type Manager interface {
 	ApplyFiltering(networkMap *mgmProto.NetworkMap)
 }
 
+type protoMatch struct {
+	ips      map[string]int
+	policyID []byte
+}
+
 // DefaultManager uses firewall manager to handle
 type DefaultManager struct {
 	firewall       firewall.Manager
@@ -240,7 +245,7 @@ func (d *DefaultManager) applyRouteACL(rule *mgmProto.RouteFirewallRule) (id.Rul
 
 	dPorts := convertPortInfo(rule.PortInfo)
 
-	addedRule, err := d.firewall.AddRouteFiltering(sources, destination, protocol, nil, dPorts, action)
+	addedRule, err := d.firewall.AddRouteFiltering(rule.PolicyID, sources, destination, protocol, nil, dPorts, action)
 	if err != nil {
 		return "", fmt.Errorf("add route rule: %w", err)
 	}
@@ -281,7 +286,7 @@ func (d *DefaultManager) protoRuleToFirewallRule(
 		}
 	}
 
-	ruleID := d.getPeerRuleID(ip, protocol, int(r.Direction), port, action, "")
+	ruleID := d.getPeerRuleID(ip, protocol, int(r.Direction), port, action)
 	if rulesPair, ok := d.peerRulesPairs[ruleID]; ok {
 		return ruleID, rulesPair, nil
 	}
@@ -289,11 +294,11 @@ func (d *DefaultManager) protoRuleToFirewallRule(
 	var rules []firewall.Rule
 	switch r.Direction {
 	case mgmProto.RuleDirection_IN:
-		rules, err = d.addInRules(ip, protocol, port, action, ipsetName, "")
+		rules, err = d.addInRules(r.PolicyID, ip, protocol, port, action, ipsetName)
 	case mgmProto.RuleDirection_OUT:
 		// TODO: Remove this soon. Outbound rules are obsolete.
 		// We only maintain this for return traffic (inbound dir) which is now handled by the stateful firewall already
-		rules, err = d.addOutRules(ip, protocol, port, action, ipsetName, "")
+		rules, err = d.addOutRules(r.PolicyID, ip, protocol, port, action, ipsetName)
 	default:
 		return "", nil, fmt.Errorf("invalid direction, skipping firewall rule")
 	}
@@ -322,14 +327,14 @@ func portInfoEmpty(portInfo *mgmProto.PortInfo) bool {
 }
 
 func (d *DefaultManager) addInRules(
+	id []byte,
 	ip net.IP,
 	protocol firewall.Protocol,
 	port *firewall.Port,
 	action firewall.Action,
 	ipsetName string,
-	comment string,
 ) ([]firewall.Rule, error) {
-	rule, err := d.firewall.AddPeerFiltering(ip, protocol, nil, port, action, ipsetName, comment)
+	rule, err := d.firewall.AddPeerFiltering(id, ip, protocol, nil, port, action, ipsetName)
 	if err != nil {
 		return nil, fmt.Errorf("add firewall rule: %w", err)
 	}
@@ -338,18 +343,18 @@ func (d *DefaultManager) addInRules(
 }
 
 func (d *DefaultManager) addOutRules(
+	id []byte,
 	ip net.IP,
 	protocol firewall.Protocol,
 	port *firewall.Port,
 	action firewall.Action,
 	ipsetName string,
-	comment string,
 ) ([]firewall.Rule, error) {
 	if shouldSkipInvertedRule(protocol, port) {
 		return nil, nil
 	}
 
-	rule, err := d.firewall.AddPeerFiltering(ip, protocol, port, nil, action, ipsetName, comment)
+	rule, err := d.firewall.AddPeerFiltering(id, ip, protocol, port, nil, action, ipsetName)
 	if err != nil {
 		return nil, fmt.Errorf("add firewall rule: %w", err)
 	}
@@ -364,9 +369,8 @@ func (d *DefaultManager) getPeerRuleID(
 	direction int,
 	port *firewall.Port,
 	action firewall.Action,
-	comment string,
 ) id.RuleID {
-	idStr := ip.String() + string(proto) + strconv.Itoa(direction) + strconv.Itoa(int(action)) + comment
+	idStr := ip.String() + string(proto) + strconv.Itoa(direction) + strconv.Itoa(int(action))
 	if port != nil {
 		idStr += port.String()
 	}
@@ -389,10 +393,8 @@ func (d *DefaultManager) squashAcceptRules(
 		}
 	}
 
-	type protoMatch map[mgmProto.RuleProtocol]map[string]int
-
-	in := protoMatch{}
-	out := protoMatch{}
+	in := map[mgmProto.RuleProtocol]*protoMatch{}
+	out := map[mgmProto.RuleProtocol]*protoMatch{}
 
 	// trace which type of protocols was squashed
 	squashedRules := []*mgmProto.FirewallRule{}
@@ -405,14 +407,18 @@ func (d *DefaultManager) squashAcceptRules(
 	// 2. Any of rule contains Port.
 	//
 	// We zeroed this to notify squash function that this protocol can't be squashed.
-	addRuleToCalculationMap := func(i int, r *mgmProto.FirewallRule, protocols protoMatch) {
+	addRuleToCalculationMap := func(i int, r *mgmProto.FirewallRule, protocols map[mgmProto.RuleProtocol]*protoMatch) {
 		drop := r.Action == mgmProto.RuleAction_DROP || r.Port != ""
 		if drop {
-			protocols[r.Protocol] = map[string]int{}
+			protocols[r.Protocol] = &protoMatch{ips: map[string]int{}}
 			return
 		}
 		if _, ok := protocols[r.Protocol]; !ok {
-			protocols[r.Protocol] = map[string]int{}
+			protocols[r.Protocol] = &protoMatch{
+				ips: map[string]int{},
+				// store the first encountered PolicyID for this protocol
+				policyID: r.PolicyID,
+			}
 		}
 
 		// special case, when we receive this all network IP address
@@ -424,7 +430,7 @@ func (d *DefaultManager) squashAcceptRules(
 			return
 		}
 
-		ipset := protocols[r.Protocol]
+		ipset := protocols[r.Protocol].ips
 
 		if _, ok := ipset[r.PeerIP]; ok {
 			return
@@ -450,9 +456,10 @@ func (d *DefaultManager) squashAcceptRules(
 		mgmProto.RuleProtocol_UDP,
 	}
 
-	squash := func(matches protoMatch, direction mgmProto.RuleDirection) {
+	squash := func(matches map[mgmProto.RuleProtocol]*protoMatch, direction mgmProto.RuleDirection) {
 		for _, protocol := range protocolOrders {
-			if ipset, ok := matches[protocol]; !ok || len(ipset) != totalIPs || len(ipset) < 2 {
+			match, ok := matches[protocol]
+			if !ok || len(match.ips) != totalIPs || len(match.ips) < 2 {
 				// don't squash if :
 				// 1. Rules not cover all peers in the network
 				// 2. Rules cover only one peer in the network.
@@ -465,6 +472,7 @@ func (d *DefaultManager) squashAcceptRules(
 				Direction: direction,
 				Action:    mgmProto.RuleAction_ACCEPT,
 				Protocol:  protocol,
+				PolicyID:  match.policyID,
 			})
 			squashedProtocols[protocol] = struct{}{}
 
@@ -493,9 +501,9 @@ func (d *DefaultManager) squashAcceptRules(
 	// if we also have other not squashed rules.
 	for i, r := range networkMap.FirewallRules {
 		if _, ok := squashedProtocols[r.Protocol]; ok {
-			if m, ok := in[r.Protocol]; ok && m[r.PeerIP] == i {
+			if m, ok := in[r.Protocol]; ok && m.ips[r.PeerIP] == i {
 				continue
-			} else if m, ok := out[r.Protocol]; ok && m[r.PeerIP] == i {
+			} else if m, ok := out[r.Protocol]; ok && m.ips[r.PeerIP] == i {
 				continue
 			}
 		}

client/internal/acl/manager_test.go

@@ -1,6 +1,7 @@
 package acl
 
 import (
+	"context"
 	"net"
 	"testing"
 
@@ -10,9 +11,12 @@ import (
 	"github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
 	"github.com/netbirdio/netbird/client/internal/acl/mocks"
+	"github.com/netbirdio/netbird/client/internal/netflow"
 	mgmProto "github.com/netbirdio/netbird/management/proto"
 )
 
+var flowLogger = netflow.NewManager(context.Background(), nil, []byte{}, nil).GetLogger()
+
 func TestDefaultManager(t *testing.T) {
 	networkMap := &mgmProto.NetworkMap{
 		FirewallRules: []*mgmProto.FirewallRule{
@@ -52,7 +56,7 @@ func TestDefaultManager(t *testing.T) {
 	ifaceMock.EXPECT().GetWGDevice().Return(nil).AnyTimes()
 
 	// we receive one rule from the management so for testing purposes ignore it
-	fw, err := firewall.NewFirewall(ifaceMock, nil, false)
+	fw, err := firewall.NewFirewall(ifaceMock, nil, flowLogger, false)
 	if err != nil {
 		t.Errorf("create firewall: %v", err)
 		return
@@ -346,7 +350,7 @@ func TestDefaultManagerEnableSSHRules(t *testing.T) {
 	ifaceMock.EXPECT().GetWGDevice().Return(nil).AnyTimes()
 
 	// we receive one rule from the management so for testing purposes ignore it
-	fw, err := firewall.NewFirewall(ifaceMock, nil, false)
+	fw, err := firewall.NewFirewall(ifaceMock, nil, flowLogger, false)
 	if err != nil {
 		t.Errorf("create firewall: %v", err)
 		return

client/internal/dns/file_repair_unix.go

@@ -22,7 +22,7 @@ var (
 	}
 )
 
-type repairConfFn func([]string, string, *resolvConf, *statemanager.Manager) error
+type repairConfFn func([]string, string, *resolvConf, statemanager.Manager) error
 
 type repair struct {
 	operationFile string
@@ -42,7 +42,7 @@ func newRepair(operationFile string, updateFn repairConfFn) *repair {
 	}
 }
 
-func (f *repair) watchFileChanges(nbSearchDomains []string, nbNameserverIP string, stateManager *statemanager.Manager) {
+func (f *repair) watchFileChanges(nbSearchDomains []string, nbNameserverIP string, stateManager statemanager.Manager) {
 	if f.inotify != nil {
 		return
 	}

client/internal/dns/file_repair_unix_test.go

@@ -105,7 +105,7 @@ nameserver 8.8.8.8`,
 
 			var changed bool
 			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-			updateFn := func([]string, string, *resolvConf, *statemanager.Manager) error {
+			updateFn := func([]string, string, *resolvConf, statemanager.Manager) error {
 				changed = true
 				cancel()
 				return nil
@@ -152,7 +152,7 @@ searchdomain netbird.cloud something`
 
 	var changed bool
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-	updateFn := func([]string, string, *resolvConf, *statemanager.Manager) error {
+	updateFn := func([]string, string, *resolvConf, statemanager.Manager) error {
 		changed = true
 		cancel()
 		return nil

client/internal/dns/file_unix.go

@@ -48,7 +48,7 @@ func (f *fileConfigurator) supportCustomPort() bool {
 	return false
 }
 
-func (f *fileConfigurator) applyDNSConfig(config HostDNSConfig, stateManager *statemanager.Manager) error {
+func (f *fileConfigurator) applyDNSConfig(config HostDNSConfig, stateManager statemanager.Manager) error {
 	backupFileExist := f.isBackupFileExist()
 	if !config.RouteAll {
 		if backupFileExist {
@@ -86,7 +86,7 @@ func (f *fileConfigurator) applyDNSConfig(config HostDNSConfig, stateManager *st
 	return nil
 }
 
-func (f *fileConfigurator) updateConfig(nbSearchDomains []string, nbNameserverIP string, cfg *resolvConf, stateManager *statemanager.Manager) error {
+func (f *fileConfigurator) updateConfig(nbSearchDomains []string, nbNameserverIP string, cfg *resolvConf, stateManager statemanager.Manager) error {
 	searchDomainList := mergeSearchDomains(nbSearchDomains, cfg.searchDomains)
 	nameServers := generateNsList(nbNameserverIP, cfg)
 

client/internal/dns/host.go

@@ -17,7 +17,7 @@ const (
 )
 
 type hostManager interface {
-	applyDNSConfig(config HostDNSConfig, stateManager *statemanager.Manager) error
+	applyDNSConfig(config HostDNSConfig, stateManager statemanager.Manager) error
 	restoreHostDNS() error
 	supportCustomPort() bool
 	string() string
@@ -43,14 +43,14 @@ type DomainConfig struct {
 }
 
 type mockHostConfigurator struct {
-	applyDNSConfigFunc            func(config HostDNSConfig, stateManager *statemanager.Manager) error
+	applyDNSConfigFunc            func(config HostDNSConfig, stateManager statemanager.Manager) error
 	restoreHostDNSFunc            func() error
 	supportCustomPortFunc         func() bool
 	restoreUncleanShutdownDNSFunc func(*netip.Addr) error
 	stringFunc                    func() string
 }
 
-func (m *mockHostConfigurator) applyDNSConfig(config HostDNSConfig, stateManager *statemanager.Manager) error {
+func (m *mockHostConfigurator) applyDNSConfig(config HostDNSConfig, stateManager statemanager.Manager) error {
 	if m.applyDNSConfigFunc != nil {
 		return m.applyDNSConfigFunc(config, stateManager)
 	}
@@ -80,7 +80,7 @@ func (m *mockHostConfigurator) string() string {
 
 func newNoopHostMocker() hostManager {
 	return &mockHostConfigurator{
-		applyDNSConfigFunc:            func(config HostDNSConfig, stateManager *statemanager.Manager) error { return nil },
+		applyDNSConfigFunc:            func(config HostDNSConfig, stateManager statemanager.Manager) error { return nil },
 		restoreHostDNSFunc:            func() error { return nil },
 		supportCustomPortFunc:         func() bool { return true },
 		restoreUncleanShutdownDNSFunc: func(*netip.Addr) error { return nil },
@@ -122,7 +122,7 @@ func dnsConfigToHostDNSConfig(dnsConfig nbdns.Config, ip string, port int) HostD
 
 type noopHostConfigurator struct{}
 
-func (n noopHostConfigurator) applyDNSConfig(HostDNSConfig, *statemanager.Manager) error {
+func (n noopHostConfigurator) applyDNSConfig(HostDNSConfig, statemanager.Manager) error {
 	return nil
 }
 

client/internal/dns/host_android.go

@@ -11,7 +11,7 @@ func newHostManager() (*androidHostManager, error) {
 	return &androidHostManager{}, nil
 }
 
-func (a androidHostManager) applyDNSConfig(HostDNSConfig, *statemanager.Manager) error {
+func (a androidHostManager) applyDNSConfig(HostDNSConfig, statemanager.Manager) error {
 	return nil
 }
 

client/internal/dns/host_darwin.go

@@ -5,6 +5,7 @@ package dns
 import (
 	"bufio"
 	"bytes"
+	"errors"
 	"fmt"
 	"io"
 	"net"
@@ -49,7 +50,7 @@ func (s *systemConfigurator) supportCustomPort() bool {
 	return true
 }
 
-func (s *systemConfigurator) applyDNSConfig(config HostDNSConfig, stateManager *statemanager.Manager) error {
+func (s *systemConfigurator) applyDNSConfig(config HostDNSConfig, stateManager statemanager.Manager) error {
 	var err error
 
 	if err := stateManager.UpdateState(&ShutdownState{}); err != nil {
@@ -200,8 +201,12 @@ func (s *systemConfigurator) recordSystemDNSSettings(force bool) error {
 func (s *systemConfigurator) getSystemDNSSettings() (SystemDNSSettings, error) {
 	primaryServiceKey, _, err := s.getPrimaryService()
 	if err != nil || primaryServiceKey == "" {
+		if err == nil {
+			err = errors.New("primary service key not found")
+		}
 		return SystemDNSSettings{}, fmt.Errorf("couldn't find the primary service key: %w", err)
 	}
+
 	dnsServiceKey := getKeyWithInput(primaryServiceStateKeyFormat, primaryServiceKey)
 	line := buildCommandLine("show", dnsServiceKey, "")
 	stdinCommands := wrapCommand(line)
@@ -379,7 +384,7 @@ func buildWriteStateOperation(operation, state, commands string) string {
 	return fmt.Sprintf("d.init\n%s %s\n%s\nset %s\n", operation, state, commands, state)
 }
 
-func runSystemConfigCommand(command string) ([]byte, error) {
+var runSystemConfigCommand = func(command string) ([]byte, error) {
 	cmd := exec.Command(scutilPath)
 	cmd.Stdin = strings.NewReader(command)
 	out, err := cmd.Output()

client/internal/dns/host_darwin_test.go

@@ -0,0 +1,210 @@
+package dns
+
+import (
+	"errors"
+	"os/exec"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"go.uber.org/mock/gomock"
+
+	"github.com/netbirdio/netbird/client/internal/statemanager/mocks"
+)
+
+// MockCommander to mock exec.Command
+type MockCommander struct {
+	mock.Mock
+}
+
+func (m *MockCommander) Command(name string, arg ...string) *exec.Cmd {
+	args := m.Called(name, arg)
+	return args.Get(0).(*exec.Cmd)
+}
+
+func TestNewHostManager(t *testing.T) {
+	tests := []struct {
+		name    string
+		wantErr bool
+	}{
+		{
+			name:    "successful creation",
+			wantErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := newHostManager()
+			if tt.wantErr {
+				assert.Error(t, err)
+				return
+			}
+			assert.NoError(t, err)
+			assert.NotNil(t, got)
+			assert.NotNil(t, got.createdKeys)
+		})
+	}
+}
+
+func TestApplyDNSConfig(t *testing.T) {
+	type mockSetup struct {
+		stateManagerError error
+		commandOutput     []byte
+		commandError      error
+	}
+
+	tests := []struct {
+		name      string
+		config    HostDNSConfig
+		mockSetup mockSetup
+		wantErr   bool
+	}{
+		{
+			name: "successful apply with search domains",
+			config: HostDNSConfig{
+				RouteAll: true,
+				Domains: []DomainConfig{
+					{Domain: "example.com", MatchOnly: false},
+					{Domain: "test.com", MatchOnly: true},
+				},
+				ServerIP:   "1.1.1.1",
+				ServerPort: 53,
+			},
+			mockSetup: mockSetup{
+				stateManagerError: nil,
+				commandOutput: []byte(`
+			PrimaryService : ABC123
+			Router : 192.168.1.1
+			DomainName : example.com
+			SearchDomains : <array> {
+			  0 : test.com
+			}
+			ServerAddresses : <array> {
+			  0 : 1.1.1.1
+			}
+			`),
+				commandError: nil,
+			},
+			wantErr: false,
+		},
+		{
+			name: "state manager error",
+			config: HostDNSConfig{
+				ServerIP: "1.1.1.1",
+			},
+			mockSetup: mockSetup{
+				stateManagerError: errors.New("state error"),
+			},
+			wantErr: false, // Function does not return an error, it only logs it.
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Setup mocks
+			s := &systemConfigurator{
+				createdKeys: make(map[string]struct{}),
+			}
+
+			ctrl := gomock.NewController(t)
+			defer ctrl.Finish() // Ensures all expectations are met
+
+			mockState := mocks.NewMockManager(ctrl)
+			mockCmd := new(MockCommander)
+
+			// Mock UpdateState
+			mockState.EXPECT().UpdateState(gomock.Any()).Return(tt.mockSetup.stateManagerError).AnyTimes()
+
+			// Mock all expected command executions
+			// mockCmd.On("Command", dscacheutilPath, "-flushcache").Return(&exec.Cmd{}).Once()
+			// mockCmd.On("Command", "killall", "-HUP", "mDNSResponder").Return(&exec.Cmd{}).Once()
+			// mockCmd.On("Command", scutilPath).Return(&exec.Cmd{}).Once() // For runSystemConfigCommand
+
+			// Mock `runSystemConfigCommand`
+			originalRunCommand := runSystemConfigCommand
+			runSystemConfigCommand = func(command string) ([]byte, error) {
+				return tt.mockSetup.commandOutput, tt.mockSetup.commandError
+			}
+			defer func() { runSystemConfigCommand = originalRunCommand }()
+
+			err := s.applyDNSConfig(tt.config, mockState)
+			if tt.wantErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+			}
+
+			mockCmd.AssertExpectations(t) // Ensure Command() is called
+		})
+	}
+}
+
+func TestGetSystemDNSSettings(t *testing.T) {
+	tests := []struct {
+		name          string
+		commandOutput []byte
+		commandError  error
+		wantSettings  SystemDNSSettings
+		wantErr       bool
+	}{
+		{
+			name: "successful retrieval",
+			commandOutput: []byte(`
+PrimaryService : ABC123
+Router : 192.168.1.1
+---
+DomainName : example.com
+SearchDomains : <array> {
+  0 : test.com
+}
+ServerAddresses : <array> {
+  0 : 1.1.1.1
+}
+`),
+			wantSettings: SystemDNSSettings{
+				Domains:    []string{"example.com", "test.com"},
+				ServerIP:   "1.1.1.1",
+				ServerPort: 53,
+			},
+			wantErr: false,
+		},
+		{
+			name:         "command error",
+			commandError: errors.New("command failed"),
+			wantErr:      true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			s := &systemConfigurator{
+				createdKeys: make(map[string]struct{}),
+			}
+
+			originalRunCommand := runSystemConfigCommand
+			runSystemConfigCommand = func(command string) ([]byte, error) {
+				return tt.commandOutput, tt.commandError
+			}
+			defer func() { runSystemConfigCommand = originalRunCommand }()
+
+			got, err := s.getSystemDNSSettings()
+			if tt.wantErr {
+				assert.Error(t, err)
+				return
+			}
+			assert.NoError(t, err)
+			assert.Equal(t, tt.wantSettings, got)
+		})
+	}
+}
+
+func TestSupportCustomPort(t *testing.T) {
+	s := &systemConfigurator{}
+	assert.True(t, s.supportCustomPort())
+}
+
+func TestString(t *testing.T) {
+	s := &systemConfigurator{}
+	assert.Equal(t, "scutil", s.string())
+}

client/internal/dns/host_ios.go

@@ -20,7 +20,7 @@ func newHostManager(dnsManager IosDnsManager) (*iosHostManager, error) {
 	}, nil
 }
 
-func (a iosHostManager) applyDNSConfig(config HostDNSConfig, _ *statemanager.Manager) error {
+func (a iosHostManager) applyDNSConfig(config HostDNSConfig, _ statemanager.Manager) error {
 	jsonData, err := json.Marshal(config)
 	if err != nil {
 		return fmt.Errorf("marshal: %w", err)

client/internal/dns/host_windows.go

@@ -74,7 +74,7 @@ func (r *registryConfigurator) supportCustomPort() bool {
 	return false
 }
 
-func (r *registryConfigurator) applyDNSConfig(config HostDNSConfig, stateManager *statemanager.Manager) error {
+func (r *registryConfigurator) applyDNSConfig(config HostDNSConfig, stateManager statemanager.Manager) error {
 	if config.RouteAll {
 		if err := r.addDNSSetupForAll(config.ServerIP); err != nil {
 			return fmt.Errorf("add dns setup: %w", err)

client/internal/dns/network_manager_unix.go

@@ -103,7 +103,7 @@ func (n *networkManagerDbusConfigurator) supportCustomPort() bool {
 	return false
 }
 
-func (n *networkManagerDbusConfigurator) applyDNSConfig(config HostDNSConfig, stateManager *statemanager.Manager) error {
+func (n *networkManagerDbusConfigurator) applyDNSConfig(config HostDNSConfig, stateManager statemanager.Manager) error {
 	connSettings, configVersion, err := n.getAppliedConnectionSettings()
 	if err != nil {
 		return fmt.Errorf("retrieving the applied connection settings, error: %w", err)

client/internal/dns/resolvconf_unix.go

@@ -84,7 +84,7 @@ func (r *resolvconf) supportCustomPort() bool {
 	return false
 }
 
-func (r *resolvconf) applyDNSConfig(config HostDNSConfig, stateManager *statemanager.Manager) error {
+func (r *resolvconf) applyDNSConfig(config HostDNSConfig, stateManager statemanager.Manager) error {
 	var err error
 	if !config.RouteAll {
 		err = r.restoreHostDNS()

client/internal/dns/server.go

@@ -75,7 +75,7 @@ type DefaultServer struct {
 	iosDnsManager        IosDnsManager
 
 	statusRecorder *peer.Status
-	stateManager   *statemanager.Manager
+	stateManager   statemanager.Manager
 }
 
 type handlerWithStop interface {
@@ -99,7 +99,7 @@ func NewDefaultServer(
 	wgInterface WGIface,
 	customAddress string,
 	statusRecorder *peer.Status,
-	stateManager *statemanager.Manager,
+	stateManager statemanager.Manager,
 	disableSys bool,
 ) (*DefaultServer, error) {
 	var addrPort *netip.AddrPort
@@ -161,7 +161,7 @@ func newDefaultServer(
 	wgInterface WGIface,
 	dnsService service,
 	statusRecorder *peer.Status,
-	stateManager *statemanager.Manager,
+	stateManager statemanager.Manager,
 	disableSys bool,
 ) *DefaultServer {
 	ctx, stop := context.WithCancel(ctx)

client/internal/dns/server_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/netbirdio/netbird/client/iface/device"
 	pfmock "github.com/netbirdio/netbird/client/iface/mocks"
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
+	"github.com/netbirdio/netbird/client/internal/netflow"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 	"github.com/netbirdio/netbird/client/internal/stdnet"
@@ -30,6 +31,8 @@ import (
 	"github.com/netbirdio/netbird/formatter"
 )
 
+var flowLogger = netflow.NewManager(context.Background(), nil, []byte{}, nil).GetLogger()
+
 type mocWGIface struct {
 	filter device.PacketFilter
 }
@@ -456,7 +459,7 @@ func TestDNSFakeResolverHandleUpdates(t *testing.T) {
 	}
 
 	packetfilter := pfmock.NewMockPacketFilter(ctrl)
-	packetfilter.EXPECT().DropOutgoing(gomock.Any()).AnyTimes()
+	packetfilter.EXPECT().DropOutgoing(gomock.Any(), gomock.Any()).AnyTimes()
 	packetfilter.EXPECT().AddUDPPacketHook(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any())
 	packetfilter.EXPECT().RemovePacketHook(gomock.Any())
 	packetfilter.EXPECT().SetNetwork(ipNet)
@@ -644,7 +647,7 @@ func TestDNSServerUpstreamDeactivateCallback(t *testing.T) {
 	}
 
 	var domainsUpdate string
-	hostManager.applyDNSConfigFunc = func(config HostDNSConfig, statemanager *statemanager.Manager) error {
+	hostManager.applyDNSConfigFunc = func(config HostDNSConfig, statemanager statemanager.Manager) error {
 		domains := []string{}
 		for _, item := range config.Domains {
 			if item.Disabled {
@@ -917,7 +920,7 @@ func createWgInterfaceWithBind(t *testing.T) (*iface.WGIface, error) {
 		return nil, err
 	}
 
-	pf, err := uspfilter.Create(wgIface, false)
+	pf, err := uspfilter.Create(wgIface, false, flowLogger)
 	if err != nil {
 		t.Fatalf("failed to create uspfilter: %v", err)
 		return nil, err

client/internal/dns/service_memory.go

@@ -2,7 +2,7 @@ package dns
 
 import (
 	"fmt"
-	"net"
+	"net/netip"
 	"sync"
 
 	"github.com/google/gopacket"
@@ -117,5 +117,10 @@ func (s *ServiceViaMemory) filterDNSTraffic() (string, error) {
 		return true
 	}
 
-	return filter.AddUDPPacketHook(false, net.ParseIP(s.runtimeIP), uint16(s.runtimePort), hook), nil
+	ip, err := netip.ParseAddr(s.runtimeIP)
+	if err != nil {
+		return "", fmt.Errorf("parse runtime ip: %w", err)
+	}
+
+	return filter.AddUDPPacketHook(false, ip, uint16(s.runtimePort), hook), nil
 }

client/internal/dns/systemd_linux.go

@@ -87,7 +87,7 @@ func (s *systemdDbusConfigurator) supportCustomPort() bool {
 	return true
 }
 
-func (s *systemdDbusConfigurator) applyDNSConfig(config HostDNSConfig, stateManager *statemanager.Manager) error {
+func (s *systemdDbusConfigurator) applyDNSConfig(config HostDNSConfig, stateManager statemanager.Manager) error {
 	parsedIP, err := netip.ParseAddr(config.ServerIP)
 	if err != nil {
 		return fmt.Errorf("unable to parse ip address, error: %w", err)

client/internal/dns/unclean_shutdown_unix.go

@@ -35,7 +35,7 @@ func (s *ShutdownState) Cleanup() error {
 }
 
 // TODO: move file contents to state manager
-func createUncleanShutdownIndicator(sourcePath string, dnsAddressStr string, stateManager *statemanager.Manager) error {
+func createUncleanShutdownIndicator(sourcePath string, dnsAddressStr string, stateManager statemanager.Manager) error {
 	dnsAddress, err := netip.ParseAddr(dnsAddressStr)
 	if err != nil {
 		return fmt.Errorf("parse dns address %s: %w", dnsAddressStr, err)

client/internal/dnsfwd/manager.go

@@ -88,7 +88,7 @@ func (h *Manager) allowDNSFirewall() error {
 		return nil
 	}
 
-	dnsRules, err := h.firewall.AddPeerFiltering(net.IP{0, 0, 0, 0}, firewall.ProtocolUDP, nil, dport, firewall.ActionAccept, "", "")
+	dnsRules, err := h.firewall.AddPeerFiltering(nil, net.IP{0, 0, 0, 0}, firewall.ProtocolUDP, nil, dport, firewall.ActionAccept, "")
 	if err != nil {
 		log.Errorf("failed to add allow DNS router rules, err: %v", err)
 		return err

client/internal/engine.go

@@ -34,6 +34,8 @@ import (
 	"github.com/netbirdio/netbird/client/internal/dns"
 	"github.com/netbirdio/netbird/client/internal/dnsfwd"
 	"github.com/netbirdio/netbird/client/internal/ingressgw"
+	"github.com/netbirdio/netbird/client/internal/netflow"
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 	"github.com/netbirdio/netbird/client/internal/networkmonitor"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/peer/guard"
@@ -182,13 +184,14 @@ type Engine struct {
 	checks []*mgmProto.Checks
 
 	relayManager *relayClient.Manager
-	stateManager *statemanager.Manager
+	stateManager statemanager.Manager
 	srWatcher    *guard.SRWatcher
 
 	// Network map persistence
 	persistNetworkMap bool
 	latestNetworkMap  *mgmProto.NetworkMap
 	connSemaphore     *semaphoregroup.SemaphoreGroup
+	flowManager       nftypes.FlowManager
 }
 
 // Peer is an instance of the Connection Peer
@@ -308,6 +311,12 @@ func (e *Engine) Stop() error {
 	time.Sleep(500 * time.Millisecond)
 
 	e.close()
+
+	// stop flow manager after wg interface is gone
+	if e.flowManager != nil {
+		e.flowManager.Close()
+	}
+
 	log.Infof("stopped Netbird Engine")
 
 	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
@@ -342,6 +351,10 @@ func (e *Engine) Start() error {
 	}
 	e.wgInterface = wgIface
 
+	// start flow manager right after interface creation
+	publicKey := e.config.WgPrivateKey.PublicKey()
+	e.flowManager = netflow.NewManager(e.ctx, e.wgInterface, publicKey[:], e.statusRecorder)
+
 	if e.config.RosenpassEnabled {
 		log.Infof("rosenpass is enabled")
 		if e.config.RosenpassPermissive {
@@ -448,7 +461,7 @@ func (e *Engine) createFirewall() error {
 	}
 
 	var err error
-	e.firewall, err = firewall.NewFirewall(e.wgInterface, e.stateManager, e.config.DisableServerRoutes)
+	e.firewall, err = firewall.NewFirewall(e.wgInterface, e.stateManager, e.flowManager.GetLogger(), e.config.DisableServerRoutes)
 	if err != nil || e.firewall == nil {
 		log.Errorf("failed creating firewall manager: %s", err)
 		return nil
@@ -482,13 +495,13 @@ func (e *Engine) initFirewall() error {
 
 	// this rule is static and will be torn down on engine down by the firewall manager
 	if _, err := e.firewall.AddPeerFiltering(
+		nil,
 		net.IP{0, 0, 0, 0},
 		firewallManager.ProtocolUDP,
 		nil,
 		&port,
 		firewallManager.ActionAccept,
 		"",
-		"",
 	); err != nil {
 		log.Errorf("failed to allow rosenpass interface traffic: %v", err)
 		return nil
@@ -512,6 +525,7 @@ func (e *Engine) blockLanAccess() {
 	v4 := netip.PrefixFrom(netip.IPv4Unspecified(), 0)
 	for _, network := range toBlock {
 		if _, err := e.firewall.AddRouteFiltering(
+			nil,
 			[]netip.Prefix{v4},
 			network,
 			firewallManager.ProtocolALL,
@@ -642,25 +656,14 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 		stunTurn = append(stunTurn, e.TURNs...)
 		e.stunTurn.Store(stunTurn)
 
-		relayMsg := wCfg.GetRelay()
-		if relayMsg != nil {
-			// when we receive token we expect valid address list too
-			c := &auth.Token{
-				Payload:   relayMsg.GetTokenPayload(),
-				Signature: relayMsg.GetTokenSignature(),
-			}
-			if err := e.relayManager.UpdateToken(c); err != nil {
-				log.Errorf("failed to update relay token: %v", err)
-				return fmt.Errorf("update relay token: %w", err)
-			}
+		err = e.handleRelayUpdate(wCfg.GetRelay())
+		if err != nil {
+			return err
+		}
 
-			e.relayManager.UpdateServerURLs(relayMsg.Urls)
-
-			// Just in case the agent started with an MGM server where the relay was disabled but was later enabled.
-			// We can ignore all errors because the guard will manage the reconnection retries.
-			_ = e.relayManager.Serve()
-		} else {
-			e.relayManager.UpdateServerURLs(nil)
+		err = e.handleFlowUpdate(wCfg.GetFlow())
+		if err != nil {
+			return fmt.Errorf("handle the flow configuration: %w", err)
 		}
 
 		// todo update signal
@@ -691,6 +694,57 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 	return nil
 }
 
+func (e *Engine) handleRelayUpdate(update *mgmProto.RelayConfig) error {
+	if update != nil {
+		// when we receive token we expect valid address list too
+		c := &auth.Token{
+			Payload:   update.GetTokenPayload(),
+			Signature: update.GetTokenSignature(),
+		}
+		if err := e.relayManager.UpdateToken(c); err != nil {
+			return fmt.Errorf("update relay token: %w", err)
+		}
+
+		e.relayManager.UpdateServerURLs(update.Urls)
+
+		// Just in case the agent started with an MGM server where the relay was disabled but was later enabled.
+		// We can ignore all errors because the guard will manage the reconnection retries.
+		_ = e.relayManager.Serve()
+	} else {
+		e.relayManager.UpdateServerURLs(nil)
+	}
+
+	return nil
+}
+
+func (e *Engine) handleFlowUpdate(config *mgmProto.FlowConfig) error {
+	if config == nil {
+		return nil
+	}
+
+	flowConfig, err := toFlowLoggerConfig(config)
+	if err != nil {
+		return err
+	}
+	return e.flowManager.Update(flowConfig)
+}
+
+func toFlowLoggerConfig(config *mgmProto.FlowConfig) (*nftypes.FlowConfig, error) {
+	if config.GetInterval() == nil {
+		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(),
+		DNSCollection:      config.GetDnsCollection(),
+		ExitNodeCollection: config.GetExitNodeCollection(),
+	}, nil
+}
+
 // updateChecksIfNew updates checks if there are changes and sync new meta with management
 func (e *Engine) updateChecksIfNew(checks []*mgmProto.Checks) error {
 	// if checks are equal, we skip the update

client/internal/engine_test.go

@@ -1435,13 +1435,13 @@ func startManagement(t *testing.T, dataDir, testFile string) (*grpc.Server, stri
 	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
 	require.NoError(t, err)
 
-	accountManager, err := server.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, ia, metrics, port_forwarding.NewControllerMock())
+	accountManager, err := server.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, ia, metrics, port_forwarding.NewControllerMock(), settings.NewManagerMock())
 	if err != nil {
 		return nil, "", err
 	}
 
 	secretsManager := server.NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig, config.Relay)
-	mgmtServer, err := server.NewServer(context.Background(), config, accountManager, settings.NewManager(store), peersUpdateManager, secretsManager, nil, nil, nil)
+	mgmtServer, err := server.NewServer(context.Background(), config, accountManager, settings.NewManagerMock(), peersUpdateManager, secretsManager, nil, nil, nil)
 	if err != nil {
 		return nil, "", err
 	}

client/internal/netflow/conntrack/conntrack.go

@@ -0,0 +1,306 @@
+//go:build linux && !android
+
+package conntrack
+
+import (
+	"encoding/binary"
+	"fmt"
+	"net/netip"
+	"sync"
+
+	"github.com/google/uuid"
+	log "github.com/sirupsen/logrus"
+	nfct "github.com/ti-mo/conntrack"
+	"github.com/ti-mo/netfilter"
+
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
+)
+
+const defaultChannelSize = 100
+
+// ConnTrack manages kernel-based conntrack events
+type ConnTrack struct {
+	flowLogger nftypes.FlowLogger
+	iface      nftypes.IFaceMapper
+
+	conn *nfct.Conn
+	mux  sync.Mutex
+
+	instanceID     uuid.UUID
+	started        bool
+	done           chan struct{}
+	sysctlModified bool
+}
+
+// New creates a new connection tracker that interfaces with the kernel's conntrack system
+func New(flowLogger nftypes.FlowLogger, iface nftypes.IFaceMapper) *ConnTrack {
+	return &ConnTrack{
+		flowLogger: flowLogger,
+		iface:      iface,
+		instanceID: uuid.New(),
+		started:    false,
+		done:       make(chan struct{}, 1),
+	}
+}
+
+// Start begins tracking connections by listening for conntrack events. This method is idempotent.
+func (c *ConnTrack) Start(enableCounters bool) error {
+	c.mux.Lock()
+	defer c.mux.Unlock()
+
+	if c.started {
+		return nil
+	}
+
+	log.Info("Starting conntrack event listening")
+
+	if enableCounters {
+		c.EnableAccounting()
+	}
+
+	conn, err := nfct.Dial(nil)
+	if err != nil {
+		return fmt.Errorf("dial conntrack: %w", err)
+	}
+	c.conn = conn
+
+	events := make(chan nfct.Event, defaultChannelSize)
+	errChan, err := conn.Listen(events, 1, []netfilter.NetlinkGroup{
+		netfilter.GroupCTNew,
+		netfilter.GroupCTDestroy,
+	})
+
+	if err != nil {
+		if err := c.conn.Close(); err != nil {
+			log.Errorf("Error closing conntrack connection: %v", err)
+		}
+		c.conn = nil
+		return fmt.Errorf("start conntrack listener: %w", err)
+	}
+
+	c.started = true
+
+	go c.receiverRoutine(events, errChan)
+
+	return nil
+}
+
+func (c *ConnTrack) receiverRoutine(events chan nfct.Event, errChan chan error) {
+	for {
+		select {
+		case event := <-events:
+			c.handleEvent(event)
+		case err := <-errChan:
+			log.Errorf("Error from conntrack event listener: %v", err)
+			if err := c.conn.Close(); err != nil {
+				log.Errorf("Error closing conntrack connection: %v", err)
+			}
+			return
+		case <-c.done:
+			return
+		}
+	}
+}
+
+// Stop stops the connection tracking. This method is idempotent.
+func (c *ConnTrack) Stop() {
+	c.mux.Lock()
+	defer c.mux.Unlock()
+
+	if !c.started {
+		return
+	}
+
+	log.Info("Stopping conntrack event listening")
+
+	select {
+	case c.done <- struct{}{}:
+	default:
+	}
+
+	if c.conn != nil {
+		if err := c.conn.Close(); err != nil {
+			log.Errorf("Error closing conntrack connection: %v", err)
+		}
+		c.conn = nil
+	}
+
+	c.started = false
+
+	c.RestoreAccounting()
+}
+
+// Close stops listening for events and cleans up resources
+func (c *ConnTrack) Close() error {
+	c.mux.Lock()
+	defer c.mux.Unlock()
+
+	if c.started {
+		select {
+		case c.done <- struct{}{}:
+		default:
+		}
+	}
+
+	if c.conn != nil {
+		err := c.conn.Close()
+		c.conn = nil
+		c.started = false
+
+		c.RestoreAccounting()
+
+		if err != nil {
+			return fmt.Errorf("close conntrack: %w", err)
+		}
+	}
+
+	return nil
+}
+
+// handleEvent processes incoming conntrack events
+func (c *ConnTrack) handleEvent(event nfct.Event) {
+	if event.Flow == nil {
+		return
+	}
+
+	if event.Type != nfct.EventNew && event.Type != nfct.EventDestroy {
+		return
+	}
+
+	flow := *event.Flow
+
+	proto := nftypes.Protocol(flow.TupleOrig.Proto.Protocol)
+	if proto == nftypes.ProtocolUnknown {
+		return
+	}
+	srcIP := flow.TupleOrig.IP.SourceAddress
+	dstIP := flow.TupleOrig.IP.DestinationAddress
+
+	if !c.relevantFlow(srcIP, dstIP) {
+		return
+	}
+
+	var srcPort, dstPort uint16
+	var icmpType, icmpCode uint8
+
+	switch proto {
+	case nftypes.TCP, nftypes.UDP, nftypes.SCTP:
+		srcPort = flow.TupleOrig.Proto.SourcePort
+		dstPort = flow.TupleOrig.Proto.DestinationPort
+	case nftypes.ICMP:
+		icmpType = flow.TupleOrig.Proto.ICMPType
+		icmpCode = flow.TupleOrig.Proto.ICMPCode
+	}
+
+	flowID := c.getFlowID(flow.ID)
+	direction := c.inferDirection(srcIP, dstIP)
+
+	eventType := nftypes.TypeStart
+	eventStr := "New"
+
+	if event.Type == nfct.EventDestroy {
+		eventType = nftypes.TypeEnd
+		eventStr = "Ended"
+	}
+
+	log.Tracef("%s %s %s connection: %s:%d -> %s:%d", eventStr, direction, proto, srcIP, srcPort, dstIP, dstPort)
+
+	c.flowLogger.StoreEvent(nftypes.EventFields{
+		FlowID:     flowID,
+		Type:       eventType,
+		Direction:  direction,
+		Protocol:   proto,
+		SourceIP:   srcIP,
+		DestIP:     dstIP,
+		SourcePort: srcPort,
+		DestPort:   dstPort,
+		ICMPType:   icmpType,
+		ICMPCode:   icmpCode,
+		RxPackets:  c.mapRxPackets(flow, direction),
+		TxPackets:  c.mapTxPackets(flow, direction),
+		RxBytes:    c.mapRxBytes(flow, direction),
+		TxBytes:    c.mapTxBytes(flow, direction),
+	})
+}
+
+// relevantFlow checks if the flow is related to the specified interface
+func (c *ConnTrack) relevantFlow(srcIP, dstIP netip.Addr) bool {
+	// TODO: filter traffic by interface
+
+	wgnet := c.iface.Address().Network
+	if !wgnet.Contains(srcIP.AsSlice()) && !wgnet.Contains(dstIP.AsSlice()) {
+		return false
+	}
+
+	return true
+}
+
+// mapRxPackets maps packet counts to RX based on flow direction
+func (c *ConnTrack) mapRxPackets(flow nfct.Flow, direction nftypes.Direction) uint64 {
+	// For Ingress: CountersOrig is from external to us (RX)
+	// For Egress: CountersReply is from external to us (RX)
+	if direction == nftypes.Ingress {
+		return flow.CountersOrig.Packets
+	}
+	return flow.CountersReply.Packets
+}
+
+// mapTxPackets maps packet counts to TX based on flow direction
+func (c *ConnTrack) mapTxPackets(flow nfct.Flow, direction nftypes.Direction) uint64 {
+	// For Ingress: CountersReply is from us to external (TX)
+	// For Egress: CountersOrig is from us to external (TX)
+	if direction == nftypes.Ingress {
+		return flow.CountersReply.Packets
+	}
+	return flow.CountersOrig.Packets
+}
+
+// mapRxBytes maps byte counts to RX based on flow direction
+func (c *ConnTrack) mapRxBytes(flow nfct.Flow, direction nftypes.Direction) uint64 {
+	// For Ingress: CountersOrig is from external to us (RX)
+	// For Egress: CountersReply is from external to us (RX)
+	if direction == nftypes.Ingress {
+		return flow.CountersOrig.Bytes
+	}
+	return flow.CountersReply.Bytes
+}
+
+// mapTxBytes maps byte counts to TX based on flow direction
+func (c *ConnTrack) mapTxBytes(flow nfct.Flow, direction nftypes.Direction) uint64 {
+	// For Ingress: CountersReply is from us to external (TX)
+	// For Egress: CountersOrig is from us to external (TX)
+	if direction == nftypes.Ingress {
+		return flow.CountersReply.Bytes
+	}
+	return flow.CountersOrig.Bytes
+}
+
+// getFlowID creates a unique UUID based on the conntrack ID and instance ID
+func (c *ConnTrack) getFlowID(conntrackID uint32) uuid.UUID {
+	var buf [4]byte
+	binary.BigEndian.PutUint32(buf[:], conntrackID)
+	return uuid.NewSHA1(c.instanceID, buf[:])
+}
+
+func (c *ConnTrack) inferDirection(srcIP, dstIP netip.Addr) nftypes.Direction {
+	wgaddr := c.iface.Address().IP
+	wgnetwork := c.iface.Address().Network
+	src, dst := srcIP.AsSlice(), dstIP.AsSlice()
+
+	switch {
+	case wgaddr.Equal(src):
+		return nftypes.Egress
+	case wgaddr.Equal(dst):
+		return nftypes.Ingress
+	case wgnetwork.Contains(src):
+		// netbird network -> resource network
+		return nftypes.Ingress
+	case wgnetwork.Contains(dst):
+		// resource network -> netbird network
+		return nftypes.Egress
+
+		// TODO: handle site2site traffic
+	}
+
+	return nftypes.DirectionUnknown
+}

client/internal/netflow/conntrack/conntrack_nonlinux.go

@@ -0,0 +1,9 @@
+//go:build !linux || android
+
+package conntrack
+
+import nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
+
+func New(flowLogger nftypes.FlowLogger, iface nftypes.IFaceMapper) nftypes.ConnTracker {
+	return nil
+}

client/internal/netflow/conntrack/sysctl.go

@@ -0,0 +1,73 @@
+//go:build linux && !android
+
+package conntrack
+
+import (
+	"fmt"
+	"os"
+	"strconv"
+	"strings"
+
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	// conntrackAcctPath is the sysctl path for conntrack accounting
+	conntrackAcctPath = "net.netfilter.nf_conntrack_acct"
+)
+
+// EnableAccounting ensures that connection tracking accounting is enabled  in the kernel.
+func (c *ConnTrack) EnableAccounting() {
+	// haven't restored yet
+	if c.sysctlModified {
+		return
+	}
+
+	modified, err := setSysctl(conntrackAcctPath, 1)
+	if err != nil {
+		log.Warnf("Failed to enable conntrack accounting: %v", err)
+		return
+	}
+	c.sysctlModified = modified
+}
+
+// RestoreAccounting restores the connection tracking accounting setting to its original value.
+func (c *ConnTrack) RestoreAccounting() {
+	if !c.sysctlModified {
+		return
+	}
+
+	if _, err := setSysctl(conntrackAcctPath, 0); err != nil {
+		log.Warnf("Failed to restore conntrack accounting: %v", err)
+		return
+	}
+
+	c.sysctlModified = false
+}
+
+// setSysctl sets a sysctl configuration and returns whether it was modified.
+func setSysctl(key string, desiredValue int) (bool, error) {
+	path := fmt.Sprintf("/proc/sys/%s", strings.ReplaceAll(key, ".", "/"))
+
+	currentValue, err := os.ReadFile(path)
+	if err != nil {
+		return false, fmt.Errorf("read sysctl %s: %w", key, err)
+	}
+
+	currentV, err := strconv.Atoi(strings.TrimSpace(string(currentValue)))
+	if err != nil && len(currentValue) > 0 {
+		return false, fmt.Errorf("convert current value to int: %w", err)
+	}
+
+	if currentV == desiredValue {
+		return false, nil
+	}
+
+	// nolint:gosec
+	if err := os.WriteFile(path, []byte(strconv.Itoa(desiredValue)), 0644); err != nil {
+		return false, fmt.Errorf("write sysctl %s: %w", key, err)
+	}
+
+	log.Debugf("Set sysctl %s from %d to %d", key, currentV, desiredValue)
+	return true, nil
+}

client/internal/netflow/logger/logger.go

@@ -0,0 +1,162 @@
+package logger
+
+import (
+	"context"
+	"net"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"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"
+)
+
+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
+	dnsCollection      atomic.Bool
+	exitNodeCollection atomic.Bool
+	Store              types.Store
+}
+
+func New(ctx context.Context, statusRecorder *peer.Status, wgIfaceIPNet net.IPNet) *Logger {
+
+	ctx, cancel := context.WithCancel(ctx)
+	return &Logger{
+		ctx:            ctx,
+		cancel:         cancel,
+		statusRecorder: statusRecorder,
+		wgIfaceIPNet:   wgIfaceIPNet,
+		Store:          store.NewMemoryStore(),
+	}
+}
+
+func (l *Logger) StoreEvent(flowEvent types.EventFields) {
+	if !l.enabled.Load() {
+		return
+	}
+
+	c := l.rcvChan.Load()
+	if c == nil {
+		return
+	}
+
+	select {
+	case *c <- &flowEvent:
+	default:
+		// todo: we should collect or log on this
+	}
+}
+
+func (l *Logger) Enable() {
+	go l.startReceiver()
+}
+
+func (l *Logger) startReceiver() {
+	if l.enabled.Load() {
+		return
+	}
+
+	l.mux.Lock()
+	ctx, cancel := context.WithCancel(l.ctx)
+	l.cancelReceiver = cancel
+	l.mux.Unlock()
+
+	c := make(rcvChan, 100)
+	l.rcvChan.Store(&c)
+	l.enabled.Store(true)
+
+	for {
+		select {
+		case <-ctx.Done():
+			log.Info("flow Memory store receiver stopped")
+			return
+		case eventFields := <-c:
+			id := uuid.New()
+			event := types.Event{
+				ID:          id,
+				EventFields: *eventFields,
+				Timestamp:   time.Now(),
+			}
+
+			var isExitNode bool
+			if event.Direction == types.Ingress {
+				if !l.wgIfaceIPNet.Contains(net.IP(event.SourceIP.AsSlice())) {
+					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, isExitNode = l.statusRecorder.CheckRoutes(event.DestIP)
+				}
+			}
+
+			if l.shouldStore(eventFields, isExitNode) {
+				l.Store.StoreEvent(&event)
+			}
+		}
+	}
+}
+
+func (l *Logger) Disable() {
+	l.stop()
+	l.Store.Close()
+}
+
+func (l *Logger) stop() {
+	if !l.enabled.Load() {
+		return
+	}
+
+	l.enabled.Store(false)
+	l.mux.Lock()
+	if l.cancelReceiver != nil {
+		l.cancelReceiver()
+		l.cancelReceiver = nil
+	}
+	l.rcvChan.Store(nil)
+	l.mux.Unlock()
+}
+
+func (l *Logger) GetEvents() []*types.Event {
+	return l.Store.GetEvents()
+}
+
+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/logger/logger_test.go

@@ -0,0 +1,68 @@
+package logger_test
+
+import (
+	"context"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+
+	"github.com/netbirdio/netbird/client/internal/netflow/logger"
+	"github.com/netbirdio/netbird/client/internal/netflow/types"
+)
+
+func TestStore(t *testing.T) {
+	logger := logger.New(context.Background(), nil, net.IPNet{})
+	logger.Enable()
+
+	event := types.EventFields{
+		FlowID:    uuid.New(),
+		Type:      types.TypeStart,
+		Direction: types.Ingress,
+		Protocol:  6,
+	}
+
+	wait := func() { time.Sleep(time.Millisecond) }
+	wait()
+	logger.StoreEvent(event)
+	wait()
+
+	allEvents := logger.GetEvents()
+	matched := false
+	for _, e := range allEvents {
+		if e.EventFields.FlowID == event.FlowID {
+			matched = true
+		}
+	}
+	if !matched {
+		t.Errorf("didn't match any event")
+	}
+
+	// test disable
+	logger.Disable()
+	wait()
+	logger.StoreEvent(event)
+	wait()
+	allEvents = logger.GetEvents()
+	if len(allEvents) != 0 {
+		t.Errorf("expected 0 events, got %d", len(allEvents))
+	}
+
+	// test re-enable
+	logger.Enable()
+	wait()
+	logger.StoreEvent(event)
+	wait()
+
+	allEvents = logger.GetEvents()
+	matched = false
+	for _, e := range allEvents {
+		if e.EventFields.FlowID == event.FlowID {
+			matched = true
+		}
+	}
+	if !matched {
+		t.Errorf("didn't match any event")
+	}
+}

client/internal/netflow/manager.go

@@ -0,0 +1,252 @@
+package netflow
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"runtime"
+	"sync"
+	"time"
+
+	"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"
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
+	"github.com/netbirdio/netbird/client/internal/peer"
+	"github.com/netbirdio/netbird/flow/client"
+	"github.com/netbirdio/netbird/flow/proto"
+)
+
+// Manager handles netflow tracking and logging
+type Manager struct {
+	mux            sync.Mutex
+	logger         nftypes.FlowLogger
+	flowConfig     *nftypes.FlowConfig
+	conntrack      nftypes.ConnTracker
+	ctx            context.Context
+	receiverClient *client.GRPCClient
+	publicKey      []byte
+}
+
+// NewManager creates a new netflow manager
+func NewManager(ctx context.Context, iface nftypes.IFaceMapper, publicKey []byte, statusRecorder *peer.Status) *Manager {
+	var ipNet net.IPNet
+	if iface != nil {
+		ipNet = *iface.Address().Network
+	}
+	flowLogger := logger.New(ctx, statusRecorder, ipNet)
+
+	var ct nftypes.ConnTracker
+	if runtime.GOOS == "linux" && iface != nil && !iface.IsUserspaceBind() {
+		ct = conntrack.New(flowLogger, iface)
+	}
+
+	return &Manager{
+		logger:    flowLogger,
+		conntrack: ct,
+		ctx:       ctx,
+		publicKey: publicKey,
+	}
+}
+
+// Update applies new flow configuration settings
+// needsNewClient checks if a new client needs to be created
+func (m *Manager) needsNewClient(previous *nftypes.FlowConfig) bool {
+	current := m.flowConfig
+	return previous == nil ||
+		!previous.Enabled ||
+		previous.TokenPayload != current.TokenPayload ||
+		previous.TokenSignature != current.TokenSignature ||
+		previous.URL != current.URL
+}
+
+// enableFlow starts components for flow tracking
+func (m *Manager) enableFlow(previous *nftypes.FlowConfig) error {
+	// first make sender ready so events don't pile up
+	if m.needsNewClient(previous) {
+		if m.receiverClient != nil {
+			if err := m.receiverClient.Close(); err != nil {
+				log.Warnf("error closing previous flow client: %s", err)
+			}
+		}
+
+		flowClient, err := client.NewClient(m.flowConfig.URL, m.flowConfig.TokenPayload, m.flowConfig.TokenSignature, m.flowConfig.Interval)
+		if err != nil {
+			return fmt.Errorf("create client: %w", err)
+		}
+		log.Infof("flow client configured to connect to %s", m.flowConfig.URL)
+
+		m.receiverClient = flowClient
+		go m.receiveACKs(flowClient)
+		go m.startSender()
+	}
+
+	m.logger.Enable()
+
+	if m.conntrack != nil {
+		if err := m.conntrack.Start(m.flowConfig.Counters); err != nil {
+			return fmt.Errorf("start conntrack: %w", err)
+		}
+	}
+
+	return nil
+}
+
+// disableFlow stops components for flow tracking
+func (m *Manager) disableFlow() error {
+	if m.conntrack != nil {
+		m.conntrack.Stop()
+	}
+
+	m.logger.Disable()
+
+	if m.receiverClient != nil {
+		return m.receiverClient.Close()
+	}
+	return nil
+}
+
+// Update applies new flow configuration settings
+func (m *Manager) Update(update *nftypes.FlowConfig) error {
+	if update == nil {
+		return nil
+	}
+
+	m.mux.Lock()
+	defer m.mux.Unlock()
+
+	previous := m.flowConfig
+	m.flowConfig = update
+
+	if previous != nil && previous.TokenPayload != "" && m.flowConfig != nil && m.flowConfig.TokenPayload == "" {
+		m.flowConfig.TokenPayload = previous.TokenPayload
+		m.flowConfig.TokenSignature = previous.TokenSignature
+	}
+
+	m.logger.UpdateConfig(update.DNSCollection, update.ExitNodeCollection)
+
+	if update.Enabled {
+		return m.enableFlow(previous)
+	}
+
+	return m.disableFlow()
+}
+
+// Close cleans up all resources
+func (m *Manager) Close() {
+	m.mux.Lock()
+	defer m.mux.Unlock()
+
+	if m.conntrack != nil {
+		m.conntrack.Close()
+	}
+
+	if m.receiverClient != nil {
+		if err := m.receiverClient.Close(); err != nil {
+			log.Warnf("failed to close receiver client: %s", err)
+		}
+	}
+
+	m.logger.Close()
+}
+
+// GetLogger returns the flow logger
+func (m *Manager) GetLogger() nftypes.FlowLogger {
+	return m.logger
+}
+
+func (m *Manager) startSender() {
+	ticker := time.NewTicker(m.flowConfig.Interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-m.ctx.Done():
+			return
+		case <-ticker.C:
+			events := m.logger.GetEvents()
+			for _, event := range events {
+				if err := m.send(event); err != nil {
+					log.Errorf("failed to send flow event to server: %s", err)
+					continue
+				}
+				log.Tracef("sent flow event: %s", event.ID)
+			}
+		}
+	}
+}
+
+func (m *Manager) receiveACKs(client *client.GRPCClient) {
+	err := client.Receive(m.ctx, m.flowConfig.Interval, func(ack *proto.FlowEventAck) error {
+		id, err := uuid.FromBytes(ack.EventId)
+		if err != nil {
+			log.Warnf("failed to convert ack event id to uuid: %s", err)
+			return nil
+		}
+		log.Tracef("received flow event ack: %s", id)
+		m.logger.DeleteEvents([]uuid.UUID{uuid.UUID(ack.EventId)})
+		return nil
+	})
+
+	if err != nil && !errors.Is(err, context.Canceled) {
+		log.Errorf("failed to receive flow event ack: %s", err)
+	}
+}
+
+func (m *Manager) send(event *nftypes.Event) error {
+	m.mux.Lock()
+	client := m.receiverClient
+	m.mux.Unlock()
+
+	if client == nil {
+		return nil
+	}
+
+	return client.Send(toProtoEvent(m.publicKey, event))
+}
+
+func toProtoEvent(publicKey []byte, event *nftypes.Event) *proto.FlowEvent {
+	protoEvent := &proto.FlowEvent{
+		EventId:   event.ID[:],
+		Timestamp: timestamppb.New(event.Timestamp),
+		PublicKey: publicKey,
+		FlowFields: &proto.FlowFields{
+			FlowId:           event.FlowID[:],
+			RuleId:           event.RuleID,
+			Type:             proto.Type(event.Type),
+			Direction:        proto.Direction(event.Direction),
+			Protocol:         uint32(event.Protocol),
+			SourceIp:         event.SourceIP.AsSlice(),
+			DestIp:           event.DestIP.AsSlice(),
+			RxPackets:        event.RxPackets,
+			TxPackets:        event.TxPackets,
+			RxBytes:          event.RxBytes,
+			TxBytes:          event.TxBytes,
+			SourceResourceId: event.SourceResourceID,
+			DestResourceId:   event.DestResourceID,
+		},
+	}
+
+	if event.Protocol == nftypes.ICMP {
+		protoEvent.FlowFields.ConnectionInfo = &proto.FlowFields_IcmpInfo{
+			IcmpInfo: &proto.ICMPInfo{
+				IcmpType: uint32(event.ICMPType),
+				IcmpCode: uint32(event.ICMPCode),
+			},
+		}
+		return protoEvent
+	}
+
+	protoEvent.FlowFields.ConnectionInfo = &proto.FlowFields_PortInfo{
+		PortInfo: &proto.PortInfo{
+			SourcePort: uint32(event.SourcePort),
+			DestPort:   uint32(event.DestPort),
+		},
+	}
+
+	return protoEvent
+}

client/internal/netflow/store/memory.go

@@ -0,0 +1,52 @@
+package store
+
+import (
+	"sync"
+
+	"golang.org/x/exp/maps"
+
+	"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
+}
+
+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()
+	maps.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)
+	}
+}

client/internal/netflow/types/types.go

@@ -0,0 +1,156 @@
+package types
+
+import (
+	"net/netip"
+	"strconv"
+	"time"
+
+	"github.com/google/uuid"
+
+	"github.com/netbirdio/netbird/client/iface/wgaddr"
+)
+
+type Protocol uint8
+
+const (
+	ProtocolUnknown = Protocol(0)
+	ICMP            = Protocol(1)
+	TCP             = Protocol(6)
+	UDP             = Protocol(17)
+	SCTP            = Protocol(132)
+)
+
+func (p Protocol) String() string {
+	switch p {
+	case 1:
+		return "ICMP"
+	case 6:
+		return "TCP"
+	case 17:
+		return "UDP"
+	case 132:
+		return "SCTP"
+	default:
+		return strconv.FormatUint(uint64(p), 10)
+	}
+}
+
+type Type int
+
+const (
+	TypeUnknown = Type(iota)
+	TypeStart
+	TypeEnd
+	TypeDrop
+)
+
+type Direction int
+
+func (d Direction) String() string {
+	switch d {
+	case Ingress:
+		return "ingress"
+	case Egress:
+		return "egress"
+	default:
+		return "unknown"
+	}
+}
+
+const (
+	DirectionUnknown = Direction(iota)
+	Ingress
+	Egress
+)
+
+type Event struct {
+	ID        uuid.UUID
+	Timestamp time.Time
+	EventFields
+}
+
+type EventFields struct {
+	FlowID           uuid.UUID
+	Type             Type
+	RuleID           []byte
+	Direction        Direction
+	Protocol         Protocol
+	SourceIP         netip.Addr
+	DestIP           netip.Addr
+	SourceResourceID []byte
+	DestResourceID   []byte
+	SourcePort       uint16
+	DestPort         uint16
+	ICMPType         uint8
+	ICMPCode         uint8
+	RxPackets        uint64
+	TxPackets        uint64
+	RxBytes          uint64
+	TxBytes          uint64
+}
+
+type FlowConfig struct {
+	URL                string
+	Interval           time.Duration
+	Enabled            bool
+	Counters           bool
+	TokenPayload       string
+	TokenSignature     string
+	DNSCollection      bool
+	ExitNodeCollection bool
+}
+
+type FlowManager interface {
+	// FlowConfig handles network map updates
+	Update(update *FlowConfig) error
+	// Close closes the manager
+	Close()
+	// GetLogger returns a flow logger
+	GetLogger() FlowLogger
+}
+
+type FlowLogger interface {
+	// 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
+	Enable()
+	// Disable disables the flow logger receiver
+	Disable()
+
+	// UpdateConfig updates the flow manager configuration
+	UpdateConfig(dnsCollection, exitNodeCollection bool)
+}
+
+type Store interface {
+	// StoreEvent stores a flow event
+	StoreEvent(event *Event)
+	// GetEvents returns all stored events
+	GetEvents() []*Event
+	// DeleteEvents deletes events from the store
+	DeleteEvents([]uuid.UUID)
+	// Close closes the store
+	Close()
+}
+
+// ConnTracker defines the interface for connection tracking functionality
+type ConnTracker interface {
+	// Start begins tracking connections by listening for conntrack events.
+	Start(bool) error
+	// Stop stops the connection tracking.
+	Stop()
+	// Close stops listening for events and cleans up resources
+	Close() error
+}
+
+// IFaceMapper provides interface to check if we're using userspace WireGuard
+type IFaceMapper interface {
+	IsUserspaceBind() bool
+	Name() string
+	Address() wgaddr.Address
+}

client/internal/peer/route.go

@@ -0,0 +1,81 @@
+package peer
+
+import (
+	"net/netip"
+	"sync"
+
+	log "github.com/sirupsen/logrus"
+)
+
+type routeIDLookup struct {
+	localMap    sync.Map
+	remoteMap   sync.Map
+	resolvedIPs sync.Map
+}
+
+func (r *routeIDLookup) AddLocalRouteID(resourceID string, route netip.Prefix) {
+	_, exists := r.localMap.LoadOrStore(route, resourceID)
+	if exists {
+		log.Tracef("resourceID %s already exists in local map", resourceID)
+	}
+}
+
+func (r *routeIDLookup) RemoveLocalRouteID(route netip.Prefix) {
+	r.localMap.Delete(route)
+}
+
+func (r *routeIDLookup) AddRemoteRouteID(resourceID string, route netip.Prefix) {
+	_, exists := r.remoteMap.LoadOrStore(route, resourceID)
+	if exists {
+		log.Tracef("resourceID %s already exists in remote map", resourceID)
+	}
+}
+
+func (r *routeIDLookup) RemoveRemoteRouteID(route netip.Prefix) {
+	r.remoteMap.Delete(route)
+}
+
+func (r *routeIDLookup) AddResolvedIP(resourceID string, route netip.Prefix) {
+	r.resolvedIPs.Store(route.Addr(), resourceID)
+}
+
+func (r *routeIDLookup) RemoveResolvedIP(route netip.Prefix) {
+	r.resolvedIPs.Delete(route.Addr())
+}
+
+// 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), false
+	}
+
+	var resourceID string
+	r.localMap.Range(func(key, value interface{}) bool {
+		pref := key.(netip.Prefix)
+		if pref.Contains(ip) {
+			resourceID = value.(string)
+			isExitNode = pref.Bits() == 0
+			return false
+
+		}
+		return true
+	})
+
+	if resourceID == "" {
+		r.remoteMap.Range(func(key, value interface{}) bool {
+			pref := key.(netip.Prefix)
+			if pref.Contains(ip) {
+				resourceID = value.(string)
+				isExitNode = pref.Bits() == 0
+				return false
+			}
+			return true
+		})
+	}
+
+	return resourceID, isExitNode
+}

client/internal/peer/status.go

@@ -176,6 +176,8 @@ type Status struct {
 	eventQueue   *EventQueue
 
 	ingressGwMgr *ingressgw.Manager
+
+	routeIDLookup routeIDLookup
 }
 
 // NewRecorder returns a new Status instance
@@ -311,7 +313,7 @@ func (d *Status) UpdatePeerState(receivedState State) error {
 	return nil
 }
 
-func (d *Status) AddPeerStateRoute(peer string, route string) error {
+func (d *Status) AddPeerStateRoute(peer string, route string, resourceId string) error {
 	d.mux.Lock()
 	defer d.mux.Unlock()
 
@@ -323,6 +325,14 @@ func (d *Status) AddPeerStateRoute(peer string, route string) error {
 	peerState.AddRoute(route)
 	d.peers[peer] = peerState
 
+	pref, err := netip.ParsePrefix(route)
+	if err != nil {
+		log.Errorf("failed to parse prefix %s: %v", route, err)
+	} else {
+
+		d.routeIDLookup.AddRemoteRouteID(resourceId, pref)
+	}
+
 	// todo: consider to make sense of this notification or not
 	d.notifyPeerListChanged()
 	return nil
@@ -340,11 +350,28 @@ func (d *Status) RemovePeerStateRoute(peer string, route string) error {
 	peerState.DeleteRoute(route)
 	d.peers[peer] = peerState
 
+	pref, err := netip.ParsePrefix(route)
+	if err != nil {
+		log.Errorf("failed to parse prefix %s: %v", route, err)
+	} else {
+		d.routeIDLookup.RemoveRemoteRouteID(pref)
+	}
+
 	// todo: consider to make sense of this notification or not
 	d.notifyPeerListChanged()
 	return nil
 }
 
+// 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) ([]byte, bool) {
+	if d == nil {
+		return nil, false
+	}
+	resId, isExitNode := d.routeIDLookup.Lookup(ip)
+	return []byte(resId), isExitNode
+}
+
 func (d *Status) UpdatePeerICEState(receivedState State) error {
 	d.mux.Lock()
 	defer d.mux.Unlock()
@@ -558,6 +585,50 @@ func (d *Status) UpdateLocalPeerState(localPeerState LocalPeerState) {
 	d.notifyAddressChanged()
 }
 
+// AddLocalPeerStateRoute adds a route to the local peer state
+func (d *Status) AddLocalPeerStateRoute(route, resourceId string) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	pref, err := netip.ParsePrefix(route)
+	if err != nil {
+		log.Errorf("failed to parse prefix %s: %v", route, err)
+		return
+	}
+
+	if d.localPeer.Routes == nil {
+		d.localPeer.Routes = map[string]struct{}{}
+	}
+
+	d.localPeer.Routes[route] = struct{}{}
+
+	d.routeIDLookup.AddLocalRouteID(resourceId, pref)
+}
+
+// RemoveLocalPeerStateRoute removes a route from the local peer state
+func (d *Status) RemoveLocalPeerStateRoute(route string) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	pref, err := netip.ParsePrefix(route)
+	if err != nil {
+		log.Errorf("failed to parse prefix %s: %v", route, err)
+		return
+	}
+
+	delete(d.localPeer.Routes, route)
+
+	d.routeIDLookup.RemoveLocalRouteID(pref)
+}
+
+// CleanLocalPeerStateRoutes cleans all routes from the local peer state
+func (d *Status) CleanLocalPeerStateRoutes() {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	d.localPeer.Routes = map[string]struct{}{}
+}
+
 // CleanLocalPeerState cleans local peer status
 func (d *Status) CleanLocalPeerState() {
 	d.mux.Lock()
@@ -641,7 +712,7 @@ func (d *Status) UpdateDNSStates(dnsStates []NSGroupState) {
 	d.nsGroupStates = dnsStates
 }
 
-func (d *Status) UpdateResolvedDomainsStates(originalDomain domain.Domain, resolvedDomain domain.Domain, prefixes []netip.Prefix) {
+func (d *Status) UpdateResolvedDomainsStates(originalDomain domain.Domain, resolvedDomain domain.Domain, prefixes []netip.Prefix, resourceId string) {
 	d.mux.Lock()
 	defer d.mux.Unlock()
 
@@ -650,6 +721,10 @@ func (d *Status) UpdateResolvedDomainsStates(originalDomain domain.Domain, resol
 		Prefixes:     prefixes,
 		ParentDomain: originalDomain,
 	}
+
+	for _, prefix := range prefixes {
+		d.routeIDLookup.AddResolvedIP(resourceId, prefix)
+	}
 }
 
 func (d *Status) DeleteResolvedDomainsStates(domain domain.Domain) {
@@ -660,6 +735,10 @@ func (d *Status) DeleteResolvedDomainsStates(domain domain.Domain) {
 	for k, v := range d.resolvedDomainsStates {
 		if v.ParentDomain == domain {
 			delete(d.resolvedDomainsStates, k)
+
+			for _, prefix := range v.Prefixes {
+				d.routeIDLookup.RemoveResolvedIP(prefix)
+			}
 		}
 	}
 }

client/internal/pkce_auth.go

@@ -37,7 +37,7 @@ type PKCEAuthProviderConfig struct {
 	RedirectURLs []string
 	// UseIDToken indicates if the id token should be used for authentication
 	UseIDToken bool
-	//ClientCertPair is used for mTLS authentication to the IDP
+	// ClientCertPair is used for mTLS authentication to the IDP
 	ClientCertPair *tls.Certificate
 }
 

client/internal/routemanager/client.go

@@ -330,7 +330,7 @@ func (c *clientNetwork) recalculateRouteAndUpdatePeerAndSystem(rsn reason) error
 		c.connectEvent()
 	}
 
-	err := c.statusRecorder.AddPeerStateRoute(c.currentChosen.Peer, c.handler.String())
+	err := c.statusRecorder.AddPeerStateRoute(c.currentChosen.Peer, c.handler.String(), c.currentChosen.GetResourceID())
 	if err != nil {
 		return fmt.Errorf("add peer state route: %w", err)
 	}

client/internal/routemanager/dnsinterceptor/handler.go

@@ -321,7 +321,7 @@ func (d *DnsInterceptor) updateDomainPrefixes(resolvedDomain, originalDomain dom
 	if len(toAdd) > 0 || len(toRemove) > 0 {
 		d.interceptedDomains[resolvedDomain] = newPrefixes
 		originalDomain = domain.Domain(strings.TrimSuffix(string(originalDomain), "."))
-		d.statusRecorder.UpdateResolvedDomainsStates(originalDomain, resolvedDomain, newPrefixes)
+		d.statusRecorder.UpdateResolvedDomainsStates(originalDomain, resolvedDomain, newPrefixes, d.route.GetResourceID())
 
 		if len(toAdd) > 0 {
 			log.Debugf("added dynamic route(s) for domain=%s (pattern: domain=%s): %s",

client/internal/routemanager/dynamic/route.go

@@ -288,7 +288,7 @@ func (r *Route) updateDynamicRoutes(ctx context.Context, newDomains domainMap) e
 		updatedPrefixes := combinePrefixes(oldPrefixes, removedPrefixes, addedPrefixes)
 		r.dynamicDomains[domain] = updatedPrefixes
 
-		r.statusRecorder.UpdateResolvedDomainsStates(domain, domain, updatedPrefixes)
+		r.statusRecorder.UpdateResolvedDomainsStates(domain, domain, updatedPrefixes, r.route.GetResourceID())
 	}
 
 	return nberrors.FormatErrorOrNil(merr)

client/internal/routemanager/manager.go

@@ -45,7 +45,7 @@ type Manager interface {
 	SetRouteChangeListener(listener listener.NetworkChangeListener)
 	InitialRouteRange() []string
 	EnableServerRouter(firewall firewall.Manager) error
-	Stop(stateManager *statemanager.Manager)
+	Stop(stateManager statemanager.Manager)
 }
 
 type ManagerConfig struct {
@@ -56,7 +56,7 @@ type ManagerConfig struct {
 	StatusRecorder      *peer.Status
 	RelayManager        *relayClient.Manager
 	InitialRoutes       []*route.Route
-	StateManager        *statemanager.Manager
+	StateManager        statemanager.Manager
 	DNSServer           dns.Server
 	PeerStore           *peerstore.Store
 	DisableClientRoutes bool
@@ -80,7 +80,7 @@ type DefaultManager struct {
 	routeRefCounter      *refcounter.RouteRefCounter
 	allowedIPsRefCounter *refcounter.AllowedIPsRefCounter
 	dnsRouteInterval     time.Duration
-	stateManager         *statemanager.Manager
+	stateManager         statemanager.Manager
 	// clientRoutes is the most recent list of clientRoutes received from the Management Service
 	clientRoutes        route.HAMap
 	dnsServer           dns.Server
@@ -234,7 +234,7 @@ func (m *DefaultManager) EnableServerRouter(firewall firewall.Manager) error {
 }
 
 // Stop stops the manager watchers and clean firewall rules
-func (m *DefaultManager) Stop(stateManager *statemanager.Manager) {
+func (m *DefaultManager) Stop(stateManager statemanager.Manager) {
 	m.stop()
 	if m.serverRouter != nil {
 		m.serverRouter.cleanUp()

client/internal/routemanager/mock.go

@@ -19,7 +19,7 @@ type MockManager struct {
 	GetRouteSelectorFunc         func() *routeselector.RouteSelector
 	GetClientRoutesFunc          func() route.HAMap
 	GetClientRoutesWithNetIDFunc func() map[route.NetID][]*route.Route
-	StopFunc                     func(manager *statemanager.Manager)
+	StopFunc                     func(manager statemanager.Manager)
 }
 
 func (m *MockManager) Init() (net.AddHookFunc, net.RemoveHookFunc, error) {
@@ -83,7 +83,7 @@ func (m *MockManager) EnableServerRouter(firewall firewall.Manager) error {
 }
 
 // Stop mock implementation of Stop from Manager interface
-func (m *MockManager) Stop(stateManager *statemanager.Manager) {
+func (m *MockManager) Stop(stateManager statemanager.Manager) {
 	if m.StopFunc != nil {
 		m.StopFunc(stateManager)
 	}

client/internal/routemanager/server_nonandroid.go

@@ -103,9 +103,7 @@ func (m *serverRouter) removeFromServerNetwork(route *route.Route) error {
 
 	delete(m.routes, route.ID)
 
-	state := m.statusRecorder.GetLocalPeerState()
-	delete(state.Routes, route.Network.String())
-	m.statusRecorder.UpdateLocalPeerState(state)
+	m.statusRecorder.RemoveLocalPeerStateRoute(route.Network.String())
 
 	return nil
 }
@@ -131,18 +129,12 @@ func (m *serverRouter) addToServerNetwork(route *route.Route) error {
 
 	m.routes[route.ID] = route
 
-	state := m.statusRecorder.GetLocalPeerState()
-	if state.Routes == nil {
-		state.Routes = map[string]struct{}{}
-	}
-
 	routeStr := route.Network.String()
 	if route.IsDynamic() {
 		routeStr = route.Domains.SafeString()
 	}
-	state.Routes[routeStr] = struct{}{}
 
-	m.statusRecorder.UpdateLocalPeerState(state)
+	m.statusRecorder.AddLocalPeerStateRoute(routeStr, route.GetResourceID())
 
 	return nil
 }
@@ -164,9 +156,7 @@ func (m *serverRouter) cleanUp() {
 
 	}
 
-	state := m.statusRecorder.GetLocalPeerState()
-	state.Routes = nil
-	m.statusRecorder.UpdateLocalPeerState(state)
+	m.statusRecorder.CleanLocalPeerStateRoutes()
 }
 
 func routeToRouterPair(route *route.Route) (firewall.RouterPair, error) {

client/internal/routemanager/systemops/systemops_android.go

@@ -13,11 +13,11 @@ import (
 	nbnet "github.com/netbirdio/netbird/util/net"
 )
 
-func (r *SysOps) SetupRouting([]net.IP, *statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
+func (r *SysOps) SetupRouting([]net.IP, statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
 	return nil, nil, nil
 }
 
-func (r *SysOps) CleanupRouting(*statemanager.Manager) error {
+func (r *SysOps) CleanupRouting(statemanager.Manager) error {
 	return nil
 }
 

client/internal/routemanager/systemops/systemops_generic.go

@@ -32,7 +32,7 @@ var splitDefaultv6_2 = netip.PrefixFrom(netip.AddrFrom16([16]byte{0x80}), 1)
 
 var ErrRoutingIsSeparate = errors.New("routing is separate")
 
-func (r *SysOps) setupRefCounter(initAddresses []net.IP, stateManager *statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
+func (r *SysOps) setupRefCounter(initAddresses []net.IP, stateManager statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
 	stateManager.RegisterState(&ShutdownState{})
 
 	initialNextHopV4, err := GetNextHop(netip.IPv4Unspecified())
@@ -80,13 +80,13 @@ func (r *SysOps) setupRefCounter(initAddresses []net.IP, stateManager *statemana
 }
 
 // updateState updates state on every change so it will be persisted regularly
-func (r *SysOps) updateState(stateManager *statemanager.Manager) {
+func (r *SysOps) updateState(stateManager statemanager.Manager) {
 	if err := stateManager.UpdateState((*ShutdownState)(r.refCounter)); err != nil {
 		log.Errorf("failed to update state: %v", err)
 	}
 }
 
-func (r *SysOps) cleanupRefCounter(stateManager *statemanager.Manager) error {
+func (r *SysOps) cleanupRefCounter(stateManager statemanager.Manager) error {
 	if r.refCounter == nil {
 		return nil
 	}
@@ -337,7 +337,7 @@ func (r *SysOps) genericRemoveVPNRoute(prefix netip.Prefix, intf *net.Interface)
 	return r.removeFromRouteTable(prefix, nextHop)
 }
 
-func (r *SysOps) setupHooks(initAddresses []net.IP, stateManager *statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
+func (r *SysOps) setupHooks(initAddresses []net.IP, stateManager statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
 	beforeHook := func(connID nbnet.ConnectionID, ip net.IP) error {
 		prefix, err := util.GetPrefixFromIP(ip)
 		if err != nil {

client/internal/routemanager/systemops/systemops_ios.go

@@ -13,14 +13,14 @@ import (
 	nbnet "github.com/netbirdio/netbird/util/net"
 )
 
-func (r *SysOps) SetupRouting([]net.IP, *statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
+func (r *SysOps) SetupRouting([]net.IP, statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	r.prefixes = make(map[netip.Prefix]struct{})
 	return nil, nil, nil
 }
 
-func (r *SysOps) CleanupRouting(*statemanager.Manager) error {
+func (r *SysOps) CleanupRouting(statemanager.Manager) error {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 

client/internal/routemanager/systemops/systemops_linux.go

@@ -72,7 +72,7 @@ func getSetupRules() []ruleParams {
 // Rule 2 (VPN Traffic Routing): Directs all remaining traffic to the 'NetbirdVPNTableID' custom routing table.
 // This table is where a default route or other specific routes received from the management server are configured,
 // enabling VPN connectivity.
-func (r *SysOps) SetupRouting(initAddresses []net.IP, stateManager *statemanager.Manager) (_ nbnet.AddHookFunc, _ nbnet.RemoveHookFunc, err error) {
+func (r *SysOps) SetupRouting(initAddresses []net.IP, stateManager statemanager.Manager) (_ nbnet.AddHookFunc, _ nbnet.RemoveHookFunc, err error) {
 	if !nbnet.AdvancedRouting() {
 		log.Infof("Using legacy routing setup")
 		return r.setupRefCounter(initAddresses, stateManager)
@@ -110,7 +110,7 @@ func (r *SysOps) SetupRouting(initAddresses []net.IP, stateManager *statemanager
 // CleanupRouting performs a thorough cleanup of the routing configuration established by 'setupRouting'.
 // It systematically removes the three rules and any associated routing table entries to ensure a clean state.
 // The function uses error aggregation to report any errors encountered during the cleanup process.
-func (r *SysOps) CleanupRouting(stateManager *statemanager.Manager) error {
+func (r *SysOps) CleanupRouting(stateManager statemanager.Manager) error {
 	if !nbnet.AdvancedRouting() {
 		return r.cleanupRefCounter(stateManager)
 	}

client/internal/routemanager/systemops/systemops_unix.go

@@ -17,11 +17,11 @@ import (
 	nbnet "github.com/netbirdio/netbird/util/net"
 )
 
-func (r *SysOps) SetupRouting(initAddresses []net.IP, stateManager *statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
+func (r *SysOps) SetupRouting(initAddresses []net.IP, stateManager statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
 	return r.setupRefCounter(initAddresses, stateManager)
 }
 
-func (r *SysOps) CleanupRouting(stateManager *statemanager.Manager) error {
+func (r *SysOps) CleanupRouting(stateManager statemanager.Manager) error {
 	return r.cleanupRefCounter(stateManager)
 }
 

client/internal/routemanager/systemops/systemops_windows.go

@@ -131,11 +131,11 @@ const (
 	RouteDeleted
 )
 
-func (r *SysOps) SetupRouting(initAddresses []net.IP, stateManager *statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
+func (r *SysOps) SetupRouting(initAddresses []net.IP, stateManager statemanager.Manager) (nbnet.AddHookFunc, nbnet.RemoveHookFunc, error) {
 	return r.setupRefCounter(initAddresses, stateManager)
 }
 
-func (r *SysOps) CleanupRouting(stateManager *statemanager.Manager) error {
+func (r *SysOps) CleanupRouting(stateManager statemanager.Manager) error {
 	return r.cleanupRefCounter(stateManager)
 }
 

client/internal/statemanager/manager.go

@@ -49,8 +49,25 @@ func (r *RawState) MarshalJSON() ([]byte, error) {
 	return r.data, nil
 }
 
+// Manager is the interface that exposes the persistence and state management methods.
+type Manager interface {
+	Start()
+	Stop(ctx context.Context) error
+	RegisterState(state State)
+	GetState(state State) State
+	UpdateState(state State) error
+	DeleteState(state State) error
+	DeleteStateByName(stateName string) error
+	DeleteAllStates() (int, error)
+	PersistState(ctx context.Context) error
+	LoadState(state State) error
+	CleanupStateByName(name string) error
+	PerformCleanup() error
+	GetSavedStateNames() ([]string, error)
+}
+
 // Manager handles the persistence and management of various states
-type Manager struct {
+type managerImpl struct {
 	mu     sync.Mutex
 	cancel context.CancelFunc
 	done   chan struct{}
@@ -65,8 +82,8 @@ type Manager struct {
 }
 
 // New creates a new Manager instance
-func New(filePath string) *Manager {
-	return &Manager{
+func New(filePath string) Manager {
+	return &managerImpl{
 		filePath:   filePath,
 		states:     make(map[string]State),
 		dirty:      make(map[string]struct{}),
@@ -75,7 +92,7 @@ func New(filePath string) *Manager {
 }
 
 // Start starts the state manager periodic save routine
-func (m *Manager) Start() {
+func (m *managerImpl) Start() {
 	if m == nil {
 		return
 	}
@@ -90,7 +107,7 @@ func (m *Manager) Start() {
 	go m.periodicStateSave(ctx)
 }
 
-func (m *Manager) Stop(ctx context.Context) error {
+func (m *managerImpl) Stop(ctx context.Context) error {
 	if m == nil {
 		return nil
 	}
@@ -114,7 +131,7 @@ func (m *Manager) Stop(ctx context.Context) error {
 
 // RegisterState registers a state with the manager but doesn't attempt to persist it.
 // Pass an uninitialized state to register it.
-func (m *Manager) RegisterState(state State) {
+func (m *managerImpl) RegisterState(state State) {
 	if m == nil {
 		return
 	}
@@ -128,7 +145,7 @@ func (m *Manager) RegisterState(state State) {
 }
 
 // GetState returns the state for the given type
-func (m *Manager) GetState(state State) State {
+func (m *managerImpl) GetState(state State) State {
 	if m == nil {
 		return nil
 	}
@@ -141,7 +158,7 @@ func (m *Manager) GetState(state State) State {
 
 // UpdateState updates the state in the manager and marks it as dirty for the next save.
 // The state will be replaced with the new one.
-func (m *Manager) UpdateState(state State) error {
+func (m *managerImpl) UpdateState(state State) error {
 	if m == nil {
 		return nil
 	}
@@ -151,7 +168,7 @@ func (m *Manager) UpdateState(state State) error {
 
 // DeleteState removes the state from the manager and marks it as dirty for the next save.
 // Pass an uninitialized state to delete it.
-func (m *Manager) DeleteState(state State) error {
+func (m *managerImpl) DeleteState(state State) error {
 	if m == nil {
 		return nil
 	}
@@ -159,7 +176,7 @@ func (m *Manager) DeleteState(state State) error {
 	return m.setState(state.Name(), nil)
 }
 
-func (m *Manager) setState(name string, state State) error {
+func (m *managerImpl) setState(name string, state State) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
@@ -175,7 +192,7 @@ func (m *Manager) setState(name string, state State) error {
 
 // DeleteStateByName handles deletion of states without cleanup.
 // It doesn't require the state to be registered.
-func (m *Manager) DeleteStateByName(stateName string) error {
+func (m *managerImpl) DeleteStateByName(stateName string) error {
 	if m == nil {
 		return nil
 	}
@@ -203,7 +220,7 @@ func (m *Manager) DeleteStateByName(stateName string) error {
 }
 
 // DeleteAllStates removes all states.
-func (m *Manager) DeleteAllStates() (int, error) {
+func (m *managerImpl) DeleteAllStates() (int, error) {
 	if m == nil {
 		return 0, nil
 	}
@@ -230,7 +247,7 @@ func (m *Manager) DeleteAllStates() (int, error) {
 	return count, nil
 }
 
-func (m *Manager) periodicStateSave(ctx context.Context) {
+func (m *managerImpl) periodicStateSave(ctx context.Context) {
 	ticker := time.NewTicker(10 * time.Second)
 	defer ticker.Stop()
 	defer close(m.done)
@@ -248,7 +265,7 @@ func (m *Manager) periodicStateSave(ctx context.Context) {
 }
 
 // PersistState persists the states that have been updated since the last save.
-func (m *Manager) PersistState(ctx context.Context) error {
+func (m *managerImpl) PersistState(ctx context.Context) error {
 	if m == nil {
 		return nil
 	}
@@ -291,7 +308,7 @@ func (m *Manager) PersistState(ctx context.Context) error {
 }
 
 // loadStateFile reads and unmarshals the state file into a map of raw JSON messages
-func (m *Manager) loadStateFile(deleteCorrupt bool) (map[string]json.RawMessage, error) {
+func (m *managerImpl) loadStateFile(deleteCorrupt bool) (map[string]json.RawMessage, error) {
 	data, err := os.ReadFile(m.filePath)
 	if err != nil {
 		if errors.Is(err, fs.ErrNotExist) {
@@ -311,7 +328,7 @@ func (m *Manager) loadStateFile(deleteCorrupt bool) (map[string]json.RawMessage,
 }
 
 // handleCorruptedState creates a backup of a corrupted state file by moving it
-func (m *Manager) handleCorruptedState(deleteCorrupt bool) {
+func (m *managerImpl) handleCorruptedState(deleteCorrupt bool) {
 	if !deleteCorrupt {
 		return
 	}
@@ -327,7 +344,7 @@ func (m *Manager) handleCorruptedState(deleteCorrupt bool) {
 }
 
 // loadSingleRawState unmarshals a raw state into a concrete state object
-func (m *Manager) loadSingleRawState(name string, rawState json.RawMessage) (State, error) {
+func (m *managerImpl) loadSingleRawState(name string, rawState json.RawMessage) (State, error) {
 	stateType, ok := m.stateTypes[name]
 	if !ok {
 		return nil, fmt.Errorf(errStateNotRegistered, name)
@@ -346,7 +363,7 @@ func (m *Manager) loadSingleRawState(name string, rawState json.RawMessage) (Sta
 }
 
 // LoadState loads a specific state from the state file
-func (m *Manager) LoadState(state State) error {
+func (m *managerImpl) LoadState(state State) error {
 	if m == nil {
 		return nil
 	}
@@ -383,7 +400,7 @@ func (m *Manager) LoadState(state State) error {
 
 // cleanupSingleState handles the cleanup of a specific state and returns any error.
 // The caller must hold the mutex.
-func (m *Manager) cleanupSingleState(name string, rawState json.RawMessage) error {
+func (m *managerImpl) cleanupSingleState(name string, rawState json.RawMessage) error {
 	// For unregistered states, preserve the raw JSON
 	if _, registered := m.stateTypes[name]; !registered {
 		m.states[name] = &RawState{data: rawState}
@@ -424,7 +441,7 @@ func (m *Manager) cleanupSingleState(name string, rawState json.RawMessage) erro
 
 // CleanupStateByName loads and cleans up a specific state by name if it implements CleanableState.
 // Returns an error if the state doesn't exist, isn't registered, or cleanup fails.
-func (m *Manager) CleanupStateByName(name string) error {
+func (m *managerImpl) CleanupStateByName(name string) error {
 	if m == nil {
 		return nil
 	}
@@ -461,7 +478,7 @@ func (m *Manager) CleanupStateByName(name string) error {
 
 // PerformCleanup retrieves all states from the state file and calls Cleanup on registered states that support it.
 // Unregistered states are preserved in their original state.
-func (m *Manager) PerformCleanup() error {
+func (m *managerImpl) PerformCleanup() error {
 	if m == nil {
 		return nil
 	}
@@ -491,7 +508,7 @@ func (m *Manager) PerformCleanup() error {
 }
 
 // GetSavedStateNames returns all state names that are currently saved in the state file.
-func (m *Manager) GetSavedStateNames() ([]string, error) {
+func (m *managerImpl) GetSavedStateNames() ([]string, error) {
 	if m == nil {
 		return nil, nil
 	}

client/internal/statemanager/mocks/manager.go

@@ -0,0 +1,312 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: client/internal/statemanager/manager.go
+//
+// Generated by this command:
+//
+//	mockgen -source client/internal/statemanager/manager.go -destination lient/internal/statemanager/manager_mock.go Manager
+//
+
+// Package mocks is a generated GoMock package.
+package mocks
+
+import (
+	context "context"
+	reflect "reflect"
+
+	statemanager "github.com/netbirdio/netbird/client/internal/statemanager"
+	gomock "go.uber.org/mock/gomock"
+)
+
+// MockState is a mock of State interface.
+type MockState struct {
+	ctrl     *gomock.Controller
+	recorder *MockStateMockRecorder
+	isgomock struct{}
+}
+
+// MockStateMockRecorder is the mock recorder for MockState.
+type MockStateMockRecorder struct {
+	mock *MockState
+}
+
+// NewMockState creates a new mock instance.
+func NewMockState(ctrl *gomock.Controller) *MockState {
+	mock := &MockState{ctrl: ctrl}
+	mock.recorder = &MockStateMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockState) EXPECT() *MockStateMockRecorder {
+	return m.recorder
+}
+
+// Name mocks base method.
+func (m *MockState) Name() string {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Name")
+	ret0, _ := ret[0].(string)
+	return ret0
+}
+
+// Name indicates an expected call of Name.
+func (mr *MockStateMockRecorder) Name() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Name", reflect.TypeOf((*MockState)(nil).Name))
+}
+
+// MockCleanableState is a mock of CleanableState interface.
+type MockCleanableState struct {
+	ctrl     *gomock.Controller
+	recorder *MockCleanableStateMockRecorder
+	isgomock struct{}
+}
+
+// MockCleanableStateMockRecorder is the mock recorder for MockCleanableState.
+type MockCleanableStateMockRecorder struct {
+	mock *MockCleanableState
+}
+
+// NewMockCleanableState creates a new mock instance.
+func NewMockCleanableState(ctrl *gomock.Controller) *MockCleanableState {
+	mock := &MockCleanableState{ctrl: ctrl}
+	mock.recorder = &MockCleanableStateMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockCleanableState) EXPECT() *MockCleanableStateMockRecorder {
+	return m.recorder
+}
+
+// Cleanup mocks base method.
+func (m *MockCleanableState) Cleanup() error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Cleanup")
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Cleanup indicates an expected call of Cleanup.
+func (mr *MockCleanableStateMockRecorder) Cleanup() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Cleanup", reflect.TypeOf((*MockCleanableState)(nil).Cleanup))
+}
+
+// Name mocks base method.
+func (m *MockCleanableState) Name() string {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Name")
+	ret0, _ := ret[0].(string)
+	return ret0
+}
+
+// Name indicates an expected call of Name.
+func (mr *MockCleanableStateMockRecorder) Name() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Name", reflect.TypeOf((*MockCleanableState)(nil).Name))
+}
+
+// MockManager is a mock of Manager interface.
+type MockManager struct {
+	ctrl     *gomock.Controller
+	recorder *MockManagerMockRecorder
+	isgomock struct{}
+}
+
+// MockManagerMockRecorder is the mock recorder for MockManager.
+type MockManagerMockRecorder struct {
+	mock *MockManager
+}
+
+// NewMockManager creates a new mock instance.
+func NewMockManager(ctrl *gomock.Controller) *MockManager {
+	mock := &MockManager{ctrl: ctrl}
+	mock.recorder = &MockManagerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockManager) EXPECT() *MockManagerMockRecorder {
+	return m.recorder
+}
+
+// CleanupStateByName mocks base method.
+func (m *MockManager) CleanupStateByName(name string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "CleanupStateByName", name)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// CleanupStateByName indicates an expected call of CleanupStateByName.
+func (mr *MockManagerMockRecorder) CleanupStateByName(name any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CleanupStateByName", reflect.TypeOf((*MockManager)(nil).CleanupStateByName), name)
+}
+
+// DeleteAllStates mocks base method.
+func (m *MockManager) DeleteAllStates() (int, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DeleteAllStates")
+	ret0, _ := ret[0].(int)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DeleteAllStates indicates an expected call of DeleteAllStates.
+func (mr *MockManagerMockRecorder) DeleteAllStates() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAllStates", reflect.TypeOf((*MockManager)(nil).DeleteAllStates))
+}
+
+// DeleteState mocks base method.
+func (m *MockManager) DeleteState(state statemanager.State) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DeleteState", state)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// DeleteState indicates an expected call of DeleteState.
+func (mr *MockManagerMockRecorder) DeleteState(state any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteState", reflect.TypeOf((*MockManager)(nil).DeleteState), state)
+}
+
+// DeleteStateByName mocks base method.
+func (m *MockManager) DeleteStateByName(stateName string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DeleteStateByName", stateName)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// DeleteStateByName indicates an expected call of DeleteStateByName.
+func (mr *MockManagerMockRecorder) DeleteStateByName(stateName any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteStateByName", reflect.TypeOf((*MockManager)(nil).DeleteStateByName), stateName)
+}
+
+// GetSavedStateNames mocks base method.
+func (m *MockManager) GetSavedStateNames() ([]string, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetSavedStateNames")
+	ret0, _ := ret[0].([]string)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetSavedStateNames indicates an expected call of GetSavedStateNames.
+func (mr *MockManagerMockRecorder) GetSavedStateNames() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSavedStateNames", reflect.TypeOf((*MockManager)(nil).GetSavedStateNames))
+}
+
+// GetState mocks base method.
+func (m *MockManager) GetState(state statemanager.State) statemanager.State {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetState", state)
+	ret0, _ := ret[0].(statemanager.State)
+	return ret0
+}
+
+// GetState indicates an expected call of GetState.
+func (mr *MockManagerMockRecorder) GetState(state any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetState", reflect.TypeOf((*MockManager)(nil).GetState), state)
+}
+
+// LoadState mocks base method.
+func (m *MockManager) LoadState(state statemanager.State) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "LoadState", state)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// LoadState indicates an expected call of LoadState.
+func (mr *MockManagerMockRecorder) LoadState(state any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LoadState", reflect.TypeOf((*MockManager)(nil).LoadState), state)
+}
+
+// PerformCleanup mocks base method.
+func (m *MockManager) PerformCleanup() error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "PerformCleanup")
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// PerformCleanup indicates an expected call of PerformCleanup.
+func (mr *MockManagerMockRecorder) PerformCleanup() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PerformCleanup", reflect.TypeOf((*MockManager)(nil).PerformCleanup))
+}
+
+// PersistState mocks base method.
+func (m *MockManager) PersistState(ctx context.Context) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "PersistState", ctx)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// PersistState indicates an expected call of PersistState.
+func (mr *MockManagerMockRecorder) PersistState(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PersistState", reflect.TypeOf((*MockManager)(nil).PersistState), ctx)
+}
+
+// RegisterState mocks base method.
+func (m *MockManager) RegisterState(state statemanager.State) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "RegisterState", state)
+}
+
+// RegisterState indicates an expected call of RegisterState.
+func (mr *MockManagerMockRecorder) RegisterState(state any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RegisterState", reflect.TypeOf((*MockManager)(nil).RegisterState), state)
+}
+
+// Start mocks base method.
+func (m *MockManager) Start() {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Start")
+}
+
+// Start indicates an expected call of Start.
+func (mr *MockManagerMockRecorder) Start() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Start", reflect.TypeOf((*MockManager)(nil).Start))
+}
+
+// Stop mocks base method.
+func (m *MockManager) Stop(ctx context.Context) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Stop", ctx)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Stop indicates an expected call of Stop.
+func (mr *MockManagerMockRecorder) Stop(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Stop", reflect.TypeOf((*MockManager)(nil).Stop), ctx)
+}
+
+// UpdateState mocks base method.
+func (m *MockManager) UpdateState(state statemanager.State) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "UpdateState", state)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// UpdateState indicates an expected call of UpdateState.
+func (mr *MockManagerMockRecorder) UpdateState(state any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateState", reflect.TypeOf((*MockManager)(nil).UpdateState), state)
+}

client/server/server_test.go

@@ -10,7 +10,6 @@ import (
 	"go.opentelemetry.io/otel"
 
 	"github.com/netbirdio/management-integrations/integrations"
-
 	log "github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/keepalive"
@@ -129,13 +128,13 @@ func startManagement(t *testing.T, signalAddr string, counter *int) (*grpc.Serve
 	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
 	require.NoError(t, err)
 
-	accountManager, err := server.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, ia, metrics, port_forwarding.NewControllerMock())
+	accountManager, err := server.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, ia, metrics, port_forwarding.NewControllerMock(), settings.NewManagerMock())
 	if err != nil {
 		return nil, "", err
 	}
 
 	secretsManager := server.NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig, config.Relay)
-	mgmtServer, err := server.NewServer(context.Background(), config, accountManager, settings.NewManager(store), peersUpdateManager, secretsManager, nil, nil, nil)
+	mgmtServer, err := server.NewServer(context.Background(), config, accountManager, settings.NewManagerMock(), peersUpdateManager, secretsManager, nil, nil, nil)
 	if err != nil {
 		return nil, "", err
 	}

client/server/state_generic.go

@@ -8,7 +8,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 )
 
-func registerStates(mgr *statemanager.Manager) {
+func registerStates(mgr statemanager.Manager) {
 	mgr.RegisterState(&dns.ShutdownState{})
 	mgr.RegisterState(&systemops.ShutdownState{})
 }

client/server/state_linux.go

@@ -10,7 +10,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 )
 
-func registerStates(mgr *statemanager.Manager) {
+func registerStates(mgr statemanager.Manager) {
 	mgr.RegisterState(&dns.ShutdownState{})
 	mgr.RegisterState(&systemops.ShutdownState{})
 	mgr.RegisterState(&nftables.ShutdownState{})

client/server/trace.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"net/netip"
 
 	fw "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/firewall/uspfilter"
@@ -41,11 +42,21 @@ func (s *Server) TracePacket(_ context.Context, req *proto.TracePacketRequest) (
 		srcIP = engine.GetWgAddr()
 	}
 
+	srcAddr, ok := netip.AddrFromSlice(srcIP)
+	if !ok {
+		return nil, fmt.Errorf("invalid source IP address")
+	}
+
 	dstIP := net.ParseIP(req.GetDestinationIp())
 	if req.GetDestinationIp() == "self" {
 		dstIP = engine.GetWgAddr()
 	}
 
+	dstAddr, ok := netip.AddrFromSlice(dstIP)
+	if !ok {
+		return nil, fmt.Errorf("invalid source IP address")
+	}
+
 	if srcIP == nil || dstIP == nil {
 		return nil, fmt.Errorf("invalid IP address")
 	}
@@ -85,8 +96,8 @@ func (s *Server) TracePacket(_ context.Context, req *proto.TracePacketRequest) (
 	}
 
 	builder := &uspfilter.PacketBuilder{
-		SrcIP:     srcIP,
-		DstIP:     dstIP,
+		SrcIP:     srcAddr,
+		DstIP:     dstAddr,
 		Protocol:  protocol,
 		SrcPort:   uint16(req.GetSourcePort()),
 		DstPort:   uint16(req.GetDestinationPort()),

flow/client/auth.go

@@ -0,0 +1,32 @@
+package client
+
+import (
+	"context"
+	"fmt"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
+)
+
+var _ credentials.PerRPCCredentials = (*authToken)(nil)
+
+type authToken struct {
+	metaMap map[string]string
+}
+
+func (t authToken) GetRequestMetadata(context.Context, ...string) (map[string]string, error) {
+	return t.metaMap, nil
+}
+
+func (authToken) RequireTransportSecurity() bool {
+	return false // Set to true if you want to require a secure connection
+}
+
+// WithAuthToken returns a DialOption which sets the receiver flow credentials and places auth state on each outbound RPC
+func withAuthToken(payload, signature string) grpc.DialOption {
+	value := fmt.Sprintf("%s.%s", signature, payload)
+	authMap := map[string]string{
+		"authorization": "Bearer " + value,
+	}
+	return grpc.WithPerRPCCredentials(authToken{metaMap: authMap})
+}

flow/client/client.go

@@ -0,0 +1,167 @@
+package client
+
+import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"errors"
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/cenkalti/backoff/v4"
+	log "github.com/sirupsen/logrus"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/keepalive"
+
+	"github.com/netbirdio/netbird/flow/proto"
+	"github.com/netbirdio/netbird/util/embeddedroots"
+	nbgrpc "github.com/netbirdio/netbird/util/grpc"
+)
+
+type GRPCClient struct {
+	realClient proto.FlowServiceClient
+	clientConn *grpc.ClientConn
+	stream     proto.FlowService_EventsClient
+	streamMu   sync.Mutex
+}
+
+func NewClient(addr, payload, signature string, interval time.Duration) (*GRPCClient, error) {
+	var opts []grpc.DialOption
+
+	if strings.Contains(addr, "443") {
+		certPool, err := x509.SystemCertPool()
+		if err != nil || certPool == nil {
+			log.Debugf("System cert pool not available; falling back to embedded cert, error: %v", err)
+			certPool = embeddedroots.Get()
+		}
+
+		opts = append(opts, grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{
+			RootCAs: certPool,
+		})))
+	} else {
+		opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	}
+
+	opts = append(opts,
+		nbgrpc.WithCustomDialer(),
+		grpc.WithIdleTimeout(interval*2),
+		grpc.WithKeepaliveParams(keepalive.ClientParameters{
+			Time:    30 * time.Second,
+			Timeout: 10 * time.Second,
+		}),
+		withAuthToken(payload, signature),
+		grpc.WithDefaultServiceConfig(`{"healthCheckConfig": {"serviceName": ""}}`),
+	)
+
+	conn, err := grpc.NewClient(addr, opts...)
+	if err != nil {
+		return nil, fmt.Errorf("creating new grpc client: %w", err)
+	}
+
+	return &GRPCClient{
+		realClient: proto.NewFlowServiceClient(conn),
+		clientConn: conn,
+	}, nil
+}
+
+func (c *GRPCClient) Close() error {
+	c.streamMu.Lock()
+	defer c.streamMu.Unlock()
+
+	c.stream = nil
+	return c.clientConn.Close()
+}
+
+func (c *GRPCClient) Receive(ctx context.Context, interval time.Duration, msgHandler func(msg *proto.FlowEventAck) error) error {
+	backOff := defaultBackoff(ctx, interval)
+	operation := func() error {
+		return c.establishStreamAndReceive(ctx, msgHandler)
+	}
+
+	if err := backoff.Retry(operation, backOff); err != nil {
+		return fmt.Errorf("receive failed permanently: %w", err)
+	}
+
+	return nil
+}
+
+func (c *GRPCClient) establishStreamAndReceive(ctx context.Context, msgHandler func(msg *proto.FlowEventAck) error) error {
+	if c.clientConn.GetState() == connectivity.Shutdown {
+		return backoff.Permanent(errors.New("connection to flow receiver has been shut down"))
+	}
+
+	stream, err := c.realClient.Events(ctx, grpc.WaitForReady(true))
+	if err != nil {
+		return fmt.Errorf("create event stream: %w", err)
+	}
+
+	if err = checkHeader(stream); err != nil {
+		return fmt.Errorf("check header: %w", err)
+	}
+
+	c.streamMu.Lock()
+	c.stream = stream
+	c.streamMu.Unlock()
+
+	return c.receive(stream, msgHandler)
+}
+
+func (c *GRPCClient) receive(stream proto.FlowService_EventsClient, msgHandler func(msg *proto.FlowEventAck) error) error {
+	for {
+		msg, err := stream.Recv()
+		if err != nil {
+			return fmt.Errorf("receive from stream: %w", err)
+		}
+
+		if err := msgHandler(msg); err != nil {
+			return fmt.Errorf("handle message: %w", err)
+		}
+	}
+}
+
+func checkHeader(stream proto.FlowService_EventsClient) error {
+	header, err := stream.Header()
+	if err != nil {
+		log.Errorf("waiting for flow receiver header: %s", err)
+		return fmt.Errorf("wait for header: %w", err)
+	}
+
+	if len(header) == 0 {
+		log.Error("flow receiver sent no headers")
+		return fmt.Errorf("should have headers")
+	}
+	return nil
+}
+
+func defaultBackoff(ctx context.Context, interval time.Duration) backoff.BackOff {
+	return backoff.WithContext(&backoff.ExponentialBackOff{
+		InitialInterval:     800 * time.Millisecond,
+		RandomizationFactor: 1,
+		Multiplier:          1.7,
+		MaxInterval:         interval / 2,
+		MaxElapsedTime:      3 * 30 * 24 * time.Hour, // 3 months
+		Stop:                backoff.Stop,
+		Clock:               backoff.SystemClock,
+	}, ctx)
+}
+
+func (c *GRPCClient) Send(event *proto.FlowEvent) error {
+	c.streamMu.Lock()
+	stream := c.stream
+	c.streamMu.Unlock()
+
+	if stream == nil {
+		return errors.New("stream not initialized")
+	}
+
+	if err := stream.Send(event); err != nil {
+		return fmt.Errorf("send flow event: %w", err)
+	}
+
+	return nil
+}

flow/proto/flow.pb.go

@@ -0,0 +1,769 @@
+// Code generated by protoc-gen-go. DO NOT EDIT.
+// versions:
+// 	protoc-gen-go v1.26.0
+// 	protoc        v4.24.3
+// source: flow.proto
+
+package proto
+
+import (
+	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
+	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
+	timestamppb "google.golang.org/protobuf/types/known/timestamppb"
+	reflect "reflect"
+	sync "sync"
+)
+
+const (
+	// Verify that this generated code is sufficiently up-to-date.
+	_ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion)
+	// Verify that runtime/protoimpl is sufficiently up-to-date.
+	_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
+)
+
+// Flow event types
+type Type int32
+
+const (
+	Type_TYPE_UNKNOWN Type = 0
+	Type_TYPE_START   Type = 1
+	Type_TYPE_END     Type = 2
+	Type_TYPE_DROP    Type = 3
+)
+
+// Enum value maps for Type.
+var (
+	Type_name = map[int32]string{
+		0: "TYPE_UNKNOWN",
+		1: "TYPE_START",
+		2: "TYPE_END",
+		3: "TYPE_DROP",
+	}
+	Type_value = map[string]int32{
+		"TYPE_UNKNOWN": 0,
+		"TYPE_START":   1,
+		"TYPE_END":     2,
+		"TYPE_DROP":    3,
+	}
+)
+
+func (x Type) Enum() *Type {
+	p := new(Type)
+	*p = x
+	return p
+}
+
+func (x Type) String() string {
+	return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x))
+}
+
+func (Type) Descriptor() protoreflect.EnumDescriptor {
+	return file_flow_proto_enumTypes[0].Descriptor()
+}
+
+func (Type) Type() protoreflect.EnumType {
+	return &file_flow_proto_enumTypes[0]
+}
+
+func (x Type) Number() protoreflect.EnumNumber {
+	return protoreflect.EnumNumber(x)
+}
+
+// Deprecated: Use Type.Descriptor instead.
+func (Type) EnumDescriptor() ([]byte, []int) {
+	return file_flow_proto_rawDescGZIP(), []int{0}
+}
+
+// Flow direction
+type Direction int32
+
+const (
+	Direction_DIRECTION_UNKNOWN Direction = 0
+	Direction_INGRESS           Direction = 1
+	Direction_EGRESS            Direction = 2
+)
+
+// Enum value maps for Direction.
+var (
+	Direction_name = map[int32]string{
+		0: "DIRECTION_UNKNOWN",
+		1: "INGRESS",
+		2: "EGRESS",
+	}
+	Direction_value = map[string]int32{
+		"DIRECTION_UNKNOWN": 0,
+		"INGRESS":           1,
+		"EGRESS":            2,
+	}
+)
+
+func (x Direction) Enum() *Direction {
+	p := new(Direction)
+	*p = x
+	return p
+}
+
+func (x Direction) String() string {
+	return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x))
+}
+
+func (Direction) Descriptor() protoreflect.EnumDescriptor {
+	return file_flow_proto_enumTypes[1].Descriptor()
+}
+
+func (Direction) Type() protoreflect.EnumType {
+	return &file_flow_proto_enumTypes[1]
+}
+
+func (x Direction) Number() protoreflect.EnumNumber {
+	return protoreflect.EnumNumber(x)
+}
+
+// Deprecated: Use Direction.Descriptor instead.
+func (Direction) EnumDescriptor() ([]byte, []int) {
+	return file_flow_proto_rawDescGZIP(), []int{1}
+}
+
+type FlowEvent struct {
+	state         protoimpl.MessageState
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+
+	// Unique client event identifier
+	EventId []byte `protobuf:"bytes,1,opt,name=event_id,json=eventId,proto3" json:"event_id,omitempty"`
+	// When the event occurred
+	Timestamp *timestamppb.Timestamp `protobuf:"bytes,2,opt,name=timestamp,proto3" json:"timestamp,omitempty"`
+	// Public key of the sending peer
+	PublicKey  []byte      `protobuf:"bytes,3,opt,name=public_key,json=publicKey,proto3" json:"public_key,omitempty"`
+	FlowFields *FlowFields `protobuf:"bytes,4,opt,name=flow_fields,json=flowFields,proto3" json:"flow_fields,omitempty"`
+}
+
+func (x *FlowEvent) Reset() {
+	*x = FlowEvent{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_flow_proto_msgTypes[0]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+func (x *FlowEvent) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*FlowEvent) ProtoMessage() {}
+
+func (x *FlowEvent) ProtoReflect() protoreflect.Message {
+	mi := &file_flow_proto_msgTypes[0]
+	if protoimpl.UnsafeEnabled && x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use FlowEvent.ProtoReflect.Descriptor instead.
+func (*FlowEvent) Descriptor() ([]byte, []int) {
+	return file_flow_proto_rawDescGZIP(), []int{0}
+}
+
+func (x *FlowEvent) GetEventId() []byte {
+	if x != nil {
+		return x.EventId
+	}
+	return nil
+}
+
+func (x *FlowEvent) GetTimestamp() *timestamppb.Timestamp {
+	if x != nil {
+		return x.Timestamp
+	}
+	return nil
+}
+
+func (x *FlowEvent) GetPublicKey() []byte {
+	if x != nil {
+		return x.PublicKey
+	}
+	return nil
+}
+
+func (x *FlowEvent) GetFlowFields() *FlowFields {
+	if x != nil {
+		return x.FlowFields
+	}
+	return nil
+}
+
+type FlowEventAck struct {
+	state         protoimpl.MessageState
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+
+	// Unique client event identifier that has been ack'ed
+	EventId []byte `protobuf:"bytes,1,opt,name=event_id,json=eventId,proto3" json:"event_id,omitempty"`
+}
+
+func (x *FlowEventAck) Reset() {
+	*x = FlowEventAck{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_flow_proto_msgTypes[1]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+func (x *FlowEventAck) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*FlowEventAck) ProtoMessage() {}
+
+func (x *FlowEventAck) ProtoReflect() protoreflect.Message {
+	mi := &file_flow_proto_msgTypes[1]
+	if protoimpl.UnsafeEnabled && x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use FlowEventAck.ProtoReflect.Descriptor instead.
+func (*FlowEventAck) Descriptor() ([]byte, []int) {
+	return file_flow_proto_rawDescGZIP(), []int{1}
+}
+
+func (x *FlowEventAck) GetEventId() []byte {
+	if x != nil {
+		return x.EventId
+	}
+	return nil
+}
+
+type FlowFields struct {
+	state         protoimpl.MessageState
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+
+	// Unique client flow session identifier
+	FlowId []byte `protobuf:"bytes,1,opt,name=flow_id,json=flowId,proto3" json:"flow_id,omitempty"`
+	// Flow type
+	Type Type `protobuf:"varint,2,opt,name=type,proto3,enum=flow.Type" json:"type,omitempty"`
+	// RuleId identifies the rule that allowed or denied the connection
+	RuleId []byte `protobuf:"bytes,3,opt,name=rule_id,json=ruleId,proto3" json:"rule_id,omitempty"`
+	// Initiating traffic direction
+	Direction Direction `protobuf:"varint,4,opt,name=direction,proto3,enum=flow.Direction" json:"direction,omitempty"`
+	// IP protocol number
+	Protocol uint32 `protobuf:"varint,5,opt,name=protocol,proto3" json:"protocol,omitempty"`
+	// Source IP address
+	SourceIp []byte `protobuf:"bytes,6,opt,name=source_ip,json=sourceIp,proto3" json:"source_ip,omitempty"`
+	// Destination IP address
+	DestIp []byte `protobuf:"bytes,7,opt,name=dest_ip,json=destIp,proto3" json:"dest_ip,omitempty"`
+	// Layer 4 -specific information
+	//
+	// Types that are assignable to ConnectionInfo:
+	//
+	//	*FlowFields_PortInfo
+	//	*FlowFields_IcmpInfo
+	ConnectionInfo isFlowFields_ConnectionInfo `protobuf_oneof:"connection_info"`
+	// Number of packets
+	RxPackets uint64 `protobuf:"varint,10,opt,name=rx_packets,json=rxPackets,proto3" json:"rx_packets,omitempty"`
+	TxPackets uint64 `protobuf:"varint,11,opt,name=tx_packets,json=txPackets,proto3" json:"tx_packets,omitempty"`
+	// Number of bytes
+	RxBytes uint64 `protobuf:"varint,12,opt,name=rx_bytes,json=rxBytes,proto3" json:"rx_bytes,omitempty"`
+	TxBytes uint64 `protobuf:"varint,13,opt,name=tx_bytes,json=txBytes,proto3" json:"tx_bytes,omitempty"`
+	// Resource ID
+	SourceResourceId []byte `protobuf:"bytes,14,opt,name=source_resource_id,json=sourceResourceId,proto3" json:"source_resource_id,omitempty"`
+	DestResourceId   []byte `protobuf:"bytes,15,opt,name=dest_resource_id,json=destResourceId,proto3" json:"dest_resource_id,omitempty"`
+}
+
+func (x *FlowFields) Reset() {
+	*x = FlowFields{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_flow_proto_msgTypes[2]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+func (x *FlowFields) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*FlowFields) ProtoMessage() {}
+
+func (x *FlowFields) ProtoReflect() protoreflect.Message {
+	mi := &file_flow_proto_msgTypes[2]
+	if protoimpl.UnsafeEnabled && x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use FlowFields.ProtoReflect.Descriptor instead.
+func (*FlowFields) Descriptor() ([]byte, []int) {
+	return file_flow_proto_rawDescGZIP(), []int{2}
+}
+
+func (x *FlowFields) GetFlowId() []byte {
+	if x != nil {
+		return x.FlowId
+	}
+	return nil
+}
+
+func (x *FlowFields) GetType() Type {
+	if x != nil {
+		return x.Type
+	}
+	return Type_TYPE_UNKNOWN
+}
+
+func (x *FlowFields) GetRuleId() []byte {
+	if x != nil {
+		return x.RuleId
+	}
+	return nil
+}
+
+func (x *FlowFields) GetDirection() Direction {
+	if x != nil {
+		return x.Direction
+	}
+	return Direction_DIRECTION_UNKNOWN
+}
+
+func (x *FlowFields) GetProtocol() uint32 {
+	if x != nil {
+		return x.Protocol
+	}
+	return 0
+}
+
+func (x *FlowFields) GetSourceIp() []byte {
+	if x != nil {
+		return x.SourceIp
+	}
+	return nil
+}
+
+func (x *FlowFields) GetDestIp() []byte {
+	if x != nil {
+		return x.DestIp
+	}
+	return nil
+}
+
+func (m *FlowFields) GetConnectionInfo() isFlowFields_ConnectionInfo {
+	if m != nil {
+		return m.ConnectionInfo
+	}
+	return nil
+}
+
+func (x *FlowFields) GetPortInfo() *PortInfo {
+	if x, ok := x.GetConnectionInfo().(*FlowFields_PortInfo); ok {
+		return x.PortInfo
+	}
+	return nil
+}
+
+func (x *FlowFields) GetIcmpInfo() *ICMPInfo {
+	if x, ok := x.GetConnectionInfo().(*FlowFields_IcmpInfo); ok {
+		return x.IcmpInfo
+	}
+	return nil
+}
+
+func (x *FlowFields) GetRxPackets() uint64 {
+	if x != nil {
+		return x.RxPackets
+	}
+	return 0
+}
+
+func (x *FlowFields) GetTxPackets() uint64 {
+	if x != nil {
+		return x.TxPackets
+	}
+	return 0
+}
+
+func (x *FlowFields) GetRxBytes() uint64 {
+	if x != nil {
+		return x.RxBytes
+	}
+	return 0
+}
+
+func (x *FlowFields) GetTxBytes() uint64 {
+	if x != nil {
+		return x.TxBytes
+	}
+	return 0
+}
+
+func (x *FlowFields) GetSourceResourceId() []byte {
+	if x != nil {
+		return x.SourceResourceId
+	}
+	return nil
+}
+
+func (x *FlowFields) GetDestResourceId() []byte {
+	if x != nil {
+		return x.DestResourceId
+	}
+	return nil
+}
+
+type isFlowFields_ConnectionInfo interface {
+	isFlowFields_ConnectionInfo()
+}
+
+type FlowFields_PortInfo struct {
+	// TCP/UDP port information
+	PortInfo *PortInfo `protobuf:"bytes,8,opt,name=port_info,json=portInfo,proto3,oneof"`
+}
+
+type FlowFields_IcmpInfo struct {
+	// ICMP type and code
+	IcmpInfo *ICMPInfo `protobuf:"bytes,9,opt,name=icmp_info,json=icmpInfo,proto3,oneof"`
+}
+
+func (*FlowFields_PortInfo) isFlowFields_ConnectionInfo() {}
+
+func (*FlowFields_IcmpInfo) isFlowFields_ConnectionInfo() {}
+
+// TCP/UDP port information
+type PortInfo struct {
+	state         protoimpl.MessageState
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+
+	SourcePort uint32 `protobuf:"varint,1,opt,name=source_port,json=sourcePort,proto3" json:"source_port,omitempty"`
+	DestPort   uint32 `protobuf:"varint,2,opt,name=dest_port,json=destPort,proto3" json:"dest_port,omitempty"`
+}
+
+func (x *PortInfo) Reset() {
+	*x = PortInfo{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_flow_proto_msgTypes[3]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+func (x *PortInfo) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*PortInfo) ProtoMessage() {}
+
+func (x *PortInfo) ProtoReflect() protoreflect.Message {
+	mi := &file_flow_proto_msgTypes[3]
+	if protoimpl.UnsafeEnabled && x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use PortInfo.ProtoReflect.Descriptor instead.
+func (*PortInfo) Descriptor() ([]byte, []int) {
+	return file_flow_proto_rawDescGZIP(), []int{3}
+}
+
+func (x *PortInfo) GetSourcePort() uint32 {
+	if x != nil {
+		return x.SourcePort
+	}
+	return 0
+}
+
+func (x *PortInfo) GetDestPort() uint32 {
+	if x != nil {
+		return x.DestPort
+	}
+	return 0
+}
+
+// ICMP message information
+type ICMPInfo struct {
+	state         protoimpl.MessageState
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+
+	IcmpType uint32 `protobuf:"varint,1,opt,name=icmp_type,json=icmpType,proto3" json:"icmp_type,omitempty"`
+	IcmpCode uint32 `protobuf:"varint,2,opt,name=icmp_code,json=icmpCode,proto3" json:"icmp_code,omitempty"`
+}
+
+func (x *ICMPInfo) Reset() {
+	*x = ICMPInfo{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_flow_proto_msgTypes[4]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+func (x *ICMPInfo) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*ICMPInfo) ProtoMessage() {}
+
+func (x *ICMPInfo) ProtoReflect() protoreflect.Message {
+	mi := &file_flow_proto_msgTypes[4]
+	if protoimpl.UnsafeEnabled && x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use ICMPInfo.ProtoReflect.Descriptor instead.
+func (*ICMPInfo) Descriptor() ([]byte, []int) {
+	return file_flow_proto_rawDescGZIP(), []int{4}
+}
+
+func (x *ICMPInfo) GetIcmpType() uint32 {
+	if x != nil {
+		return x.IcmpType
+	}
+	return 0
+}
+
+func (x *ICMPInfo) GetIcmpCode() uint32 {
+	if x != nil {
+		return x.IcmpCode
+	}
+	return 0
+}
+
+var File_flow_proto protoreflect.FileDescriptor
+
+var file_flow_proto_rawDesc = []byte{
+	0x0a, 0x0a, 0x66, 0x6c, 0x6f, 0x77, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x04, 0x66, 0x6c,
+	0x6f, 0x77, 0x1a, 0x1f, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f,
+	0x62, 0x75, 0x66, 0x2f, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x2e, 0x70, 0x72,
+	0x6f, 0x74, 0x6f, 0x22, 0xb2, 0x01, 0x0a, 0x09, 0x46, 0x6c, 0x6f, 0x77, 0x45, 0x76, 0x65, 0x6e,
+	0x74, 0x12, 0x19, 0x0a, 0x08, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20,
+	0x01, 0x28, 0x0c, 0x52, 0x07, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x49, 0x64, 0x12, 0x38, 0x0a, 0x09,
+	0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32,
+	0x1a, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75,
+	0x66, 0x2e, 0x54, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x52, 0x09, 0x74, 0x69, 0x6d,
+	0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x12, 0x1d, 0x0a, 0x0a, 0x70, 0x75, 0x62, 0x6c, 0x69, 0x63,
+	0x5f, 0x6b, 0x65, 0x79, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x09, 0x70, 0x75, 0x62, 0x6c,
+	0x69, 0x63, 0x4b, 0x65, 0x79, 0x12, 0x31, 0x0a, 0x0b, 0x66, 0x6c, 0x6f, 0x77, 0x5f, 0x66, 0x69,
+	0x65, 0x6c, 0x64, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x10, 0x2e, 0x66, 0x6c, 0x6f,
+	0x77, 0x2e, 0x46, 0x6c, 0x6f, 0x77, 0x46, 0x69, 0x65, 0x6c, 0x64, 0x73, 0x52, 0x0a, 0x66, 0x6c,
+	0x6f, 0x77, 0x46, 0x69, 0x65, 0x6c, 0x64, 0x73, 0x22, 0x29, 0x0a, 0x0c, 0x46, 0x6c, 0x6f, 0x77,
+	0x45, 0x76, 0x65, 0x6e, 0x74, 0x41, 0x63, 0x6b, 0x12, 0x19, 0x0a, 0x08, 0x65, 0x76, 0x65, 0x6e,
+	0x74, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x65, 0x76, 0x65, 0x6e,
+	0x74, 0x49, 0x64, 0x22, 0x9c, 0x04, 0x0a, 0x0a, 0x46, 0x6c, 0x6f, 0x77, 0x46, 0x69, 0x65, 0x6c,
+	0x64, 0x73, 0x12, 0x17, 0x0a, 0x07, 0x66, 0x6c, 0x6f, 0x77, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20,
+	0x01, 0x28, 0x0c, 0x52, 0x06, 0x66, 0x6c, 0x6f, 0x77, 0x49, 0x64, 0x12, 0x1e, 0x0a, 0x04, 0x74,
+	0x79, 0x70, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x0a, 0x2e, 0x66, 0x6c, 0x6f, 0x77,
+	0x2e, 0x54, 0x79, 0x70, 0x65, 0x52, 0x04, 0x74, 0x79, 0x70, 0x65, 0x12, 0x17, 0x0a, 0x07, 0x72,
+	0x75, 0x6c, 0x65, 0x5f, 0x69, 0x64, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x06, 0x72, 0x75,
+	0x6c, 0x65, 0x49, 0x64, 0x12, 0x2d, 0x0a, 0x09, 0x64, 0x69, 0x72, 0x65, 0x63, 0x74, 0x69, 0x6f,
+	0x6e, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0e, 0x32, 0x0f, 0x2e, 0x66, 0x6c, 0x6f, 0x77, 0x2e, 0x44,
+	0x69, 0x72, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x09, 0x64, 0x69, 0x72, 0x65, 0x63, 0x74,
+	0x69, 0x6f, 0x6e, 0x12, 0x1a, 0x0a, 0x08, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x18,
+	0x05, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x08, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x12,
+	0x1b, 0x0a, 0x09, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x5f, 0x69, 0x70, 0x18, 0x06, 0x20, 0x01,
+	0x28, 0x0c, 0x52, 0x08, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x49, 0x70, 0x12, 0x17, 0x0a, 0x07,
+	0x64, 0x65, 0x73, 0x74, 0x5f, 0x69, 0x70, 0x18, 0x07, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x06, 0x64,
+	0x65, 0x73, 0x74, 0x49, 0x70, 0x12, 0x2d, 0x0a, 0x09, 0x70, 0x6f, 0x72, 0x74, 0x5f, 0x69, 0x6e,
+	0x66, 0x6f, 0x18, 0x08, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x66, 0x6c, 0x6f, 0x77, 0x2e,
+	0x50, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x48, 0x00, 0x52, 0x08, 0x70, 0x6f, 0x72, 0x74,
+	0x49, 0x6e, 0x66, 0x6f, 0x12, 0x2d, 0x0a, 0x09, 0x69, 0x63, 0x6d, 0x70, 0x5f, 0x69, 0x6e, 0x66,
+	0x6f, 0x18, 0x09, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x66, 0x6c, 0x6f, 0x77, 0x2e, 0x49,
+	0x43, 0x4d, 0x50, 0x49, 0x6e, 0x66, 0x6f, 0x48, 0x00, 0x52, 0x08, 0x69, 0x63, 0x6d, 0x70, 0x49,
+	0x6e, 0x66, 0x6f, 0x12, 0x1d, 0x0a, 0x0a, 0x72, 0x78, 0x5f, 0x70, 0x61, 0x63, 0x6b, 0x65, 0x74,
+	0x73, 0x18, 0x0a, 0x20, 0x01, 0x28, 0x04, 0x52, 0x09, 0x72, 0x78, 0x50, 0x61, 0x63, 0x6b, 0x65,
+	0x74, 0x73, 0x12, 0x1d, 0x0a, 0x0a, 0x74, 0x78, 0x5f, 0x70, 0x61, 0x63, 0x6b, 0x65, 0x74, 0x73,
+	0x18, 0x0b, 0x20, 0x01, 0x28, 0x04, 0x52, 0x09, 0x74, 0x78, 0x50, 0x61, 0x63, 0x6b, 0x65, 0x74,
+	0x73, 0x12, 0x19, 0x0a, 0x08, 0x72, 0x78, 0x5f, 0x62, 0x79, 0x74, 0x65, 0x73, 0x18, 0x0c, 0x20,
+	0x01, 0x28, 0x04, 0x52, 0x07, 0x72, 0x78, 0x42, 0x79, 0x74, 0x65, 0x73, 0x12, 0x19, 0x0a, 0x08,
+	0x74, 0x78, 0x5f, 0x62, 0x79, 0x74, 0x65, 0x73, 0x18, 0x0d, 0x20, 0x01, 0x28, 0x04, 0x52, 0x07,
+	0x74, 0x78, 0x42, 0x79, 0x74, 0x65, 0x73, 0x12, 0x2c, 0x0a, 0x12, 0x73, 0x6f, 0x75, 0x72, 0x63,
+	0x65, 0x5f, 0x72, 0x65, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x5f, 0x69, 0x64, 0x18, 0x0e, 0x20,
+	0x01, 0x28, 0x0c, 0x52, 0x10, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x52, 0x65, 0x73, 0x6f, 0x75,
+	0x72, 0x63, 0x65, 0x49, 0x64, 0x12, 0x28, 0x0a, 0x10, 0x64, 0x65, 0x73, 0x74, 0x5f, 0x72, 0x65,
+	0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x5f, 0x69, 0x64, 0x18, 0x0f, 0x20, 0x01, 0x28, 0x0c, 0x52,
+	0x0e, 0x64, 0x65, 0x73, 0x74, 0x52, 0x65, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x49, 0x64, 0x42,
+	0x11, 0x0a, 0x0f, 0x63, 0x6f, 0x6e, 0x6e, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x6e,
+	0x66, 0x6f, 0x22, 0x48, 0x0a, 0x08, 0x50, 0x6f, 0x72, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x12, 0x1f,
+	0x0a, 0x0b, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x5f, 0x70, 0x6f, 0x72, 0x74, 0x18, 0x01, 0x20,
+	0x01, 0x28, 0x0d, 0x52, 0x0a, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x50, 0x6f, 0x72, 0x74, 0x12,
+	0x1b, 0x0a, 0x09, 0x64, 0x65, 0x73, 0x74, 0x5f, 0x70, 0x6f, 0x72, 0x74, 0x18, 0x02, 0x20, 0x01,
+	0x28, 0x0d, 0x52, 0x08, 0x64, 0x65, 0x73, 0x74, 0x50, 0x6f, 0x72, 0x74, 0x22, 0x44, 0x0a, 0x08,
+	0x49, 0x43, 0x4d, 0x50, 0x49, 0x6e, 0x66, 0x6f, 0x12, 0x1b, 0x0a, 0x09, 0x69, 0x63, 0x6d, 0x70,
+	0x5f, 0x74, 0x79, 0x70, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x08, 0x69, 0x63, 0x6d,
+	0x70, 0x54, 0x79, 0x70, 0x65, 0x12, 0x1b, 0x0a, 0x09, 0x69, 0x63, 0x6d, 0x70, 0x5f, 0x63, 0x6f,
+	0x64, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x08, 0x69, 0x63, 0x6d, 0x70, 0x43, 0x6f,
+	0x64, 0x65, 0x2a, 0x45, 0x0a, 0x04, 0x54, 0x79, 0x70, 0x65, 0x12, 0x10, 0x0a, 0x0c, 0x54, 0x59,
+	0x50, 0x45, 0x5f, 0x55, 0x4e, 0x4b, 0x4e, 0x4f, 0x57, 0x4e, 0x10, 0x00, 0x12, 0x0e, 0x0a, 0x0a,
+	0x54, 0x59, 0x50, 0x45, 0x5f, 0x53, 0x54, 0x41, 0x52, 0x54, 0x10, 0x01, 0x12, 0x0c, 0x0a, 0x08,
+	0x54, 0x59, 0x50, 0x45, 0x5f, 0x45, 0x4e, 0x44, 0x10, 0x02, 0x12, 0x0d, 0x0a, 0x09, 0x54, 0x59,
+	0x50, 0x45, 0x5f, 0x44, 0x52, 0x4f, 0x50, 0x10, 0x03, 0x2a, 0x3b, 0x0a, 0x09, 0x44, 0x69, 0x72,
+	0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x15, 0x0a, 0x11, 0x44, 0x49, 0x52, 0x45, 0x43, 0x54,
+	0x49, 0x4f, 0x4e, 0x5f, 0x55, 0x4e, 0x4b, 0x4e, 0x4f, 0x57, 0x4e, 0x10, 0x00, 0x12, 0x0b, 0x0a,
+	0x07, 0x49, 0x4e, 0x47, 0x52, 0x45, 0x53, 0x53, 0x10, 0x01, 0x12, 0x0a, 0x0a, 0x06, 0x45, 0x47,
+	0x52, 0x45, 0x53, 0x53, 0x10, 0x02, 0x32, 0x42, 0x0a, 0x0b, 0x46, 0x6c, 0x6f, 0x77, 0x53, 0x65,
+	0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x33, 0x0a, 0x06, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x73, 0x12,
+	0x0f, 0x2e, 0x66, 0x6c, 0x6f, 0x77, 0x2e, 0x46, 0x6c, 0x6f, 0x77, 0x45, 0x76, 0x65, 0x6e, 0x74,
+	0x1a, 0x12, 0x2e, 0x66, 0x6c, 0x6f, 0x77, 0x2e, 0x46, 0x6c, 0x6f, 0x77, 0x45, 0x76, 0x65, 0x6e,
+	0x74, 0x41, 0x63, 0x6b, 0x22, 0x00, 0x28, 0x01, 0x30, 0x01, 0x42, 0x08, 0x5a, 0x06, 0x2f, 0x70,
+	0x72, 0x6f, 0x74, 0x6f, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
+}
+
+var (
+	file_flow_proto_rawDescOnce sync.Once
+	file_flow_proto_rawDescData = file_flow_proto_rawDesc
+)
+
+func file_flow_proto_rawDescGZIP() []byte {
+	file_flow_proto_rawDescOnce.Do(func() {
+		file_flow_proto_rawDescData = protoimpl.X.CompressGZIP(file_flow_proto_rawDescData)
+	})
+	return file_flow_proto_rawDescData
+}
+
+var file_flow_proto_enumTypes = make([]protoimpl.EnumInfo, 2)
+var file_flow_proto_msgTypes = make([]protoimpl.MessageInfo, 5)
+var file_flow_proto_goTypes = []interface{}{
+	(Type)(0),                     // 0: flow.Type
+	(Direction)(0),                // 1: flow.Direction
+	(*FlowEvent)(nil),             // 2: flow.FlowEvent
+	(*FlowEventAck)(nil),          // 3: flow.FlowEventAck
+	(*FlowFields)(nil),            // 4: flow.FlowFields
+	(*PortInfo)(nil),              // 5: flow.PortInfo
+	(*ICMPInfo)(nil),              // 6: flow.ICMPInfo
+	(*timestamppb.Timestamp)(nil), // 7: google.protobuf.Timestamp
+}
+var file_flow_proto_depIdxs = []int32{
+	7, // 0: flow.FlowEvent.timestamp:type_name -> google.protobuf.Timestamp
+	4, // 1: flow.FlowEvent.flow_fields:type_name -> flow.FlowFields
+	0, // 2: flow.FlowFields.type:type_name -> flow.Type
+	1, // 3: flow.FlowFields.direction:type_name -> flow.Direction
+	5, // 4: flow.FlowFields.port_info:type_name -> flow.PortInfo
+	6, // 5: flow.FlowFields.icmp_info:type_name -> flow.ICMPInfo
+	2, // 6: flow.FlowService.Events:input_type -> flow.FlowEvent
+	3, // 7: flow.FlowService.Events:output_type -> flow.FlowEventAck
+	7, // [7:8] is the sub-list for method output_type
+	6, // [6:7] is the sub-list for method input_type
+	6, // [6:6] is the sub-list for extension type_name
+	6, // [6:6] is the sub-list for extension extendee
+	0, // [0:6] is the sub-list for field type_name
+}
+
+func init() { file_flow_proto_init() }
+func file_flow_proto_init() {
+	if File_flow_proto != nil {
+		return
+	}
+	if !protoimpl.UnsafeEnabled {
+		file_flow_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} {
+			switch v := v.(*FlowEvent); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
+		file_flow_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} {
+			switch v := v.(*FlowEventAck); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
+		file_flow_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} {
+			switch v := v.(*FlowFields); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
+		file_flow_proto_msgTypes[3].Exporter = func(v interface{}, i int) interface{} {
+			switch v := v.(*PortInfo); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
+		file_flow_proto_msgTypes[4].Exporter = func(v interface{}, i int) interface{} {
+			switch v := v.(*ICMPInfo); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
+	}
+	file_flow_proto_msgTypes[2].OneofWrappers = []interface{}{
+		(*FlowFields_PortInfo)(nil),
+		(*FlowFields_IcmpInfo)(nil),
+	}
+	type x struct{}
+	out := protoimpl.TypeBuilder{
+		File: protoimpl.DescBuilder{
+			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
+			RawDescriptor: file_flow_proto_rawDesc,
+			NumEnums:      2,
+			NumMessages:   5,
+			NumExtensions: 0,
+			NumServices:   1,
+		},
+		GoTypes:           file_flow_proto_goTypes,
+		DependencyIndexes: file_flow_proto_depIdxs,
+		EnumInfos:         file_flow_proto_enumTypes,
+		MessageInfos:      file_flow_proto_msgTypes,
+	}.Build()
+	File_flow_proto = out.File
+	file_flow_proto_rawDesc = nil
+	file_flow_proto_goTypes = nil
+	file_flow_proto_depIdxs = nil
+}

flow/proto/flow.proto

@@ -0,0 +1,102 @@
+syntax = "proto3";
+
+import "google/protobuf/timestamp.proto";
+
+option go_package = "/proto";
+
+package flow;
+
+service FlowService {
+  // Client to receiver streams of events and acknowledgements
+  rpc Events(stream FlowEvent) returns (stream FlowEventAck) {}
+}
+
+message FlowEvent {
+  // Unique client event identifier
+  bytes event_id = 1;
+
+  // When the event occurred
+  google.protobuf.Timestamp timestamp = 2;
+
+  // Public key of the sending peer
+  bytes public_key = 3;
+
+  FlowFields flow_fields = 4;
+}
+
+message FlowEventAck {
+  // Unique client event identifier that has been ack'ed
+  bytes event_id = 1;
+}
+
+message FlowFields {
+  // Unique client flow session identifier
+  bytes flow_id = 1;
+
+  // Flow type
+  Type type = 2;
+
+  // RuleId identifies the rule that allowed or denied the connection
+  bytes rule_id = 3;
+
+  // Initiating traffic direction
+  Direction direction = 4;
+
+  // IP protocol number
+  uint32 protocol = 5;
+
+  // Source IP address
+  bytes source_ip = 6;
+
+  // Destination IP address
+  bytes dest_ip = 7;
+
+  // Layer 4 -specific information
+  oneof connection_info {
+    // TCP/UDP port information
+    PortInfo port_info = 8;
+
+    // ICMP type and code
+    ICMPInfo icmp_info = 9;
+  }
+
+  // Number of packets
+  uint64 rx_packets = 10;
+  uint64 tx_packets = 11;
+
+  // Number of bytes
+  uint64 rx_bytes = 12;
+  uint64 tx_bytes = 13;
+
+  // Resource ID
+  bytes source_resource_id = 14;
+  bytes dest_resource_id = 15;
+
+}
+
+// Flow event types
+enum Type {
+  TYPE_UNKNOWN = 0;
+  TYPE_START = 1;
+  TYPE_END = 2;
+  TYPE_DROP = 3;
+}
+
+// Flow direction
+enum Direction {
+  DIRECTION_UNKNOWN = 0;
+  INGRESS = 1;
+  EGRESS = 2;
+}
+
+// TCP/UDP port information
+message PortInfo {
+  uint32 source_port = 1;
+  uint32 dest_port = 2;
+}
+
+// ICMP message information
+message ICMPInfo {
+  uint32 icmp_type = 1;
+  uint32 icmp_code = 2;
+}

flow/proto/flow_grpc.pb.go

@@ -0,0 +1,135 @@
+// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
+
+package proto
+
+import (
+	context "context"
+	grpc "google.golang.org/grpc"
+	codes "google.golang.org/grpc/codes"
+	status "google.golang.org/grpc/status"
+)
+
+// This is a compile-time assertion to ensure that this generated file
+// is compatible with the grpc package it is being compiled against.
+// Requires gRPC-Go v1.32.0 or later.
+const _ = grpc.SupportPackageIsVersion7
+
+// FlowServiceClient is the client API for FlowService service.
+//
+// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream.
+type FlowServiceClient interface {
+	// Client to receiver streams of events and acknowledgements
+	Events(ctx context.Context, opts ...grpc.CallOption) (FlowService_EventsClient, error)
+}
+
+type flowServiceClient struct {
+	cc grpc.ClientConnInterface
+}
+
+func NewFlowServiceClient(cc grpc.ClientConnInterface) FlowServiceClient {
+	return &flowServiceClient{cc}
+}
+
+func (c *flowServiceClient) Events(ctx context.Context, opts ...grpc.CallOption) (FlowService_EventsClient, error) {
+	stream, err := c.cc.NewStream(ctx, &FlowService_ServiceDesc.Streams[0], "/flow.FlowService/Events", opts...)
+	if err != nil {
+		return nil, err
+	}
+	x := &flowServiceEventsClient{stream}
+	return x, nil
+}
+
+type FlowService_EventsClient interface {
+	Send(*FlowEvent) error
+	Recv() (*FlowEventAck, error)
+	grpc.ClientStream
+}
+
+type flowServiceEventsClient struct {
+	grpc.ClientStream
+}
+
+func (x *flowServiceEventsClient) Send(m *FlowEvent) error {
+	return x.ClientStream.SendMsg(m)
+}
+
+func (x *flowServiceEventsClient) Recv() (*FlowEventAck, error) {
+	m := new(FlowEventAck)
+	if err := x.ClientStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
+// FlowServiceServer is the server API for FlowService service.
+// All implementations must embed UnimplementedFlowServiceServer
+// for forward compatibility
+type FlowServiceServer interface {
+	// Client to receiver streams of events and acknowledgements
+	Events(FlowService_EventsServer) error
+	mustEmbedUnimplementedFlowServiceServer()
+}
+
+// UnimplementedFlowServiceServer must be embedded to have forward compatible implementations.
+type UnimplementedFlowServiceServer struct {
+}
+
+func (UnimplementedFlowServiceServer) Events(FlowService_EventsServer) error {
+	return status.Errorf(codes.Unimplemented, "method Events not implemented")
+}
+func (UnimplementedFlowServiceServer) mustEmbedUnimplementedFlowServiceServer() {}
+
+// UnsafeFlowServiceServer may be embedded to opt out of forward compatibility for this service.
+// Use of this interface is not recommended, as added methods to FlowServiceServer will
+// result in compilation errors.
+type UnsafeFlowServiceServer interface {
+	mustEmbedUnimplementedFlowServiceServer()
+}
+
+func RegisterFlowServiceServer(s grpc.ServiceRegistrar, srv FlowServiceServer) {
+	s.RegisterService(&FlowService_ServiceDesc, srv)
+}
+
+func _FlowService_Events_Handler(srv interface{}, stream grpc.ServerStream) error {
+	return srv.(FlowServiceServer).Events(&flowServiceEventsServer{stream})
+}
+
+type FlowService_EventsServer interface {
+	Send(*FlowEventAck) error
+	Recv() (*FlowEvent, error)
+	grpc.ServerStream
+}
+
+type flowServiceEventsServer struct {
+	grpc.ServerStream
+}
+
+func (x *flowServiceEventsServer) Send(m *FlowEventAck) error {
+	return x.ServerStream.SendMsg(m)
+}
+
+func (x *flowServiceEventsServer) Recv() (*FlowEvent, error) {
+	m := new(FlowEvent)
+	if err := x.ServerStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
+// FlowService_ServiceDesc is the grpc.ServiceDesc for FlowService service.
+// It's only intended for direct use with grpc.RegisterService,
+// and not to be introspected or modified (even as a copy)
+var FlowService_ServiceDesc = grpc.ServiceDesc{
+	ServiceName: "flow.FlowService",
+	HandlerType: (*FlowServiceServer)(nil),
+	Methods:     []grpc.MethodDesc{},
+	Streams: []grpc.StreamDesc{
+		{
+			StreamName:    "Events",
+			Handler:       _FlowService_Events_Handler,
+			ServerStreams: true,
+			ClientStreams: true,
+		},
+	},
+	Metadata: "flow.proto",
+}

flow/proto/generate.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+set -e
+
+if ! which realpath > /dev/null 2>&1
+then
+  echo realpath is not installed
+  echo run: brew install coreutils
+  exit 1
+fi
+
+old_pwd=$(pwd)
+script_path=$(dirname $(realpath "$0"))
+cd "$script_path"
+go install google.golang.org/protobuf/cmd/protoc-gen-go@v1.26
+go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@v1.1
+protoc -I ./ ./flow.proto --go_out=../ --go-grpc_out=../
+cd "$old_pwd"