Revert "[client] Fix race between WG watcher initial handshake read and endpoint creation (#6626)"

Restore the pre-#6626 WG watcher shape
(single EnableWgWatcher reading the baseline inside the goroutine, no
PrepareInitialHandshake, no ctx-recheck before onDisconnected) to test whether
#6626 causes the missed handshake progression and subsequent reconnect loop.

Kept the WGW-DIAG logs.

This reverts commit 06839a4731.

client/internal/peer/conn.go

@@ -803,17 +803,15 @@ func (conn *Conn) isConnectedOnAllWay() (status guard.ConnStatus) {
 }
 
 func (conn *Conn) enableWgWatcherIfNeeded(enabledTime time.Time) {
-	if !conn.wgWatcher.PrepareInitialHandshake() {
-		return
+	if !conn.wgWatcher.IsEnabled() {
+		wgWatcherCtx, wgWatcherCancel := context.WithCancel(conn.ctx)
+		conn.wgWatcherCancel = wgWatcherCancel
+		conn.wgWatcherWg.Add(1)
+		go func() {
+			defer conn.wgWatcherWg.Done()
+			conn.wgWatcher.EnableWgWatcher(wgWatcherCtx, enabledTime, conn.onWGDisconnected, conn.onWGHandshakeSuccess)
+		}()
 	}
-
-	wgWatcherCtx, wgWatcherCancel := context.WithCancel(conn.ctx)
-	conn.wgWatcherCancel = wgWatcherCancel
-	conn.wgWatcherWg.Add(1)
-	go func() {
-		defer conn.wgWatcherWg.Done()
-		conn.wgWatcher.EnableWgWatcher(wgWatcherCtx, enabledTime, conn.onWGDisconnected, conn.onWGHandshakeSuccess)
-	}()
 }
 
 func (conn *Conn) disableWgWatcherIfNeeded() {
@@ -932,18 +930,22 @@ func (conn *Conn) AgentVersionString() string {
 
 func (conn *Conn) presharedKey(remoteRosenpassKey []byte) *wgtypes.Key {
 	if conn.config.RosenpassConfig.PubKey == nil {
+		conn.Log.Warnf("PSK-DIAG: rosenpass off -> static PSK (present=%v)", conn.config.WgConfig.PreSharedKey != nil)
 		return conn.config.WgConfig.PreSharedKey
 	}
 
 	if remoteRosenpassKey == nil && conn.config.RosenpassConfig.PermissiveMode {
+		conn.Log.Warnf("PSK-DIAG: rosenpass permissive + no remote RP key -> static PSK bridge (present=%v)", conn.config.WgConfig.PreSharedKey != nil)
 		return conn.config.WgConfig.PreSharedKey
 	}
 
 	// If Rosenpass has already set a PSK for this peer, return nil to prevent
 	// UpdatePeer from overwriting the Rosenpass-managed key.
 	if conn.rosenpassInitializedPresharedKeyValidator != nil && conn.rosenpassInitializedPresharedKeyValidator(conn.config.Key) {
+		conn.Log.Warnf("PSK-DIAG: rosenpass initialized -> returning nil (keep existing on-wire PSK; NOT re-set)")
 		return nil
 	}
+	conn.Log.Warnf("PSK-DIAG: rosenpass strict, not yet initialized -> seeding PSK (staticPresent=%v)", conn.config.WgConfig.PreSharedKey != nil)
 
 	// Use NetBird PSK as the seed for Rosenpass. This same PSK is passed to
 	// Rosenpass as PeerConfig.PresharedKey, ensuring the derived post-quantum

client/internal/peer/endpoint.go

@@ -38,6 +38,8 @@ func (e *EndpointUpdater) ConfigureWGEndpoint(addr *net.UDPAddr, presharedKey *w
 	e.mu.Lock()
 	defer e.mu.Unlock()
 
+	e.log.Warnf("PSK-DIAG: ConfigureWGEndpoint endpoint=%s psk_present=%v", addr, presharedKey != nil)
+
 	if e.initiator {
 		e.log.Debugf("configure up WireGuard as initiator")
 		return e.configureAsInitiator(addr, presharedKey)
@@ -51,6 +53,8 @@ func (e *EndpointUpdater) SwitchWGEndpoint(addr *net.UDPAddr, presharedKey *wgty
 	e.mu.Lock()
 	defer e.mu.Unlock()
 
+	e.log.Warnf("PSK-DIAG: SwitchWGEndpoint endpoint=%s psk_present=%v", addr, presharedKey != nil)
+
 	// prevent to run new update while cancel the previous update
 	e.waitForCloseTheDelayedUpdate()
 
@@ -69,6 +73,8 @@ func (e *EndpointUpdater) RemoveEndpointAddress() error {
 	e.mu.Lock()
 	defer e.mu.Unlock()
 
+	e.log.Warnf("PSK-DIAG: RemoveEndpointAddress -> peer re-added with only PublicKey+AllowedIPs; on-wire PSK is dropped")
+
 	e.waitForCloseTheDelayedUpdate()
 	return e.wgConfig.WgInterface.RemoveEndpointAddress(e.wgConfig.RemoteKey)
 }

client/internal/peer/wg_watcher.go

@@ -31,9 +31,7 @@ type WGWatcher struct {
 	stateDump     *stateDump
 
 	enabled   bool
-	muEnabled sync.Mutex
-	// initialHandshake is not thread-safe; never call PrepareInitialHandshake and EnableWgWatcher concurrently.
-	initialHandshake time.Time
+	muEnabled sync.RWMutex
 
 	resetCh chan struct{}
 }
@@ -48,38 +46,40 @@ func NewWGWatcher(log *log.Entry, wgIfaceStater WGInterfaceStater, peerKey strin
 	}
 }
 
-// PrepareInitialHandshake reserves the watcher and reads the peer's current WireGuard
-// handshake time. It must be called before the peer is (re)configured on the WireGuard
-// interface, so the captured baseline reflects the state prior to this connection attempt
-// instead of racing with that configuration. Returns ok=false if the watcher is already
-// running, in which case EnableWgWatcher must not be called.
-func (w *WGWatcher) PrepareInitialHandshake() (ok bool) {
+// EnableWgWatcher starts the WireGuard watcher. If it is already enabled, it will return immediately and do nothing.
+// The watcher runs until ctx is cancelled. Caller is responsible for context lifecycle management.
+// NOTE: reverted to the pre-#6626 shape for bisecting the NHN issue.
+func (w *WGWatcher) EnableWgWatcher(ctx context.Context, enabledTime time.Time, onDisconnectedFn func(), onHandshakeSuccessFn func(when time.Time)) {
 	w.muEnabled.Lock()
 	if w.enabled {
 		w.muEnabled.Unlock()
-		return false
+		return
 	}
 
 	w.log.Debugf("enable WireGuard watcher")
 	w.enabled = true
 	w.muEnabled.Unlock()
 
-	handshake, _ := w.wgState()
-	w.initialHandshake = handshake
-	return true
-}
+	initialHandshake, err := w.wgState()
+	if err != nil {
+		w.log.Warnf("failed to read initial wg stats: %v", err)
+	}
+	w.log.Warnf("PSK-DIAG: watcher baseline handshake=%v (zero=%v) [pre-6626 revert]", initialHandshake, initialHandshake.IsZero())
 
-// EnableWgWatcher runs the WireGuard watcher loop using the handshake baseline captured by
-// PrepareInitialHandshake. The watcher runs until ctx is cancelled. Caller is responsible
-// for context lifecycle management.
-func (w *WGWatcher) EnableWgWatcher(ctx context.Context, enabledTime time.Time, onDisconnectedFn func(), onHandshakeSuccessFn func(when time.Time)) {
-	w.periodicHandshakeCheck(ctx, onDisconnectedFn, onHandshakeSuccessFn, enabledTime, w.initialHandshake)
+	w.periodicHandshakeCheck(ctx, onDisconnectedFn, onHandshakeSuccessFn, enabledTime, initialHandshake)
 
 	w.muEnabled.Lock()
 	w.enabled = false
 	w.muEnabled.Unlock()
 }
 
+// IsEnabled returns true if the WireGuard watcher is currently enabled
+func (w *WGWatcher) IsEnabled() bool {
+	w.muEnabled.RLock()
+	defer w.muEnabled.RUnlock()
+	return w.enabled
+}
+
 // Reset signals the watcher that the WireGuard peer has been reset and a new
 // handshake is expected. This restarts the handshake timeout from scratch.
 func (w *WGWatcher) Reset() {
@@ -92,6 +92,7 @@ func (w *WGWatcher) Reset() {
 // wgStateCheck help to check the state of the WireGuard handshake and relay connection
 func (w *WGWatcher) periodicHandshakeCheck(ctx context.Context, onDisconnectedFn func(), onHandshakeSuccessFn func(when time.Time), enabledTime time.Time, initialHandshake time.Time) {
 	w.log.Infof("WireGuard watcher started")
+	w.log.Warnf("WGW-DIAG: watcher started id=%d baseline=%v (zero=%v) firstCheckIn=%v", enabledTime.UnixNano(), initialHandshake, initialHandshake.IsZero(), wgHandshakeOvertime)
 
 	timer := time.NewTimer(wgHandshakeOvertime)
 	defer timer.Stop()
@@ -101,18 +102,17 @@ func (w *WGWatcher) periodicHandshakeCheck(ctx context.Context, onDisconnectedFn
 	for {
 		select {
 		case <-timer.C:
+			w.log.Warnf("WGW-DIAG: check fire id=%d lastHandshake=%v", enabledTime.UnixNano(), lastHandshake)
 			handshake, ok := w.handshakeCheck(lastHandshake)
 			if !ok {
-				if ctx.Err() != nil {
-					return
-				}
+				w.log.Warnf("WGW-DIAG: check failed -> firing onDisconnected (TEARDOWN, pre-6626 no ctx-recheck) id=%d", enabledTime.UnixNano())
 				onDisconnectedFn()
 				return
 			}
 			if lastHandshake.IsZero() {
 				elapsed := calcElapsed(enabledTime, *handshake)
 				w.log.Infof("first wg handshake detected within: %.2fsec, (%s)", elapsed, handshake)
-				if onHandshakeSuccessFn != nil && ctx.Err() == nil {
+				if onHandshakeSuccessFn != nil {
 					onHandshakeSuccessFn(*handshake)
 				}
 			}
@@ -123,15 +123,15 @@ func (w *WGWatcher) periodicHandshakeCheck(ctx context.Context, onDisconnectedFn
 			timer.Reset(resetTime)
 			w.stateDump.WGcheckSuccess()
 
-			w.log.Debugf("WireGuard watcher reset timer: %v", resetTime)
+			w.log.Warnf("WGW-DIAG: check ok id=%d handshake=%v nextCheckIn=%v", enabledTime.UnixNano(), handshake, resetTime)
 		case <-w.resetCh:
-			w.log.Infof("WireGuard watcher received peer reset, restarting handshake timeout")
+			w.log.Warnf("WGW-DIAG: peer reset received, restarting timeout id=%d", enabledTime.UnixNano())
 			lastHandshake = time.Time{}
 			enabledTime = time.Now()
 			timer.Stop()
 			timer.Reset(wgHandshakeOvertime)
 		case <-ctx.Done():
-			w.log.Infof("WireGuard watcher stopped")
+			w.log.Warnf("WGW-DIAG: watcher stopped (ctx done) id=%d", enabledTime.UnixNano())
 			return
 		}
 	}

client/internal/peer/wg_watcher_test.go

@@ -7,7 +7,6 @@ import (
 	"time"
 
 	log "github.com/sirupsen/logrus"
-	"github.com/stretchr/testify/require"
 
 	"github.com/netbirdio/netbird/client/iface/configurer"
 )
@@ -35,9 +34,6 @@ func TestWGWatcher_EnableWgWatcher(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	ok := watcher.PrepareInitialHandshake()
-	require.True(t, ok, "watcher should not be enabled yet")
-
 	onDisconnected := make(chan struct{}, 1)
 	go watcher.EnableWgWatcher(ctx, time.Now(), func() {
 		mlog.Infof("onDisconnectedFn")
@@ -66,9 +62,6 @@ func TestWGWatcher_ReEnable(t *testing.T) {
 	watcher := NewWGWatcher(mlog, mocWgIface, "", newStateDump("peer", mlog, &Status{}))
 
 	ctx, cancel := context.WithCancel(context.Background())
-	ok := watcher.PrepareInitialHandshake()
-	require.True(t, ok, "watcher should not be enabled yet")
-
 	wg := &sync.WaitGroup{}
 	wg.Add(1)
 	go func() {
@@ -83,9 +76,6 @@ func TestWGWatcher_ReEnable(t *testing.T) {
 	ctx, cancel = context.WithCancel(context.Background())
 	defer cancel()
 
-	ok = watcher.PrepareInitialHandshake()
-	require.True(t, ok, "watcher should be re-enabled after the previous run stopped")
-
 	onDisconnected := make(chan struct{}, 1)
 	go watcher.EnableWgWatcher(ctx, time.Now(), func() {
 		onDisconnected <- struct{}{}

client/internal/rosenpass/manager.go

@@ -280,13 +280,15 @@ func (m *Manager) OnConnected(remoteWireGuardKey string, remoteRosenpassPubKey [
 	}
 
 	rpKeyHash := hashRosenpassKey(remoteRosenpassPubKey)
-	log.Debugf("received remote rosenpass key %s, my key %s", rpKeyHash, m.rpKeyHash)
+	initiator := bytes.Compare(m.spk, remoteRosenpassPubKey) == 1
+	log.Warnf("PSK-DIAG: remote peer %s advertises rosenpass (key %s, addr %s); starting RP negotiation as initiator=%v", remoteWireGuardKey, rpKeyHash, remoteRosenpassAddr, initiator)
 
 	err := m.addPeer(remoteRosenpassPubKey, remoteRosenpassAddr, wireGuardIP, remoteWireGuardKey)
 	if err != nil {
 		log.Errorf("failed to add rosenpass peer: %s", err)
 		return
 	}
+	log.Warnf("PSK-DIAG: rosenpass peer added for %s (RP negotiation started; watch for 'set PSK on WG' to confirm completion)", remoteWireGuardKey)
 }
 
 // IsPresharedKeyInitialized returns true if Rosenpass has completed a handshake

client/internal/rosenpass/netbird_handler.go

@@ -100,6 +100,7 @@ func (h *NetbirdHandler) outputKey(_ rp.KeyOutputReason, pid rp.PeerID, psk rp.K
 		log.Errorf("Failed to apply rosenpass key: %v", err)
 		return
 	}
+	log.Warnf("PSK-DIAG: rosenpass set PSK on WG for peer %s (updateOnly=%v)", peerKey, isInitialized)
 
 	// Mark peer as isInitialized after the successful first rotation
 	if !isInitialized {

client/internal/routemanager/client/client.go

@@ -231,11 +231,15 @@ func (w *Watcher) getBestRouteFromStatuses(routePeerStatuses map[route.ID]router
 	switch {
 	case chosen == "":
 		var peers []string
-		for _, r := range w.routes {
-			peers = append(peers, r.Peer)
+		for id, r := range w.routes {
+			st := "unknown(no status)"
+			if ps, ok := routePeerStatuses[id]; ok {
+				st = fmt.Sprintf("%s relayed=%v lat=%v", ps.status, ps.relayed, ps.latency)
+			}
+			peers = append(peers, fmt.Sprintf("%s{%s}", r.Peer, st))
 		}
 
-		log.Infof("network [%v] has not been assigned a routing peer as no peers from the list %s are currently available", w.handler, peers)
+		log.Warnf("ROUTE-DIAG: network [%v] no routing peer available yet (all candidates connecting/absent); candidates: %s", w.handler, peers)
 	case chosen != currID:
 		// we compare the current score + 10ms to the chosen score to avoid flapping between routes
 		if currScore != 0 && currScore+0.01 > chosenScore {
@@ -247,7 +251,7 @@ func (w *Watcher) getBestRouteFromStatuses(routePeerStatuses map[route.ID]router
 		if rt := w.routes[chosen]; rt != nil {
 			p = rt.Peer
 		}
-		log.Infof("New chosen route is %s with peer %s with score %f for network [%v]", chosen, p, chosenScore, w.handler)
+		log.Warnf("ROUTE-DIAG: new chosen route %s peer %s score %f for network [%v] (routing peer now usable)", chosen, p, chosenScore, w.handler)
 	}
 
 	return chosen, chosenStatus