From 06839a4731a64fea2ef7fb0d4857c356ca0eb60f Mon Sep 17 00:00:00 2001 From: Zoltan Papp Date: Wed, 1 Jul 2026 22:08:23 +0200 Subject: [PATCH] [client] Fix race between WG watcher initial handshake read and endpoint creation (#6626) * [client] Fix race between WG watcher initial handshake read and endpoint config The watcher's initial handshake read ran in a separate goroutine with no ordering guarantee relative to the WireGuard endpoint configuration, so it would sometimes race with the peer being added to the interface. Split enabling into a synchronous PrepareInitialHandshake, called before the endpoint is configured, and an EnableWgWatcher that only runs the monitoring loop, making the baseline read deterministic and keeping it correct for reconnects where the peer's WireGuard entry survives. * [client] Skip WG watcher disconnect callback when context is cancelled A superseded or cancelled watcher whose handshake-check timer fires before it observes ctx.Done() would still invoke onDisconnectedFn, tearing down a now-healthy connection. Re-check ctx before firing the disconnect and handshake-success callbacks and stand down silently if it was cancelled. --- client/internal/peer/conn.go | 18 ++++++----- client/internal/peer/wg_watcher.go | 41 ++++++++++++++----------- client/internal/peer/wg_watcher_test.go | 10 ++++++ 3 files changed, 43 insertions(+), 26 deletions(-) diff --git a/client/internal/peer/conn.go b/client/internal/peer/conn.go index 85e54ba5f..fb468696f 100644 --- a/client/internal/peer/conn.go +++ b/client/internal/peer/conn.go @@ -803,15 +803,17 @@ func (conn *Conn) isConnectedOnAllWay() (status guard.ConnStatus) { } func (conn *Conn) enableWgWatcherIfNeeded(enabledTime time.Time) { - 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) - }() + if !conn.wgWatcher.PrepareInitialHandshake() { + return } + + 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() { diff --git a/client/internal/peer/wg_watcher.go b/client/internal/peer/wg_watcher.go index 805a6f24a..4fc883d17 100644 --- a/client/internal/peer/wg_watcher.go +++ b/client/internal/peer/wg_watcher.go @@ -31,7 +31,9 @@ type WGWatcher struct { stateDump *stateDump enabled bool - muEnabled sync.RWMutex + muEnabled sync.Mutex + // initialHandshake is not thread-safe; never call PrepareInitialHandshake and EnableWgWatcher concurrently. + initialHandshake time.Time resetCh chan struct{} } @@ -46,38 +48,38 @@ func NewWGWatcher(log *log.Entry, wgIfaceStater WGInterfaceStater, peerKey strin } } -// 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. -func (w *WGWatcher) EnableWgWatcher(ctx context.Context, enabledTime time.Time, onDisconnectedFn func(), onHandshakeSuccessFn func(when time.Time)) { +// 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) { w.muEnabled.Lock() if w.enabled { w.muEnabled.Unlock() - return + return false } w.log.Debugf("enable WireGuard watcher") w.enabled = true w.muEnabled.Unlock() - initialHandshake, err := w.wgState() - if err != nil { - w.log.Warnf("failed to read initial wg stats: %v", err) - } + handshake, _ := w.wgState() + w.initialHandshake = handshake + return true +} - w.periodicHandshakeCheck(ctx, onDisconnectedFn, onHandshakeSuccessFn, enabledTime, initialHandshake) +// 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.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() { @@ -101,13 +103,16 @@ func (w *WGWatcher) periodicHandshakeCheck(ctx context.Context, onDisconnectedFn case <-timer.C: handshake, ok := w.handshakeCheck(lastHandshake) if !ok { + if ctx.Err() != nil { + return + } onDisconnectedFn() return } if lastHandshake.IsZero() { elapsed := calcElapsed(enabledTime, *handshake) w.log.Infof("first wg handshake detected within: %.2fsec, (%s)", elapsed, handshake) - if onHandshakeSuccessFn != nil { + if onHandshakeSuccessFn != nil && ctx.Err() == nil { onHandshakeSuccessFn(*handshake) } } diff --git a/client/internal/peer/wg_watcher_test.go b/client/internal/peer/wg_watcher_test.go index 3ce91cd46..634d7974f 100644 --- a/client/internal/peer/wg_watcher_test.go +++ b/client/internal/peer/wg_watcher_test.go @@ -7,6 +7,7 @@ import ( "time" log "github.com/sirupsen/logrus" + "github.com/stretchr/testify/require" "github.com/netbirdio/netbird/client/iface/configurer" ) @@ -34,6 +35,9 @@ 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") @@ -62,6 +66,9 @@ 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() { @@ -76,6 +83,9 @@ 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{}{}