[misc] Update FreeBSD workflow (#3638)

Update FreeBSD release to 14.2 and download Go package directly since port wasn't finding the package to install

.github/workflows/golang-test-freebsd.yml

@@ -22,14 +22,20 @@ jobs:
         with:
           usesh: true
           copyback: false
-          release: "14.1"
+          release: "14.2"
           prepare: |
-            pkg install -y go pkgconf xorg
+            pkg install -y curl pkgconf xorg
+            LATEST_VERSION=$(curl -s https://go.dev/VERSION?m=text|head -n 1)
+            GO_TARBALL="$LATEST_VERSION.freebsd-amd64.tar.gz"
+            GO_URL="https://go.dev/dl/$GO_TARBALL"
+            curl -vLO "$GO_URL"
+            tar -C /usr/local -vxzf "$GO_TARBALL"            
 
           # -x - to print all executed commands
           # -e - to faile on first error
           run: |
             set -e -x
+            export PATH=$PATH:/usr/local/go/bin:$HOME/go/bin
             time go build -o netbird client/main.go
             # check all component except management, since we do not support management server on freebsd
             time go test -timeout 1m -failfast ./base62/...

.github/workflows/golangci-lint.yml

@@ -19,7 +19,7 @@ jobs:
       - name: codespell
         uses: codespell-project/actions-codespell@v2
         with:
-          ignore_words_list: erro,clienta,hastable,iif,groupd,testin
+          ignore_words_list: erro,clienta,hastable,iif,groupd,testin,groupe
           skip: go.mod,go.sum
           only_warn: 1
   golangci:

.github/workflows/release.yml

@@ -87,25 +87,25 @@ jobs:
         with:
           name: release
           path: dist/
-          retention-days: 3
+          retention-days: 7
       - name: upload linux packages
         uses: actions/upload-artifact@v4
         with:
           name: linux-packages
           path: dist/netbird_linux**
-          retention-days: 3
+          retention-days: 7
       - name: upload windows packages
         uses: actions/upload-artifact@v4
         with:
           name: windows-packages
           path: dist/netbird_windows**
-          retention-days: 3
+          retention-days: 7
       - name: upload macos packages
         uses: actions/upload-artifact@v4
         with:
           name: macos-packages
           path: dist/netbird_darwin**
-          retention-days: 3
+          retention-days: 7
 
   release_ui:
     runs-on: ubuntu-latest

CONTRIBUTOR_LICENSE_AGREEMENT.md

@@ -1,148 +1,64 @@
-#  Contributor License Agreement
+##  Contributor License Agreement
 
-We are incredibly thankful for the contributions we receive from the community.
-We require our external contributors to sign a Contributor License Agreement ("CLA") in
-order to ensure that our projects remain licensed under Free and Open Source licenses such
-as BSD-3 while allowing NetBird to build a sustainable business.
-
-NetBird is committed to having a true Open Source Software ("OSS") license for
-our software. A CLA enables NetBird to safely commercialize our products
-while keeping a standard OSS license with all the rights that license grants to users: the
-ability to use the project in their own projects or businesses, to republish modified
-source, or to completely fork the project.
-
-This page gives a human-friendly summary of our CLA, details on why we require a CLA, how
-contributors can sign our CLA, and more. You may view the full legal CLA document (below).
-
-# Human-friendly summary
-
-This is a human-readable summary of (and not a substitute for) the full agreement (below).
-This highlights only some of key terms of the CLA. It has no legal value and you should
-carefully review all the terms of the actual CLA before agreeing.
-
-<li>Grant of copyright license. You give NetBird permission to use your copyrighted work
-in commercial products.
-</li>
-
-<li>Grant of patent license. If your contributed work uses a patent, you give NetBird a
-license to use that patent including within commercial products. You also agree that you
-have permission to grant this license.
-</li>
-
-<li>No Warranty or Support Obligations.
-By making a contribution, you are not obligating yourself to provide support for the
-contribution, and you are not taking on any warranty obligations or providing any
-assurances about how it will perform.
-</li>
-
-The CLA does not change the terms of the standard open source license used by our software
-such as BSD-3 or MIT.
-You are still free to use our projects within your own projects or businesses, republish
-modified source, and more.
-Please reference the appropriate license for the project you're contributing to to learn
-more.
-
-# Why require a CLA?
-
-Agreeing to a CLA explicitly states that you are entitled to provide a contribution, that you cannot withdraw permission
-to use your contribution at a later date, and that NetBird has permission to use your contribution in our commercial
-products.
-
-This removes any ambiguities or uncertainties caused by not having a CLA and allows users and customers to confidently
-adopt our projects. At the same time, the CLA ensures that all contributions to our open source projects are licensed
-under the project's respective open source license, such as BSD-3.
-
-Requiring a CLA is a common and well-accepted practice in open source. Major open source projects require CLAs such as
-Apache Software Foundation projects, Facebook projects (such as React), Google projects (including Go), Python, Django,
-and more. Each of these projects remains licensed under permissive OSS licenses such as MIT, Apache, BSD, and more.
-
-# Signing the CLA
-
-Open a pull request ("PR") to any of our open source projects to sign the CLA. A bot will comment on the PR asking you
-to sign the CLA if you haven't already.
-
-Follow the steps given by the bot to sign the CLA. This will require you to log in with GitHub (we only request public
-information from your account) and to fill in a few additional details such as your name and email address. We will only
-use this information for CLA tracking; none of your submitted information will be used for marketing purposes.
-
-You only have to sign the CLA once. Once you've signed the CLA, future contributions to any NetBird project will not
-require you to sign again.
-
-# Legal Terms and Agreement
-
-In order to clarify the intellectual property license granted with Contributions from any person or entity, NetBird
-GmbH ("NetBird") must have a Contributor License Agreement ("CLA") on file that has been signed
-by each Contributor, indicating agreement to the license terms below. This license does not change your rights to use
-your own Contributions for any other purpose.
-
-You accept and agree to the following terms and conditions for Your present and future Contributions submitted to
-NetBird. Except for the license granted herein to NetBird and recipients of software distributed by NetBird,
-You reserve all right, title, and interest in and to Your Contributions.
-
-1. Definitions.
-
-    ```
-   "You" (or "Your") shall mean the copyright owner or legal entity authorized by the copyright owner 
-   that is making this Agreement with NetBird. For legal entities, the entity making a Contribution and all other 
-   entities that control, are controlled by, or are under common control with that entity are considered 
-   to be a single Contributor. For the purposes of this definition, "control" means (i) the power, direct or indirect,
-    to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty 
-   percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
-   ```
-   ```
-    "Contribution" shall mean any original work of authorship, including any modifications or additions to 
-   an existing work, that is or previously has been intentionally submitted by You to NetBird for inclusion in, 
-   or documentation of, any of the products owned or managed by NetBird (the "Work"). 
-   For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication 
-   sent to NetBird or its representatives, including but not limited to communication on electronic mailing lists, 
-   source code control systems, and issue tracking systems that are managed by, or on behalf of, 
-   NetBird for the purpose of discussing and improving the Work, but excluding communication that is conspicuously 
-   marked or otherwise designated in writing by You as "Not a Contribution."
-   ```
-
-2. Grant of Copyright License. Subject to the terms and conditions of this Agreement, You hereby grant to NetBird
-   and to recipients of software distributed by NetBird a perpetual, worldwide, non-exclusive, no-charge,
-   royalty-free, irrevocable copyright license to reproduce, prepare derivative works of, publicly display, publicly
-   perform, sublicense, and distribute Your Contributions and such derivative works.
+This Contributor License Agreement (referred to as the "Agreement") is entered into by the individual 
+submitting this Agreement and NetBird GmbH, c/o Max-Beer-Straße 2-4 Münzstraße 12 10178 Berlin, Germany, 
+referred to as "NetBird" (collectively, the "Parties"). The Agreement outlines the terms and conditions 
+under which NetBird may utilize software contributions provided by the Contributor for inclusion in 
+its software development projects. By submitting this Agreement, the Contributor confirms their acceptance 
+of the terms and conditions outlined below. The Contributor further represents that they are authorized to 
+complete this process as described herein.
 
 
-3. Grant of Patent License. Subject to the terms and conditions of this Agreement, You hereby grant to NetBird and
-   to recipients of software distributed by NetBird a perpetual, worldwide, non-exclusive, no-charge, royalty-free,
-   irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import,
-   and otherwise transfer the Work, where such license applies only to those patent claims licensable by You that are
-   necessarily infringed by Your Contribution(s) alone or by combination of Your Contribution(s) with the Work to which
-   such Contribution(s) was submitted. If any entity institutes patent litigation against You or any other entity (
-   including a cross-claim or counterclaim in a lawsuit) alleging that your Contribution, or the Work to which you have
-   contributed, constitutes direct or contributory patent infringement, then any patent licenses granted to that entity
-   under this Agreement for that Contribution or Work shall terminate as of the date such litigation is filed.
+## 1 Preamble
+In order to clarify the IP Rights situation with regard to Contributions from any person or entity, NetBird 
+must have a contributor license agreement on file to be signed by each Contributor, containing the license 
+terms below. This license serves as protection for both the Contributor as well as NetBird and its software users; 
+it does not change Contributor’s rights to use his/her own Contributions for any other purpose.
 
+## 2 Definitions
+2.1 “IP Rights” shall mean all industrial and intellectual property rights, whether registered or not registered, whether created by Contributor or acquired by Contributor from third parties, and similar rights, including (but not limited to) semiconductor property rights, design rights, copyrights (including in the form of database rights and rights to software), all neighbouring rights (Leistungsschutzrechte), trademarks, service marks, titles, internet domain names, trade names and other labelling rights, rights deriving from corresponding applications and registrations of such rights as well as any licenses (Nutzungsrechte) under and entitlements to any such intellectual and industrial property rights.
 
-4. You represent that you are legally entitled to grant the above license. If your employer(s) has rights to
-   intellectual property that you create that includes your Contributions, you represent that you have received
-   permission to make Contributions on behalf of that employer, that you will have received permission from your current
-   and future employers for all future Contributions, that your applicable employer has waived such rights for all of
-   your current and future Contributions to NetBird, or that your employer has executed a separate Corporate CLA
-   with NetBird.
+2.2 "Contribution" shall mean any original work of authorship, including any modifications or additions to an existing work, that is or previously has been intentionally Submitted by Contributor to NetBird for inclusion in, or documentation of any Work. 
 
+2.3 "Contributor" shall mean the copyright owner or legal entity authorized by the copyright owner that is concluding this Agreement with NetBird. For legal entities, the entity making a Contribution and all other entities that control, are controlled by, or are under common control with that entity are considered to be a single Contributor. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
 
-5. You represent that each of Your Contributions is Your original creation (see section 7 for submissions on behalf of
-   others). You represent that Your Contribution submissions include complete details of any third-party license or
-   other restriction (including, but not limited to, related patents and trademarks) of which you are personally aware
-   and which are associated with any part of Your Contributions.
+2.4 "Submitted" shall mean any form of electronic, verbal, or written communication sent to NetBird or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, NetBird for the purpose of discussing and improving the Work, but excluding communication that is marked or otherwise designated in writing by Contributor as "Not a Contribution".
 
+2.5 "Work" means any of the products owned or managed by NetBird, in particular, but not exclusively, software.
 
-6. You are not expected to provide support for Your Contributions, except to the extent You desire to provide support.
-   You may provide support for free, for a fee, or not at all. Unless required by applicable law or agreed to in
-   writing, You provide Your Contributions on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
-   express or implied, including, without limitation, any warranties or conditions of TITLE, NON- INFRINGEMENT,
-   MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE.
+## 3 Licenses 
+3.1 Subject to the terms and conditions of this agreement, Contributor hereby grants to NetBird and to recipients of software distributed by NetBird a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable license to reproduce by any means and in any form, in whole or in part, permanently or temporarily, the Contributions (including loading, displaying, executing, transmitting or storing works for the purpose of executing and processing data or transferring them to video, audio and other data carriers), including the right to distribute, display and present such Contributions and make them available to the public (e.g. via the internet) and to transmit and display such Contributions by any means. The license also includes the right to modify, translate, adapt, edit and otherwise alter the Contributions and to use these results in the same manner as the original Contributions and derivative works. Except for licenses in patents acc. to Sec. 3, such license refers to any IP Rights in the Contributions and derivative works. The Contributor acknowledges that NetBird is not required to credit them by name for their Contribution and agrees to waive any moral rights associated with their Contribution in relation to NetBird or its sublicensees.
 
+3.2 Subject to the terms and conditions of this agreement, Contributor hereby grants to NetBird and to recipients of software distributed by NetBird a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license in the Contributions to make, have made, use, sell, offer to sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by the Contributor which are necessarily infringed by Contributor‘s Contribution(s) alone or by combination of Contributor’s Contribution(s) with the Work to which such Contribution(s) was Submitted. 
 
-7. Should You wish to submit work that is not Your original creation, You may submit it to NetBird separately from
-   any Contribution, identifying the complete details of its source and of any license or other restriction (including,
-   but not limited to, related patents, trademarks, and license agreements) of which you are personally aware, and
-   conspicuously marking the work as "Submitted on behalf of a third-party: [named here]".
+3.3 NetBird hereby accepts such licenses. 
 
+## 4 Contributor’s Representations
+4.1 Contributor represents that Contributor is legally entitled to grant the above license. If Contributor’s employer has IP Rights to Contributor’s Contributions, Contributor represent that he/she has received permission to make Contributions on behalf of such employer, that such employer has waived such IP Rights to the Contributions of Contributor to NetBird, or that such employer has executed a separate contributor license agreement with NetBird.
+
+4.2 Contributor represents that any Contribution is his/her original creation. 
+
+4.3 Contributor represents to his/her best knowledge that any Contribution does not violate any third party IP Rights.
+
+4.4 Contributor represents that any Contribution submission includes complete details of any third-party license or other restriction (including, but not limited to, related patents and trademarks) of which Contributor is personally aware and which are associated with any part of the Contribution.
+
+4.5 The Contributor represents that their Contribution does not include any work distributed under a copyleft license.
+
+## 5 Information obligation
+Contributor agrees to notify NetBird of any facts or circumstances of which Contributor become aware that would make these representations inaccurate in any respect.
+
+## 6 Submission of Third-Party works
+Should Contributor wish to submit work that is not Contributor’s original creation, Contributor may submit it to NetBird separately from any Contribution, identifying the complete details of its source and of any license or other restriction (including, but not limited to, related patents, trademarks, and license agreements) of which Contributor are personally aware, and conspicuously marking the work as "Submitted on behalf of a third-party: [named here]".
+
+## 7 No Consideration
+Unless compensation is mandatory under statutory law, no compensation for any license under this agreement shall be payable. 
+
+## 8 Final Provisions
+8.1 Laws. This Agreement is governed by the laws of the Federal Republic of Germany. 
+
+8.2 Venue. Place of jurisdiction shall, to the extent legally permissible, be Berlin, Germany. 
+
+8.3 Severability. If any provision in this agreement is unlawful, invalid or ineffective, it shall not affect the enforceability or effectiveness of the remainder of this agreement. The parties agree to replace any unlawful, invalid or ineffective provision with a provision that comes as close as possible to the commercial intent and purpose of the original provision. This section also applies accordingly to any gaps in the contract. 
+
+8.4 Variations. Any variations, amendments or supplements to this Agreement must be in writing. This also applies to any variation of this Section 8.4.
 
-8. You agree to notify NetBird of any facts or circumstances of which you become aware that would make these
-   representations inaccurate in any respect.

README.md

@@ -12,7 +12,7 @@
        <img src="https://img.shields.io/badge/license-BSD--3-blue" />
      </a> 
     <br>
-    <a href="https://join.slack.com/t/netbirdio/shared_invite/zt-2utg2ncdz-W7LEB6toRBLE1Jca37dYpg">
+    <a href="https://join.slack.com/t/netbirdio/shared_invite/zt-31rofwmxc-27akKd0Le0vyRpBcwXkP0g">
         <img src="https://img.shields.io/badge/slack-@netbird-red.svg?logo=slack"/>
      </a>  
      <br>
@@ -29,13 +29,13 @@
   <br/>
   See <a href="https://netbird.io/docs/">Documentation</a>
   <br/>
-   Join our <a href="https://join.slack.com/t/netbirdio/shared_invite/zt-2utg2ncdz-W7LEB6toRBLE1Jca37dYpg">Slack channel</a>
+   Join our <a href="https://join.slack.com/t/netbirdio/shared_invite/zt-31rofwmxc-27akKd0Le0vyRpBcwXkP0g">Slack channel</a>
   <br/>
  
 </strong>
 <br>
-<a href="https://netbird.io/webinars/achieve-zero-trust-access-to-k8s?utm_source=github&utm_campaign=2502%20-%20webinar%20-%20How%20to%20Achieve%20Zero%20Trust%20Access%20to%20Kubernetes%20-%20Effortlessly&utm_medium=github">
-    Webinar: Securely Access Kubernetes without Port Forwarding and Jump Hosts
+<a href="https://github.com/netbirdio/kubernetes-operator">
+    New: NetBird Kubernetes Operator
   </a> 
 </p>
 

client/anonymize/anonymize.go

@@ -26,7 +26,7 @@ type Anonymizer struct {
 }
 
 func DefaultAddresses() (netip.Addr, netip.Addr) {
-	// 192.51.100.0, 100::
+	// 198.51.100.0, 100::
 	return netip.AddrFrom4([4]byte{198, 51, 100, 0}), netip.AddrFrom16([16]byte{0x01})
 }
 

client/cmd/login.go

@@ -19,6 +19,10 @@ import (
 	"github.com/netbirdio/netbird/util"
 )
 
+func init() {
+	loginCmd.PersistentFlags().BoolVar(&noBrowser, noBrowserFlag, false, noBrowserDesc)
+}
+
 var loginCmd = &cobra.Command{
 	Use:   "login",
 	Short: "login to the Netbird Management Service (first run)",
@@ -127,7 +131,7 @@ var loginCmd = &cobra.Command{
 		}
 
 		if loginResp.NeedsSSOLogin {
-			openURL(cmd, loginResp.VerificationURIComplete, loginResp.UserCode)
+			openURL(cmd, loginResp.VerificationURIComplete, loginResp.UserCode, noBrowser)
 
 			_, err = client.WaitSSOLogin(ctx, &proto.WaitSSOLoginRequest{UserCode: loginResp.UserCode, Hostname: hostName})
 			if err != nil {
@@ -198,7 +202,7 @@ func foregroundGetTokenInfo(ctx context.Context, cmd *cobra.Command, config *int
 		return nil, fmt.Errorf("getting a request OAuth flow info failed: %v", err)
 	}
 
-	openURL(cmd, flowInfo.VerificationURIComplete, flowInfo.UserCode)
+	openURL(cmd, flowInfo.VerificationURIComplete, flowInfo.UserCode, noBrowser)
 
 	waitTimeout := time.Duration(flowInfo.ExpiresIn) * time.Second
 	waitCTX, c := context.WithTimeout(context.TODO(), waitTimeout)
@@ -212,19 +216,27 @@ func foregroundGetTokenInfo(ctx context.Context, cmd *cobra.Command, config *int
 	return &tokenInfo, nil
 }
 
-func openURL(cmd *cobra.Command, verificationURIComplete, userCode string) {
+func openURL(cmd *cobra.Command, verificationURIComplete, userCode string, noBrowser bool) {
 	var codeMsg string
 	if userCode != "" && !strings.Contains(verificationURIComplete, userCode) {
 		codeMsg = fmt.Sprintf("and enter the code %s to authenticate.", userCode)
 	}
 
-	cmd.Println("Please do the SSO login in your browser. \n" +
-		"If your browser didn't open automatically, use this URL to log in:\n\n" +
-		verificationURIComplete + " " + codeMsg)
+	if noBrowser {
+		cmd.Println("Use this URL to log in:\n\n" + verificationURIComplete + " " + codeMsg)
+	} else {
+		cmd.Println("Please do the SSO login in your browser. \n" +
+			"If your browser didn't open automatically, use this URL to log in:\n\n" +
+			verificationURIComplete + " " + codeMsg)
+	}
+
 	cmd.Println("")
-	if err := open.Run(verificationURIComplete); err != nil {
-		cmd.Println("\nAlternatively, you may want to use a setup key, see:\n\n" +
-			"https://docs.netbird.io/how-to/register-machines-using-setup-keys")
+
+	if !noBrowser {
+		if err := open.Run(verificationURIComplete); err != nil {
+			cmd.Println("\nAlternatively, you may want to use a setup key, see:\n\n" +
+				"https://docs.netbird.io/how-to/register-machines-using-setup-keys")
+		}
 	}
 }
 

client/cmd/root.go

@@ -180,7 +180,7 @@ func init() {
 	upCmd.PersistentFlags().BoolVar(&serverSSHAllowed, serverSSHAllowedFlag, false, "Allow SSH server on peer. If enabled, the SSH server will be permitted")
 	upCmd.PersistentFlags().BoolVar(&autoConnectDisabled, disableAutoConnectFlag, false, "Disables auto-connect feature. If enabled, then the client won't connect automatically when the service starts.")
 
-	debugCmd.PersistentFlags().BoolVarP(&debugSystemInfoFlag, systemInfoFlag, "S", false, "Adds system information to the debug bundle")
+	debugCmd.PersistentFlags().BoolVarP(&debugSystemInfoFlag, systemInfoFlag, "S", true, "Adds system information to the debug bundle")
 }
 
 // SetupCloseHandler handles SIGTERM signal and exits with success

client/cmd/testutil_test.go

@@ -6,14 +6,17 @@ import (
 	"testing"
 	"time"
 
+	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel"
 
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
+	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/settings"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/telemetry"
+	"github.com/netbirdio/netbird/management/server/types"
 
 	"github.com/netbirdio/netbird/util"
 
@@ -31,7 +34,7 @@ import (
 
 func startTestingServices(t *testing.T) string {
 	t.Helper()
-	config := &mgmt.Config{}
+	config := &types.Config{}
 	_, err := util.ReadJson("../testdata/management.json", config)
 	if err != nil {
 		t.Fatal(err)
@@ -66,7 +69,7 @@ func startSignal(t *testing.T) (*grpc.Server, net.Listener) {
 	return s, lis
 }
 
-func startManagement(t *testing.T, config *mgmt.Config, testFile string) (*grpc.Server, net.Listener) {
+func startManagement(t *testing.T, config *types.Config, testFile string) (*grpc.Server, net.Listener) {
 	t.Helper()
 
 	lis, err := net.Listen("tcp", ":0")
@@ -89,14 +92,19 @@ func startManagement(t *testing.T, config *mgmt.Config, testFile string) (*grpc.
 
 	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
 	require.NoError(t, err)
+	permissionsManagerMock := permissions.NewManagerMock()
+	ctrl := gomock.NewController(t)
+	t.Cleanup(ctrl.Finish)
 
-	accountManager, err := mgmt.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, iv, metrics, port_forwarding.NewControllerMock())
+	settingsMockManager := settings.NewMockManager(ctrl)
+
+	accountManager, err := mgmt.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, iv, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManagerMock)
 	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)
+	secretsManager := mgmt.NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig, config.Relay, settingsMockManager)
+	mgmtServer, err := mgmt.NewServer(context.Background(), config, accountManager, settingsMockManager, peersUpdateManager, secretsManager, nil, nil, nil)
 	if err != nil {
 		t.Fatal(err)
 	}

client/cmd/up.go

@@ -32,12 +32,16 @@ const (
 
 const (
 	dnsLabelsFlag = "extra-dns-labels"
+
+	noBrowserFlag = "no-browser"
+	noBrowserDesc = "do not open the browser for SSO login"
 )
 
 var (
 	foregroundMode     bool
 	dnsLabels          []string
 	dnsLabelsValidated domain.List
+	noBrowser          bool
 
 	upCmd = &cobra.Command{
 		Use:   "up",
@@ -65,6 +69,9 @@ func init() {
 			`E.g. --extra-dns-labels vpc1 or --extra-dns-labels vpc1,mgmt1 `+
 			`or --extra-dns-labels ""`,
 	)
+
+	upCmd.PersistentFlags().BoolVar(&noBrowser, noBrowserFlag, false, noBrowserDesc)
+
 }
 
 func upFunc(cmd *cobra.Command, args []string) error {
@@ -349,7 +356,7 @@ func runInDaemonMode(ctx context.Context, cmd *cobra.Command) error {
 
 	if loginResp.NeedsSSOLogin {
 
-		openURL(cmd, loginResp.VerificationURIComplete, loginResp.UserCode)
+		openURL(cmd, loginResp.VerificationURIComplete, loginResp.UserCode, noBrowser)
 
 		_, err = client.WaitSSOLogin(ctx, &proto.WaitSSOLoginRequest{UserCode: loginResp.UserCode, Hostname: hostName})
 		if err != nil {

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,7 +48,7 @@ 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) {
@@ -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

@@ -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

@@ -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
@@ -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"
@@ -38,10 +38,12 @@ const (
 	routingFinalForwardJump = "ACCEPT"
 	routingFinalNatJump     = "MASQUERADE"
 
-	jumpManglePre = "jump-mangle-pre"
-	jumpNatPre    = "jump-nat-pre"
-	jumpNatPost   = "jump-nat-post"
-	matchSet      = "--match-set"
+	jumpManglePre  = "jump-mangle-pre"
+	jumpNatPre     = "jump-nat-pre"
+	jumpNatPost    = "jump-nat-post"
+	markManglePre  = "mark-mangle-pre"
+	markManglePost = "mark-mangle-post"
+	matchSet       = "--match-set"
 
 	dnatSuffix = "_dnat"
 	snatSuffix = "_snat"
@@ -115,12 +117,17 @@ func (r *router) init(stateManager *statemanager.Manager) error {
 		return fmt.Errorf("create containers: %w", err)
 	}
 
+	if err := r.setupDataPlaneMark(); err != nil {
+		log.Errorf("failed to set up data plane mark: %v", err)
+	}
+
 	r.updateState()
 
 	return nil
 }
 
 func (r *router) AddRouteFiltering(
+	id []byte,
 	sources []netip.Prefix,
 	destination netip.Prefix,
 	proto firewall.Protocol,
@@ -128,7 +135,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
 	}
@@ -347,12 +354,16 @@ func (r *router) Reset() error {
 	if err := r.cleanUpDefaultForwardRules(); err != nil {
 		merr = multierror.Append(merr, err)
 	}
-	r.rules = make(map[string][]string)
 
 	if err := r.ipsetCounter.Flush(); err != nil {
 		merr = multierror.Append(merr, err)
 	}
 
+	if err := r.cleanupDataPlaneMark(); err != nil {
+		merr = multierror.Append(merr, err)
+	}
+
+	r.rules = make(map[string][]string)
 	r.updateState()
 
 	return nberrors.FormatErrorOrNil(merr)
@@ -422,6 +433,57 @@ func (r *router) createContainers() error {
 	return nil
 }
 
+// setupDataPlaneMark configures the fwmark for the data plane
+func (r *router) setupDataPlaneMark() error {
+	var merr *multierror.Error
+	preRule := []string{
+		"-i", r.wgIface.Name(),
+		"-m", "conntrack", "--ctstate", "NEW",
+		"-j", "CONNMARK", "--set-mark", fmt.Sprintf("%#x", nbnet.DataPlaneMarkIn),
+	}
+
+	if err := r.iptablesClient.AppendUnique(tableMangle, chainPREROUTING, preRule...); err != nil {
+		merr = multierror.Append(merr, fmt.Errorf("add mangle prerouting rule: %w", err))
+	} else {
+		r.rules[markManglePre] = preRule
+	}
+
+	postRule := []string{
+		"-o", r.wgIface.Name(),
+		"-m", "conntrack", "--ctstate", "NEW",
+		"-j", "CONNMARK", "--set-mark", fmt.Sprintf("%#x", nbnet.DataPlaneMarkOut),
+	}
+
+	if err := r.iptablesClient.AppendUnique(tableMangle, chainPOSTROUTING, postRule...); err != nil {
+		merr = multierror.Append(merr, fmt.Errorf("add mangle postrouting rule: %w", err))
+	} else {
+		r.rules[markManglePost] = postRule
+	}
+
+	return nberrors.FormatErrorOrNil(merr)
+}
+
+func (r *router) cleanupDataPlaneMark() error {
+	var merr *multierror.Error
+	if preRule, exists := r.rules[markManglePre]; exists {
+		if err := r.iptablesClient.DeleteIfExists(tableMangle, chainPREROUTING, preRule...); err != nil {
+			merr = multierror.Append(merr, fmt.Errorf("remove mangle prerouting rule: %w", err))
+		} else {
+			delete(r.rules, markManglePre)
+		}
+	}
+
+	if postRule, exists := r.rules[markManglePost]; exists {
+		if err := r.iptablesClient.DeleteIfExists(tableMangle, chainPOSTROUTING, postRule...); err != nil {
+			merr = multierror.Append(merr, fmt.Errorf("remove mangle postrouting rule: %w", err))
+		} else {
+			delete(r.rules, markManglePost)
+		}
+	}
+
+	return nberrors.FormatErrorOrNil(merr)
+}
+
 func (r *router) addPostroutingRules() error {
 	// First rule for outbound masquerade
 	rule1 := []string{
@@ -463,7 +525,7 @@ func (r *router) insertEstablishedRule(chain string) error {
 }
 
 func (r *router) addJumpRules() error {
-	// Jump to NAT chain
+	// Jump to nat chain
 	natRule := []string{"-j", chainRTNAT}
 	if err := r.iptablesClient.Insert(tableNat, chainPOSTROUTING, 1, natRule...); err != nil {
 		return fmt.Errorf("add nat postrouting jump rule: %v", err)

client/firewall/iptables/router_linux_test.go

@@ -46,7 +46,9 @@ func TestIptablesManager_RestoreOrCreateContainers(t *testing.T) {
 	// 5. jump rule to PRE nat chain
 	// 6. static outbound masquerade rule
 	// 7. static return masquerade rule
-	require.Len(t, manager.rules, 7, "should have created rules map")
+	// 8. mangle prerouting mark rule
+	// 9. mangle postrouting mark rule
+	require.Len(t, manager.rules, 9, "should have created rules map")
 
 	exists, err := manager.iptablesClient.Exists(tableNat, chainPOSTROUTING, "-j", chainRTNAT)
 	require.NoError(t, err, "should be able to query the iptables %s table and %s chain", tableNat, chainPOSTROUTING)
@@ -330,7 +332,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

@@ -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

client/firewall/nftables/acl_linux.go

@@ -25,9 +25,10 @@ const (
 	chainNameInputRules = "netbird-acl-input-rules"
 
 	// filter chains contains the rules that jump to the rules chains
-	chainNameInputFilter   = "netbird-acl-input-filter"
-	chainNameForwardFilter = "netbird-acl-forward-filter"
-	chainNamePrerouting    = "netbird-rt-prerouting"
+	chainNameInputFilter       = "netbird-acl-input-filter"
+	chainNameForwardFilter     = "netbird-acl-forward-filter"
+	chainNameManglePrerouting  = "netbird-mangle-prerouting"
+	chainNameManglePostrouting = "netbird-mangle-postrouting"
 
 	allowNetbirdInputRuleID = "allow Netbird incoming traffic"
 )
@@ -84,13 +85,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 +103,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 +257,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 +338,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{
@@ -463,13 +463,15 @@ func (m *AclManager) createDefaultChains() (err error) {
 // go through the input filter as well. This will enable e.g. Docker services to keep working by accessing the
 // netbird peer IP.
 func (m *AclManager) allowRedirectedTraffic(chainFwFilter *nftables.Chain) error {
-	m.chainPrerouting = m.rConn.AddChain(&nftables.Chain{
-		Name:     chainNamePrerouting,
+	// Chain is created by route manager
+	// TODO: move creation to a common place
+	m.chainPrerouting = &nftables.Chain{
+		Name:     chainNameManglePrerouting,
 		Table:    m.workTable,
 		Type:     nftables.ChainTypeFilter,
 		Hooknum:  nftables.ChainHookPrerouting,
 		Priority: nftables.ChainPriorityMangle,
-	})
+	}
 
 	m.addFwmarkToForward(chainFwFilter)
 

client/firewall/nftables/manager_linux.go

@@ -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

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"
@@ -100,6 +100,10 @@ func (r *router) init(workTable *nftables.Table) error {
 		return fmt.Errorf("create containers: %w", err)
 	}
 
+	if err := r.setupDataPlaneMark(); err != nil {
+		log.Errorf("failed to set up data plane mark: %v", err)
+	}
+
 	return nil
 }
 
@@ -196,15 +200,21 @@ func (r *router) createContainers() error {
 		Type:     nftables.ChainTypeNAT,
 	})
 
-	// Chain is created by acl manager
-	// TODO: move creation to a common place
-	r.chains[chainNamePrerouting] = &nftables.Chain{
-		Name:     chainNamePrerouting,
+	r.chains[chainNameManglePostrouting] = r.conn.AddChain(&nftables.Chain{
+		Name:     chainNameManglePostrouting,
+		Table:    r.workTable,
+		Hooknum:  nftables.ChainHookPostrouting,
+		Priority: nftables.ChainPriorityMangle,
+		Type:     nftables.ChainTypeFilter,
+	})
+
+	r.chains[chainNameManglePrerouting] = r.conn.AddChain(&nftables.Chain{
+		Name:     chainNameManglePrerouting,
 		Table:    r.workTable,
 		Hooknum:  nftables.ChainHookPrerouting,
 		Priority: nftables.ChainPriorityMangle,
 		Type:     nftables.ChainTypeFilter,
-	}
+	})
 
 	// Add the single NAT rule that matches on mark
 	if err := r.addPostroutingRules(); err != nil {
@@ -220,7 +230,83 @@ func (r *router) createContainers() error {
 	}
 
 	if err := r.conn.Flush(); err != nil {
-		return fmt.Errorf("nftables: unable to initialize table: %v", err)
+		return fmt.Errorf("initialize tables: %v", err)
+	}
+
+	return nil
+}
+
+// setupDataPlaneMark configures the fwmark for the data plane
+func (r *router) setupDataPlaneMark() error {
+	if r.chains[chainNameManglePrerouting] == nil || r.chains[chainNameManglePostrouting] == nil {
+		return errors.New("no mangle chains found")
+	}
+
+	ctNew := getCtNewExprs()
+	preExprs := []expr.Any{
+		&expr.Meta{
+			Key:      expr.MetaKeyIIFNAME,
+			Register: 1,
+		},
+		&expr.Cmp{
+			Op:       expr.CmpOpEq,
+			Register: 1,
+			Data:     ifname(r.wgIface.Name()),
+		},
+	}
+	preExprs = append(preExprs, ctNew...)
+	preExprs = append(preExprs,
+		&expr.Immediate{
+			Register: 1,
+			Data:     binaryutil.NativeEndian.PutUint32(nbnet.DataPlaneMarkIn),
+		},
+		&expr.Ct{
+			Key:            expr.CtKeyMARK,
+			Register:       1,
+			SourceRegister: true,
+		},
+	)
+
+	preNftRule := &nftables.Rule{
+		Table: r.workTable,
+		Chain: r.chains[chainNameManglePrerouting],
+		Exprs: preExprs,
+	}
+	r.conn.AddRule(preNftRule)
+
+	postExprs := []expr.Any{
+		&expr.Meta{
+			Key:      expr.MetaKeyOIFNAME,
+			Register: 1,
+		},
+		&expr.Cmp{
+			Op:       expr.CmpOpEq,
+			Register: 1,
+			Data:     ifname(r.wgIface.Name()),
+		},
+	}
+	postExprs = append(postExprs, ctNew...)
+	postExprs = append(postExprs,
+		&expr.Immediate{
+			Register: 1,
+			Data:     binaryutil.NativeEndian.PutUint32(nbnet.DataPlaneMarkOut),
+		},
+		&expr.Ct{
+			Key:            expr.CtKeyMARK,
+			Register:       1,
+			SourceRegister: true,
+		},
+	)
+
+	postNftRule := &nftables.Rule{
+		Table: r.workTable,
+		Chain: r.chains[chainNameManglePostrouting],
+		Exprs: postExprs,
+	}
+	r.conn.AddRule(postNftRule)
+
+	if err := r.conn.Flush(); err != nil {
+		return fmt.Errorf("flush: %w", err)
 	}
 
 	return nil
@@ -228,6 +314,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 +323,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
 	}
@@ -515,26 +602,10 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 		op = expr.CmpOpNeq
 	}
 
-	exprs := []expr.Any{
-		// We only care about NEW connections to mark them and later identify them in the postrouting chain for masquerading.
-		// Masquerading will take care of the conntrack state, which means we won't need to mark established connections.
-		&expr.Ct{
-			Key:      expr.CtKeySTATE,
-			Register: 1,
-		},
-		&expr.Bitwise{
-			SourceRegister: 1,
-			DestRegister:   1,
-			Len:            4,
-			Mask:           binaryutil.NativeEndian.PutUint32(expr.CtStateBitNEW),
-			Xor:            binaryutil.NativeEndian.PutUint32(0),
-		},
-		&expr.Cmp{
-			Op:       expr.CmpOpNeq,
-			Register: 1,
-			Data:     []byte{0, 0, 0, 0},
-		},
-
+	// We only care about NEW connections to mark them and later identify them in the postrouting chain for masquerading.
+	// Masquerading will take care of the conntrack state, which means we won't need to mark established connections.
+	exprs := getCtNewExprs()
+	exprs = append(exprs,
 		// interface matching
 		&expr.Meta{
 			Key:      expr.MetaKeyIIFNAME,
@@ -545,7 +616,7 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 			Register: 1,
 			Data:     ifname(r.wgIface.Name()),
 		},
-	}
+	)
 
 	exprs = append(exprs, sourceExp...)
 	exprs = append(exprs, destExp...)
@@ -577,7 +648,7 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 
 	r.rules[ruleKey] = r.conn.AddRule(&nftables.Rule{
 		Table:    r.workTable,
-		Chain:    r.chains[chainNamePrerouting],
+		Chain:    r.chains[chainNameManglePrerouting],
 		Exprs:    exprs,
 		UserData: []byte(ruleKey),
 	})
@@ -1323,3 +1394,24 @@ func applyPort(port *firewall.Port, isSource bool) []expr.Any {
 
 	return exprs
 }
+
+func getCtNewExprs() []expr.Any {
+	return []expr.Any{
+		&expr.Ct{
+			Key:      expr.CtKeySTATE,
+			Register: 1,
+		},
+		&expr.Bitwise{
+			SourceRegister: 1,
+			DestRegister:   1,
+			Len:            4,
+			Mask:           binaryutil.NativeEndian.PutUint32(expr.CtStateBitNEW),
+			Xor:            binaryutil.NativeEndian.PutUint32(0),
+		},
+		&expr.Cmp{
+			Op:       expr.CmpOpNeq,
+			Register: 1,
+			Data:     []byte{0, 0, 0, 0},
+		},
+	}
+}

client/firewall/nftables/router_linux_test.go

@@ -100,7 +100,7 @@ func TestNftablesManager_AddNatRule(t *testing.T) {
 				natRuleKey := firewall.GenKey(firewall.PreroutingFormat, testCase.InputPair)
 				found := 0
 				for _, chain := range rtr.chains {
-					if chain.Name == chainNamePrerouting {
+					if chain.Name == chainNameManglePrerouting {
 						rules, err := nftablesTestingClient.GetRules(chain.Table, chain)
 						require.NoError(t, err, "should list rules for %s table and %s chain", chain.Table.Name, chain.Name)
 						for _, rule := range rules {
@@ -141,7 +141,7 @@ func TestNftablesManager_RemoveNatRule(t *testing.T) {
 			// Verify the rule was added
 			natRuleKey := firewall.GenKey(firewall.PreroutingFormat, testCase.InputPair)
 			found := false
-			rules, err := rtr.conn.GetRules(rtr.workTable, rtr.chains[chainNamePrerouting])
+			rules, err := rtr.conn.GetRules(rtr.workTable, rtr.chains[chainNameManglePrerouting])
 			require.NoError(t, err, "should list rules")
 			for _, rule := range rules {
 				if len(rule.UserData) > 0 && string(rule.UserData) == natRuleKey {
@@ -157,7 +157,7 @@ func TestNftablesManager_RemoveNatRule(t *testing.T) {
 
 			// Verify the rule was removed
 			found = false
-			rules, err = rtr.conn.GetRules(rtr.workTable, rtr.chains[chainNamePrerouting])
+			rules, err = rtr.conn.GetRules(rtr.workTable, rtr.chains[chainNameManglePrerouting])
 			require.NoError(t, err, "should list rules after removal")
 			for _, rule := range rules {
 				if len(rule.UserData) > 0 && string(rule.UserData) == natRuleKey {
@@ -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"
@@ -16,8 +17,8 @@ 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
+// 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,66 @@
 package conntrack
 
 import (
-	"net"
+	"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(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,108 @@ 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()
+	conn.UpdateCounters(direction, size)
+
+	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 +180,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.Trace("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 +201,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,39 @@
 package conntrack
 
 import (
-	"net"
+	"net/netip"
 	"testing"
 )
 
 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 (
@@ -20,11 +23,11 @@ const (
 )
 
 const (
-	TCPSyn  uint8 = 0x02
-	TCPAck  uint8 = 0x10
 	TCPFin  uint8 = 0x01
+	TCPSyn  uint8 = 0x02
 	TCPRst  uint8 = 0x04
 	TCPPush uint8 = 0x08
+	TCPAck  uint8 = 0x10
 	TCPUrg  uint8 = 0x20
 )
 
@@ -38,7 +41,36 @@ const (
 )
 
 // TCPState represents the state of a TCP connection
-type TCPState int
+type TCPState int32
+
+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
@@ -54,30 +86,38 @@ 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
-	State       TCPState
-	established atomic.Bool
-	sync.RWMutex
+	SourcePort uint16
+	DestPort   uint16
+	state      atomic.Int32
+	tombstone  atomic.Bool
 }
 
-// IsEstablished safely checks if connection is established
-func (t *TCPConnTrack) IsEstablished() bool {
-	return t.established.Load()
+// GetState safely retrieves the current state
+func (t *TCPConnTrack) GetState() TCPState {
+	return TCPState(t.state.Load())
 }
 
-// SetEstablished safely sets the established state
-func (t *TCPConnTrack) SetEstablished(state bool) {
-	t.established.Store(state)
+// SetState safely updates the current state
+func (t *TCPConnTrack) SetState(state TCPState) {
+	t.state.Store(int32(state))
+}
+
+// CompareAndSwapState atomically changes the state from old to new if current == old
+func (t *TCPConnTrack) CompareAndSwapState(old, newState TCPState) bool {
+	return t.state.CompareAndSwap(int32(old), int32(newState))
+}
+
+// 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
@@ -88,11 +128,18 @@ type TCPTracker struct {
 	cleanupTicker *time.Ticker
 	tickerCancel  context.CancelFunc
 	timeout       time.Duration
-	ipPool        *PreallocatedIPs
+	waitTimeout   time.Duration
+	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 {
+	waitTimeout := TimeWaitTimeout
+	if timeout == 0 {
+		timeout = DefaultTCPTimeout
+	} else {
+		waitTimeout = timeout / 45
+	}
 
 	ctx, cancel := context.WithCancel(context.Background())
 
@@ -102,179 +149,211 @@ func NewTCPTracker(timeout time.Duration, logger *nblog.Logger) *TCPTracker {
 		cleanupTicker: time.NewTicker(TCPCleanupInterval),
 		tickerCancel:  cancel,
 		timeout:       timeout,
-		ipPool:        NewPreallocatedIPs(),
+		waitTimeout:   waitTimeout,
+		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)
-
-	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)
+func (t *TCPTracker) updateIfExists(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, direction nftypes.Direction, size int) (ConnKey, bool) {
+	key := ConnKey{
+		SrcIP:   srcIP,
+		DstIP:   dstIP,
+		SrcPort: srcPort,
+		DstPort: dstPort,
 	}
-	t.mutex.Unlock()
-
-	// Lock individual connection for state update
-	conn.Lock()
-	t.updateState(conn, flags, true)
-	conn.Unlock()
-	conn.UpdateLastSeen()
-}
-
-// 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
-	}
-
-	key := makeConnKey(dstIP, srcIP, dstPort, srcPort)
 
 	t.mutex.RLock()
 	conn, exists := t.connections[key]
 	t.mutex.RUnlock()
 
-	if !exists {
-		return false
+	if exists {
+		t.updateState(key, conn, flags, direction, size)
+		return key, true
 	}
 
-	// Handle RST packets
-	if flags&TCPRst != 0 {
-		conn.Lock()
-		if conn.IsEstablished() || conn.State == TCPStateSynSent || conn.State == TCPStateSynReceived {
-			conn.State = TCPStateClosed
-			conn.SetEstablished(false)
-			conn.Unlock()
-			return true
-		}
-		conn.Unlock()
-		return false
-	}
-
-	conn.Lock()
-	t.updateState(conn, flags, false)
-	conn.UpdateLastSeen()
-	isEstablished := conn.IsEstablished()
-	isValidState := t.isValidStateForFlags(conn.State, flags)
-	conn.Unlock()
-
-	return isEstablished || isValidState
+	return key, false
 }
 
-// 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)
+// TrackOutbound records an outbound TCP connection
+func (t *TCPTracker) TrackOutbound(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, size int) {
+	if _, exists := t.updateIfExists(dstIP, srcIP, dstPort, srcPort, flags, nftypes.Egress, size); !exists {
+		// if (inverted direction) conn is not tracked, track this direction
+		t.track(srcIP, dstIP, srcPort, dstPort, flags, nftypes.Egress, nil, size)
+	}
+}
 
-		t.logger.Trace("TCP connection reset: %s:%d -> %s:%d",
-			conn.SourceIP, conn.SourcePort, conn.DestIP, conn.DestPort)
+// TrackInbound processes an inbound TCP packet and updates connection state
+func (t *TCPTracker) TrackInbound(srcIP, dstIP netip.Addr, srcPort, 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, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, direction nftypes.Direction, ruleID []byte, size int) {
+	key, exists := t.updateIfExists(srcIP, dstIP, srcPort, dstPort, flags, direction, size)
+	if exists || flags&TCPSyn == 0 {
 		return
 	}
 
-	switch conn.State {
+	conn := &TCPConnTrack{
+		BaseConnTrack: BaseConnTrack{
+			FlowId:    uuid.New(),
+			Direction: direction,
+			SourceIP:  srcIP,
+			DestIP:    dstIP,
+		},
+		SourcePort: srcPort,
+		DestPort:   dstPort,
+	}
+
+	conn.tombstone.Store(false)
+	conn.state.Store(int32(TCPStateNew))
+
+	t.logger.Trace("New %s TCP connection: %s", direction, key)
+	t.updateState(key, conn, flags, direction, size)
+
+	t.mutex.Lock()
+	t.connections[key] = conn
+	t.mutex.Unlock()
+
+	t.sendEvent(nftypes.TypeStart, conn, ruleID)
+}
+
+// IsValidInbound checks if an inbound TCP packet matches a tracked connection
+func (t *TCPTracker) IsValidInbound(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, 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.IsTombstone() {
+		return false
+	}
+
+	currentState := conn.GetState()
+	if !t.isValidStateForFlags(currentState, flags) {
+		t.logger.Warn("TCP state %s is not valid with flags %x for connection %s", currentState, flags, key)
+		// allow all flags for established for now
+		if currentState == TCPStateEstablished {
+			return true
+		}
+		return false
+	}
+
+	t.updateState(key, conn, flags, nftypes.Ingress, size)
+	return true
+}
+
+// updateState updates the TCP connection state based on flags
+func (t *TCPTracker) updateState(key ConnKey, conn *TCPConnTrack, flags uint8, packetDir nftypes.Direction, size int) {
+	conn.UpdateLastSeen()
+	conn.UpdateCounters(packetDir, size)
+
+	currentState := conn.GetState()
+
+	if flags&TCPRst != 0 {
+		if conn.CompareAndSwapState(currentState, TCPStateClosed) {
+			conn.SetTombstone()
+			t.logger.Trace("TCP connection reset: %s (dir: %s) [in: %d Pkts/%d B, out: %d Pkts/%d B]",
+				key, packetDir, conn.PacketsRx.Load(), conn.BytesRx.Load(), conn.PacketsTx.Load(), conn.BytesTx.Load())
+			t.sendEvent(nftypes.TypeEnd, conn, nil)
+		}
+		return
+	}
+
+	var newState TCPState
+	switch currentState {
 	case TCPStateNew:
 		if flags&TCPSyn != 0 && flags&TCPAck == 0 {
-			conn.State = TCPStateSynSent
+			if conn.Direction == nftypes.Egress {
+				newState = TCPStateSynSent
+			} else {
+				newState = TCPStateSynReceived
+			}
 		}
 
 	case TCPStateSynSent:
 		if flags&TCPSyn != 0 && flags&TCPAck != 0 {
-			if isOutbound {
-				conn.State = TCPStateSynReceived
+			if packetDir != conn.Direction {
+				newState = TCPStateEstablished
 			} else {
 				// Simultaneous open
-				conn.State = TCPStateEstablished
-				conn.SetEstablished(true)
+				newState = TCPStateSynReceived
 			}
 		}
 
 	case TCPStateSynReceived:
 		if flags&TCPAck != 0 && flags&TCPSyn == 0 {
-			conn.State = TCPStateEstablished
-			conn.SetEstablished(true)
+			if packetDir == conn.Direction {
+				newState = TCPStateEstablished
+			}
 		}
 
 	case TCPStateEstablished:
 		if flags&TCPFin != 0 {
-			if isOutbound {
-				conn.State = TCPStateFinWait1
+			if packetDir == conn.Direction {
+				newState = TCPStateFinWait1
 			} else {
-				conn.State = TCPStateCloseWait
+				newState = TCPStateCloseWait
 			}
-			conn.SetEstablished(false)
 		}
 
 	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
+		if packetDir != conn.Direction {
+			switch {
+			case flags&TCPFin != 0 && flags&TCPAck != 0:
+				newState = TCPStateClosing
+			case flags&TCPFin != 0:
+				newState = TCPStateClosing
+			case flags&TCPAck != 0:
+				newState = TCPStateFinWait2
+			}
 		}
 
 	case TCPStateFinWait2:
 		if flags&TCPFin != 0 {
-			conn.State = TCPStateTimeWait
+			newState = TCPStateTimeWait
 		}
 
 	case TCPStateClosing:
 		if flags&TCPAck != 0 {
-			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)
+			newState = TCPStateTimeWait
 		}
 
 	case TCPStateCloseWait:
 		if flags&TCPFin != 0 {
-			conn.State = TCPStateLastAck
+			newState = TCPStateLastAck
 		}
 
 	case TCPStateLastAck:
 		if flags&TCPAck != 0 {
-			conn.State = TCPStateClosed
-
-			t.logger.Trace("TCP connection gracefully closed: %s:%d -> %s:%d",
-				conn.SourceIP, conn.SourcePort, conn.DestIP, conn.DestPort)
+			newState = TCPStateClosed
 		}
+	}
 
-	case TCPStateTimeWait:
-		// Stay in TIME-WAIT for 2MSL before transitioning to closed
-		// This is handled by the cleanup routine
+	if newState != 0 && conn.CompareAndSwapState(currentState, newState) {
+		t.logger.Trace("TCP connection %s transitioned from %s to %s (dir: %s)", key, currentState, newState, packetDir)
 
-		t.logger.Trace("TCP connection completed - %s:%d -> %s:%d",
-			conn.SourceIP, conn.SourcePort, conn.DestIP, conn.DestPort)
+		switch newState {
+		case TCPStateTimeWait:
+			t.logger.Trace("TCP connection %s completed [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)
+
+		case TCPStateClosed:
+			conn.SetTombstone()
+			t.logger.Trace("TCP connection %s closed gracefully [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)
+		}
 	}
 }
 
@@ -283,18 +362,22 @@ func (t *TCPTracker) isValidStateForFlags(state TCPState, flags uint8) bool {
 	if !isValidFlagCombination(flags) {
 		return false
 	}
+	if flags&TCPRst != 0 {
+		if state == TCPStateSynSent {
+			return flags&TCPAck != 0
+		}
+		return true
+	}
 
 	switch state {
 	case TCPStateNew:
 		return flags&TCPSyn != 0 && flags&TCPAck == 0
 	case TCPStateSynSent:
+		// TODO: support simultaneous open
 		return flags&TCPSyn != 0 && flags&TCPAck != 0
 	case TCPStateSynReceived:
 		return flags&TCPAck != 0
 	case TCPStateEstablished:
-		if flags&TCPRst != 0 {
-			return true
-		}
 		return flags&TCPAck != 0
 	case TCPStateFinWait1:
 		return flags&TCPFin != 0 || flags&TCPAck != 0
@@ -311,9 +394,7 @@ func (t *TCPTracker) isValidStateForFlags(state TCPState, flags uint8) bool {
 	case TCPStateLastAck:
 		return flags&TCPAck != 0
 	case TCPStateClosed:
-		// Accept retransmitted ACKs in closed state
-		// This is important because the final ACK might be lost
-		// and the peer will retransmit their FIN-ACK
+		// Accept retransmitted ACKs in closed state, the final ACK might be lost and the peer will retransmit their FIN-ACK
 		return flags&TCPAck != 0
 	}
 	return false
@@ -337,24 +418,33 @@ 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:
-			timeout = TimeWaitTimeout
-		case conn.IsEstablished():
+		currentState := conn.GetState()
+		switch currentState {
+		case TCPStateTimeWait:
+			timeout = t.waitTimeout
+		case TCPStateEstablished:
 			timeout = t.timeout
 		default:
 			timeout = TCPHandshakeTimeout
 		}
 
-		lastSeen := conn.GetLastSeen()
-		if time.Since(lastSeen) > timeout {
-			// Return IPs to pool
-			t.ipPool.Put(conn.SourceIP)
-			t.ipPool.Put(conn.DestIP)
+		if conn.timeoutExceeded(timeout) {
 			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 (%s) [in: %d Pkts/%d, B out: %d Pkts/%d B]",
+				key, conn.GetState(), conn.PacketsRx.Load(), conn.BytesRx.Load(), conn.PacketsTx.Load(), conn.BytesTx.Load())
+
+			// event already handled by state change
+			if currentState != TCPStateTimeWait {
+				t.sendEvent(nftypes.TypeEnd, conn, nil)
+			}
 		}
 	}
 }
@@ -365,10 +455,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 +472,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_bench_test.go

@@ -0,0 +1,83 @@
+package conntrack
+
+import (
+	"net/netip"
+	"testing"
+	"time"
+)
+
+func BenchmarkTCPTracker(b *testing.B) {
+	b.Run("TrackOutbound", func(b *testing.B) {
+		tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
+		defer tracker.Close()
+
+		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, 0)
+		}
+	})
+
+	b.Run("IsValidInbound", func(b *testing.B) {
+		tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
+		defer tracker.Close()
+
+		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, 0)
+		}
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			tracker.IsValidInbound(dstIP, srcIP, 80, uint16(i%1000), TCPAck|TCPSyn, 0)
+		}
+	})
+
+	b.Run("ConcurrentAccess", func(b *testing.B) {
+		tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
+		defer tracker.Close()
+
+		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, 0)
+				} else {
+					tracker.IsValidInbound(dstIP, srcIP, 80, uint16(i%65535), TCPAck|TCPSyn, 0)
+				}
+				i++
+			}
+		})
+	})
+}
+
+// Benchmark connection cleanup
+func BenchmarkCleanup(b *testing.B) {
+	b.Run("TCPCleanup", func(b *testing.B) {
+		tracker := NewTCPTracker(100*time.Millisecond, logger, flowLogger)
+		defer tracker.Close()
+
+		// Pre-populate with expired connections
+		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, 0)
+		}
+
+		// Wait for connections to expire
+		time.Sleep(200 * time.Millisecond)
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			tracker.cleanup()
+		}
+	})
+}

client/firewall/uspfilter/conntrack/tcp_test.go

@@ -1,19 +1,20 @@
 package conntrack
 
 import (
-	"net"
+	"net/netip"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 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 +59,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 +77,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 +100,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,11 +123,8 @@ 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
-					// The connection will be cleaned up by timeout
 				},
 			},
 			{
@@ -138,13 +136,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 +152,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 +160,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 +179,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 +193,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,101 +206,455 @@ 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)
-				require.Equal(t, TCPStateClosed, conn.State)
-				require.False(t, conn.IsEstablished())
+				require.Equal(t, TCPStateClosed, conn.GetState())
 			}
 		})
 	}
 }
 
+func TestTCPRetransmissions(t *testing.T) {
+	tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
+	defer tracker.Close()
+
+	srcIP := netip.MustParseAddr("100.64.0.1")
+	dstIP := netip.MustParseAddr("100.64.0.2")
+	srcPort := uint16(12345)
+	dstPort := uint16(80)
+
+	// Test SYN retransmission
+	t.Run("SYN Retransmission", func(t *testing.T) {
+		// Initial SYN
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPSyn, 0)
+
+		// Retransmit SYN (should not affect the state machine)
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPSyn, 0)
+
+		// Verify we're still in SYN-SENT state
+		key := ConnKey{
+			SrcIP:   srcIP,
+			DstIP:   dstIP,
+			SrcPort: srcPort,
+			DstPort: dstPort,
+		}
+		conn := tracker.connections[key]
+		require.NotNil(t, conn)
+		require.Equal(t, TCPStateSynSent, conn.GetState())
+
+		// Complete the handshake
+		valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPSyn|TCPAck, 0)
+		require.True(t, valid)
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, 0)
+
+		// Verify we're in ESTABLISHED state
+		require.Equal(t, TCPStateEstablished, conn.GetState())
+	})
+
+	// Test ACK retransmission in established state
+	t.Run("ACK Retransmission", func(t *testing.T) {
+		tracker = NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
+
+		// Establish connection
+		establishConnection(t, tracker, srcIP, dstIP, srcPort, dstPort)
+
+		// Get connection object
+		key := ConnKey{
+			SrcIP:   srcIP,
+			DstIP:   dstIP,
+			SrcPort: srcPort,
+			DstPort: dstPort,
+		}
+		conn := tracker.connections[key]
+		require.NotNil(t, conn)
+		require.Equal(t, TCPStateEstablished, conn.GetState())
+
+		// Retransmit ACK
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, 0)
+
+		// State should remain ESTABLISHED
+		require.Equal(t, TCPStateEstablished, conn.GetState())
+	})
+
+	// Test FIN retransmission
+	t.Run("FIN Retransmission", func(t *testing.T) {
+		tracker = NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
+
+		// Establish connection
+		establishConnection(t, tracker, srcIP, dstIP, srcPort, dstPort)
+
+		// Get connection object
+		key := ConnKey{
+			SrcIP:   srcIP,
+			DstIP:   dstIP,
+			SrcPort: srcPort,
+			DstPort: dstPort,
+		}
+		conn := tracker.connections[key]
+		require.NotNil(t, conn)
+
+		// Send FIN
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPFin|TCPAck, 0)
+		require.Equal(t, TCPStateFinWait1, conn.GetState())
+
+		// Retransmit FIN (should not change state)
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPFin|TCPAck, 0)
+		require.Equal(t, TCPStateFinWait1, conn.GetState())
+
+		// Receive ACK for FIN
+		valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPAck, 0)
+		require.True(t, valid)
+		require.Equal(t, TCPStateFinWait2, conn.GetState())
+	})
+}
+
+func TestTCPDataTransfer(t *testing.T) {
+	tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
+	defer tracker.Close()
+
+	srcIP := netip.MustParseAddr("100.64.0.1")
+	dstIP := netip.MustParseAddr("100.64.0.2")
+	srcPort := uint16(12345)
+	dstPort := uint16(80)
+
+	t.Run("Data Transfer", func(t *testing.T) {
+		establishConnection(t, tracker, srcIP, dstIP, srcPort, dstPort)
+
+		// Get connection object
+		key := ConnKey{
+			SrcIP:   srcIP,
+			DstIP:   dstIP,
+			SrcPort: srcPort,
+			DstPort: dstPort,
+		}
+		conn := tracker.connections[key]
+		require.NotNil(t, conn)
+
+		// Send data
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPPush|TCPAck, 1000)
+
+		// Receive ACK for data
+		valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPAck, 100)
+		require.True(t, valid)
+
+		// Receive data
+		valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPPush|TCPAck, 1500)
+		require.True(t, valid)
+
+		// Send ACK for received data
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, 100)
+
+		// State should remain ESTABLISHED
+		require.Equal(t, TCPStateEstablished, conn.GetState())
+
+		assert.Equal(t, uint64(1300), conn.BytesTx.Load())
+		assert.Equal(t, uint64(1700), conn.BytesRx.Load())
+		assert.Equal(t, uint64(4), conn.PacketsTx.Load())
+		assert.Equal(t, uint64(3), conn.PacketsRx.Load())
+	})
+}
+
+func TestTCPHalfClosedConnections(t *testing.T) {
+	tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
+	defer tracker.Close()
+
+	srcIP := netip.MustParseAddr("100.64.0.1")
+	dstIP := netip.MustParseAddr("100.64.0.2")
+	srcPort := uint16(12345)
+	dstPort := uint16(80)
+
+	// Test half-closed connection: local end closes, remote end continues sending data
+	t.Run("Local Close, Remote Data", func(t *testing.T) {
+		establishConnection(t, tracker, srcIP, dstIP, srcPort, dstPort)
+
+		key := ConnKey{
+			SrcIP:   srcIP,
+			DstIP:   dstIP,
+			SrcPort: srcPort,
+			DstPort: dstPort,
+		}
+		conn := tracker.connections[key]
+		require.NotNil(t, conn)
+
+		// Send FIN
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPFin|TCPAck, 0)
+		require.Equal(t, TCPStateFinWait1, conn.GetState())
+
+		// Receive ACK for FIN
+		valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPAck, 0)
+		require.True(t, valid)
+		require.Equal(t, TCPStateFinWait2, conn.GetState())
+
+		// Remote end can still send data
+		valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPPush|TCPAck, 1000)
+		require.True(t, valid)
+
+		// We can still ACK their data
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, 0)
+
+		// Receive FIN from remote end
+		valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPFin|TCPAck, 0)
+		require.True(t, valid)
+		require.Equal(t, TCPStateTimeWait, conn.GetState())
+
+		// Send final ACK
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, 0)
+
+		// State should remain TIME-WAIT (waiting for possible retransmissions)
+		require.Equal(t, TCPStateTimeWait, conn.GetState())
+	})
+
+	// Test half-closed connection: remote end closes, local end continues sending data
+	t.Run("Remote Close, Local Data", func(t *testing.T) {
+		tracker = NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
+
+		// Establish connection
+		establishConnection(t, tracker, srcIP, dstIP, srcPort, dstPort)
+
+		// Get connection object
+		key := ConnKey{
+			SrcIP:   srcIP,
+			DstIP:   dstIP,
+			SrcPort: srcPort,
+			DstPort: dstPort,
+		}
+		conn := tracker.connections[key]
+		require.NotNil(t, conn)
+
+		// Receive FIN from remote
+		valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPFin|TCPAck, 0)
+		require.True(t, valid)
+		require.Equal(t, TCPStateCloseWait, conn.GetState())
+
+		// We can still send data
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPPush|TCPAck, 1000)
+
+		// Remote can still ACK our data
+		valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPAck, 0)
+		require.True(t, valid)
+
+		// Send our FIN
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPFin|TCPAck, 0)
+		require.Equal(t, TCPStateLastAck, conn.GetState())
+
+		// Receive final ACK
+		valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPAck, 0)
+		require.True(t, valid)
+		require.Equal(t, TCPStateClosed, conn.GetState())
+	})
+}
+
+func TestTCPAbnormalSequences(t *testing.T) {
+	tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
+	defer tracker.Close()
+
+	srcIP := netip.MustParseAddr("100.64.0.1")
+	dstIP := netip.MustParseAddr("100.64.0.2")
+	srcPort := uint16(12345)
+	dstPort := uint16(80)
+
+	// Test handling of unsolicited RST in various states
+	t.Run("Unsolicited RST in SYN-SENT", func(t *testing.T) {
+		// Send SYN
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPSyn, 0)
+
+		// Receive unsolicited RST (without proper ACK)
+		valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPRst, 0)
+		require.False(t, valid, "RST without proper ACK in SYN-SENT should be rejected")
+
+		// Receive RST with proper ACK
+		valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPRst|TCPAck, 0)
+		require.True(t, valid, "RST with proper ACK in SYN-SENT should be accepted")
+
+		key := ConnKey{
+			SrcIP:   srcIP,
+			DstIP:   dstIP,
+			SrcPort: srcPort,
+			DstPort: dstPort,
+		}
+		conn := tracker.connections[key]
+		require.Equal(t, TCPStateClosed, conn.GetState())
+		require.True(t, conn.IsTombstone())
+	})
+}
+
+func TestTCPTimeoutHandling(t *testing.T) {
+	// Create tracker with a very short timeout for testing
+	shortTimeout := 100 * time.Millisecond
+	tracker := NewTCPTracker(shortTimeout, logger, flowLogger)
+	defer tracker.Close()
+
+	srcIP := netip.MustParseAddr("100.64.0.1")
+	dstIP := netip.MustParseAddr("100.64.0.2")
+	srcPort := uint16(12345)
+	dstPort := uint16(80)
+
+	t.Run("Connection Timeout", func(t *testing.T) {
+		// Establish a connection
+		establishConnection(t, tracker, srcIP, dstIP, srcPort, dstPort)
+
+		// Get connection object
+		key := ConnKey{
+			SrcIP:   srcIP,
+			DstIP:   dstIP,
+			SrcPort: srcPort,
+			DstPort: dstPort,
+		}
+		conn := tracker.connections[key]
+		require.NotNil(t, conn)
+		require.Equal(t, TCPStateEstablished, conn.GetState())
+
+		// Wait for the connection to timeout
+		time.Sleep(2 * shortTimeout)
+
+		// Force cleanup
+		tracker.cleanup()
+
+		// Connection should be removed
+		_, exists := tracker.connections[key]
+		require.False(t, exists, "Connection should be removed after timeout")
+	})
+
+	t.Run("TIME_WAIT Timeout", func(t *testing.T) {
+		tracker = NewTCPTracker(shortTimeout, logger, flowLogger)
+
+		establishConnection(t, tracker, srcIP, dstIP, srcPort, dstPort)
+
+		key := ConnKey{
+			SrcIP:   srcIP,
+			DstIP:   dstIP,
+			SrcPort: srcPort,
+			DstPort: dstPort,
+		}
+		conn := tracker.connections[key]
+		require.NotNil(t, conn)
+
+		// Complete the connection close to enter TIME_WAIT
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPFin|TCPAck, 0)
+		tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPAck, 0)
+		tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPFin|TCPAck, 0)
+		tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, 0)
+
+		require.Equal(t, TCPStateTimeWait, conn.GetState())
+
+		// TIME_WAIT should have its own timeout value (usually 2*MSL)
+		// For the test, we're using a short timeout
+		time.Sleep(2 * shortTimeout)
+
+		tracker.cleanup()
+
+		// Connection should be removed
+		_, exists := tracker.connections[key]
+		require.False(t, exists, "Connection should be removed after TIME_WAIT timeout")
+	})
+}
+
+func TestSynFlood(t *testing.T) {
+	tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
+	defer tracker.Close()
+
+	srcIP := netip.MustParseAddr("100.64.0.1")
+	dstIP := netip.MustParseAddr("100.64.0.2")
+	basePort := uint16(10000)
+	dstPort := uint16(80)
+
+	// Create a large number of SYN packets to simulate a SYN flood
+	for i := uint16(0); i < 1000; i++ {
+		tracker.TrackOutbound(srcIP, dstIP, basePort+i, dstPort, TCPSyn, 0)
+	}
+
+	// Check that we're tracking all connections
+	require.Equal(t, 1000, len(tracker.connections))
+
+	// Now simulate SYN timeout
+	var oldConns int
+	tracker.mutex.Lock()
+	for _, conn := range tracker.connections {
+		if conn.GetState() == TCPStateSynSent {
+			// Make the connection appear old
+			conn.lastSeen.Store(time.Now().Add(-TCPHandshakeTimeout - time.Second).UnixNano())
+			oldConns++
+		}
+	}
+	tracker.mutex.Unlock()
+	require.Equal(t, 1000, oldConns)
+
+	// Run cleanup
+	tracker.cleanup()
+
+	// Check that stale connections were cleaned up
+	require.Equal(t, 0, len(tracker.connections))
+}
+
+func TestTCPInboundInitiatedConnection(t *testing.T) {
+	tracker := NewTCPTracker(DefaultTCPTimeout, logger, flowLogger)
+	defer tracker.Close()
+
+	clientIP := netip.MustParseAddr("100.64.0.1")
+	serverIP := netip.MustParseAddr("100.64.0.2")
+	clientPort := uint16(12345)
+	serverPort := uint16(80)
+
+	// 1. Client sends SYN (we receive it as inbound)
+	tracker.TrackInbound(clientIP, serverIP, clientPort, serverPort, TCPSyn, nil, 100)
+
+	key := ConnKey{
+		SrcIP:   clientIP,
+		DstIP:   serverIP,
+		SrcPort: clientPort,
+		DstPort: serverPort,
+	}
+
+	tracker.mutex.RLock()
+	conn := tracker.connections[key]
+	tracker.mutex.RUnlock()
+
+	require.NotNil(t, conn)
+	require.Equal(t, TCPStateSynReceived, conn.GetState(), "Connection should be in SYN-RECEIVED state after inbound SYN")
+
+	// 2. Server sends SYN-ACK response
+	tracker.TrackOutbound(serverIP, clientIP, serverPort, clientPort, TCPSyn|TCPAck, 100)
+
+	// 3. Client sends ACK to complete handshake
+	tracker.TrackInbound(clientIP, serverIP, clientPort, serverPort, TCPAck, nil, 100)
+	require.Equal(t, TCPStateEstablished, conn.GetState(), "Connection should be ESTABLISHED after handshake completion")
+
+	// 4. Test data transfer
+	// Client sends data
+	tracker.TrackInbound(clientIP, serverIP, clientPort, serverPort, TCPPush|TCPAck, nil, 1000)
+
+	// Server sends ACK for data
+	tracker.TrackOutbound(serverIP, clientIP, serverPort, clientPort, TCPAck, 100)
+
+	// Server sends data
+	tracker.TrackOutbound(serverIP, clientIP, serverPort, clientPort, TCPPush|TCPAck, 1500)
+
+	// Client sends ACK for data
+	tracker.TrackInbound(clientIP, serverIP, clientPort, serverPort, TCPAck, nil, 100)
+
+	// Verify state and counters
+	require.Equal(t, TCPStateEstablished, conn.GetState())
+	assert.Equal(t, uint64(1300), conn.BytesRx.Load()) // 3 packets * 100 + 1000 data
+	assert.Equal(t, uint64(1700), conn.BytesTx.Load()) // 2 packets * 100 + 1500 data
+	assert.Equal(t, uint64(4), conn.PacketsRx.Load())  // SYN, ACK, Data
+	assert.Equal(t, uint64(3), conn.PacketsTx.Load())  // SYN-ACK, Data
+}
+
 // 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, 100)
 
-	valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPSyn|TCPAck)
+	valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPSyn|TCPAck, 100)
 	require.True(t, valid, "SYN-ACK should be allowed")
 
-	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck)
-}
-
-func BenchmarkTCPTracker(b *testing.B) {
-	b.Run("TrackOutbound", func(b *testing.B) {
-		tracker := NewTCPTracker(DefaultTCPTimeout, logger)
-		defer tracker.Close()
-
-		srcIP := net.ParseIP("192.168.1.1")
-		dstIP := net.ParseIP("192.168.1.2")
-
-		b.ResetTimer()
-		for i := 0; i < b.N; i++ {
-			tracker.TrackOutbound(srcIP, dstIP, uint16(i%65535), 80, TCPSyn)
-		}
-	})
-
-	b.Run("IsValidInbound", func(b *testing.B) {
-		tracker := NewTCPTracker(DefaultTCPTimeout, logger)
-		defer tracker.Close()
-
-		srcIP := net.ParseIP("192.168.1.1")
-		dstIP := net.ParseIP("192.168.1.2")
-
-		// Pre-populate some connections
-		for i := 0; i < 1000; i++ {
-			tracker.TrackOutbound(srcIP, dstIP, uint16(i), 80, TCPSyn)
-		}
-
-		b.ResetTimer()
-		for i := 0; i < b.N; i++ {
-			tracker.IsValidInbound(dstIP, srcIP, 80, uint16(i%1000), TCPAck)
-		}
-	})
-
-	b.Run("ConcurrentAccess", func(b *testing.B) {
-		tracker := NewTCPTracker(DefaultTCPTimeout, logger)
-		defer tracker.Close()
-
-		srcIP := net.ParseIP("192.168.1.1")
-		dstIP := net.ParseIP("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)
-				} else {
-					tracker.IsValidInbound(dstIP, srcIP, 80, uint16(i%65535), TCPAck)
-				}
-				i++
-			}
-		})
-	})
-}
-
-// 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
-		defer tracker.Close()
-
-		// Pre-populate with expired connections
-		srcIP := net.ParseIP("192.168.1.1")
-		dstIP := net.ParseIP("192.168.1.2")
-		for i := 0; i < 10000; i++ {
-			tracker.TrackOutbound(srcIP, dstIP, uint16(i), 80, TCPSyn)
-		}
-
-		// Wait for connections to expire
-		time.Sleep(200 * time.Millisecond)
-
-		b.ResetTimer()
-		for i := 0; i < b.N; i++ {
-			tracker.cleanup()
-		}
-	})
+	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, 100)
 }

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,88 @@ 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()
+	conn.UpdateCounters(direction, size)
+
+	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 +163,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 +177,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"
@@ -13,8 +14,13 @@ import (
 type localIPManager struct {
 	mu sync.RWMutex
 
-	// Use bitmap for IPv4 (32 bits * 2^16 = 256KB memory)
-	ipv4Bitmap [1 << 16]uint32
+	// fixed-size high array for upper byte of a IPv4 address
+	ipv4Bitmap [256]*ipv4LowBitmap
+}
+
+// ipv4LowBitmap is a map for the low 16 bits of a IPv4 address
+type ipv4LowBitmap struct {
+	bitmap [8192]uint32
 }
 
 func newLocalIPManager() *localIPManager {
@@ -26,39 +32,59 @@ func (m *localIPManager) setBitmapBit(ip net.IP) {
 	if ipv4 == nil {
 		return
 	}
-	high := (uint16(ipv4[0]) << 8) | uint16(ipv4[1])
-	low := (uint16(ipv4[2]) << 8) | uint16(ipv4[3])
-	m.ipv4Bitmap[high] |= 1 << (low % 32)
-}
+	high := uint16(ipv4[0])
+	low := (uint16(ipv4[1]) << 8) | (uint16(ipv4[2]) << 4) | uint16(ipv4[3])
 
-func (m *localIPManager) checkBitmapBit(ip net.IP) bool {
-	ipv4 := ip.To4()
-	if ipv4 == nil {
-		return false
+	index := low / 32
+	bit := low % 32
+
+	if m.ipv4Bitmap[high] == nil {
+		m.ipv4Bitmap[high] = &ipv4LowBitmap{}
 	}
-	high := (uint16(ipv4[0]) << 8) | uint16(ipv4[1])
-	low := (uint16(ipv4[2]) << 8) | uint16(ipv4[3])
-	return (m.ipv4Bitmap[high] & (1 << (low % 32))) != 0
+
+	m.ipv4Bitmap[high].bitmap[index] |= 1 << bit
 }
 
-func (m *localIPManager) processIP(ip net.IP, newIPv4Bitmap *[1 << 16]uint32, ipv4Set map[string]struct{}, ipv4Addresses *[]string) error {
+func (m *localIPManager) setBitInBitmap(ip net.IP, bitmap *[256]*ipv4LowBitmap, ipv4Set map[string]struct{}, ipv4Addresses *[]string) {
 	if ipv4 := ip.To4(); ipv4 != nil {
-		high := (uint16(ipv4[0]) << 8) | uint16(ipv4[1])
-		low := (uint16(ipv4[2]) << 8) | uint16(ipv4[3])
-		if int(high) >= len(*newIPv4Bitmap) {
-			return fmt.Errorf("invalid IPv4 address: %s", ip)
+		high := uint16(ipv4[0])
+		low := (uint16(ipv4[1]) << 8) | (uint16(ipv4[2]) << 4) | uint16(ipv4[3])
+
+		if bitmap[high] == nil {
+			bitmap[high] = &ipv4LowBitmap{}
 		}
-		ipStr := ip.String()
+
+		index := low / 32
+		bit := low % 32
+		bitmap[high].bitmap[index] |= 1 << bit
+
+		ipStr := ipv4.String()
 		if _, exists := ipv4Set[ipStr]; !exists {
 			ipv4Set[ipStr] = struct{}{}
 			*ipv4Addresses = append(*ipv4Addresses, ipStr)
-			newIPv4Bitmap[high] |= 1 << (low % 32)
 		}
 	}
+}
+
+func (m *localIPManager) checkBitmapBit(ip []byte) bool {
+	high := uint16(ip[0])
+	low := (uint16(ip[1]) << 8) | (uint16(ip[2]) << 4) | uint16(ip[3])
+
+	if m.ipv4Bitmap[high] == nil {
+		return false
+	}
+
+	index := low / 32
+	bit := low % 32
+	return (m.ipv4Bitmap[high].bitmap[index] & (1 << bit)) != 0
+}
+
+func (m *localIPManager) processIP(ip net.IP, bitmap *[256]*ipv4LowBitmap, ipv4Set map[string]struct{}, ipv4Addresses *[]string) error {
+	m.setBitInBitmap(ip, bitmap, ipv4Set, ipv4Addresses)
 	return nil
 }
 
-func (m *localIPManager) processInterface(iface net.Interface, newIPv4Bitmap *[1 << 16]uint32, ipv4Set map[string]struct{}, ipv4Addresses *[]string) {
+func (m *localIPManager) processInterface(iface net.Interface, bitmap *[256]*ipv4LowBitmap, ipv4Set map[string]struct{}, ipv4Addresses *[]string) {
 	addrs, err := iface.Addrs()
 	if err != nil {
 		log.Debugf("get addresses for interface %s failed: %v", iface.Name, err)
@@ -76,7 +102,7 @@ func (m *localIPManager) processInterface(iface net.Interface, newIPv4Bitmap *[1
 			continue
 		}
 
-		if err := m.processIP(ip, newIPv4Bitmap, ipv4Set, ipv4Addresses); err != nil {
+		if err := m.processIP(ip, bitmap, ipv4Set, ipv4Addresses); err != nil {
 			log.Debugf("process IP failed: %v", err)
 		}
 	}
@@ -89,14 +115,14 @@ func (m *localIPManager) UpdateLocalIPs(iface common.IFaceMapper) (err error) {
 		}
 	}()
 
-	var newIPv4Bitmap [1 << 16]uint32
+	var newIPv4Bitmap [256]*ipv4LowBitmap
 	ipv4Set := make(map[string]struct{})
 	var ipv4Addresses []string
 
 	// 127.0.0.0/8
-	high := uint16(127) << 8
-	for i := uint16(0); i < 256; i++ {
-		newIPv4Bitmap[high|i] = 0xffffffff
+	newIPv4Bitmap[127] = &ipv4LowBitmap{}
+	for i := 0; i < 8192; i++ {
+		newIPv4Bitmap[127].bitmap[i] = 0xFFFFFFFF
 	}
 
 	if iface != nil {
@@ -122,13 +148,13 @@ 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 {
+	if !ip.Is4() {
+		return false
+	}
+
 	m.mu.RLock()
 	defer m.mu.RUnlock()
 
-	if ipv4 := ip.To4(); ipv4 != nil {
-		return m.checkBitmapBit(ipv4)
-	}
-
-	return false
+	return m.checkBitmapBit(ip.AsSlice())
 }

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,19 @@ 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,
+		},
+		{
+			name: "Local IP doesn't match - addresses 32 apart",
+			setupAddr: wgaddr.Address{
+				IP: net.ParseIP("192.168.1.1"),
+				Network: &net.IPNet{
+					IP:   net.ParseIP("192.168.1.0"),
+					Mask: net.CIDRMask(24, 32),
+				},
+			},
+			testIP:   netip.MustParseAddr("192.168.1.33"),
 			expected: false,
 		},
 		{
@@ -85,7 +98,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 +187,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)
 		})
 	}
@@ -191,10 +204,8 @@ func BenchmarkIPChecks(b *testing.B) {
 		interfaces[i] = net.IPv4(10, 0, byte(i>>8), byte(i))
 	}
 
-	// Setup bitmap version
-	bitmapManager := &localIPManager{
-		ipv4Bitmap: [1 << 16]uint32{},
-	}
+	// Setup bitmap
+	bitmapManager := newLocalIPManager()
 	for _, ip := range interfaces[:8] { // Add half of IPs
 		bitmapManager.setBitmapBit(ip)
 	}
@@ -247,7 +258,7 @@ func BenchmarkWGPosition(b *testing.B) {
 
 	// Create two managers - one checks WG IP first, other checks it last
 	b.Run("WG_First", func(b *testing.B) {
-		bm := &localIPManager{ipv4Bitmap: [1 << 16]uint32{}}
+		bm := newLocalIPManager()
 		bm.setBitmapBit(wgIP)
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
@@ -256,7 +267,7 @@ func BenchmarkWGPosition(b *testing.B) {
 	})
 
 	b.Run("WG_Last", func(b *testing.B) {
-		bm := &localIPManager{ipv4Bitmap: [1 << 16]uint32{}}
+		bm := newLocalIPManager()
 		// Fill with other IPs first
 		for i := 0; i < 15; i++ {
 			bm.setBitmapBit(net.IPv4(10, 0, byte(i>>8), byte(i)))

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,24 +298,24 @@ 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")
 
@@ -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,53 +663,73 @@ 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)
+	return m.handleRoutedTraffic(d, srcIP, dstIP, packetData, size)
 }
 
 // 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
 	}
 
 	// if running in netstack mode we need to pass this to the forwarder
-	if m.netstack {
+	if m.netstack && m.localForwarding {
 		return m.handleNetstackLocalTraffic(packetData)
 	}
 
+	// track inbound packets to get the correct direction and session id for flows
+	m.trackInbound(d, srcIP, dstIP, ruleID, size)
+
+	// pass to either native or virtual stack (to be picked up by listeners)
 	return false
 }
 
 func (m *Manager) handleNetstackLocalTraffic(packetData []byte) bool {
-	if !m.localForwarding {
-		// pass to virtual tcp/ip stack to be picked up by listeners
-		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 +739,68 @@ 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, size int) 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() {
+		m.trackInbound(d, srcIP, dstIP, nil, size)
 		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
+			RxPackets: 1,
+			RxBytes:   uint64(size),
+		})
 		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

@@ -3,6 +3,7 @@ package uspfilter
 import (
 	"fmt"
 	"net"
+	"net/netip"
 	"sync"
 	"testing"
 	"time"
@@ -18,9 +19,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(nil, []byte{}, nil).GetLogger()
 
 type IFaceMock struct {
 	SetFilterFunc   func(device.PacketFilter) error
@@ -62,7 +65,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 +85,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 +95,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 +118,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 +150,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 +161,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 +170,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 +178,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 +188,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 +207,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 +237,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 +247,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 +279,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 +292,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 +327,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 +346,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 +356,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 +392,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 +400,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 +422,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 +457,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 +467,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 +478,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 +493,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 +505,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 +514,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 +533,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 +542,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 +568,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 +584,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 +635,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 +684,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 +706,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/bind/udp_mux.go

@@ -150,7 +150,7 @@ func isZeros(ip net.IP) bool {
 // NewUDPMuxDefault creates an implementation of UDPMux
 func NewUDPMuxDefault(params UDPMuxParams) *UDPMuxDefault {
 	if params.Logger == nil {
-		params.Logger = logging.NewDefaultLoggerFactory().NewLogger("ice")
+		params.Logger = getLogger()
 	}
 
 	mux := &UDPMuxDefault{
@@ -455,3 +455,9 @@ func newBufferHolder(size int) *bufferHolder {
 		buf: make([]byte, size),
 	}
 }
+
+func getLogger() logging.LeveledLogger {
+	fac := logging.NewDefaultLoggerFactory()
+	fac.Writer = log.StandardLogger().Writer()
+	return fac.NewLogger("ice")
+}

client/iface/bind/udp_mux_universal.go

@@ -49,7 +49,7 @@ type UniversalUDPMuxParams struct {
 // NewUniversalUDPMuxDefault creates an implementation of UniversalUDPMux embedding UDPMux
 func NewUniversalUDPMuxDefault(params UniversalUDPMuxParams) *UniversalUDPMuxDefault {
 	if params.Logger == nil {
-		params.Logger = logging.NewDefaultLoggerFactory().NewLogger("ice")
+		params.Logger = getLogger()
 	}
 	if params.XORMappedAddrCacheTTL == 0 {
 		params.XORMappedAddrCacheTTL = time.Second * 25

client/iface/configurer/usp.go

@@ -357,7 +357,7 @@ func toWgUserspaceString(wgCfg wgtypes.Config) string {
 
 func getFwmark() int {
 	if nbnet.AdvancedRouting() {
-		return nbnet.NetbirdFwmark
+		return nbnet.ControlPlaneMark
 	}
 	return 0
 }

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

@@ -10,9 +10,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(nil, []byte{}, nil).GetLogger()
+
 func TestDefaultManager(t *testing.T) {
 	networkMap := &mgmProto.NetworkMap{
 		FirewallRules: []*mgmProto.FirewallRule{
@@ -52,7 +55,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 +349,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/auth/pkce_flow.go

@@ -99,6 +99,7 @@ func (p *PKCEAuthorizationFlow) RequestAuthInfo(ctx context.Context) (AuthFlowIn
 		oauth2.SetAuthURLParam("code_challenge_method", "S256"),
 		oauth2.SetAuthURLParam("code_challenge", codeChallenge),
 		oauth2.SetAuthURLParam("audience", p.providerConfig.Audience),
+		oauth2.SetAuthURLParam("prompt", "login"),
 	)
 
 	return AuthFlowInfo{

client/internal/dns/file_unix.go

@@ -239,7 +239,7 @@ func searchDomains(config HostDNSConfig) []string {
 			continue
 		}
 
-		listOfDomains = append(listOfDomains, dConf.Domain)
+		listOfDomains = append(listOfDomains, strings.TrimSuffix(dConf.Domain, "."))
 	}
 	return listOfDomains
 }

client/internal/dns/handler_chain.go

@@ -75,12 +75,7 @@ func (c *HandlerChain) AddHandler(pattern string, handler dns.Handler, priority
 	}
 
 	// First remove any existing handler with same pattern (case-insensitive) and priority
-	for i := len(c.handlers) - 1; i >= 0; i-- {
-		if strings.EqualFold(c.handlers[i].OrigPattern, origPattern) && c.handlers[i].Priority == priority {
-			c.handlers = append(c.handlers[:i], c.handlers[i+1:]...)
-			break
-		}
-	}
+	c.removeEntry(origPattern, priority)
 
 	// Check if handler implements SubdomainMatcher interface
 	matchSubdomains := false
@@ -133,30 +128,20 @@ func (c *HandlerChain) RemoveHandler(pattern string, priority int) {
 
 	pattern = dns.Fqdn(pattern)
 
+	c.removeEntry(pattern, priority)
+}
+
+func (c *HandlerChain) removeEntry(pattern string, priority int) {
 	// Find and remove handlers matching both original pattern (case-insensitive) and priority
 	for i := len(c.handlers) - 1; i >= 0; i-- {
 		entry := c.handlers[i]
 		if strings.EqualFold(entry.OrigPattern, pattern) && entry.Priority == priority {
 			c.handlers = append(c.handlers[:i], c.handlers[i+1:]...)
-			return
+			break
 		}
 	}
 }
 
-// HasHandlers returns true if there are any handlers remaining for the given pattern
-func (c *HandlerChain) HasHandlers(pattern string) bool {
-	c.mu.RLock()
-	defer c.mu.RUnlock()
-
-	pattern = strings.ToLower(dns.Fqdn(pattern))
-	for _, entry := range c.handlers {
-		if strings.EqualFold(entry.Pattern, pattern) {
-			return true
-		}
-	}
-	return false
-}
-
 func (c *HandlerChain) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 	if len(r.Question) == 0 {
 		return

client/internal/dns/handler_chain_test.go

@@ -443,14 +443,6 @@ func TestHandlerChain_PriorityDeregistration(t *testing.T) {
 			for _, handler := range handlers {
 				handler.AssertExpectations(t)
 			}
-
-			// Verify handler exists check
-			for priority, shouldExist := range tt.expectedCalls {
-				if shouldExist {
-					assert.True(t, chain.HasHandlers(tt.ops[0].pattern),
-						"Handler chain should have handlers for pattern after removing priority %d", priority)
-				}
-			}
 		})
 	}
 }
@@ -470,45 +462,69 @@ func TestHandlerChain_MultiPriorityHandling(t *testing.T) {
 	r := new(dns.Msg)
 	r.SetQuestion(testQuery, dns.TypeA)
 
+	// Keep track of mocks for the final assertion in Step 4
+	mocks := []*nbdns.MockSubdomainHandler{routeHandler, matchHandler, defaultHandler}
+
 	// Add handlers in mixed order
 	chain.AddHandler(testDomain, defaultHandler, nbdns.PriorityDefault)
 	chain.AddHandler(testDomain, routeHandler, nbdns.PriorityDNSRoute)
 	chain.AddHandler(testDomain, matchHandler, nbdns.PriorityMatchDomain)
 
-	// Test 1: Initial state with all three handlers
-	w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	// Test 1: Initial state
+	w1 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 	// Highest priority handler (routeHandler) should be called
 	routeHandler.On("ServeDNS", mock.Anything, r).Return().Once()
+	matchHandler.On("ServeDNS", mock.Anything, r).Maybe()   // Ensure others are not expected yet
+	defaultHandler.On("ServeDNS", mock.Anything, r).Maybe() // Ensure others are not expected yet
 
-	chain.ServeDNS(w, r)
+	chain.ServeDNS(w1, r)
 	routeHandler.AssertExpectations(t)
 
+	routeHandler.ExpectedCalls = nil
+	routeHandler.Calls = nil
+	matchHandler.ExpectedCalls = nil
+	matchHandler.Calls = nil
+	defaultHandler.ExpectedCalls = nil
+	defaultHandler.Calls = nil
+
 	// Test 2: Remove highest priority handler
 	chain.RemoveHandler(testDomain, nbdns.PriorityDNSRoute)
-	assert.True(t, chain.HasHandlers(testDomain))
 
-	w = &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	w2 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 	// Now middle priority handler (matchHandler) should be called
 	matchHandler.On("ServeDNS", mock.Anything, r).Return().Once()
+	defaultHandler.On("ServeDNS", mock.Anything, r).Maybe() // Ensure default is not expected yet
 
-	chain.ServeDNS(w, r)
+	chain.ServeDNS(w2, r)
 	matchHandler.AssertExpectations(t)
 
+	matchHandler.ExpectedCalls = nil
+	matchHandler.Calls = nil
+	defaultHandler.ExpectedCalls = nil
+	defaultHandler.Calls = nil
+
 	// Test 3: Remove middle priority handler
 	chain.RemoveHandler(testDomain, nbdns.PriorityMatchDomain)
-	assert.True(t, chain.HasHandlers(testDomain))
 
-	w = &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	w3 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 	// Now lowest priority handler (defaultHandler) should be called
 	defaultHandler.On("ServeDNS", mock.Anything, r).Return().Once()
 
-	chain.ServeDNS(w, r)
+	chain.ServeDNS(w3, r)
 	defaultHandler.AssertExpectations(t)
 
+	defaultHandler.ExpectedCalls = nil
+	defaultHandler.Calls = nil
+
 	// Test 4: Remove last handler
 	chain.RemoveHandler(testDomain, nbdns.PriorityDefault)
 
-	assert.False(t, chain.HasHandlers(testDomain))
+	w4 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	chain.ServeDNS(w4, r) // Call ServeDNS on the now empty chain for this domain
+
+	for _, m := range mocks {
+		m.AssertNumberOfCalls(t, "ServeDNS", 0)
+	}
 }
 
 func TestHandlerChain_CaseSensitivity(t *testing.T) {
@@ -830,3 +846,165 @@ func TestHandlerChain_DomainSpecificityOrdering(t *testing.T) {
 		})
 	}
 }
+
+func TestHandlerChain_AddRemoveRoundtrip(t *testing.T) {
+	tests := []struct {
+		name            string
+		addPattern      string
+		removePattern   string
+		queryPattern    string
+		shouldBeRemoved bool
+		description     string
+	}{
+		{
+			name:            "exact same pattern",
+			addPattern:      "example.com.",
+			removePattern:   "example.com.",
+			queryPattern:    "example.com.",
+			shouldBeRemoved: true,
+			description:     "Adding and removing with identical patterns",
+		},
+		{
+			name:            "case difference",
+			addPattern:      "Example.Com.",
+			removePattern:   "EXAMPLE.COM.",
+			queryPattern:    "example.com.",
+			shouldBeRemoved: true,
+			description:     "Adding with mixed case, removing with uppercase",
+		},
+		{
+			name:            "reversed case difference",
+			addPattern:      "EXAMPLE.ORG.",
+			removePattern:   "example.org.",
+			queryPattern:    "example.org.",
+			shouldBeRemoved: true,
+			description:     "Adding with uppercase, removing with lowercase",
+		},
+		{
+			name:            "add wildcard, remove wildcard",
+			addPattern:      "*.example.com.",
+			removePattern:   "*.example.com.",
+			queryPattern:    "sub.example.com.",
+			shouldBeRemoved: true,
+			description:     "Adding and removing with identical wildcard patterns",
+		},
+		{
+			name:            "add wildcard, remove transformed pattern",
+			addPattern:      "*.example.net.",
+			removePattern:   "example.net.",
+			queryPattern:    "sub.example.net.",
+			shouldBeRemoved: false,
+			description:     "Adding with wildcard, removing with non-wildcard pattern",
+		},
+		{
+			name:            "add transformed pattern, remove wildcard",
+			addPattern:      "example.io.",
+			removePattern:   "*.example.io.",
+			queryPattern:    "example.io.",
+			shouldBeRemoved: false,
+			description:     "Adding with non-wildcard pattern, removing with wildcard pattern",
+		},
+		{
+			name:            "trailing dot difference",
+			addPattern:      "example.dev",
+			removePattern:   "example.dev.",
+			queryPattern:    "example.dev.",
+			shouldBeRemoved: true,
+			description:     "Adding without trailing dot, removing with trailing dot",
+		},
+		{
+			name:            "reversed trailing dot difference",
+			addPattern:      "example.app.",
+			removePattern:   "example.app",
+			queryPattern:    "example.app.",
+			shouldBeRemoved: true,
+			description:     "Adding with trailing dot, removing without trailing dot",
+		},
+		{
+			name:            "mixed case and wildcard",
+			addPattern:      "*.Example.Site.",
+			removePattern:   "*.EXAMPLE.SITE.",
+			queryPattern:    "sub.example.site.",
+			shouldBeRemoved: true,
+			description:     "Adding mixed case wildcard, removing uppercase wildcard",
+		},
+		{
+			name:            "root zone",
+			addPattern:      ".",
+			removePattern:   ".",
+			queryPattern:    "random.domain.",
+			shouldBeRemoved: true,
+			description:     "Adding and removing root zone",
+		},
+		{
+			name:            "wrong domain",
+			addPattern:      "example.com.",
+			removePattern:   "different.com.",
+			queryPattern:    "example.com.",
+			shouldBeRemoved: false,
+			description:     "Adding one domain, trying to remove a different domain",
+		},
+		{
+			name:            "subdomain mismatch",
+			addPattern:      "sub.example.com.",
+			removePattern:   "example.com.",
+			queryPattern:    "sub.example.com.",
+			shouldBeRemoved: false,
+			description:     "Adding subdomain, trying to remove parent domain",
+		},
+		{
+			name:            "parent domain mismatch",
+			addPattern:      "example.com.",
+			removePattern:   "sub.example.com.",
+			queryPattern:    "example.com.",
+			shouldBeRemoved: false,
+			description:     "Adding parent domain, trying to remove subdomain",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			chain := nbdns.NewHandlerChain()
+
+			handler := &nbdns.MockHandler{}
+			r := new(dns.Msg)
+			r.SetQuestion(tt.queryPattern, dns.TypeA)
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+
+			// First verify no handler is called before adding any
+			chain.ServeDNS(w, r)
+			handler.AssertNotCalled(t, "ServeDNS")
+
+			// Add handler
+			chain.AddHandler(tt.addPattern, handler, nbdns.PriorityDefault)
+
+			// Verify handler is called after adding
+			handler.On("ServeDNS", mock.Anything, r).Once()
+			chain.ServeDNS(w, r)
+			handler.AssertExpectations(t)
+
+			// Reset mock for the next test
+			handler.ExpectedCalls = nil
+
+			// Remove handler
+			chain.RemoveHandler(tt.removePattern, nbdns.PriorityDefault)
+
+			// Set up expectations based on whether removal should succeed
+			if !tt.shouldBeRemoved {
+				handler.On("ServeDNS", mock.Anything, r).Once()
+			}
+
+			// Test if handler is still called after removal attempt
+			chain.ServeDNS(w, r)
+
+			if tt.shouldBeRemoved {
+				handler.AssertNotCalled(t, "ServeDNS",
+					"Handler should not be called after successful removal with pattern %q",
+					tt.removePattern)
+			} else {
+				handler.AssertExpectations(t)
+				handler.ExpectedCalls = nil
+			}
+		})
+	}
+}

client/internal/dns/host.go

@@ -5,6 +5,8 @@ import (
 	"net/netip"
 	"strings"
 
+	"github.com/miekg/dns"
+
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 	nbdns "github.com/netbirdio/netbird/dns"
 )
@@ -12,8 +14,8 @@ import (
 var ErrRouteAllWithoutNameserverGroup = fmt.Errorf("unable to configure DNS for this peer using file manager without a nameserver group with all domains configured")
 
 const (
-	ipv4ReverseZone = ".in-addr.arpa"
-	ipv6ReverseZone = ".ip6.arpa"
+	ipv4ReverseZone = ".in-addr.arpa."
+	ipv6ReverseZone = ".ip6.arpa."
 )
 
 type hostManager interface {
@@ -103,7 +105,7 @@ func dnsConfigToHostDNSConfig(dnsConfig nbdns.Config, ip string, port int) HostD
 
 		for _, domain := range nsConfig.Domains {
 			config.Domains = append(config.Domains, DomainConfig{
-				Domain:    strings.TrimSuffix(domain, "."),
+				Domain:    strings.ToLower(dns.Fqdn(domain)),
 				MatchOnly: !nsConfig.SearchDomainsEnabled,
 			})
 		}
@@ -112,7 +114,7 @@ func dnsConfigToHostDNSConfig(dnsConfig nbdns.Config, ip string, port int) HostD
 	for _, customZone := range dnsConfig.CustomZones {
 		matchOnly := strings.HasSuffix(customZone.Domain, ipv4ReverseZone) || strings.HasSuffix(customZone.Domain, ipv6ReverseZone)
 		config.Domains = append(config.Domains, DomainConfig{
-			Domain:    strings.TrimSuffix(customZone.Domain, "."),
+			Domain:    strings.ToLower(dns.Fqdn(customZone.Domain)),
 			MatchOnly: matchOnly,
 		})
 	}

client/internal/dns/host_darwin.go

@@ -79,10 +79,10 @@ func (s *systemConfigurator) applyDNSConfig(config HostDNSConfig, stateManager *
 			continue
 		}
 		if dConf.MatchOnly {
-			matchDomains = append(matchDomains, dConf.Domain)
+			matchDomains = append(matchDomains, strings.TrimSuffix(dConf.Domain, "."))
 			continue
 		}
-		searchDomains = append(searchDomains, dConf.Domain)
+		searchDomains = append(searchDomains, strings.TrimSuffix(""+dConf.Domain, "."))
 	}
 
 	matchKey := getKeyWithInput(netbirdDNSStateKeyFormat, matchSuffix)

client/internal/dns/host_windows.go

@@ -17,15 +17,18 @@ import (
 
 var (
 	userenv = syscall.NewLazyDLL("userenv.dll")
+	dnsapi  = syscall.NewLazyDLL("dnsapi.dll")
 
 	// https://learn.microsoft.com/en-us/windows/win32/api/userenv/nf-userenv-refreshpolicyex
 	refreshPolicyExFn = userenv.NewProc("RefreshPolicyEx")
+
+	dnsFlushResolverCacheFn = dnsapi.NewProc("DnsFlushResolverCache")
 )
 
 const (
 	dnsPolicyConfigMatchPath    = `SYSTEM\CurrentControlSet\Services\Dnscache\Parameters\DnsPolicyConfig\NetBird-Match`
-	gpoDnsPolicyRoot            = `SOFTWARE\Policies\Microsoft\Windows NT\DNSClient`
-	gpoDnsPolicyConfigMatchPath = gpoDnsPolicyRoot + `\DnsPolicyConfig\NetBird-Match`
+	gpoDnsPolicyRoot            = `SOFTWARE\Policies\Microsoft\Windows NT\DNSClient\DnsPolicyConfig`
+	gpoDnsPolicyConfigMatchPath = gpoDnsPolicyRoot + `\NetBird-Match`
 
 	dnsPolicyConfigVersionKey           = "Version"
 	dnsPolicyConfigVersionValue         = 2
@@ -97,9 +100,9 @@ func (r *registryConfigurator) applyDNSConfig(config HostDNSConfig, stateManager
 			continue
 		}
 		if !dConf.MatchOnly {
-			searchDomains = append(searchDomains, dConf.Domain)
+			searchDomains = append(searchDomains, strings.TrimSuffix(dConf.Domain, "."))
 		}
-		matchDomains = append(matchDomains, "."+dConf.Domain)
+		matchDomains = append(matchDomains, "."+strings.TrimSuffix(dConf.Domain, "."))
 	}
 
 	if len(matchDomains) != 0 {
@@ -116,6 +119,10 @@ func (r *registryConfigurator) applyDNSConfig(config HostDNSConfig, stateManager
 		return fmt.Errorf("update search domains: %w", err)
 	}
 
+	if err := r.flushDNSCache(); err != nil {
+		log.Errorf("failed to flush DNS cache: %v", err)
+	}
+
 	return nil
 }
 
@@ -136,10 +143,6 @@ func (r *registryConfigurator) addDNSMatchPolicy(domains []string, ip string) er
 			return fmt.Errorf("configure GPO DNS policy: %w", err)
 		}
 
-		if err := r.configureDNSPolicy(dnsPolicyConfigMatchPath, domains, ip); err != nil {
-			return fmt.Errorf("configure local DNS policy: %w", err)
-		}
-
 		if err := refreshGroupPolicy(); err != nil {
 			log.Warnf("failed to refresh group policy: %v", err)
 		}
@@ -188,6 +191,26 @@ func (r *registryConfigurator) string() string {
 	return "registry"
 }
 
+func (r *registryConfigurator) flushDNSCache() error {
+	// dnsFlushResolverCacheFn.Call() may panic if the func is not found
+	defer func() {
+		if rec := recover(); rec != nil {
+			log.Errorf("Recovered from panic in flushDNSCache: %v", rec)
+		}
+	}()
+
+	ret, _, err := dnsFlushResolverCacheFn.Call()
+	if ret == 0 {
+		if err != nil && !errors.Is(err, syscall.Errno(0)) {
+			return fmt.Errorf("DnsFlushResolverCache failed: %w", err)
+		}
+		return fmt.Errorf("DnsFlushResolverCache failed")
+	}
+
+	log.Info("flushed DNS cache")
+	return nil
+}
+
 func (r *registryConfigurator) updateSearchDomains(domains []string) error {
 	if err := r.setInterfaceRegistryKeyStringValue(interfaceConfigSearchListKey, strings.Join(domains, ",")); err != nil {
 		return fmt.Errorf("update search domains: %w", err)
@@ -240,6 +263,10 @@ func (r *registryConfigurator) restoreHostDNS() error {
 		return fmt.Errorf("remove interface registry key: %w", err)
 	}
 
+	if err := r.flushDNSCache(); err != nil {
+		log.Errorf("failed to flush DNS cache: %v", err)
+	}
+
 	return nil
 }
 

client/internal/dns/mock_server.go

@@ -6,6 +6,7 @@ import (
 	"github.com/miekg/dns"
 
 	nbdns "github.com/netbirdio/netbird/dns"
+	"github.com/netbirdio/netbird/management/domain"
 )
 
 // MockServer is the mock instance of a dns server
@@ -13,17 +14,17 @@ type MockServer struct {
 	InitializeFunc        func() error
 	StopFunc              func()
 	UpdateDNSServerFunc   func(serial uint64, update nbdns.Config) error
-	RegisterHandlerFunc   func([]string, dns.Handler, int)
-	DeregisterHandlerFunc func([]string, int)
+	RegisterHandlerFunc   func(domain.List, dns.Handler, int)
+	DeregisterHandlerFunc func(domain.List, int)
 }
 
-func (m *MockServer) RegisterHandler(domains []string, handler dns.Handler, priority int) {
+func (m *MockServer) RegisterHandler(domains domain.List, handler dns.Handler, priority int) {
 	if m.RegisterHandlerFunc != nil {
 		m.RegisterHandlerFunc(domains, handler, priority)
 	}
 }
 
-func (m *MockServer) DeregisterHandler(domains []string, priority int) {
+func (m *MockServer) DeregisterHandler(domains domain.List, priority int) {
 	if m.DeregisterHandlerFunc != nil {
 		m.DeregisterHandlerFunc(domains, priority)
 	}

client/internal/dns/network_manager_unix.go

@@ -13,7 +13,6 @@ import (
 
 	"github.com/godbus/dbus/v5"
 	"github.com/hashicorp/go-version"
-	"github.com/miekg/dns"
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/client/internal/statemanager"
@@ -126,10 +125,10 @@ func (n *networkManagerDbusConfigurator) applyDNSConfig(config HostDNSConfig, st
 			continue
 		}
 		if dConf.MatchOnly {
-			matchDomains = append(matchDomains, "~."+dns.Fqdn(dConf.Domain))
+			matchDomains = append(matchDomains, "~."+dConf.Domain)
 			continue
 		}
-		searchDomains = append(searchDomains, dns.Fqdn(dConf.Domain))
+		searchDomains = append(searchDomains, dConf.Domain)
 	}
 
 	newDomainList := append(searchDomains, matchDomains...) //nolint:gocritic

client/internal/dns/server.go

@@ -6,11 +6,13 @@ import (
 	"fmt"
 	"net/netip"
 	"runtime"
+	"strings"
 	"sync"
 
 	"github.com/miekg/dns"
 	"github.com/mitchellh/hashstructure/v2"
 	log "github.com/sirupsen/logrus"
+	"golang.org/x/exp/maps"
 
 	"github.com/netbirdio/netbird/client/iface/netstack"
 	"github.com/netbirdio/netbird/client/internal/listener"
@@ -18,6 +20,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 	cProto "github.com/netbirdio/netbird/client/proto"
 	nbdns "github.com/netbirdio/netbird/dns"
+	"github.com/netbirdio/netbird/management/domain"
 )
 
 // ReadyListener is a notification mechanism what indicate the server is ready to handle host dns address changes
@@ -32,8 +35,8 @@ type IosDnsManager interface {
 
 // Server is a dns server interface
 type Server interface {
-	RegisterHandler(domains []string, handler dns.Handler, priority int)
-	DeregisterHandler(domains []string, priority int)
+	RegisterHandler(domains domain.List, handler dns.Handler, priority int)
+	DeregisterHandler(domains domain.List, priority int)
 	Initialize() error
 	Stop()
 	DnsIP() string
@@ -65,6 +68,7 @@ type DefaultServer struct {
 	previousConfigHash uint64
 	currentConfig      HostDNSConfig
 	handlerChain       *HandlerChain
+	extraDomains       map[domain.Domain]int
 
 	// permanent related properties
 	permanent      bool
@@ -164,13 +168,15 @@ func newDefaultServer(
 	stateManager *statemanager.Manager,
 	disableSys bool,
 ) *DefaultServer {
+	handlerChain := NewHandlerChain()
 	ctx, stop := context.WithCancel(ctx)
 	defaultServer := &DefaultServer{
 		ctx:          ctx,
 		ctxCancel:    stop,
 		disableSys:   disableSys,
 		service:      dnsService,
-		handlerChain: NewHandlerChain(),
+		handlerChain: handlerChain,
+		extraDomains: make(map[domain.Domain]int),
 		dnsMuxMap:    make(registeredHandlerMap),
 		localResolver: &localResolver{
 			registeredMap: make(registrationMap),
@@ -181,14 +187,26 @@ func newDefaultServer(
 		hostsDNSHolder: newHostsDNSHolder(),
 	}
 
+	// register with root zone, handler chain takes care of the routing
+	dnsService.RegisterMux(".", handlerChain)
+
 	return defaultServer
 }
 
-func (s *DefaultServer) RegisterHandler(domains []string, handler dns.Handler, priority int) {
+// RegisterHandler registers a handler for the given domains with the given priority.
+// Any previously registered handler for the same domain and priority will be replaced.
+func (s *DefaultServer) RegisterHandler(domains domain.List, handler dns.Handler, priority int) {
 	s.mux.Lock()
 	defer s.mux.Unlock()
 
-	s.registerHandler(domains, handler, priority)
+	s.registerHandler(domains.ToPunycodeList(), handler, priority)
+
+	// TODO: This will take over zones for non-wildcard domains, for which we might not have a handler in the chain
+	for _, domain := range domains {
+		// convert to zone with simple ref counter
+		s.extraDomains[toZone(domain)]++
+	}
+	s.applyHostConfig()
 }
 
 func (s *DefaultServer) registerHandler(domains []string, handler dns.Handler, priority int) {
@@ -200,15 +218,23 @@ func (s *DefaultServer) registerHandler(domains []string, handler dns.Handler, p
 			continue
 		}
 		s.handlerChain.AddHandler(domain, handler, priority)
-		s.service.RegisterMux(nbdns.NormalizeZone(domain), s.handlerChain)
 	}
 }
 
-func (s *DefaultServer) DeregisterHandler(domains []string, priority int) {
+// DeregisterHandler deregisters the handler for the given domains with the given priority.
+func (s *DefaultServer) DeregisterHandler(domains domain.List, priority int) {
 	s.mux.Lock()
 	defer s.mux.Unlock()
 
-	s.deregisterHandler(domains, priority)
+	s.deregisterHandler(domains.ToPunycodeList(), priority)
+	for _, domain := range domains {
+		zone := toZone(domain)
+		s.extraDomains[zone]--
+		if s.extraDomains[zone] <= 0 {
+			delete(s.extraDomains, zone)
+		}
+	}
+	s.applyHostConfig()
 }
 
 func (s *DefaultServer) deregisterHandler(domains []string, priority int) {
@@ -221,11 +247,6 @@ func (s *DefaultServer) deregisterHandler(domains []string, priority int) {
 		}
 
 		s.handlerChain.RemoveHandler(domain, priority)
-
-		// Only deregister from service if no handlers remain
-		if !s.handlerChain.HasHandlers(domain) {
-			s.service.DeregisterMux(nbdns.NormalizeZone(domain))
-		}
 	}
 }
 
@@ -286,6 +307,8 @@ func (s *DefaultServer) Stop() {
 	}
 
 	s.service.Stop()
+
+	maps.Clear(s.extraDomains)
 }
 
 // OnUpdatedHostDNSServer update the DNS servers addresses for root zones
@@ -390,7 +413,9 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 	// is the service should be Disabled, we stop the listener or fake resolver
 	// and proceed with a regular update to clean up the handlers and records
 	if update.ServiceEnable {
-		_ = s.service.Listen()
+		if err := s.service.Listen(); err != nil {
+			log.Errorf("failed to start DNS service: %v", err)
+		}
 	} else if !s.permanent {
 		s.service.Stop()
 	}
@@ -413,17 +438,13 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 
 	s.currentConfig = dnsConfigToHostDNSConfig(update, s.service.RuntimeIP(), s.service.RuntimePort())
 
-	hostUpdate := s.currentConfig
 	if s.service.RuntimePort() != defaultPort && !s.hostManager.supportCustomPort() {
 		log.Warnf("the DNS manager of this peer doesn't support custom port. Disabling primary DNS setup. " +
 			"Learn more at: https://docs.netbird.io/how-to/manage-dns-in-your-network#local-resolver")
-		hostUpdate.RouteAll = false
+		s.currentConfig.RouteAll = false
 	}
 
-	if err = s.hostManager.applyDNSConfig(hostUpdate, s.stateManager); err != nil {
-		log.Error(err)
-		s.handleErrNoGroupaAll(err)
-	}
+	s.applyHostConfig()
 
 	go func() {
 		// persist dns state right away
@@ -441,6 +462,38 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 	return nil
 }
 
+func (s *DefaultServer) applyHostConfig() {
+	if s.hostManager == nil {
+		return
+	}
+
+	config := s.currentConfig
+
+	existingDomains := make(map[string]struct{})
+	for _, d := range config.Domains {
+		existingDomains[d.Domain] = struct{}{}
+	}
+
+	// add extra domains only if they're not already in the config
+	for domain := range s.extraDomains {
+		domainStr := domain.PunycodeString()
+
+		if _, exists := existingDomains[domainStr]; !exists {
+			config.Domains = append(config.Domains, DomainConfig{
+				Domain:    domainStr,
+				MatchOnly: true,
+			})
+		}
+	}
+
+	log.Debugf("extra match domains: %v", s.extraDomains)
+
+	if err := s.hostManager.applyDNSConfig(config, s.stateManager); err != nil {
+		log.Errorf("failed to apply DNS host manager update: %v", err)
+		s.handleErrNoGroupaAll(err)
+	}
+}
+
 func (s *DefaultServer) handleErrNoGroupaAll(err error) {
 	if !errors.Is(ErrRouteAllWithoutNameserverGroup, err) {
 		return
@@ -690,10 +743,7 @@ func (s *DefaultServer) upstreamCallbacks(
 			}
 		}
 
-		if err := s.hostManager.applyDNSConfig(s.currentConfig, s.stateManager); err != nil {
-			s.handleErrNoGroupaAll(err)
-			l.Errorf("Failed to apply nameserver deactivation on the host: %v", err)
-		}
+		s.applyHostConfig()
 
 		go func() {
 			if err := s.stateManager.PersistState(s.ctx); err != nil {
@@ -728,12 +778,7 @@ func (s *DefaultServer) upstreamCallbacks(
 			s.registerHandler([]string{nbdns.RootZone}, handler, priority)
 		}
 
-		if s.hostManager != nil {
-			if err := s.hostManager.applyDNSConfig(s.currentConfig, s.stateManager); err != nil {
-				s.handleErrNoGroupaAll(err)
-				l.WithError(err).Error("reactivate temporary disabled nameserver group, DNS update apply")
-			}
-		}
+		s.applyHostConfig()
 
 		s.updateNSState(nsGroup, nil, true)
 	}
@@ -836,3 +881,13 @@ func groupNSGroupsByDomain(nsGroups []*nbdns.NameServerGroup) []nsGroupsByDomain
 
 	return result
 }
+
+func toZone(d domain.Domain) domain.Domain {
+	return domain.Domain(
+		nbdns.NormalizeZone(
+			dns.Fqdn(
+				strings.ToLower(d.PunycodeString()),
+			),
+		),
+	)
+}

client/internal/dns/server_test.go

@@ -23,19 +23,23 @@ 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"
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/formatter"
+	"github.com/netbirdio/netbird/management/domain"
 )
 
+var flowLogger = netflow.NewManager(nil, []byte{}, nil).GetLogger()
+
 type mocWGIface struct {
 	filter device.PacketFilter
 }
 
 func (w *mocWGIface) Name() string {
-	panic("implement me")
+	return "utun2301"
 }
 
 func (w *mocWGIface) Address() wgaddr.Address {
@@ -456,7 +460,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)
@@ -917,7 +921,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
@@ -1445,3 +1449,497 @@ func TestDefaultServer_UpdateMux(t *testing.T) {
 		})
 	}
 }
+
+func TestExtraDomains(t *testing.T) {
+	tests := []struct {
+		name                string
+		initialConfig       nbdns.Config
+		registerDomains     []domain.List
+		deregisterDomains   []domain.List
+		finalConfig         nbdns.Config
+		expectedDomains     []string
+		expectedMatchOnly   []string
+		applyHostConfigCall int
+	}{
+		{
+			name: "Register domains before config update",
+			registerDomains: []domain.List{
+				{"extra1.example.com", "extra2.example.com"},
+			},
+			initialConfig: nbdns.Config{
+				ServiceEnable: true,
+				CustomZones: []nbdns.CustomZone{
+					{Domain: "config.example.com"},
+				},
+			},
+			expectedDomains: []string{
+				"config.example.com.",
+				"extra1.example.com.",
+				"extra2.example.com.",
+			},
+			expectedMatchOnly: []string{
+				"extra1.example.com.",
+				"extra2.example.com.",
+			},
+			applyHostConfigCall: 2,
+		},
+		{
+			name: "Register domains after config update",
+			initialConfig: nbdns.Config{
+				ServiceEnable: true,
+				CustomZones: []nbdns.CustomZone{
+					{Domain: "config.example.com"},
+				},
+			},
+			registerDomains: []domain.List{
+				{"extra1.example.com", "extra2.example.com"},
+			},
+			expectedDomains: []string{
+				"config.example.com.",
+				"extra1.example.com.",
+				"extra2.example.com.",
+			},
+			expectedMatchOnly: []string{
+				"extra1.example.com.",
+				"extra2.example.com.",
+			},
+			applyHostConfigCall: 2,
+		},
+		{
+			name: "Register overlapping domains",
+			initialConfig: nbdns.Config{
+				ServiceEnable: true,
+				CustomZones: []nbdns.CustomZone{
+					{Domain: "config.example.com"},
+					{Domain: "overlap.example.com"},
+				},
+			},
+			registerDomains: []domain.List{
+				{"extra.example.com", "overlap.example.com"},
+			},
+			expectedDomains: []string{
+				"config.example.com.",
+				"overlap.example.com.",
+				"extra.example.com.",
+			},
+			expectedMatchOnly: []string{
+				"extra.example.com.",
+			},
+			applyHostConfigCall: 2,
+		},
+		{
+			name: "Register and deregister domains",
+			initialConfig: nbdns.Config{
+				ServiceEnable: true,
+				CustomZones: []nbdns.CustomZone{
+					{Domain: "config.example.com"},
+				},
+			},
+			registerDomains: []domain.List{
+				{"extra1.example.com", "extra2.example.com"},
+				{"extra3.example.com", "extra4.example.com"},
+			},
+			deregisterDomains: []domain.List{
+				{"extra1.example.com", "extra3.example.com"},
+			},
+			expectedDomains: []string{
+				"config.example.com.",
+				"extra2.example.com.",
+				"extra4.example.com.",
+			},
+			expectedMatchOnly: []string{
+				"extra2.example.com.",
+				"extra4.example.com.",
+			},
+			applyHostConfigCall: 4,
+		},
+		{
+			name: "Register domains with ref counter",
+			initialConfig: nbdns.Config{
+				ServiceEnable: true,
+				CustomZones: []nbdns.CustomZone{
+					{Domain: "config.example.com"},
+				},
+			},
+			registerDomains: []domain.List{
+				{"extra.example.com", "duplicate.example.com"},
+				{"other.example.com", "duplicate.example.com"},
+			},
+			deregisterDomains: []domain.List{
+				{"duplicate.example.com"},
+			},
+			expectedDomains: []string{
+				"config.example.com.",
+				"extra.example.com.",
+				"other.example.com.",
+				"duplicate.example.com.",
+			},
+			expectedMatchOnly: []string{
+				"extra.example.com.",
+				"other.example.com.",
+				"duplicate.example.com.",
+			},
+			applyHostConfigCall: 4,
+		},
+		{
+			name: "Config update with new domains after registration",
+			initialConfig: nbdns.Config{
+				ServiceEnable: true,
+				CustomZones: []nbdns.CustomZone{
+					{Domain: "config.example.com"},
+				},
+			},
+			registerDomains: []domain.List{
+				{"extra.example.com", "duplicate.example.com"},
+			},
+			finalConfig: nbdns.Config{
+				ServiceEnable: true,
+				CustomZones: []nbdns.CustomZone{
+					{Domain: "config.example.com"},
+					{Domain: "newconfig.example.com"},
+				},
+			},
+			expectedDomains: []string{
+				"config.example.com.",
+				"newconfig.example.com.",
+				"extra.example.com.",
+				"duplicate.example.com.",
+			},
+			expectedMatchOnly: []string{
+				"extra.example.com.",
+				"duplicate.example.com.",
+			},
+			applyHostConfigCall: 3,
+		},
+		{
+			name: "Deregister domain that is part of customZones",
+			initialConfig: nbdns.Config{
+				ServiceEnable: true,
+				CustomZones: []nbdns.CustomZone{
+					{Domain: "config.example.com"},
+					{Domain: "protected.example.com"},
+				},
+			},
+			registerDomains: []domain.List{
+				{"extra.example.com", "protected.example.com"},
+			},
+			deregisterDomains: []domain.List{
+				{"protected.example.com"},
+			},
+			expectedDomains: []string{
+				"extra.example.com.",
+				"config.example.com.",
+				"protected.example.com.",
+			},
+			expectedMatchOnly: []string{
+				"extra.example.com.",
+			},
+			applyHostConfigCall: 3,
+		},
+		{
+			name: "Register domain that is part of nameserver group",
+			initialConfig: nbdns.Config{
+				ServiceEnable: true,
+				NameServerGroups: []*nbdns.NameServerGroup{
+					{
+						Domains: []string{"ns.example.com", "overlap.ns.example.com"},
+						NameServers: []nbdns.NameServer{
+							{
+								IP:     netip.MustParseAddr("8.8.8.8"),
+								NSType: nbdns.UDPNameServerType,
+								Port:   53,
+							},
+						},
+					},
+				},
+			},
+			registerDomains: []domain.List{
+				{"extra.example.com", "overlap.ns.example.com"},
+			},
+			expectedDomains: []string{
+				"ns.example.com.",
+				"overlap.ns.example.com.",
+				"extra.example.com.",
+			},
+			expectedMatchOnly: []string{
+				"ns.example.com.",
+				"overlap.ns.example.com.",
+				"extra.example.com.",
+			},
+			applyHostConfigCall: 2,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var capturedConfigs []HostDNSConfig
+			mockHostConfig := &mockHostConfigurator{
+				applyDNSConfigFunc: func(config HostDNSConfig, _ *statemanager.Manager) error {
+					capturedConfigs = append(capturedConfigs, config)
+					return nil
+				},
+				restoreHostDNSFunc: func() error {
+					return nil
+				},
+				supportCustomPortFunc: func() bool {
+					return true
+				},
+				stringFunc: func() string {
+					return "mock"
+				},
+			}
+
+			mockSvc := &mockService{}
+
+			server := &DefaultServer{
+				ctx:            context.Background(),
+				handlerChain:   NewHandlerChain(),
+				wgInterface:    &mocWGIface{},
+				hostManager:    mockHostConfig,
+				localResolver:  &localResolver{},
+				service:        mockSvc,
+				statusRecorder: peer.NewRecorder("test"),
+				extraDomains:   make(map[domain.Domain]int),
+			}
+
+			// Apply initial configuration
+			if tt.initialConfig.ServiceEnable {
+				err := server.applyConfiguration(tt.initialConfig)
+				assert.NoError(t, err)
+			}
+
+			// Register domains
+			for _, domains := range tt.registerDomains {
+				server.RegisterHandler(domains, &MockHandler{}, PriorityDefault)
+			}
+
+			// Deregister domains if specified
+			for _, domains := range tt.deregisterDomains {
+				server.DeregisterHandler(domains, PriorityDefault)
+			}
+
+			// Apply final configuration if specified
+			if tt.finalConfig.ServiceEnable {
+				err := server.applyConfiguration(tt.finalConfig)
+				assert.NoError(t, err)
+			}
+
+			// Verify number of calls
+			assert.Equal(t, tt.applyHostConfigCall, len(capturedConfigs),
+				"Expected %d calls to applyDNSConfig, got %d", tt.applyHostConfigCall, len(capturedConfigs))
+
+			// Get the last applied config
+			lastConfig := capturedConfigs[len(capturedConfigs)-1]
+
+			// Check all expected domains are present
+			domainMap := make(map[string]bool)
+			matchOnlyMap := make(map[string]bool)
+
+			for _, d := range lastConfig.Domains {
+				domainMap[d.Domain] = true
+				if d.MatchOnly {
+					matchOnlyMap[d.Domain] = true
+				}
+			}
+
+			// Verify expected domains
+			for _, d := range tt.expectedDomains {
+				assert.True(t, domainMap[d], "Expected domain %s not found in final config", d)
+			}
+
+			// Verify match-only domains
+			for _, d := range tt.expectedMatchOnly {
+				assert.True(t, matchOnlyMap[d], "Expected match-only domain %s not found in final config", d)
+			}
+
+			// Verify no unexpected domains
+			assert.Equal(t, len(tt.expectedDomains), len(domainMap), "Unexpected number of domains in final config")
+			assert.Equal(t, len(tt.expectedMatchOnly), len(matchOnlyMap), "Unexpected number of match-only domains in final config")
+		})
+	}
+}
+
+func TestExtraDomainsRefCounting(t *testing.T) {
+	mockHostConfig := &mockHostConfigurator{
+		applyDNSConfigFunc: func(config HostDNSConfig, _ *statemanager.Manager) error {
+			return nil
+		},
+		restoreHostDNSFunc: func() error {
+			return nil
+		},
+		supportCustomPortFunc: func() bool {
+			return true
+		},
+		stringFunc: func() string {
+			return "mock"
+		},
+	}
+
+	mockSvc := &mockService{}
+
+	server := &DefaultServer{
+		ctx:            context.Background(),
+		handlerChain:   NewHandlerChain(),
+		hostManager:    mockHostConfig,
+		localResolver:  &localResolver{},
+		service:        mockSvc,
+		statusRecorder: peer.NewRecorder("test"),
+		extraDomains:   make(map[domain.Domain]int),
+	}
+
+	// Register domains from different handlers with same domain
+	server.RegisterHandler(domain.List{"*.shared.example.com"}, &MockHandler{}, PriorityDNSRoute)
+	server.RegisterHandler(domain.List{"shared.example.com."}, &MockHandler{}, PriorityMatchDomain)
+
+	// Verify refcount is 2
+	zoneKey := toZone("shared.example.com")
+	assert.Equal(t, 2, server.extraDomains[zoneKey], "Refcount should be 2 after registering same domain twice")
+
+	// Deregister one handler
+	server.DeregisterHandler(domain.List{"shared.example.com"}, PriorityMatchDomain)
+
+	// Verify refcount is 1
+	assert.Equal(t, 1, server.extraDomains[zoneKey], "Refcount should be 1 after deregistering one handler")
+
+	// Deregister the other handler
+	server.DeregisterHandler(domain.List{"shared.example.com"}, PriorityDNSRoute)
+
+	// Verify domain is removed
+	_, exists := server.extraDomains[zoneKey]
+	assert.False(t, exists, "Domain should be removed after deregistering all handlers")
+}
+
+func TestUpdateConfigWithExistingExtraDomains(t *testing.T) {
+	var capturedConfig HostDNSConfig
+	mockHostConfig := &mockHostConfigurator{
+		applyDNSConfigFunc: func(config HostDNSConfig, _ *statemanager.Manager) error {
+			capturedConfig = config
+			return nil
+		},
+		restoreHostDNSFunc: func() error {
+			return nil
+		},
+		supportCustomPortFunc: func() bool {
+			return true
+		},
+		stringFunc: func() string {
+			return "mock"
+		},
+	}
+
+	mockSvc := &mockService{}
+
+	server := &DefaultServer{
+		ctx:            context.Background(),
+		handlerChain:   NewHandlerChain(),
+		hostManager:    mockHostConfig,
+		localResolver:  &localResolver{},
+		service:        mockSvc,
+		statusRecorder: peer.NewRecorder("test"),
+		extraDomains:   make(map[domain.Domain]int),
+	}
+
+	server.RegisterHandler(domain.List{"extra.example.com"}, &MockHandler{}, PriorityDefault)
+
+	initialConfig := nbdns.Config{
+		ServiceEnable: true,
+		CustomZones: []nbdns.CustomZone{
+			{Domain: "config.example.com"},
+		},
+	}
+	err := server.applyConfiguration(initialConfig)
+	assert.NoError(t, err)
+
+	var domains []string
+	for _, d := range capturedConfig.Domains {
+		domains = append(domains, d.Domain)
+	}
+	assert.Contains(t, domains, "config.example.com.")
+	assert.Contains(t, domains, "extra.example.com.")
+
+	// Now apply a new configuration with overlapping domain
+	updatedConfig := nbdns.Config{
+		ServiceEnable: true,
+		CustomZones: []nbdns.CustomZone{
+			{Domain: "config.example.com"},
+			{Domain: "extra.example.com"},
+		},
+	}
+	err = server.applyConfiguration(updatedConfig)
+	assert.NoError(t, err)
+
+	// Verify both domains are in config, but no duplicates
+	domains = []string{}
+	matchOnlyCount := 0
+	for _, d := range capturedConfig.Domains {
+		domains = append(domains, d.Domain)
+		if d.MatchOnly {
+			matchOnlyCount++
+		}
+	}
+
+	assert.Contains(t, domains, "config.example.com.")
+	assert.Contains(t, domains, "extra.example.com.")
+	assert.Equal(t, 2, len(domains), "Should have exactly 2 domains with no duplicates")
+
+	// Extra domain should no longer be marked as match-only when in config
+	matchOnlyDomain := ""
+	for _, d := range capturedConfig.Domains {
+		if d.Domain == "extra.example.com." && d.MatchOnly {
+			matchOnlyDomain = d.Domain
+			break
+		}
+	}
+	assert.Empty(t, matchOnlyDomain, "Domain should not be match-only when included in config")
+}
+
+func TestDomainCaseHandling(t *testing.T) {
+	var capturedConfig HostDNSConfig
+	mockHostConfig := &mockHostConfigurator{
+		applyDNSConfigFunc: func(config HostDNSConfig, _ *statemanager.Manager) error {
+			capturedConfig = config
+			return nil
+		},
+		restoreHostDNSFunc: func() error {
+			return nil
+		},
+		supportCustomPortFunc: func() bool {
+			return true
+		},
+		stringFunc: func() string {
+			return "mock"
+		},
+	}
+
+	mockSvc := &mockService{}
+	server := &DefaultServer{
+		ctx:            context.Background(),
+		handlerChain:   NewHandlerChain(),
+		hostManager:    mockHostConfig,
+		localResolver:  &localResolver{},
+		service:        mockSvc,
+		statusRecorder: peer.NewRecorder("test"),
+		extraDomains:   make(map[domain.Domain]int),
+	}
+
+	server.RegisterHandler(domain.List{"MIXED.example.com"}, &MockHandler{}, PriorityDefault)
+	server.RegisterHandler(domain.List{"mixed.EXAMPLE.com"}, &MockHandler{}, PriorityMatchDomain)
+
+	assert.Equal(t, 1, len(server.extraDomains), "Case differences should be normalized")
+
+	config := nbdns.Config{
+		ServiceEnable: true,
+		CustomZones: []nbdns.CustomZone{
+			{Domain: "config.example.com"},
+		},
+	}
+	err := server.applyConfiguration(config)
+	assert.NoError(t, err)
+
+	var domains []string
+	for _, d := range capturedConfig.Domains {
+		domains = append(domains, d.Domain)
+	}
+	assert.Contains(t, domains, "config.example.com.", "Mixed case domain should be normalized and pre.sent")
+	assert.Contains(t, domains, "mixed.example.com.", "Mixed case domain should be normalized and present")
+}

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

@@ -11,7 +11,6 @@ import (
 	"time"
 
 	"github.com/godbus/dbus/v5"
-	"github.com/miekg/dns"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/sys/unix"
 
@@ -38,7 +37,6 @@ const (
 
 type systemdDbusConfigurator struct {
 	dbusLinkObject dbus.ObjectPath
-	routingAll     bool
 	ifaceName      string
 }
 
@@ -112,7 +110,7 @@ func (s *systemdDbusConfigurator) applyDNSConfig(config HostDNSConfig, stateMana
 			continue
 		}
 		domainsInput = append(domainsInput, systemdDbusLinkDomainsInput{
-			Domain:    dns.Fqdn(dConf.Domain),
+			Domain:    dConf.Domain,
 			MatchOnly: dConf.MatchOnly,
 		})
 
@@ -124,18 +122,19 @@ func (s *systemdDbusConfigurator) applyDNSConfig(config HostDNSConfig, stateMana
 	}
 
 	if config.RouteAll {
-		log.Infof("configured %s:%d as main DNS forwarder for this peer", config.ServerIP, config.ServerPort)
 		err = s.callLinkMethod(systemdDbusSetDefaultRouteMethodSuffix, true)
 		if err != nil {
-			return fmt.Errorf("setting link as default dns router, failed with error: %w", err)
+			return fmt.Errorf("set link as default dns router: %w", err)
 		}
 		domainsInput = append(domainsInput, systemdDbusLinkDomainsInput{
 			Domain:    nbdns.RootZone,
 			MatchOnly: true,
 		})
-		s.routingAll = true
-	} else if s.routingAll {
-		log.Infof("removing %s:%d as main DNS forwarder for this peer", config.ServerIP, config.ServerPort)
+		log.Infof("configured %s:%d as main DNS forwarder for this peer", config.ServerIP, config.ServerPort)
+	} else {
+		if err = s.callLinkMethod(systemdDbusSetDefaultRouteMethodSuffix, false); err != nil {
+			return fmt.Errorf("remove link as default dns router: %w", err)
+		}
 	}
 
 	state := &ShutdownState{
@@ -151,6 +150,11 @@ func (s *systemdDbusConfigurator) applyDNSConfig(config HostDNSConfig, stateMana
 	if err != nil {
 		log.Error(err)
 	}
+
+	if err := s.flushDNSCache(); err != nil {
+		log.Errorf("failed to flush DNS cache: %v", err)
+	}
+
 	return nil
 }
 
@@ -163,7 +167,8 @@ func (s *systemdDbusConfigurator) setDomainsForInterface(domainsInput []systemdD
 	if err != nil {
 		return fmt.Errorf("setting domains configuration failed with error: %w", err)
 	}
-	return s.flushCaches()
+
+	return nil
 }
 
 func (s *systemdDbusConfigurator) restoreHostDNS() error {
@@ -183,10 +188,14 @@ func (s *systemdDbusConfigurator) restoreHostDNS() error {
 		return fmt.Errorf("unable to revert link configuration, got error: %w", err)
 	}
 
-	return s.flushCaches()
+	if err := s.flushDNSCache(); err != nil {
+		log.Errorf("failed to flush DNS cache: %v", err)
+	}
+
+	return nil
 }
 
-func (s *systemdDbusConfigurator) flushCaches() error {
+func (s *systemdDbusConfigurator) flushDNSCache() error {
 	obj, closeConn, err := getDbusObject(systemdResolvedDest, systemdDbusObjectNode)
 	if err != nil {
 		return fmt.Errorf("attempting to retrieve the object %s, err: %w", systemdDbusObjectNode, err)

client/internal/dns/upstream.go

@@ -23,9 +23,10 @@ import (
 )
 
 const (
+	UpstreamTimeout = 15 * time.Second
+
 	failsTillDeact   = int32(5)
 	reactivatePeriod = 30 * time.Second
-	upstreamTimeout  = 15 * time.Second
 	probeTimeout     = 2 * time.Second
 )
 
@@ -66,7 +67,7 @@ func newUpstreamResolverBase(ctx context.Context, statusRecorder *peer.Status, d
 		ctx:              ctx,
 		cancel:           cancel,
 		domain:           domain,
-		upstreamTimeout:  upstreamTimeout,
+		upstreamTimeout:  UpstreamTimeout,
 		reactivatePeriod: reactivatePeriod,
 		failsTillDeact:   failsTillDeact,
 		statusRecorder:   statusRecorder,

client/internal/dns/upstream_android.go

@@ -55,7 +55,7 @@ func (u *upstreamResolver) exchangeWithinVPN(ctx context.Context, upstream strin
 
 // exchangeWithoutVPN protect the UDP socket by Android SDK to avoid to goes through the VPN
 func (u *upstreamResolver) exchangeWithoutVPN(ctx context.Context, upstream string, r *dns.Msg) (rm *dns.Msg, t time.Duration, err error) {
-	timeout := upstreamTimeout
+	timeout := UpstreamTimeout
 	if deadline, ok := ctx.Deadline(); ok {
 		timeout = time.Until(deadline)
 	}

client/internal/dns/upstream_ios.go

@@ -52,7 +52,7 @@ func (u *upstreamResolverIOS) exchange(ctx context.Context, upstream string, r *
 		return nil, 0, fmt.Errorf("error while parsing upstream host: %s", err)
 	}
 
-	timeout := upstreamTimeout
+	timeout := UpstreamTimeout
 	if deadline, ok := ctx.Deadline(); ok {
 		timeout = time.Until(deadline)
 	}

client/internal/dns/upstream_test.go

@@ -26,7 +26,7 @@ func TestUpstreamResolver_ServeDNS(t *testing.T) {
 			name:           "Should Resolve A Record",
 			inputMSG:       new(dns.Msg).SetQuestion("one.one.one.one.", dns.TypeA),
 			InputServers:   []string{"8.8.8.8:53", "8.8.4.4:53"},
-			timeout:        upstreamTimeout,
+			timeout:        UpstreamTimeout,
 			expectedAnswer: "1.1.1.1",
 		},
 		{
@@ -48,7 +48,7 @@ func TestUpstreamResolver_ServeDNS(t *testing.T) {
 			inputMSG:            new(dns.Msg).SetQuestion("one.one.one.one.", dns.TypeA),
 			InputServers:        []string{"8.0.0.0:53", "8.8.4.4:53"},
 			cancelCTX:           true,
-			timeout:             upstreamTimeout,
+			timeout:             UpstreamTimeout,
 			responseShouldBeNil: true,
 		},
 	}
@@ -122,7 +122,7 @@ func TestUpstreamResolver_DeactivationReactivation(t *testing.T) {
 			r:   new(dns.Msg),
 			rtt: time.Millisecond,
 		},
-		upstreamTimeout:  upstreamTimeout,
+		upstreamTimeout:  UpstreamTimeout,
 		reactivatePeriod: reactivatePeriod,
 		failsTillDeact:   failsTillDeact,
 	}

client/internal/dnsfwd/forwarder.go

@@ -4,33 +4,43 @@ import (
 	"context"
 	"errors"
 	"net"
+	"net/netip"
+	"strings"
+	"sync"
+	"time"
 
 	"github.com/miekg/dns"
 	log "github.com/sirupsen/logrus"
 
+	"github.com/netbirdio/netbird/client/internal/peer"
 	nbdns "github.com/netbirdio/netbird/dns"
 )
 
 const errResolveFailed = "failed to resolve query for domain=%s: %v"
+const upstreamTimeout = 15 * time.Second
 
 type DNSForwarder struct {
-	listenAddress string
-	ttl           uint32
-	domains       []string
+	listenAddress  string
+	ttl            uint32
+	domains        []string
+	statusRecorder *peer.Status
 
 	dnsServer *dns.Server
 	mux       *dns.ServeMux
+
+	resId sync.Map
 }
 
-func NewDNSForwarder(listenAddress string, ttl uint32) *DNSForwarder {
+func NewDNSForwarder(listenAddress string, ttl uint32, statusRecorder *peer.Status) *DNSForwarder {
 	log.Debugf("creating DNS forwarder with listen_address=%s ttl=%d", listenAddress, ttl)
 	return &DNSForwarder{
-		listenAddress: listenAddress,
-		ttl:           ttl,
+		listenAddress:  listenAddress,
+		ttl:            ttl,
+		statusRecorder: statusRecorder,
 	}
 }
 
-func (f *DNSForwarder) Listen(domains []string) error {
+func (f *DNSForwarder) Listen(domains []string, resIds map[string]string) error {
 	log.Infof("listen DNS forwarder on address=%s", f.listenAddress)
 	mux := dns.NewServeMux()
 
@@ -42,22 +52,31 @@ func (f *DNSForwarder) Listen(domains []string) error {
 	f.dnsServer = dnsServer
 	f.mux = mux
 
-	f.UpdateDomains(domains)
+	f.UpdateDomains(domains, resIds)
 
 	return dnsServer.ListenAndServe()
 }
 
-func (f *DNSForwarder) UpdateDomains(domains []string) {
+func (f *DNSForwarder) UpdateDomains(domains []string, resIds map[string]string) {
 	log.Debugf("Updating domains from %v to %v", f.domains, domains)
 
 	for _, d := range f.domains {
 		f.mux.HandleRemove(d)
+		f.statusRecorder.RemoveResolvedIPLookupEntry(d)
 	}
+	f.resId.Clear()
 
 	newDomains := filterDomains(domains)
 	for _, d := range newDomains {
 		f.mux.HandleFunc(d, f.handleDNSQuery)
 	}
+
+	for domain, resId := range resIds {
+		if domain != "" {
+			f.resId.Store(domain, resId)
+		}
+	}
+
 	f.domains = newDomains
 }
 
@@ -79,41 +98,87 @@ func (f *DNSForwarder) handleDNSQuery(w dns.ResponseWriter, query *dns.Msg) {
 	domain := question.Name
 
 	resp := query.SetReply(query)
+	var network string
+	switch question.Qtype {
+	case dns.TypeA:
+		network = "ip4"
+	case dns.TypeAAAA:
+		network = "ip6"
+	default:
+		// TODO: Handle other types
 
-	ips, err := net.LookupIP(domain)
-	if err != nil {
-		var dnsErr *net.DNSError
-
-		switch {
-		case errors.As(err, &dnsErr):
-			resp.Rcode = dns.RcodeServerFailure
-			if dnsErr.IsNotFound {
-				// Pass through NXDOMAIN
-				resp.Rcode = dns.RcodeNameError
-			}
-
-			if dnsErr.Server != "" {
-				log.Warnf("failed to resolve query for domain=%s server=%s: %v", domain, dnsErr.Server, err)
-			} else {
-				log.Warnf(errResolveFailed, domain, err)
-			}
-		default:
-			resp.Rcode = dns.RcodeServerFailure
-			log.Warnf(errResolveFailed, domain, err)
-		}
-
+		resp.Rcode = dns.RcodeNotImplemented
 		if err := w.WriteMsg(resp); err != nil {
-			log.Errorf("failed to write failure DNS response: %v", err)
+			log.Errorf("failed to write DNS response: %v", err)
 		}
 		return
 	}
 
+	ctx, cancel := context.WithTimeout(context.Background(), upstreamTimeout)
+	defer cancel()
+	ips, err := net.DefaultResolver.LookupNetIP(ctx, network, domain)
+	if err != nil {
+		f.handleDNSError(w, resp, domain, err)
+		return
+	}
+
+	resId, ok := f.resId.Load(strings.TrimSuffix(domain, "."))
+	if ok {
+		for _, ip := range ips {
+			var ipWithSuffix string
+			if ip.Is4() {
+				ipWithSuffix = ip.String() + "/32"
+				log.Tracef("resolved domain=%s to IPv4=%s", domain, ipWithSuffix)
+			} else {
+				ipWithSuffix = ip.String() + "/128"
+				log.Tracef("resolved domain=%s to IPv6=%s", domain, ipWithSuffix)
+			}
+			f.statusRecorder.AddResolvedIPLookupEntry(ipWithSuffix, resId.(string))
+		}
+	}
+
+	f.addIPsToResponse(resp, domain, ips)
+
+	if err := w.WriteMsg(resp); err != nil {
+		log.Errorf("failed to write DNS response: %v", err)
+	}
+}
+
+// handleDNSError processes DNS lookup errors and sends an appropriate error response
+func (f *DNSForwarder) handleDNSError(w dns.ResponseWriter, resp *dns.Msg, domain string, err error) {
+	var dnsErr *net.DNSError
+
+	switch {
+	case errors.As(err, &dnsErr):
+		resp.Rcode = dns.RcodeServerFailure
+		if dnsErr.IsNotFound {
+			// Pass through NXDOMAIN
+			resp.Rcode = dns.RcodeNameError
+		}
+
+		if dnsErr.Server != "" {
+			log.Warnf("failed to resolve query for domain=%s server=%s: %v", domain, dnsErr.Server, err)
+		} else {
+			log.Warnf(errResolveFailed, domain, err)
+		}
+	default:
+		resp.Rcode = dns.RcodeServerFailure
+		log.Warnf(errResolveFailed, domain, err)
+	}
+
+	if err := w.WriteMsg(resp); err != nil {
+		log.Errorf("failed to write failure DNS response: %v", err)
+	}
+}
+
+// addIPsToResponse adds IP addresses to the DNS response as appropriate A or AAAA records
+func (f *DNSForwarder) addIPsToResponse(resp *dns.Msg, domain string, ips []netip.Addr) {
 	for _, ip := range ips {
 		var respRecord dns.RR
-		if ip.To4() == nil {
+		if ip.Is6() {
 			log.Tracef("resolved domain=%s to IPv6=%s", domain, ip)
 			rr := dns.AAAA{
-				AAAA: ip,
+				AAAA: ip.AsSlice(),
 				Hdr: dns.RR_Header{
 					Name:   domain,
 					Rrtype: dns.TypeAAAA,
@@ -125,7 +190,7 @@ func (f *DNSForwarder) handleDNSQuery(w dns.ResponseWriter, query *dns.Msg) {
 		} else {
 			log.Tracef("resolved domain=%s to IPv4=%s", domain, ip)
 			rr := dns.A{
-				A: ip,
+				A: ip.AsSlice(),
 				Hdr: dns.RR_Header{
 					Name:   domain,
 					Rrtype: dns.TypeA,
@@ -137,10 +202,6 @@ func (f *DNSForwarder) handleDNSQuery(w dns.ResponseWriter, query *dns.Msg) {
 		}
 		resp.Answer = append(resp.Answer, respRecord)
 	}
-
-	if err := w.WriteMsg(resp); err != nil {
-		log.Errorf("failed to write DNS response: %v", err)
-	}
 }
 
 // filterDomains returns a list of normalized domains

client/internal/dnsfwd/manager.go

@@ -10,6 +10,7 @@ import (
 
 	nberrors "github.com/netbirdio/netbird/client/errors"
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
+	"github.com/netbirdio/netbird/client/internal/peer"
 )
 
 const (
@@ -19,19 +20,21 @@ const (
 )
 
 type Manager struct {
-	firewall firewall.Manager
+	firewall       firewall.Manager
+	statusRecorder *peer.Status
 
 	fwRules      []firewall.Rule
 	dnsForwarder *DNSForwarder
 }
 
-func NewManager(fw firewall.Manager) *Manager {
+func NewManager(fw firewall.Manager, statusRecorder *peer.Status) *Manager {
 	return &Manager{
-		firewall: fw,
+		firewall:       fw,
+		statusRecorder: statusRecorder,
 	}
 }
 
-func (m *Manager) Start(domains []string) error {
+func (m *Manager) Start(domains []string, resIds map[string]string) error {
 	log.Infof("starting DNS forwarder")
 	if m.dnsForwarder != nil {
 		return nil
@@ -41,9 +44,9 @@ func (m *Manager) Start(domains []string) error {
 		return err
 	}
 
-	m.dnsForwarder = NewDNSForwarder(fmt.Sprintf(":%d", ListenPort), dnsTTL)
+	m.dnsForwarder = NewDNSForwarder(fmt.Sprintf(":%d", ListenPort), dnsTTL, m.statusRecorder)
 	go func() {
-		if err := m.dnsForwarder.Listen(domains); err != nil {
+		if err := m.dnsForwarder.Listen(domains, resIds); err != nil {
 			// todo handle close error if it is exists
 			log.Errorf("failed to start DNS forwarder, err: %v", err)
 		}
@@ -52,12 +55,12 @@ func (m *Manager) Start(domains []string) error {
 	return nil
 }
 
-func (m *Manager) UpdateDomains(domains []string) {
+func (m *Manager) UpdateDomains(domains []string, resIds map[string]string) {
 	if m.dnsForwarder == nil {
 		return
 	}
 
-	m.dnsForwarder.UpdateDomains(domains)
+	m.dnsForwarder.UpdateDomains(domains, resIds)
 }
 
 func (m *Manager) Stop(ctx context.Context) error {
@@ -88,7 +91,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"
@@ -189,6 +191,7 @@ type Engine struct {
 	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.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
@@ -898,11 +952,6 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
 		return nil
 	}
 
-	// Apply ACLs in the beginning to avoid security leaks
-	if e.acl != nil {
-		e.acl.ApplyFiltering(networkMap)
-	}
-
 	if e.firewall != nil {
 		if localipfw, ok := e.firewall.(localIpUpdater); ok {
 			if err := localipfw.UpdateLocalIPs(); err != nil {
@@ -913,14 +962,19 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
 
 	// DNS forwarder
 	dnsRouteFeatureFlag := toDNSFeatureFlag(networkMap)
-	dnsRouteDomains := toRouteDomains(e.config.WgPrivateKey.PublicKey().String(), networkMap.GetRoutes())
-	e.updateDNSForwarder(dnsRouteFeatureFlag, dnsRouteDomains)
+	dnsRouteDomains, resourceIds := toRouteDomains(e.config.WgPrivateKey.PublicKey().String(), networkMap.GetRoutes())
+	e.updateDNSForwarder(dnsRouteFeatureFlag, dnsRouteDomains, resourceIds)
 
 	routes := toRoutes(networkMap.GetRoutes())
 	if err := e.routeManager.UpdateRoutes(serial, routes, dnsRouteFeatureFlag); err != nil {
 		log.Errorf("failed to update clientRoutes, err: %v", err)
 	}
 
+	// acls might need routing to be enabled, so we apply after routes
+	if e.acl != nil {
+		e.acl.ApplyFiltering(networkMap)
+	}
+
 	// Ingress forward rules
 	if err := e.updateForwardRules(networkMap.GetForwardingRules()); err != nil {
 		log.Errorf("failed to update forward rules, err: %v", err)
@@ -1025,21 +1079,29 @@ func toRoutes(protoRoutes []*mgmProto.Route) []*route.Route {
 	return routes
 }
 
-func toRouteDomains(myPubKey string, protoRoutes []*mgmProto.Route) []string {
+func toRouteDomains(myPubKey string, protoRoutes []*mgmProto.Route) ([]string, map[string]string) {
 	if protoRoutes == nil {
 		protoRoutes = []*mgmProto.Route{}
 	}
 
 	var dnsRoutes []string
+	resIds := make(map[string]string)
 	for _, protoRoute := range protoRoutes {
 		if len(protoRoute.Domains) == 0 {
 			continue
 		}
 		if protoRoute.Peer == myPubKey {
 			dnsRoutes = append(dnsRoutes, protoRoute.Domains...)
+			// resource ID is the first part of the ID
+			resId := strings.Split(protoRoute.ID, ":")
+			for _, domain := range protoRoute.Domains {
+				if len(resId) > 0 {
+					resIds[domain] = resId[0]
+				}
+			}
 		}
 	}
-	return dnsRoutes
+	return dnsRoutes, resIds
 }
 
 func toDNSConfig(protoDNSConfig *mgmProto.DNSConfig, network *net.IPNet) nbdns.Config {
@@ -1706,7 +1768,7 @@ func (e *Engine) GetWgAddr() net.IP {
 }
 
 // updateDNSForwarder start or stop the DNS forwarder based on the domains and the feature flag
-func (e *Engine) updateDNSForwarder(enabled bool, domains []string) {
+func (e *Engine) updateDNSForwarder(enabled bool, domains []string, resIds map[string]string) {
 	if !enabled {
 		if e.dnsForwardMgr == nil {
 			return
@@ -1720,15 +1782,15 @@ func (e *Engine) updateDNSForwarder(enabled bool, domains []string) {
 	if len(domains) > 0 {
 		log.Infof("enable domain router service for domains: %v", domains)
 		if e.dnsForwardMgr == nil {
-			e.dnsForwardMgr = dnsfwd.NewManager(e.firewall)
+			e.dnsForwardMgr = dnsfwd.NewManager(e.firewall, e.statusRecorder)
 
-			if err := e.dnsForwardMgr.Start(domains); err != nil {
+			if err := e.dnsForwardMgr.Start(domains, resIds); err != nil {
 				log.Errorf("failed to start DNS forward: %v", err)
 				e.dnsForwardMgr = nil
 			}
 		} else {
 			log.Infof("update domain router service for domains: %v", domains)
-			e.dnsForwardMgr.UpdateDomains(domains)
+			e.dnsForwardMgr.UpdateDomains(domains, resIds)
 		}
 	} else if e.dnsForwardMgr != nil {
 		log.Infof("disable domain router service")

client/internal/engine_test.go

@@ -12,6 +12,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/golang/mock/gomock"
 	"github.com/google/uuid"
 	"github.com/pion/transport/v3/stdnet"
 	log "github.com/sirupsen/logrus"
@@ -27,6 +28,8 @@ import (
 
 	"github.com/netbirdio/management-integrations/integrations"
 
+	"github.com/netbirdio/netbird/management/server/types"
+
 	"github.com/netbirdio/netbird/client/iface"
 	"github.com/netbirdio/netbird/client/iface/bind"
 	"github.com/netbirdio/netbird/client/iface/configurer"
@@ -46,6 +49,7 @@ import (
 	"github.com/netbirdio/netbird/management/server"
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
+	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/settings"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/telemetry"
@@ -1397,15 +1401,15 @@ func startSignal(t *testing.T) (*grpc.Server, string, error) {
 func startManagement(t *testing.T, dataDir, testFile string) (*grpc.Server, string, error) {
 	t.Helper()
 
-	config := &server.Config{
-		Stuns:      []*server.Host{},
-		TURNConfig: &server.TURNConfig{},
-		Relay: &server.Relay{
+	config := &types.Config{
+		Stuns:      []*types.Host{},
+		TURNConfig: &types.TURNConfig{},
+		Relay: &types.Relay{
 			Addresses:      []string{"127.0.0.1:1234"},
 			CredentialsTTL: util.Duration{Duration: time.Hour},
 			Secret:         "222222222222222222",
 		},
-		Signal: &server.Host{
+		Signal: &types.Host{
 			Proto: "http",
 			URI:   "localhost:10000",
 		},
@@ -1435,13 +1439,23 @@ 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())
+	permissionsManagerMock := permissions.NewManagerMock()
+
+	ctrl := gomock.NewController(t)
+	t.Cleanup(ctrl.Finish)
+	settingsMockManager := settings.NewMockManager(ctrl)
+	settingsMockManager.EXPECT().
+		GetSettings(gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(&types.Settings{}, nil).
+		AnyTimes()
+
+	accountManager, err := server.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, ia, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManagerMock)
 	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)
+	secretsManager := server.NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig, config.Relay, settingsMockManager)
+	mgmtServer, err := server.NewServer(context.Background(), config, accountManager, settingsMockManager, peersUpdateManager, secretsManager, nil, nil, nil)
 	if err != nil {
 		return nil, "", err
 	}

client/internal/netflow/conntrack/conntrack.go

@@ -0,0 +1,312 @@
+//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"
+	nbnet "github.com/netbirdio/netbird/util/net"
+)
+
+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(flow.Mark, 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(flow.Mark, 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(mark uint32, srcIP, dstIP netip.Addr) bool {
+	if nbnet.IsDataPlaneMark(mark) {
+		return true
+	}
+
+	// fallback if mark rules are not in place
+	wgnet := c.iface.Address().Network
+	return wgnet.Contains(srcIP.AsSlice()) || wgnet.Contains(dstIP.AsSlice())
+}
+
+// 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(mark uint32, srcIP, dstIP netip.Addr) nftypes.Direction {
+	switch mark {
+	case nbnet.DataPlaneMarkIn:
+		return nftypes.Ingress
+	case nbnet.DataPlaneMarkOut:
+		return nftypes.Egress
+	}
+
+	// fallback if marks are not set
+	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
+	}
+
+	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,152 @@
+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
+	enabled            atomic.Bool
+	rcvChan            atomic.Pointer[rcvChan]
+	cancel             context.CancelFunc
+	statusRecorder     *peer.Status
+	wgIfaceIPNet       net.IPNet
+	dnsCollection      atomic.Bool
+	exitNodeCollection atomic.Bool
+	Store              types.Store
+}
+
+func New(statusRecorder *peer.Status, wgIfaceIPNet net.IPNet) *Logger {
+
+	return &Logger{
+		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(context.Background())
+	l.cancel = 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().UTC(),
+			}
+
+			var isSrcExitNode bool
+			var isDestExitNode bool
+
+			if !l.wgIfaceIPNet.Contains(net.IP(event.SourceIP.AsSlice())) {
+				event.SourceResourceID, isSrcExitNode = l.statusRecorder.CheckRoutes(event.SourceIP)
+			}
+
+			if !l.wgIfaceIPNet.Contains(net.IP(event.DestIP.AsSlice())) {
+				event.DestResourceID, isDestExitNode = l.statusRecorder.CheckRoutes(event.DestIP)
+			}
+
+			if l.shouldStore(eventFields, isSrcExitNode || isDestExitNode) {
+				l.Store.StoreEvent(&event)
+			}
+		}
+	}
+}
+
+func (l *Logger) Close() {
+	l.stop()
+	l.Store.Close()
+}
+
+func (l *Logger) stop() {
+	if !l.enabled.Load() {
+		return
+	}
+
+	l.enabled.Store(false)
+	l.mux.Lock()
+	if l.cancel != nil {
+		l.cancel()
+		l.cancel = 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) 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,67 @@
+package logger_test
+
+import (
+	"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(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.Close()
+	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,275 @@
+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
+	receiverClient *client.GRPCClient
+	publicKey      []byte
+	cancel         context.CancelFunc
+}
+
+// NewManager creates a new netflow manager
+func NewManager(iface nftypes.IFaceMapper, publicKey []byte, statusRecorder *peer.Status) *Manager {
+	var ipNet net.IPNet
+	if iface != nil {
+		ipNet = *iface.Address().Network
+	}
+	flowLogger := logger.New(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,
+		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 err := m.resetClient(); err != nil {
+			return fmt.Errorf("reset client: %w", err)
+		}
+	}
+
+	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
+}
+
+func (m *Manager) resetClient() error {
+	if m.receiverClient != nil {
+		if err := m.receiverClient.Close(); err != nil {
+			log.Warnf("error closing previous flow client: %v", 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
+
+	if m.cancel != nil {
+		m.cancel()
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	m.cancel = cancel
+
+	go m.receiveACKs(ctx, flowClient)
+	go m.startSender(ctx)
+
+	return nil
+}
+
+// disableFlow stops components for flow tracking
+func (m *Manager) disableFlow() error {
+	if m.cancel != nil {
+		m.cancel()
+	}
+
+	if m.conntrack != nil {
+		m.conntrack.Stop()
+	}
+
+	m.logger.Close()
+
+	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 {
+		log.Debug("no update provided; skipping update")
+		return nil
+	}
+
+	log.Tracef("updating flow configuration with new settings: url -> %s, interval -> %s, enabled? %t", update.URL, update.Interval, update.Enabled)
+
+	m.mux.Lock()
+	defer m.mux.Unlock()
+
+	previous := m.flowConfig
+	m.flowConfig = update
+
+	// Preserve TokenPayload and TokenSignature if they were set previously
+	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)
+
+	changed := previous != nil && update.Enabled != previous.Enabled
+	if update.Enabled {
+		if changed {
+			log.Infof("netflow manager enabled; starting netflow manager")
+		}
+		return m.enableFlow(previous)
+	}
+
+	if changed {
+		log.Infof("netflow manager disabled; stopping netflow manager")
+	}
+	return m.disableFlow()
+}
+
+// Close cleans up all resources
+func (m *Manager) Close() {
+	m.mux.Lock()
+	defer m.mux.Unlock()
+
+	if err := m.disableFlow(); err != nil {
+		log.Warnf("failed to disable flow manager: %v", err)
+	}
+}
+
+// GetLogger returns the flow logger
+func (m *Manager) GetLogger() nftypes.FlowLogger {
+	return m.logger
+}
+
+func (m *Manager) startSender(ctx context.Context) {
+	ticker := time.NewTicker(m.flowConfig.Interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-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: %v", err)
+					continue
+				}
+				log.Tracef("sent flow event: %s", event.ID)
+			}
+		}
+	}
+}
+
+func (m *Manager) receiveACKs(ctx context.Context, client *client.GRPCClient) {
+	err := client.Receive(ctx, m.flowConfig.Interval, func(ack *proto.FlowEventAck) error {
+		id, err := uuid.FromBytes(ack.EventId)
+		if err != nil {
+			log.Warnf("failed to convert ack event id to uuid: %v", err)
+			return nil
+		}
+		log.Tracef("received flow event ack: %s", id)
+		m.logger.DeleteEvents([]uuid.UUID{id})
+		return nil
+	})
+
+	if err != nil && !errors.Is(err, context.Canceled) {
+		log.Errorf("failed to receive flow event ack: %v", 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/manager_test.go

@@ -0,0 +1,200 @@
+package netflow
+
+import (
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/client/iface/wgaddr"
+	"github.com/netbirdio/netbird/client/internal/netflow/types"
+	"github.com/netbirdio/netbird/client/internal/peer"
+)
+
+type mockIFaceMapper struct {
+	address         wgaddr.Address
+	isUserspaceBind bool
+}
+
+func (m *mockIFaceMapper) Name() string {
+	return "wt0"
+}
+
+func (m *mockIFaceMapper) Address() wgaddr.Address {
+	return m.address
+}
+
+func (m *mockIFaceMapper) IsUserspaceBind() bool {
+	return m.isUserspaceBind
+}
+
+func TestManager_Update(t *testing.T) {
+	mockIFace := &mockIFaceMapper{
+		address: wgaddr.Address{
+			Network: &net.IPNet{
+				IP:   net.ParseIP("192.168.1.1"),
+				Mask: net.CIDRMask(24, 32),
+			},
+		},
+		isUserspaceBind: true,
+	}
+
+	publicKey := []byte("test-public-key")
+	statusRecorder := peer.NewRecorder("")
+
+	manager := NewManager(mockIFace, publicKey, statusRecorder)
+
+	tests := []struct {
+		name   string
+		config *types.FlowConfig
+	}{
+		{
+			name:   "nil config",
+			config: nil,
+		},
+		{
+			name: "disabled config",
+			config: &types.FlowConfig{
+				Enabled: false,
+			},
+		},
+		{
+			name: "enabled config with minimal valid settings",
+			config: &types.FlowConfig{
+				Enabled:        true,
+				URL:            "https://example.com",
+				TokenPayload:   "test-payload",
+				TokenSignature: "test-signature",
+				Interval:       30 * time.Second,
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			err := manager.Update(tc.config)
+
+			assert.NoError(t, err)
+
+			if tc.config == nil {
+				return
+			}
+
+			require.NotNil(t, manager.flowConfig)
+
+			if tc.config.Enabled {
+				assert.Equal(t, tc.config.Enabled, manager.flowConfig.Enabled)
+			}
+
+			if tc.config.URL != "" {
+				assert.Equal(t, tc.config.URL, manager.flowConfig.URL)
+			}
+
+			if tc.config.TokenPayload != "" {
+				assert.Equal(t, tc.config.TokenPayload, manager.flowConfig.TokenPayload)
+			}
+		})
+	}
+}
+
+func TestManager_Update_TokenPreservation(t *testing.T) {
+	mockIFace := &mockIFaceMapper{
+		address: wgaddr.Address{
+			Network: &net.IPNet{
+				IP:   net.ParseIP("192.168.1.1"),
+				Mask: net.CIDRMask(24, 32),
+			},
+		},
+		isUserspaceBind: true,
+	}
+
+	publicKey := []byte("test-public-key")
+	manager := NewManager(mockIFace, publicKey, nil)
+
+	// First update with tokens
+	initialConfig := &types.FlowConfig{
+		Enabled:        false,
+		TokenPayload:   "initial-payload",
+		TokenSignature: "initial-signature",
+	}
+
+	err := manager.Update(initialConfig)
+	require.NoError(t, err)
+
+	// Second update without tokens should preserve them
+	updatedConfig := &types.FlowConfig{
+		Enabled: false,
+		URL:     "https://example.com",
+	}
+
+	err = manager.Update(updatedConfig)
+	require.NoError(t, err)
+
+	// Verify tokens were preserved
+	assert.Equal(t, "initial-payload", manager.flowConfig.TokenPayload)
+	assert.Equal(t, "initial-signature", manager.flowConfig.TokenSignature)
+}
+
+func TestManager_NeedsNewClient(t *testing.T) {
+	manager := &Manager{}
+
+	tests := []struct {
+		name     string
+		previous *types.FlowConfig
+		current  *types.FlowConfig
+		expected bool
+	}{
+		{
+			name:     "nil previous config",
+			previous: nil,
+			current:  &types.FlowConfig{},
+			expected: true,
+		},
+		{
+			name:     "previous disabled",
+			previous: &types.FlowConfig{Enabled: false},
+			current:  &types.FlowConfig{Enabled: true},
+			expected: true,
+		},
+		{
+			name:     "different URL",
+			previous: &types.FlowConfig{Enabled: true, URL: "old-url"},
+			current:  &types.FlowConfig{Enabled: true, URL: "new-url"},
+			expected: true,
+		},
+		{
+			name:     "different TokenPayload",
+			previous: &types.FlowConfig{Enabled: true, TokenPayload: "old-payload"},
+			current:  &types.FlowConfig{Enabled: true, TokenPayload: "new-payload"},
+			expected: true,
+		},
+		{
+			name:     "different TokenSignature",
+			previous: &types.FlowConfig{Enabled: true, TokenSignature: "old-signature"},
+			current:  &types.FlowConfig{Enabled: true, TokenSignature: "new-signature"},
+			expected: true,
+		},
+		{
+			name:     "same config",
+			previous: &types.FlowConfig{Enabled: true, URL: "url", TokenPayload: "payload", TokenSignature: "signature"},
+			current:  &types.FlowConfig{Enabled: true, URL: "url", TokenPayload: "payload", TokenSignature: "signature"},
+			expected: false,
+		},
+		{
+			name:     "only interval changed",
+			previous: &types.FlowConfig{Enabled: true, URL: "url", TokenPayload: "payload", TokenSignature: "signature", Interval: 30 * time.Second},
+			current:  &types.FlowConfig{Enabled: true, URL: "url", TokenPayload: "payload", TokenSignature: "signature", Interval: 60 * time.Second},
+			expected: false,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			manager.flowConfig = tc.current
+			result := manager.needsNewClient(tc.previous)
+			assert.Equal(t, tc.expected, result)
+		})
+	}
+}

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,155 @@
+package types
+
+import (
+	"net/netip"
+	"strconv"
+	"time"
+
+	"github.com/google/uuid"
+
+	"github.com/netbirdio/netbird/client/iface/wgaddr"
+)
+
+const ZoneID = 0x1BD0
+
+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()
+	// 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/conn.go

@@ -103,6 +103,7 @@ type Conn struct {
 
 	workerICE   *WorkerICE
 	workerRelay *WorkerRelay
+	wgWatcherWg sync.WaitGroup
 
 	connIDRelay          nbnet.ConnectionID
 	connIDICE            nbnet.ConnectionID
@@ -114,6 +115,9 @@ type Conn struct {
 
 	guard     *guard.Guard
 	semaphore *semaphoregroup.SemaphoreGroup
+
+	// debug purpose
+	dumpState *stateDump
 }
 
 // NewConn creates a new not opened Conn to the remote peer.
@@ -137,10 +141,11 @@ func NewConn(engineCtx context.Context, config ConnConfig, statusRecorder *Statu
 		statusRelay:    NewAtomicConnStatus(),
 		statusICE:      NewAtomicConnStatus(),
 		semaphore:      semaphore,
+		dumpState:      newStateDump(config.Key, connLog, statusRecorder),
 	}
 
 	ctrl := isController(config)
-	conn.workerRelay = NewWorkerRelay(connLog, ctrl, config, conn, relayManager)
+	conn.workerRelay = NewWorkerRelay(connLog, ctrl, config, conn, relayManager, conn.dumpState)
 
 	relayIsSupportedLocally := conn.workerRelay.RelayIsSupportedLocally()
 	workerICE, err := NewWorkerICE(ctx, connLog, config, conn, signaler, iFaceDiscover, statusRecorder, relayIsSupportedLocally)
@@ -160,6 +165,7 @@ func NewConn(engineCtx context.Context, config ConnConfig, statusRecorder *Statu
 
 	go conn.handshaker.Listen()
 
+	go conn.dumpState.Start(ctx)
 	return conn, nil
 }
 
@@ -193,6 +199,7 @@ func (conn *Conn) startHandshakeAndReconnect(ctx context.Context) {
 	defer conn.semaphore.Done(conn.ctx)
 	conn.waitInitialRandomSleepTime(ctx)
 
+	conn.dumpState.SendOffer()
 	err := conn.handshaker.sendOffer()
 	if err != nil {
 		conn.log.Errorf("failed to send initial offer: %v", err)
@@ -205,6 +212,7 @@ func (conn *Conn) startHandshakeAndReconnect(ctx context.Context) {
 // Close closes this peer Conn issuing a close event to the Conn closeCh
 func (conn *Conn) Close() {
 	conn.mu.Lock()
+	defer conn.wgWatcherWg.Wait()
 	defer conn.mu.Unlock()
 
 	conn.log.Infof("close peer connection")
@@ -246,17 +254,20 @@ func (conn *Conn) Close() {
 	}
 
 	conn.setStatusToDisconnected()
+	conn.log.Infof("peer connection has been closed")
 }
 
 // OnRemoteAnswer handles an offer from the remote peer and returns true if the message was accepted, false otherwise
 // doesn't block, discards the message if connection wasn't ready
 func (conn *Conn) OnRemoteAnswer(answer OfferAnswer) bool {
-	conn.log.Debugf("OnRemoteAnswer, status ICE: %s, status relay: %s", conn.statusICE, conn.statusRelay)
+	conn.dumpState.RemoteAnswer()
+	conn.log.Infof("OnRemoteAnswer, priority: %s, status ICE: %s, status relay: %s", conn.currentConnPriority, conn.statusICE, conn.statusRelay)
 	return conn.handshaker.OnRemoteAnswer(answer)
 }
 
 // OnRemoteCandidate Handles ICE connection Candidate provided by the remote peer.
 func (conn *Conn) OnRemoteCandidate(candidate ice.Candidate, haRoutes route.HAMap) {
+	conn.dumpState.RemoteCandidate()
 	conn.workerICE.OnRemoteCandidate(candidate, haRoutes)
 }
 
@@ -278,7 +289,8 @@ func (conn *Conn) SetOnDisconnected(handler func(remotePeer string)) {
 }
 
 func (conn *Conn) OnRemoteOffer(offer OfferAnswer) bool {
-	conn.log.Debugf("OnRemoteOffer, on status ICE: %s, status Relay: %s", conn.statusICE, conn.statusRelay)
+	conn.dumpState.RemoteOffer()
+	conn.log.Infof("OnRemoteOffer, on status ICE: %s, status Relay: %s", conn.statusICE, conn.statusRelay)
 	return conn.handshaker.OnRemoteOffer(offer)
 }
 
@@ -322,6 +334,7 @@ func (conn *Conn) onICEConnectionIsReady(priority ConnPriority, iceConnInfo ICEC
 	}
 
 	conn.log.Infof("set ICE to active connection")
+	conn.dumpState.P2PConnected()
 
 	var (
 		ep      *net.UDPAddr
@@ -329,6 +342,7 @@ func (conn *Conn) onICEConnectionIsReady(priority ConnPriority, iceConnInfo ICEC
 		err     error
 	)
 	if iceConnInfo.RelayedOnLocal {
+		conn.dumpState.NewLocalProxy()
 		wgProxy, err = conn.newProxy(iceConnInfo.RemoteConn)
 		if err != nil {
 			conn.log.Errorf("failed to add turn net.Conn to local proxy: %v", err)
@@ -351,6 +365,7 @@ func (conn *Conn) onICEConnectionIsReady(priority ConnPriority, iceConnInfo ICEC
 	}
 
 	conn.workerRelay.DisableWgWatcher()
+	// todo consider to run conn.wgWatcherWg.Wait() here
 
 	if conn.wgProxyRelay != nil {
 		conn.wgProxyRelay.Pause()
@@ -390,12 +405,18 @@ func (conn *Conn) onICEStateDisconnected() {
 	// switch back to relay connection
 	if conn.isReadyToUpgrade() {
 		conn.log.Infof("ICE disconnected, set Relay to active connection")
+		conn.dumpState.SwitchToRelay()
 		conn.wgProxyRelay.Work()
 
 		if err := conn.configureWGEndpoint(conn.wgProxyRelay.EndpointAddr()); err != nil {
 			conn.log.Errorf("failed to switch to relay conn: %v", err)
 		}
-		conn.workerRelay.EnableWgWatcher(conn.ctx)
+
+		conn.wgWatcherWg.Add(1)
+		go func() {
+			defer conn.wgWatcherWg.Done()
+			conn.workerRelay.EnableWgWatcher(conn.ctx)
+		}()
 		conn.currentConnPriority = connPriorityRelay
 	} else {
 		conn.log.Infof("ICE disconnected, do not switch to Relay. Reset priority to: %s", connPriorityNone.String())
@@ -432,6 +453,7 @@ func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) {
 		return
 	}
 
+	conn.dumpState.RelayConnected()
 	conn.log.Debugf("Relay connection has been established, setup the WireGuard")
 
 	wgProxy, err := conn.newProxy(rci.relayedConn)
@@ -439,6 +461,7 @@ func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) {
 		conn.log.Errorf("failed to add relayed net.Conn to local proxy: %v", err)
 		return
 	}
+	conn.dumpState.NewLocalProxy()
 
 	conn.log.Infof("created new wgProxy for relay connection: %s", wgProxy.EndpointAddr().String())
 
@@ -462,7 +485,12 @@ func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) {
 		conn.log.Errorf("Failed to update WireGuard peer configuration: %v", err)
 		return
 	}
-	conn.workerRelay.EnableWgWatcher(conn.ctx)
+
+	conn.wgWatcherWg.Add(1)
+	go func() {
+		defer conn.wgWatcherWg.Done()
+		conn.workerRelay.EnableWgWatcher(conn.ctx)
+	}()
 
 	wgConfigWorkaround()
 	conn.currentConnPriority = connPriorityRelay
@@ -481,13 +509,14 @@ func (conn *Conn) onRelayDisconnected() {
 		return
 	}
 
-	conn.log.Debugf("relay connection is disconnected")
+	conn.log.Infof("relay connection is disconnected")
 
 	if conn.currentConnPriority == connPriorityRelay {
-		conn.log.Debugf("clean up WireGuard config")
+		conn.log.Infof("clean up WireGuard config")
 		if err := conn.removeWgPeer(); err != nil {
 			conn.log.Errorf("failed to remove wg endpoint: %v", err)
 		}
+		conn.currentConnPriority = connPriorityNone
 	}
 
 	if conn.wgProxyRelay != nil {
@@ -516,7 +545,8 @@ func (conn *Conn) listenGuardEvent(ctx context.Context) {
 	for {
 		select {
 		case <-conn.guard.Reconnect:
-			conn.log.Debugf("send offer to peer")
+			conn.log.Infof("send offer to peer")
+			conn.dumpState.SendOffer()
 			if err := conn.handshaker.SendOffer(); err != nil {
 				conn.log.Errorf("failed to send offer: %v", err)
 			}

client/internal/peer/handshaker.go

@@ -76,19 +76,19 @@ func (h *Handshaker) AddOnNewOfferListener(offer func(remoteOfferAnswer *OfferAn
 
 func (h *Handshaker) Listen() {
 	for {
-		h.log.Debugf("wait for remote offer confirmation")
+		h.log.Info("wait for remote offer confirmation")
 		remoteOfferAnswer, err := h.waitForRemoteOfferConfirmation()
 		if err != nil {
 			var connectionClosedError *ConnectionClosedError
 			if errors.As(err, &connectionClosedError) {
-				h.log.Tracef("stop handshaker")
+				h.log.Info("exit from handshaker")
 				return
 			}
 			h.log.Errorf("failed to received remote offer confirmation: %s", err)
 			continue
 		}
 
-		h.log.Debugf("received connection confirmation, running version %s and with remote WireGuard listen port %d", remoteOfferAnswer.Version, remoteOfferAnswer.WgListenPort)
+		h.log.Infof("received connection confirmation, running version %s and with remote WireGuard listen port %d", remoteOfferAnswer.Version, remoteOfferAnswer.WgListenPort)
 		for _, listener := range h.onNewOfferListeners {
 			go listener(remoteOfferAnswer)
 		}
@@ -108,7 +108,7 @@ func (h *Handshaker) OnRemoteOffer(offer OfferAnswer) bool {
 	case h.remoteOffersCh <- offer:
 		return true
 	default:
-		h.log.Debugf("OnRemoteOffer skipping message because is not ready")
+		h.log.Warnf("OnRemoteOffer skipping message because is not ready")
 		// connection might not be ready yet to receive so we ignore the message
 		return false
 	}
@@ -131,8 +131,7 @@ func (h *Handshaker) waitForRemoteOfferConfirmation() (*OfferAnswer, error) {
 	select {
 	case remoteOfferAnswer := <-h.remoteOffersCh:
 		// received confirmation from the remote peer -> ready to proceed
-		err := h.sendAnswer()
-		if err != nil {
+		if err := h.sendAnswer(); err != nil {
 			return nil, err
 		}
 		return &remoteOfferAnswer, nil
@@ -168,7 +167,7 @@ func (h *Handshaker) sendOffer() error {
 }
 
 func (h *Handshaker) sendAnswer() error {
-	h.log.Debugf("sending answer")
+	h.log.Infof("sending answer")
 	uFrag, pwd := h.ice.GetLocalUserCredentials()
 
 	answer := OfferAnswer{

client/internal/peer/ice/agent.go

@@ -4,6 +4,7 @@ import (
 	"time"
 
 	"github.com/pion/ice/v3"
+	"github.com/pion/logging"
 	"github.com/pion/randutil"
 	log "github.com/sirupsen/logrus"
 
@@ -35,6 +36,9 @@ func NewAgent(iFaceDiscover stdnet.ExternalIFaceDiscover, config Config, candida
 		log.Errorf("failed to create pion's stdnet: %s", err)
 	}
 
+	fac := logging.NewDefaultLoggerFactory()
+	fac.Writer = log.StandardLogger().Writer()
+
 	agentConfig := &ice.AgentConfig{
 		MulticastDNSMode:       ice.MulticastDNSModeDisabled,
 		NetworkTypes:           []ice.NetworkType{ice.NetworkTypeUDP4, ice.NetworkTypeUDP6},
@@ -51,6 +55,7 @@ func NewAgent(iFaceDiscover stdnet.ExternalIFaceDiscover, config Config, candida
 		RelayAcceptanceMinWait: &iceRelayAcceptanceMinWait,
 		LocalUfrag:             ufrag,
 		LocalPwd:               pwd,
+		LoggerFactory:          fac,
 	}
 
 	if config.DisableIPv6Discovery {

client/internal/peer/route.go

@@ -0,0 +1,131 @@
+package peer
+
+import (
+	"net/netip"
+	"sort"
+	"sync"
+
+	log "github.com/sirupsen/logrus"
+)
+
+// routeEntry holds the route prefix and the corresponding resource ID.
+type routeEntry struct {
+	prefix     netip.Prefix
+	resourceID string
+}
+
+type routeIDLookup struct {
+	localRoutes []routeEntry
+	localLock   sync.RWMutex
+
+	remoteRoutes []routeEntry
+	remoteLock   sync.RWMutex
+
+	resolvedIPs sync.Map
+}
+
+func (r *routeIDLookup) AddLocalRouteID(resourceID string, route netip.Prefix) {
+	r.localLock.Lock()
+	defer r.localLock.Unlock()
+
+	// update the resource id if the route already exists.
+	for i, entry := range r.localRoutes {
+		if entry.prefix == route {
+			r.localRoutes[i].resourceID = resourceID
+			log.Tracef("resourceID for route %v updated to %s in local routes", route, resourceID)
+			return
+		}
+	}
+
+	// append and sort descending by prefix bits (more specific first)
+	r.localRoutes = append(r.localRoutes, routeEntry{prefix: route, resourceID: resourceID})
+	sort.Slice(r.localRoutes, func(i, j int) bool {
+		return r.localRoutes[i].prefix.Bits() > r.localRoutes[j].prefix.Bits()
+	})
+}
+
+func (r *routeIDLookup) RemoveLocalRouteID(route netip.Prefix) {
+	r.localLock.Lock()
+	defer r.localLock.Unlock()
+
+	for i, entry := range r.localRoutes {
+		if entry.prefix == route {
+			r.localRoutes = append(r.localRoutes[:i], r.localRoutes[i+1:]...)
+			return
+		}
+	}
+}
+
+func (r *routeIDLookup) AddRemoteRouteID(resourceID string, route netip.Prefix) {
+	r.remoteLock.Lock()
+	defer r.remoteLock.Unlock()
+
+	for i, entry := range r.remoteRoutes {
+		if entry.prefix == route {
+			r.remoteRoutes[i].resourceID = resourceID
+			log.Tracef("resourceID for route %v updated to %s in remote routes", route, resourceID)
+			return
+		}
+	}
+
+	// append and sort descending by prefix bits.
+	r.remoteRoutes = append(r.remoteRoutes, routeEntry{prefix: route, resourceID: resourceID})
+	sort.Slice(r.remoteRoutes, func(i, j int) bool {
+		return r.remoteRoutes[i].prefix.Bits() > r.remoteRoutes[j].prefix.Bits()
+	})
+}
+
+func (r *routeIDLookup) RemoveRemoteRouteID(route netip.Prefix) {
+	r.remoteLock.Lock()
+	defer r.remoteLock.Unlock()
+
+	for i, entry := range r.remoteRoutes {
+		if entry.prefix == route {
+			r.remoteRoutes = append(r.remoteRoutes[:i], r.remoteRoutes[i+1:]...)
+			return
+		}
+	}
+}
+
+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) {
+	if res, ok := r.resolvedIPs.Load(ip); ok {
+		return res.(string), false
+	}
+
+	var resourceID string
+	var isExitNode bool
+
+	r.localLock.RLock()
+	for _, entry := range r.localRoutes {
+		if entry.prefix.Contains(ip) {
+			resourceID = entry.resourceID
+			isExitNode = (entry.prefix.Bits() == 0)
+			break
+		}
+	}
+	r.localLock.RUnlock()
+
+	if resourceID == "" {
+		r.remoteLock.RLock()
+		for _, entry := range r.remoteRoutes {
+			if entry.prefix.Contains(ip) {
+				resourceID = entry.resourceID
+				isExitNode = (entry.prefix.Bits() == 0)
+				break
+			}
+		}
+		r.remoteLock.RUnlock()
+	}
+
+	return resourceID, isExitNode
+}

client/internal/peer/state_dump.go

@@ -0,0 +1,122 @@
+package peer
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+)
+
+type stateDump struct {
+	log    *log.Entry
+	status *Status
+	key    string
+
+	sentOffer       int
+	remoteOffer     int
+	remoteAnswer    int
+	remoteCandidate int
+	p2pConnected    int
+	switchToRelay   int
+	wgCheckSuccess  int
+	relayConnected  int
+	localProxies    int
+
+	mu sync.Mutex
+}
+
+func newStateDump(key string, log *log.Entry, statusRecorder *Status) *stateDump {
+	return &stateDump{
+		log:    log,
+		status: statusRecorder,
+		key:    key,
+	}
+}
+
+func (s *stateDump) Start(ctx context.Context) {
+	ticker := time.NewTicker(10 * time.Minute)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			s.dumpState()
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
+func (s *stateDump) RemoteOffer() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.remoteOffer++
+}
+
+func (s *stateDump) RemoteCandidate() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.remoteCandidate++
+}
+
+func (s *stateDump) SendOffer() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.sentOffer++
+}
+
+func (s *stateDump) dumpState() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	status := "unknown"
+	state, e := s.status.GetPeer(s.key)
+	if e == nil {
+		status = state.ConnStatus.String()
+	}
+
+	s.log.Infof("Dump stat: Status: %s, SentOffer: %d, RemoteOffer: %d, RemoteAnswer: %d, RemoteCandidate: %d, P2PConnected: %d, SwitchToRelay: %d, WGCheckSuccess: %d, RelayConnected: %d, LocalProxies: %d",
+		status, s.sentOffer, s.remoteOffer, s.remoteAnswer, s.remoteCandidate, s.p2pConnected, s.switchToRelay, s.wgCheckSuccess, s.relayConnected, s.localProxies)
+}
+
+func (s *stateDump) RemoteAnswer() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.remoteAnswer++
+}
+
+func (s *stateDump) P2PConnected() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.p2pConnected++
+}
+
+func (s *stateDump) SwitchToRelay() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.switchToRelay++
+}
+
+func (s *stateDump) WGcheckSuccess() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.wgCheckSuccess++
+}
+
+func (s *stateDump) RelayConnected() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.relayConnected++
+}
+
+func (s *stateDump) NewLocalProxy() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.localProxies++
+}

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,11 @@ func (d *Status) AddPeerStateRoute(peer string, route string) error {
 	peerState.AddRoute(route)
 	d.peers[peer] = peerState
 
+	pref, err := netip.ParsePrefix(route)
+	if err == nil {
+		d.routeIDLookup.AddRemoteRouteID(resourceId, pref)
+	}
+
 	// todo: consider to make sense of this notification or not
 	d.notifyPeerListChanged()
 	return nil
@@ -340,11 +347,26 @@ func (d *Status) RemovePeerStateRoute(peer string, route string) error {
 	peerState.DeleteRoute(route)
 	d.peers[peer] = peerState
 
+	pref, err := netip.ParsePrefix(route)
+	if err == nil {
+		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 +580,66 @@ 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 {
+		d.routeIDLookup.AddLocalRouteID(resourceId, pref)
+	}
+
+	if d.localPeer.Routes == nil {
+		d.localPeer.Routes = map[string]struct{}{}
+	}
+
+	d.localPeer.Routes[route] = struct{}{}
+}
+
+// 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 {
+		d.routeIDLookup.RemoveLocalRouteID(pref)
+	}
+
+	delete(d.localPeer.Routes, route)
+}
+
+// AddResolvedIPLookupEntry adds a resolved IP lookup entry
+func (d *Status) AddResolvedIPLookupEntry(route, resourceId string) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	pref, err := netip.ParsePrefix(route)
+	if err == nil {
+		d.routeIDLookup.AddResolvedIP(resourceId, pref)
+	}
+}
+
+// RemoveResolvedIPLookupEntry removes a resolved IP lookup entry
+func (d *Status) RemoveResolvedIPLookupEntry(route string) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	pref, err := netip.ParsePrefix(route)
+	if err == nil {
+		d.routeIDLookup.RemoveResolvedIP(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 +723,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 +732,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 +746,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/peer/wg_watcher.go

@@ -27,18 +27,19 @@ type WGWatcher struct {
 	log           *log.Entry
 	wgIfaceStater WGInterfaceStater
 	peerKey       string
+	stateDump     *stateDump
 
 	ctx       context.Context
 	ctxCancel context.CancelFunc
 	ctxLock   sync.Mutex
-	waitGroup sync.WaitGroup
 }
 
-func NewWGWatcher(log *log.Entry, wgIfaceStater WGInterfaceStater, peerKey string) *WGWatcher {
+func NewWGWatcher(log *log.Entry, wgIfaceStater WGInterfaceStater, peerKey string, stateDump *stateDump) *WGWatcher {
 	return &WGWatcher{
 		log:           log,
 		wgIfaceStater: wgIfaceStater,
 		peerKey:       peerKey,
+		stateDump:     stateDump,
 	}
 }
 
@@ -46,24 +47,24 @@ func NewWGWatcher(log *log.Entry, wgIfaceStater WGInterfaceStater, peerKey strin
 func (w *WGWatcher) EnableWgWatcher(parentCtx context.Context, onDisconnectedFn func()) {
 	w.log.Debugf("enable WireGuard watcher")
 	w.ctxLock.Lock()
-	defer w.ctxLock.Unlock()
 
 	if w.ctx != nil && w.ctx.Err() == nil {
 		w.log.Errorf("WireGuard watcher already enabled")
+		w.ctxLock.Unlock()
 		return
 	}
 
 	ctx, ctxCancel := context.WithCancel(parentCtx)
 	w.ctx = ctx
 	w.ctxCancel = ctxCancel
+	w.ctxLock.Unlock()
 
 	initialHandshake, err := w.wgState()
 	if err != nil {
 		w.log.Warnf("failed to read initial wg stats: %v", err)
 	}
 
-	w.waitGroup.Add(1)
-	go w.periodicHandshakeCheck(ctx, ctxCancel, onDisconnectedFn, initialHandshake)
+	w.periodicHandshakeCheck(ctx, ctxCancel, onDisconnectedFn, initialHandshake)
 }
 
 // DisableWgWatcher stops the WireGuard watcher and wait for the watcher to exit
@@ -79,13 +80,11 @@ func (w *WGWatcher) DisableWgWatcher() {
 
 	w.ctxCancel()
 	w.ctxCancel = nil
-	w.waitGroup.Wait()
 }
 
 // wgStateCheck help to check the state of the WireGuard handshake and relay connection
 func (w *WGWatcher) periodicHandshakeCheck(ctx context.Context, ctxCancel context.CancelFunc, onDisconnectedFn func(), initialHandshake time.Time) {
 	w.log.Infof("WireGuard watcher started")
-	defer w.waitGroup.Done()
 
 	timer := time.NewTimer(wgHandshakeOvertime)
 	defer timer.Stop()
@@ -105,6 +104,7 @@ func (w *WGWatcher) periodicHandshakeCheck(ctx context.Context, ctxCancel contex
 
 			resetTime := time.Until(handshake.Add(checkPeriod))
 			timer.Reset(resetTime)
+			w.stateDump.WGcheckSuccess()
 
 			w.log.Debugf("WireGuard watcher reset timer: %v", resetTime)
 		case <-ctx.Done():

client/internal/peer/wg_watcher_test.go

@@ -43,13 +43,13 @@ func TestWGWatcher_EnableWgWatcher(t *testing.T) {
 
 	mlog := log.WithField("peer", "tet")
 	mocWgIface := &MocWgIface{}
-	watcher := NewWGWatcher(mlog, mocWgIface, "")
+	watcher := NewWGWatcher(mlog, mocWgIface, "", newStateDump("peer", mlog, &Status{}))
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
 	onDisconnected := make(chan struct{}, 1)
-	watcher.EnableWgWatcher(ctx, func() {
+	go watcher.EnableWgWatcher(ctx, func() {
 		mlog.Infof("onDisconnectedFn")
 		onDisconnected <- struct{}{}
 	})
@@ -72,17 +72,18 @@ func TestWGWatcher_ReEnable(t *testing.T) {
 
 	mlog := log.WithField("peer", "tet")
 	mocWgIface := &MocWgIface{}
-	watcher := NewWGWatcher(mlog, mocWgIface, "")
+	watcher := NewWGWatcher(mlog, mocWgIface, "", newStateDump("peer", mlog, &Status{}))
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
 	onDisconnected := make(chan struct{}, 1)
 
-	watcher.EnableWgWatcher(ctx, func() {})
+	go watcher.EnableWgWatcher(ctx, func() {})
+	time.Sleep(1 * time.Second)
 	watcher.DisableWgWatcher()
 
-	watcher.EnableWgWatcher(ctx, func() {
+	go watcher.EnableWgWatcher(ctx, func() {
 		onDisconnected <- struct{}{}
 	})
 

client/internal/peer/worker_ice.go

@@ -65,6 +65,7 @@ func NewWorkerICE(ctx context.Context, log *log.Entry, config ConnConfig, conn *
 		iFaceDiscover:     ifaceDiscover,
 		statusRecorder:    statusRecorder,
 		hasRelayOnLocally: hasRelayOnLocally,
+		lastKnownState:    ice.ConnectionStateDisconnected,
 	}
 
 	localUfrag, localPwd, err := icemaker.GenerateICECredentials()
@@ -213,7 +214,7 @@ func (w *WorkerICE) reCreateAgent(agentCancel context.CancelFunc, candidates []i
 			w.lastKnownState = ice.ConnectionStateConnected
 			return
 		case ice.ConnectionStateFailed, ice.ConnectionStateDisconnected:
-			if w.lastKnownState != ice.ConnectionStateDisconnected {
+			if w.lastKnownState == ice.ConnectionStateConnected {
 				w.lastKnownState = ice.ConnectionStateDisconnected
 				w.conn.onICEStateDisconnected()
 			}

client/internal/peer/worker_relay.go

@@ -33,14 +33,14 @@ type WorkerRelay struct {
 	wgWatcher *WGWatcher
 }
 
-func NewWorkerRelay(log *log.Entry, ctrl bool, config ConnConfig, conn *Conn, relayManager relayClient.ManagerService) *WorkerRelay {
+func NewWorkerRelay(log *log.Entry, ctrl bool, config ConnConfig, conn *Conn, relayManager relayClient.ManagerService, stateDump *stateDump) *WorkerRelay {
 	r := &WorkerRelay{
 		log:          log,
 		isController: ctrl,
 		config:       config,
 		conn:         conn,
 		relayManager: relayManager,
-		wgWatcher:    NewWGWatcher(log, config.WgConfig.WgInterface, config.Key),
+		wgWatcher:    NewWGWatcher(log, config.WgConfig.WgInterface, config.Key, stateDump),
 	}
 	return r
 }

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
 }