Merge branch 'main' into embedded-vnc

client/internal/dns/host_windows.go

@@ -6,8 +6,10 @@ import (
 	"fmt"
 	"io"
 	"net/netip"
+	"os"
 	"os/exec"
 	"slices"
+	"strconv"
 	"strings"
 	"syscall"
 	"time"
@@ -34,10 +36,16 @@ var (
 // Registry locations of the host DNS configuration this package programs,
 // exported so a diagnostic reader reports the same locations that are written.
 const (
-	// NRPTKeyPrefix starts the name of every NRPT rule key this client creates.
-	// Older versions used different layouts under the same prefix: a single
-	// unsuffixed key, then one key per domain, now one key per batch of domains.
-	NRPTKeyPrefix = "NetBird-Match"
+	// NRPTKeyPrefix starts the name of every NRPT rule key this client creates:
+	// the match rules, the catch-all, and the .local exemption. Cleanup
+	// enumerates by this prefix, so a new kind of rule is removed by existing
+	// code as long as its key starts here.
+	NRPTKeyPrefix = "NetBird-"
+
+	// nrptMatchKeyName names the match-domain rules. Older versions used
+	// different layouts under the same name: a single unsuffixed key, then one
+	// key per domain, now one key per batch of domains.
+	nrptMatchKeyName = NRPTKeyPrefix + "Match"
 
 	// DNSPolicyConfigRoot holds the NRPT rules of the local policy store.
 	DNSPolicyConfigRoot = `SYSTEM\CurrentControlSet\Services\Dnscache\Parameters\DnsPolicyConfig`
@@ -53,8 +61,24 @@ const (
 )
 
 const (
-	dnsPolicyConfigMatchPath    = DNSPolicyConfigRoot + `\` + NRPTKeyPrefix
-	gpoDnsPolicyConfigMatchPath = GPODNSPolicyConfigRoot + `\` + NRPTKeyPrefix
+	dnsPolicyConfigMatchPath    = DNSPolicyConfigRoot + `\` + nrptMatchKeyName
+	gpoDnsPolicyConfigMatchPath = GPODNSPolicyConfigRoot + `\` + nrptMatchKeyName
+
+	dnsPolicyConfigExemptLocalPath    = DNSPolicyConfigRoot + `\` + NRPTKeyPrefix + `ExemptLocal`
+	gpoDnsPolicyConfigExemptLocalPath = GPODNSPolicyConfigRoot + `\` + NRPTKeyPrefix + `ExemptLocal`
+
+	nrptCatchAllNamespace = "."
+	// nrptLocalNamespace is reserved for multicast DNS by RFC 6762: a unicast
+	// resolver must not answer for it. The catch-all rule would hand it to us
+	// anyway, so it gets an exemption rule of its own.
+	nrptLocalNamespace = ".local"
+
+	// envLegacyDNSResolution restores the pre-catch-all behaviour: the adapter's
+	// NameServer alone, leaving the OS free to query other adapters' resolvers in
+	// parallel. An escape hatch for setups that depend on a resolver of theirs
+	// still being reachable while connected, at the cost of the leak and of the
+	// race the catch-all rule exists to close.
+	envLegacyDNSResolution = "NB_USE_LEGACY_DNS_RESOLUTION"
 
 	dnsPolicyConfigVersionKey           = "Version"
 	dnsPolicyConfigVersionValue         = 2
@@ -293,6 +317,13 @@ func (r *registryConfigurator) disableWINSForInterface() error {
 }
 
 func (r *registryConfigurator) applyDNSConfig(config HostDNSConfig, stateManager *statemanager.Manager) error {
+	// Clear every rule the previous apply installed before installing any new
+	// one, including a leftover catch-all: removal is unconditional so a rule
+	// from an earlier run cannot survive into a config that no longer wants it.
+	if err := r.removeDNSMatchPolicies(); err != nil {
+		log.Errorf("cleanup old dns match policies: %s", err)
+	}
+
 	if config.RouteAll {
 		if err := r.addDNSSetupForAll(config.ServerIP); err != nil {
 			return fmt.Errorf("add dns setup: %w", err)
@@ -318,8 +349,22 @@ func (r *registryConfigurator) applyDNSConfig(config HostDNSConfig, stateManager
 		matchDomains = append(matchDomains, "."+strings.TrimSuffix(dConf.Domain, "."))
 	}
 
-	if err := r.removeDNSMatchPolicies(); err != nil {
-		log.Errorf("cleanup old dns match policies: %s", err)
+	// The root namespace is a match domain like any other: it just happens to
+	// match every name. Without it the adapter's NameServer only adds one more
+	// resolver to the set Windows queries in parallel, keeping whichever answer
+	// comes back first — which leaks every query to the local network and lets a
+	// resolver other than ours answer for a name we are authoritative for.
+	if config.RouteAll {
+		if parseBoolEnv(envLegacyDNSResolution) {
+			log.Infof("%s is set, leaving DNS resolution shared with the other adapters' resolvers instead of forcing it through %s", envLegacyDNSResolution, config.ServerIP)
+		} else {
+			matchDomains = append(matchDomains, nrptCatchAllNamespace)
+			log.Infof("routing every namespace through %s: DNS resolution is now exclusive to NetBird", config.ServerIP)
+
+			if err := r.addDNSExemptLocalPolicy(); err != nil {
+				return fmt.Errorf("add dns exempt policy: %w", err)
+			}
+		}
 	}
 
 	if len(matchDomains) != 0 {
@@ -397,6 +442,42 @@ func (r *registryConfigurator) addDNSMatchPolicy(domains []string, ip netip.Addr
 	return nil
 }
 
+// addDNSExemptLocalPolicy carves .local back out of the catch-all. RFC 6762
+// reserves it for multicast DNS, so forwarding those names to a unicast
+// upstream answers NXDOMAIN for hosts that do exist - printers, NAS boxes, and
+// anything else announcing itself on the link - and the answer is authoritative
+// enough that Windows stops looking. A rule naming the namespace with no
+// servers hands it back to the DNS client untouched. A more specific rule still
+// wins, so a match domain under .local keeps going through us.
+func (r *registryConfigurator) addDNSExemptLocalPolicy() error {
+	var noServers netip.Addr
+
+	if err := r.configureDNSPolicy(dnsPolicyConfigExemptLocalPath, []string{nrptLocalNamespace}, noServers); err != nil {
+		return fmt.Errorf("configure exempt policy for %s: %w", nrptLocalNamespace, err)
+	}
+
+	if r.gpo {
+		if err := r.configureDNSPolicy(gpoDnsPolicyConfigExemptLocalPath, []string{nrptLocalNamespace}, noServers); err != nil {
+			return fmt.Errorf("configure gpo exempt policy for %s: %w", nrptLocalNamespace, err)
+		}
+		if err := refreshGroupPolicy(); err != nil {
+			log.Warnf("failed to refresh group policy: %v", err)
+		}
+	}
+
+	log.Infof("added NRPT exemption for %s, leaving it to the OS resolver", nrptLocalNamespace)
+	return nil
+}
+
+// configureDNSPolicy writes one NRPT rule. An invalid ip writes an exemption
+// rule: the namespace with an empty server list, which tells the DNS client to
+// resolve those names the way it would without any rule at all.
+//
+// The empty string is the whole difference, and it has to be written: dropping
+// the value and clearing ConfigOptions instead produces a rule Windows treats
+// as a no-op, keeps out of Get-DnsClientNrptPolicy -Effective, and ignores in
+// favour of the catch-all. 0x8 says the server list is the meaningful part of
+// the rule, and an empty list then means "no server, resolve normally".
 func (r *registryConfigurator) configureDNSPolicy(policyPath string, domains []string, ip netip.Addr) error {
 	if err := removeRegistryKeyFromDNSPolicyConfig(policyPath); err != nil {
 		return fmt.Errorf("remove existing dns policy: %w", err)
@@ -416,7 +497,11 @@ func (r *registryConfigurator) configureDNSPolicy(policyPath string, domains []s
 		return fmt.Errorf("set %s: %w", dnsPolicyConfigNameKey, err)
 	}
 
-	if err := regKey.SetStringValue(dnsPolicyConfigGenericDNSServersKey, ip.String()); err != nil {
+	var servers string
+	if ip.IsValid() {
+		servers = ip.String()
+	}
+	if err := regKey.SetStringValue(dnsPolicyConfigGenericDNSServersKey, servers); err != nil {
 		return fmt.Errorf("set %s: %w", dnsPolicyConfigGenericDNSServersKey, err)
 	}
 
@@ -514,8 +599,11 @@ func (r *registryConfigurator) getInterfaceRegistryKey() (registry.Key, error) {
 }
 
 func (r *registryConfigurator) restoreHostDNS() error {
+	// Propagated, unlike in applyDNSConfig: there we are about to write fresh
+	// rules over whatever survived, here we are leaving, and a rule left behind
+	// keeps sending every query to an address that is about to disappear.
 	if err := r.removeDNSMatchPolicies(); err != nil {
-		log.Errorf("remove dns match policies: %s", err)
+		return fmt.Errorf("remove dns match policies: %w", err)
 	}
 
 	if err := r.deleteInterfaceRegistryKeyProperty(interfaceConfigSearchListKey); err != nil {
@@ -598,9 +686,17 @@ func listNRPTRuleKeys(root string) ([]string, error) {
 
 func removeRegistryKeyFromDNSPolicyConfig(regKeyPath string) error {
 	k, err := registry.OpenKey(registry.LOCAL_MACHINE, regKeyPath, registry.QUERY_VALUE)
-	if err != nil {
-		log.Debugf("failed to open HKEY_LOCAL_MACHINE\\%s: %v", regKeyPath, err)
+	switch {
+	case errors.Is(err, registry.ErrNotExist), errors.Is(err, syscall.ERROR_PATH_NOT_FOUND):
+		// nothing to remove, which is the normal case for a rule this config
+		// never installed
+		log.Debugf("HKEY_LOCAL_MACHINE\\%s does not exist", regKeyPath)
 		return nil
+	case err != nil:
+		// anything else has to reach the caller: reporting success here would
+		// leave the rule in force while claiming it was removed, which is how a
+		// stale rule outlives the interface it points at
+		return fmt.Errorf("open HKEY_LOCAL_MACHINE\\%s: %w", regKeyPath, err)
 	}
 
 	closer(k)
@@ -636,6 +732,20 @@ func refreshGroupPolicy() error {
 	return nil
 }
 
+func parseBoolEnv(key string) bool {
+	val := os.Getenv(key)
+	if val == "" {
+		return false
+	}
+
+	parsed, err := strconv.ParseBool(val)
+	if err != nil {
+		log.Warnf("failed to parse %s=%q: %v", key, val, err)
+		return false
+	}
+	return parsed
+}
+
 func closer(closer io.Closer) {
 	if err := closer.Close(); err != nil {
 		log.Errorf("failed to close: %s", err)

client/internal/dns/host_windows_test.go

@@ -94,6 +94,145 @@ func TestNRPTEntriesCleanupOnConfigChange(t *testing.T) {
 	assert.False(t, exists, "NRPT rule 2 should NOT exist after reducing to 75 domains")
 }
 
+// TestNRPTCatchAllRule verifies that RouteAll adds the root namespace to the
+// match rule instead of a rule of its own, that .local is carved back out with
+// an empty server list, and that both go away when RouteAll is cleared or the
+// host DNS is restored.
+func TestNRPTCatchAllRule(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping registry integration test in short mode")
+	}
+
+	defer cleanupRegistryKeys(t)
+	cleanupRegistryKeys(t)
+
+	testIP := netip.MustParseAddr("100.64.0.1")
+	testGUID := "{12345678-1234-1234-1234-123456789ABC}"
+	interfacePath := InterfaceConfigPath + `\` + testGUID
+	testKey, _, err := registry.CreateKey(registry.LOCAL_MACHINE, interfacePath, registry.SET_VALUE)
+	require.NoError(t, err, "Should create test interface registry key")
+	require.NoError(t, testKey.Close(), "close test interface registry key")
+	defer func() {
+		assert.NoError(t, registry.DeleteKey(registry.LOCAL_MACHINE, interfacePath), "delete test interface registry key")
+	}()
+
+	cfg := &registryConfigurator{guid: testGUID}
+
+	matchOnly := HostDNSConfig{
+		ServerIP: testIP,
+		Domains:  []DomainConfig{{Domain: "example.com", MatchOnly: true}},
+	}
+	primary := HostDNSConfig{
+		ServerIP: testIP,
+		RouteAll: true,
+		Domains:  []DomainConfig{{Domain: "example.com", MatchOnly: true}},
+	}
+	firstRule := fmt.Sprintf("%s-0", dnsPolicyConfigMatchPath)
+
+	// The root namespace is not a rule of its own: it rides in the match rule,
+	// which is the point of it not being a special case.
+	require.NoError(t, cfg.applyDNSConfig(matchOnly, nil))
+	names := ruleNamespaces(t, firstRule)
+	assert.Contains(t, names, ".example.com")
+	assert.NotContains(t, names, nrptCatchAllNamespace, "a match-only config must not claim every namespace")
+
+	require.NoError(t, cfg.applyDNSConfig(primary, nil))
+	names = ruleNamespaces(t, firstRule)
+	assert.Contains(t, names, ".example.com")
+	assert.Contains(t, names, nrptCatchAllNamespace, "RouteAll should add the root namespace to the match rule")
+
+	k, err := registry.OpenKey(registry.LOCAL_MACHINE, firstRule, registry.QUERY_VALUE)
+	require.NoError(t, err)
+	servers, _, err := k.GetStringValue(dnsPolicyConfigGenericDNSServersKey)
+	require.NoError(t, err)
+	assert.Equal(t, testIP.String(), servers, "every namespace in the rule resolves through our resolver")
+	require.NoError(t, k.Close(), "close match rule key")
+
+	// .local is carved back out: RFC 6762 reserves it for mDNS, so it needs a
+	// rule of its own — it is the one rule with a different server list.
+	ek, err := registry.OpenKey(registry.LOCAL_MACHINE, dnsPolicyConfigExemptLocalPath, registry.QUERY_VALUE)
+	require.NoError(t, err, "exemption rule should exist once the root namespace is claimed")
+
+	exemptNames, _, err := ek.GetStringsValue(dnsPolicyConfigNameKey)
+	require.NoError(t, err)
+	assert.Equal(t, []string{nrptLocalNamespace}, exemptNames, "the exemption should name only the mDNS namespace")
+
+	exemptServers, _, err := ek.GetStringValue(dnsPolicyConfigGenericDNSServersKey)
+	require.NoError(t, err, "the value has to be present, empty: without it Windows drops the rule")
+	assert.Empty(t, exemptServers, "an exemption rule lists no servers")
+
+	exemptOpts, _, err := ek.GetIntegerValue(dnsPolicyConfigConfigOptionsKey)
+	require.NoError(t, err)
+	assert.EqualValues(t, dnsPolicyConfigConfigOptionsValue, exemptOpts, "same options as a normal rule; the empty server list is what makes it an exemption")
+	require.NoError(t, ek.Close(), "close exemption rule key")
+
+	require.NoError(t, cfg.applyDNSConfig(matchOnly, nil))
+	names = ruleNamespaces(t, firstRule)
+	assert.NotContains(t, names, nrptCatchAllNamespace, "clearing RouteAll should drop the root namespace")
+
+	exists, err := registryKeyExists(dnsPolicyConfigExemptLocalPath)
+	require.NoError(t, err)
+	assert.False(t, exists, "exemption rule should go with the namespace it carves out of")
+
+	require.NoError(t, cfg.applyDNSConfig(primary, nil))
+	require.NoError(t, cfg.restoreHostDNS())
+	exists, err = registryKeyExists(firstRule)
+	require.NoError(t, err)
+	assert.False(t, exists, "restore should leave no rule behind")
+}
+
+// ruleNamespaces returns the namespaces an NRPT rule key claims.
+func ruleNamespaces(t *testing.T, path string) []string {
+	t.Helper()
+	k, err := registry.OpenKey(registry.LOCAL_MACHINE, path, registry.QUERY_VALUE)
+	require.NoError(t, err, "rule key %s should exist", path)
+	defer k.Close()
+
+	names, _, err := k.GetStringsValue(dnsPolicyConfigNameKey)
+	require.NoError(t, err)
+	return names
+}
+
+// TestNRPTCatchAllRuleLegacyEnv verifies that NB_USE_LEGACY_DNS_RESOLUTION
+// leaves the root namespace unclaimed, so no rule is written for a RouteAll
+// config that carries no match domains.
+func TestNRPTCatchAllRuleLegacyEnv(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping registry integration test in short mode")
+	}
+
+	defer cleanupRegistryKeys(t)
+	cleanupRegistryKeys(t)
+
+	t.Setenv(envLegacyDNSResolution, "true")
+
+	testGUID := "{12345678-1234-1234-1234-123456789ABC}"
+	interfacePath := InterfaceConfigPath + `\` + testGUID
+	testKey, _, err := registry.CreateKey(registry.LOCAL_MACHINE, interfacePath, registry.SET_VALUE)
+	require.NoError(t, err, "Should create test interface registry key")
+	require.NoError(t, testKey.Close(), "close test interface registry key")
+	defer func() {
+		assert.NoError(t, registry.DeleteKey(registry.LOCAL_MACHINE, interfacePath), "delete test interface registry key")
+	}()
+
+	cfg := &registryConfigurator{guid: testGUID}
+	config := HostDNSConfig{
+		ServerIP: netip.MustParseAddr("100.64.0.1"),
+		RouteAll: true,
+	}
+
+	require.NoError(t, cfg.applyDNSConfig(config, nil))
+
+	// RouteAll with no match domains and the switch set leaves nothing to write.
+	exists, err := registryKeyExists(fmt.Sprintf("%s-0", dnsPolicyConfigMatchPath))
+	require.NoError(t, err)
+	assert.False(t, exists, "no rule should be written when the legacy env var is set")
+
+	exists, err = registryKeyExists(dnsPolicyConfigExemptLocalPath)
+	require.NoError(t, err)
+	assert.False(t, exists, "no exemption without a claimed root namespace")
+}
+
 func registryKeyExists(path string) (bool, error) {
 	k, err := registry.OpenKey(registry.LOCAL_MACHINE, path, registry.QUERY_VALUE)
 	if err != nil {

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,15 +78,13 @@ 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.
@@ -90,9 +92,16 @@ type Client struct {
 	// preloadedConfig holds config loaded from JSON (used on tvOS where file writes are blocked)
 	preloadedConfig *profilemanager.Config
 
+	// 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,7 +116,6 @@ 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),
@@ -156,17 +164,21 @@ func (c *Client) Run(fd int32, interfaceName string, envList *EnvList) error {
 	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 +227,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.
@@ -376,16 +412,14 @@ 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
@@ -433,17 +467,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
@@ -470,7 +509,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 +528,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 +759,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.