Merge branch 'main' into ssh-rewrite

client/internal/acl/manager_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/netbirdio/netbird/client/firewall"
+	"github.com/netbirdio/netbird/client/iface"
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
 	"github.com/netbirdio/netbird/client/internal/acl/mocks"
 	"github.com/netbirdio/netbird/client/internal/netflow"
@@ -52,7 +53,7 @@ func TestDefaultManager(t *testing.T) {
 	}).AnyTimes()
 	ifaceMock.EXPECT().GetWGDevice().Return(nil).AnyTimes()
 
-	fw, err := firewall.NewFirewall(ifaceMock, nil, flowLogger, false)
+	fw, err := firewall.NewFirewall(ifaceMock, nil, flowLogger, false, iface.DefaultMTU)
 	require.NoError(t, err)
 	defer func() {
 		err = fw.Close(nil)
@@ -170,7 +171,7 @@ func TestDefaultManagerStateless(t *testing.T) {
 	}).AnyTimes()
 	ifaceMock.EXPECT().GetWGDevice().Return(nil).AnyTimes()
 
-	fw, err := firewall.NewFirewall(ifaceMock, nil, flowLogger, false)
+	fw, err := firewall.NewFirewall(ifaceMock, nil, flowLogger, false, iface.DefaultMTU)
 	require.NoError(t, err)
 	defer func() {
 		err = fw.Close(nil)

client/internal/connect.go

@@ -289,15 +289,18 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
 		}
 
 		<-engineCtx.Done()
+
 		c.engineMutex.Lock()
-		if c.engine != nil && c.engine.wgInterface != nil {
-			log.Infof("ensuring %s is removed, Netbird engine context cancelled", c.engine.wgInterface.Name())
-			if err := c.engine.Stop(); err != nil {
+		engine := c.engine
+		c.engine = nil
+		c.engineMutex.Unlock()
+
+		if engine != nil && engine.wgInterface != nil {
+			log.Infof("ensuring %s is removed, Netbird engine context cancelled", engine.wgInterface.Name())
+			if err := engine.Stop(); err != nil {
 				log.Errorf("Failed to stop engine: %v", err)
 			}
-			c.engine = nil
 		}
-		c.engineMutex.Unlock()
 		c.statusRecorder.ClientTeardown()
 
 		backOff.Reset()
@@ -382,19 +385,12 @@ func (c *ConnectClient) Status() StatusType {
 }
 
 func (c *ConnectClient) Stop() error {
-	if c == nil {
-		return nil
+	engine := c.Engine()
+	if engine != nil {
+		if err := engine.Stop(); err != nil {
+			return fmt.Errorf("stop engine: %w", err)
+		}
 	}
-	c.engineMutex.Lock()
-	defer c.engineMutex.Unlock()
-
-	if c.engine == nil {
-		return nil
-	}
-	if err := c.engine.Stop(); err != nil {
-		return fmt.Errorf("stop engine: %w", err)
-	}
-
 	return nil
 }
 

client/internal/dns/server.go

@@ -65,8 +65,9 @@ type hostManagerWithOriginalNS interface {
 
 // DefaultServer dns server object
 type DefaultServer struct {
-	ctx       context.Context
-	ctxCancel context.CancelFunc
+	ctx        context.Context
+	ctxCancel  context.CancelFunc
+	shutdownWg sync.WaitGroup
 	// disableSys disables system DNS management (e.g., /etc/resolv.conf updates) while keeping the DNS service running.
 	// This is different from ServiceEnable=false from management which completely disables the DNS service.
 	disableSys         bool
@@ -318,6 +319,7 @@ func (s *DefaultServer) DnsIP() netip.Addr {
 // Stop stops the server
 func (s *DefaultServer) Stop() {
 	s.ctxCancel()
+	s.shutdownWg.Wait()
 
 	s.mux.Lock()
 	defer s.mux.Unlock()
@@ -507,8 +509,9 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 
 	s.applyHostConfig()
 
+	s.shutdownWg.Add(1)
 	go func() {
-		// persist dns state right away
+		defer s.shutdownWg.Done()
 		if err := s.stateManager.PersistState(s.ctx); err != nil {
 			log.Errorf("Failed to persist dns state: %v", err)
 		}

client/internal/dns/server_test.go

@@ -944,7 +944,7 @@ func createWgInterfaceWithBind(t *testing.T) (*iface.WGIface, error) {
 		return nil, err
 	}
 
-	pf, err := uspfilter.Create(wgIface, false, flowLogger)
+	pf, err := uspfilter.Create(wgIface, false, flowLogger, iface.DefaultMTU)
 	if err != nil {
 		t.Fatalf("failed to create uspfilter: %v", err)
 		return nil, err

client/internal/dnsfwd/manager.go

@@ -15,6 +15,7 @@ import (
 	nberrors "github.com/netbirdio/netbird/client/errors"
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/iface/wgaddr"
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/route"
@@ -134,6 +135,8 @@ func (m *Manager) Stop(ctx context.Context) error {
 		}
 	}
 
+	m.unregisterNetstackServices()
+
 	if err := m.dropDNSFirewall(); err != nil {
 		mErr = multierror.Append(mErr, err)
 	}
@@ -158,21 +161,50 @@ func (m *Manager) allowDNSFirewall() error {
 
 	dnsRules, err := m.firewall.AddPeerFiltering(nil, net.IP{0, 0, 0, 0}, firewall.ProtocolUDP, nil, dport, firewall.ActionAccept, "")
 	if err != nil {
-		log.Errorf("failed to add allow DNS router rules, err: %v", err)
-		return err
+		return fmt.Errorf("add udp firewall rule: %w", err)
 	}
-	m.fwRules = dnsRules
 
 	tcpRules, err := m.firewall.AddPeerFiltering(nil, net.IP{0, 0, 0, 0}, firewall.ProtocolTCP, nil, dport, firewall.ActionAccept, "")
 	if err != nil {
-		log.Errorf("failed to add allow DNS router rules, err: %v", err)
-		return err
+		return fmt.Errorf("add tcp firewall rule: %w", err)
 	}
+
+	if err := m.firewall.Flush(); err != nil {
+		return fmt.Errorf("flush: %w", err)
+	}
+
+	m.fwRules = dnsRules
 	m.tcpRules = tcpRules
 
+	m.registerNetstackServices()
+
 	return nil
 }
 
+func (m *Manager) registerNetstackServices() {
+	if netstackNet := m.wgIface.GetNet(); netstackNet != nil {
+		if registrar, ok := m.firewall.(interface {
+			RegisterNetstackService(protocol nftypes.Protocol, port uint16)
+		}); ok {
+			registrar.RegisterNetstackService(nftypes.TCP, m.serverPort)
+			registrar.RegisterNetstackService(nftypes.UDP, m.serverPort)
+			log.Debugf("registered DNS forwarder service with netstack for UDP/TCP:%d", m.serverPort)
+		}
+	}
+}
+
+func (m *Manager) unregisterNetstackServices() {
+	if netstackNet := m.wgIface.GetNet(); netstackNet != nil {
+		if registrar, ok := m.firewall.(interface {
+			UnregisterNetstackService(protocol nftypes.Protocol, port uint16)
+		}); ok {
+			registrar.UnregisterNetstackService(nftypes.TCP, m.serverPort)
+			registrar.UnregisterNetstackService(nftypes.UDP, m.serverPort)
+			log.Debugf("unregistered DNS forwarder service with netstack for UDP/TCP:%d", m.serverPort)
+		}
+	}
+}
+
 func (m *Manager) dropDNSFirewall() error {
 	var mErr *multierror.Error
 	for _, rule := range m.fwRules {

client/internal/engine.go

@@ -202,8 +202,10 @@ type Engine struct {
 	flowManager         nftypes.FlowManager
 
 	// WireGuard interface monitor
-	wgIfaceMonitor   *WGIfaceMonitor
-	wgIfaceMonitorWg sync.WaitGroup
+	wgIfaceMonitor *WGIfaceMonitor
+
+	// shutdownWg tracks all long-running goroutines to ensure clean shutdown
+	shutdownWg sync.WaitGroup
 
 	probeStunTurn *relay.StunTurnProbe
 }
@@ -306,17 +308,12 @@ func (e *Engine) Stop() error {
 		e.ingressGatewayMgr = nil
 	}
 
+	e.stopDNSForwarder()
+
 	if e.routeManager != nil {
 		e.routeManager.Stop(e.stateManager)
 	}
 
-	if e.dnsForwardMgr != nil {
-		if err := e.dnsForwardMgr.Stop(context.Background()); err != nil {
-			log.Errorf("failed to stop DNS forward: %v", err)
-		}
-		e.dnsForwardMgr = nil
-	}
-
 	if e.srWatcher != nil {
 		e.srWatcher.Close()
 	}
@@ -333,10 +330,6 @@ func (e *Engine) Stop() error {
 		e.cancel()
 	}
 
-	// very ugly but we want to remove peers from the WireGuard interface first before removing interface.
-	// Removing peers happens in the conn.Close() asynchronously
-	time.Sleep(500 * time.Millisecond)
-
 	e.close()
 
 	// stop flow manager after wg interface is gone
@@ -344,8 +337,6 @@ func (e *Engine) Stop() error {
 		e.flowManager.Close()
 	}
 
-	log.Infof("stopped Netbird Engine")
-
 	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
 	defer cancel()
 
@@ -356,12 +347,52 @@ func (e *Engine) Stop() error {
 		log.Errorf("failed to persist state: %v", err)
 	}
 
-	// Stop WireGuard interface monitor and wait for it to exit
-	e.wgIfaceMonitorWg.Wait()
+	timeout := e.calculateShutdownTimeout()
+	log.Debugf("waiting for goroutines to finish with timeout: %v", timeout)
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	if err := waitWithContext(shutdownCtx, &e.shutdownWg); err != nil {
+		log.Warnf("shutdown timeout exceeded after %v, some goroutines may still be running", timeout)
+	}
+
+	log.Infof("stopped Netbird Engine")
 
 	return nil
 }
 
+// calculateShutdownTimeout returns shutdown timeout: 10s base + 100ms per peer, capped at 30s.
+func (e *Engine) calculateShutdownTimeout() time.Duration {
+	peerCount := len(e.peerStore.PeersPubKey())
+
+	baseTimeout := 10 * time.Second
+	perPeerTimeout := time.Duration(peerCount) * 100 * time.Millisecond
+	timeout := baseTimeout + perPeerTimeout
+
+	maxTimeout := 30 * time.Second
+	if timeout > maxTimeout {
+		timeout = maxTimeout
+	}
+
+	return timeout
+}
+
+// waitWithContext waits for WaitGroup with timeout, returns ctx.Err() on timeout.
+func waitWithContext(ctx context.Context, wg *sync.WaitGroup) error {
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
 // Start creates a new WireGuard tunnel interface and listens to events from Signal and Management services
 // Connections to remote peers are not established here.
 // However, they will be established once an event with a list of peers to connect to will be received from Management Service
@@ -491,14 +522,14 @@ func (e *Engine) Start(netbirdConfig *mgmProto.NetbirdConfig, mgmtURL *url.URL)
 
 	// monitor WireGuard interface lifecycle and restart engine on changes
 	e.wgIfaceMonitor = NewWGIfaceMonitor()
-	e.wgIfaceMonitorWg.Add(1)
+	e.shutdownWg.Add(1)
 
 	go func() {
-		defer e.wgIfaceMonitorWg.Done()
+		defer e.shutdownWg.Done()
 
 		if shouldRestart, err := e.wgIfaceMonitor.Start(e.ctx, e.wgInterface.Name()); shouldRestart {
 			log.Infof("WireGuard interface monitor: %s, restarting engine", err)
-			e.restartEngine()
+			e.triggerClientRestart()
 		} else if err != nil {
 			log.Warnf("WireGuard interface monitor: %s", err)
 		}
@@ -514,7 +545,7 @@ func (e *Engine) createFirewall() error {
 	}
 
 	var err error
-	e.firewall, err = firewall.NewFirewall(e.wgInterface, e.stateManager, e.flowManager.GetLogger(), e.config.DisableServerRoutes)
+	e.firewall, err = firewall.NewFirewall(e.wgInterface, e.stateManager, e.flowManager.GetLogger(), e.config.DisableServerRoutes, e.config.MTU)
 	if err != nil || e.firewall == nil {
 		log.Errorf("failed creating firewall manager: %s", err)
 		return nil
@@ -898,7 +929,9 @@ func (e *Engine) updateConfig(conf *mgmProto.PeerConfig) error {
 // receiveManagementEvents connects to the Management Service event stream to receive updates from the management service
 // E.g. when a new peer has been registered and we are allowed to connect to it.
 func (e *Engine) receiveManagementEvents() {
+	e.shutdownWg.Add(1)
 	go func() {
+		defer e.shutdownWg.Done()
 		info, err := system.GetInfoWithChecks(e.ctx, e.checks)
 		if err != nil {
 			log.Warnf("failed to get system info with checks: %v", err)
@@ -1325,7 +1358,9 @@ func (e *Engine) createPeerConn(pubKey string, allowedIPs []netip.Prefix, agentV
 
 // receiveSignalEvents connects to the Signal Service event stream to negotiate connection with remote peers
 func (e *Engine) receiveSignalEvents() {
+	e.shutdownWg.Add(1)
 	go func() {
+		defer e.shutdownWg.Done()
 		// connect to a stream of messages coming from the signal server
 		err := e.signal.Receive(e.ctx, func(msg *sProto.Message) error {
 			e.syncMsgMux.Lock()
@@ -1676,8 +1711,10 @@ func (e *Engine) RunHealthProbes(waitForResult bool) bool {
 	return allHealthy
 }
 
-// restartEngine restarts the engine by cancelling the client context
-func (e *Engine) restartEngine() {
+// triggerClientRestart triggers a full client restart by cancelling the client context.
+// Note: This does NOT just restart the engine - it cancels the entire client context,
+// which causes the connect client's retry loop to create a completely new engine.
+func (e *Engine) triggerClientRestart() {
 	e.syncMsgMux.Lock()
 	defer e.syncMsgMux.Unlock()
 
@@ -1699,7 +1736,9 @@ func (e *Engine) startNetworkMonitor() {
 	}
 
 	e.networkMonitor = networkmonitor.New()
+	e.shutdownWg.Add(1)
 	go func() {
+		defer e.shutdownWg.Done()
 		if err := e.networkMonitor.Listen(e.ctx); err != nil {
 			if errors.Is(err, context.Canceled) {
 				log.Infof("network monitor stopped")
@@ -1709,8 +1748,8 @@ func (e *Engine) startNetworkMonitor() {
 			return
 		}
 
-		log.Infof("Network monitor: detected network change, restarting engine")
-		e.restartEngine()
+		log.Infof("Network monitor: detected network change, triggering client restart")
+		e.triggerClientRestart()
 	}()
 }
 
@@ -1819,7 +1858,6 @@ func (e *Engine) updateDNSForwarder(
 
 func (e *Engine) startDNSForwarder(fwdEntries []*dnsfwd.ForwarderEntry) {
 	e.dnsForwardMgr = dnsfwd.NewManager(e.firewall, e.statusRecorder, e.wgInterface)
-	e.registerDNSServices()
 
 	if err := e.dnsForwardMgr.Start(fwdEntries); err != nil {
 		log.Errorf("failed to start DNS forward: %v", err)
@@ -1839,34 +1877,9 @@ func (e *Engine) stopDNSForwarder() {
 		log.Errorf("failed to stop DNS forward: %v", err)
 	}
 
-	e.unregisterDNSServices()
 	e.dnsForwardMgr = nil
 }
 
-func (e *Engine) registerDNSServices() {
-	if netstackNet := e.wgInterface.GetNet(); netstackNet != nil {
-		if registrar, ok := e.firewall.(interface {
-			RegisterNetstackService(protocol nftypes.Protocol, port uint16)
-		}); ok {
-			registrar.RegisterNetstackService(nftypes.UDP, nbdns.ForwarderServerPort)
-			registrar.RegisterNetstackService(nftypes.TCP, nbdns.ForwarderServerPort)
-			log.Debugf("registered DNS forwarder service with netstack for UDP/TCP:%d", nbdns.ForwarderServerPort)
-		}
-	}
-}
-
-func (e *Engine) unregisterDNSServices() {
-	if netstackNet := e.wgInterface.GetNet(); netstackNet != nil {
-		if registrar, ok := e.firewall.(interface {
-			UnregisterNetstackService(protocol nftypes.Protocol, port uint16)
-		}); ok {
-			registrar.UnregisterNetstackService(nftypes.UDP, nbdns.ForwarderServerPort)
-			registrar.UnregisterNetstackService(nftypes.TCP, nbdns.ForwarderServerPort)
-			log.Debugf("unregistered DNS forwarder service with netstack for UDP/TCP:%d", nbdns.ForwarderServerPort)
-		}
-	}
-}
-
 func (e *Engine) GetNet() (*netstack.Net, error) {
 	e.syncMsgMux.Lock()
 	intf := e.wgInterface

client/internal/netflow/manager.go

@@ -24,6 +24,7 @@ import (
 // Manager handles netflow tracking and logging
 type Manager struct {
 	mux            sync.Mutex
+	shutdownWg     sync.WaitGroup
 	logger         nftypes.FlowLogger
 	flowConfig     *nftypes.FlowConfig
 	conntrack      nftypes.ConnTracker
@@ -105,8 +106,15 @@ func (m *Manager) resetClient() error {
 	ctx, cancel := context.WithCancel(context.Background())
 	m.cancel = cancel
 
-	go m.receiveACKs(ctx, flowClient)
-	go m.startSender(ctx)
+	m.shutdownWg.Add(2)
+	go func() {
+		defer m.shutdownWg.Done()
+		m.receiveACKs(ctx, flowClient)
+	}()
+	go func() {
+		defer m.shutdownWg.Done()
+		m.startSender(ctx)
+	}()
 
 	return nil
 }
@@ -176,11 +184,12 @@ func (m *Manager) Update(update *nftypes.FlowConfig) error {
 // Close cleans up all resources
 func (m *Manager) Close() {
 	m.mux.Lock()
-	defer m.mux.Unlock()
-
 	if err := m.disableFlow(); err != nil {
 		log.Warnf("failed to disable flow manager: %v", err)
 	}
+	m.mux.Unlock()
+
+	m.shutdownWg.Wait()
 }
 
 // GetLogger returns the flow logger

client/internal/peer/guard/sr_watcher.go

@@ -19,11 +19,10 @@ type SRWatcher struct {
 	signalClient chNotifier
 	relayManager chNotifier
 
-	listeners     map[chan struct{}]struct{}
-	mu            sync.Mutex
-	iFaceDiscover stdnet.ExternalIFaceDiscover
-	iceConfig     ice.Config
-
+	listeners        map[chan struct{}]struct{}
+	mu               sync.Mutex
+	iFaceDiscover    stdnet.ExternalIFaceDiscover
+	iceConfig        ice.Config
 	cancelIceMonitor context.CancelFunc
 }
 

client/internal/routemanager/manager.go

@@ -81,6 +81,7 @@ type DefaultManager struct {
 	ctx                  context.Context
 	stop                 context.CancelFunc
 	mux                  sync.Mutex
+	shutdownWg           sync.WaitGroup
 	clientNetworks       map[route.HAUniqueID]*client.Watcher
 	routeSelector        *routeselector.RouteSelector
 	serverRouter         *server.Router
@@ -283,6 +284,7 @@ func (m *DefaultManager) SetDNSForwarderPort(port uint16) {
 // Stop stops the manager watchers and clean firewall rules
 func (m *DefaultManager) Stop(stateManager *statemanager.Manager) {
 	m.stop()
+	m.shutdownWg.Wait()
 	if m.serverRouter != nil {
 		m.serverRouter.CleanUp()
 	}
@@ -485,7 +487,11 @@ func (m *DefaultManager) TriggerSelection(networks route.HAMap) {
 		}
 		clientNetworkWatcher := client.NewWatcher(config)
 		m.clientNetworks[id] = clientNetworkWatcher
-		go clientNetworkWatcher.Start()
+		m.shutdownWg.Add(1)
+		go func() {
+			defer m.shutdownWg.Done()
+			clientNetworkWatcher.Start()
+		}()
 		clientNetworkWatcher.SendUpdate(client.RoutesUpdate{Routes: routes})
 	}
 
@@ -527,7 +533,11 @@ func (m *DefaultManager) updateClientNetworks(updateSerial uint64, networks rout
 			}
 			clientNetworkWatcher = client.NewWatcher(config)
 			m.clientNetworks[id] = clientNetworkWatcher
-			go clientNetworkWatcher.Start()
+			m.shutdownWg.Add(1)
+			go func() {
+				defer m.shutdownWg.Done()
+				clientNetworkWatcher.Start()
+			}()
 		}
 		update := client.RoutesUpdate{
 			UpdateSerial: updateSerial,

client/internal/routeselector/routeselector.go

@@ -9,8 +9,6 @@ import (
 	"github.com/hashicorp/go-multierror"
 	"golang.org/x/exp/maps"
 
-	log "github.com/sirupsen/logrus"
-
 	"github.com/netbirdio/netbird/client/errors"
 	"github.com/netbirdio/netbird/route"
 )
@@ -128,13 +126,11 @@ func (rs *RouteSelector) IsSelected(routeID route.NetID) bool {
 	defer rs.mu.RUnlock()
 
 	if rs.deselectAll {
-		log.Debugf("Route %s not selected (deselect all)", routeID)
 		return false
 	}
 
 	_, deselected := rs.deselectedRoutes[routeID]
 	isSelected := !deselected
-	log.Debugf("Route %s selection status: %v (deselected: %v)", routeID, isSelected, deselected)
 	return isSelected
 }