mirror of
https://github.com/netbirdio/netbird.git
synced 2026-07-10 00:29:56 +00:00
Compare commits
3 Commits
update-pro
...
rosenpass-
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
90eefcc7d4 | ||
|
|
7cd5c1732b | ||
|
|
816d80602f |
@@ -522,7 +522,7 @@ func (e *Engine) Start(netbirdConfig *mgmProto.NetbirdConfig, mgmtURL *url.URL)
|
||||
} else {
|
||||
log.Infof("running rosenpass in strict mode")
|
||||
}
|
||||
e.rpManager, err = rosenpass.NewManager(e.config.PreSharedKey, e.config.WgIfaceName)
|
||||
e.rpManager, err = rosenpass.NewManager(e.config.PreSharedKey, e.config.WgIfaceName, publicKey)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create rosenpass manager: %w", err)
|
||||
}
|
||||
|
||||
@@ -30,6 +30,11 @@ import (
|
||||
relayClient "github.com/netbirdio/netbird/shared/relay/client"
|
||||
)
|
||||
|
||||
// wgTimeoutEscalationThreshold is the number of consecutive WireGuard
|
||||
// handshake timeouts after which the rosenpass state for the peer is
|
||||
// considered desynced and gets reset.
|
||||
const wgTimeoutEscalationThreshold = 3
|
||||
|
||||
// MetricsRecorder is an interface for recording peer connection metrics
|
||||
type MetricsRecorder interface {
|
||||
RecordConnectionStages(
|
||||
@@ -118,6 +123,9 @@ type Conn struct {
|
||||
wgWatcher *WGWatcher
|
||||
wgWatcherWg sync.WaitGroup
|
||||
wgWatcherCancel context.CancelFunc
|
||||
// wgTimeouts counts consecutive WireGuard handshake timeouts without a
|
||||
// successful handshake in between. Guarded by mu.
|
||||
wgTimeouts int
|
||||
|
||||
// used to store the remote Rosenpass key for Relayed connection in case of connection update from ice
|
||||
rosenpassRemoteKey []byte
|
||||
@@ -683,6 +691,29 @@ func (conn *Conn) onWGDisconnected() {
|
||||
default:
|
||||
conn.Log.Debugf("No active connection to close on WG timeout")
|
||||
}
|
||||
|
||||
conn.escalateWGTimeoutLocked()
|
||||
}
|
||||
|
||||
// escalateWGTimeoutLocked resets the peer's rosenpass state after repeated
|
||||
// handshake timeouts. With rosenpass enabled, persistent timeouts mean the
|
||||
// preshared keys have desynced; the renewal exchange runs over the dead
|
||||
// tunnel and cannot resync them. Reporting the peer disconnected drops its
|
||||
// rosenpass state, so the next connection configuration programs the
|
||||
// rendezvous key and the tunnel can bootstrap again. Callers must hold mu.
|
||||
func (conn *Conn) escalateWGTimeoutLocked() {
|
||||
if conn.config.RosenpassConfig.PubKey == nil {
|
||||
return
|
||||
}
|
||||
|
||||
conn.wgTimeouts++
|
||||
if conn.wgTimeouts < wgTimeoutEscalationThreshold || conn.onDisconnected == nil {
|
||||
return
|
||||
}
|
||||
conn.wgTimeouts = 0
|
||||
|
||||
conn.Log.Warnf("%d consecutive WireGuard handshake timeouts, resetting rosenpass state for peer", wgTimeoutEscalationThreshold)
|
||||
conn.onDisconnected(conn.config.WgConfig.RemoteKey)
|
||||
}
|
||||
|
||||
func (conn *Conn) updateRelayStatus(relayServerAddr string, rosenpassPubKey []byte, updateTime time.Time) {
|
||||
@@ -812,7 +843,7 @@ func (conn *Conn) enableWgWatcherIfNeeded(enabledTime time.Time) {
|
||||
conn.wgWatcherWg.Add(1)
|
||||
go func() {
|
||||
defer conn.wgWatcherWg.Done()
|
||||
conn.wgWatcher.EnableWgWatcher(wgWatcherCtx, enabledTime, conn.onWGDisconnected, conn.onWGHandshakeSuccess)
|
||||
conn.wgWatcher.EnableWgWatcher(wgWatcherCtx, enabledTime, conn.onWGDisconnected, conn.onWGHandshakeSuccess, conn.onWGCheckSuccess)
|
||||
}()
|
||||
}
|
||||
|
||||
@@ -892,6 +923,15 @@ func (conn *Conn) onWGHandshakeSuccess(when time.Time) {
|
||||
conn.recordConnectionMetrics()
|
||||
}
|
||||
|
||||
// onWGCheckSuccess is called for every watcher check that observed a fresh
|
||||
// handshake, including handshakes of connections that were already up when
|
||||
// the watcher started.
|
||||
func (conn *Conn) onWGCheckSuccess() {
|
||||
conn.mu.Lock()
|
||||
conn.wgTimeouts = 0
|
||||
conn.mu.Unlock()
|
||||
}
|
||||
|
||||
// recordConnectionMetrics records connection stage timestamps as metrics
|
||||
func (conn *Conn) recordConnectionMetrics() {
|
||||
if conn.metricsRecorder == nil {
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/netbirdio/netbird/client/iface"
|
||||
@@ -304,3 +305,84 @@ func TestConn_presharedKey_RosenpassManaged(t *testing.T) {
|
||||
t.Fatalf("expected non-nil presharedKey before Rosenpass manages PSK")
|
||||
}
|
||||
}
|
||||
|
||||
func newWGTimeoutTestConn(rosenpassEnabled bool, disconnected *[]string) *Conn {
|
||||
cfg := ConnConfig{
|
||||
Key: "LLHf3Ma6z6mdLbriAJbqhX7+nM/B71lgw2+91q3LfhU=",
|
||||
LocalKey: "RRHf3Ma6z6mdLbriAJbqhX7+nM/B71lgw2+91q3LfhU=",
|
||||
WgConfig: WgConfig{RemoteKey: "LLHf3Ma6z6mdLbriAJbqhX7+nM/B71lgw2+91q3LfhU="},
|
||||
}
|
||||
if rosenpassEnabled {
|
||||
cfg.RosenpassConfig = RosenpassConfig{PubKey: []byte("dummykey")}
|
||||
}
|
||||
|
||||
conn := &Conn{
|
||||
ctx: context.Background(),
|
||||
config: cfg,
|
||||
Log: log.WithField("peer", cfg.Key),
|
||||
metricsStages: &MetricsStages{},
|
||||
}
|
||||
conn.SetOnDisconnected(func(remotePeer string) {
|
||||
*disconnected = append(*disconnected, remotePeer)
|
||||
})
|
||||
return conn
|
||||
}
|
||||
|
||||
// TestConn_onWGDisconnected_EscalatesToRosenpassReset: repeated handshake
|
||||
// timeouts with rosenpass enabled mean the preshared keys have desynced. The
|
||||
// renewal exchange runs over the dead tunnel and cannot resync them, so after
|
||||
// wgTimeoutEscalationThreshold consecutive timeouts the conn must report the
|
||||
// peer disconnected, dropping its rosenpass state so the next configuration
|
||||
// programs the rendezvous key.
|
||||
func TestConn_onWGDisconnected_EscalatesToRosenpassReset(t *testing.T) {
|
||||
var disconnected []string
|
||||
conn := newWGTimeoutTestConn(true, &disconnected)
|
||||
|
||||
for i := 0; i < wgTimeoutEscalationThreshold-1; i++ {
|
||||
conn.onWGDisconnected()
|
||||
}
|
||||
assert.Empty(t, disconnected, "escalation must not fire below the threshold")
|
||||
|
||||
conn.onWGDisconnected()
|
||||
assert.Equal(t, []string{conn.config.WgConfig.RemoteKey}, disconnected,
|
||||
"reaching the threshold must report the peer disconnected once")
|
||||
|
||||
for i := 0; i < wgTimeoutEscalationThreshold-1; i++ {
|
||||
conn.onWGDisconnected()
|
||||
}
|
||||
assert.Len(t, disconnected, 1, "escalation must restart counting after firing")
|
||||
|
||||
conn.onWGDisconnected()
|
||||
assert.Len(t, disconnected, 2, "continued timeouts must escalate again")
|
||||
}
|
||||
|
||||
// TestConn_onWGDisconnected_CheckSuccessResetsEscalation: a successful
|
||||
// handshake between timeouts means the tunnel recovered; the counter must
|
||||
// start over.
|
||||
func TestConn_onWGDisconnected_CheckSuccessResetsEscalation(t *testing.T) {
|
||||
var disconnected []string
|
||||
conn := newWGTimeoutTestConn(true, &disconnected)
|
||||
|
||||
for i := 0; i < wgTimeoutEscalationThreshold-1; i++ {
|
||||
conn.onWGDisconnected()
|
||||
}
|
||||
conn.onWGCheckSuccess()
|
||||
|
||||
for i := 0; i < wgTimeoutEscalationThreshold-1; i++ {
|
||||
conn.onWGDisconnected()
|
||||
}
|
||||
assert.Empty(t, disconnected, "handshake success must reset the timeout count")
|
||||
}
|
||||
|
||||
// TestConn_onWGDisconnected_NoEscalationWithoutRosenpass: without rosenpass
|
||||
// there is no per-peer key state to reset; repeated timeouts must not report
|
||||
// disconnects.
|
||||
func TestConn_onWGDisconnected_NoEscalationWithoutRosenpass(t *testing.T) {
|
||||
var disconnected []string
|
||||
conn := newWGTimeoutTestConn(false, &disconnected)
|
||||
|
||||
for i := 0; i < wgTimeoutEscalationThreshold*3; i++ {
|
||||
conn.onWGDisconnected()
|
||||
}
|
||||
assert.Empty(t, disconnected, "escalation must be limited to rosenpass connections")
|
||||
}
|
||||
|
||||
@@ -71,9 +71,11 @@ func (w *WGWatcher) PrepareInitialHandshake() (ok bool) {
|
||||
|
||||
// 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)
|
||||
// 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)
|
||||
|
||||
w.muEnabled.Lock()
|
||||
w.enabled = false
|
||||
@@ -90,7 +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) {
|
||||
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)
|
||||
@@ -117,6 +119,10 @@ func (w *WGWatcher) periodicHandshakeCheck(ctx context.Context, onDisconnectedFn
|
||||
}
|
||||
}
|
||||
|
||||
if onCheckSuccessFn != nil && ctx.Err() == nil {
|
||||
onCheckSuccessFn()
|
||||
}
|
||||
|
||||
lastHandshake = *handshake
|
||||
|
||||
resetTime := time.Until(handshake.Add(checkPeriod))
|
||||
|
||||
@@ -24,6 +24,65 @@ 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 = 5 * time.Second
|
||||
wgHandshakeOvertime = 1 * time.Second
|
||||
|
||||
mlog := log.WithField("peer", "tet")
|
||||
stats := &mockHandshakeStats{handshake: time.Now()}
|
||||
watcher := NewWGWatcher(mlog, stats, "", newStateDump("peer", mlog, &Status{}))
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
require.True(t, watcher.PrepareInitialHandshake())
|
||||
|
||||
firstHandshake := make(chan struct{}, 1)
|
||||
checkSuccess := make(chan struct{}, 1)
|
||||
go 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:
|
||||
}
|
||||
}
|
||||
|
||||
func TestWGWatcher_EnableWgWatcher(t *testing.T) {
|
||||
checkPeriod = 5 * time.Second
|
||||
wgHandshakeOvertime = 1 * time.Second
|
||||
@@ -44,7 +103,7 @@ func TestWGWatcher_EnableWgWatcher(t *testing.T) {
|
||||
onDisconnected <- struct{}{}
|
||||
}, func(when time.Time) {
|
||||
mlog.Infof("onHandshakeSuccess: %v", when)
|
||||
})
|
||||
}, nil)
|
||||
|
||||
// wait for initial reading
|
||||
time.Sleep(2 * time.Second)
|
||||
@@ -73,7 +132,7 @@ func TestWGWatcher_ReEnable(t *testing.T) {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
watcher.EnableWgWatcher(ctx, time.Now(), func() {}, func(when time.Time) {})
|
||||
watcher.EnableWgWatcher(ctx, time.Now(), func() {}, func(when time.Time) {}, nil)
|
||||
}()
|
||||
cancel()
|
||||
|
||||
@@ -89,7 +148,7 @@ func TestWGWatcher_ReEnable(t *testing.T) {
|
||||
onDisconnected := make(chan struct{}, 1)
|
||||
go watcher.EnableWgWatcher(ctx, time.Now(), func() {
|
||||
onDisconnected <- struct{}{}
|
||||
}, func(when time.Time) {})
|
||||
}, func(when time.Time) {}, nil)
|
||||
|
||||
time.Sleep(2 * time.Second)
|
||||
mocWgIface.disconnect()
|
||||
|
||||
@@ -39,6 +39,7 @@ type rpServer interface {
|
||||
|
||||
type Manager struct {
|
||||
ifaceName string
|
||||
localWgKey wgtypes.Key
|
||||
spk []byte
|
||||
ssk []byte
|
||||
rpKeyHash string
|
||||
@@ -51,8 +52,9 @@ type Manager struct {
|
||||
wgIface PresharedKeySetter
|
||||
}
|
||||
|
||||
// NewManager creates a new Rosenpass manager
|
||||
func NewManager(preSharedKey *wgtypes.Key, wgIfaceName string) (*Manager, error) {
|
||||
// NewManager creates a new Rosenpass manager. localWgKey is the local
|
||||
// WireGuard public key, used to derive the per-peer rendezvous key.
|
||||
func NewManager(preSharedKey *wgtypes.Key, wgIfaceName string, localWgKey wgtypes.Key) (*Manager, error) {
|
||||
public, secret, err := rp.GenerateKeyPair()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -62,6 +64,7 @@ func NewManager(preSharedKey *wgtypes.Key, wgIfaceName string) (*Manager, error)
|
||||
log.Tracef("generated new rosenpass key pair with public key %s", rpKeyHash)
|
||||
return &Manager{
|
||||
ifaceName: wgIfaceName,
|
||||
localWgKey: localWgKey,
|
||||
rpKeyHash: rpKeyHash,
|
||||
spk: public,
|
||||
ssk: secret,
|
||||
@@ -73,7 +76,7 @@ func NewManager(preSharedKey *wgtypes.Key, wgIfaceName string) (*Manager, error)
|
||||
// nil receiver in addPeer -> m.rpWgHandler.AddPeer. generateConfig will
|
||||
// replace it with a fresh handler on each Run() to clear stale peer
|
||||
// state from previous engine sessions.
|
||||
rpWgHandler: NewNetbirdHandler(),
|
||||
rpWgHandler: NewNetbirdHandler((*[32]byte)(preSharedKey), localWgKey),
|
||||
lock: sync.Mutex{},
|
||||
}, nil
|
||||
}
|
||||
@@ -161,7 +164,7 @@ func (m *Manager) generateConfig() (rp.Config, error) {
|
||||
cfg.Peers = []rp.PeerConfig{}
|
||||
|
||||
m.lock.Lock()
|
||||
m.rpWgHandler = NewNetbirdHandler()
|
||||
m.rpWgHandler = NewNetbirdHandler(m.preSharedKey, m.localWgKey)
|
||||
if m.wgIface != nil {
|
||||
m.rpWgHandler.SetInterface(m.wgIface)
|
||||
}
|
||||
|
||||
@@ -85,7 +85,7 @@ func newTestManager(spkFirstByte byte, mock *mockServer) *Manager {
|
||||
ssk: make([]byte, 32),
|
||||
rpKeyHash: "test-hash",
|
||||
rpPeerIDs: make(map[string]*rp.PeerID),
|
||||
rpWgHandler: NewNetbirdHandler(),
|
||||
rpWgHandler: NewNetbirdHandler(nil, wgtypes.Key{0x01}),
|
||||
server: mock,
|
||||
}
|
||||
}
|
||||
@@ -255,7 +255,7 @@ func TestAddPeer_NilServer_ReturnsErrorNoCrash(t *testing.T) {
|
||||
// issue #4341 cannot occur in the window between NewManager and Run().
|
||||
func TestNewManager_PreInitializesHandler(t *testing.T) {
|
||||
psk := wgtypes.Key{}
|
||||
m, err := NewManager(&psk, "wt0")
|
||||
m, err := NewManager(&psk, "wt0", wgtypes.Key{0x01})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, m.rpWgHandler, "rpWgHandler must be initialized in NewManager")
|
||||
}
|
||||
@@ -332,7 +332,7 @@ func TestIsPresharedKeyInitialized_AddedButNotHandshaken_ReturnsFalse(t *testing
|
||||
// --- NetbirdHandler.outputKey ----------------------------------------------
|
||||
|
||||
func TestHandler_OutputKey_FirstCallUsesUpdateOnlyFalse(t *testing.T) {
|
||||
h := NewNetbirdHandler()
|
||||
h := NewNetbirdHandler(nil, wgtypes.Key{0x01})
|
||||
iface := &mockIface{}
|
||||
h.SetInterface(iface)
|
||||
|
||||
@@ -349,7 +349,7 @@ func TestHandler_OutputKey_FirstCallUsesUpdateOnlyFalse(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestHandler_OutputKey_SubsequentCallsUseUpdateOnlyTrue(t *testing.T) {
|
||||
h := NewNetbirdHandler()
|
||||
h := NewNetbirdHandler(nil, wgtypes.Key{0x01})
|
||||
iface := &mockIface{}
|
||||
h.SetInterface(iface)
|
||||
|
||||
@@ -365,7 +365,7 @@ func TestHandler_OutputKey_SubsequentCallsUseUpdateOnlyTrue(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestHandler_OutputKey_NilInterface_NoCrashNoCall(t *testing.T) {
|
||||
h := NewNetbirdHandler()
|
||||
h := NewNetbirdHandler(nil, wgtypes.Key{0x01})
|
||||
// no SetInterface — iface remains nil
|
||||
pid := rp.PeerID{0x03}
|
||||
h.AddPeer(pid, "wt0", rp.Key(wgtypes.Key{}))
|
||||
@@ -375,7 +375,7 @@ func TestHandler_OutputKey_NilInterface_NoCrashNoCall(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestHandler_OutputKey_UnknownPeer_NoCall(t *testing.T) {
|
||||
h := NewNetbirdHandler()
|
||||
h := NewNetbirdHandler(nil, wgtypes.Key{0x01})
|
||||
iface := &mockIface{}
|
||||
h.SetInterface(iface)
|
||||
|
||||
@@ -384,7 +384,7 @@ func TestHandler_OutputKey_UnknownPeer_NoCall(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestHandler_RemovePeer_ClearsInitializedState(t *testing.T) {
|
||||
h := NewNetbirdHandler()
|
||||
h := NewNetbirdHandler(nil, wgtypes.Key{0x01})
|
||||
iface := &mockIface{}
|
||||
h.SetInterface(iface)
|
||||
|
||||
@@ -398,7 +398,7 @@ func TestHandler_RemovePeer_ClearsInitializedState(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestHandler_SetInterfaceAfterAddPeer_StillReceivesKey(t *testing.T) {
|
||||
h := NewNetbirdHandler()
|
||||
h := NewNetbirdHandler(nil, wgtypes.Key{0x01})
|
||||
pid := rp.PeerID{0x05}
|
||||
wgKey := wgtypes.Key{0xEE}
|
||||
h.AddPeer(pid, "wt0", rp.Key(wgKey))
|
||||
|
||||
@@ -18,19 +18,34 @@ type PresharedKeySetter interface {
|
||||
type wireGuardPeer struct {
|
||||
Interface string
|
||||
PublicKey rp.Key
|
||||
// initialized is true once a completed exchange has set a
|
||||
// Rosenpass-managed PSK for this peer.
|
||||
initialized bool
|
||||
// chainKey is the key output by the last completed exchange, advanced by
|
||||
// one ratchet step on expiry. Nil until the first exchange completes and
|
||||
// after the peer has fallen back to the rendezvous key.
|
||||
chainKey *wgtypes.Key
|
||||
// expiries counts failed renewals since the last completed exchange.
|
||||
expiries int
|
||||
}
|
||||
|
||||
type NetbirdHandler struct {
|
||||
mu sync.Mutex
|
||||
iface PresharedKeySetter
|
||||
peers map[rp.PeerID]wireGuardPeer
|
||||
initializedPeers map[rp.PeerID]bool
|
||||
mu sync.Mutex
|
||||
iface PresharedKeySetter
|
||||
// preSharedKey is the account-level preshared key, used as the rendezvous
|
||||
// key when set. Nil means the deterministic seed key is used instead.
|
||||
preSharedKey *[32]byte
|
||||
// localWgKey is the local WireGuard public key, one of the two inputs to
|
||||
// the deterministic seed key.
|
||||
localWgKey wgtypes.Key
|
||||
peers map[rp.PeerID]*wireGuardPeer
|
||||
}
|
||||
|
||||
func NewNetbirdHandler() *NetbirdHandler {
|
||||
func NewNetbirdHandler(preSharedKey *[32]byte, localWgKey wgtypes.Key) *NetbirdHandler {
|
||||
return &NetbirdHandler{
|
||||
peers: map[rp.PeerID]wireGuardPeer{},
|
||||
initializedPeers: map[rp.PeerID]bool{},
|
||||
preSharedKey: preSharedKey,
|
||||
localWgKey: localWgKey,
|
||||
peers: map[rp.PeerID]*wireGuardPeer{},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -42,10 +57,16 @@ func (h *NetbirdHandler) SetInterface(iface PresharedKeySetter) {
|
||||
h.iface = iface
|
||||
}
|
||||
|
||||
// AddPeer registers a peer with the handler. Re-adding a known peer (every
|
||||
// reconnection does) keeps its key recovery state.
|
||||
func (h *NetbirdHandler) AddPeer(pid rp.PeerID, intf string, pk rp.Key) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
h.peers[pid] = wireGuardPeer{
|
||||
if existing, ok := h.peers[pid]; ok && existing.PublicKey == pk {
|
||||
existing.Interface = intf
|
||||
return
|
||||
}
|
||||
h.peers[pid] = &wireGuardPeer{
|
||||
Interface: intf,
|
||||
PublicKey: pk,
|
||||
}
|
||||
@@ -55,7 +76,6 @@ func (h *NetbirdHandler) RemovePeer(pid rp.PeerID) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
delete(h.peers, pid)
|
||||
delete(h.initializedPeers, pid)
|
||||
}
|
||||
|
||||
// IsPeerInitialized returns true if Rosenpass has completed a handshake
|
||||
@@ -63,50 +83,131 @@ func (h *NetbirdHandler) RemovePeer(pid rp.PeerID) {
|
||||
func (h *NetbirdHandler) IsPeerInitialized(pid rp.PeerID) bool {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
return h.initializedPeers[pid]
|
||||
peer, ok := h.peers[pid]
|
||||
return ok && peer.initialized
|
||||
}
|
||||
|
||||
// HandshakeCompleted programs the freshly exchanged output key and resets the
|
||||
// peer's key recovery state.
|
||||
func (h *NetbirdHandler) HandshakeCompleted(pid rp.PeerID, key rp.Key) {
|
||||
h.outputKey(rp.KeyOutputReasonStale, pid, key)
|
||||
psk := wgtypes.Key(key)
|
||||
|
||||
h.mu.Lock()
|
||||
peer, ok := h.peers[pid]
|
||||
if !ok {
|
||||
h.mu.Unlock()
|
||||
return
|
||||
}
|
||||
if peer.expiries > 0 {
|
||||
log.Infof("rosenpass exchange completed for peer %s after %d expired renewals", wgtypes.Key(peer.PublicKey), peer.expiries)
|
||||
}
|
||||
peer.chainKey = &psk
|
||||
peer.expiries = 0
|
||||
updateOnly := peer.initialized
|
||||
h.mu.Unlock()
|
||||
|
||||
if !h.applyKey(pid, psk, updateOnly) {
|
||||
return
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
if peer, ok := h.peers[pid]; ok {
|
||||
peer.initialized = true
|
||||
}
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
// HandshakeExpired replaces the expired key. The renewal exchange runs over
|
||||
// the tunnel keyed by the PSK itself, so the replacement must be derivable on
|
||||
// both ends without communication: the first expiry ratchets the last shared
|
||||
// key forward, repeated expiries (and expiries without a completed exchange)
|
||||
// fall back to the rendezvous key and drop the peer out of the initialized
|
||||
// state so connection reconfigurations reprogram the rendezvous key as well.
|
||||
func (h *NetbirdHandler) HandshakeExpired(pid rp.PeerID) {
|
||||
key, _ := rp.GeneratePresharedKey()
|
||||
h.outputKey(rp.KeyOutputReasonStale, pid, key)
|
||||
h.mu.Lock()
|
||||
peer, ok := h.peers[pid]
|
||||
if !ok {
|
||||
h.mu.Unlock()
|
||||
return
|
||||
}
|
||||
|
||||
peer.expiries++
|
||||
|
||||
var psk wgtypes.Key
|
||||
if peer.chainKey != nil && peer.expiries == 1 {
|
||||
log.Infof("rosenpass key for peer %s expired without renewal, advancing to ratcheted key", wgtypes.Key(peer.PublicKey))
|
||||
psk = RatchetKey(*peer.chainKey)
|
||||
peer.chainKey = &psk
|
||||
} else {
|
||||
rendezvous, err := h.rendezvousKey(peer)
|
||||
if err != nil {
|
||||
h.mu.Unlock()
|
||||
// Fail closed: without a rendezvous key the expired key must
|
||||
// still be rotated out, even if the replacement is unusable.
|
||||
log.Errorf("failed to derive rendezvous key, replacing expired key with a random one: %v", err)
|
||||
h.applyRandomKey(pid)
|
||||
return
|
||||
}
|
||||
log.Warnf("rosenpass key for peer %s expired %d times without renewal, falling back to the rendezvous key", wgtypes.Key(peer.PublicKey), peer.expiries)
|
||||
psk = rendezvous
|
||||
peer.chainKey = nil
|
||||
peer.initialized = false
|
||||
}
|
||||
h.mu.Unlock()
|
||||
|
||||
h.applyKey(pid, psk, true)
|
||||
}
|
||||
|
||||
func (h *NetbirdHandler) outputKey(_ rp.KeyOutputReason, pid rp.PeerID, psk rp.Key) {
|
||||
// rendezvousKey returns the key both ends converge on without communication:
|
||||
// the account-level preshared key when configured, the deterministic seed key
|
||||
// otherwise. It mirrors the key that peer connections program when Rosenpass
|
||||
// does not manage the peer yet.
|
||||
func (h *NetbirdHandler) rendezvousKey(peer *wireGuardPeer) (wgtypes.Key, error) {
|
||||
if h.preSharedKey != nil {
|
||||
return *h.preSharedKey, nil
|
||||
}
|
||||
|
||||
seed, err := DeterministicSeedKey(h.localWgKey.String(), wgtypes.Key(peer.PublicKey).String())
|
||||
if err != nil {
|
||||
return wgtypes.Key{}, err
|
||||
}
|
||||
return *seed, nil
|
||||
}
|
||||
|
||||
// applyKey writes the preshared key for the peer to the WireGuard interface
|
||||
// and reports whether the write succeeded.
|
||||
func (h *NetbirdHandler) applyKey(pid rp.PeerID, psk wgtypes.Key, updateOnly bool) bool {
|
||||
h.mu.Lock()
|
||||
iface := h.iface
|
||||
wg, ok := h.peers[pid]
|
||||
isInitialized := h.initializedPeers[pid]
|
||||
peer, ok := h.peers[pid]
|
||||
var peerKey string
|
||||
if ok {
|
||||
peerKey = wgtypes.Key(peer.PublicKey).String()
|
||||
}
|
||||
h.mu.Unlock()
|
||||
|
||||
if iface == nil {
|
||||
log.Warn("rosenpass: interface not set, cannot update preshared key")
|
||||
return
|
||||
return false
|
||||
}
|
||||
|
||||
if !ok {
|
||||
return
|
||||
return false
|
||||
}
|
||||
|
||||
peerKey := wgtypes.Key(wg.PublicKey).String()
|
||||
pskKey := wgtypes.Key(psk)
|
||||
|
||||
// Use updateOnly=true for later rotations (peer already has Rosenpass PSK)
|
||||
// Use updateOnly=false for first rotation (peer has original/empty PSK)
|
||||
if err := iface.SetPresharedKey(peerKey, pskKey, isInitialized); err != nil {
|
||||
if err := iface.SetPresharedKey(peerKey, psk, updateOnly); err != nil {
|
||||
log.Errorf("Failed to apply rosenpass key: %v", err)
|
||||
return
|
||||
return false
|
||||
}
|
||||
|
||||
// Mark peer as isInitialized after the successful first rotation
|
||||
if !isInitialized {
|
||||
h.mu.Lock()
|
||||
if _, exists := h.peers[pid]; exists {
|
||||
h.initializedPeers[pid] = true
|
||||
}
|
||||
h.mu.Unlock()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (h *NetbirdHandler) applyRandomKey(pid rp.PeerID) {
|
||||
key, err := rp.GeneratePresharedKey()
|
||||
if err != nil {
|
||||
log.Errorf("failed to generate random preshared key: %v", err)
|
||||
return
|
||||
}
|
||||
h.applyKey(pid, wgtypes.Key(key), true)
|
||||
}
|
||||
|
||||
248
client/internal/rosenpass/netbird_handler_test.go
Normal file
248
client/internal/rosenpass/netbird_handler_test.go
Normal file
@@ -0,0 +1,248 @@
|
||||
package rosenpass
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
rp "cunicu.li/go-rosenpass"
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
|
||||
)
|
||||
|
||||
// handlerTestLink wires two NetbirdHandlers as the two ends of a single
|
||||
// tunnel: handler A manages the rosenpass peer B and vice versa, the way two
|
||||
// NetBird clients see each other.
|
||||
type handlerTestLink struct {
|
||||
handlerA, handlerB *NetbirdHandler
|
||||
ifaceA, ifaceB *mockIface
|
||||
pidA, pidB rp.PeerID
|
||||
wgKeyA, wgKeyB wgtypes.Key
|
||||
}
|
||||
|
||||
func newHandlerTestLink(t *testing.T, preSharedKey *[32]byte) *handlerTestLink {
|
||||
t.Helper()
|
||||
|
||||
link := &handlerTestLink{
|
||||
ifaceA: &mockIface{},
|
||||
ifaceB: &mockIface{},
|
||||
}
|
||||
link.pidA[0] = 0xaa
|
||||
link.pidB[0] = 0xbb
|
||||
link.wgKeyA[31] = 1
|
||||
link.wgKeyB[31] = 2
|
||||
|
||||
link.handlerA = NewNetbirdHandler(preSharedKey, link.wgKeyA)
|
||||
link.handlerB = NewNetbirdHandler(preSharedKey, link.wgKeyB)
|
||||
|
||||
link.handlerA.SetInterface(link.ifaceA)
|
||||
link.handlerB.SetInterface(link.ifaceB)
|
||||
|
||||
link.handlerA.AddPeer(link.pidB, "wt0", rp.Key(link.wgKeyB))
|
||||
link.handlerB.AddPeer(link.pidA, "wt0", rp.Key(link.wgKeyA))
|
||||
|
||||
return link
|
||||
}
|
||||
|
||||
// complete simulates a completed rosenpass exchange: both ends derive the
|
||||
// same output key.
|
||||
func (l *handlerTestLink) complete(osk rp.Key) {
|
||||
l.handlerA.HandshakeCompleted(l.pidB, osk)
|
||||
l.handlerB.HandshakeCompleted(l.pidA, osk)
|
||||
}
|
||||
|
||||
// expire simulates a failed key renewal on both ends.
|
||||
func (l *handlerTestLink) expire() {
|
||||
l.handlerA.HandshakeExpired(l.pidB)
|
||||
l.handlerB.HandshakeExpired(l.pidA)
|
||||
}
|
||||
|
||||
func lastPSK(t *testing.T, m *mockIface) wgtypes.Key {
|
||||
t.Helper()
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
require.NotEmpty(t, m.calls, "expected at least one SetPresharedKey call")
|
||||
return m.calls[len(m.calls)-1].psk
|
||||
}
|
||||
|
||||
func TestHandshakeCompleted_SetsKeyAndInitializes(t *testing.T) {
|
||||
link := newHandlerTestLink(t, nil)
|
||||
|
||||
var osk rp.Key
|
||||
osk[0] = 0x42
|
||||
link.complete(osk)
|
||||
|
||||
require.Equal(t, wgtypes.Key(osk), lastPSK(t, link.ifaceA), "completed exchange must program the osk")
|
||||
require.False(t, link.ifaceA.calls[0].updateOnly, "first rotation must not be update-only")
|
||||
require.True(t, link.handlerA.IsPeerInitialized(link.pidB), "peer must be initialized after first completed exchange")
|
||||
|
||||
link.complete(osk)
|
||||
require.True(t, link.ifaceA.calls[1].updateOnly, "later rotations must be update-only")
|
||||
}
|
||||
|
||||
// TestHandshakeExpired_BothSidesConverge encodes the core recovery invariant:
|
||||
// rosenpass renewals run over the tunnel that the PSK itself keys, so when a
|
||||
// renewal fails on both ends, both ends must fall back to the same key or the
|
||||
// tunnel can never handshake again.
|
||||
func TestHandshakeExpired_BothSidesConverge(t *testing.T) {
|
||||
link := newHandlerTestLink(t, nil)
|
||||
|
||||
var osk rp.Key
|
||||
osk[0] = 0x42
|
||||
link.complete(osk)
|
||||
|
||||
link.expire()
|
||||
keyA := lastPSK(t, link.ifaceA)
|
||||
keyB := lastPSK(t, link.ifaceB)
|
||||
require.NotEqual(t, wgtypes.Key(osk), keyA, "expired key must be rotated out")
|
||||
require.Equal(t, keyA, keyB, "both ends must converge on the same key after expiry")
|
||||
|
||||
link.expire()
|
||||
require.Equal(t, lastPSK(t, link.ifaceA), lastPSK(t, link.ifaceB),
|
||||
"both ends must still converge after repeated expiries")
|
||||
}
|
||||
|
||||
// TestHandshakeExpired_ExpiryWithoutCompletionConverges covers the bootstrap
|
||||
// case: the initial exchange never completed (the tunnel ran on the rendezvous
|
||||
// key), so an expiry must not replace the working key with an unrecoverable
|
||||
// one on either end.
|
||||
func TestHandshakeExpired_ExpiryWithoutCompletionConverges(t *testing.T) {
|
||||
link := newHandlerTestLink(t, nil)
|
||||
|
||||
link.expire()
|
||||
require.Equal(t, lastPSK(t, link.ifaceA), lastPSK(t, link.ifaceB),
|
||||
"both ends must converge when the exchange never completed")
|
||||
}
|
||||
|
||||
// TestHandshakeExpired_RepeatedExpiryClearsInitialized: once renewals keep
|
||||
// failing, the peer must drop out of the initialized state so the next
|
||||
// connection reconfiguration reprograms the rendezvous key instead of
|
||||
// preserving a poisoned rosenpass-managed key.
|
||||
func TestHandshakeExpired_RepeatedExpiryClearsInitialized(t *testing.T) {
|
||||
link := newHandlerTestLink(t, nil)
|
||||
|
||||
var osk rp.Key
|
||||
osk[0] = 0x42
|
||||
link.complete(osk)
|
||||
|
||||
link.expire()
|
||||
link.expire()
|
||||
|
||||
require.False(t, link.handlerA.IsPeerInitialized(link.pidB),
|
||||
"repeated expiries must clear the initialized state")
|
||||
require.False(t, link.handlerB.IsPeerInitialized(link.pidA),
|
||||
"repeated expiries must clear the initialized state")
|
||||
}
|
||||
|
||||
// TestHandshakeCompleted_AfterExpiryRecovers: a completed exchange after a
|
||||
// desync must fully reset the recovery state.
|
||||
func TestHandshakeCompleted_AfterExpiryRecovers(t *testing.T) {
|
||||
link := newHandlerTestLink(t, nil)
|
||||
|
||||
var osk1, osk2 rp.Key
|
||||
osk1[0] = 1
|
||||
osk2[0] = 2
|
||||
|
||||
link.complete(osk1)
|
||||
link.expire()
|
||||
link.expire()
|
||||
|
||||
link.complete(osk2)
|
||||
require.Equal(t, wgtypes.Key(osk2), lastPSK(t, link.ifaceA), "new exchange must program the fresh osk")
|
||||
require.True(t, link.handlerA.IsPeerInitialized(link.pidB), "peer must be initialized again after recovery")
|
||||
|
||||
link.expire()
|
||||
require.Equal(t, lastPSK(t, link.ifaceA), lastPSK(t, link.ifaceB),
|
||||
"recovered link must converge again on the next expiry")
|
||||
require.NotEqual(t, wgtypes.Key(osk2), lastPSK(t, link.ifaceA), "expired key must be rotated out")
|
||||
}
|
||||
|
||||
// TestHandshakeExpired_FirstExpiryRatchetsLastKey: the first expiry must
|
||||
// derive the replacement from the last shared key, so an attacker who only
|
||||
// blocks the renewal exchange gains nothing over the previous key.
|
||||
func TestHandshakeExpired_FirstExpiryRatchetsLastKey(t *testing.T) {
|
||||
link := newHandlerTestLink(t, nil)
|
||||
|
||||
var osk rp.Key
|
||||
osk[0] = 0x42
|
||||
link.complete(osk)
|
||||
|
||||
link.expire()
|
||||
require.Equal(t, RatchetKey(wgtypes.Key(osk)), lastPSK(t, link.ifaceA),
|
||||
"first expiry must program the ratcheted key")
|
||||
require.True(t, link.handlerA.IsPeerInitialized(link.pidB),
|
||||
"ratchet step must keep the peer initialized so reconfigurations preserve the key")
|
||||
}
|
||||
|
||||
// TestHandshakeExpired_RepeatedExpiryFallsBackToSeed: once the ratchet key
|
||||
// also fails, both ends must land on the same key that peer connections
|
||||
// program for uninitialized peers, so a reconnect completes the recovery.
|
||||
func TestHandshakeExpired_RepeatedExpiryFallsBackToSeed(t *testing.T) {
|
||||
link := newHandlerTestLink(t, nil)
|
||||
|
||||
var osk rp.Key
|
||||
osk[0] = 0x42
|
||||
link.complete(osk)
|
||||
|
||||
link.expire()
|
||||
link.expire()
|
||||
|
||||
seed, err := DeterministicSeedKey(link.wgKeyA.String(), link.wgKeyB.String())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, *seed, lastPSK(t, link.ifaceA), "repeated expiry must fall back to the seed key")
|
||||
require.Equal(t, *seed, lastPSK(t, link.ifaceB), "repeated expiry must fall back to the seed key")
|
||||
}
|
||||
|
||||
// TestHandshakeExpired_ConfiguredPSKUsedAsRendezvous: with an account-level
|
||||
// preshared key configured, the fallback must be that key, matching what peer
|
||||
// connections program for uninitialized peers.
|
||||
func TestHandshakeExpired_ConfiguredPSKUsedAsRendezvous(t *testing.T) {
|
||||
psk := &[32]byte{0x77}
|
||||
link := newHandlerTestLink(t, psk)
|
||||
|
||||
var osk rp.Key
|
||||
osk[0] = 0x42
|
||||
link.complete(osk)
|
||||
|
||||
link.expire()
|
||||
link.expire()
|
||||
|
||||
require.Equal(t, wgtypes.Key(*psk), lastPSK(t, link.ifaceA),
|
||||
"fallback must be the configured preshared key")
|
||||
}
|
||||
|
||||
// TestHandshakeExpired_ExpiryWritesAreUpdateOnly: expiry replacements must
|
||||
// never create a WireGuard peer that connection management has removed.
|
||||
func TestHandshakeExpired_ExpiryWritesAreUpdateOnly(t *testing.T) {
|
||||
link := newHandlerTestLink(t, nil)
|
||||
|
||||
var osk rp.Key
|
||||
osk[0] = 0x42
|
||||
link.complete(osk)
|
||||
|
||||
link.expire()
|
||||
link.expire()
|
||||
|
||||
for _, call := range link.ifaceA.calls[1:] {
|
||||
require.True(t, call.updateOnly, "expiry writes must be update-only")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAddPeer_ReAddKeepsRecoveryState: reconnections re-add the peer on every
|
||||
// OnConnected; that must not reset the expiry chain state.
|
||||
func TestAddPeer_ReAddKeepsRecoveryState(t *testing.T) {
|
||||
link := newHandlerTestLink(t, nil)
|
||||
|
||||
var osk rp.Key
|
||||
osk[0] = 0x42
|
||||
link.complete(osk)
|
||||
link.expire()
|
||||
|
||||
link.handlerA.AddPeer(link.pidB, "wt0", rp.Key(link.wgKeyB))
|
||||
require.True(t, link.handlerA.IsPeerInitialized(link.pidB),
|
||||
"re-adding a known peer must keep its state")
|
||||
|
||||
link.expire()
|
||||
seed, err := DeterministicSeedKey(link.wgKeyA.String(), link.wgKeyB.String())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, *seed, lastPSK(t, link.ifaceA),
|
||||
"second expiry after re-add must continue to the seed fallback")
|
||||
}
|
||||
@@ -1,11 +1,28 @@
|
||||
package rosenpass
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"fmt"
|
||||
|
||||
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
|
||||
)
|
||||
|
||||
// ratchetLabel domain-separates the expiry ratchet from other uses of the
|
||||
// rosenpass output key.
|
||||
const ratchetLabel = "netbird-rosenpass-expiry-ratchet"
|
||||
|
||||
// RatchetKey derives the successor preshared key from the previous Rosenpass
|
||||
// output key. When a key expires without a completed renewal, both peers
|
||||
// advance their last shared key by one ratchet step: the expired key is
|
||||
// rotated out while both ends still converge on an identical, non-public
|
||||
// replacement without communicating.
|
||||
func RatchetKey(prev wgtypes.Key) wgtypes.Key {
|
||||
input := make([]byte, 0, len(ratchetLabel)+len(prev))
|
||||
input = append(input, ratchetLabel...)
|
||||
input = append(input, prev[:]...)
|
||||
return sha256.Sum256(input)
|
||||
}
|
||||
|
||||
// DeterministicSeedKey derives a 32-byte WireGuard preshared key from a pair
|
||||
// of peer public keys. Both peers, given the same key pair, produce the same
|
||||
// output regardless of which side runs the function: the inputs are ordered
|
||||
|
||||
@@ -7,7 +7,7 @@ import (
|
||||
"os"
|
||||
"slices"
|
||||
|
||||
"github.com/shirou/gopsutil/v3/process"
|
||||
"github.com/shirou/gopsutil/v4/process"
|
||||
)
|
||||
|
||||
// getRunningProcesses returns a list of running process paths. The context bounds the work:
|
||||
|
||||
@@ -4,7 +4,7 @@ import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/shirou/gopsutil/v3/process"
|
||||
"github.com/shirou/gopsutil/v4/process"
|
||||
)
|
||||
|
||||
func Benchmark_getRunningProcesses(b *testing.B) {
|
||||
|
||||
2
go.mod
2
go.mod
@@ -104,6 +104,7 @@ require (
|
||||
github.com/redis/go-redis/v9 v9.7.3
|
||||
github.com/rs/xid v1.3.0
|
||||
github.com/shirou/gopsutil/v3 v3.24.4
|
||||
github.com/shirou/gopsutil/v4 v4.25.8
|
||||
github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966
|
||||
github.com/songgao/water v0.0.0-20200317203138-2b4b6d7c09d8
|
||||
github.com/stretchr/testify v1.11.1
|
||||
@@ -308,7 +309,6 @@ require (
|
||||
github.com/russellhaering/goxmldsig v1.6.0 // indirect
|
||||
github.com/ryanuber/go-glob v1.0.0 // indirect
|
||||
github.com/rymdport/portal v0.4.2 // indirect
|
||||
github.com/shirou/gopsutil/v4 v4.25.8 // indirect
|
||||
github.com/shoenig/go-m1cpu v0.2.1 // indirect
|
||||
github.com/shopspring/decimal v1.4.0 // indirect
|
||||
github.com/spf13/cast v1.7.0 // indirect
|
||||
|
||||
@@ -30,11 +30,16 @@ type RelayTrack struct {
|
||||
relayClient *Client
|
||||
err error
|
||||
created time.Time
|
||||
// ready is closed once the dial started by openConnVia finishes (relayClient
|
||||
// or err is set). Callers reusing a track wait on this instead of the track
|
||||
// lock, so the dial never runs under rt.Lock.
|
||||
ready chan struct{}
|
||||
}
|
||||
|
||||
func NewRelayTrack() *RelayTrack {
|
||||
return &RelayTrack{
|
||||
created: time.Now(),
|
||||
ready: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -326,34 +331,24 @@ func (m *Manager) openConnVia(ctx context.Context, serverAddress, peerKey string
|
||||
// check if already has a connection to the desired relay server
|
||||
m.relayClientsMutex.RLock()
|
||||
rt, ok := m.relayClients[serverAddress]
|
||||
if ok {
|
||||
rt.RLock()
|
||||
m.relayClientsMutex.RUnlock()
|
||||
defer rt.RUnlock()
|
||||
if rt.err != nil {
|
||||
return nil, rt.err
|
||||
}
|
||||
return rt.relayClient.OpenConn(ctx, peerKey)
|
||||
}
|
||||
m.relayClientsMutex.RUnlock()
|
||||
if ok {
|
||||
return m.openConnOnTrack(ctx, rt, peerKey)
|
||||
}
|
||||
|
||||
// if not, establish a new connection but check it again (because changed the lock type) before starting the
|
||||
// connection
|
||||
m.relayClientsMutex.Lock()
|
||||
rt, ok = m.relayClients[serverAddress]
|
||||
if ok {
|
||||
rt.RLock()
|
||||
m.relayClientsMutex.Unlock()
|
||||
defer rt.RUnlock()
|
||||
if rt.err != nil {
|
||||
return nil, rt.err
|
||||
}
|
||||
return rt.relayClient.OpenConn(ctx, peerKey)
|
||||
return m.openConnOnTrack(ctx, rt, peerKey)
|
||||
}
|
||||
|
||||
// create a new relay client and store it in the relayClients map
|
||||
// Publish the track and release the map lock BEFORE dialing, so the dial does
|
||||
// not run under rt.Lock (which would block RelayStates and the cleanup loop
|
||||
// for the full dial). Concurrent callers find this track and wait on rt.ready.
|
||||
rt = NewRelayTrack()
|
||||
rt.Lock()
|
||||
m.relayClients[serverAddress] = rt
|
||||
m.relayClientsMutex.Unlock()
|
||||
|
||||
@@ -361,8 +356,10 @@ func (m *Manager) openConnVia(ctx context.Context, serverAddress, peerKey string
|
||||
relayClient.SetTransportFallback(m.transportFallback)
|
||||
err := relayClient.Connect(m.ctx)
|
||||
if err != nil {
|
||||
rt.Lock()
|
||||
rt.err = err
|
||||
rt.Unlock()
|
||||
close(rt.ready)
|
||||
m.relayClientsMutex.Lock()
|
||||
delete(m.relayClients, serverAddress)
|
||||
m.relayClientsMutex.Unlock()
|
||||
@@ -370,14 +367,34 @@ func (m *Manager) openConnVia(ctx context.Context, serverAddress, peerKey string
|
||||
}
|
||||
// if connection closed then delete the relay client from the list
|
||||
relayClient.SetOnDisconnectListener(m.onServerDisconnected)
|
||||
rt.Lock()
|
||||
rt.relayClient = relayClient
|
||||
rt.Unlock()
|
||||
close(rt.ready)
|
||||
|
||||
conn, err := relayClient.OpenConn(ctx, peerKey)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return relayClient.OpenConn(ctx, peerKey)
|
||||
}
|
||||
|
||||
// openConnOnTrack opens a peer connection through an existing relay track,
|
||||
// waiting for the dial started by another openConnVia call to finish. It waits
|
||||
// on rt.ready rather than the track lock, so it neither holds nor contends the
|
||||
// track lock across the dial.
|
||||
func (m *Manager) openConnOnTrack(ctx context.Context, rt *RelayTrack, peerKey string) (net.Conn, error) {
|
||||
select {
|
||||
case <-rt.ready:
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
return conn, nil
|
||||
|
||||
rt.RLock()
|
||||
defer rt.RUnlock()
|
||||
if rt.err != nil {
|
||||
return nil, rt.err
|
||||
}
|
||||
if rt.relayClient == nil {
|
||||
return nil, ErrRelayClientNotConnected
|
||||
}
|
||||
return rt.relayClient.OpenConn(ctx, peerKey)
|
||||
}
|
||||
|
||||
func (m *Manager) onServerConnected() {
|
||||
@@ -476,6 +493,13 @@ func (m *Manager) cleanUpUnusedRelays() {
|
||||
continue
|
||||
}
|
||||
|
||||
// dial still in progress (openConnVia publishes the track before Connect
|
||||
// completes and no longer holds rt.Lock during it), nothing to clean up.
|
||||
if rt.relayClient == nil {
|
||||
rt.Unlock()
|
||||
continue
|
||||
}
|
||||
|
||||
if time.Since(rt.created) <= m.keepUnusedServerTime {
|
||||
rt.Unlock()
|
||||
continue
|
||||
|
||||
60
shared/relay/client/manager_cleanup_test.go
Normal file
60
shared/relay/client/manager_cleanup_test.go
Normal file
@@ -0,0 +1,60 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/netip"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestCleanUpUnusedRelays_DoesNotBlockOnRealHangingDial drives a real, hanging foreign
|
||||
// relay dial and asserts cleanUpUnusedRelays does not stall behind it.
|
||||
func TestCleanUpUnusedRelays_DoesNotBlockOnRealHangingDial(t *testing.T) {
|
||||
serverAddr := stallingRelayListener(t)
|
||||
|
||||
mCtx, mCancel := context.WithCancel(context.Background())
|
||||
t.Cleanup(mCancel)
|
||||
|
||||
m := NewManager(mCtx, nil, "alice", 1280)
|
||||
|
||||
dialDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(dialDone)
|
||||
_, _ = m.openConnVia(mCtx, serverAddr, "peerKey", netip.Addr{})
|
||||
}()
|
||||
|
||||
// The track appears in the map once the dial is in flight.
|
||||
require.Eventually(t, func() bool {
|
||||
m.relayClientsMutex.RLock()
|
||||
defer m.relayClientsMutex.RUnlock()
|
||||
_, ok := m.relayClients[serverAddr]
|
||||
return ok
|
||||
}, 5*time.Second, 5*time.Millisecond, "relay dial did not start")
|
||||
|
||||
cleanupDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(cleanupDone)
|
||||
m.cleanUpUnusedRelays()
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-cleanupDone:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("cleanUpUnusedRelays blocked on an in-progress relay dial while holding the relay map lock")
|
||||
}
|
||||
|
||||
m.relayClientsMutex.RLock()
|
||||
_, stillTracked := m.relayClients[serverAddr]
|
||||
m.relayClientsMutex.RUnlock()
|
||||
require.True(t, stillTracked, "an in-progress relay dial must not be evicted by cleanup")
|
||||
|
||||
// Release the hanging dial so the goroutine can exit cleanly.
|
||||
mCancel()
|
||||
select {
|
||||
case <-dialDone:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("openConnVia did not return after context cancellation")
|
||||
}
|
||||
}
|
||||
91
shared/relay/client/manager_relaystates_test.go
Normal file
91
shared/relay/client/manager_relaystates_test.go
Normal file
@@ -0,0 +1,91 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"net/netip"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// stallingRelayListener accepts TCP connections and holds them open without ever
|
||||
// responding, so a relay handshake dialed against it blocks until its context is
|
||||
// cancelled. It returns the "rel://host:port" URL to dial.
|
||||
func stallingRelayListener(t *testing.T) string {
|
||||
t.Helper()
|
||||
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
|
||||
var mu sync.Mutex
|
||||
var conns []net.Conn
|
||||
go func() {
|
||||
for {
|
||||
c, err := ln.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
mu.Lock()
|
||||
conns = append(conns, c)
|
||||
mu.Unlock()
|
||||
}
|
||||
}()
|
||||
t.Cleanup(func() {
|
||||
_ = ln.Close()
|
||||
mu.Lock()
|
||||
for _, c := range conns {
|
||||
_ = c.Close()
|
||||
}
|
||||
mu.Unlock()
|
||||
})
|
||||
|
||||
return "rel://" + ln.Addr().String()
|
||||
}
|
||||
|
||||
// TestRelayStates_DoesNotBlockOnRealHangingDial is a regression test for
|
||||
// RelayStates() called by a "status -d command" hanging behind an in-progress
|
||||
// relay dial.
|
||||
func TestRelayStates_DoesNotBlockOnRealHangingDial(t *testing.T) {
|
||||
serverAddr := stallingRelayListener(t)
|
||||
|
||||
mCtx, mCancel := context.WithCancel(context.Background())
|
||||
t.Cleanup(mCancel)
|
||||
|
||||
m := NewManager(mCtx, nil, "alice", 1280)
|
||||
|
||||
dialDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(dialDone)
|
||||
_, _ = m.openConnVia(mCtx, serverAddr, "peerKey", netip.Addr{})
|
||||
}()
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
m.relayClientsMutex.RLock()
|
||||
defer m.relayClientsMutex.RUnlock()
|
||||
_, ok := m.relayClients[serverAddr]
|
||||
return ok
|
||||
}, 5*time.Second, 5*time.Millisecond, "relay dial did not start")
|
||||
|
||||
done := make(chan []RelayConnState, 1)
|
||||
go func() {
|
||||
done <- m.RelayStates()
|
||||
}()
|
||||
|
||||
select {
|
||||
case states := <-done:
|
||||
require.Empty(t, states, "a relay still being dialed carries no state and must be omitted")
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("RelayStates blocked on a foreign relay whose Connect() is in progress")
|
||||
}
|
||||
|
||||
// Release the hanging dial so the goroutine can exit cleanly.
|
||||
mCancel()
|
||||
select {
|
||||
case <-dialDone:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("openConnVia did not return after context cancellation")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user