fix test expectation

.git-branches.toml

@@ -1,27 +0,0 @@
-# More info around this file at https://www.git-town.com/configuration-file
-
-[branches]
-main = "main"
-perennials = []
-perennial-regex = ""
-
-[create]
-new-branch-type = "feature"
-push-new-branches = false
-
-[hosting]
-dev-remote = "origin"
-# platform = ""
-# origin-hostname = ""
-
-[ship]
-delete-tracking-branch = false
-strategy = "squash-merge"
-
-[sync]
-feature-strategy = "merge"
-perennial-strategy = "rebase"
-prototype-strategy = "merge"
-push-hook = true
-tags = true
-upstream = false

.github/pull_request_template.md

@@ -2,10 +2,6 @@
 
 ## Issue ticket number and link
 
-## Stack
-
-<!-- branch-stack -->
-
 ### Checklist
 - [ ] Is it a bug fix
 - [ ] Is a typo/documentation fix

.github/workflows/git-town.yml

@@ -1,21 +0,0 @@
-name: Git Town
-
-on:
-  pull_request:
-    branches:
-      - '**'
-
-jobs:
-  git-town:
-    name: Display the branch stack
-    runs-on: ubuntu-latest
-
-    permissions:
-      contents: read
-      pull-requests: write
-
-    steps:
-      - uses: actions/checkout@v4
-      - uses: git-town/action@v1
-        with:
-          skip-single-stacks: true

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

@@ -22,20 +22,14 @@ jobs:
         with:
           usesh: true
           copyback: false
-          release: "14.2"
+          release: "14.1"
           prepare: |
-            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"            
+            pkg install -y go pkgconf xorg
 
           # -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/test-infrastructure-files.yml

@@ -178,7 +178,6 @@ jobs:
           grep -A 10 'relay:' docker-compose.yml | egrep 'NB_AUTH_SECRET=.+$'
           grep -A 7 Relay management.json | grep "rel://$CI_NETBIRD_DOMAIN:33445"
           grep -A 7 Relay management.json | egrep '"Secret": ".+"'
-          grep DisablePromptLogin management.json | grep 'true'
 
       - name: Install modules
         run: go mod tidy

README.md

@@ -57,16 +57,16 @@
 
 ### Key features
 
-| Connectivity | Management | Security | Automation| Platforms |
-|----|----|----|----|----|
-| <ul><li>- \[x] Kernel WireGuard</ul></li> | <ul><li>- \[x] [Admin Web UI](https://github.com/netbirdio/dashboard)</ul></li> | <ul><li>- \[x] [SSO & MFA support](https://docs.netbird.io/how-to/installation#running-net-bird-with-sso-login)</ul></li> | <ul><li>- \[x] [Public API](https://docs.netbird.io/api)</ul></li> | <ul><li>- \[x] Linux</ul></li> |
-| <ul><li>- \[x] Peer-to-peer connections</ul></li> | <ul><li>- \[x] Auto peer discovery and configuration</ui></li> | <ul><li>- \[x] [Access control - groups & rules](https://docs.netbird.io/how-to/manage-network-access)</ui></li> | <ul><li>- \[x] [Setup keys for bulk network provisioning](https://docs.netbird.io/how-to/register-machines-using-setup-keys)</ui></li> | <ul><li>- \[x] Mac</ui></li> |
-| <ul><li>- \[x] Connection relay fallback</ui></li> | <ul><li>- \[x] [IdP integrations](https://docs.netbird.io/selfhosted/identity-providers)</ui></li> | <ul><li>- \[x] [Activity logging](https://docs.netbird.io/how-to/monitor-system-and-network-activity)</ui></li> | <ul><li>- \[x] [Self-hosting quickstart script](https://docs.netbird.io/selfhosted/selfhosted-quickstart)</ui></li> | <ul><li>- \[x] Windows</ui></li> |
-| <ul><li>- \[x] [Routes to external networks](https://docs.netbird.io/how-to/routing-traffic-to-private-networks)</ui></li> | <ul><li>- \[x] [Private DNS](https://docs.netbird.io/how-to/manage-dns-in-your-network)</ui></li> | <ul><li>- \[x] [Device posture checks](https://docs.netbird.io/how-to/manage-posture-checks)</ui></li> | <ul><li>- \[x] IdP groups sync with JWT</ui></li> | <ul><li>- \[x] Android</ui></li> |
-| <ul><li>- \[x] NAT traversal with BPF</ui></li> | <ul><li>- \[x] [Multiuser support](https://docs.netbird.io/how-to/add-users-to-your-network)</ui></li> | <ul><li>- \[x] Peer-to-peer encryption</ui></li> || <ul><li>- \[x] iOS</ui></li> |
-||| <ul><li>- \[x] [Quantum-resistance with Rosenpass](https://netbird.io/knowledge-hub/the-first-quantum-resistant-mesh-vpn)</ui></li> || <ul><li>- \[x] OpenWRT</ui></li> |
-||| <ul><li>- \[x] [Periodic re-authentication](https://docs.netbird.io/how-to/enforce-periodic-user-authentication)</ui></li> || <ul><li>- \[x] [Serverless](https://docs.netbird.io/how-to/netbird-on-faas)</ui></li> |
-||||| <ul><li>- \[x] Docker</ui></li> |
+| Connectivity                                                                                                                 | Management                                                                                               | Security                                                                                                                              | Automation                                                                                                                               | Platforms                                                                               |
+|------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------|
+| <ul><li> - \[x] Kernel WireGuard </ul></li>                                                                                  | <ul><li> - \[x] [Admin Web UI](https://github.com/netbirdio/dashboard) </ul></li>                        | <ul><li> - \[x] [SSO & MFA support](https://docs.netbird.io/how-to/installation#running-net-bird-with-sso-login) </ul></li>           | <ul><li> - \[x] [Public API](https://docs.netbird.io/api) </ul></li>                                                                     | <ul><li> - \[x] Linux </ul></li>                                                        |
+| <ul><li> - \[x] Peer-to-peer connections </ul></li>                                                                          | <ul><li> - \[x] Auto peer discovery and configuration </ul></li>                                         | <ul><li> - \[x] [Access control - groups & rules](https://docs.netbird.io/how-to/manage-network-access) </ul></li>                    | <ul><li> - \[x] [Setup keys for bulk network provisioning](https://docs.netbird.io/how-to/register-machines-using-setup-keys) </ul></li> | <ul><li> - \[x] Mac </ul></li>                                                          |
+| <ul><li> - \[x] Connection relay fallback </ul></li>                                                                         | <ul><li> - \[x] [IdP integrations](https://docs.netbird.io/selfhosted/identity-providers) </ul></li>     | <ul><li> - \[x] [Activity logging](https://docs.netbird.io/how-to/monitor-system-and-network-activity) </ul></li>                     | <ul><li> - \[x] [Self-hosting quickstart script](https://docs.netbird.io/selfhosted/selfhosted-quickstart) </ul></li>                    | <ul><li> - \[x] Windows </ul></li>                                                      |
+| <ul><li> - \[x] [Routes to external networks](https://docs.netbird.io/how-to/routing-traffic-to-private-networks) </ul></li> | <ul><li> - \[x] [Private DNS](https://docs.netbird.io/how-to/manage-dns-in-your-network) </ul></li>      | <ul><li> - \[x] [Device posture checks](https://docs.netbird.io/how-to/manage-posture-checks) </ul></li>                              | <ul><li> - \[x] IdP groups sync with JWT </ul></li>                                                                                      | <ul><li> - \[x] Android </ul></li>                                                      |
+| <ul><li> - \[x] NAT traversal with BPF </ul></li>                                                                            | <ul><li> - \[x] [Multiuser support](https://docs.netbird.io/how-to/add-users-to-your-network) </ul></li> | <ul><li> -  \[x] Peer-to-peer encryption </ul></li>                                                                                   |                                                                                                                                          | <ul><li> - \[x] iOS </ul></li>                                                          |
+|                                                                                                                              |                                                                                                          | <ul><li> - \[x] [Quantum-resistance with Rosenpass](https://netbird.io/knowledge-hub/the-first-quantum-resistant-mesh-vpn) </ul></li> |                                                                                                                                          | <ul><li> - \[x] OpenWRT </ul></li>                                                      |
+|                                                                                                                              |                                                                                                          | <ui><li> - \[x] [Periodic re-authentication](https://docs.netbird.io/how-to/enforce-periodic-user-authentication)</ul></li>           |                                                                                                                                          | <ul><li> - \[x] [Serverless](https://docs.netbird.io/how-to/netbird-on-faas) </ul></li> |
+|                                                                                                                              |                                                                                                          |                                                                                                                                       |                                                                                                                                          | <ul><li> - \[x] Docker </ul></li>                                                       |
 
 ### Quickstart with NetBird Cloud
 

client/cmd/login.go

@@ -19,10 +19,6 @@ 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)",
@@ -131,7 +127,7 @@ var loginCmd = &cobra.Command{
 		}
 
 		if loginResp.NeedsSSOLogin {
-			openURL(cmd, loginResp.VerificationURIComplete, loginResp.UserCode, noBrowser)
+			openURL(cmd, loginResp.VerificationURIComplete, loginResp.UserCode)
 
 			_, err = client.WaitSSOLogin(ctx, &proto.WaitSSOLoginRequest{UserCode: loginResp.UserCode, Hostname: hostName})
 			if err != nil {
@@ -202,7 +198,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, noBrowser)
+	openURL(cmd, flowInfo.VerificationURIComplete, flowInfo.UserCode)
 
 	waitTimeout := time.Duration(flowInfo.ExpiresIn) * time.Second
 	waitCTX, c := context.WithTimeout(context.TODO(), waitTimeout)
@@ -216,27 +212,19 @@ func foregroundGetTokenInfo(ctx context.Context, cmd *cobra.Command, config *int
 	return &tokenInfo, nil
 }
 
-func openURL(cmd *cobra.Command, verificationURIComplete, userCode string, noBrowser bool) {
+func openURL(cmd *cobra.Command, verificationURIComplete, userCode string) {
 	var codeMsg string
 	if userCode != "" && !strings.Contains(verificationURIComplete, userCode) {
 		codeMsg = fmt.Sprintf("and enter the code %s to authenticate.", userCode)
 	}
 
-	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("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 !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")
-		}
+	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/testutil_test.go

@@ -92,11 +92,11 @@ func startManagement(t *testing.T, config *types.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)
 
 	settingsMockManager := settings.NewMockManager(ctrl)
-	permissionsManagerMock := permissions.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 {

client/cmd/up.go

@@ -32,16 +32,12 @@ 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",
@@ -69,9 +65,6 @@ 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 {
@@ -356,7 +349,7 @@ func runInDaemonMode(ctx context.Context, cmd *cobra.Command) error {
 
 	if loginResp.NeedsSSOLogin {
 
-		openURL(cmd, loginResp.VerificationURIComplete, loginResp.UserCode, noBrowser)
+		openURL(cmd, loginResp.VerificationURIComplete, loginResp.UserCode)
 
 		_, err = client.WaitSSOLogin(ctx, &proto.WaitSSOLoginRequest{UserCode: loginResp.UserCode, Hostname: hostName})
 		if err != nil {

client/firewall/iptables/router_linux.go

@@ -38,12 +38,10 @@ const (
 	routingFinalForwardJump = "ACCEPT"
 	routingFinalNatJump     = "MASQUERADE"
 
-	jumpManglePre  = "jump-mangle-pre"
-	jumpNatPre     = "jump-nat-pre"
-	jumpNatPost    = "jump-nat-post"
-	markManglePre  = "mark-mangle-pre"
-	markManglePost = "mark-mangle-post"
-	matchSet       = "--match-set"
+	jumpManglePre = "jump-mangle-pre"
+	jumpNatPre    = "jump-nat-pre"
+	jumpNatPost   = "jump-nat-post"
+	matchSet      = "--match-set"
 
 	dnatSuffix = "_dnat"
 	snatSuffix = "_snat"
@@ -117,10 +115,6 @@ 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
@@ -354,16 +348,12 @@ 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)
@@ -433,57 +423,6 @@ 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{
@@ -525,7 +464,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,9 +46,7 @@ func TestIptablesManager_RestoreOrCreateContainers(t *testing.T) {
 	// 5. jump rule to PRE nat chain
 	// 6. static outbound masquerade rule
 	// 7. static return masquerade rule
-	// 8. mangle prerouting mark rule
-	// 9. mangle postrouting mark rule
-	require.Len(t, manager.rules, 9, "should have created rules map")
+	require.Len(t, manager.rules, 7, "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)

client/firewall/nftables/acl_linux.go

@@ -25,10 +25,9 @@ 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"
-	chainNameManglePrerouting  = "netbird-mangle-prerouting"
-	chainNameManglePostrouting = "netbird-mangle-postrouting"
+	chainNameInputFilter   = "netbird-acl-input-filter"
+	chainNameForwardFilter = "netbird-acl-forward-filter"
+	chainNamePrerouting    = "netbird-rt-prerouting"
 
 	allowNetbirdInputRuleID = "allow Netbird incoming traffic"
 )
@@ -463,15 +462,13 @@ 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 {
-	// Chain is created by route manager
-	// TODO: move creation to a common place
-	m.chainPrerouting = &nftables.Chain{
-		Name:     chainNameManglePrerouting,
+	m.chainPrerouting = m.rConn.AddChain(&nftables.Chain{
+		Name:     chainNamePrerouting,
 		Table:    m.workTable,
 		Type:     nftables.ChainTypeFilter,
 		Hooknum:  nftables.ChainHookPrerouting,
 		Priority: nftables.ChainPriorityMangle,
-	}
+	})
 
 	m.addFwmarkToForward(chainFwFilter)
 

client/firewall/nftables/router_linux.go

@@ -100,10 +100,6 @@ 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
 }
 
@@ -200,21 +196,15 @@ func (r *router) createContainers() error {
 		Type:     nftables.ChainTypeNAT,
 	})
 
-	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,
+	// Chain is created by acl manager
+	// TODO: move creation to a common place
+	r.chains[chainNamePrerouting] = &nftables.Chain{
+		Name:     chainNamePrerouting,
 		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 {
@@ -230,83 +220,7 @@ func (r *router) createContainers() error {
 	}
 
 	if err := r.conn.Flush(); err != nil {
-		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 fmt.Errorf("nftables: unable to initialize table: %v", err)
 	}
 
 	return nil
@@ -602,10 +516,26 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 		op = expr.CmpOpNeq
 	}
 
-	// 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,
+	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},
+		},
+
 		// interface matching
 		&expr.Meta{
 			Key:      expr.MetaKeyIIFNAME,
@@ -616,7 +546,7 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 			Register: 1,
 			Data:     ifname(r.wgIface.Name()),
 		},
-	)
+	}
 
 	exprs = append(exprs, sourceExp...)
 	exprs = append(exprs, destExp...)
@@ -648,7 +578,7 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 
 	r.rules[ruleKey] = r.conn.AddRule(&nftables.Rule{
 		Table:    r.workTable,
-		Chain:    r.chains[chainNameManglePrerouting],
+		Chain:    r.chains[chainNamePrerouting],
 		Exprs:    exprs,
 		UserData: []byte(ruleKey),
 	})
@@ -1394,24 +1324,3 @@ 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 == chainNameManglePrerouting {
+					if chain.Name == chainNamePrerouting {
 						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[chainNameManglePrerouting])
+			rules, err := rtr.conn.GetRules(rtr.workTable, rtr.chains[chainNamePrerouting])
 			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[chainNameManglePrerouting])
+			rules, err = rtr.conn.GetRules(rtr.workTable, rtr.chains[chainNamePrerouting])
 			require.NoError(t, err, "should list rules after removal")
 			for _, rule := range rules {
 				if len(rule.UserData) > 0 && string(rule.UserData) == natRuleKey {

client/firewall/uspfilter/conntrack/icmp.go

@@ -189,7 +189,7 @@ func (t *ICMPTracker) cleanup() {
 		if conn.timeoutExceeded(t.timeout) {
 			delete(t.connections, key)
 
-			t.logger.Trace("Removed ICMP connection %s (timeout) [in: %d Pkts/%d B out: %d Pkts/%d B]",
+			t.logger.Debug("Removed ICMP connection %s (timeout) [in: %d Pkts/%d B out: %d Pkts/%d B]",
 				key, conn.PacketsRx.Load(), conn.BytesRx.Load(), conn.PacketsTx.Load(), conn.BytesTx.Load())
 			t.sendEvent(nftypes.TypeEnd, conn, nil)
 		}

client/firewall/uspfilter/conntrack/tcp.go

@@ -23,11 +23,11 @@ const (
 )
 
 const (
-	TCPFin  uint8 = 0x01
 	TCPSyn  uint8 = 0x02
+	TCPAck  uint8 = 0x10
+	TCPFin  uint8 = 0x01
 	TCPRst  uint8 = 0x04
 	TCPPush uint8 = 0x08
-	TCPAck  uint8 = 0x10
 	TCPUrg  uint8 = 0x20
 )
 
@@ -41,7 +41,7 @@ const (
 )
 
 // TCPState represents the state of a TCP connection
-type TCPState int32
+type TCPState int
 
 func (s TCPState) String() string {
 	switch s {
@@ -89,25 +89,22 @@ const (
 // TCPConnTrack represents a TCP connection state
 type TCPConnTrack struct {
 	BaseConnTrack
-	SourcePort uint16
-	DestPort   uint16
-	state      atomic.Int32
-	tombstone  atomic.Bool
+	SourcePort  uint16
+	DestPort    uint16
+	State       TCPState
+	established atomic.Bool
+	tombstone   atomic.Bool
+	sync.RWMutex
 }
 
-// GetState safely retrieves the current state
-func (t *TCPConnTrack) GetState() TCPState {
-	return TCPState(t.state.Load())
+// IsEstablished safely checks if connection is established
+func (t *TCPConnTrack) IsEstablished() bool {
+	return t.established.Load()
 }
 
-// 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))
+// SetEstablished safely sets the established state
+func (t *TCPConnTrack) SetEstablished(state bool) {
+	t.established.Store(state)
 }
 
 // IsTombstone safely checks if the connection is marked for deletion
@@ -128,17 +125,13 @@ type TCPTracker struct {
 	cleanupTicker *time.Ticker
 	tickerCancel  context.CancelFunc
 	timeout       time.Duration
-	waitTimeout   time.Duration
 	flowLogger    nftypes.FlowLogger
 }
 
 // NewTCPTracker creates a new TCP connection tracker
 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())
@@ -149,7 +142,6 @@ func NewTCPTracker(timeout time.Duration, logger *nblog.Logger, flowLogger nftyp
 		cleanupTicker: time.NewTicker(TCPCleanupInterval),
 		tickerCancel:  cancel,
 		timeout:       timeout,
-		waitTimeout:   waitTimeout,
 		flowLogger:    flowLogger,
 	}
 
@@ -157,7 +149,7 @@ func NewTCPTracker(timeout time.Duration, logger *nblog.Logger, flowLogger nftyp
 	return tracker
 }
 
-func (t *TCPTracker) updateIfExists(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, direction nftypes.Direction, size int) (ConnKey, bool) {
+func (t *TCPTracker) updateIfExists(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, flags uint8, direction nftypes.Direction, size int) (ConnKey, bool) {
 	key := ConnKey{
 		SrcIP:   srcIP,
 		DstIP:   dstIP,
@@ -170,7 +162,12 @@ func (t *TCPTracker) updateIfExists(srcIP, dstIP netip.Addr, srcPort, dstPort ui
 	t.mutex.RUnlock()
 
 	if exists {
-		t.updateState(key, conn, flags, direction, size)
+		conn.Lock()
+		t.updateState(key, conn, flags, conn.Direction == nftypes.Egress)
+		conn.Unlock()
+
+		conn.UpdateCounters(direction, size)
+
 		return key, true
 	}
 
@@ -178,7 +175,7 @@ func (t *TCPTracker) updateIfExists(srcIP, dstIP netip.Addr, srcPort, dstPort ui
 }
 
 // TrackOutbound records an outbound TCP connection
-func (t *TCPTracker) TrackOutbound(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, size int) {
+func (t *TCPTracker) TrackOutbound(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, flags uint8, size int) {
 	if _, exists := t.updateIfExists(dstIP, srcIP, dstPort, srcPort, flags, 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)
@@ -186,14 +183,14 @@ func (t *TCPTracker) TrackOutbound(srcIP, dstIP netip.Addr, srcPort, dstPort uin
 }
 
 // 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) {
+func (t *TCPTracker) TrackInbound(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, flags uint8, ruleID []byte, size int) {
 	t.track(srcIP, dstIP, srcPort, dstPort, flags, nftypes.Ingress, ruleID, size)
 }
 
 // track is the common implementation for tracking both inbound and outbound connections
-func (t *TCPTracker) track(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, flags uint8, direction nftypes.Direction, ruleID []byte, size int) {
+func (t *TCPTracker) track(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, flags uint8, direction nftypes.Direction, ruleID []byte, size int) {
 	key, exists := t.updateIfExists(srcIP, dstIP, srcPort, dstPort, flags, direction, size)
-	if exists || flags&TCPSyn == 0 {
+	if exists {
 		return
 	}
 
@@ -208,11 +205,12 @@ func (t *TCPTracker) track(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, fla
 		DestPort:   dstPort,
 	}
 
+	conn.established.Store(false)
 	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.updateState(key, conn, flags, direction == nftypes.Egress)
+	conn.UpdateCounters(direction, size)
 
 	t.mutex.Lock()
 	t.connections[key] = conn
@@ -222,7 +220,7 @@ func (t *TCPTracker) track(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, fla
 }
 
 // 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 {
+func (t *TCPTracker) IsValidInbound(srcIP netip.Addr, dstIP netip.Addr, srcPort uint16, dstPort uint16, flags uint8, size int) bool {
 	key := ConnKey{
 		SrcIP:   dstIP,
 		DstIP:   srcIP,
@@ -234,125 +232,134 @@ func (t *TCPTracker) IsValidInbound(srcIP, dstIP netip.Addr, srcPort, dstPort ui
 	conn, exists := t.connections[key]
 	t.mutex.RUnlock()
 
-	if !exists || conn.IsTombstone() {
+	if !exists {
 		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
+	// Handle RST flag specially - it always causes transition to closed
+	if flags&TCPRst != 0 {
+		return t.handleRst(key, conn, size)
 	}
 
-	t.updateState(key, conn, flags, nftypes.Ingress, size)
+	conn.Lock()
+	t.updateState(key, conn, flags, false)
+	isEstablished := conn.IsEstablished()
+	isValidState := t.isValidStateForFlags(conn.State, flags)
+	conn.Unlock()
+	conn.UpdateCounters(nftypes.Ingress, size)
+
+	return isEstablished || isValidState
+}
+
+func (t *TCPTracker) handleRst(key ConnKey, conn *TCPConnTrack, size int) bool {
+	if conn.IsTombstone() {
+		return true
+	}
+
+	conn.Lock()
+	conn.SetTombstone()
+	conn.State = TCPStateClosed
+	conn.SetEstablished(false)
+	conn.Unlock()
+	conn.UpdateCounters(nftypes.Ingress, size)
+
+	t.logger.Trace("TCP connection reset: %s", key)
+	t.sendEvent(nftypes.TypeEnd, conn, nil)
 	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) {
+func (t *TCPTracker) updateState(key ConnKey, conn *TCPConnTrack, flags uint8, isOutbound bool) {
 	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)
+	state := conn.State
+	defer func() {
+		if state != conn.State {
+			t.logger.Trace("TCP connection %s transitioned from %s to %s", key, state, conn.State)
 		}
-		return
-	}
+	}()
 
-	var newState TCPState
-	switch currentState {
+	switch state {
 	case TCPStateNew:
 		if flags&TCPSyn != 0 && flags&TCPAck == 0 {
-			if conn.Direction == nftypes.Egress {
-				newState = TCPStateSynSent
-			} else {
-				newState = TCPStateSynReceived
-			}
+			conn.State = TCPStateSynSent
 		}
 
 	case TCPStateSynSent:
 		if flags&TCPSyn != 0 && flags&TCPAck != 0 {
-			if packetDir != conn.Direction {
-				newState = TCPStateEstablished
+			if isOutbound {
+				conn.State = TCPStateEstablished
+				conn.SetEstablished(true)
 			} else {
 				// Simultaneous open
-				newState = TCPStateSynReceived
+				conn.State = TCPStateSynReceived
 			}
 		}
 
 	case TCPStateSynReceived:
 		if flags&TCPAck != 0 && flags&TCPSyn == 0 {
-			if packetDir == conn.Direction {
-				newState = TCPStateEstablished
-			}
+			conn.State = TCPStateEstablished
+			conn.SetEstablished(true)
 		}
 
 	case TCPStateEstablished:
 		if flags&TCPFin != 0 {
-			if packetDir == conn.Direction {
-				newState = TCPStateFinWait1
+			if isOutbound {
+				conn.State = TCPStateFinWait1
 			} else {
-				newState = TCPStateCloseWait
+				conn.State = TCPStateCloseWait
 			}
+			conn.SetEstablished(false)
+		} else if flags&TCPRst != 0 {
+			conn.State = TCPStateClosed
+			conn.SetTombstone()
+			t.sendEvent(nftypes.TypeEnd, conn, nil)
 		}
 
 	case TCPStateFinWait1:
-		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
-			}
+		switch {
+		case flags&TCPFin != 0 && flags&TCPAck != 0:
+			conn.State = TCPStateClosing
+		case flags&TCPFin != 0:
+			conn.State = TCPStateFinWait2
+		case flags&TCPAck != 0:
+			conn.State = TCPStateFinWait2
+		case flags&TCPRst != 0:
+			conn.State = TCPStateClosed
+			conn.SetTombstone()
+			t.sendEvent(nftypes.TypeEnd, conn, nil)
 		}
 
 	case TCPStateFinWait2:
 		if flags&TCPFin != 0 {
-			newState = TCPStateTimeWait
+			conn.State = TCPStateTimeWait
+
+			t.logger.Trace("TCP connection %s completed", key)
+			t.sendEvent(nftypes.TypeEnd, conn, nil)
 		}
 
 	case TCPStateClosing:
 		if flags&TCPAck != 0 {
-			newState = TCPStateTimeWait
+			conn.State = TCPStateTimeWait
+			// Keep established = false from previous state
+
+			t.logger.Trace("TCP connection %s closed (simultaneous)", key)
+			t.sendEvent(nftypes.TypeEnd, conn, nil)
 		}
 
 	case TCPStateCloseWait:
 		if flags&TCPFin != 0 {
-			newState = TCPStateLastAck
+			conn.State = TCPStateLastAck
 		}
 
 	case TCPStateLastAck:
 		if flags&TCPAck != 0 {
-			newState = TCPStateClosed
-		}
-	}
-
-	if newState != 0 && conn.CompareAndSwapState(currentState, newState) {
-		t.logger.Trace("TCP connection %s transitioned from %s to %s (dir: %s)", key, currentState, newState, packetDir)
-
-		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.State = 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())
+
+			// Send close event for gracefully closed connections
 			t.sendEvent(nftypes.TypeEnd, conn, nil)
+			t.logger.Trace("TCP connection %s closed gracefully", key)
 		}
 	}
 }
@@ -362,22 +369,18 @@ 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
@@ -394,7 +397,9 @@ func (t *TCPTracker) isValidStateForFlags(state TCPState, flags uint8) bool {
 	case TCPStateLastAck:
 		return flags&TCPAck != 0
 	case TCPStateClosed:
-		// Accept retransmitted ACKs in closed state, the final ACK might be lost and the peer will retransmit their FIN-ACK
+		// Accept retransmitted ACKs in closed state
+		// This is important because the final ACK might be lost
+		// and the peer will retransmit their FIN-ACK
 		return flags&TCPAck != 0
 	}
 	return false
@@ -425,24 +430,23 @@ func (t *TCPTracker) cleanup() {
 		}
 
 		var timeout time.Duration
-		currentState := conn.GetState()
-		switch currentState {
-		case TCPStateTimeWait:
-			timeout = t.waitTimeout
-		case TCPStateEstablished:
+		switch {
+		case conn.State == TCPStateTimeWait:
+			timeout = TimeWaitTimeout
+		case conn.IsEstablished():
 			timeout = t.timeout
 		default:
 			timeout = TCPHandshakeTimeout
 		}
 
 		if conn.timeoutExceeded(timeout) {
+			// Return IPs to pool
 			delete(t.connections, key)
 
-			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())
+			t.logger.Trace("Cleaned up timed-out TCP connection %s", key)
 
 			// event already handled by state change
-			if currentState != TCPStateTimeWait {
+			if conn.State != TCPStateTimeWait {
 				t.sendEvent(nftypes.TypeEnd, conn, nil)
 			}
 		}

client/firewall/uspfilter/conntrack/tcp_bench_test.go

@@ -1,83 +0,0 @@
-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

@@ -5,7 +5,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -125,6 +124,9 @@ func TestTCPStateMachine(t *testing.T) {
 					// Receive RST
 					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
 				},
 			},
 			{
@@ -215,446 +217,97 @@ func TestRSTHandling(t *testing.T) {
 			conn := tracker.connections[key]
 			if tt.wantValid {
 				require.NotNil(t, conn)
-				require.Equal(t, TCPStateClosed, conn.GetState())
+				require.Equal(t, TCPStateClosed, conn.State)
+				require.False(t, conn.IsEstablished())
 			}
 		})
 	}
 }
 
-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 netip.Addr, srcPort, dstPort uint16) {
 	t.Helper()
 
-	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPSyn, 100)
+	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPSyn, 0)
 
-	valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPSyn|TCPAck, 100)
+	valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPSyn|TCPAck, 0)
 	require.True(t, valid, "SYN-ACK should be allowed")
 
-	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, 100)
+	tracker.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPAck, 0)
+}
+
+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, 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, 0)
+				}
+				i++
+			}
+		})
+	})
+}
+
+// Benchmark connection cleanup
+func BenchmarkCleanup(b *testing.B) {
+	b.Run("TCPCleanup", func(b *testing.B) {
+		tracker := NewTCPTracker(100*time.Millisecond, logger, flowLogger) // Short timeout for testing
+		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/udp.go

@@ -165,7 +165,7 @@ func (t *UDPTracker) cleanup() {
 		if conn.timeoutExceeded(t.timeout) {
 			delete(t.connections, key)
 
-			t.logger.Trace("Removed UDP connection %s (timeout) [in: %d Pkts/%d B, out: %d Pkts/%d B]",
+			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)
 		}

client/firewall/uspfilter/uspfilter.go

@@ -658,8 +658,7 @@ func (m *Manager) dropFilter(packetData []byte, size int) bool {
 	d := m.decoders.Get().(*decoder)
 	defer m.decoders.Put(d)
 
-	valid, fragment := m.isValidPacket(d, packetData)
-	if !valid {
+	if !m.isValidPacket(d, packetData) {
 		return true
 	}
 
@@ -669,13 +668,6 @@ func (m *Manager) dropFilter(packetData []byte, size int) bool {
 		return true
 	}
 
-	// TODO: pass fragments of routed packets to forwarder
-	if fragment {
-		m.logger.Trace("packet is a fragment: src=%v dst=%v id=%v flags=%v",
-			srcIP, dstIP, d.ip4.Id, d.ip4.Flags)
-		return false
-	}
-
 	// 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, size) {
@@ -686,7 +678,7 @@ func (m *Manager) dropFilter(packetData []byte, size int) bool {
 		return m.handleLocalTraffic(d, srcIP, dstIP, packetData, size)
 	}
 
-	return m.handleRoutedTraffic(d, srcIP, dstIP, packetData, size)
+	return m.handleRoutedTraffic(d, srcIP, dstIP, packetData)
 }
 
 // handleLocalTraffic handles local traffic.
@@ -747,7 +739,7 @@ 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 netip.Addr, packetData []byte, size int) bool {
+func (m *Manager) handleRoutedTraffic(d *decoder, srcIP, dstIP netip.Addr, packetData []byte) bool {
 	// Drop if routing is disabled
 	if !m.routingEnabled.Load() {
 		m.logger.Trace("Dropping routed packet (routing disabled): src=%s dst=%s",
@@ -757,7 +749,6 @@ func (m *Manager) handleRoutedTraffic(d *decoder, srcIP, dstIP netip.Addr, packe
 
 	// Pass to native stack if native router is enabled or forced
 	if m.nativeRouter.Load() {
-		m.trackInbound(d, srcIP, dstIP, nil, size)
 		return false
 	}
 
@@ -779,8 +770,6 @@ func (m *Manager) handleRoutedTraffic(d *decoder, srcIP, dstIP netip.Addr, packe
 			SourcePort: srcPort,
 			DestPort:   dstPort,
 			// TODO: icmp type/code
-			RxPackets: 1,
-			RxBytes:   uint64(size),
 		})
 		return true
 	}
@@ -823,32 +812,17 @@ func getPortsFromPacket(d *decoder) (srcPort, dstPort uint16) {
 	}
 }
 
-// isValidPacket checks if the packet is valid.
-// It returns true, false if the packet is valid and not a fragment.
-// It returns true, true if the packet is a fragment and valid.
-func (m *Manager) isValidPacket(d *decoder, packetData []byte) (bool, bool) {
+func (m *Manager) isValidPacket(d *decoder, packetData []byte) bool {
 	if err := d.parser.DecodeLayers(packetData, &d.decoded); err != nil {
 		m.logger.Trace("couldn't decode packet, err: %s", err)
-		return false, false
+		return false
 	}
 
-	l := len(d.decoded)
-
-	// L3 and L4 are mandatory
-	if l >= 2 {
-		return true, false
+	if len(d.decoded) < 2 {
+		m.logger.Trace("packet doesn't have network and transport layers")
+		return false
 	}
-
-	// Fragments are also valid
-	if l == 1 && d.decoded[0] == layers.LayerTypeIPv4 {
-		ip4 := d.ip4
-		if ip4.Flags&layers.IPv4MoreFragments != 0 || ip4.FragOffset != 0 {
-			return true, true
-		}
-	}
-
-	m.logger.Trace("packet doesn't have network and transport layers")
-	return false, false
+	return true
 }
 
 func (m *Manager) isValidTrackedConnection(d *decoder, srcIP, dstIP netip.Addr, size int) bool {

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 = getLogger()
+		params.Logger = logging.NewDefaultLoggerFactory().NewLogger("ice")
 	}
 
 	mux := &UDPMuxDefault{
@@ -455,9 +455,3 @@ 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 = getLogger()
+		params.Logger = logging.NewDefaultLoggerFactory().NewLogger("ice")
 	}
 	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.ControlPlaneMark
+		return nbnet.NetbirdFwmark
 	}
 	return 0
 }

client/internal/auth/pkce_flow.go

@@ -94,17 +94,13 @@ func (p *PKCEAuthorizationFlow) RequestAuthInfo(ctx context.Context) (AuthFlowIn
 	p.codeVerifier = codeVerifier
 
 	codeChallenge := createCodeChallenge(codeVerifier)
-
-	params := []oauth2.AuthCodeOption{
+	authURL := p.oAuthConfig.AuthCodeURL(
+		state,
 		oauth2.SetAuthURLParam("code_challenge_method", "S256"),
 		oauth2.SetAuthURLParam("code_challenge", codeChallenge),
 		oauth2.SetAuthURLParam("audience", p.providerConfig.Audience),
-	}
-	if !p.providerConfig.DisablePromptLogin {
-		params = append(params, oauth2.SetAuthURLParam("prompt", "login"))
-	}
-
-	authURL := p.oAuthConfig.AuthCodeURL(state, params...)
+		oauth2.SetAuthURLParam("prompt", "login"),
+	)
 
 	return AuthFlowInfo{
 		VerificationURIComplete: authURL,

client/internal/auth/pkce_flow_test.go

@@ -1,49 +0,0 @@
-package auth
-
-import (
-	"context"
-	"testing"
-
-	"github.com/stretchr/testify/require"
-
-	"github.com/netbirdio/netbird/client/internal"
-)
-
-func TestPromptLogin(t *testing.T) {
-	tt := []struct {
-		name   string
-		prompt bool
-	}{
-		{"PromptLogin", true},
-		{"NoPromptLogin", false},
-	}
-
-	for _, tc := range tt {
-		t.Run(tc.name, func(t *testing.T) {
-			config := internal.PKCEAuthProviderConfig{
-				ClientID:              "test-client-id",
-				Audience:              "test-audience",
-				TokenEndpoint:         "https://test-token-endpoint.com/token",
-				Scope:                 "openid email profile",
-				AuthorizationEndpoint: "https://test-auth-endpoint.com/authorize",
-				RedirectURLs:          []string{"http://127.0.0.1:33992/"},
-				UseIDToken:            true,
-				DisablePromptLogin:    !tc.prompt,
-			}
-			pkce, err := NewPKCEAuthorizationFlow(config)
-			if err != nil {
-				t.Fatalf("Failed to create PKCEAuthorizationFlow: %v", err)
-			}
-			authInfo, err := pkce.RequestAuthInfo(context.Background())
-			if err != nil {
-				t.Fatalf("Failed to request auth info: %v", err)
-			}
-			pattern := "prompt=login"
-			if tc.prompt {
-				require.Contains(t, authInfo.VerificationURIComplete, pattern)
-			} else {
-				require.NotContains(t, authInfo.VerificationURIComplete, pattern)
-			}
-		})
-	}
-}

client/internal/dns/file_unix.go

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

client/internal/dns/handler_chain.go

@@ -75,7 +75,12 @@ func (c *HandlerChain) AddHandler(pattern string, handler dns.Handler, priority
 	}
 
 	// First remove any existing handler with same pattern (case-insensitive) and priority
-	c.removeEntry(origPattern, 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
+		}
+	}
 
 	// Check if handler implements SubdomainMatcher interface
 	matchSubdomains := false
@@ -128,20 +133,30 @@ 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:]...)
-			break
+			return
 		}
 	}
 }
 
+// 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,6 +443,14 @@ 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)
+				}
+			}
 		})
 	}
 }
@@ -462,69 +470,45 @@ 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
-	w1 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	// Test 1: Initial state with all three handlers
+	w := &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(w1, r)
+	chain.ServeDNS(w, 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))
 
-	w2 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	w = &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(w2, r)
+	chain.ServeDNS(w, 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))
 
-	w3 := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	w = &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
 	// Now lowest priority handler (defaultHandler) should be called
 	defaultHandler.On("ServeDNS", mock.Anything, r).Return().Once()
 
-	chain.ServeDNS(w3, r)
+	chain.ServeDNS(w, r)
 	defaultHandler.AssertExpectations(t)
 
-	defaultHandler.ExpectedCalls = nil
-	defaultHandler.Calls = nil
-
 	// Test 4: Remove last handler
 	chain.RemoveHandler(testDomain, nbdns.PriorityDefault)
 
-	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)
-	}
+	assert.False(t, chain.HasHandlers(testDomain))
 }
 
 func TestHandlerChain_CaseSensitivity(t *testing.T) {
@@ -846,165 +830,3 @@ 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,8 +5,6 @@ import (
 	"net/netip"
 	"strings"
 
-	"github.com/miekg/dns"
-
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 	nbdns "github.com/netbirdio/netbird/dns"
 )
@@ -14,8 +12,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 {
@@ -105,7 +103,7 @@ func dnsConfigToHostDNSConfig(dnsConfig nbdns.Config, ip string, port int) HostD
 
 		for _, domain := range nsConfig.Domains {
 			config.Domains = append(config.Domains, DomainConfig{
-				Domain:    strings.ToLower(dns.Fqdn(domain)),
+				Domain:    strings.TrimSuffix(domain, "."),
 				MatchOnly: !nsConfig.SearchDomainsEnabled,
 			})
 		}
@@ -114,7 +112,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.ToLower(dns.Fqdn(customZone.Domain)),
+			Domain:    strings.TrimSuffix(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, strings.TrimSuffix(dConf.Domain, "."))
+			matchDomains = append(matchDomains, dConf.Domain)
 			continue
 		}
-		searchDomains = append(searchDomains, strings.TrimSuffix(""+dConf.Domain, "."))
+		searchDomains = append(searchDomains, dConf.Domain)
 	}
 
 	matchKey := getKeyWithInput(netbirdDNSStateKeyFormat, matchSuffix)

client/internal/dns/host_windows.go

@@ -17,12 +17,9 @@ 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 (
@@ -100,9 +97,9 @@ func (r *registryConfigurator) applyDNSConfig(config HostDNSConfig, stateManager
 			continue
 		}
 		if !dConf.MatchOnly {
-			searchDomains = append(searchDomains, strings.TrimSuffix(dConf.Domain, "."))
+			searchDomains = append(searchDomains, dConf.Domain)
 		}
-		matchDomains = append(matchDomains, "."+strings.TrimSuffix(dConf.Domain, "."))
+		matchDomains = append(matchDomains, "."+dConf.Domain)
 	}
 
 	if len(matchDomains) != 0 {
@@ -119,10 +116,6 @@ 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
 }
 
@@ -191,26 +184,6 @@ 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)
@@ -263,10 +236,6 @@ 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/local.go

@@ -71,12 +71,6 @@ func (d *localResolver) lookupRecords(r *dns.Msg) []dns.RR {
 
 	value, found := d.records.Load(key)
 	if !found {
-		// alternatively check if we have a cname
-		if question.Qtype != dns.TypeCNAME {
-			r.Question[0].Qtype = dns.TypeCNAME
-			return d.lookupRecords(r)
-		}
-
 		return nil
 	}
 

client/internal/dns/mock_server.go

@@ -6,7 +6,6 @@ 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
@@ -14,17 +13,17 @@ type MockServer struct {
 	InitializeFunc        func() error
 	StopFunc              func()
 	UpdateDNSServerFunc   func(serial uint64, update nbdns.Config) error
-	RegisterHandlerFunc   func(domain.List, dns.Handler, int)
-	DeregisterHandlerFunc func(domain.List, int)
+	RegisterHandlerFunc   func([]string, dns.Handler, int)
+	DeregisterHandlerFunc func([]string, int)
 }
 
-func (m *MockServer) RegisterHandler(domains domain.List, handler dns.Handler, priority int) {
+func (m *MockServer) RegisterHandler(domains []string, handler dns.Handler, priority int) {
 	if m.RegisterHandlerFunc != nil {
 		m.RegisterHandlerFunc(domains, handler, priority)
 	}
 }
 
-func (m *MockServer) DeregisterHandler(domains domain.List, priority int) {
+func (m *MockServer) DeregisterHandler(domains []string, priority int) {
 	if m.DeregisterHandlerFunc != nil {
 		m.DeregisterHandlerFunc(domains, priority)
 	}

client/internal/dns/network_manager_unix.go

@@ -13,6 +13,7 @@ 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"
@@ -125,10 +126,10 @@ func (n *networkManagerDbusConfigurator) applyDNSConfig(config HostDNSConfig, st
 			continue
 		}
 		if dConf.MatchOnly {
-			matchDomains = append(matchDomains, "~."+dConf.Domain)
+			matchDomains = append(matchDomains, "~."+dns.Fqdn(dConf.Domain))
 			continue
 		}
-		searchDomains = append(searchDomains, dConf.Domain)
+		searchDomains = append(searchDomains, dns.Fqdn(dConf.Domain))
 	}
 
 	newDomainList := append(searchDomains, matchDomains...) //nolint:gocritic

client/internal/dns/server.go

@@ -6,13 +6,11 @@ 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"
@@ -20,7 +18,6 @@ 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
@@ -35,8 +32,8 @@ type IosDnsManager interface {
 
 // Server is a dns server interface
 type Server interface {
-	RegisterHandler(domains domain.List, handler dns.Handler, priority int)
-	DeregisterHandler(domains domain.List, priority int)
+	RegisterHandler(domains []string, handler dns.Handler, priority int)
+	DeregisterHandler(domains []string, priority int)
 	Initialize() error
 	Stop()
 	DnsIP() string
@@ -68,7 +65,6 @@ type DefaultServer struct {
 	previousConfigHash uint64
 	currentConfig      HostDNSConfig
 	handlerChain       *HandlerChain
-	extraDomains       map[domain.Domain]int
 
 	// permanent related properties
 	permanent      bool
@@ -168,15 +164,13 @@ 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: handlerChain,
-		extraDomains: make(map[domain.Domain]int),
+		handlerChain: NewHandlerChain(),
 		dnsMuxMap:    make(registeredHandlerMap),
 		localResolver: &localResolver{
 			registeredMap: make(registrationMap),
@@ -187,26 +181,14 @@ func newDefaultServer(
 		hostsDNSHolder: newHostsDNSHolder(),
 	}
 
-	// register with root zone, handler chain takes care of the routing
-	dnsService.RegisterMux(".", handlerChain)
-
 	return defaultServer
 }
 
-// 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) {
+func (s *DefaultServer) RegisterHandler(domains []string, handler dns.Handler, priority int) {
 	s.mux.Lock()
 	defer s.mux.Unlock()
 
-	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()
+	s.registerHandler(domains, handler, priority)
 }
 
 func (s *DefaultServer) registerHandler(domains []string, handler dns.Handler, priority int) {
@@ -218,23 +200,15 @@ 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)
 	}
 }
 
-// DeregisterHandler deregisters the handler for the given domains with the given priority.
-func (s *DefaultServer) DeregisterHandler(domains domain.List, priority int) {
+func (s *DefaultServer) DeregisterHandler(domains []string, priority int) {
 	s.mux.Lock()
 	defer s.mux.Unlock()
 
-	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()
+	s.deregisterHandler(domains, priority)
 }
 
 func (s *DefaultServer) deregisterHandler(domains []string, priority int) {
@@ -247,6 +221,11 @@ 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))
+		}
 	}
 }
 
@@ -307,8 +286,6 @@ func (s *DefaultServer) Stop() {
 	}
 
 	s.service.Stop()
-
-	maps.Clear(s.extraDomains)
 }
 
 // OnUpdatedHostDNSServer update the DNS servers addresses for root zones
@@ -413,9 +390,7 @@ 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 {
-		if err := s.service.Listen(); err != nil {
-			log.Errorf("failed to start DNS service: %v", err)
-		}
+		_ = s.service.Listen()
 	} else if !s.permanent {
 		s.service.Stop()
 	}
@@ -438,13 +413,17 @@ 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")
-		s.currentConfig.RouteAll = false
+		hostUpdate.RouteAll = false
 	}
 
-	s.applyHostConfig()
+	if err = s.hostManager.applyDNSConfig(hostUpdate, s.stateManager); err != nil {
+		log.Error(err)
+		s.handleErrNoGroupaAll(err)
+	}
 
 	go func() {
 		// persist dns state right away
@@ -462,43 +441,6 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 	return nil
 }
 
-func (s *DefaultServer) applyHostConfig() {
-	if s.hostManager == nil {
-		return
-	}
-
-	// prevent reapplying config if we're shutting down
-	if s.ctx.Err() != 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
@@ -748,7 +690,10 @@ func (s *DefaultServer) upstreamCallbacks(
 			}
 		}
 
-		s.applyHostConfig()
+		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)
+		}
 
 		go func() {
 			if err := s.stateManager.PersistState(s.ctx); err != nil {
@@ -783,7 +728,12 @@ func (s *DefaultServer) upstreamCallbacks(
 			s.registerHandler([]string{nbdns.RootZone}, handler, priority)
 		}
 
-		s.applyHostConfig()
+		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.updateNSState(nsGroup, nil, true)
 	}
@@ -886,13 +836,3 @@ 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

@@ -29,7 +29,6 @@ import (
 	"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()
@@ -39,7 +38,7 @@ type mocWGIface struct {
 }
 
 func (w *mocWGIface) Name() string {
-	return "utun2301"
+	panic("implement me")
 }
 
 func (w *mocWGIface) Address() wgaddr.Address {
@@ -1449,497 +1448,3 @@ 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/systemd_linux.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/godbus/dbus/v5"
+	"github.com/miekg/dns"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/sys/unix"
 
@@ -110,7 +111,7 @@ func (s *systemdDbusConfigurator) applyDNSConfig(config HostDNSConfig, stateMana
 			continue
 		}
 		domainsInput = append(domainsInput, systemdDbusLinkDomainsInput{
-			Domain:    dConf.Domain,
+			Domain:    dns.Fqdn(dConf.Domain),
 			MatchOnly: dConf.MatchOnly,
 		})
 
@@ -150,11 +151,6 @@ 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
 }
 
@@ -167,8 +163,7 @@ func (s *systemdDbusConfigurator) setDomainsForInterface(domainsInput []systemdD
 	if err != nil {
 		return fmt.Errorf("setting domains configuration failed with error: %w", err)
 	}
-
-	return nil
+	return s.flushCaches()
 }
 
 func (s *systemdDbusConfigurator) restoreHostDNS() error {
@@ -188,14 +183,10 @@ func (s *systemdDbusConfigurator) restoreHostDNS() error {
 		return fmt.Errorf("unable to revert link configuration, got error: %w", err)
 	}
 
-	if err := s.flushDNSCache(); err != nil {
-		log.Errorf("failed to flush DNS cache: %v", err)
-	}
-
-	return nil
+	return s.flushCaches()
 }
 
-func (s *systemdDbusConfigurator) flushDNSCache() error {
+func (s *systemdDbusConfigurator) flushCaches() 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

@@ -18,16 +18,14 @@ import (
 	"github.com/miekg/dns"
 	log "github.com/sirupsen/logrus"
 
-	"github.com/netbirdio/netbird/client/iface"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/proto"
 )
 
 const (
-	UpstreamTimeout = 15 * time.Second
-
 	failsTillDeact   = int32(5)
 	reactivatePeriod = 30 * time.Second
+	upstreamTimeout  = 15 * time.Second
 	probeTimeout     = 2 * time.Second
 )
 
@@ -68,7 +66,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,
@@ -108,7 +106,9 @@ func (u *upstreamResolverBase) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 	}()
 
 	log.Tracef("received upstream question: domain=%s type=%v class=%v", r.Question[0].Name, r.Question[0].Qtype, r.Question[0].Qclass)
+	// set the AuthenticatedData flag and the EDNS0 buffer size to 4096 bytes to support larger dns records
 	if r.Extra == nil {
+		r.SetEdns0(4096, false)
 		r.MsgHdr.AuthenticatedData = true
 	}
 
@@ -336,51 +336,3 @@ func (u *upstreamResolverBase) testNameserver(server string, timeout time.Durati
 	_, _, err := u.upstreamClient.exchange(ctx, server, r)
 	return err
 }
-
-// ExchangeWithFallback exchanges a DNS message with the upstream server.
-// It first tries to use UDP, and if it is truncated, it falls back to TCP.
-// If the passed context is nil, this will use Exchange instead of ExchangeContext.
-func ExchangeWithFallback(ctx context.Context, client *dns.Client, r *dns.Msg, upstream string) (*dns.Msg, time.Duration, error) {
-	// MTU - ip + udp headers
-	// Note: this could be sent out on an interface that is not ours, but our MTU should always be lower.
-	client.UDPSize = iface.DefaultMTU - (60 + 8)
-
-	var (
-		rm  *dns.Msg
-		t   time.Duration
-		err error
-	)
-
-	if ctx == nil {
-		rm, t, err = client.Exchange(r, upstream)
-	} else {
-		rm, t, err = client.ExchangeContext(ctx, r, upstream)
-	}
-
-	if err != nil {
-		return nil, t, fmt.Errorf("with udp: %w", err)
-	}
-
-	if rm == nil || !rm.MsgHdr.Truncated {
-		return rm, t, nil
-	}
-
-	log.Tracef("udp response for domain=%s type=%v class=%v is truncated, trying TCP.",
-		r.Question[0].Name, r.Question[0].Qtype, r.Question[0].Qclass)
-
-	client.Net = "tcp"
-
-	if ctx == nil {
-		rm, t, err = client.Exchange(r, upstream)
-	} else {
-		rm, t, err = client.ExchangeContext(ctx, r, upstream)
-	}
-
-	if err != nil {
-		return nil, t, fmt.Errorf("with tcp: %w", err)
-	}
-
-	// TODO: once TCP is implemented, rm.Truncate() if the request came in over UDP
-
-	return rm, t, nil
-}

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_general.go

@@ -34,5 +34,6 @@ func newUpstreamResolver(
 }
 
 func (u *upstreamResolver) exchange(ctx context.Context, upstream string, r *dns.Msg) (rm *dns.Msg, t time.Duration, err error) {
-	return ExchangeWithFallback(ctx, &dns.Client{}, r, upstream)
+	upstreamExchangeClient := &dns.Client{}
+	return upstreamExchangeClient.ExchangeContext(ctx, r, upstream)
 }

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)
 	}
@@ -68,7 +68,7 @@ func (u *upstreamResolverIOS) exchange(ctx context.Context, upstream string, r *
 	}
 
 	// Cannot use client.ExchangeContext because it overwrites our Dialer
-	return ExchangeWithFallback(nil, client, r, upstream)
+	return client.Exchange(r, upstream)
 }
 
 // GetClientPrivate returns a new DNS client bound to the local IP address of the Netbird interface

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

@@ -3,45 +3,34 @@ package dnsfwd
 import (
 	"context"
 	"errors"
-	"math"
 	"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
-	statusRecorder *peer.Status
+	listenAddress string
+	ttl           uint32
+	domains       []string
 
 	dnsServer *dns.Server
 	mux       *dns.ServeMux
-
-	resId sync.Map
 }
 
-func NewDNSForwarder(listenAddress string, ttl uint32, statusRecorder *peer.Status) *DNSForwarder {
+func NewDNSForwarder(listenAddress string, ttl uint32) *DNSForwarder {
 	log.Debugf("creating DNS forwarder with listen_address=%s ttl=%d", listenAddress, ttl)
 	return &DNSForwarder{
-		listenAddress:  listenAddress,
-		ttl:            ttl,
-		statusRecorder: statusRecorder,
+		listenAddress: listenAddress,
+		ttl:           ttl,
 	}
 }
 
-func (f *DNSForwarder) Listen(domains []string, resIds map[string]string) error {
+func (f *DNSForwarder) Listen(domains []string) error {
 	log.Infof("listen DNS forwarder on address=%s", f.listenAddress)
 	mux := dns.NewServeMux()
 
@@ -53,30 +42,22 @@ func (f *DNSForwarder) Listen(domains []string, resIds map[string]string) error
 	f.dnsServer = dnsServer
 	f.mux = mux
 
-	f.UpdateDomains(domains, resIds)
+	f.UpdateDomains(domains)
 
 	return dnsServer.ListenAndServe()
 }
 
-func (f *DNSForwarder) UpdateDomains(domains []string, resIds map[string]string) {
+func (f *DNSForwarder) UpdateDomains(domains []string) {
 	log.Debugf("Updating domains from %v to %v", f.domains, domains)
 
 	for _, d := range f.domains {
 		f.mux.HandleRemove(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
 }
 
@@ -98,87 +79,41 @@ 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
 
-		resp.Rcode = dns.RcodeNotImplemented
-		if err := w.WriteMsg(resp); err != nil {
-			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)
+	ips, err := net.LookupIP(domain)
 	if err != nil {
-		f.handleDNSError(w, resp, domain, err)
-		return
-	}
+		var dnsErr *net.DNSError
 
-	resId := f.getResIdForDomain(strings.TrimSuffix(domain, "."))
-	if resId != "" {
-		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)
+		switch {
+		case errors.As(err, &dnsErr):
+			resp.Rcode = dns.RcodeServerFailure
+			if dnsErr.IsNotFound {
+				// Pass through NXDOMAIN
+				resp.Rcode = dns.RcodeNameError
 			}
-			f.statusRecorder.AddResolvedIPLookupEntry(ipWithSuffix, resId)
-		}
-	}
 
-	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 {
+			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)
 		}
-	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)
+		}
+		return
 	}
 
-	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.Is6() {
+		if ip.To4() == nil {
 			log.Tracef("resolved domain=%s to IPv6=%s", domain, ip)
 			rr := dns.AAAA{
-				AAAA: ip.AsSlice(),
+				AAAA: ip,
 				Hdr: dns.RR_Header{
 					Name:   domain,
 					Rrtype: dns.TypeAAAA,
@@ -190,7 +125,7 @@ func (f *DNSForwarder) addIPsToResponse(resp *dns.Msg, domain string, ips []neti
 		} else {
 			log.Tracef("resolved domain=%s to IPv4=%s", domain, ip)
 			rr := dns.A{
-				A: ip.AsSlice(),
+				A: ip,
 				Hdr: dns.RR_Header{
 					Name:   domain,
 					Rrtype: dns.TypeA,
@@ -202,36 +137,10 @@ func (f *DNSForwarder) addIPsToResponse(resp *dns.Msg, domain string, ips []neti
 		}
 		resp.Answer = append(resp.Answer, respRecord)
 	}
-}
 
-func (f *DNSForwarder) getResIdForDomain(domain string) string {
-	var selectedResId string
-	var bestScore int
-
-	f.resId.Range(func(key, value interface{}) bool {
-		var score int
-		pattern := key.(string)
-
-		switch {
-		case strings.HasPrefix(pattern, "*."):
-			baseDomain := strings.TrimPrefix(pattern, "*.")
-			if domain == baseDomain || strings.HasSuffix(domain, "."+baseDomain) {
-				score = len(baseDomain)
-			}
-		case domain == pattern:
-			score = math.MaxInt
-		default:
-			return true
-		}
-
-		if score > bestScore {
-			bestScore = score
-			selectedResId = value.(string)
-		}
-		return true
-	})
-
-	return selectedResId
+	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/forwarder_test.go

@@ -1,95 +0,0 @@
-package dnsfwd
-
-import (
-	"sync"
-	"testing"
-)
-
-func TestGetResIdForDomain(t *testing.T) {
-	testCases := []struct {
-		name           string
-		storedMappings map[string]string // key: domain pattern, value: resId
-		queryDomain    string
-		expectedResId  string
-	}{
-		{
-			name:           "Empty map returns empty string",
-			storedMappings: map[string]string{},
-			queryDomain:    "example.com",
-			expectedResId:  "",
-		},
-		{
-			name:           "Exact match returns stored resId",
-			storedMappings: map[string]string{"example.com": "res1"},
-			queryDomain:    "example.com",
-			expectedResId:  "res1",
-		},
-		{
-			name:           "Wildcard pattern matches base domain",
-			storedMappings: map[string]string{"*.example.com": "res2"},
-			queryDomain:    "example.com",
-			expectedResId:  "res2",
-		},
-		{
-			name:           "Wildcard pattern matches subdomain",
-			storedMappings: map[string]string{"*.example.com": "res3"},
-			queryDomain:    "foo.example.com",
-			expectedResId:  "res3",
-		},
-		{
-			name:           "Wildcard pattern does not match different domain",
-			storedMappings: map[string]string{"*.example.com": "res4"},
-			queryDomain:    "foo.notexample.com",
-			expectedResId:  "",
-		},
-		{
-			name:           "Non-wildcard pattern does not match subdomain",
-			storedMappings: map[string]string{"example.com": "res5"},
-			queryDomain:    "foo.example.com",
-			expectedResId:  "",
-		},
-		{
-			name: "Exact match over overlapping wildcard",
-			storedMappings: map[string]string{
-				"*.example.com":   "resWildcard",
-				"foo.example.com": "resExact",
-			},
-			queryDomain:   "foo.example.com",
-			expectedResId: "resExact",
-		},
-		{
-			name: "Overlapping wildcards: Select more specific wildcard",
-			storedMappings: map[string]string{
-				"*.example.com":     "resA",
-				"*.sub.example.com": "resB",
-			},
-			queryDomain:   "bar.sub.example.com",
-			expectedResId: "resB",
-		},
-		{
-			name: "Wildcard multi-level subdomain match",
-			storedMappings: map[string]string{
-				"*.example.com": "resMulti",
-			},
-			queryDomain:   "a.b.example.com",
-			expectedResId: "resMulti",
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.name, func(t *testing.T) {
-			fwd := &DNSForwarder{
-				resId: sync.Map{},
-			}
-
-			for domainPattern, resId := range tc.storedMappings {
-				fwd.resId.Store(domainPattern, resId)
-			}
-
-			got := fwd.getResIdForDomain(tc.queryDomain)
-			if got != tc.expectedResId {
-				t.Errorf("For query domain %q, expected resId %q, but got %q", tc.queryDomain, tc.expectedResId, got)
-			}
-		})
-	}
-}

client/internal/dnsfwd/manager.go

@@ -10,7 +10,6 @@ import (
 
 	nberrors "github.com/netbirdio/netbird/client/errors"
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
-	"github.com/netbirdio/netbird/client/internal/peer"
 )
 
 const (
@@ -20,21 +19,19 @@ const (
 )
 
 type Manager struct {
-	firewall       firewall.Manager
-	statusRecorder *peer.Status
+	firewall firewall.Manager
 
 	fwRules      []firewall.Rule
 	dnsForwarder *DNSForwarder
 }
 
-func NewManager(fw firewall.Manager, statusRecorder *peer.Status) *Manager {
+func NewManager(fw firewall.Manager) *Manager {
 	return &Manager{
-		firewall:       fw,
-		statusRecorder: statusRecorder,
+		firewall: fw,
 	}
 }
 
-func (m *Manager) Start(domains []string, resIds map[string]string) error {
+func (m *Manager) Start(domains []string) error {
 	log.Infof("starting DNS forwarder")
 	if m.dnsForwarder != nil {
 		return nil
@@ -44,9 +41,9 @@ func (m *Manager) Start(domains []string, resIds map[string]string) error {
 		return err
 	}
 
-	m.dnsForwarder = NewDNSForwarder(fmt.Sprintf(":%d", ListenPort), dnsTTL, m.statusRecorder)
+	m.dnsForwarder = NewDNSForwarder(fmt.Sprintf(":%d", ListenPort), dnsTTL)
 	go func() {
-		if err := m.dnsForwarder.Listen(domains, resIds); err != nil {
+		if err := m.dnsForwarder.Listen(domains); err != nil {
 			// todo handle close error if it is exists
 			log.Errorf("failed to start DNS forwarder, err: %v", err)
 		}
@@ -55,12 +52,12 @@ func (m *Manager) Start(domains []string, resIds map[string]string) error {
 	return nil
 }
 
-func (m *Manager) UpdateDomains(domains []string, resIds map[string]string) {
+func (m *Manager) UpdateDomains(domains []string) {
 	if m.dnsForwarder == nil {
 		return
 	}
 
-	m.dnsForwarder.UpdateDomains(domains, resIds)
+	m.dnsForwarder.UpdateDomains(domains)
 }
 
 func (m *Manager) Stop(ctx context.Context) error {

client/internal/engine.go

@@ -952,6 +952,11 @@ 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 {
@@ -962,19 +967,14 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
 
 	// DNS forwarder
 	dnsRouteFeatureFlag := toDNSFeatureFlag(networkMap)
-	dnsRouteDomains, resourceIds := toRouteDomains(e.config.WgPrivateKey.PublicKey().String(), networkMap.GetRoutes())
-	e.updateDNSForwarder(dnsRouteFeatureFlag, dnsRouteDomains, resourceIds)
+	dnsRouteDomains := toRouteDomains(e.config.WgPrivateKey.PublicKey().String(), networkMap.GetRoutes())
+	e.updateDNSForwarder(dnsRouteFeatureFlag, dnsRouteDomains)
 
 	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)
@@ -1079,29 +1079,21 @@ func toRoutes(protoRoutes []*mgmProto.Route) []*route.Route {
 	return routes
 }
 
-func toRouteDomains(myPubKey string, protoRoutes []*mgmProto.Route) ([]string, map[string]string) {
+func toRouteDomains(myPubKey string, protoRoutes []*mgmProto.Route) []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, resIds
+	return dnsRoutes
 }
 
 func toDNSConfig(protoDNSConfig *mgmProto.DNSConfig, network *net.IPNet) nbdns.Config {
@@ -1768,7 +1760,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, resIds map[string]string) {
+func (e *Engine) updateDNSForwarder(enabled bool, domains []string) {
 	if !enabled {
 		if e.dnsForwardMgr == nil {
 			return
@@ -1782,15 +1774,15 @@ func (e *Engine) updateDNSForwarder(enabled bool, domains []string, resIds map[s
 	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.statusRecorder)
+			e.dnsForwardMgr = dnsfwd.NewManager(e.firewall)
 
-			if err := e.dnsForwardMgr.Start(domains, resIds); err != nil {
+			if err := e.dnsForwardMgr.Start(domains); 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, resIds)
+			e.dnsForwardMgr.UpdateDomains(domains)
 		}
 	} else if e.dnsForwardMgr != nil {
 		log.Infof("disable domain router service")

client/internal/engine_test.go

@@ -1439,6 +1439,8 @@ func startManagement(t *testing.T, dataDir, testFile string) (*grpc.Server, stri
 	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
 	require.NoError(t, err)
 
+	permissionsManagerMock := permissions.NewManagerMock()
+
 	ctrl := gomock.NewController(t)
 	t.Cleanup(ctrl.Finish)
 	settingsMockManager := settings.NewMockManager(ctrl)
@@ -1447,9 +1449,7 @@ func startManagement(t *testing.T, dataDir, testFile string) (*grpc.Server, stri
 		Return(&types.Settings{}, nil).
 		AnyTimes()
 
-	permissionsManager := permissions.NewManager(store)
-
-	accountManager, err := server.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, ia, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager)
+	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
 	}

client/internal/netflow/conntrack/conntrack.go

@@ -14,7 +14,6 @@ import (
 	"github.com/ti-mo/netfilter"
 
 	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
-	nbnet "github.com/netbirdio/netbird/util/net"
 )
 
 const defaultChannelSize = 100
@@ -177,7 +176,7 @@ func (c *ConnTrack) handleEvent(event nfct.Event) {
 	srcIP := flow.TupleOrig.IP.SourceAddress
 	dstIP := flow.TupleOrig.IP.DestinationAddress
 
-	if !c.relevantFlow(flow.Mark, srcIP, dstIP) {
+	if !c.relevantFlow(srcIP, dstIP) {
 		return
 	}
 
@@ -194,7 +193,7 @@ func (c *ConnTrack) handleEvent(event nfct.Event) {
 	}
 
 	flowID := c.getFlowID(flow.ID)
-	direction := c.inferDirection(flow.Mark, srcIP, dstIP)
+	direction := c.inferDirection(srcIP, dstIP)
 
 	eventType := nftypes.TypeStart
 	eventStr := "New"
@@ -225,14 +224,15 @@ func (c *ConnTrack) handleEvent(event nfct.Event) {
 }
 
 // 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
+func (c *ConnTrack) relevantFlow(srcIP, dstIP netip.Addr) bool {
+	// TODO: filter traffic by interface
+
+	wgnet := c.iface.Address().Network
+	if !wgnet.Contains(srcIP.AsSlice()) && !wgnet.Contains(dstIP.AsSlice()) {
+		return false
 	}
 
-	// fallback if mark rules are not in place
-	wgnet := c.iface.Address().Network
-	return wgnet.Contains(srcIP.AsSlice()) || wgnet.Contains(dstIP.AsSlice())
+	return true
 }
 
 // mapRxPackets maps packet counts to RX based on flow direction
@@ -282,15 +282,7 @@ func (c *ConnTrack) getFlowID(conntrackID uint32) uuid.UUID {
 	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
+func (c *ConnTrack) inferDirection(srcIP, dstIP netip.Addr) nftypes.Direction {
 	wgaddr := c.iface.Address().IP
 	wgnetwork := c.iface.Address().Network
 	src, dst := srcIP.AsSlice(), dstIP.AsSlice()
@@ -306,6 +298,8 @@ func (c *ConnTrack) inferDirection(mark uint32, srcIP, dstIP netip.Addr) nftypes
 	case wgnetwork.Contains(dst):
 		// resource network -> netbird network
 		return nftypes.Egress
+
+		// TODO: handle site2site traffic
 	}
 
 	return nftypes.DirectionUnknown

client/internal/netflow/logger/logger.go

@@ -86,18 +86,18 @@ func (l *Logger) startReceiver() {
 				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)
+			var isExitNode bool
+			if event.Direction == types.Ingress {
+				if !l.wgIfaceIPNet.Contains(net.IP(event.SourceIP.AsSlice())) {
+					event.SourceResourceID, isExitNode = l.statusRecorder.CheckRoutes(event.SourceIP)
+				}
+			} else if event.Direction == types.Egress {
+				if !l.wgIfaceIPNet.Contains(net.IP(event.DestIP.AsSlice())) {
+					event.DestResourceID, isExitNode = l.statusRecorder.CheckRoutes(event.DestIP)
+				}
 			}
 
-			if !l.wgIfaceIPNet.Contains(net.IP(event.DestIP.AsSlice())) {
-				event.DestResourceID, isDestExitNode = l.statusRecorder.CheckRoutes(event.DestIP)
-			}
-
-			if l.shouldStore(eventFields, isSrcExitNode || isDestExitNode) {
+			if l.shouldStore(eventFields, isExitNode) {
 				l.Store.StoreEvent(&event)
 			}
 		}

client/internal/netflow/types/types.go

@@ -10,8 +10,6 @@ import (
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
 )
 
-const ZoneID = 0x1BD0
-
 type Protocol uint8
 
 const (

client/internal/peer/ice/agent.go

@@ -4,7 +4,6 @@ import (
 	"time"
 
 	"github.com/pion/ice/v3"
-	"github.com/pion/logging"
 	"github.com/pion/randutil"
 	log "github.com/sirupsen/logrus"
 
@@ -36,9 +35,6 @@ 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},
@@ -55,7 +51,6 @@ func NewAgent(iFaceDiscover stdnet.ExternalIFaceDiscover, config Config, candida
 		RelayAcceptanceMinWait: &iceRelayAcceptanceMinWait,
 		LocalUfrag:             ufrag,
 		LocalPwd:               pwd,
-		LoggerFactory:          fac,
 	}
 
 	if config.DisableIPv6Discovery {

client/internal/peer/route.go

@@ -2,89 +2,37 @@ 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
-
+	localMap    sync.Map
+	remoteMap   sync.Map
 	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
-		}
+	_, exists := r.localMap.LoadOrStore(route, resourceID)
+	if exists {
+		log.Tracef("resourceID %s already exists in local map", resourceID)
 	}
-
-	// 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
-		}
-	}
+	r.localMap.Delete(route)
 }
 
 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
-		}
+	_, exists := r.remoteMap.LoadOrStore(route, resourceID)
+	if exists {
+		log.Tracef("resourceID %s already exists in remote map", resourceID)
 	}
-
-	// 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
-		}
-	}
+	r.remoteMap.Delete(route)
 }
 
 func (r *routeIDLookup) AddResolvedIP(resourceID string, route netip.Prefix) {
@@ -96,35 +44,37 @@ func (r *routeIDLookup) RemoveResolvedIP(route netip.Prefix) {
 }
 
 // Lookup returns the resource ID for the given IP address
-// and a bool indicating if the IP is an exit node.
+// 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 isExitNode bool
+
+	resId, ok := r.resolvedIPs.Load(ip)
+	if ok {
+		return resId.(string), false
 	}
 
 	var resourceID string
-	var isExitNode bool
+	r.localMap.Range(func(key, value interface{}) bool {
+		pref := key.(netip.Prefix)
+		if pref.Contains(ip) {
+			resourceID = value.(string)
+			isExitNode = pref.Bits() == 0
+			return false
 
-	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()
+		return true
+	})
 
 	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.remoteMap.Range(func(key, value interface{}) bool {
+			pref := key.(netip.Prefix)
+			if pref.Contains(ip) {
+				resourceID = value.(string)
+				isExitNode = pref.Bits() == 0
+				return false
 			}
-		}
-		r.remoteLock.RUnlock()
+			return true
+		})
 	}
 
 	return resourceID, isExitNode

client/internal/peer/status.go

@@ -610,28 +610,6 @@ func (d *Status) RemoveLocalPeerStateRoute(route string) {
 	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()

client/internal/pkce_auth.go

@@ -39,8 +39,6 @@ type PKCEAuthProviderConfig struct {
 	UseIDToken bool
 	// ClientCertPair is used for mTLS authentication to the IDP
 	ClientCertPair *tls.Certificate
-	// DisablePromptLogin makes the PKCE flow to not prompt the user for login
-	DisablePromptLogin bool
 }
 
 // GetPKCEAuthorizationFlowInfo initialize a PKCEAuthorizationFlow instance and return with it
@@ -99,7 +97,6 @@ func GetPKCEAuthorizationFlowInfo(ctx context.Context, privateKey string, mgmURL
 			RedirectURLs:          protoPKCEAuthorizationFlow.GetProviderConfig().GetRedirectURLs(),
 			UseIDToken:            protoPKCEAuthorizationFlow.GetProviderConfig().GetUseIDToken(),
 			ClientCertPair:        clientCert,
-			DisablePromptLogin:    protoPKCEAuthorizationFlow.GetProviderConfig().GetDisablePromptLogin(),
 		},
 	}
 

client/internal/routemanager/dnsinterceptor/handler.go

@@ -6,6 +6,7 @@ import (
 	"net/netip"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/hashicorp/go-multierror"
 	"github.com/miekg/dns"
@@ -59,7 +60,7 @@ func (d *DnsInterceptor) String() string {
 }
 
 func (d *DnsInterceptor) AddRoute(context.Context) error {
-	d.dnsServer.RegisterHandler(d.route.Domains, d, nbdns.PriorityDNSRoute)
+	d.dnsServer.RegisterHandler(d.route.Domains.ToPunycodeList(), d, nbdns.PriorityDNSRoute)
 	return nil
 }
 
@@ -88,7 +89,7 @@ func (d *DnsInterceptor) RemoveRoute() error {
 	clear(d.interceptedDomains)
 	d.mu.Unlock()
 
-	d.dnsServer.DeregisterHandler(d.route.Domains, nbdns.PriorityDNSRoute)
+	d.dnsServer.DeregisterHandler(d.route.Domains.ToPunycodeList(), nbdns.PriorityDNSRoute)
 
 	return nberrors.FormatErrorOrNil(merr)
 }
@@ -141,67 +142,57 @@ func (d *DnsInterceptor) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 	log.Tracef("received DNS request for domain=%s type=%v class=%v",
 		r.Question[0].Name, r.Question[0].Qtype, r.Question[0].Qclass)
 
-	// pass if non A/AAAA query
-	if r.Question[0].Qtype != dns.TypeA && r.Question[0].Qtype != dns.TypeAAAA {
-		d.continueToNextHandler(w, r, "non A/AAAA query")
-		return
-	}
-
 	d.mu.RLock()
 	peerKey := d.currentPeerKey
 	d.mu.RUnlock()
 
 	if peerKey == "" {
-		d.writeDNSError(w, r, "no current peer key")
+		log.Tracef("no current peer key set, letting next handler try for domain=%s", r.Question[0].Name)
+
+		d.continueToNextHandler(w, r, "no current peer key")
 		return
 	}
 
 	upstreamIP, err := d.getUpstreamIP(peerKey)
 	if err != nil {
-		d.writeDNSError(w, r, fmt.Sprintf("get upstream IP: %v", err))
+		log.Errorf("failed to get upstream IP: %v", err)
+		d.continueToNextHandler(w, r, fmt.Sprintf("failed to get upstream IP: %v", err))
 		return
 	}
 
+	// set the AuthenticatedData flag and the EDNS0 buffer size to 4096 bytes to support larger dns records
 	if r.Extra == nil {
+		r.SetEdns0(4096, false)
 		r.MsgHdr.AuthenticatedData = true
 	}
+
 	client := &dns.Client{
-		Timeout: nbdns.UpstreamTimeout,
+		Timeout: 5 * time.Second,
 		Net:     "udp",
 	}
 	upstream := fmt.Sprintf("%s:%d", upstreamIP.String(), dnsfwd.ListenPort)
-	reply, _, err := nbdns.ExchangeWithFallback(context.TODO(), client, r, upstream)
+	reply, _, err := client.ExchangeContext(context.Background(), r, upstream)
+
+	var answer []dns.RR
+	if reply != nil {
+		answer = reply.Answer
+	}
+	log.Tracef("upstream %s (%s) DNS response for domain=%s answers=%v", upstreamIP.String(), peerKey, r.Question[0].Name, answer)
+
 	if err != nil {
-		log.Errorf("failed to exchange DNS request with %s (%s) for domain=%s: %v", upstreamIP.String(), peerKey, r.Question[0].Name, err)
+		log.Errorf("failed to exchange DNS request with %s: %v", upstream, err)
 		if err := w.WriteMsg(&dns.Msg{MsgHdr: dns.MsgHdr{Rcode: dns.RcodeServerFailure, Id: r.Id}}); err != nil {
 			log.Errorf("failed writing DNS response: %v", err)
 		}
 		return
 	}
 
-	var answer []dns.RR
-	if reply != nil {
-		answer = reply.Answer
-	}
-
-	log.Tracef("upstream %s (%s) DNS response for domain=%s answers=%v", upstreamIP.String(), peerKey, r.Question[0].Name, answer)
-
 	reply.Id = r.Id
 	if err := d.writeMsg(w, reply); err != nil {
 		log.Errorf("failed writing DNS response: %v", err)
 	}
 }
 
-func (d *DnsInterceptor) writeDNSError(w dns.ResponseWriter, r *dns.Msg, reason string) {
-	log.Warnf("failed to query upstream for domain=%s: %s", r.Question[0].Name, reason)
-
-	resp := new(dns.Msg)
-	resp.SetRcode(r, dns.RcodeServerFailure)
-	if err := w.WriteMsg(resp); err != nil {
-		log.Errorf("failed to write DNS error response: %v", err)
-	}
-}
-
 // continueToNextHandler signals the handler chain to try the next handler
 func (d *DnsInterceptor) continueToNextHandler(w dns.ResponseWriter, r *dns.Msg, reason string) {
 	log.Tracef("continuing to next handler for domain=%s reason=%s", r.Question[0].Name, reason)

client/internal/routemanager/dynamic/route.go

@@ -235,7 +235,7 @@ func (r *Route) resolve(results chan resolveResult) {
 			ips, err := r.getIPsFromResolver(domain)
 			if err != nil {
 				log.Tracef("Failed to resolve domain %s with private resolver: %v", domain.SafeString(), err)
-				ips, err = net.LookupIP(domain.PunycodeString())
+				ips, err = net.LookupIP(string(domain))
 				if err != nil {
 					results <- resolveResult{domain: domain, err: fmt.Errorf("resolve d %s: %w", domain.SafeString(), err)}
 					return

client/internal/routemanager/dynamic/route_generic.go

@@ -9,5 +9,5 @@ import (
 )
 
 func (r *Route) getIPsFromResolver(domain domain.Domain) ([]net.IP, error) {
-	return net.LookupIP(domain.PunycodeString())
+	return net.LookupIP(string(domain))
 }

client/internal/routemanager/dynamic/route_ios.go

@@ -23,11 +23,11 @@ func (r *Route) getIPsFromResolver(domain domain.Domain) ([]net.IP, error) {
 	}
 
 	msg := new(dns.Msg)
-	msg.SetQuestion(dns.Fqdn(domain.PunycodeString()), dns.TypeA)
+	msg.SetQuestion(dns.Fqdn(string(domain)), dns.TypeA)
 
 	startTime := time.Now()
 
-	response, _, err := nbdns.ExchangeWithFallback(nil, privateClient, msg, r.resolverAddr)
+	response, _, err := privateClient.Exchange(msg, r.resolverAddr)
 	if err != nil {
 		return nil, fmt.Errorf("DNS query for %s failed after %s: %s ", domain.SafeString(), time.Since(startTime), err)
 	}

client/internal/routemanager/server_nonandroid.go

@@ -55,18 +55,6 @@ func (m *serverRouter) updateRoutes(routesMap map[route.ID]*route.Route) error {
 		delete(m.routes, routeID)
 	}
 
-	// If routing is to be disabled, do it after routes have been removed
-	// If routing is to be enabled, do it before adding new routes; addToServerNetwork needs routing to be enabled
-	if len(routesMap) > 0 {
-		if err := m.firewall.EnableRouting(); err != nil {
-			return fmt.Errorf("enable routing: %w", err)
-		}
-	} else {
-		if err := m.firewall.DisableRouting(); err != nil {
-			return fmt.Errorf("disable routing: %w", err)
-		}
-	}
-
 	for id, newRoute := range routesMap {
 		_, found := m.routes[id]
 		if found {
@@ -81,6 +69,16 @@ func (m *serverRouter) updateRoutes(routesMap map[route.ID]*route.Route) error {
 		m.routes[id] = newRoute
 	}
 
+	if len(m.routes) > 0 {
+		if err := m.firewall.EnableRouting(); err != nil {
+			return fmt.Errorf("enable routing: %w", err)
+		}
+	} else {
+		if err := m.firewall.DisableRouting(); err != nil {
+			return fmt.Errorf("disable routing: %w", err)
+		}
+	}
+
 	return nil
 }
 
@@ -105,11 +103,7 @@ func (m *serverRouter) removeFromServerNetwork(route *route.Route) error {
 
 	delete(m.routes, route.ID)
 
-	routeStr := route.Network.String()
-	if route.IsDynamic() {
-		routeStr = route.Domains.SafeString()
-	}
-	m.statusRecorder.RemoveLocalPeerStateRoute(routeStr)
+	m.statusRecorder.RemoveLocalPeerStateRoute(route.Network.String())
 
 	return nil
 }

client/internal/routemanager/systemops/systemops_linux.go

@@ -57,8 +57,8 @@ func getSetupRules() []ruleParams {
 	return []ruleParams{
 		{100, -1, syscall.RT_TABLE_MAIN, netlink.FAMILY_V4, false, 0, "rule with suppress prefixlen v4"},
 		{100, -1, syscall.RT_TABLE_MAIN, netlink.FAMILY_V6, false, 0, "rule with suppress prefixlen v6"},
-		{110, nbnet.ControlPlaneMark, NetbirdVPNTableID, netlink.FAMILY_V4, true, -1, "rule v4 netbird"},
-		{110, nbnet.ControlPlaneMark, NetbirdVPNTableID, netlink.FAMILY_V6, true, -1, "rule v6 netbird"},
+		{110, nbnet.NetbirdFwmark, NetbirdVPNTableID, netlink.FAMILY_V4, true, -1, "rule v4 netbird"},
+		{110, nbnet.NetbirdFwmark, NetbirdVPNTableID, netlink.FAMILY_V6, true, -1, "rule v6 netbird"},
 	}
 }
 

client/server/debug.go

@@ -34,8 +34,7 @@ import (
 )
 
 const readmeContent = `Netbird debug bundle
-This debug bundle contains the following files.
-If the --anonymize flag is set, the files are anonymized to protect sensitive information.
+This debug bundle contains the following files:
 
 status.txt: Anonymized status information of the NetBird client.
 client.log: Most recent, anonymized client log file of the NetBird client.
@@ -51,9 +50,6 @@ state.json: Anonymized client state dump containing netbird states.
 mutex.prof: Mutex profiling information.
 goroutine.prof: Goroutine profiling information.
 block.prof: Block profiling information.
-heap.prof: Heap profiling information (snapshot of memory allocations).
-allocs.prof: Allocations profiling information.
-threadcreate.prof: Thread creation profiling information.
 
 
 Anonymization Process
@@ -96,14 +92,11 @@ The state file follows the same anonymization rules as other files:
 - Domain names are consistently anonymized
 - Technical identifiers and non-sensitive data remain unchanged
 
-Mutex, Goroutines, Block, and Heap Profiling Files
-The goroutine, block, mutex, and heap profiling files contain process information that might help the NetBird team diagnose performance or memory issues. The information in these files doesn't contain personal data.
+Mutex, Goroutines, and Block Profiling Files
+The goroutine, block, and mutex profiling files contains process information that might help the NetBird team diagnose performance issues. The information in these files don't contain personal data.
 You can check each using the following go command:
 
-go tool pprof -http=:8088 <profile_name>.prof
-
-For example, to view the heap profile:
-go tool pprof -http=:8088 heap.prof
+go tool pprof -http=:8088 mutex.prof
 
 This will open a web browser tab with the profiling information.
 
@@ -187,7 +180,7 @@ func (s *Server) DebugBundle(_ context.Context, req *proto.DebugBundleRequest) (
 
 func (s *Server) createArchive(bundlePath *os.File, req *proto.DebugBundleRequest) error {
 	archive := zip.NewWriter(bundlePath)
-	if err := s.addReadme(archive); err != nil {
+	if err := s.addReadme(req, archive); err != nil {
 		return fmt.Errorf("add readme: %w", err)
 	}
 
@@ -207,8 +200,8 @@ func (s *Server) createArchive(bundlePath *os.File, req *proto.DebugBundleReques
 		s.addSystemInfo(req, anonymizer, archive)
 	}
 
-	if err := s.addProf(archive); err != nil {
-		log.Errorf("Failed to add profiles to debug bundle: %v", err)
+	if err := s.addProf(req, anonymizer, archive); err != nil {
+		log.Errorf("Failed to add goroutines rules to debug bundle: %v", err)
 	}
 
 	if err := s.addNetworkMap(req, anonymizer, archive); err != nil {
@@ -249,10 +242,12 @@ func (s *Server) addSystemInfo(req *proto.DebugBundleRequest, anonymizer *anonym
 	}
 }
 
-func (s *Server) addReadme(archive *zip.Writer) error {
-	readmeReader := strings.NewReader(readmeContent)
-	if err := addFileToZip(archive, readmeReader, "README.txt"); err != nil {
-		return fmt.Errorf("add README file to zip: %w", err)
+func (s *Server) addReadme(req *proto.DebugBundleRequest, archive *zip.Writer) error {
+	if req.GetAnonymize() {
+		readmeReader := strings.NewReader(readmeContent)
+		if err := addFileToZip(archive, readmeReader, "README.txt"); err != nil {
+			return fmt.Errorf("add README file to zip: %w", err)
+		}
 	}
 	return nil
 }
@@ -331,13 +326,7 @@ func (s *Server) addCommonConfigFields(configContent *strings.Builder) {
 	configContent.WriteString(fmt.Sprintf("BlockLANAccess: %v\n", s.config.BlockLANAccess))
 }
 
-func (s *Server) addProf(archive *zip.Writer) (err error) {
-	defer func() {
-		if r := recover(); r != nil {
-			err = fmt.Errorf("panic while profiling: %v", r)
-		}
-	}()
-
+func (s *Server) addProf(req *proto.DebugBundleRequest, anonymizer *anonymize.Anonymizer, archive *zip.Writer) error {
 	runtime.SetBlockProfileRate(1)
 	_ = runtime.SetMutexProfileFraction(1)
 	defer runtime.SetBlockProfileRate(0)
@@ -345,7 +334,7 @@ func (s *Server) addProf(archive *zip.Writer) (err error) {
 
 	time.Sleep(5 * time.Second)
 
-	for _, profile := range []string{"goroutine", "block", "mutex", "heap", "allocs", "threadcreate"} {
+	for _, profile := range []string{"goroutine", "block", "mutex"} {
 		var buff []byte
 		myBuff := bytes.NewBuffer(buff)
 		err := pprof.Lookup(profile).WriteTo(myBuff, 0)

client/server/network.go

@@ -100,7 +100,7 @@ func (s *Server) ListNetworks(context.Context, *proto.ListNetworksRequest) (*pro
 
 		// Convert to proto format
 		for domain, ips := range domainMap {
-			pbRoute.ResolvedIPs[domain.PunycodeString()] = &proto.IPList{
+			pbRoute.ResolvedIPs[string(domain)] = &proto.IPList{
 				Ips: ips,
 			}
 		}

client/server/server_test.go

@@ -8,11 +8,10 @@ import (
 	"time"
 
 	"github.com/golang/mock/gomock"
+	"github.com/netbirdio/management-integrations/integrations"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel"
 
-	"github.com/netbirdio/management-integrations/integrations"
-
 	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/grpc"
@@ -201,10 +200,10 @@ func startManagement(t *testing.T, signalAddr string, counter *int) (*grpc.Serve
 	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
 	require.NoError(t, err)
 
+	permissionsManagerMock := permissions.NewManagerMock()
 	ctrl := gomock.NewController(t)
 	t.Cleanup(ctrl.Finish)
 	settingsMockManager := settings.NewMockManager(ctrl)
-	permissionsManagerMock := permissions.NewMockManager(ctrl)
 
 	accountManager, err := server.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, ia, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManagerMock)
 	if err != nil {

formatter/hook/hook.go

@@ -3,7 +3,6 @@ package hook
 import (
 	"fmt"
 	"path"
-	"runtime"
 	"runtime/debug"
 	"strings"
 
@@ -41,12 +40,8 @@ func (hook ContextHook) Levels() []logrus.Level {
 
 // Fire extend with the source information the entry.Data
 func (hook ContextHook) Fire(entry *logrus.Entry) error {
-	caller := &runtime.Frame{Line: 0, File: "caller_not_available"}
-	if entry.Caller != nil {
-		caller = entry.Caller
-	}
-	src := hook.parseSrc(caller.File)
-	entry.Data[EntryKeySource] = fmt.Sprintf("%s:%v", src, caller.Line)
+	src := hook.parseSrc(entry.Caller.File)
+	entry.Data[EntryKeySource] = fmt.Sprintf("%s:%v", src, entry.Caller.Line)
 	additionalEntries(entry)
 
 	if entry.Context == nil {

infrastructure_files/base.setup.env

@@ -58,7 +58,6 @@ NETBIRD_TOKEN_SOURCE=${NETBIRD_TOKEN_SOURCE:-accessToken}
 # PKCE authorization flow
 NETBIRD_AUTH_PKCE_REDIRECT_URL_PORTS=${NETBIRD_AUTH_PKCE_REDIRECT_URL_PORTS:-"53000"}
 NETBIRD_AUTH_PKCE_USE_ID_TOKEN=${NETBIRD_AUTH_PKCE_USE_ID_TOKEN:-false}
-NETBIRD_AUTH_PKCE_DISABLE_PROMPT_LOGIN=${NETBIRD_AUTH_PKCE_DISABLE_PROMPT_LOGIN:-false}
 NETBIRD_AUTH_PKCE_AUDIENCE=$NETBIRD_AUTH_AUDIENCE
 
 # Dashboard
@@ -121,7 +120,6 @@ export NETBIRD_AUTH_DEVICE_AUTH_SCOPE
 export NETBIRD_AUTH_DEVICE_AUTH_USE_ID_TOKEN
 export NETBIRD_AUTH_PKCE_AUTHORIZATION_ENDPOINT
 export NETBIRD_AUTH_PKCE_USE_ID_TOKEN
-export NETBIRD_AUTH_PKCE_DISABLE_PROMPT_LOGIN
 export NETBIRD_AUTH_PKCE_AUDIENCE
 export NETBIRD_DASH_AUTH_USE_AUDIENCE
 export NETBIRD_DASH_AUTH_AUDIENCE

infrastructure_files/management.json.tmpl

@@ -94,8 +94,7 @@
             "TokenEndpoint": "$NETBIRD_AUTH_TOKEN_ENDPOINT",
             "Scope": "$NETBIRD_AUTH_SUPPORTED_SCOPES",
             "RedirectURLs": [$NETBIRD_AUTH_PKCE_REDIRECT_URLS],
-            "UseIDToken": $NETBIRD_AUTH_PKCE_USE_ID_TOKEN,
-            "DisablePromptLogin": $NETBIRD_AUTH_PKCE_DISABLE_PROMPT_LOGIN
+            "UseIDToken": $NETBIRD_AUTH_PKCE_USE_ID_TOKEN
         }
     }
 }

infrastructure_files/tests/setup.env

@@ -27,4 +27,3 @@ NETBIRD_STORE_CONFIG_ENGINE=$CI_NETBIRD_STORE_CONFIG_ENGINE
 NETBIRD_MGMT_IDP_SIGNKEY_REFRESH=$CI_NETBIRD_MGMT_IDP_SIGNKEY_REFRESH
 NETBIRD_TURN_EXTERNAL_IP=1.2.3.4
 NETBIRD_RELAY_PORT=33445
-NETBIRD_AUTH_PKCE_DISABLE_PROMPT_LOGIN=true

management/client/client_test.go

@@ -75,6 +75,7 @@ func startManagement(t *testing.T) (*grpc.Server, net.Listener) {
 	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
 	require.NoError(t, err)
 
+	permissionsManagerMock := permissions.NewManagerMock()
 	ctrl := gomock.NewController(t)
 	t.Cleanup(ctrl.Finish)
 	settingsMockManager := settings.NewMockManager(ctrl)
@@ -87,18 +88,6 @@ func startManagement(t *testing.T) (*grpc.Server, net.Listener) {
 		).
 		Return(&types.Settings{}, nil).
 		AnyTimes()
-	permissionsManagerMock := permissions.NewMockManager(ctrl)
-	permissionsManagerMock.
-		EXPECT().
-		ValidateUserPermissions(
-			gomock.Any(),
-			gomock.Any(),
-			gomock.Any(),
-			gomock.Any(),
-			gomock.Any(),
-		).
-		Return(true, nil).
-		AnyTimes()
 
 	accountManager, err := mgmt.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, ia, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManagerMock)
 	if err != nil {

management/client/rest/users.go

@@ -80,16 +80,3 @@ func (a *UsersAPI) ResendInvitation(ctx context.Context, userID string) error {
 
 	return nil
 }
-
-// Current gets the current user info
-// See more: https://docs.netbird.io/api/resources/users#retrieve-current-user
-func (a *UsersAPI) Current(ctx context.Context) (*api.User, error) {
-	resp, err := a.c.newRequest(ctx, "GET", "/api/users/current", nil)
-	if err != nil {
-		return nil, err
-	}
-	defer resp.Body.Close()
-
-	ret, err := parseResponse[api.User](resp)
-	return &ret, err
-}

management/client/rest/users_test.go

@@ -196,42 +196,8 @@ func TestUsers_ResendInvitation_Err(t *testing.T) {
 	})
 }
 
-func TestUsers_Current_200(t *testing.T) {
-	withMockClient(func(c *rest.Client, mux *http.ServeMux) {
-		mux.HandleFunc("/api/users/current", func(w http.ResponseWriter, r *http.Request) {
-			retBytes, _ := json.Marshal(testUser)
-			_, err := w.Write(retBytes)
-			require.NoError(t, err)
-		})
-		ret, err := c.Users.Current(context.Background())
-		require.NoError(t, err)
-		assert.Equal(t, testUser, *ret)
-	})
-}
-
-func TestUsers_Current_Err(t *testing.T) {
-	withMockClient(func(c *rest.Client, mux *http.ServeMux) {
-		mux.HandleFunc("/api/users/current", func(w http.ResponseWriter, r *http.Request) {
-			retBytes, _ := json.Marshal(util.ErrorResponse{Message: "No", Code: 400})
-			w.WriteHeader(400)
-			_, err := w.Write(retBytes)
-			require.NoError(t, err)
-		})
-		ret, err := c.Users.Current(context.Background())
-		assert.Error(t, err)
-		assert.Equal(t, "No", err.Error())
-		assert.Empty(t, ret)
-	})
-}
-
 func TestUsers_Integration(t *testing.T) {
 	withBlackBoxServer(t, func(c *rest.Client) {
-		// rest client PAT is owner's
-		current, err := c.Users.Current(context.Background())
-		require.NoError(t, err)
-		assert.Equal(t, "a23efe53-63fb-11ec-90d6-0242ac120003", current.Id)
-		assert.Equal(t, "owner", current.Role)
-
 		user, err := c.Users.Create(context.Background(), api.UserCreateRequest{
 			AutoGroups:    []string{},
 			Email:         ptr("test@example.com"),

management/domain/domain.go

@@ -24,11 +24,6 @@ func (d Domain) SafeString() string {
 	return str
 }
 
-// PunycodeString returns the punycode representation of the Domain.
-func (d Domain) PunycodeString() string {
-	return string(d)
-}
-
 // FromString creates a Domain from a string, converting it to punycode.
 func FromString(s string) (Domain, error) {
 	ascii, err := idna.ToASCII(s)

management/proto/management.pb.go

@@ -1,7 +1,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.26.0
-// 	protoc        v3.21.9
+// 	protoc        v4.24.3
 // source: management.proto
 
 package proto
@@ -1447,11 +1447,11 @@ type FlowConfig struct {
 	TokenSignature string               `protobuf:"bytes,3,opt,name=tokenSignature,proto3" json:"tokenSignature,omitempty"`
 	Interval       *durationpb.Duration `protobuf:"bytes,4,opt,name=interval,proto3" json:"interval,omitempty"`
 	Enabled        bool                 `protobuf:"varint,5,opt,name=enabled,proto3" json:"enabled,omitempty"`
-	// counters determines if flow packets and bytes counters should be sent
+	// Counters determines if flow packets and bytes counters should be sent
 	Counters bool `protobuf:"varint,6,opt,name=counters,proto3" json:"counters,omitempty"`
-	// exitNodeCollection determines if event collection on exit nodes should be enabled
+	// ExitNodeCollection determines if event collection on exit nodes should be enabled
 	ExitNodeCollection bool `protobuf:"varint,7,opt,name=exitNodeCollection,proto3" json:"exitNodeCollection,omitempty"`
-	// dnsCollection determines if DNS event collection should be enabled
+	// DnsCollection determines if DNS event collection should be enabled
 	DnsCollection bool `protobuf:"varint,8,opt,name=dnsCollection,proto3" json:"dnsCollection,omitempty"`
 }
 
@@ -2192,8 +2192,6 @@ type ProviderConfig struct {
 	AuthorizationEndpoint string `protobuf:"bytes,9,opt,name=AuthorizationEndpoint,proto3" json:"AuthorizationEndpoint,omitempty"`
 	// RedirectURLs handles authorization code from IDP manager
 	RedirectURLs []string `protobuf:"bytes,10,rep,name=RedirectURLs,proto3" json:"RedirectURLs,omitempty"`
-	// DisablePromptLogin makes the PKCE flow to not prompt the user for login
-	DisablePromptLogin bool `protobuf:"varint,11,opt,name=DisablePromptLogin,proto3" json:"DisablePromptLogin,omitempty"`
 }
 
 func (x *ProviderConfig) Reset() {
@@ -2298,13 +2296,6 @@ func (x *ProviderConfig) GetRedirectURLs() []string {
 	return nil
 }
 
-func (x *ProviderConfig) GetDisablePromptLogin() bool {
-	if x != nil {
-		return x.DisablePromptLogin
-	}
-	return false
-}
-
 // Route represents a route.Route object
 type Route struct {
 	state         protoimpl.MessageState
@@ -3587,7 +3578,7 @@ var file_management_proto_rawDesc = []byte{
 	0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x6d,
 	0x61, 0x6e, 0x61, 0x67, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2e, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64,
 	0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x52, 0x0e, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64,
-	0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x22, 0x9a, 0x03, 0x0a, 0x0e, 0x50, 0x72, 0x6f,
+	0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x22, 0xea, 0x02, 0x0a, 0x0e, 0x50, 0x72, 0x6f,
 	0x76, 0x69, 0x64, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x12, 0x1a, 0x0a, 0x08, 0x43,
 	0x6c, 0x69, 0x65, 0x6e, 0x74, 0x49, 0x44, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x43,
 	0x6c, 0x69, 0x65, 0x6e, 0x74, 0x49, 0x44, 0x12, 0x22, 0x0a, 0x0c, 0x43, 0x6c, 0x69, 0x65, 0x6e,
@@ -3610,10 +3601,7 @@ var file_management_proto_rawDesc = []byte{
 	0x6f, 0x72, 0x69, 0x7a, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x70, 0x6f, 0x69, 0x6e,
 	0x74, 0x12, 0x22, 0x0a, 0x0c, 0x52, 0x65, 0x64, 0x69, 0x72, 0x65, 0x63, 0x74, 0x55, 0x52, 0x4c,
 	0x73, 0x18, 0x0a, 0x20, 0x03, 0x28, 0x09, 0x52, 0x0c, 0x52, 0x65, 0x64, 0x69, 0x72, 0x65, 0x63,
-	0x74, 0x55, 0x52, 0x4c, 0x73, 0x12, 0x2e, 0x0a, 0x12, 0x44, 0x69, 0x73, 0x61, 0x62, 0x6c, 0x65,
-	0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x4c, 0x6f, 0x67, 0x69, 0x6e, 0x18, 0x0b, 0x20, 0x01, 0x28,
-	0x08, 0x52, 0x12, 0x44, 0x69, 0x73, 0x61, 0x62, 0x6c, 0x65, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74,
-	0x4c, 0x6f, 0x67, 0x69, 0x6e, 0x22, 0xed, 0x01, 0x0a, 0x05, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x12,
+	0x74, 0x55, 0x52, 0x4c, 0x73, 0x22, 0xed, 0x01, 0x0a, 0x05, 0x52, 0x6f, 0x75, 0x74, 0x65, 0x12,
 	0x0e, 0x0a, 0x02, 0x49, 0x44, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x49, 0x44, 0x12,
 	0x18, 0x0a, 0x07, 0x4e, 0x65, 0x74, 0x77, 0x6f, 0x72, 0x6b, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09,
 	0x52, 0x07, 0x4e, 0x65, 0x74, 0x77, 0x6f, 0x72, 0x6b, 0x12, 0x20, 0x0a, 0x0b, 0x4e, 0x65, 0x74,

management/proto/management.proto

@@ -372,8 +372,6 @@ message ProviderConfig {
   string AuthorizationEndpoint = 9;
   // RedirectURLs handles authorization code from IDP manager
   repeated string RedirectURLs = 10;
-  // DisablePromptLogin makes the PKCE flow to not prompt the user for login
-  bool DisablePromptLogin = 11;
 }
 
 // Route represents a route.Route object

management/server/account.go

@@ -6,14 +6,11 @@ import (
 	"fmt"
 	"math/rand"
 	"net"
-	"os"
 	"reflect"
 	"regexp"
 	"slices"
-	"strconv"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	cacheStore "github.com/eko/gocache/lib/v4/store"
@@ -33,8 +30,6 @@ import (
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/permissions"
-	"github.com/netbirdio/netbird/management/server/permissions/modules"
-	"github.com/netbirdio/netbird/management/server/permissions/operations"
 	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/netbirdio/netbird/management/server/settings"
 	"github.com/netbirdio/netbird/management/server/status"
@@ -97,9 +92,6 @@ type DefaultAccountManager struct {
 	metrics telemetry.AppMetrics
 
 	permissionsManager permissions.Manager
-
-	accountUpdateLocks               sync.Map
-	updateAccountPeersBufferInterval atomic.Int64
 }
 
 // getJWTGroupsChanges calculates the changes needed to sync a user's JWT groups.
@@ -194,23 +186,6 @@ func BuildManager(
 		settingsManager:          settingsManager,
 		permissionsManager:       permissionsManager,
 	}
-
-	var initialInterval int64
-	intervalStr := os.Getenv("PEER_UPDATE_INTERVAL_MS")
-	interval, err := strconv.Atoi(intervalStr)
-	if err != nil {
-		initialInterval = 1
-	} else {
-		initialInterval = int64(interval) * 10
-		go func() {
-			time.Sleep(30 * time.Second)
-			am.updateAccountPeersBufferInterval.Store(int64(time.Duration(interval) * time.Millisecond))
-			log.WithContext(ctx).Infof("set peer update buffer interval to %dms", interval)
-		}()
-	}
-	am.updateAccountPeersBufferInterval.Store(initialInterval)
-	log.WithContext(ctx).Infof("set peer update buffer interval to %dms", initialInterval)
-
 	accountsCounter, err := store.GetAccountsCounter(ctx)
 	if err != nil {
 		log.WithContext(ctx).Error(err)
@@ -283,7 +258,7 @@ func (am *DefaultAccountManager) UpdateAccountSettings(ctx context.Context, acco
 		return nil, err
 	}
 
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Settings, operations.Write)
+	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, permissions.Settings, permissions.Write)
 	if err != nil {
 		return nil, fmt.Errorf("failed to validate user permissions: %w", err)
 	}
@@ -533,7 +508,7 @@ func (am *DefaultAccountManager) DeleteAccount(ctx context.Context, accountID, u
 		return err
 	}
 
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Accounts, operations.Write)
+	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, permissions.Accounts, permissions.Write)
 	if err != nil {
 		return fmt.Errorf("failed to validate user permissions: %w", err)
 	}
@@ -1046,12 +1021,13 @@ func (am *DefaultAccountManager) GetAccount(ctx context.Context, accountID strin
 
 // GetAccountByID returns an account associated with this account ID.
 func (am *DefaultAccountManager) GetAccountByID(ctx context.Context, accountID string, userID string) (*types.Account, error) {
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Accounts, operations.Read)
+	user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthShare, userID)
 	if err != nil {
-		return nil, status.NewPermissionValidationError(err)
+		return nil, err
 	}
-	if !allowed {
-		return nil, status.NewPermissionDeniedError()
+
+	if err := am.permissionsManager.ValidateAccountAccess(ctx, accountID, user, false); err != nil {
+		return nil, err
 	}
 
 	return am.Store.GetAccount(ctx, accountID)
@@ -1247,7 +1223,7 @@ func (am *DefaultAccountManager) SyncUserJWTGroups(ctx context.Context, userAuth
 
 		if removedGroupAffectsPeers || newGroupsAffectsPeers {
 			log.WithContext(ctx).Tracef("user %s: JWT group membership changed, updating account peers", userAuth.UserId)
-			am.BufferUpdateAccountPeers(ctx, userAuth.AccountId)
+			am.UpdateAccountPeers(ctx, userAuth.AccountId)
 		}
 	}
 
@@ -1486,7 +1462,7 @@ func (am *DefaultAccountManager) GetDNSDomain() string {
 
 func (am *DefaultAccountManager) onPeersInvalidated(ctx context.Context, accountID string) {
 	log.WithContext(ctx).Debugf("validated peers has been invalidated for account %s", accountID)
-	am.BufferUpdateAccountPeers(ctx, accountID)
+	am.UpdateAccountPeers(ctx, accountID)
 }
 
 func (am *DefaultAccountManager) FindExistingPostureCheck(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error) {
@@ -1539,13 +1515,19 @@ func (am *DefaultAccountManager) getFreeDNSLabel(ctx context.Context, s store.St
 }
 
 func (am *DefaultAccountManager) GetAccountSettings(ctx context.Context, accountID string, userID string) (*types.Settings, error) {
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Settings, operations.Read)
+	user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthShare, userID)
 	if err != nil {
-		return nil, status.NewPermissionValidationError(err)
+		return nil, err
 	}
-	if !allowed {
-		return nil, status.NewPermissionDeniedError()
+
+	if err := am.permissionsManager.ValidateAccountAccess(ctx, accountID, user, false); err != nil {
+		return nil, err
 	}
+
+	if !user.HasAdminPower() && !user.IsServiceUser {
+		return nil, status.Errorf(status.PermissionDenied, "the user has no permission to access account data")
+	}
+
 	return am.Store.GetAccountSettings(ctx, store.LockingStrengthShare, accountID)
 }
 

management/server/account/manager.go

@@ -114,5 +114,4 @@ type Manager interface {
 	CreateAccountByPrivateDomain(ctx context.Context, initiatorId, domain string) (*types.Account, error)
 	UpdateToPrimaryAccount(ctx context.Context, accountId string) (*types.Account, error)
 	GetOwnerInfo(ctx context.Context, accountId string) (*types.UserInfo, error)
-	GetCurrentUserInfo(ctx context.Context, accountID, userID string) (*types.UserInfo, error)
 }

management/server/account_test.go

@@ -2794,13 +2794,13 @@ func TestAccount_UserGroupsRemoveFromPeers(t *testing.T) {
 	})
 }
 
-// type TB interface {
+//type TB interface {
 //	Cleanup(func())
 //	Helper()
 //	TempDir() string
 //	Errorf(format string, args ...interface{})
 //	Fatalf(format string, args ...interface{})
-// }
+//}
 
 func createManager(t testing.TB) (*DefaultAccountManager, error) {
 	t.Helper()
@@ -2816,6 +2816,8 @@ func createManager(t testing.TB) (*DefaultAccountManager, error) {
 		return nil, err
 	}
 
+	permissionsManagerMock := permissions.NewManagerMock()
+
 	ctrl := gomock.NewController(t)
 	t.Cleanup(ctrl.Finish)
 
@@ -2829,9 +2831,7 @@ func createManager(t testing.TB) (*DefaultAccountManager, error) {
 		Return(false, nil).
 		AnyTimes()
 
-	permissionsManager := permissions.NewManager(store)
-
-	manager, err := BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.cloud", eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager)
+	manager, err := BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.cloud", eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManagerMock)
 	if err != nil {
 		return nil, err
 	}

management/server/dns.go

@@ -10,8 +10,6 @@ import (
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/proto"
 	"github.com/netbirdio/netbird/management/server/activity"
-	"github.com/netbirdio/netbird/management/server/permissions/modules"
-	"github.com/netbirdio/netbird/management/server/permissions/operations"
 	"github.com/netbirdio/netbird/management/server/status"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/types"
@@ -64,12 +62,17 @@ func (c *DNSConfigCache) SetNameServerGroup(key string, value *proto.NameServerG
 
 // GetDNSSettings validates a user role and returns the DNS settings for the provided account ID
 func (am *DefaultAccountManager) GetDNSSettings(ctx context.Context, accountID string, userID string) (*types.DNSSettings, error) {
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Dns, operations.Read)
+	user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthShare, userID)
 	if err != nil {
-		return nil, status.NewPermissionValidationError(err)
+		return nil, err
 	}
-	if !allowed {
-		return nil, status.NewPermissionDeniedError()
+
+	if err := am.permissionsManager.ValidateAccountAccess(ctx, accountID, user, false); err != nil {
+		return nil, err
+	}
+
+	if user.IsRegularUser() {
+		return nil, status.NewAdminPermissionError()
 	}
 
 	return am.Store.GetAccountDNSSettings(ctx, store.LockingStrengthShare, accountID)
@@ -81,12 +84,17 @@ func (am *DefaultAccountManager) SaveDNSSettings(ctx context.Context, accountID
 		return status.Errorf(status.InvalidArgument, "the dns settings provided are nil")
 	}
 
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Dns, operations.Write)
+	user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthShare, userID)
 	if err != nil {
-		return status.NewPermissionValidationError(err)
+		return err
 	}
-	if !allowed {
-		return status.NewPermissionDeniedError()
+
+	if err := am.permissionsManager.ValidateAccountAccess(ctx, accountID, user, false); err != nil {
+		return err
+	}
+
+	if !user.HasAdminPower() {
+		return status.NewAdminPermissionError()
 	}
 
 	var updateAccountPeers bool

management/server/dns_test.go

@@ -211,13 +211,14 @@ func createDNSManager(t *testing.T) (*DefaultAccountManager, error) {
 
 	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
 	require.NoError(t, err)
+	permissionsManagerMock := permissions.NewManagerMock()
 
 	ctrl := gomock.NewController(t)
 	t.Cleanup(ctrl.Finish)
 
 	settingsMockManager := settings.NewMockManager(ctrl)
-	permissionsManager := permissions.NewManager(store)
-	return BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.test", eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager)
+
+	return BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.test", eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManagerMock)
 }
 
 func createDNSStore(t *testing.T) (store.Store, error) {

management/server/event.go

@@ -9,8 +9,6 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/management/server/activity"
-	"github.com/netbirdio/netbird/management/server/permissions/modules"
-	"github.com/netbirdio/netbird/management/server/permissions/operations"
 	"github.com/netbirdio/netbird/management/server/status"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/types"
@@ -23,12 +21,17 @@ func isEnabled() bool {
 
 // GetEvents returns a list of activity events of an account
 func (am *DefaultAccountManager) GetEvents(ctx context.Context, accountID, userID string) ([]*activity.Event, error) {
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Events, operations.Read)
+	user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthShare, userID)
 	if err != nil {
-		return nil, status.NewPermissionValidationError(err)
+		return nil, err
 	}
-	if !allowed {
-		return nil, status.NewPermissionDeniedError()
+
+	if err := am.permissionsManager.ValidateAccountAccess(ctx, accountID, user, false); err != nil {
+		return nil, err
+	}
+
+	if !(user.HasAdminPower() || user.IsServiceUser) {
+		return nil, status.Errorf(status.PermissionDenied, "only users with admin power can view events")
 	}
 
 	events, err := am.eventStore.Get(ctx, accountID, 0, 10000, true)
@@ -53,7 +56,7 @@ func (am *DefaultAccountManager) GetEvents(ctx context.Context, accountID, userI
 		filtered = append(filtered, event)
 	}
 
-	err = am.fillEventsWithUserInfo(ctx, events, accountID, userID)
+	err = am.fillEventsWithUserInfo(ctx, events, accountID, user)
 	if err != nil {
 		return nil, err
 	}
@@ -86,8 +89,8 @@ type eventUserInfo struct {
 	accountId string
 }
 
-func (am *DefaultAccountManager) fillEventsWithUserInfo(ctx context.Context, events []*activity.Event, accountId string, userId string) error {
-	eventUserInfo, err := am.getEventsUserInfo(ctx, events, accountId, userId)
+func (am *DefaultAccountManager) fillEventsWithUserInfo(ctx context.Context, events []*activity.Event, accountId string, user *types.User) error {
+	eventUserInfo, err := am.getEventsUserInfo(ctx, events, accountId, user)
 	if err != nil {
 		return err
 	}
@@ -102,14 +105,14 @@ func (am *DefaultAccountManager) fillEventsWithUserInfo(ctx context.Context, eve
 	return nil
 }
 
-func (am *DefaultAccountManager) getEventsUserInfo(ctx context.Context, events []*activity.Event, accountId string, userId string) (map[string]eventUserInfo, error) {
+func (am *DefaultAccountManager) getEventsUserInfo(ctx context.Context, events []*activity.Event, accountId string, user *types.User) (map[string]eventUserInfo, error) {
 	accountUsers, err := am.Store.GetAccountUsers(ctx, store.LockingStrengthShare, accountId)
 	if err != nil {
 		return nil, err
 	}
 
 	// @note check whether using a external initiator user here is an issue
-	userInfos, err := am.BuildUserInfosForAccount(ctx, accountId, userId, accountUsers)
+	userInfos, err := am.BuildUserInfosForAccount(ctx, accountId, user.Id, accountUsers)
 	if err != nil {
 		return nil, err
 	}
@@ -143,10 +146,10 @@ func (am *DefaultAccountManager) getEventsUserInfo(ctx context.Context, events [
 		return eventUserInfos, nil
 	}
 
-	return am.getEventsExternalUserInfo(ctx, externalUserIds, eventUserInfos, userId)
+	return am.getEventsExternalUserInfo(ctx, externalUserIds, eventUserInfos, user)
 }
 
-func (am *DefaultAccountManager) getEventsExternalUserInfo(ctx context.Context, externalUserIds []string, eventUserInfos map[string]eventUserInfo, userId string) (map[string]eventUserInfo, error) {
+func (am *DefaultAccountManager) getEventsExternalUserInfo(ctx context.Context, externalUserIds []string, eventUserInfos map[string]eventUserInfo, user *types.User) (map[string]eventUserInfo, error) {
 	externalAccountId := ""
 	fetched := make(map[string]struct{})
 	externalUsers := []*types.User{}
@@ -179,7 +182,7 @@ func (am *DefaultAccountManager) getEventsExternalUserInfo(ctx context.Context,
 		return eventUserInfos, nil
 	}
 
-	externalUserInfos, err := am.BuildUserInfosForAccount(ctx, externalAccountId, userId, externalUsers)
+	externalUserInfos, err := am.BuildUserInfosForAccount(ctx, externalAccountId, user.Id, externalUsers)
 	if err != nil {
 		return nil, err
 	}

management/server/group.go

@@ -12,8 +12,6 @@ import (
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/server/activity"
 	routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
-	"github.com/netbirdio/netbird/management/server/permissions/modules"
-	"github.com/netbirdio/netbird/management/server/permissions/operations"
 	"github.com/netbirdio/netbird/management/server/status"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/types"
@@ -32,13 +30,17 @@ func (e *GroupLinkError) Error() string {
 
 // CheckGroupPermissions validates if a user has the necessary permissions to view groups
 func (am *DefaultAccountManager) CheckGroupPermissions(ctx context.Context, accountID, userID string) error {
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Groups, operations.Read)
+	user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthShare, userID)
 	if err != nil {
 		return err
 	}
 
-	if !allowed {
-		return status.NewPermissionDeniedError()
+	if err := am.permissionsManager.ValidateAccountAccess(ctx, accountID, user, false); err != nil {
+		return err
+	}
+
+	if user.IsRegularUser() {
+		return status.NewAdminPermissionError()
 	}
 
 	return nil
@@ -76,12 +78,17 @@ func (am *DefaultAccountManager) SaveGroup(ctx context.Context, accountID, userI
 // Note: This function does not acquire the global lock.
 // It is the caller's responsibility to ensure proper locking is in place before invoking this method.
 func (am *DefaultAccountManager) SaveGroups(ctx context.Context, accountID, userID string, groups []*types.Group) error {
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Groups, operations.Write)
+	user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthShare, userID)
 	if err != nil {
-		return status.NewPermissionValidationError(err)
+		return err
 	}
-	if !allowed {
-		return status.NewPermissionDeniedError()
+
+	if err := am.permissionsManager.ValidateAccountAccess(ctx, accountID, user, false); err != nil {
+		return err
+	}
+
+	if user.IsRegularUser() {
+		return status.NewAdminPermissionError()
 	}
 
 	var eventsToStore []func()
@@ -203,12 +210,17 @@ func (am *DefaultAccountManager) DeleteGroup(ctx context.Context, accountID, use
 // If an error occurs while deleting a group, the function skips it and continues deleting other groups.
 // Errors are collected and returned at the end.
 func (am *DefaultAccountManager) DeleteGroups(ctx context.Context, accountID, userID string, groupIDs []string) error {
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Groups, operations.Write)
+	user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthShare, userID)
 	if err != nil {
-		return status.NewPermissionValidationError(err)
+		return err
 	}
-	if !allowed {
-		return status.NewPermissionDeniedError()
+
+	if err := am.permissionsManager.ValidateAccountAccess(ctx, accountID, user, false); err != nil {
+		return err
+	}
+
+	if user.IsRegularUser() {
+		return status.NewAdminPermissionError()
 	}
 
 	var allErrors error

management/server/groups/manager.go

@@ -8,8 +8,6 @@ import (
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/http/api"
 	"github.com/netbirdio/netbird/management/server/permissions"
-	"github.com/netbirdio/netbird/management/server/permissions/modules"
-	"github.com/netbirdio/netbird/management/server/permissions/operations"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/types"
 )
@@ -41,7 +39,7 @@ func NewManager(store store.Store, permissionsManager permissions.Manager, accou
 }
 
 func (m *managerImpl) GetAllGroups(ctx context.Context, accountID, userID string) ([]*types.Group, error) {
-	ok, err := m.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Groups, operations.Read)
+	ok, err := m.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, permissions.Groups, permissions.Read)
 	if err != nil {
 		return nil, err
 	}
@@ -72,7 +70,7 @@ func (m *managerImpl) GetAllGroupsMap(ctx context.Context, accountID, userID str
 }
 
 func (m *managerImpl) AddResourceToGroup(ctx context.Context, accountID, userID, groupID string, resource *types.Resource) error {
-	ok, err := m.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Groups, operations.Write)
+	ok, err := m.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, permissions.Groups, permissions.Write)
 	if err != nil {
 		return err
 	}

management/server/grpcserver.go

@@ -828,7 +828,6 @@ func (s *GRPCServer) GetPKCEAuthorizationFlow(ctx context.Context, req *proto.En
 			Scope:                 s.config.PKCEAuthorizationFlow.ProviderConfig.Scope,
 			RedirectURLs:          s.config.PKCEAuthorizationFlow.ProviderConfig.RedirectURLs,
 			UseIDToken:            s.config.PKCEAuthorizationFlow.ProviderConfig.UseIDToken,
-			DisablePromptLogin:    s.config.PKCEAuthorizationFlow.ProviderConfig.DisablePromptLogin,
 		},
 	}
 

management/server/http/api/openapi.yml

@@ -2397,29 +2397,6 @@ paths:
           "$ref": "#/components/responses/forbidden"
         '500':
           "$ref": "#/components/responses/internal_error"
-  /api/users/current:
-    get:
-      summary: Retrieve current user
-      description: Get information about the current user
-      tags: [ Users ]
-      security:
-        - BearerAuth: [ ]
-        - TokenAuth: [ ]
-      responses:
-        '200':
-          description: A User object
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/User'
-        '400':
-          "$ref": "#/components/responses/bad_request"
-        '401':
-          "$ref": "#/components/responses/requires_authentication"
-        '403':
-          "$ref": "#/components/responses/forbidden"
-        '500':
-          "$ref": "#/components/responses/internal_error"
   /api/peers:
     get:
       summary: List all Peers

management/server/http/api/types.gen.go

@@ -230,7 +230,7 @@ type Account struct {
 
 // AccountExtraSettings defines model for AccountExtraSettings.
 type AccountExtraSettings struct {
-	// NetworkTrafficLogsEnabled Enables or disables network traffic logging. If enabled, all network traffic events from peers will be stored.
+	// NetworkTrafficLogsEnabled Enables or disables network traffic logs. If enabled, all network traffic logs from peers will be stored.
 	NetworkTrafficLogsEnabled bool `json:"network_traffic_logs_enabled"`
 
 	// NetworkTrafficPacketCounterEnabled Enables or disables network traffic packet counter. If enabled, network packets and their size will be counted and reported. (This can have an slight impact on performance)

management/server/http/handler.go

@@ -66,13 +66,15 @@ func NewAPIHandler(
 
 	corsMiddleware := cors.AllowAll()
 
+	acMiddleware := middleware.NewAccessControl(accountManager.GetUserFromUserAuth)
+
 	rootRouter := mux.NewRouter()
 	metricsMiddleware := appMetrics.HTTPMiddleware()
 
 	prefix := apiPrefix
 	router := rootRouter.PathPrefix(prefix).Subrouter()
 
-	router.Use(metricsMiddleware.Handler, corsMiddleware.Handler, authMiddleware.Handler)
+	router.Use(metricsMiddleware.Handler, corsMiddleware.Handler, authMiddleware.Handler, acMiddleware.Handler)
 
 	if _, err := integrations.RegisterHandlers(ctx, prefix, router, accountManager, integratedValidator, appMetrics.GetMeter(), permissionsManager, peersManager, proxyController, settingsManager); err != nil {
 		return nil, fmt.Errorf("register integrations endpoints: %w", err)
@@ -83,8 +85,6 @@ func NewAPIHandler(
 	users.AddEndpoints(accountManager, router)
 	setup_keys.AddEndpoints(accountManager, router)
 	policies.AddEndpoints(accountManager, LocationManager, router)
-	policies.AddPostureCheckEndpoints(accountManager, LocationManager, router)
-	policies.AddLocationsEndpoints(accountManager, LocationManager, permissionsManager, router)
 	groups.AddEndpoints(accountManager, router)
 	routes.AddEndpoints(accountManager, router)
 	dns.AddEndpoints(accountManager, router)

management/server/http/handlers/peers/peers_handler.go

@@ -10,7 +10,6 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/management/server/account"
-	"github.com/netbirdio/netbird/management/server/activity"
 	nbcontext "github.com/netbirdio/netbird/management/server/context"
 	"github.com/netbirdio/netbird/management/server/groups"
 	"github.com/netbirdio/netbird/management/server/http/api"
@@ -245,7 +244,7 @@ func (h *Handler) GetAccessiblePeers(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	account, err := h.accountManager.GetAccountByID(r.Context(), accountID, activity.SystemInitiator)
+	account, err := h.accountManager.GetAccountByID(r.Context(), accountID, userID)
 	if err != nil {
 		util.WriteError(r.Context(), err, w)
 		return

management/server/http/handlers/policies/geolocation_handler_test.go

@@ -10,7 +10,6 @@ import (
 	"path/filepath"
 	"testing"
 
-	"github.com/golang/mock/gomock"
 	"github.com/gorilla/mux"
 	"github.com/stretchr/testify/assert"
 
@@ -18,9 +17,6 @@ import (
 	"github.com/netbirdio/netbird/management/server/geolocation"
 	"github.com/netbirdio/netbird/management/server/http/api"
 	"github.com/netbirdio/netbird/management/server/mock_server"
-	"github.com/netbirdio/netbird/management/server/permissions"
-	"github.com/netbirdio/netbird/management/server/permissions/modules"
-	"github.com/netbirdio/netbird/management/server/permissions/operations"
 	"github.com/netbirdio/netbird/management/server/types"
 	"github.com/netbirdio/netbird/util"
 )
@@ -45,14 +41,6 @@ func initGeolocationTestData(t *testing.T) *geolocationsHandler {
 	assert.NoError(t, err)
 	t.Cleanup(func() { _ = geo.Stop() })
 
-	ctrl := gomock.NewController(t)
-	permissionsManagerMock := permissions.NewMockManager(ctrl)
-	permissionsManagerMock.
-		EXPECT().
-		ValidateUserPermissions(gomock.Any(), gomock.Any(), gomock.Any(), modules.Policies, operations.Read).
-		Return(true, nil).
-		AnyTimes()
-
 	return &geolocationsHandler{
 		accountManager: &mock_server.MockAccountManager{
 			GetUserByIDFunc: func(ctx context.Context, id string) (*types.User, error) {
@@ -60,7 +48,6 @@ func initGeolocationTestData(t *testing.T) *geolocationsHandler {
 			},
 		},
 		geolocationManager: geo,
-		permissionsManager: permissionsManagerMock,
 	}
 }
 

management/server/http/handlers/policies/geolocations_handler.go

@@ -11,9 +11,6 @@ import (
 	"github.com/netbirdio/netbird/management/server/geolocation"
 	"github.com/netbirdio/netbird/management/server/http/api"
 	"github.com/netbirdio/netbird/management/server/http/util"
-	"github.com/netbirdio/netbird/management/server/permissions"
-	"github.com/netbirdio/netbird/management/server/permissions/modules"
-	"github.com/netbirdio/netbird/management/server/permissions/operations"
 	"github.com/netbirdio/netbird/management/server/status"
 )
 
@@ -25,21 +22,19 @@ var (
 type geolocationsHandler struct {
 	accountManager     account.Manager
 	geolocationManager geolocation.Geolocation
-	permissionsManager permissions.Manager
 }
 
-func AddLocationsEndpoints(accountManager account.Manager, locationManager geolocation.Geolocation, permissionsManager permissions.Manager, router *mux.Router) {
-	locationHandler := newGeolocationsHandlerHandler(accountManager, locationManager, permissionsManager)
+func addLocationsEndpoint(accountManager account.Manager, locationManager geolocation.Geolocation, router *mux.Router) {
+	locationHandler := newGeolocationsHandlerHandler(accountManager, locationManager)
 	router.HandleFunc("/locations/countries", locationHandler.getAllCountries).Methods("GET", "OPTIONS")
 	router.HandleFunc("/locations/countries/{country}/cities", locationHandler.getCitiesByCountry).Methods("GET", "OPTIONS")
 }
 
 // newGeolocationsHandlerHandler creates a new Geolocations handler
-func newGeolocationsHandlerHandler(accountManager account.Manager, geolocationManager geolocation.Geolocation, permissionsManager permissions.Manager) *geolocationsHandler {
+func newGeolocationsHandlerHandler(accountManager account.Manager, geolocationManager geolocation.Geolocation) *geolocationsHandler {
 	return &geolocationsHandler{
 		accountManager:     accountManager,
 		geolocationManager: geolocationManager,
-		permissionsManager: permissionsManager,
 	}
 }
 
@@ -103,22 +98,20 @@ func (l *geolocationsHandler) getCitiesByCountry(w http.ResponseWriter, r *http.
 }
 
 func (l *geolocationsHandler) authenticateUser(r *http.Request) error {
-	ctx := r.Context()
-
-	userAuth, err := nbcontext.GetUserAuthFromContext(ctx)
+	userAuth, err := nbcontext.GetUserAuthFromContext(r.Context())
 	if err != nil {
 		return err
 	}
 
-	accountID, userID := userAuth.AccountId, userAuth.UserId
+	_, userID := userAuth.AccountId, userAuth.UserId
 
-	allowed, err := l.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Policies, operations.Read)
+	user, err := l.accountManager.GetUserByID(r.Context(), userID)
 	if err != nil {
-		return status.NewPermissionValidationError(err)
+		return err
 	}
 
-	if !allowed {
-		return status.NewPermissionDeniedError()
+	if !user.HasAdminPower() {
+		return status.Errorf(status.PermissionDenied, "user is not allowed to perform this action")
 	}
 	return nil
 }

management/server/http/handlers/policies/policies_handler.go

@@ -28,6 +28,7 @@ func AddEndpoints(accountManager account.Manager, locationManager geolocation.Ge
 	router.HandleFunc("/policies/{policyId}", policiesHandler.updatePolicy).Methods("PUT", "OPTIONS")
 	router.HandleFunc("/policies/{policyId}", policiesHandler.getPolicy).Methods("GET", "OPTIONS")
 	router.HandleFunc("/policies/{policyId}", policiesHandler.deletePolicy).Methods("DELETE", "OPTIONS")
+	addPostureCheckEndpoint(accountManager, locationManager, router)
 }
 
 // newHandler creates a new policies handler

management/server/http/handlers/policies/posture_checks_handler.go

@@ -21,13 +21,14 @@ type postureChecksHandler struct {
 	geolocationManager geolocation.Geolocation
 }
 
-func AddPostureCheckEndpoints(accountManager account.Manager, locationManager geolocation.Geolocation, router *mux.Router) {
+func addPostureCheckEndpoint(accountManager account.Manager, locationManager geolocation.Geolocation, router *mux.Router) {
 	postureCheckHandler := newPostureChecksHandler(accountManager, locationManager)
 	router.HandleFunc("/posture-checks", postureCheckHandler.getAllPostureChecks).Methods("GET", "OPTIONS")
 	router.HandleFunc("/posture-checks", postureCheckHandler.createPostureCheck).Methods("POST", "OPTIONS")
 	router.HandleFunc("/posture-checks/{postureCheckId}", postureCheckHandler.updatePostureCheck).Methods("PUT", "OPTIONS")
 	router.HandleFunc("/posture-checks/{postureCheckId}", postureCheckHandler.getPostureCheck).Methods("GET", "OPTIONS")
 	router.HandleFunc("/posture-checks/{postureCheckId}", postureCheckHandler.deletePostureCheck).Methods("DELETE", "OPTIONS")
+	addLocationsEndpoint(accountManager, locationManager, router)
 }
 
 // newPostureChecksHandler creates a new PostureChecks handler

management/server/http/handlers/users/users_handler.go

@@ -25,7 +25,6 @@ type handler struct {
 func AddEndpoints(accountManager account.Manager, router *mux.Router) {
 	userHandler := newHandler(accountManager)
 	router.HandleFunc("/users", userHandler.getAllUsers).Methods("GET", "OPTIONS")
-	router.HandleFunc("/users/current", userHandler.getCurrentUser).Methods("GET", "OPTIONS")
 	router.HandleFunc("/users/{userId}", userHandler.updateUser).Methods("PUT", "OPTIONS")
 	router.HandleFunc("/users/{userId}", userHandler.deleteUser).Methods("DELETE", "OPTIONS")
 	router.HandleFunc("/users", userHandler.createUser).Methods("POST", "OPTIONS")
@@ -260,29 +259,6 @@ func (h *handler) inviteUser(w http.ResponseWriter, r *http.Request) {
 	util.WriteJSONObject(r.Context(), w, util.EmptyObject{})
 }
 
-func (h *handler) getCurrentUser(w http.ResponseWriter, r *http.Request) {
-	if r.Method != http.MethodGet {
-		util.WriteErrorResponse("wrong HTTP method", http.StatusMethodNotAllowed, w)
-		return
-	}
-	ctx := r.Context()
-	userAuth, err := nbcontext.GetUserAuthFromContext(ctx)
-	if err != nil {
-		util.WriteError(r.Context(), err, w)
-		return
-	}
-
-	accountID, userID := userAuth.AccountId, userAuth.UserId
-
-	user, err := h.accountManager.GetCurrentUserInfo(ctx, accountID, userID)
-	if err != nil {
-		util.WriteError(r.Context(), err, w)
-		return
-	}
-
-	util.WriteJSONObject(r.Context(), w, toUserResponse(user, userID))
-}
-
 func toUserResponse(user *types.UserInfo, currenUserID string) *api.User {
 	autoGroups := user.AutoGroups
 	if autoGroups == nil {

management/server/http/handlers/users/users_handler_test.go

@@ -9,7 +9,6 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
-	"time"
 
 	"github.com/gorilla/mux"
 	"github.com/stretchr/testify/assert"
@@ -124,64 +123,6 @@ func initUsersTestData() *handler {
 
 				return nil
 			},
-			GetCurrentUserInfoFunc: func(ctx context.Context, accountID, userID string) (*types.UserInfo, error) {
-				switch userID {
-				case "not-found":
-					return nil, status.NewUserNotFoundError("not-found")
-				case "not-of-account":
-					return nil, status.NewUserNotPartOfAccountError()
-				case "blocked-user":
-					return nil, status.NewUserBlockedError()
-				case "service-user":
-					return nil, status.NewPermissionDeniedError()
-				case "owner":
-					return &types.UserInfo{
-						ID:            "owner",
-						Name:          "",
-						Role:          "owner",
-						Status:        "active",
-						IsServiceUser: false,
-						IsBlocked:     false,
-						NonDeletable:  false,
-						Issued:        "api",
-						Permissions: types.UserPermissions{
-							DashboardView: "full",
-						},
-					}, nil
-				case "regular-user":
-					return &types.UserInfo{
-						ID:            "regular-user",
-						Name:          "",
-						Role:          "user",
-						Status:        "active",
-						IsServiceUser: false,
-						IsBlocked:     false,
-						NonDeletable:  false,
-						Issued:        "api",
-						Permissions: types.UserPermissions{
-							DashboardView: "limited",
-						},
-					}, nil
-
-				case "admin-user":
-					return &types.UserInfo{
-						ID:            "admin-user",
-						Name:          "",
-						Role:          "admin",
-						Status:        "active",
-						IsServiceUser: false,
-						IsBlocked:     false,
-						NonDeletable:  false,
-						LastLogin:     time.Time{},
-						Issued:        "api",
-						Permissions: types.UserPermissions{
-							DashboardView: "full",
-						},
-					}, nil
-				}
-
-				return nil, fmt.Errorf("user id %s not handled", userID)
-			},
 		},
 	}
 }
@@ -540,73 +481,3 @@ func TestDeleteUser(t *testing.T) {
 		})
 	}
 }
-
-func TestCurrentUser(t *testing.T) {
-	tt := []struct {
-		name           string
-		expectedStatus int
-		requestAuth    nbcontext.UserAuth
-	}{
-		{
-			name:           "without auth",
-			expectedStatus: http.StatusInternalServerError,
-		},
-		{
-			name:           "user not found",
-			requestAuth:    nbcontext.UserAuth{UserId: "not-found"},
-			expectedStatus: http.StatusNotFound,
-		},
-		{
-			name:           "not of account",
-			requestAuth:    nbcontext.UserAuth{UserId: "not-of-account"},
-			expectedStatus: http.StatusForbidden,
-		},
-		{
-			name:           "blocked user",
-			requestAuth:    nbcontext.UserAuth{UserId: "blocked-user"},
-			expectedStatus: http.StatusForbidden,
-		},
-		{
-			name:           "service user",
-			requestAuth:    nbcontext.UserAuth{UserId: "service-user"},
-			expectedStatus: http.StatusForbidden,
-		},
-		{
-			name:           "owner",
-			requestAuth:    nbcontext.UserAuth{UserId: "owner"},
-			expectedStatus: http.StatusOK,
-		},
-		{
-			name:           "regular user",
-			requestAuth:    nbcontext.UserAuth{UserId: "regular-user"},
-			expectedStatus: http.StatusOK,
-		},
-		{
-			name:           "admin user",
-			requestAuth:    nbcontext.UserAuth{UserId: "admin-user"},
-			expectedStatus: http.StatusOK,
-		},
-	}
-
-	userHandler := initUsersTestData()
-	for _, tc := range tt {
-		t.Run(tc.name, func(t *testing.T) {
-			req := httptest.NewRequest(http.MethodGet, "/api/users/current", nil)
-			if tc.requestAuth.UserId != "" {
-				req = nbcontext.SetUserAuthInRequest(req, tc.requestAuth)
-			}
-
-			rr := httptest.NewRecorder()
-
-			userHandler.getCurrentUser(rr, req)
-
-			res := rr.Result()
-			defer res.Body.Close()
-
-			if status := rr.Code; status != tc.expectedStatus {
-				t.Fatalf("handler returned wrong status code: got %v want %v",
-					status, tc.expectedStatus)
-			}
-		})
-	}
-}

management/server/http/middleware/access_control.go

@@ -0,0 +1,77 @@
+package middleware
+
+import (
+	"context"
+	"net/http"
+	"regexp"
+
+	log "github.com/sirupsen/logrus"
+
+	nbcontext "github.com/netbirdio/netbird/management/server/context"
+	"github.com/netbirdio/netbird/management/server/http/middleware/bypass"
+	"github.com/netbirdio/netbird/management/server/http/util"
+	"github.com/netbirdio/netbird/management/server/status"
+	"github.com/netbirdio/netbird/management/server/types"
+)
+
+// GetUser function defines a function to fetch user from Account by jwtclaims.AuthorizationClaims
+type GetUser func(ctx context.Context, userAuth nbcontext.UserAuth) (*types.User, error)
+
+// AccessControl middleware to restrict to make POST/PUT/DELETE requests by admin only
+type AccessControl struct {
+	getUser GetUser
+}
+
+// NewAccessControl instance constructor
+func NewAccessControl(getUser GetUser) *AccessControl {
+	return &AccessControl{
+		getUser: getUser,
+	}
+}
+
+var tokenPathRegexp = regexp.MustCompile(`^.*/api/users/.*/tokens.*$`)
+
+// Handler method of the middleware which forbids all modify requests for non admin users
+func (a *AccessControl) Handler(h http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+
+		if bypass.ShouldBypass(r.URL.Path, h, w, r) {
+			return
+		}
+
+		userAuth, err := nbcontext.GetUserAuthFromRequest(r)
+		if err != nil {
+			log.WithContext(r.Context()).Errorf("failed to get user auth from request: %s", err)
+			util.WriteError(r.Context(), status.Errorf(status.Unauthorized, "invalid user auth"), w)
+		}
+
+		user, err := a.getUser(r.Context(), userAuth)
+		if err != nil {
+			log.WithContext(r.Context()).Errorf("failed to get user: %s", err)
+			util.WriteError(r.Context(), status.Errorf(status.Unauthorized, "invalid user auth"), w)
+			return
+		}
+
+		if user.IsBlocked() {
+			util.WriteError(r.Context(), status.Errorf(status.PermissionDenied, "the user has no access to the API or is blocked"), w)
+			return
+		}
+
+		if !user.HasAdminPower() {
+			switch r.Method {
+			case http.MethodDelete, http.MethodPost, http.MethodPatch, http.MethodPut:
+
+				if tokenPathRegexp.MatchString(r.URL.Path) {
+					log.WithContext(r.Context()).Debugf("valid Path")
+					h.ServeHTTP(w, r)
+					return
+				}
+
+				util.WriteError(r.Context(), status.Errorf(status.PermissionDenied, "only users with admin power can perform this operation"), w)
+				return
+			}
+		}
+
+		h.ServeHTTP(w, r)
+	})
+}

management/server/http/middleware/auth_middleware.go

@@ -142,6 +142,12 @@ func (m *AuthMiddleware) checkPATFromRequest(r *http.Request, auth []string) (*h
 		return r, fmt.Errorf("token expired")
 	}
 
+	if impersonate, ok := r.URL.Query()["account"]; ok && len(impersonate) == 1 {
+		if user.AccountID != impersonate[0] {
+			return r, fmt.Errorf("token is not valid for this account")
+		}
+	}
+
 	err = m.authManager.MarkPATUsed(ctx, pat.ID)
 	if err != nil {
 		return r, err

management/server/http/middleware/auth_middleware_test.go

@@ -239,16 +239,9 @@ func TestAuthMiddleware_Handler_Child(t *testing.T) {
 			},
 		},
 		{
-			name:       "Valid PAT Token ignores child",
+			name:       "Valid PAT Token returns forbidden for child",
 			path:       "/test?account=xyz",
 			authHeader: "Token " + PAT,
-			expectedUserAuth: &nbcontext.UserAuth{
-				AccountId:      accountID,
-				UserId:         userID,
-				Domain:         testAccount.Domain,
-				DomainCategory: testAccount.DomainCategory,
-				IsPAT:          true,
-			},
 		},
 		{
 			name:       "Valid JWT Token",

management/server/http/testing/testing_tools/tools.go

@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"github.com/golang-jwt/jwt"
+
 	"github.com/stretchr/testify/assert"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 
@@ -122,9 +123,9 @@ func BuildApiBlackBoxWithDBState(t TB, sqlFile string, expectedPeerUpdate *serve
 	validatorMock := server.MocIntegratedValidator{}
 	proxyController := integrations.NewController(store)
 	userManager := users.NewManager(store)
-	permissionsManager := permissions.NewManager(store)
-	settingsManager := settings.NewManager(store, userManager, integrations.NewManager(&activity.InMemoryEventStore{}), permissionsManager)
-	am, err := server.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "", &activity.InMemoryEventStore{}, geoMock, false, validatorMock, metrics, proxyController, settingsManager, permissionsManager)
+	permissionsManagerMock := permissions.NewManagerMock()
+	settingsManager := settings.NewManager(store, userManager, integrations.NewManager(&activity.InMemoryEventStore{}), permissionsManagerMock)
+	am, err := server.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "", &activity.InMemoryEventStore{}, geoMock, false, validatorMock, metrics, proxyController, settingsManager, permissionsManagerMock)
 	if err != nil {
 		t.Fatalf("Failed to create manager: %v", err)
 	}
@@ -142,9 +143,9 @@ func BuildApiBlackBoxWithDBState(t TB, sqlFile string, expectedPeerUpdate *serve
 	resourcesManagerMock := resources.NewManagerMock()
 	routersManagerMock := routers.NewManagerMock()
 	groupsManagerMock := groups.NewManagerMock()
-	peersManager := peers.NewManager(store, permissionsManager)
+	peersManager := peers.NewManager(store, permissionsManagerMock)
 
-	apiHandler, err := nbhttp.NewAPIHandler(context.Background(), am, networksManagerMock, resourcesManagerMock, routersManagerMock, groupsManagerMock, geoMock, authManagerMock, metrics, validatorMock, proxyController, permissionsManager, peersManager, settingsManager)
+	apiHandler, err := nbhttp.NewAPIHandler(context.Background(), am, networksManagerMock, resourcesManagerMock, routersManagerMock, groupsManagerMock, geoMock, authManagerMock, metrics, validatorMock, proxyController, permissionsManagerMock, peersManager, settingsManager)
 	if err != nil {
 		t.Fatalf("Failed to create API handler: %v", err)
 	}

management/server/management_proto_test.go

@@ -432,6 +432,8 @@ func startManagementForTest(t *testing.T, testFile string, config *types.Config)
 	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
 	require.NoError(t, err)
 
+	permissionsManagerMock := permissions.NewManagerMock()
+
 	ctrl := gomock.NewController(t)
 	t.Cleanup(ctrl.Finish)
 	settingsMockManager := settings.NewMockManager(ctrl)
@@ -441,10 +443,8 @@ func startManagementForTest(t *testing.T, testFile string, config *types.Config)
 		AnyTimes().
 		Return(&types.Settings{}, nil)
 
-	permissionsManager := permissions.NewManager(store)
-
 	accountManager, err := BuildManager(ctx, store, peersUpdateManager, nil, "", "netbird.selfhosted",
-		eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManager)
+		eventStore, nil, false, MocIntegratedValidator{}, metrics, port_forwarding.NewControllerMock(), settingsMockManager, permissionsManagerMock)
 
 	if err != nil {
 		cleanup()

management/server/management_test.go

@@ -195,7 +195,7 @@ func startServer(
 		Return(&types.Settings{}, nil).
 		AnyTimes()
 
-	permissionsManager := permissions.NewManager(str)
+	permissionsManagerMock := permissions.NewManagerMock()
 	accountManager, err := server.BuildManager(
 		context.Background(),
 		str,
@@ -210,7 +210,7 @@ func startServer(
 		metrics,
 		port_forwarding.NewControllerMock(),
 		settingsMockManager,
-		permissionsManager,
+		permissionsManagerMock,
 	)
 	if err != nil {
 		t.Fatalf("failed creating an account manager: %v", err)

management/server/mock_server/account_mock.go

@@ -115,7 +115,6 @@ type MockAccountManager struct {
 	CreateAccountByPrivateDomainFunc    func(ctx context.Context, initiatorId, domain string) (*types.Account, error)
 	UpdateToPrimaryAccountFunc          func(ctx context.Context, accountId string) (*types.Account, error)
 	GetOwnerInfoFunc                    func(ctx context.Context, accountID string) (*types.UserInfo, error)
-	GetCurrentUserInfoFunc              func(ctx context.Context, accountID, userID string) (*types.UserInfo, error)
 }
 
 func (am *MockAccountManager) UpdateAccountPeers(ctx context.Context, accountID string) {
@@ -872,10 +871,3 @@ func (am *MockAccountManager) GetOwnerInfo(ctx context.Context, accountId string
 	}
 	return nil, status.Errorf(codes.Unimplemented, "method GetOwnerInfo is not implemented")
 }
-
-func (am *MockAccountManager) GetCurrentUserInfo(ctx context.Context, accountID, userID string) (*types.UserInfo, error) {
-	if am.GetCurrentUserInfoFunc != nil {
-		return am.GetCurrentUserInfoFunc(ctx, accountID, userID)
-	}
-	return nil, status.Errorf(codes.Unimplemented, "method GetCurrentUserInfo is not implemented")
-}

management/server/nameserver.go

@@ -11,8 +11,6 @@ import (
 
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/server/activity"
-	"github.com/netbirdio/netbird/management/server/permissions/modules"
-	"github.com/netbirdio/netbird/management/server/permissions/operations"
 	"github.com/netbirdio/netbird/management/server/status"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/types"
@@ -22,12 +20,17 @@ const domainPattern = `^(?i)[a-z0-9]+([\-\.]{1}[a-z0-9]+)*\.[a-z]{2,}$`
 
 // GetNameServerGroup gets a nameserver group object from account and nameserver group IDs
 func (am *DefaultAccountManager) GetNameServerGroup(ctx context.Context, accountID, userID, nsGroupID string) (*nbdns.NameServerGroup, error) {
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Nameservers, operations.Read)
+	user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthShare, userID)
 	if err != nil {
-		return nil, status.NewPermissionValidationError(err)
+		return nil, err
 	}
-	if !allowed {
-		return nil, status.NewPermissionDeniedError()
+
+	if err := am.permissionsManager.ValidateAccountAccess(ctx, accountID, user, false); err != nil {
+		return nil, err
+	}
+
+	if user.IsRegularUser() {
+		return nil, status.NewAdminPermissionError()
 	}
 
 	return am.Store.GetNameServerGroupByID(ctx, store.LockingStrengthShare, accountID, nsGroupID)
@@ -38,12 +41,13 @@ func (am *DefaultAccountManager) CreateNameServerGroup(ctx context.Context, acco
 	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Nameservers, operations.Write)
+	user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthShare, userID)
 	if err != nil {
-		return nil, status.NewPermissionValidationError(err)
+		return nil, err
 	}
-	if !allowed {
-		return nil, status.NewPermissionDeniedError()
+
+	if err := am.permissionsManager.ValidateAccountAccess(ctx, accountID, user, false); err != nil {
+		return nil, err
 	}
 
 	newNSGroup := &nbdns.NameServerGroup{
@@ -99,12 +103,13 @@ func (am *DefaultAccountManager) SaveNameServerGroup(ctx context.Context, accoun
 		return status.Errorf(status.InvalidArgument, "nameserver group provided is nil")
 	}
 
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Nameservers, operations.Write)
+	user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthShare, userID)
 	if err != nil {
-		return status.NewPermissionValidationError(err)
+		return err
 	}
-	if !allowed {
-		return status.NewPermissionDeniedError()
+
+	if err := am.permissionsManager.ValidateAccountAccess(ctx, accountID, user, false); err != nil {
+		return err
 	}
 
 	var updateAccountPeers bool
@@ -149,12 +154,13 @@ func (am *DefaultAccountManager) DeleteNameServerGroup(ctx context.Context, acco
 	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Nameservers, operations.Write)
+	user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthShare, userID)
 	if err != nil {
-		return status.NewPermissionValidationError(err)
+		return err
 	}
-	if !allowed {
-		return status.NewPermissionDeniedError()
+
+	if err := am.permissionsManager.ValidateAccountAccess(ctx, accountID, user, false); err != nil {
+		return err
 	}
 
 	var nsGroup *nbdns.NameServerGroup
@@ -192,12 +198,17 @@ func (am *DefaultAccountManager) DeleteNameServerGroup(ctx context.Context, acco
 
 // ListNameServerGroups returns a list of nameserver groups from account
 func (am *DefaultAccountManager) ListNameServerGroups(ctx context.Context, accountID string, userID string) ([]*nbdns.NameServerGroup, error) {
-	allowed, err := am.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Nameservers, operations.Read)
+	user, err := am.Store.GetUserByUserID(ctx, store.LockingStrengthShare, userID)
 	if err != nil {
-		return nil, status.NewPermissionValidationError(err)
+		return nil, err
 	}
-	if !allowed {
-		return nil, status.NewPermissionDeniedError()
+
+	if err := am.permissionsManager.ValidateAccountAccess(ctx, accountID, user, false); err != nil {
+		return nil, err
+	}
+
+	if user.IsRegularUser() {
+		return nil, status.NewAdminPermissionError()
 	}
 
 	return am.Store.GetAccountNameServerGroups(ctx, store.LockingStrengthShare, accountID)