Merge branch 'main' into fix/pkce-flow-session-extend

# Conflicts:
#	client/ios/NetBirdSDK/login.go
#	client/server/server.go
#	shared/management/proto/management.pb.go

client/ios/NetBirdSDK/client.go

@@ -4,12 +4,14 @@ package NetBirdSDK
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net/netip"
 	"os"
 	"sort"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -37,6 +39,8 @@ const (
 	AnonymizeLevelStrict  = nbAnonymize.LevelStrictString
 )
 
+var errClientAlreadyRunning = errors.New("client is already running")
+
 // RouteListener export internal RouteListener for mobile
 type NetworkChangeListener interface {
 	listener.NetworkChangeListener
@@ -74,25 +78,36 @@ type Client struct {
 	cacheDir              string
 	logFilePath           string
 	recorder              *peer.Status
-	ctxCancel             context.CancelFunc
-	ctxCancelLock         *sync.Mutex
 	deviceName            string
 	osName                string
 	osVersion             string
 	networkChangeListener listener.NetworkChangeListener
 	onHostDnsFn           func([]string)
 	dnsManager            dns.IosDnsManager
-	loginComplete         bool
+	loginComplete         atomic.Bool
 	// netMgr outlives engine restarts: it mirrors the OS connectivity, not
 	// the engine lifecycle. Run injects its state and sweeper into each new
 	// ConnectClient.
-	netMgr *netevents.Manager
-	// preloadedConfig holds config loaded from JSON (used on tvOS where file writes are blocked)
-	preloadedConfig *profilemanager.Config
+	netMgr              *netevents.Manager
+	preloadedConfigJSON atomic.Pointer[string]
 
+	// mdmSource holds the per-Client MDM policy source and its change
+	// detector as one unit. Set by SetMDMPolicyFetcher (called from the
+	// Swift side at extension init). Each Run passes the loader to the
+	// resolved Config so applyMDMPolicy picks up the active overlay. Nil
+	// means "MDM enforcement off for this Client".
+	mdmSource atomic.Pointer[mdmSource]
+
+	// stateMu guards the run lifecycle as one unit: the cancel installed by
+	// the current run, the channel it closes on exit, and the state it
+	// published. One run at a time: startRun refuses a second Run while the
+	// previous one has not exited, and the platform serializes Stop before
+	// Start, so no generation tracking is needed.
 	stateMu       sync.RWMutex
 	connectClient *internal.ConnectClient
 	config        *profilemanager.Config
+	runDone       chan struct{}
+	ctxCancel     context.CancelFunc
 }
 
 // NewClient instantiate a new Client
@@ -107,66 +122,69 @@ func NewClient(cfgFile, stateFile, cacheDir, logFilePath, deviceName string, osV
 		osName:                osName,
 		osVersion:             osVersion,
 		recorder:              recorder,
-		ctxCancelLock:         &sync.Mutex{},
 		networkChangeListener: networkChangeListener,
 		dnsManager:            dnsManager,
 		netMgr:                netevents.NewManager(recorder),
 	}
 }
 
-// SetConfigFromJSON loads config from a JSON string into memory.
-// This is used on tvOS where file writes to App Group containers are blocked.
-// When set, IsLoginRequired() and Run() will use this preloaded config instead of reading from file.
+// SetConfigFromJSON stores the JSON config that later loads resolve instead of the config file (tvOS).
 func (c *Client) SetConfigFromJSON(jsonStr string) error {
-	cfg, err := profilemanager.ConfigFromJSON(jsonStr)
-	if err != nil {
+	if _, err := profilemanager.ConfigFromJSON(jsonStr); err != nil {
 		log.Errorf("SetConfigFromJSON: failed to parse config JSON: %v", err)
 		return err
 	}
-	c.preloadedConfig = cfg
+	c.preloadedConfigJSON.Store(&jsonStr)
 	log.Infof("SetConfigFromJSON: config loaded successfully from JSON")
 	return nil
 }
 
+func (c *Client) loadConfig(input profilemanager.ConfigInput) (*profilemanager.Config, error) {
+	var cfg *profilemanager.Config
+	var err error
+	if preloaded := c.preloadedConfigJSON.Load(); preloaded != nil {
+		cfg, err = profilemanager.ConfigFromJSON(*preloaded)
+	} else {
+		cfg, err = profilemanager.DirectUpdateOrCreateConfig(input)
+	}
+	if err != nil {
+		return nil, err
+	}
+	c.applyMDMOverlay(cfg)
+	return cfg, nil
+}
+
 // Run start the internal client. It is a blocker function
 func (c *Client) Run(fd int32, interfaceName string, envList *EnvList) error {
 	exportEnvList(envList)
 	log.Infof("Starting NetBird client")
 	log.Debugf("Tunnel uses interface: %s", interfaceName)
 
-	var cfg *profilemanager.Config
-	var err error
-
-	// Use preloaded config if available (tvOS where file writes are blocked)
-	if c.preloadedConfig != nil {
-		log.Infof("Run: using preloaded config from memory")
-		cfg = c.preloadedConfig
-	} else {
-		log.Infof("Run: loading config from file")
-		// Use DirectUpdateOrCreateConfig to avoid atomic file operations (temp file + rename)
-		// which are blocked by the tvOS sandbox in App Group containers
-		cfg, err = profilemanager.DirectUpdateOrCreateConfig(profilemanager.ConfigInput{
-			ConfigPath:    c.cfgFile,
-			StateFilePath: c.stateFile,
-		})
-		if err != nil {
-			return err
-		}
+	cfg, err := c.loadConfig(profilemanager.ConfigInput{
+		ConfigPath:    c.cfgFile,
+		StateFilePath: c.stateFile,
+	})
+	if err != nil {
+		return err
 	}
 	c.recorder.UpdateManagementAddress(cfg.ManagementURL.String())
 	c.recorder.UpdateRosenpass(cfg.RosenpassEnabled, cfg.RosenpassPermissive)
 
-	var ctx context.Context
 	//nolint
 	ctxWithValues := context.WithValue(context.Background(), system.DeviceNameCtxKey, c.deviceName)
 	//nolint
 	ctxWithValues = context.WithValue(ctxWithValues, system.OsNameCtxKey, c.osName)
 	//nolint
 	ctxWithValues = context.WithValue(ctxWithValues, system.OsVersionCtxKey, c.osVersion)
-	c.ctxCancelLock.Lock()
-	ctx, c.ctxCancel = context.WithCancel(ctxWithValues)
-	defer c.ctxCancel()
-	c.ctxCancelLock.Unlock()
+	runCtx, runCancel := context.WithCancel(ctxWithValues)
+	defer runCancel()
+
+	done, err := c.startRun(runCancel)
+	if err != nil {
+		return err
+	}
+	defer c.finishRun(done)
+	ctx := runCtx
 
 	// No login pre-flight here. The engine's own loginToManagement (connect.go) performs
 	// the authoritative Login immediately before the first Sync, so a LoginSync() call at
@@ -215,16 +233,40 @@ func (c *Client) NotifyNetworkChange() {
 	c.netMgr.NotifyNetworkChange()
 }
 
-// Stop the internal client and free the resources
+// Stop cancels the running client and waits for the run loop to exit, so a
+// caller that restarts immediately cannot race the outgoing teardown.
 func (c *Client) Stop() {
-	c.ctxCancelLock.Lock()
-	defer c.ctxCancelLock.Unlock()
-	if c.ctxCancel == nil {
+	done := c.cancelRun()
+	if done == nil {
 		return
 	}
 
-	c.ctxCancel()
-	c.setState(nil, nil)
+	select {
+	case <-done:
+	case <-time.After(stopRunWaitTimeout):
+		log.Warnf("Stop: timed out waiting for the run loop to exit")
+	}
+}
+
+// StopWithoutWait cancels the running client without waiting for the run loop.
+// Use it where the caller is on a deadline the wait could overrun, such as
+// NEPacketTunnelProvider.stopTunnel, which iOS gives only a few seconds
+// before it kills the extension.
+func (c *Client) StopWithoutWait() {
+	c.cancelRun()
+}
+
+func (c *Client) cancelRun() chan struct{} {
+	c.stateMu.RLock()
+	done := c.runDone
+	cancel := c.ctxCancel
+	c.stateMu.RUnlock()
+
+	if cancel != nil {
+		cancel()
+	}
+
+	return done
 }
 
 // DebugBundle generates a debug bundle, uploads it and returns the upload key.
@@ -238,19 +280,13 @@ func (c *Client) DebugBundle(anonymize bool, anonymizeLevel string) (string, err
 
 	// If the engine hasn't been started, load config so we can reach management.
 	if cfg == nil {
-		if c.preloadedConfig != nil {
-			cfg = c.preloadedConfig
-		} else {
-			var err error
-			// Use DirectUpdateOrCreateConfig to avoid atomic file operations
-			// (temp file + rename) blocked by the tvOS sandbox.
-			cfg, err = profilemanager.DirectUpdateOrCreateConfig(profilemanager.ConfigInput{
-				ConfigPath:    c.cfgFile,
-				StateFilePath: c.stateFile,
-			})
-			if err != nil {
-				return "", fmt.Errorf("load config: %w", err)
-			}
+		var err error
+		cfg, err = c.loadConfig(profilemanager.ConfigInput{
+			ConfigPath:    c.cfgFile,
+			StateFilePath: c.stateFile,
+		})
+		if err != nil {
+			return "", fmt.Errorf("load config: %w", err)
 		}
 	}
 
@@ -376,40 +412,18 @@ func (c *Client) IsLoginRequiredCached() bool {
 }
 
 func (c *Client) IsLoginRequired() bool {
-	var ctx context.Context
 	//nolint
 	ctxWithValues := context.WithValue(context.Background(), system.DeviceNameCtxKey, c.deviceName)
 	//nolint
 	ctxWithValues = context.WithValue(ctxWithValues, system.OsNameCtxKey, c.osName)
 	//nolint
 	ctxWithValues = context.WithValue(ctxWithValues, system.OsVersionCtxKey, c.osVersion)
-	c.ctxCancelLock.Lock()
-	defer c.ctxCancelLock.Unlock()
-	ctx, c.ctxCancel = context.WithCancel(ctxWithValues)
+	ctx, cancel := context.WithCancel(ctxWithValues)
+	defer cancel()
 
-	var cfg *profilemanager.Config
-	var err error
-
-	// Use preloaded config if available (tvOS where file writes are blocked)
-	if c.preloadedConfig != nil {
-		log.Infof("IsLoginRequired: using preloaded config from memory")
-		cfg = c.preloadedConfig
-	} else {
-		log.Infof("IsLoginRequired: loading config from file")
-		// Use DirectUpdateOrCreateConfig to avoid atomic file operations (temp file + rename)
-		// which are blocked by the tvOS sandbox in App Group containers
-		cfg, err = profilemanager.DirectUpdateOrCreateConfig(profilemanager.ConfigInput{
-			ConfigPath: c.cfgFile,
-		})
-		if err != nil {
-			log.Errorf("IsLoginRequired: failed to load config: %v", err)
-			// If we can't load config, assume login is required
-			return true
-		}
-	}
-
-	if cfg == nil {
-		log.Errorf("IsLoginRequired: config is nil")
+	cfg, err := c.loadConfig(profilemanager.ConfigInput{ConfigPath: c.cfgFile})
+	if err != nil {
+		log.Errorf("IsLoginRequired: failed to load config: %v", err)
 		return true
 	}
 
@@ -433,17 +447,22 @@ func (c *Client) IsLoginRequired() bool {
 // loginForMobileAuthTimeout is the timeout for requesting auth info from the server
 const loginForMobileAuthTimeout = 30 * time.Second
 
+const stopRunWaitTimeout = 20 * time.Second
+
 func (c *Client) LoginForMobile() string {
-	var ctx context.Context
 	//nolint
 	ctxWithValues := context.WithValue(context.Background(), system.DeviceNameCtxKey, c.deviceName)
 	//nolint
 	ctxWithValues = context.WithValue(ctxWithValues, system.OsNameCtxKey, c.osName)
 	//nolint
 	ctxWithValues = context.WithValue(ctxWithValues, system.OsVersionCtxKey, c.osVersion)
-	c.ctxCancelLock.Lock()
-	defer c.ctxCancelLock.Unlock()
-	ctx, c.ctxCancel = context.WithCancel(ctxWithValues)
+	ctx, cancel := context.WithCancel(ctxWithValues)
+	loginDone := false
+	defer func() {
+		if !loginDone {
+			cancel()
+		}
+	}()
 
 	// Use DirectUpdateOrCreateConfig to avoid atomic file operations (temp file + rename)
 	// which are blocked by the tvOS sandbox in App Group containers
@@ -454,6 +473,7 @@ func (c *Client) LoginForMobile() string {
 		log.Errorf("LoginForMobile: failed to load config: %v", err)
 		return fmt.Sprintf("failed to load config: %v", err)
 	}
+	c.applyMDMOverlay(cfg)
 
 	oAuthFlow, err := auth.NewOAuthFlow(ctx, cfg, false, false, "", false)
 	if err != nil {
@@ -470,7 +490,9 @@ func (c *Client) LoginForMobile() string {
 	}
 
 	// This could cause a potential race condition with loading the extension which need to be handled on swift side
+	loginDone = true
 	go func() {
+		defer cancel()
 		tokenInfo, err := oAuthFlow.WaitToken(ctx, flowInfo)
 		if err != nil {
 			log.Errorf("LoginForMobile: WaitToken failed: %v", err)
@@ -487,18 +509,18 @@ func (c *Client) LoginForMobile() string {
 			log.Errorf("LoginForMobile: Login failed: %v", err)
 			return
 		}
-		c.loginComplete = true
+		c.loginComplete.Store(true)
 	}()
 
 	return flowInfo.VerificationURIComplete
 }
 
 func (c *Client) IsLoginComplete() bool {
-	return c.loginComplete
+	return c.loginComplete.Load()
 }
 
 func (c *Client) ClearLoginComplete() {
-	c.loginComplete = false
+	c.loginComplete.Store(false)
 }
 
 func (c *Client) GetRoutesSelectionDetails() (*RoutesSelectionDetails, error) {
@@ -718,13 +740,36 @@ func (c *Client) DeselectRoute(id string) error {
 	return nil
 }
 
-// setState stores the running engine state so DebugBundle can reuse the live
-// config and ConnectClient. It is cleared on Stop.
-func (c *Client) setState(cfg *profilemanager.Config, cc *internal.ConnectClient) {
+func (c *Client) startRun(cancel context.CancelFunc) (chan struct{}, error) {
 	c.stateMu.Lock()
 	defer c.stateMu.Unlock()
+
+	if c.runDone != nil {
+		return nil, errClientAlreadyRunning
+	}
+
+	done := make(chan struct{})
+	c.runDone = done
+	c.ctxCancel = cancel
+	return done, nil
+}
+
+func (c *Client) finishRun(done chan struct{}) {
+	c.stateMu.Lock()
+	c.connectClient = nil
+	c.config = nil
+	c.runDone = nil
+	c.ctxCancel = nil
+	c.stateMu.Unlock()
+
+	close(done)
+}
+
+func (c *Client) setState(cfg *profilemanager.Config, cc *internal.ConnectClient) {
+	c.stateMu.Lock()
 	c.config = cfg
 	c.connectClient = cc
+	c.stateMu.Unlock()
 }
 
 // stateSnapshot returns the current config and ConnectClient under the lock.