mirror of
https://github.com/netbirdio/netbird.git
synced 2026-07-18 04:29:54 +00:00
* [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.
103 lines
2.4 KiB
Go
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")
|
|
}
|
|
}
|