[client] Extract WGWatcher into its own wg_watcher package

Move the WireGuard handshake watcher out of the peer package into client/internal/peer/wg_watcher. The test stays an internal test since it drives the unexported checkPeriod. Update conn.go to reference the watcher through the package.
2026-10-24 22:39:09 +02:00 · 2026-08-05 16:14:11 +02:00
parent a84302ac71
commit a3aabda1e9
3 changed files with 5 additions and 4 deletions
@@ -0,0 +1,183 @@
+package wg_watcher
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/client/iface/configurer"
+	"github.com/netbirdio/netbird/client/internal/peer/state_dump"
+)
+
+const (
+	wgHandshakePeriod = 3 * time.Minute
+)
+
+var (
+	wgHandshakeOvertime = 30 * time.Second // allowed delay in network
+	checkPeriod         = wgHandshakePeriod + wgHandshakeOvertime
+)
+
+type WGInterfaceStater interface {
+	GetStats() (map[string]configurer.WGStats, error)
+}
+
+// WGWatcher is single-shot: one instance per connection attempt, run once, then discarded.
+// Lifecycle is owned by Conn under conn.mu, so it keeps no "enabled" state to go stale.
+type WGWatcher struct {
+	log           *log.Entry
+	wgIfaceStater WGInterfaceStater
+	peerKey       string
+	stateDump     *state_dump.StateDump
+
+	// initialHandshake is not thread-safe; never call PrepareInitialHandshake and EnableWgWatcher concurrently.
+	initialHandshake time.Time
+
+	resetCh chan struct{}
+}
+
+func NewWGWatcher(log *log.Entry, wgIfaceStater WGInterfaceStater, peerKey string, stateDump *state_dump.StateDump) *WGWatcher {
+	return &WGWatcher{
+		log:           log,
+		wgIfaceStater: wgIfaceStater,
+		peerKey:       peerKey,
+		stateDump:     stateDump,
+		resetCh:       make(chan struct{}, 1),
+	}
+}
+
+// PrepareInitialHandshake 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.
+func (w *WGWatcher) PrepareInitialHandshake() {
+	w.log.Debugf("enable WireGuard watcher")
+	handshake, _ := w.wgState()
+	w.initialHandshake = handshake
+}
+
+// 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. onHandshakeSuccessFn is called only for the first
+// handshake observed by this run, onCheckSuccessFn for every check that observed a fresh
+// handshake, including the first.
+func (w *WGWatcher) EnableWgWatcher(ctx context.Context, enabledTime time.Time, onDisconnectedFn func(), onHandshakeSuccessFn func(when time.Time), onCheckSuccessFn func()) {
+	w.periodicHandshakeCheck(ctx, onDisconnectedFn, onHandshakeSuccessFn, onCheckSuccessFn, enabledTime, w.initialHandshake)
+}
+
+// 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() {
+	select {
+	case w.resetCh <- struct{}{}:
+	default:
+	}
+}
+
+// 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), onCheckSuccessFn func(), enabledTime time.Time, initialHandshake time.Time) {
+	w.log.Infof("WireGuard watcher started")
+
+	timer := time.NewTimer(wgHandshakeOvertime)
+	defer timer.Stop()
+
+	lastHandshake := initialHandshake
+
+	for {
+		select {
+		case <-timer.C:
+			handshake, ok := w.handshakeCheck(lastHandshake)
+			if !ok {
+				// early ctx cancel check return
+				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 && ctx.Err() == nil {
+					onHandshakeSuccessFn(*handshake)
+				}
+			}
+
+			if onCheckSuccessFn != nil && ctx.Err() == nil {
+				onCheckSuccessFn()
+			}
+
+			lastHandshake = *handshake
+
+			resetTime := time.Until(handshake.Add(checkPeriod))
+			timer.Reset(resetTime)
+			w.stateDump.WGcheckSuccess()
+
+			w.log.Debugf("WireGuard watcher reset timer: %v", resetTime)
+		case <-w.resetCh:
+			w.log.Infof("WireGuard watcher received peer reset, restarting handshake timeout")
+			lastHandshake = time.Time{}
+			enabledTime = time.Now()
+			timer.Stop()
+			timer.Reset(wgHandshakeOvertime)
+		case <-ctx.Done():
+			w.log.Infof("WireGuard watcher stopped")
+			return
+		}
+	}
+}
+
+// handshakeCheck checks the WireGuard handshake and return the new handshake time if it is different from the previous one
+func (w *WGWatcher) handshakeCheck(lastHandshake time.Time) (*time.Time, bool) {
+	handshake, err := w.wgState()
+	if err != nil {
+		w.log.Errorf("failed to read wg stats: %v", err)
+		return nil, false
+	}
+
+	w.log.Tracef("previous handshake, handshake: %v, %v", lastHandshake, handshake)
+
+	// the current known handshake did not change
+	if handshake.Equal(lastHandshake) {
+		w.log.Warnf("WireGuard handshake not updated: %v", handshake)
+		return nil, false
+	}
+
+	// in case if the machine is suspended, the handshake time will be in the past
+	if handshake.Add(checkPeriod).Before(time.Now()) {
+		w.log.Warnf("WireGuard handshake timed out: %v", handshake)
+		return nil, false
+	}
+
+	// error handling for handshake time in the future
+	if handshake.After(time.Now()) {
+		w.log.Warnf("WireGuard handshake is in the future: %v", handshake)
+		return nil, false
+	}
+
+	return &handshake, true
+}
+
+func (w *WGWatcher) wgState() (time.Time, error) {
+	wgStates, err := w.wgIfaceStater.GetStats()
+	if err != nil {
+		return time.Time{}, err
+	}
+	wgState, ok := wgStates[w.peerKey]
+	if !ok {
+		return time.Time{}, fmt.Errorf("peer %s not found in WireGuard endpoints", w.peerKey)
+	}
+	return wgState.LastHandshake, nil
+}
+
+// calcElapsed calculates elapsed time since watcher was enabled.
+// The watcher started after the wg configuration happens, because of this need to normalise the negative value
+func calcElapsed(enabledTime, handshake time.Time) float64 {
+	elapsed := handshake.Sub(enabledTime).Seconds()
+	if elapsed < 0 {
+		elapsed = 0
+	}
+	return elapsed
+}
@@ -0,0 +1,175 @@
+package wg_watcher
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/client/iface/configurer"
+	"github.com/netbirdio/netbird/client/internal/peer/state_dump"
+	"github.com/netbirdio/netbird/client/internal/peer/status"
+)
+
+type MocWgIface struct {
+	stop bool
+}
+
+func (m *MocWgIface) GetStats() (map[string]configurer.WGStats, error) {
+	return map[string]configurer.WGStats{}, nil
+}
+
+func (m *MocWgIface) disconnect() {
+	m.stop = true
+}
+
+type mockHandshakeStats struct {
+	mu        sync.Mutex
+	handshake time.Time
+}
+
+func (m *mockHandshakeStats) GetStats() (map[string]configurer.WGStats, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return map[string]configurer.WGStats{"": {LastHandshake: m.handshake}}, nil
+}
+
+func (m *mockHandshakeStats) advance() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.handshake = time.Now()
+}
+
+// TestWGWatcher_CheckSuccessCallback: onCheckSuccessFn must fire for a fresh
+// handshake even when the watcher started with an existing handshake baseline,
+// the case where onHandshakeSuccessFn stays silent.
+func TestWGWatcher_CheckSuccessCallback(t *testing.T) {
+	// checkPeriod bounds how stale a handshake may be before the watcher treats it
+	// as a suspended-machine timeout. The first check fires after wgHandshakeOvertime,
+	// so keep checkPeriod well above any scheduling jitter to avoid a false timeout
+	// converting the expected success into a disconnect on a loaded runner.
+	checkPeriod = 1 * time.Minute
+	wgHandshakeOvertime = 1 * time.Second
+
+	mlog := log.WithField("peer", "tet")
+	// Use an old baseline so advance() yields a strictly newer handshake even on
+	// platforms with coarse clock resolution (Windows), where two time.Now() calls
+	// microseconds apart can return the same instant and read as a timed-out handshake.
+	stats := &mockHandshakeStats{handshake: time.Now().Add(-time.Hour)}
+	watcher := NewWGWatcher(mlog, stats, "", state_dump.NewStateDump("peer", mlog, &status.Recorder{}))
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	watcher.PrepareInitialHandshake()
+
+	firstHandshake := make(chan struct{}, 1)
+	checkSuccess := make(chan struct{}, 1)
+	watcherDone := make(chan struct{})
+	go func() {
+		defer close(watcherDone)
+		watcher.EnableWgWatcher(ctx, time.Now(), func() {}, func(when time.Time) {
+			firstHandshake <- struct{}{}
+		}, func() {
+			select {
+			case checkSuccess <- struct{}{}:
+			default:
+			}
+		})
+	}()
+
+	stats.advance()
+
+	select {
+	case <-checkSuccess:
+	case <-time.After(10 * time.Second):
+		t.Errorf("timeout waiting for check success callback")
+	}
+
+	select {
+	case <-firstHandshake:
+		t.Errorf("first-handshake callback must not fire for a non-zero baseline")
+	default:
+	}
+
+	// Wait for the watcher goroutine to exit so it cannot race with other
+	// tests mutating the package-level check timing variables.
+	cancel()
+	<-watcherDone
+}
+
+func TestWGWatcher_EnableWgWatcher(t *testing.T) {
+	checkPeriod = 5 * time.Second
+	wgHandshakeOvertime = 1 * time.Second
+
+	mlog := log.WithField("peer", "tet")
+	mocWgIface := &MocWgIface{}
+	watcher := NewWGWatcher(mlog, mocWgIface, "", state_dump.NewStateDump("peer", mlog, &status.Recorder{}))
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	watcher.PrepareInitialHandshake()
+
+	onDisconnected := make(chan struct{}, 1)
+	go watcher.EnableWgWatcher(ctx, time.Now(), func() {
+		mlog.Infof("onDisconnectedFn")
+		onDisconnected <- struct{}{}
+	}, func(when time.Time) {
+		mlog.Infof("onHandshakeSuccess: %v", when)
+	}, nil)
+
+	// wait for initial reading
+	time.Sleep(2 * time.Second)
+	mocWgIface.disconnect()
+
+	select {
+	case <-onDisconnected:
+	case <-time.After(10 * time.Second):
+		t.Errorf("timeout")
+	}
+}
+
+func TestWGWatcher_ReEnable(t *testing.T) {
+	checkPeriod = 5 * time.Second
+	wgHandshakeOvertime = 1 * time.Second
+
+	mlog := log.WithField("peer", "tet")
+	mocWgIface := &MocWgIface{}
+	watcher := NewWGWatcher(mlog, mocWgIface, "", state_dump.NewStateDump("peer", mlog, &status.Recorder{}))
+
+	ctx, cancel := context.WithCancel(context.Background())
+	watcher.PrepareInitialHandshake()
+
+	wg := &sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		watcher.EnableWgWatcher(ctx, time.Now(), func() {}, func(when time.Time) {}, nil)
+	}()
+	cancel()
+
+	wg.Wait()
+
+	// Re-enable with a new context
+	ctx, cancel = context.WithCancel(context.Background())
+	defer cancel()
+
+	watcher.PrepareInitialHandshake()
+
+	onDisconnected := make(chan struct{}, 1)
+	go watcher.EnableWgWatcher(ctx, time.Now(), func() {
+		onDisconnected <- struct{}{}
+	}, func(when time.Time) {}, nil)
+
+	time.Sleep(2 * time.Second)
+	mocWgIface.disconnect()
+
+	select {
+	case <-onDisconnected:
+	case <-time.After(10 * time.Second):
+		t.Errorf("timeout")
+	}
+}