Files
netbird/client/internal/peer/wg_watcher_test.go
Zoltan Papp 06839a4731 [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.
2026-07-01 22:08:23 +02:00

103 lines
2.4 KiB
Go

package peer
import (
"context"
"sync"
"testing"
"time"
log "github.com/sirupsen/logrus"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/client/iface/configurer"
)
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
}
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, "", newStateDump("peer", mlog, &Status{}))
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")
onDisconnected <- struct{}{}
}, func(when time.Time) {
mlog.Infof("onHandshakeSuccess: %v", when)
})
// 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, "", 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() {
defer wg.Done()
watcher.EnableWgWatcher(ctx, time.Now(), func() {}, func(when time.Time) {})
}()
cancel()
wg.Wait()
// Re-enable with a new context
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{}{}
}, func(when time.Time) {})
time.Sleep(2 * time.Second)
mocWgIface.disconnect()
select {
case <-onDisconnected:
case <-time.After(10 * time.Second):
t.Errorf("timeout")
}
}