Merge remote-tracking branch 'origin/main' into fix/userspace-native-firewall

client/cmd/debug.go

@@ -199,11 +199,9 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 		cmd.Println("Log level set to trace.")
 	}
 
-	needsRestoreUp := false
 	if _, err := client.Down(cmd.Context(), &proto.DownRequest{}); err != nil {
 		cmd.PrintErrf("Failed to bring service down: %v\n", status.Convert(err).Message())
 	} else {
-		needsRestoreUp = !stateWasDown
 		cmd.Println("netbird down")
 	}
 
@@ -219,7 +217,6 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 	if _, err := client.Up(cmd.Context(), &proto.UpRequest{}); err != nil {
 		cmd.PrintErrf("Failed to bring service up: %v\n", status.Convert(err).Message())
 	} else {
-		needsRestoreUp = false
 		cmd.Println("netbird up")
 	}
 
@@ -267,14 +264,6 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to bundle debug: %v", status.Convert(err).Message())
 	}
 
-	if needsRestoreUp {
-		if _, err := client.Up(cmd.Context(), &proto.UpRequest{}); err != nil {
-			cmd.PrintErrf("Failed to restore service up state: %v\n", status.Convert(err).Message())
-		} else {
-			cmd.Println("netbird up (restored)")
-		}
-	}
-
 	if stateWasDown {
 		if _, err := client.Down(cmd.Context(), &proto.DownRequest{}); err != nil {
 			cmd.PrintErrf("Failed to restore service down state: %v\n", status.Convert(err).Message())

client/firewall/create_linux.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"strconv"
 
 	"github.com/coreos/go-iptables/iptables"
 	"github.com/google/nftables"
@@ -35,20 +36,27 @@ const SKIP_NFTABLES_ENV = "NB_SKIP_NFTABLES_CHECK"
 type FWType int
 
 func NewFirewall(iface IFaceMapper, stateManager *statemanager.Manager, flowLogger nftypes.FlowLogger, disableServerRoutes bool, mtu uint16) (firewall.Manager, error) {
-	// on the linux system we try to user nftables or iptables
-	// in any case, because we need to allow netbird interface traffic
-	// so we use AllowNetbird traffic from these firewall managers
-	// for the userspace packet filtering firewall
+	// We run in userspace mode and force userspace firewall was requested. We don't attempt native firewall.
+	if iface.IsUserspaceBind() && forceUserspaceFirewall() {
+		log.Info("forcing userspace firewall")
+		return createUserspaceFirewall(iface, nil, disableServerRoutes, flowLogger, mtu)
+	}
+
+	// Use native firewall for either kernel or userspace, the interface appears identical to netfilter
 	fm, err := createNativeFirewall(iface, stateManager, disableServerRoutes, mtu)
 
+	// Kernel cannot fall back to anything else, need to return error
 	if !iface.IsUserspaceBind() {
 		return fm, err
 	}
 
+	// Fall back to the userspace packet filter if native is unavailable
 	if err != nil {
 		log.Warnf("failed to create native firewall: %v. Proceeding with userspace", err)
+		return createUserspaceFirewall(iface, nil, disableServerRoutes, flowLogger, mtu)
 	}
-	return createUserspaceFirewall(iface, fm, disableServerRoutes, flowLogger, mtu)
+
+	return fm, nil
 }
 
 func createNativeFirewall(iface IFaceMapper, stateManager *statemanager.Manager, routes bool, mtu uint16) (firewall.Manager, error) {
@@ -160,3 +168,17 @@ func isIptablesClientAvailable(client *iptables.IPTables) bool {
 	_, err := client.ListChains("filter")
 	return err == nil
 }
+
+func forceUserspaceFirewall() bool {
+	val := os.Getenv(EnvForceUserspaceFirewall)
+	if val == "" {
+		return false
+	}
+
+	force, err := strconv.ParseBool(val)
+	if err != nil {
+		log.Warnf("failed to parse %s: %v", EnvForceUserspaceFirewall, err)
+		return false
+	}
+	return force
+}

client/firewall/iface.go

@@ -7,6 +7,12 @@ import (
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
 )
 
+// EnvForceUserspaceFirewall forces the use of the userspace packet filter even when
+// native iptables/nftables is available. This only applies when the WireGuard interface
+// runs in userspace mode. When set, peer ACLs are handled by USPFilter instead of
+// kernel netfilter rules.
+const EnvForceUserspaceFirewall = "NB_FORCE_USERSPACE_FIREWALL"
+
 // IFaceMapper defines subset methods of interface required for manager
 type IFaceMapper interface {
 	Name() string

client/firewall/iptables/manager_linux.go

@@ -33,7 +33,6 @@ type Manager struct {
 type iFaceMapper interface {
 	Name() string
 	Address() wgaddr.Address
-	IsUserspaceBind() bool
 }
 
 // Create iptables firewall manager
@@ -64,10 +63,9 @@ func Create(wgIface iFaceMapper, mtu uint16) (*Manager, error) {
 func (m *Manager) Init(stateManager *statemanager.Manager) error {
 	state := &ShutdownState{
 		InterfaceState: &InterfaceState{
-			NameStr:       m.wgIface.Name(),
-			WGAddress:     m.wgIface.Address(),
-			UserspaceBind: m.wgIface.IsUserspaceBind(),
-			MTU:           m.router.mtu,
+			NameStr:   m.wgIface.Name(),
+			WGAddress: m.wgIface.Address(),
+			MTU:       m.router.mtu,
 		},
 	}
 	stateManager.RegisterState(state)
@@ -203,12 +201,10 @@ func (m *Manager) Close(stateManager *statemanager.Manager) error {
 	return nberrors.FormatErrorOrNil(merr)
 }
 
-// AllowNetbird allows netbird interface traffic
+// AllowNetbird allows netbird interface traffic.
+// This is called when USPFilter wraps the native firewall, adding blanket accept
+// rules so that packet filtering is handled in userspace instead of by netfilter.
 func (m *Manager) AllowNetbird() error {
-	if !m.wgIface.IsUserspaceBind() {
-		return nil
-	}
-
 	_, err := m.AddPeerFiltering(
 		nil,
 		net.IP{0, 0, 0, 0},

client/firewall/iptables/manager_linux_test.go

@@ -47,8 +47,6 @@ func (i *iFaceMock) Address() wgaddr.Address {
 	panic("AddressFunc is not set")
 }
 
-func (i *iFaceMock) IsUserspaceBind() bool { return false }
-
 func TestIptablesManager(t *testing.T) {
 	ipv4Client, err := iptables.NewWithProtocol(iptables.ProtocolIPv4)
 	require.NoError(t, err)

client/firewall/iptables/state_linux.go

@@ -9,10 +9,9 @@ import (
 )
 
 type InterfaceState struct {
-	NameStr       string         `json:"name"`
-	WGAddress     wgaddr.Address `json:"wg_address"`
-	UserspaceBind bool           `json:"userspace_bind"`
-	MTU           uint16         `json:"mtu"`
+	NameStr   string         `json:"name"`
+	WGAddress wgaddr.Address `json:"wg_address"`
+	MTU       uint16         `json:"mtu"`
 }
 
 func (i *InterfaceState) Name() string {
@@ -23,10 +22,6 @@ func (i *InterfaceState) Address() wgaddr.Address {
 	return i.WGAddress
 }
 
-func (i *InterfaceState) IsUserspaceBind() bool {
-	return i.UserspaceBind
-}
-
 type ShutdownState struct {
 	sync.Mutex
 

client/firewall/nftables/manager_linux.go

@@ -40,7 +40,6 @@ func getTableName() string {
 type iFaceMapper interface {
 	Name() string
 	Address() wgaddr.Address
-	IsUserspaceBind() bool
 }
 
 // Manager of iptables firewall
@@ -106,10 +105,9 @@ func (m *Manager) Init(stateManager *statemanager.Manager) error {
 	// cleanup using Close() without needing to store specific rules.
 	if err := stateManager.UpdateState(&ShutdownState{
 		InterfaceState: &InterfaceState{
-			NameStr:       m.wgIface.Name(),
-			WGAddress:     m.wgIface.Address(),
-			UserspaceBind: m.wgIface.IsUserspaceBind(),
-			MTU:           m.router.mtu,
+			NameStr:   m.wgIface.Name(),
+			WGAddress: m.wgIface.Address(),
+			MTU:       m.router.mtu,
 		},
 	}); err != nil {
 		log.Errorf("failed to update state: %v", err)
@@ -205,12 +203,10 @@ func (m *Manager) RemoveNatRule(pair firewall.RouterPair) error {
 	return m.router.RemoveNatRule(pair)
 }
 
-// AllowNetbird allows netbird interface traffic
+// AllowNetbird allows netbird interface traffic.
+// This is called when USPFilter wraps the native firewall, adding blanket accept
+// rules so that packet filtering is handled in userspace instead of by netfilter.
 func (m *Manager) AllowNetbird() error {
-	if !m.wgIface.IsUserspaceBind() {
-		return nil
-	}
-
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 

client/firewall/nftables/manager_linux_test.go

@@ -52,8 +52,6 @@ func (i *iFaceMock) Address() wgaddr.Address {
 	panic("AddressFunc is not set")
 }
 
-func (i *iFaceMock) IsUserspaceBind() bool { return false }
-
 func TestNftablesManager(t *testing.T) {
 
 	// just check on the local interface

client/firewall/nftables/state_linux.go

@@ -8,10 +8,9 @@ import (
 )
 
 type InterfaceState struct {
-	NameStr       string         `json:"name"`
-	WGAddress     wgaddr.Address `json:"wg_address"`
-	UserspaceBind bool           `json:"userspace_bind"`
-	MTU           uint16         `json:"mtu"`
+	NameStr   string         `json:"name"`
+	WGAddress wgaddr.Address `json:"wg_address"`
+	MTU       uint16         `json:"mtu"`
 }
 
 func (i *InterfaceState) Name() string {
@@ -22,10 +21,6 @@ func (i *InterfaceState) Address() wgaddr.Address {
 	return i.WGAddress
 }
 
-func (i *InterfaceState) IsUserspaceBind() bool {
-	return i.UserspaceBind
-}
-
 type ShutdownState struct {
 	InterfaceState *InterfaceState `json:"interface_state,omitempty"`
 }

client/internal/acl/manager_test.go

@@ -19,6 +19,9 @@ import (
 var flowLogger = netflow.NewManager(nil, []byte{}, nil).GetLogger()
 
 func TestDefaultManager(t *testing.T) {
+	t.Setenv("NB_WG_KERNEL_DISABLED", "true")
+	t.Setenv(firewall.EnvForceUserspaceFirewall, "true")
+
 	networkMap := &mgmProto.NetworkMap{
 		FirewallRules: []*mgmProto.FirewallRule{
 			{
@@ -135,6 +138,7 @@ func TestDefaultManager(t *testing.T) {
 func TestDefaultManagerStateless(t *testing.T) {
 	// stateless currently only in userspace, so we have to disable kernel
 	t.Setenv("NB_WG_KERNEL_DISABLED", "true")
+	t.Setenv(firewall.EnvForceUserspaceFirewall, "true")
 	t.Setenv("NB_DISABLE_CONNTRACK", "true")
 
 	networkMap := &mgmProto.NetworkMap{
@@ -194,6 +198,7 @@ func TestDefaultManagerStateless(t *testing.T) {
 // This tests the full ACL manager -> uspfilter integration.
 func TestDenyRulesNotAccumulatedOnRepeatedApply(t *testing.T) {
 	t.Setenv("NB_WG_KERNEL_DISABLED", "true")
+	t.Setenv(firewall.EnvForceUserspaceFirewall, "true")
 
 	networkMap := &mgmProto.NetworkMap{
 		FirewallRules: []*mgmProto.FirewallRule{
@@ -258,6 +263,7 @@ func TestDenyRulesNotAccumulatedOnRepeatedApply(t *testing.T) {
 // up when they're removed from the network map in a subsequent update.
 func TestDenyRulesCleanedUpOnRemoval(t *testing.T) {
 	t.Setenv("NB_WG_KERNEL_DISABLED", "true")
+	t.Setenv(firewall.EnvForceUserspaceFirewall, "true")
 
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
@@ -339,6 +345,7 @@ func TestDenyRulesCleanedUpOnRemoval(t *testing.T) {
 // one added without leaking.
 func TestRuleUpdateChangingAction(t *testing.T) {
 	t.Setenv("NB_WG_KERNEL_DISABLED", "true")
+	t.Setenv(firewall.EnvForceUserspaceFirewall, "true")
 
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()

client/internal/connect.go

@@ -111,7 +111,6 @@ func (c *ConnectClient) RunOniOS(
 	fileDescriptor int32,
 	networkChangeListener listener.NetworkChangeListener,
 	dnsManager dns.IosDnsManager,
-	dnsAddresses []netip.AddrPort,
 	stateFilePath string,
 ) error {
 	// Set GC percent to 5% to reduce memory usage as iOS only allows 50MB of memory for the extension.
@@ -121,7 +120,6 @@ func (c *ConnectClient) RunOniOS(
 		FileDescriptor:        fileDescriptor,
 		NetworkChangeListener: networkChangeListener,
 		DnsManager:            dnsManager,
-		HostDNSAddresses:      dnsAddresses,
 		StateFilePath:         stateFilePath,
 	}
 	return c.run(mobileDependency, nil, "")

client/internal/dns/local/local_test.go

@@ -1263,9 +1263,9 @@ func TestLocalResolver_AuthoritativeFlag(t *testing.T) {
 	})
 }
 
-// TestLocalResolver_Stop tests cleanup on GracefullyStop
+// TestLocalResolver_Stop tests cleanup on Stop
 func TestLocalResolver_Stop(t *testing.T) {
-	t.Run("GracefullyStop clears all state", func(t *testing.T) {
+	t.Run("Stop clears all state", func(t *testing.T) {
 		resolver := NewResolver()
 		resolver.Update([]nbdns.CustomZone{{
 			Domain: "example.com.",
@@ -1285,7 +1285,7 @@ func TestLocalResolver_Stop(t *testing.T) {
 		assert.False(t, resolver.isInManagedZone("host.example.com."))
 	})
 
-	t.Run("GracefullyStop is safe to call multiple times", func(t *testing.T) {
+	t.Run("Stop is safe to call multiple times", func(t *testing.T) {
 		resolver := NewResolver()
 		resolver.Update([]nbdns.CustomZone{{
 			Domain: "example.com.",
@@ -1299,7 +1299,7 @@ func TestLocalResolver_Stop(t *testing.T) {
 		resolver.Stop()
 	})
 
-	t.Run("GracefullyStop cancels in-flight external resolution", func(t *testing.T) {
+	t.Run("Stop cancels in-flight external resolution", func(t *testing.T) {
 		resolver := NewResolver()
 
 		lookupStarted := make(chan struct{})

client/internal/dns/server.go

@@ -187,16 +187,11 @@ func NewDefaultServerIos(
 	ctx context.Context,
 	wgInterface WGIface,
 	iosDnsManager IosDnsManager,
-	hostsDnsList []netip.AddrPort,
 	statusRecorder *peer.Status,
 	disableSys bool,
 ) *DefaultServer {
-	log.Debugf("iOS host dns address list is: %v", hostsDnsList)
 	ds := newDefaultServer(ctx, wgInterface, NewServiceViaMemory(wgInterface), statusRecorder, nil, disableSys)
 	ds.iosDnsManager = iosDnsManager
-	ds.hostsDNSHolder.set(hostsDnsList)
-	ds.permanent = true
-	ds.addHostRootZone()
 	return ds
 }
 

client/internal/engine.go

@@ -46,7 +46,6 @@ import (
 	"github.com/netbirdio/netbird/client/internal/peer/guard"
 	icemaker "github.com/netbirdio/netbird/client/internal/peer/ice"
 	"github.com/netbirdio/netbird/client/internal/peerstore"
-	"github.com/netbirdio/netbird/client/internal/portforward"
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
 	"github.com/netbirdio/netbird/client/internal/relay"
 	"github.com/netbirdio/netbird/client/internal/rosenpass"
@@ -211,10 +210,9 @@ type Engine struct {
 	// checks are the client-applied posture checks that need to be evaluated on the client
 	checks []*mgmProto.Checks
 
-	relayManager       *relayClient.Manager
-	stateManager       *statemanager.Manager
-	portForwardManager *portforward.Manager
-	srWatcher          *guard.SRWatcher
+	relayManager *relayClient.Manager
+	stateManager *statemanager.Manager
+	srWatcher    *guard.SRWatcher
 
 	// Sync response persistence (protected by syncRespMux)
 	syncRespMux         sync.RWMutex
@@ -261,27 +259,26 @@ func NewEngine(
 	mobileDep MobileDependency,
 ) *Engine {
 	engine := &Engine{
-		clientCtx:          clientCtx,
-		clientCancel:       clientCancel,
-		signal:             services.SignalClient,
-		signaler:           peer.NewSignaler(services.SignalClient, config.WgPrivateKey),
-		mgmClient:          services.MgmClient,
-		relayManager:       services.RelayManager,
-		peerStore:          peerstore.NewConnStore(),
-		syncMsgMux:         &sync.Mutex{},
-		config:             config,
-		mobileDep:          mobileDep,
-		STUNs:              []*stun.URI{},
-		TURNs:              []*stun.URI{},
-		networkSerial:      0,
-		statusRecorder:     services.StatusRecorder,
-		stateManager:       services.StateManager,
-		portForwardManager: portforward.NewManager(),
-		checks:             services.Checks,
-		probeStunTurn:      relay.NewStunTurnProbe(relay.DefaultCacheTTL),
-		jobExecutor:        jobexec.NewExecutor(),
-		clientMetrics:      services.ClientMetrics,
-		updateManager:      services.UpdateManager,
+		clientCtx:      clientCtx,
+		clientCancel:   clientCancel,
+		signal:         services.SignalClient,
+		signaler:       peer.NewSignaler(services.SignalClient, config.WgPrivateKey),
+		mgmClient:      services.MgmClient,
+		relayManager:   services.RelayManager,
+		peerStore:      peerstore.NewConnStore(),
+		syncMsgMux:     &sync.Mutex{},
+		config:         config,
+		mobileDep:      mobileDep,
+		STUNs:          []*stun.URI{},
+		TURNs:          []*stun.URI{},
+		networkSerial:  0,
+		statusRecorder: services.StatusRecorder,
+		stateManager:   services.StateManager,
+		checks:         services.Checks,
+		probeStunTurn:  relay.NewStunTurnProbe(relay.DefaultCacheTTL),
+		jobExecutor:    jobexec.NewExecutor(),
+		clientMetrics:  services.ClientMetrics,
+		updateManager:  services.UpdateManager,
 	}
 
 	log.Infof("I am: %s", config.WgPrivateKey.PublicKey().String())
@@ -540,13 +537,6 @@ func (e *Engine) Start(netbirdConfig *mgmProto.NetbirdConfig, mgmtURL *url.URL)
 	// conntrack entries from being created before the rules are in place
 	e.setupWGProxyNoTrack()
 
-	// Start after interface is up since port may have been resolved from 0 or changed if occupied
-	e.shutdownWg.Add(1)
-	go func() {
-		defer e.shutdownWg.Done()
-		e.portForwardManager.Start(e.ctx, uint16(e.config.WgPort))
-	}()
-
 	// Set the WireGuard interface for rosenpass after interface is up
 	if e.rpManager != nil {
 		e.rpManager.SetInterface(e.wgInterface)
@@ -1550,13 +1540,12 @@ func (e *Engine) createPeerConn(pubKey string, allowedIPs []netip.Prefix, agentV
 	}
 
 	serviceDependencies := peer.ServiceDependencies{
-		StatusRecorder:     e.statusRecorder,
-		Signaler:           e.signaler,
-		IFaceDiscover:      e.mobileDep.IFaceDiscover,
-		RelayManager:       e.relayManager,
-		SrWatcher:          e.srWatcher,
-		PortForwardManager: e.portForwardManager,
-		MetricsRecorder:    e.clientMetrics,
+		StatusRecorder:  e.statusRecorder,
+		Signaler:        e.signaler,
+		IFaceDiscover:   e.mobileDep.IFaceDiscover,
+		RelayManager:    e.relayManager,
+		SrWatcher:       e.srWatcher,
+		MetricsRecorder: e.clientMetrics,
 	}
 	peerConn, err := peer.NewConn(config, serviceDependencies)
 	if err != nil {
@@ -1713,12 +1702,6 @@ func (e *Engine) close() {
 	if e.rpManager != nil {
 		_ = e.rpManager.Close()
 	}
-
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
-	if err := e.portForwardManager.GracefullyStop(ctx); err != nil {
-		log.Warnf("failed to gracefully stop port forwarding manager: %s", err)
-	}
 }
 
 func (e *Engine) readInitialSettings() ([]*route.Route, *nbdns.Config, bool, error) {
@@ -1822,7 +1805,7 @@ func (e *Engine) newDnsServer(dnsConfig *nbdns.Config) (dns.Server, error) {
 		return dnsServer, nil
 
 	case "ios":
-		dnsServer := dns.NewDefaultServerIos(e.ctx, e.wgInterface, e.mobileDep.DnsManager, e.mobileDep.HostDNSAddresses, e.statusRecorder, e.config.DisableDNS)
+		dnsServer := dns.NewDefaultServerIos(e.ctx, e.wgInterface, e.mobileDep.DnsManager, e.statusRecorder, e.config.DisableDNS)
 		return dnsServer, nil
 
 	default:

client/internal/peer/conn.go

@@ -22,7 +22,6 @@ import (
 	icemaker "github.com/netbirdio/netbird/client/internal/peer/ice"
 	"github.com/netbirdio/netbird/client/internal/peer/id"
 	"github.com/netbirdio/netbird/client/internal/peer/worker"
-	"github.com/netbirdio/netbird/client/internal/portforward"
 	"github.com/netbirdio/netbird/client/internal/stdnet"
 	"github.com/netbirdio/netbird/route"
 	relayClient "github.com/netbirdio/netbird/shared/relay/client"
@@ -46,7 +45,6 @@ type ServiceDependencies struct {
 	RelayManager       *relayClient.Manager
 	SrWatcher          *guard.SRWatcher
 	PeerConnDispatcher *dispatcher.ConnectionDispatcher
-	PortForwardManager *portforward.Manager
 	MetricsRecorder    MetricsRecorder
 }
 
@@ -89,17 +87,16 @@ type ConnConfig struct {
 }
 
 type Conn struct {
-	Log                *log.Entry
-	mu                 sync.Mutex
-	ctx                context.Context
-	ctxCancel          context.CancelFunc
-	config             ConnConfig
-	statusRecorder     *Status
-	signaler           *Signaler
-	iFaceDiscover      stdnet.ExternalIFaceDiscover
-	relayManager       *relayClient.Manager
-	srWatcher          *guard.SRWatcher
-	portForwardManager *portforward.Manager
+	Log            *log.Entry
+	mu             sync.Mutex
+	ctx            context.Context
+	ctxCancel      context.CancelFunc
+	config         ConnConfig
+	statusRecorder *Status
+	signaler       *Signaler
+	iFaceDiscover  stdnet.ExternalIFaceDiscover
+	relayManager   *relayClient.Manager
+	srWatcher      *guard.SRWatcher
 
 	onConnected                               func(remoteWireGuardKey string, remoteRosenpassPubKey []byte, wireGuardIP string, remoteRosenpassAddr string)
 	onDisconnected                            func(remotePeer string)
@@ -148,20 +145,19 @@ func NewConn(config ConnConfig, services ServiceDependencies) (*Conn, error) {
 
 	dumpState := newStateDump(config.Key, connLog, services.StatusRecorder)
 	var conn = &Conn{
-		Log:                connLog,
-		config:             config,
-		statusRecorder:     services.StatusRecorder,
-		signaler:           services.Signaler,
-		iFaceDiscover:      services.IFaceDiscover,
-		relayManager:       services.RelayManager,
-		srWatcher:          services.SrWatcher,
-		portForwardManager: services.PortForwardManager,
-		statusRelay:        worker.NewAtomicStatus(),
-		statusICE:          worker.NewAtomicStatus(),
-		dumpState:          dumpState,
-		endpointUpdater:    NewEndpointUpdater(connLog, config.WgConfig, isController(config)),
-		wgWatcher:          NewWGWatcher(connLog, config.WgConfig.WgInterface, config.Key, dumpState),
-		metricsRecorder:    services.MetricsRecorder,
+		Log:             connLog,
+		config:          config,
+		statusRecorder:  services.StatusRecorder,
+		signaler:        services.Signaler,
+		iFaceDiscover:   services.IFaceDiscover,
+		relayManager:    services.RelayManager,
+		srWatcher:       services.SrWatcher,
+		statusRelay:     worker.NewAtomicStatus(),
+		statusICE:       worker.NewAtomicStatus(),
+		dumpState:       dumpState,
+		endpointUpdater: NewEndpointUpdater(connLog, config.WgConfig, isController(config)),
+		wgWatcher:       NewWGWatcher(connLog, config.WgConfig.WgInterface, config.Key, dumpState),
+		metricsRecorder: services.MetricsRecorder,
 	}
 
 	return conn, nil

client/internal/peer/worker_ice.go

@@ -16,7 +16,6 @@ import (
 	"github.com/netbirdio/netbird/client/iface/udpmux"
 	"github.com/netbirdio/netbird/client/internal/peer/conntype"
 	icemaker "github.com/netbirdio/netbird/client/internal/peer/ice"
-	"github.com/netbirdio/netbird/client/internal/portforward"
 	"github.com/netbirdio/netbird/client/internal/stdnet"
 	"github.com/netbirdio/netbird/route"
 )
@@ -62,9 +61,6 @@ type WorkerICE struct {
 
 	// we record the last known state of the ICE agent to avoid duplicate on disconnected events
 	lastKnownState ice.ConnectionState
-
-	// portForwardAttempted tracks if we've already tried port forwarding this session
-	portForwardAttempted bool
 }
 
 func NewWorkerICE(ctx context.Context, log *log.Entry, config ConnConfig, conn *Conn, signaler *Signaler, ifaceDiscover stdnet.ExternalIFaceDiscover, statusRecorder *Status, hasRelayOnLocally bool) (*WorkerICE, error) {
@@ -218,8 +214,6 @@ func (w *WorkerICE) Close() {
 }
 
 func (w *WorkerICE) reCreateAgent(dialerCancel context.CancelFunc, candidates []ice.CandidateType) (*icemaker.ThreadSafeAgent, error) {
-	w.portForwardAttempted = false
-
 	agent, err := icemaker.NewAgent(w.ctx, w.iFaceDiscover, w.config.ICEConfig, candidates, w.localUfrag, w.localPwd)
 	if err != nil {
 		return nil, fmt.Errorf("create agent: %w", err)
@@ -376,93 +370,6 @@ func (w *WorkerICE) onICECandidate(candidate ice.Candidate) {
 			w.log.Errorf("failed signaling candidate to the remote peer %s %s", w.config.Key, err)
 		}
 	}()
-
-	if candidate.Type() == ice.CandidateTypeServerReflexive {
-		w.injectPortForwardedCandidate(candidate)
-	}
-}
-
-// injectPortForwardedCandidate signals an additional candidate using the pre-created port mapping.
-func (w *WorkerICE) injectPortForwardedCandidate(srflxCandidate ice.Candidate) {
-	pfManager := w.conn.portForwardManager
-	if pfManager == nil {
-		return
-	}
-
-	mapping := pfManager.GetMapping()
-	if mapping == nil {
-		return
-	}
-
-	w.muxAgent.Lock()
-	if w.portForwardAttempted {
-		w.muxAgent.Unlock()
-		return
-	}
-	w.portForwardAttempted = true
-	w.muxAgent.Unlock()
-
-	forwardedCandidate, err := w.createForwardedCandidate(srflxCandidate, mapping)
-	if err != nil {
-		w.log.Warnf("create forwarded candidate: %v", err)
-		return
-	}
-
-	w.log.Debugf("injecting port-forwarded candidate: %s (mapping: %d -> %d via %s, priority: %d)",
-		forwardedCandidate.String(), mapping.InternalPort, mapping.ExternalPort, mapping.NATType, forwardedCandidate.Priority())
-
-	go func() {
-		if err := w.signaler.SignalICECandidate(forwardedCandidate, w.config.Key); err != nil {
-			w.log.Errorf("signal port-forwarded candidate: %v", err)
-		}
-	}()
-}
-
-// createForwardedCandidate creates a new server reflexive candidate with the forwarded port.
-// It uses the NAT gateway's external IP with the forwarded port.
-func (w *WorkerICE) createForwardedCandidate(srflxCandidate ice.Candidate, mapping *portforward.Mapping) (ice.Candidate, error) {
-	var externalIP string
-	if mapping.ExternalIP != nil && !mapping.ExternalIP.IsUnspecified() {
-		externalIP = mapping.ExternalIP.String()
-	} else {
-		// Fallback to STUN-discovered address if NAT didn't provide external IP
-		externalIP = srflxCandidate.Address()
-	}
-
-	// Per RFC 8445, the related address for srflx is the base (host candidate address).
-	// If the original srflx has unspecified related address, use its own address as base.
-	relAddr := srflxCandidate.RelatedAddress().Address
-	if relAddr == "" || relAddr == "0.0.0.0" || relAddr == "::" {
-		relAddr = srflxCandidate.Address()
-	}
-
-	// Arbitrary +1000 boost on top of RFC 8445 priority to favor port-forwarded candidates
-	// over regular srflx during ICE connectivity checks.
-	priority := srflxCandidate.Priority() + 1000
-
-	candidate, err := ice.NewCandidateServerReflexive(&ice.CandidateServerReflexiveConfig{
-		Network:   srflxCandidate.NetworkType().String(),
-		Address:   externalIP,
-		Port:      int(mapping.ExternalPort),
-		Component: srflxCandidate.Component(),
-		Priority:  priority,
-		RelAddr:   relAddr,
-		RelPort:   int(mapping.InternalPort),
-	})
-	if err != nil {
-		return nil, fmt.Errorf("create candidate: %w", err)
-	}
-
-	for _, e := range srflxCandidate.Extensions() {
-		if e.Key == ice.ExtensionKeyCandidateID {
-			e.Value = srflxCandidate.ID()
-		}
-		if err := candidate.AddExtension(e); err != nil {
-			return nil, fmt.Errorf("add extension: %w", err)
-		}
-	}
-
-	return candidate, nil
 }
 
 func (w *WorkerICE) onICESelectedCandidatePair(agent *icemaker.ThreadSafeAgent, c1, c2 ice.Candidate) {
@@ -504,10 +411,10 @@ func (w *WorkerICE) logSuccessfulPaths(agent *icemaker.ThreadSafeAgent) {
 			if !lok || !rok {
 				continue
 			}
-			w.log.Debugf("successful ICE path %s: [%s %s %s:%d] <-> [%s %s %s:%d] rtt=%.3fms",
+			w.log.Debugf("successful ICE path %s: [%s %s %s] <-> [%s %s %s] rtt=%.3fms",
 				sessionID,
-				local.NetworkType(), local.Type(), local.Address(), local.Port(),
-				remote.NetworkType(), remote.Type(), remote.Address(), remote.Port(),
+				local.NetworkType(), local.Type(), local.Address(),
+				remote.NetworkType(), remote.Type(), remote.Address(),
 				stat.CurrentRoundTripTime*1000)
 		}
 	}

client/internal/portforward/env.go

@@ -1,26 +0,0 @@
-package portforward
-
-import (
-	"os"
-	"strconv"
-
-	log "github.com/sirupsen/logrus"
-)
-
-const (
-	envDisableNATMapper = "NB_DISABLE_NAT_MAPPER"
-)
-
-func isDisabledByEnv() bool {
-	val := os.Getenv(envDisableNATMapper)
-	if val == "" {
-		return false
-	}
-
-	disabled, err := strconv.ParseBool(val)
-	if err != nil {
-		log.Warnf("failed to parse %s: %v", envDisableNATMapper, err)
-		return false
-	}
-	return disabled
-}

client/internal/portforward/manager.go

@@ -1,250 +0,0 @@
-//go:build !js
-
-package portforward
-
-import (
-	"context"
-	"fmt"
-	"net"
-	"sync"
-	"time"
-
-	"github.com/libp2p/go-nat"
-	log "github.com/sirupsen/logrus"
-)
-
-const (
-	defaultMappingTTL  = 2 * time.Hour
-	renewalInterval    = defaultMappingTTL / 2
-	discoveryTimeout   = 10 * time.Second
-	mappingDescription = "NetBird"
-)
-
-type Mapping struct {
-	Protocol     string
-	InternalPort uint16
-	ExternalPort uint16
-	ExternalIP   net.IP
-	NATType      string
-}
-
-type Manager struct {
-	cancel context.CancelFunc
-
-	mapping     *Mapping
-	mappingLock sync.Mutex
-
-	wgPort uint16
-
-	done    chan struct{}
-	stopCtx chan context.Context
-
-	// protect exported functions
-	mu sync.Mutex
-}
-
-func NewManager() *Manager {
-	return &Manager{
-		stopCtx: make(chan context.Context, 1),
-	}
-}
-
-func (m *Manager) Start(ctx context.Context, wgPort uint16) {
-	m.mu.Lock()
-	if m.cancel != nil {
-		m.mu.Unlock()
-		return
-	}
-
-	if isDisabledByEnv() {
-		log.Infof("NAT port mapper disabled via %s", envDisableNATMapper)
-		m.mu.Unlock()
-		return
-	}
-
-	if wgPort == 0 {
-		log.Warnf("invalid WireGuard port 0; NAT mapping disabled")
-		m.mu.Unlock()
-		return
-	}
-	m.wgPort = wgPort
-
-	m.done = make(chan struct{})
-	defer close(m.done)
-
-	ctx, m.cancel = context.WithCancel(ctx)
-	m.mu.Unlock()
-
-	gateway, mapping, err := m.setup(ctx)
-	if err != nil {
-		log.Errorf("failed to setup NAT port mapping: %v", err)
-
-		return
-	}
-
-	m.mappingLock.Lock()
-	m.mapping = mapping
-	m.mappingLock.Unlock()
-
-	m.renewLoop(ctx, gateway)
-
-	select {
-	case cleanupCtx := <-m.stopCtx:
-		// block the Start while cleaned up gracefully
-		m.cleanup(cleanupCtx, gateway)
-	default:
-		// return Start immediately and cleanup in background
-		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
-		go func() {
-			defer cleanupCancel()
-			m.cleanup(cleanupCtx, gateway)
-		}()
-	}
-}
-
-// GetMapping returns the current mapping if ready, nil otherwise
-func (m *Manager) GetMapping() *Mapping {
-	m.mappingLock.Lock()
-	defer m.mappingLock.Unlock()
-
-	if m.mapping == nil {
-		return nil
-	}
-
-	mapping := *m.mapping
-	return &mapping
-}
-
-// GracefullyStop cancels the manager and attempts to delete the port mapping.
-// After GracefullyStop returns, the manager cannot be restarted.
-func (m *Manager) GracefullyStop(ctx context.Context) error {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	if m.cancel == nil {
-		return nil
-	}
-
-	// Send cleanup context before cancelling, so Start picks it up after renewLoop exits.
-	m.startTearDown(ctx)
-
-	m.cancel()
-	m.cancel = nil
-
-	select {
-	case <-ctx.Done():
-		return ctx.Err()
-	case <-m.done:
-		return nil
-	}
-}
-
-func (m *Manager) setup(ctx context.Context) (nat.NAT, *Mapping, error) {
-	discoverCtx, discoverCancel := context.WithTimeout(ctx, discoveryTimeout)
-	defer discoverCancel()
-
-	gateway, err := nat.DiscoverGateway(discoverCtx)
-	if err != nil {
-		log.Infof("NAT gateway discovery failed: %v (port forwarding disabled)", err)
-		return nil, nil, err
-	}
-
-	log.Infof("discovered NAT gateway: %s", gateway.Type())
-
-	mapping, err := m.createMapping(ctx, gateway)
-	if err != nil {
-		log.Warnf("failed to create port mapping: %v", err)
-		return nil, nil, err
-	}
-	return gateway, mapping, nil
-}
-
-func (m *Manager) createMapping(ctx context.Context, gateway nat.NAT) (*Mapping, error) {
-	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
-	defer cancel()
-
-	externalPort, err := gateway.AddPortMapping(ctx, "udp", int(m.wgPort), mappingDescription, defaultMappingTTL)
-	if err != nil {
-		return nil, err
-	}
-
-	externalIP, err := gateway.GetExternalAddress()
-	if err != nil {
-		log.Debugf("failed to get external address: %v", err)
-		// todo return with err?
-	}
-
-	mapping := &Mapping{
-		Protocol:     "udp",
-		InternalPort: m.wgPort,
-		ExternalPort: uint16(externalPort),
-		ExternalIP:   externalIP,
-		NATType:      gateway.Type(),
-	}
-
-	log.Infof("created port mapping: %d -> %d via %s (external IP: %s)",
-		m.wgPort, externalPort, gateway.Type(), externalIP)
-	return mapping, nil
-}
-
-func (m *Manager) renewLoop(ctx context.Context, gateway nat.NAT) {
-	ticker := time.NewTicker(renewalInterval)
-	defer ticker.Stop()
-
-	for {
-		select {
-		case <-ctx.Done():
-			return
-		case <-ticker.C:
-			if err := m.renewMapping(ctx, gateway); err != nil {
-				log.Warnf("failed to renew port mapping: %v", err)
-				continue
-			}
-		}
-	}
-}
-
-func (m *Manager) renewMapping(ctx context.Context, gateway nat.NAT) error {
-	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
-	defer cancel()
-
-	externalPort, err := gateway.AddPortMapping(ctx, m.mapping.Protocol, int(m.mapping.InternalPort), mappingDescription, defaultMappingTTL)
-	if err != nil {
-		return fmt.Errorf("add port mapping: %w", err)
-	}
-
-	if uint16(externalPort) != m.mapping.ExternalPort {
-		log.Warnf("external port changed on renewal: %d -> %d (candidate may be stale)", m.mapping.ExternalPort, externalPort)
-		m.mappingLock.Lock()
-		m.mapping.ExternalPort = uint16(externalPort)
-		m.mappingLock.Unlock()
-	}
-
-	log.Debugf("renewed port mapping: %d -> %d", m.mapping.InternalPort, m.mapping.ExternalPort)
-	return nil
-}
-
-func (m *Manager) cleanup(ctx context.Context, gateway nat.NAT) {
-	m.mappingLock.Lock()
-	mapping := m.mapping
-	m.mapping = nil
-	m.mappingLock.Unlock()
-
-	if mapping == nil {
-		return
-	}
-
-	if err := gateway.DeletePortMapping(ctx, mapping.Protocol, int(mapping.InternalPort)); err != nil {
-		log.Warnf("delete port mapping on stop: %v", err)
-		return
-	}
-
-	log.Infof("deleted port mapping for port %d", mapping.InternalPort)
-}
-
-func (m *Manager) startTearDown(ctx context.Context) {
-	select {
-	case m.stopCtx <- ctx:
-	default:
-	}
-}

client/internal/portforward/manager_js.go

@@ -1,36 +0,0 @@
-package portforward
-
-import (
-	"context"
-	"net"
-)
-
-// Mapping represents port mapping information.
-type Mapping struct {
-	Protocol     string
-	InternalPort uint16
-	ExternalPort uint16
-	ExternalIP   net.IP
-	NATType      string
-}
-
-// Manager is a stub for js/wasm builds where NAT-PMP/UPnP is not supported.
-type Manager struct{}
-
-// NewManager returns a stub manager for js/wasm builds.
-func NewManager() *Manager {
-	return &Manager{}
-}
-
-// Start is a no-op on js/wasm: NAT-PMP/UPnP is not available in browser environments.
-func (m *Manager) Start(context.Context, uint16) {
-	// no NAT traversal in wasm
-}
-
-// GracefullyStop is a no-op on js/wasm.
-func (m *Manager) GracefullyStop(context.Context) error { return nil }
-
-// GetMapping always returns nil on js/wasm.
-func (m *Manager) GetMapping() *Mapping {
-	return nil
-}

client/internal/portforward/manager_test.go

@@ -1,159 +0,0 @@
-//go:build !js
-
-package portforward
-
-import (
-	"context"
-	"net"
-	"testing"
-	"time"
-
-	"github.com/libp2p/go-nat"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-)
-
-type mockNAT struct {
-	natType          string
-	deviceAddr       net.IP
-	externalAddr     net.IP
-	internalAddr     net.IP
-	mappings         map[int]int
-	addMappingErr    error
-	deleteMappingErr error
-}
-
-func newMockNAT() *mockNAT {
-	return &mockNAT{
-		natType:      "Mock-NAT",
-		deviceAddr:   net.ParseIP("192.168.1.1"),
-		externalAddr: net.ParseIP("203.0.113.50"),
-		internalAddr: net.ParseIP("192.168.1.100"),
-		mappings:     make(map[int]int),
-	}
-}
-
-func (m *mockNAT) Type() string {
-	return m.natType
-}
-
-func (m *mockNAT) GetDeviceAddress() (net.IP, error) {
-	return m.deviceAddr, nil
-}
-
-func (m *mockNAT) GetExternalAddress() (net.IP, error) {
-	return m.externalAddr, nil
-}
-
-func (m *mockNAT) GetInternalAddress() (net.IP, error) {
-	return m.internalAddr, nil
-}
-
-func (m *mockNAT) AddPortMapping(ctx context.Context, protocol string, internalPort int, description string, timeout time.Duration) (int, error) {
-	if m.addMappingErr != nil {
-		return 0, m.addMappingErr
-	}
-	externalPort := internalPort
-	m.mappings[internalPort] = externalPort
-	return externalPort, nil
-}
-
-func (m *mockNAT) DeletePortMapping(ctx context.Context, protocol string, internalPort int) error {
-	if m.deleteMappingErr != nil {
-		return m.deleteMappingErr
-	}
-	delete(m.mappings, internalPort)
-	return nil
-}
-
-func TestManager_CreateMapping(t *testing.T) {
-	m := NewManager()
-	m.wgPort = 51820
-
-	gateway := newMockNAT()
-	mapping, err := m.createMapping(context.Background(), gateway)
-	require.NoError(t, err)
-	require.NotNil(t, mapping)
-
-	assert.Equal(t, "udp", mapping.Protocol)
-	assert.Equal(t, uint16(51820), mapping.InternalPort)
-	assert.Equal(t, uint16(51820), mapping.ExternalPort)
-	assert.Equal(t, "Mock-NAT", mapping.NATType)
-	assert.Equal(t, net.ParseIP("203.0.113.50").To4(), mapping.ExternalIP.To4())
-}
-
-func TestManager_GetMapping_ReturnsNilWhenNotReady(t *testing.T) {
-	m := NewManager()
-	assert.Nil(t, m.GetMapping())
-}
-
-func TestManager_GetMapping_ReturnsCopy(t *testing.T) {
-	m := NewManager()
-	m.mapping = &Mapping{
-		Protocol:     "udp",
-		InternalPort: 51820,
-		ExternalPort: 51820,
-	}
-
-	mapping := m.GetMapping()
-	require.NotNil(t, mapping)
-	assert.Equal(t, uint16(51820), mapping.InternalPort)
-
-	// Mutating the returned copy should not affect the manager's mapping.
-	mapping.ExternalPort = 9999
-	assert.Equal(t, uint16(51820), m.GetMapping().ExternalPort)
-}
-
-func TestManager_Cleanup_DeletesMapping(t *testing.T) {
-	m := NewManager()
-	m.mapping = &Mapping{
-		Protocol:     "udp",
-		InternalPort: 51820,
-		ExternalPort: 51820,
-	}
-
-	gateway := newMockNAT()
-	// Seed the mock so we can verify deletion.
-	gateway.mappings[51820] = 51820
-
-	m.cleanup(context.Background(), gateway)
-
-	_, exists := gateway.mappings[51820]
-	assert.False(t, exists, "mapping should be deleted from gateway")
-	assert.Nil(t, m.GetMapping(), "in-memory mapping should be cleared")
-}
-
-func TestManager_Cleanup_NilMapping(t *testing.T) {
-	m := NewManager()
-	gateway := newMockNAT()
-
-	// Should not panic or call gateway.
-	m.cleanup(context.Background(), gateway)
-}
-
-func TestState_Cleanup(t *testing.T) {
-	origDiscover := discoverGateway
-	defer func() { discoverGateway = origDiscover }()
-
-	mockGateway := newMockNAT()
-	mockGateway.mappings[51820] = 51820
-	discoverGateway = func(ctx context.Context) (nat.NAT, error) {
-		return mockGateway, nil
-	}
-
-	state := &State{
-		Protocol:     "udp",
-		InternalPort: 51820,
-	}
-
-	err := state.Cleanup()
-	assert.NoError(t, err)
-
-	_, exists := mockGateway.mappings[51820]
-	assert.False(t, exists, "mapping should be deleted after cleanup")
-}
-
-func TestState_Name(t *testing.T) {
-	state := &State{}
-	assert.Equal(t, "port_forward_state", state.Name())
-}

client/internal/portforward/state.go

@@ -1,50 +0,0 @@
-//go:build !js
-
-package portforward
-
-import (
-	"context"
-	"fmt"
-
-	"github.com/libp2p/go-nat"
-	log "github.com/sirupsen/logrus"
-)
-
-// discoverGateway is the function used for NAT gateway discovery.
-// It can be replaced in tests to avoid real network operations.
-var discoverGateway = nat.DiscoverGateway
-
-// State is persisted only for crash recovery cleanup
-type State struct {
-	InternalPort uint16 `json:"internal_port,omitempty"`
-	Protocol     string `json:"protocol,omitempty"`
-}
-
-func (s *State) Name() string {
-	return "port_forward_state"
-}
-
-// Cleanup implements statemanager.CleanableState for crash recovery
-func (s *State) Cleanup() error {
-	if s.InternalPort == 0 {
-		return nil
-	}
-
-	log.Infof("cleaning up stale port mapping for port %d", s.InternalPort)
-
-	ctx, cancel := context.WithTimeout(context.Background(), discoveryTimeout)
-	defer cancel()
-
-	gateway, err := discoverGateway(ctx)
-	if err != nil {
-		// Discovery failure is not an error - gateway may not exist
-		log.Debugf("cleanup: no gateway found: %v", err)
-		return nil
-	}
-
-	if err := gateway.DeletePortMapping(ctx, s.Protocol, int(s.InternalPort)); err != nil {
-		return fmt.Errorf("delete port mapping: %w", err)
-	}
-
-	return nil
-}

client/internal/routemanager/notifier/notifier_ios.go

@@ -53,6 +53,7 @@ func (n *Notifier) OnNewPrefixes(prefixes []netip.Prefix) {
 	n.currentPrefixes = newNets
 	n.notify()
 }
+
 func (n *Notifier) notify() {
 	n.listenerMux.Lock()
 	defer n.listenerMux.Unlock()

client/ios/NetBirdSDK/client.go

@@ -161,11 +161,7 @@ func (c *Client) Run(fd int32, interfaceName string, envList *EnvList) error {
 	cfg.WgIface = interfaceName
 
 	c.connectClient = internal.NewConnectClient(ctx, cfg, c.recorder)
-	hostDNS := []netip.AddrPort{
-		netip.MustParseAddrPort("9.9.9.9:53"),
-		netip.MustParseAddrPort("149.112.112.112:53"),
-	}
-	return c.connectClient.RunOniOS(fd, c.networkChangeListener, c.dnsManager, hostDNS, c.stateFile)
+	return c.connectClient.RunOniOS(fd, c.networkChangeListener, c.dnsManager, c.stateFile)
 }
 
 // Stop the internal client and free the resources

client/server/state_generic.go

@@ -9,11 +9,6 @@ import (
 	"github.com/netbirdio/netbird/client/ssh/config"
 )
 
-// registerStates registers all states that need crash recovery cleanup.
-// Note: portforward.State is intentionally NOT registered here to avoid blocking startup
-// for up to 10 seconds during NAT gateway discovery when no gateway is present.
-// The gateway reference cannot be persisted across restarts, so cleanup requires re-discovery.
-// Port forward cleanup is handled by the Manager during normal operation instead.
 func registerStates(mgr *statemanager.Manager) {
 	mgr.RegisterState(&dns.ShutdownState{})
 	mgr.RegisterState(&systemops.ShutdownState{})

client/server/state_linux.go

@@ -11,11 +11,6 @@ import (
 	"github.com/netbirdio/netbird/client/ssh/config"
 )
 
-// registerStates registers all states that need crash recovery cleanup.
-// Note: portforward.State is intentionally NOT registered here to avoid blocking startup
-// for up to 10 seconds during NAT gateway discovery when no gateway is present.
-// The gateway reference cannot be persisted across restarts, so cleanup requires re-discovery.
-// Port forward cleanup is handled by the Manager during normal operation instead.
 func registerStates(mgr *statemanager.Manager) {
 	mgr.RegisterState(&dns.ShutdownState{})
 	mgr.RegisterState(&systemops.ShutdownState{})

client/ssh/proxy/proxy.go

@@ -141,7 +141,7 @@ func (p *SSHProxy) runProxySSHServer(jwtToken string) error {
 
 func (p *SSHProxy) handleSSHSession(session ssh.Session) {
 	ptyReq, winCh, isPty := session.Pty()
-	hasCommand := session.RawCommand() != ""
+	hasCommand := len(session.Command()) > 0
 
 	sshClient, err := p.getOrCreateBackendClient(session.Context(), session.User())
 	if err != nil {
@@ -180,7 +180,7 @@ func (p *SSHProxy) handleSSHSession(session ssh.Session) {
 	}
 
 	if hasCommand {
-		if err := serverSession.Run(session.RawCommand()); err != nil {
+		if err := serverSession.Run(strings.Join(session.Command(), " ")); err != nil {
 			log.Debugf("run command: %v", err)
 			p.handleProxyExitCode(session, err)
 		}

client/ssh/proxy/proxy_test.go

@@ -1,7 +1,6 @@
 package proxy
 
 import (
-	"bytes"
 	"context"
 	"crypto/rand"
 	"crypto/rsa"
@@ -246,191 +245,6 @@ func TestSSHProxy_Connect(t *testing.T) {
 	cancel()
 }
 
-// TestSSHProxy_CommandQuoting verifies that the proxy preserves shell quoting
-// when forwarding commands to the backend. This is critical for tools like
-// Ansible that send commands such as:
-//
-//	/bin/sh -c '( umask 77 && mkdir -p ... ) && sleep 0'
-//
-// The single quotes must be preserved so the backend shell receives the
-// subshell expression as a single argument to -c.
-func TestSSHProxy_CommandQuoting(t *testing.T) {
-	if testing.Short() {
-		t.Skip("Skipping integration test in short mode")
-	}
-
-	sshClient, cleanup := setupProxySSHClient(t)
-	defer cleanup()
-
-	// These commands simulate what the SSH protocol delivers as exec payloads.
-	// When a user types: ssh host '/bin/sh -c "( echo hello )"'
-	// the local shell strips the outer single quotes, and the SSH exec request
-	// contains the raw string: /bin/sh -c "( echo hello )"
-	//
-	// The proxy must forward this string verbatim. Using session.Command()
-	// (shlex.Split + strings.Join) strips the inner double quotes, breaking
-	// the command on the backend.
-	tests := []struct {
-		name    string
-		command string
-		expect  string
-	}{
-		{
-			name:    "subshell_in_double_quotes",
-			command: `/bin/sh -c "( echo from-subshell ) && echo outer"`,
-			expect:  "from-subshell\nouter\n",
-		},
-		{
-			name:    "printf_with_special_chars",
-			command: `/bin/sh -c "printf '%s\n' 'hello world'"`,
-			expect:  "hello world\n",
-		},
-		{
-			name:    "nested_command_substitution",
-			command: `/bin/sh -c "echo $(echo nested)"`,
-			expect:  "nested\n",
-		},
-	}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			session, err := sshClient.NewSession()
-			require.NoError(t, err)
-			defer func() { _ = session.Close() }()
-
-			var stderrBuf bytes.Buffer
-			session.Stderr = &stderrBuf
-
-			outputCh := make(chan []byte, 1)
-			errCh := make(chan error, 1)
-			go func() {
-				output, err := session.Output(tc.command)
-				outputCh <- output
-				errCh <- err
-			}()
-
-			select {
-			case output := <-outputCh:
-				err := <-errCh
-				if stderrBuf.Len() > 0 {
-					t.Logf("stderr: %s", stderrBuf.String())
-				}
-				require.NoError(t, err, "command should succeed: %s", tc.command)
-				assert.Equal(t, tc.expect, string(output), "output mismatch for: %s", tc.command)
-			case <-time.After(5 * time.Second):
-				t.Fatalf("command timed out: %s", tc.command)
-			}
-		})
-	}
-}
-
-// setupProxySSHClient creates a full proxy test environment and returns
-// an SSH client connected through the proxy to a backend NetBird SSH server.
-func setupProxySSHClient(t *testing.T) (*cryptossh.Client, func()) {
-	t.Helper()
-
-	const (
-		issuer   = "https://test-issuer.example.com"
-		audience = "test-audience"
-	)
-
-	jwksServer, privateKey, jwksURL := setupJWKSServer(t)
-
-	hostKey, err := nbssh.GeneratePrivateKey(nbssh.ED25519)
-	require.NoError(t, err)
-	hostPubKey, err := nbssh.GeneratePublicKey(hostKey)
-	require.NoError(t, err)
-
-	serverConfig := &server.Config{
-		HostKeyPEM: hostKey,
-		JWT: &server.JWTConfig{
-			Issuer:       issuer,
-			Audiences:    []string{audience},
-			KeysLocation: jwksURL,
-		},
-	}
-	sshServer := server.New(serverConfig)
-	sshServer.SetAllowRootLogin(true)
-
-	testUsername := testutil.GetTestUsername(t)
-	testJWTUser := "test-username"
-	testUserHash, err := sshuserhash.HashUserID(testJWTUser)
-	require.NoError(t, err)
-
-	authConfig := &sshauth.Config{
-		UserIDClaim:     sshauth.DefaultUserIDClaim,
-		AuthorizedUsers: []sshuserhash.UserIDHash{testUserHash},
-		MachineUsers: map[string][]uint32{
-			testUsername: {0},
-		},
-	}
-	sshServer.UpdateSSHAuth(authConfig)
-
-	sshServerAddr := server.StartTestServer(t, sshServer)
-
-	mockDaemon := startMockDaemon(t)
-
-	host, portStr, err := net.SplitHostPort(sshServerAddr)
-	require.NoError(t, err)
-	port, err := strconv.Atoi(portStr)
-	require.NoError(t, err)
-
-	mockDaemon.setHostKey(host, hostPubKey)
-
-	validToken := generateValidJWT(t, privateKey, issuer, audience, testJWTUser)
-	mockDaemon.setJWTToken(validToken)
-
-	proxyInstance, err := New(mockDaemon.addr, host, port, io.Discard, nil)
-	require.NoError(t, err)
-
-	origStdin := os.Stdin
-	origStdout := os.Stdout
-
-	stdinReader, stdinWriter, err := os.Pipe()
-	require.NoError(t, err)
-	stdoutReader, stdoutWriter, err := os.Pipe()
-	require.NoError(t, err)
-
-	os.Stdin = stdinReader
-	os.Stdout = stdoutWriter
-
-	clientConn, proxyConn := net.Pipe()
-
-	go func() { _, _ = io.Copy(stdinWriter, proxyConn) }()
-	go func() { _, _ = io.Copy(proxyConn, stdoutReader) }()
-
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-
-	go func() {
-		_ = proxyInstance.Connect(ctx)
-	}()
-
-	sshConfig := &cryptossh.ClientConfig{
-		User:            testutil.GetTestUsername(t),
-		Auth:            []cryptossh.AuthMethod{},
-		HostKeyCallback: cryptossh.InsecureIgnoreHostKey(),
-		Timeout:         5 * time.Second,
-	}
-
-	sshClientConn, chans, reqs, err := cryptossh.NewClientConn(clientConn, "test", sshConfig)
-	require.NoError(t, err)
-
-	client := cryptossh.NewClient(sshClientConn, chans, reqs)
-
-	cleanupFn := func() {
-		_ = client.Close()
-		_ = clientConn.Close()
-		cancel()
-		os.Stdin = origStdin
-		os.Stdout = origStdout
-		_ = sshServer.Stop()
-		mockDaemon.stop()
-		jwksServer.Close()
-	}
-
-	return client, cleanupFn
-}
-
 type mockDaemonServer struct {
 	proto.UnimplementedDaemonServiceServer
 	hostKeys map[string][]byte

client/ssh/server/session_handlers.go

@@ -60,7 +60,7 @@ func (s *Server) sessionHandler(session ssh.Session) {
 	}
 
 	ptyReq, winCh, isPty := session.Pty()
-	hasCommand := session.RawCommand() != ""
+	hasCommand := len(session.Command()) > 0
 
 	if isPty && !hasCommand {
 		// ssh <host> - PTY interactive session (login)

client/ui/debug.go

@@ -24,10 +24,9 @@ import (
 
 // Initial state for the debug collection
 type debugInitialState struct {
-	wasDown        bool
-	needsRestoreUp bool
-	logLevel       proto.LogLevel
-	isLevelTrace   bool
+	wasDown      bool
+	logLevel     proto.LogLevel
+	isLevelTrace bool
 }
 
 // Debug collection parameters
@@ -372,51 +371,46 @@ func (s *serviceClient) configureServiceForDebug(
 	conn proto.DaemonServiceClient,
 	state *debugInitialState,
 	enablePersistence bool,
-) {
+) error {
 	if state.wasDown {
 		if _, err := conn.Up(s.ctx, &proto.UpRequest{}); err != nil {
-			log.Warnf("failed to bring service up: %v", err)
-		} else {
-			log.Info("Service brought up for debug")
-			time.Sleep(time.Second * 10)
+			return fmt.Errorf("bring service up: %v", err)
 		}
+		log.Info("Service brought up for debug")
+		time.Sleep(time.Second * 10)
 	}
 
 	if !state.isLevelTrace {
 		if _, err := conn.SetLogLevel(s.ctx, &proto.SetLogLevelRequest{Level: proto.LogLevel_TRACE}); err != nil {
-			log.Warnf("failed to set log level to TRACE: %v", err)
-		} else {
-			log.Info("Log level set to TRACE for debug")
+			return fmt.Errorf("set log level to TRACE: %v", err)
 		}
+		log.Info("Log level set to TRACE for debug")
 	}
 
 	if _, err := conn.Down(s.ctx, &proto.DownRequest{}); err != nil {
-		log.Warnf("failed to bring service down: %v", err)
-	} else {
-		state.needsRestoreUp = !state.wasDown
-		time.Sleep(time.Second)
+		return fmt.Errorf("bring service down: %v", err)
 	}
+	time.Sleep(time.Second)
 
 	if enablePersistence {
 		if _, err := conn.SetSyncResponsePersistence(s.ctx, &proto.SetSyncResponsePersistenceRequest{
 			Enabled: true,
 		}); err != nil {
-			log.Warnf("failed to enable sync response persistence: %v", err)
-		} else {
-			log.Info("Sync response persistence enabled for debug")
+			return fmt.Errorf("enable sync response persistence: %v", err)
 		}
+		log.Info("Sync response persistence enabled for debug")
 	}
 
 	if _, err := conn.Up(s.ctx, &proto.UpRequest{}); err != nil {
-		log.Warnf("failed to bring service back up: %v", err)
-	} else {
-		state.needsRestoreUp = false
-		time.Sleep(time.Second * 3)
+		return fmt.Errorf("bring service back up: %v", err)
 	}
+	time.Sleep(time.Second * 3)
 
 	if _, err := conn.StartCPUProfile(s.ctx, &proto.StartCPUProfileRequest{}); err != nil {
 		log.Warnf("failed to start CPU profiling: %v", err)
 	}
+
+	return nil
 }
 
 func (s *serviceClient) collectDebugData(
@@ -430,7 +424,9 @@ func (s *serviceClient) collectDebugData(
 	var wg sync.WaitGroup
 	startProgressTracker(ctx, &wg, params.duration, progress)
 
-	s.configureServiceForDebug(conn, state, params.enablePersistence)
+	if err := s.configureServiceForDebug(conn, state, params.enablePersistence); err != nil {
+		return err
+	}
 
 	wg.Wait()
 	progress.progressBar.Hide()
@@ -486,17 +482,9 @@ func (s *serviceClient) createDebugBundleFromCollection(
 
 // Restore service to original state
 func (s *serviceClient) restoreServiceState(conn proto.DaemonServiceClient, state *debugInitialState) {
-	if state.needsRestoreUp {
-		if _, err := conn.Up(s.ctx, &proto.UpRequest{}); err != nil {
-			log.Warnf("failed to restore up state: %v", err)
-		} else {
-			log.Info("Service state restored to up")
-		}
-	}
-
 	if state.wasDown {
 		if _, err := conn.Down(s.ctx, &proto.DownRequest{}); err != nil {
-			log.Warnf("failed to restore down state: %v", err)
+			log.Errorf("Failed to restore down state: %v", err)
 		} else {
 			log.Info("Service state restored to down")
 		}
@@ -504,7 +492,7 @@ func (s *serviceClient) restoreServiceState(conn proto.DaemonServiceClient, stat
 
 	if !state.isLevelTrace {
 		if _, err := conn.SetLogLevel(s.ctx, &proto.SetLogLevelRequest{Level: state.logLevel}); err != nil {
-			log.Warnf("failed to restore log level: %v", err)
+			log.Errorf("Failed to restore log level: %v", err)
 		} else {
 			log.Info("Log level restored to original setting")
 		}

combined/cmd/root.go

@@ -29,7 +29,6 @@ import (
 	"github.com/netbirdio/netbird/management/server/telemetry"
 	"github.com/netbirdio/netbird/relay/healthcheck"
 	relayServer "github.com/netbirdio/netbird/relay/server"
-	"github.com/netbirdio/netbird/relay/server/listener"
 	"github.com/netbirdio/netbird/relay/server/listener/ws"
 	sharedMetrics "github.com/netbirdio/netbird/shared/metrics"
 	"github.com/netbirdio/netbird/shared/relay/auth"
@@ -524,7 +523,7 @@ func createManagementServer(cfg *CombinedConfig, mgmtConfig *nbconfig.Config) (*
 func createCombinedHandler(grpcServer *grpc.Server, httpHandler http.Handler, relaySrv *relayServer.Server, meter metric.Meter, cfg *CombinedConfig) http.Handler {
 	wsProxy := wsproxyserver.New(grpcServer, wsproxyserver.WithOTelMeter(meter))
 
-	var relayAcceptFn func(conn listener.Conn)
+	var relayAcceptFn func(conn net.Conn)
 	if relaySrv != nil {
 		relayAcceptFn = relaySrv.RelayAccept()
 	}
@@ -564,7 +563,7 @@ func createCombinedHandler(grpcServer *grpc.Server, httpHandler http.Handler, re
 }
 
 // handleRelayWebSocket handles incoming WebSocket connections for the relay service
-func handleRelayWebSocket(w http.ResponseWriter, r *http.Request, acceptFn func(conn listener.Conn), cfg *CombinedConfig) {
+func handleRelayWebSocket(w http.ResponseWriter, r *http.Request, acceptFn func(conn net.Conn), cfg *CombinedConfig) {
 	acceptOptions := &websocket.AcceptOptions{
 		OriginPatterns: []string{"*"},
 	}
@@ -586,9 +585,15 @@ func handleRelayWebSocket(w http.ResponseWriter, r *http.Request, acceptFn func(
 		return
 	}
 
+	lAddr, err := net.ResolveTCPAddr("tcp", cfg.Server.ListenAddress)
+	if err != nil {
+		_ = wsConn.Close(websocket.StatusInternalError, "internal error")
+		return
+	}
+
 	log.Debugf("Relay WS client connected from: %s", rAddr)
 
-	conn := ws.NewConn(wsConn, rAddr)
+	conn := ws.NewConn(wsConn, lAddr, rAddr)
 	acceptFn(conn)
 }
 

go.mod

@@ -63,7 +63,6 @@ require (
 	github.com/hashicorp/go-version v1.6.0
 	github.com/jackc/pgx/v5 v5.5.5
 	github.com/libdns/route53 v1.5.0
-	github.com/libp2p/go-nat v0.2.0
 	github.com/libp2p/go-netroute v0.2.1
 	github.com/lrh3321/ipset-go v0.0.0-20250619021614-54a0a98ace81
 	github.com/mdlayher/socket v0.5.1
@@ -201,12 +200,10 @@ require (
 	github.com/hashicorp/errwrap v1.1.0 // indirect
 	github.com/hashicorp/go-uuid v1.0.3 // indirect
 	github.com/huandu/xstrings v1.5.0 // indirect
-	github.com/huin/goupnp v1.2.0 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
 	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
 	github.com/jackc/puddle/v2 v2.2.1 // indirect
-	github.com/jackpal/go-nat-pmp v1.0.2 // indirect
 	github.com/jeandeaual/go-locale v0.0.0-20250612000132-0ef82f21eade // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/jinzhu/now v1.1.5 // indirect
@@ -216,7 +213,6 @@ require (
 	github.com/kelseyhightower/envconfig v1.4.0 // indirect
 	github.com/klauspost/compress v1.18.0 // indirect
 	github.com/klauspost/cpuid/v2 v2.2.7 // indirect
-	github.com/koron/go-ssdp v0.0.4 // indirect
 	github.com/kr/fs v0.1.0 // indirect
 	github.com/lib/pq v1.10.9 // indirect
 	github.com/libdns/libdns v0.2.2 // indirect

go.sum

@@ -281,8 +281,6 @@ github.com/hashicorp/go-version v1.6.0/go.mod h1:fltr4n8CU8Ke44wwGCBoEymUuxUHl09
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
 github.com/huandu/xstrings v1.5.0 h1:2ag3IFq9ZDANvthTwTiqSSZLjDc+BedvHPAp5tJy2TI=
 github.com/huandu/xstrings v1.5.0/go.mod h1:y5/lhBue+AyNmUVz9RLU9xbLR0o4KIIExikq4ovT0aE=
-github.com/huin/goupnp v1.2.0 h1:uOKW26NG1hsSSbXIZ1IR7XP9Gjd1U8pnLaCMgntmkmY=
-github.com/huin/goupnp v1.2.0/go.mod h1:gnGPsThkYa7bFi/KWmEysQRf48l2dvR5bxr2OFckNX8=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
 github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
@@ -293,8 +291,6 @@ github.com/jackc/pgx/v5 v5.5.5 h1:amBjrZVmksIdNjxGW/IiIMzxMKZFelXbUoPNb+8sjQw=
 github.com/jackc/pgx/v5 v5.5.5/go.mod h1:ez9gk+OAat140fv9ErkZDYFWmXLfV+++K0uAOiwgm1A=
 github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk=
 github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
-github.com/jackpal/go-nat-pmp v1.0.2 h1:KzKSgb7qkJvOUTqYl9/Hg/me3pWgBmERKrTGD7BdWus=
-github.com/jackpal/go-nat-pmp v1.0.2/go.mod h1:QPH045xvCAeXUZOxsnwmrtiCoxIr9eob+4orBN1SBKc=
 github.com/jcmturner/aescts/v2 v2.0.0 h1:9YKLH6ey7H4eDBXW8khjYslgyqG2xZikXP0EQFKrle8=
 github.com/jcmturner/aescts/v2 v2.0.0/go.mod h1:AiaICIRyfYg35RUkr8yESTqvSy7csK90qZ5xfvvsoNs=
 github.com/jcmturner/dnsutils/v2 v2.0.0 h1:lltnkeZGL0wILNvrNiVCR6Ro5PGU/SeBvVO/8c/iPbo=
@@ -332,8 +328,6 @@ github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYW
 github.com/klauspost/cpuid/v2 v2.0.12/go.mod h1:g2LTdtYhdyuGPqyWyv7qRAmj1WBqxuObKfj5c0PQa7c=
 github.com/klauspost/cpuid/v2 v2.2.7 h1:ZWSB3igEs+d0qvnxR/ZBzXVmxkgt8DdzP6m9pfuVLDM=
 github.com/klauspost/cpuid/v2 v2.2.7/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws=
-github.com/koron/go-ssdp v0.0.4 h1:1IDwrghSKYM7yLf7XCzbByg2sJ/JcNOZRXS2jczTwz0=
-github.com/koron/go-ssdp v0.0.4/go.mod h1:oDXq+E5IL5q0U8uSBcoAXzTzInwy5lEgC91HoKtbmZk=
 github.com/kr/fs v0.1.0 h1:Jskdu9ieNAYnjxsi0LbQp1ulIKZV1LAFgK1tWhpZgl8=
 github.com/kr/fs v0.1.0/go.mod h1:FFnZGqtBN9Gxj7eW1uZ42v5BccTP0vu6NEaFoC2HwRg=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
@@ -352,8 +346,6 @@ github.com/libdns/libdns v0.2.2 h1:O6ws7bAfRPaBsgAYt8MDe2HcNBGC29hkZ9MX2eUSX3s=
 github.com/libdns/libdns v0.2.2/go.mod h1:4Bj9+5CQiNMVGf87wjX4CY3HQJypUHRuLvlsfsZqLWQ=
 github.com/libdns/route53 v1.5.0 h1:2SKdpPFl/qgWsXQvsLNJJAoX7rSxlk7zgoL4jnWdXVA=
 github.com/libdns/route53 v1.5.0/go.mod h1:joT4hKmaTNKHEwb7GmZ65eoDz1whTu7KKYPS8ZqIh6Q=
-github.com/libp2p/go-nat v0.2.0 h1:Tyz+bUFAYqGyJ/ppPPymMGbIgNRH+WqC5QrT5fKrrGk=
-github.com/libp2p/go-nat v0.2.0/go.mod h1:3MJr+GRpRkyT65EpVPBstXLvOlAPzUVlG6Pwg9ohLJk=
 github.com/lrh3321/ipset-go v0.0.0-20250619021614-54a0a98ace81 h1:J56rFEfUTFT9j9CiRXhi1r8lUJ4W5idG3CiaBZGojNU=
 github.com/lrh3321/ipset-go v0.0.0-20250619021614-54a0a98ace81/go.mod h1:RD8ML/YdXctQ7qbcizZkw5mZ6l8Ogrl1dodBzVJduwI=
 github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I=

relay/server/handshake.go

@@ -1,13 +1,11 @@
 package server
 
 import (
-	"context"
 	"fmt"
-	"time"
+	"net"
 
 	log "github.com/sirupsen/logrus"
 
-	"github.com/netbirdio/netbird/relay/server/listener"
 	"github.com/netbirdio/netbird/shared/relay/messages"
 	//nolint:staticcheck
 	"github.com/netbirdio/netbird/shared/relay/messages/address"
@@ -15,12 +13,6 @@ import (
 	authmsg "github.com/netbirdio/netbird/shared/relay/messages/auth"
 )
 
-const (
-	// handshakeTimeout bounds how long a connection may remain in the
-	// pre-authentication handshake phase before being closed.
-	handshakeTimeout = 10 * time.Second
-)
-
 type Validator interface {
 	Validate(any) error
 	// Deprecated: Use Validate instead.
@@ -66,7 +58,7 @@ func marshalResponseHelloMsg(instanceURL string) ([]byte, error) {
 }
 
 type handshake struct {
-	conn        listener.Conn
+	conn        net.Conn
 	validator   Validator
 	preparedMsg *preparedMsg
 
@@ -74,9 +66,9 @@ type handshake struct {
 	peerID              *messages.PeerID
 }
 
-func (h *handshake) handshakeReceive(ctx context.Context) (*messages.PeerID, error) {
+func (h *handshake) handshakeReceive() (*messages.PeerID, error) {
 	buf := make([]byte, messages.MaxHandshakeSize)
-	n, err := h.conn.Read(ctx, buf)
+	n, err := h.conn.Read(buf)
 	if err != nil {
 		return nil, fmt.Errorf("read from %s: %w", h.conn.RemoteAddr(), err)
 	}
@@ -111,7 +103,7 @@ func (h *handshake) handshakeReceive(ctx context.Context) (*messages.PeerID, err
 	return peerID, nil
 }
 
-func (h *handshake) handshakeResponse(ctx context.Context) error {
+func (h *handshake) handshakeResponse() error {
 	var responseMsg []byte
 	if h.handshakeMethodAuth {
 		responseMsg = h.preparedMsg.responseAuthMsg
@@ -119,7 +111,7 @@ func (h *handshake) handshakeResponse(ctx context.Context) error {
 		responseMsg = h.preparedMsg.responseHelloMsg
 	}
 
-	if _, err := h.conn.Write(ctx, responseMsg); err != nil {
+	if _, err := h.conn.Write(responseMsg); err != nil {
 		return fmt.Errorf("handshake response write to %s (%s): %w", h.peerID, h.conn.RemoteAddr(), err)
 	}
 

relay/server/listener/conn.go

@@ -1,14 +0,0 @@
-package listener
-
-import (
-	"context"
-	"net"
-)
-
-// Conn is the relay connection contract implemented by WS and QUIC transports.
-type Conn interface {
-	Read(ctx context.Context, b []byte) (n int, err error)
-	Write(ctx context.Context, b []byte) (n int, err error)
-	RemoteAddr() net.Addr
-	Close() error
-}

relay/server/listener/listener.go

@@ -0,0 +1,14 @@
+package listener
+
+import (
+	"context"
+	"net"
+
+	"github.com/netbirdio/netbird/relay/protocol"
+)
+
+type Listener interface {
+	Listen(func(conn net.Conn)) error
+	Shutdown(ctx context.Context) error
+	Protocol() protocol.Protocol
+}

relay/server/listener/quic/conn.go

@@ -3,26 +3,33 @@ package quic
 import (
 	"context"
 	"errors"
+	"fmt"
 	"net"
 	"sync"
+	"time"
 
 	"github.com/quic-go/quic-go"
 )
 
 type Conn struct {
-	session  *quic.Conn
-	closed   bool
-	closedMu sync.Mutex
+	session   *quic.Conn
+	closed    bool
+	closedMu  sync.Mutex
+	ctx       context.Context
+	ctxCancel context.CancelFunc
 }
 
 func NewConn(session *quic.Conn) *Conn {
+	ctx, cancel := context.WithCancel(context.Background())
 	return &Conn{
-		session: session,
+		session:   session,
+		ctx:       ctx,
+		ctxCancel: cancel,
 	}
 }
 
-func (c *Conn) Read(ctx context.Context, b []byte) (n int, err error) {
-	dgram, err := c.session.ReceiveDatagram(ctx)
+func (c *Conn) Read(b []byte) (n int, err error) {
+	dgram, err := c.session.ReceiveDatagram(c.ctx)
 	if err != nil {
 		return 0, c.remoteCloseErrHandling(err)
 	}
@@ -31,17 +38,33 @@ func (c *Conn) Read(ctx context.Context, b []byte) (n int, err error) {
 	return n, nil
 }
 
-func (c *Conn) Write(_ context.Context, b []byte) (int, error) {
+func (c *Conn) Write(b []byte) (int, error) {
 	if err := c.session.SendDatagram(b); err != nil {
 		return 0, c.remoteCloseErrHandling(err)
 	}
 	return len(b), nil
 }
 
+func (c *Conn) LocalAddr() net.Addr {
+	return c.session.LocalAddr()
+}
+
 func (c *Conn) RemoteAddr() net.Addr {
 	return c.session.RemoteAddr()
 }
 
+func (c *Conn) SetReadDeadline(t time.Time) error {
+	return nil
+}
+
+func (c *Conn) SetWriteDeadline(t time.Time) error {
+	return fmt.Errorf("SetWriteDeadline is not implemented")
+}
+
+func (c *Conn) SetDeadline(t time.Time) error {
+	return fmt.Errorf("SetDeadline is not implemented")
+}
+
 func (c *Conn) Close() error {
 	c.closedMu.Lock()
 	if c.closed {
@@ -51,6 +74,8 @@ func (c *Conn) Close() error {
 	c.closed = true
 	c.closedMu.Unlock()
 
+	c.ctxCancel() // Cancel the context
+
 	sessionErr := c.session.CloseWithError(0, "normal closure")
 	return sessionErr
 }

relay/server/listener/quic/listener.go

@@ -5,12 +5,12 @@ import (
 	"crypto/tls"
 	"errors"
 	"fmt"
+	"net"
 
 	"github.com/quic-go/quic-go"
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/relay/protocol"
-	relaylistener "github.com/netbirdio/netbird/relay/server/listener"
 	nbRelay "github.com/netbirdio/netbird/shared/relay"
 )
 
@@ -25,7 +25,7 @@ type Listener struct {
 	listener *quic.Listener
 }
 
-func (l *Listener) Listen(acceptFn func(conn relaylistener.Conn)) error {
+func (l *Listener) Listen(acceptFn func(conn net.Conn)) error {
 	quicCfg := &quic.Config{
 		EnableDatagrams:   true,
 		InitialPacketSize: nbRelay.QUICInitialPacketSize,

relay/server/listener/ws/conn.go

@@ -18,21 +18,25 @@ const (
 
 type Conn struct {
 	*websocket.Conn
+	lAddr *net.TCPAddr
 	rAddr *net.TCPAddr
 
 	closed   bool
 	closedMu sync.Mutex
+	ctx      context.Context
 }
 
-func NewConn(wsConn *websocket.Conn, rAddr *net.TCPAddr) *Conn {
+func NewConn(wsConn *websocket.Conn, lAddr, rAddr *net.TCPAddr) *Conn {
 	return &Conn{
 		Conn:  wsConn,
+		lAddr: lAddr,
 		rAddr: rAddr,
+		ctx:   context.Background(),
 	}
 }
 
-func (c *Conn) Read(ctx context.Context, b []byte) (n int, err error) {
-	t, r, err := c.Reader(ctx)
+func (c *Conn) Read(b []byte) (n int, err error) {
+	t, r, err := c.Reader(c.ctx)
 	if err != nil {
 		return 0, c.ioErrHandling(err)
 	}
@@ -52,18 +56,34 @@ func (c *Conn) Read(ctx context.Context, b []byte) (n int, err error) {
 // Write writes a binary message with the given payload.
 // It does not block until fill the internal buffer.
 // If the buffer filled up, wait until the buffer is drained or timeout.
-func (c *Conn) Write(ctx context.Context, b []byte) (int, error) {
-	ctx, ctxCancel := context.WithTimeout(ctx, writeTimeout)
+func (c *Conn) Write(b []byte) (int, error) {
+	ctx, ctxCancel := context.WithTimeout(c.ctx, writeTimeout)
 	defer ctxCancel()
 
 	err := c.Conn.Write(ctx, websocket.MessageBinary, b)
 	return len(b), err
 }
 
+func (c *Conn) LocalAddr() net.Addr {
+	return c.lAddr
+}
+
 func (c *Conn) RemoteAddr() net.Addr {
 	return c.rAddr
 }
 
+func (c *Conn) SetReadDeadline(t time.Time) error {
+	return fmt.Errorf("SetReadDeadline is not implemented")
+}
+
+func (c *Conn) SetWriteDeadline(t time.Time) error {
+	return fmt.Errorf("SetWriteDeadline is not implemented")
+}
+
+func (c *Conn) SetDeadline(t time.Time) error {
+	return fmt.Errorf("SetDeadline is not implemented")
+}
+
 func (c *Conn) Close() error {
 	c.closedMu.Lock()
 	c.closed = true

relay/server/listener/ws/listener.go

@@ -7,13 +7,11 @@ import (
 	"fmt"
 	"net"
 	"net/http"
-	"time"
 
 	"github.com/coder/websocket"
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/relay/protocol"
-	relaylistener "github.com/netbirdio/netbird/relay/server/listener"
 	"github.com/netbirdio/netbird/shared/relay"
 )
 
@@ -29,19 +27,18 @@ type Listener struct {
 	TLSConfig *tls.Config
 
 	server   *http.Server
-	acceptFn func(conn relaylistener.Conn)
+	acceptFn func(conn net.Conn)
 }
 
-func (l *Listener) Listen(acceptFn func(conn relaylistener.Conn)) error {
+func (l *Listener) Listen(acceptFn func(conn net.Conn)) error {
 	l.acceptFn = acceptFn
 	mux := http.NewServeMux()
 	mux.HandleFunc(URLPath, l.onAccept)
 
 	l.server = &http.Server{
-		Addr:              l.Address,
-		Handler:           mux,
-		TLSConfig:         l.TLSConfig,
-		ReadHeaderTimeout: 5 * time.Second,
+		Addr:      l.Address,
+		Handler:   mux,
+		TLSConfig: l.TLSConfig,
 	}
 
 	log.Infof("WS server listening address: %s", l.Address)
@@ -96,9 +93,18 @@ func (l *Listener) onAccept(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	lAddr, err := net.ResolveTCPAddr("tcp", l.server.Addr)
+	if err != nil {
+		err = wsConn.Close(websocket.StatusInternalError, "internal error")
+		if err != nil {
+			log.Errorf("failed to close ws connection: %s", err)
+		}
+		return
+	}
+
 	log.Infof("WS client connected from: %s", rAddr)
 
-	conn := NewConn(wsConn, rAddr)
+	conn := NewConn(wsConn, lAddr, rAddr)
 	l.acceptFn(conn)
 }
 

relay/server/peer.go

@@ -10,7 +10,6 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/relay/metrics"
-	"github.com/netbirdio/netbird/relay/server/listener"
 	"github.com/netbirdio/netbird/relay/server/store"
 	"github.com/netbirdio/netbird/shared/relay/healthcheck"
 	"github.com/netbirdio/netbird/shared/relay/messages"
@@ -27,14 +26,11 @@ type Peer struct {
 	metrics  *metrics.Metrics
 	log      *log.Entry
 	id       messages.PeerID
-	conn     listener.Conn
+	conn     net.Conn
 	connMu   sync.RWMutex
 	store    *store.Store
 	notifier *store.PeerNotifier
 
-	ctx       context.Context
-	ctxCancel context.CancelFunc
-
 	peersListener *store.Listener
 
 	// between the online peer collection step and the notification sending should not be sent offline notifications from another thread
@@ -42,17 +38,14 @@ type Peer struct {
 }
 
 // NewPeer creates a new Peer instance and prepare custom logging
-func NewPeer(metrics *metrics.Metrics, id messages.PeerID, conn listener.Conn, store *store.Store, notifier *store.PeerNotifier) *Peer {
-	ctx, cancel := context.WithCancel(context.Background())
+func NewPeer(metrics *metrics.Metrics, id messages.PeerID, conn net.Conn, store *store.Store, notifier *store.PeerNotifier) *Peer {
 	p := &Peer{
-		metrics:   metrics,
-		log:       log.WithField("peer_id", id.String()),
-		id:        id,
-		conn:      conn,
-		store:     store,
-		notifier:  notifier,
-		ctx:       ctx,
-		ctxCancel: cancel,
+		metrics:  metrics,
+		log:      log.WithField("peer_id", id.String()),
+		id:       id,
+		conn:     conn,
+		store:    store,
+		notifier: notifier,
 	}
 
 	return p
@@ -64,7 +57,6 @@ func NewPeer(metrics *metrics.Metrics, id messages.PeerID, conn listener.Conn, s
 func (p *Peer) Work() {
 	p.peersListener = p.notifier.NewListener(p.sendPeersOnline, p.sendPeersWentOffline)
 	defer func() {
-		p.ctxCancel()
 		p.notifier.RemoveListener(p.peersListener)
 
 		if err := p.conn.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
@@ -72,7 +64,8 @@ func (p *Peer) Work() {
 		}
 	}()
 
-	ctx := p.ctx
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
 	hc := healthcheck.NewSender(p.log)
 	go hc.StartHealthCheck(ctx)
@@ -80,7 +73,7 @@ func (p *Peer) Work() {
 
 	buf := make([]byte, bufferSize)
 	for {
-		n, err := p.conn.Read(ctx, buf)
+		n, err := p.conn.Read(buf)
 		if err != nil {
 			if !errors.Is(err, net.ErrClosed) {
 				p.log.Errorf("failed to read message: %s", err)
@@ -138,10 +131,10 @@ func (p *Peer) handleMsgType(ctx context.Context, msgType messages.MsgType, hc *
 }
 
 // Write writes data to the connection
-func (p *Peer) Write(ctx context.Context, b []byte) (int, error) {
+func (p *Peer) Write(b []byte) (int, error) {
 	p.connMu.RLock()
 	defer p.connMu.RUnlock()
-	return p.conn.Write(ctx, b)
+	return p.conn.Write(b)
 }
 
 // CloseGracefully closes the connection with the peer gracefully. Send a close message to the client and close the
@@ -154,7 +147,6 @@ func (p *Peer) CloseGracefully(ctx context.Context) {
 		p.log.Errorf("failed to send close message to peer: %s", p.String())
 	}
 
-	p.ctxCancel()
 	if err := p.conn.Close(); err != nil {
 		p.log.Errorf(errCloseConn, err)
 	}
@@ -164,7 +156,6 @@ func (p *Peer) Close() {
 	p.connMu.Lock()
 	defer p.connMu.Unlock()
 
-	p.ctxCancel()
 	if err := p.conn.Close(); err != nil {
 		p.log.Errorf(errCloseConn, err)
 	}
@@ -179,15 +170,26 @@ func (p *Peer) writeWithTimeout(ctx context.Context, buf []byte) error {
 	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
 	defer cancel()
 
-	_, err := p.conn.Write(ctx, buf)
-	return err
+	writeDone := make(chan struct{})
+	var err error
+	go func() {
+		_, err = p.conn.Write(buf)
+		close(writeDone)
+	}()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-writeDone:
+		return err
+	}
 }
 
 func (p *Peer) handleHealthcheckEvents(ctx context.Context, hc *healthcheck.Sender) {
 	for {
 		select {
 		case <-hc.HealthCheck:
-			_, err := p.Write(ctx, messages.MarshalHealthcheck())
+			_, err := p.Write(messages.MarshalHealthcheck())
 			if err != nil {
 				p.log.Errorf("failed to send healthcheck message: %s", err)
 				return
@@ -226,12 +228,12 @@ func (p *Peer) handleTransportMsg(msg []byte) {
 		return
 	}
 
-	n, err := dp.Write(dp.ctx, msg)
+	n, err := dp.Write(msg)
 	if err != nil {
 		p.log.Errorf("failed to write transport message to: %s", dp.String())
 		return
 	}
-	p.metrics.TransferBytesSent.Add(p.ctx, int64(n))
+	p.metrics.TransferBytesSent.Add(context.Background(), int64(n))
 }
 
 func (p *Peer) handleSubscribePeerState(msg []byte) {
@@ -274,7 +276,7 @@ func (p *Peer) sendPeersOnline(peers []messages.PeerID) {
 	}
 
 	for n, msg := range msgs {
-		if _, err := p.Write(p.ctx, msg); err != nil {
+		if _, err := p.Write(msg); err != nil {
 			p.log.Errorf("failed to write %d. peers offline message: %s", n, err)
 		}
 	}
@@ -291,7 +293,7 @@ func (p *Peer) sendPeersWentOffline(peers []messages.PeerID) {
 	}
 
 	for n, msg := range msgs {
-		if _, err := p.Write(p.ctx, msg); err != nil {
+		if _, err := p.Write(msg); err != nil {
 			p.log.Errorf("failed to write %d. peers offline message: %s", n, err)
 		}
 	}

relay/server/relay.go

@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"fmt"
+	"net"
 	"net/url"
 	"sync"
 	"time"
@@ -12,20 +13,11 @@ import (
 	"go.opentelemetry.io/otel/metric"
 
 	"github.com/netbirdio/netbird/relay/healthcheck/peerid"
-	"github.com/netbirdio/netbird/relay/protocol"
-	"github.com/netbirdio/netbird/relay/server/listener"
-
 	//nolint:staticcheck
 	"github.com/netbirdio/netbird/relay/metrics"
 	"github.com/netbirdio/netbird/relay/server/store"
 )
 
-type Listener interface {
-	Listen(func(conn listener.Conn)) error
-	Shutdown(ctx context.Context) error
-	Protocol() protocol.Protocol
-}
-
 type Config struct {
 	Meter          metric.Meter
 	ExposedAddress string
@@ -117,7 +109,7 @@ func NewRelay(config Config) (*Relay, error) {
 }
 
 // Accept start to handle a new peer connection
-func (r *Relay) Accept(conn listener.Conn) {
+func (r *Relay) Accept(conn net.Conn) {
 	acceptTime := time.Now()
 	r.closeMu.RLock()
 	defer r.closeMu.RUnlock()
@@ -125,15 +117,12 @@ func (r *Relay) Accept(conn listener.Conn) {
 		return
 	}
 
-	hsCtx, hsCancel := context.WithTimeout(context.Background(), handshakeTimeout)
-	defer hsCancel()
-
 	h := handshake{
 		conn:        conn,
 		validator:   r.validator,
 		preparedMsg: r.preparedMsg,
 	}
-	peerID, err := h.handshakeReceive(hsCtx)
+	peerID, err := h.handshakeReceive()
 	if err != nil {
 		if peerid.IsHealthCheck(peerID) {
 			log.Debugf("health check connection from %s", conn.RemoteAddr())
@@ -165,7 +154,7 @@ func (r *Relay) Accept(conn listener.Conn) {
 		r.metrics.PeerDisconnected(peer.String())
 	}()
 
-	if err := h.handshakeResponse(hsCtx); err != nil {
+	if err := h.handshakeResponse(); err != nil {
 		log.Errorf("failed to send handshake response, close peer: %s", err)
 		peer.Close()
 	}

relay/server/server.go

@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"crypto/tls"
+	"net"
 	"net/url"
 	"sync"
 
@@ -30,7 +31,7 @@ type ListenerConfig struct {
 // In a new HTTP connection, the server will accept the connection and pass it to the Relay server via the Accept method.
 type Server struct {
 	relay       *Relay
-	listeners   []Listener
+	listeners   []listener.Listener
 	listenerMux sync.Mutex
 }
 
@@ -55,7 +56,7 @@ func NewServer(config Config) (*Server, error) {
 	}
 	return &Server{
 		relay:     relay,
-		listeners: make([]Listener, 0, 2),
+		listeners: make([]listener.Listener, 0, 2),
 	}, nil
 }
 
@@ -85,7 +86,7 @@ func (r *Server) Listen(cfg ListenerConfig) error {
 	wg := sync.WaitGroup{}
 	for _, l := range r.listeners {
 		wg.Add(1)
-		go func(listener Listener) {
+		go func(listener listener.Listener) {
 			defer wg.Done()
 			errChan <- listener.Listen(r.relay.Accept)
 		}(l)
@@ -138,6 +139,6 @@ func (r *Server) InstanceURL() url.URL {
 // RelayAccept returns the relay's Accept function for handling incoming connections.
 // This allows external HTTP handlers to route connections to the relay without
 // starting the relay's own listeners.
-func (r *Server) RelayAccept() func(conn listener.Conn) {
+func (r *Server) RelayAccept() func(conn net.Conn) {
 	return r.relay.Accept
 }