Refactor WG endpoint setup with role-based proxy activation

Move wgProxy.Work() and wgConfigWorkaround() into EndpointUpdater to
centralize the timing of proxy activation relative to WireGuard peer
updates. The initiator calls Work() before updating the peer (so the
proxy is ready for the immediate handshake), while the responder calls
Work() after (to avoid premature activation).

Add SwitchWGEndpoint for relay fallback on ICE disconnect, which skips
the initiator/responder distinction since WireGuard roaming wil
overwrite the endpoint regardless.

Author: pappz

client/internal/peer/conn.go

@@ -385,19 +385,14 @@ func (conn *Conn) onICEConnectionIsReady(priority conntype.ConnPriority, iceConn
 		conn.wgProxyRelay.Pause()
 	}
 
-	if wgProxy != nil {
-		wgProxy.Work()
-	}
-
 	conn.Log.Infof("configure WireGuard endpoint to: %s", ep.String())
 	conn.enableWgWatcherIfNeeded()
 
 	presharedKey := conn.presharedKey(iceConnInfo.RosenpassPubKey)
-	if err = conn.endpointUpdater.ConfigureWGEndpoint(ep, presharedKey); err != nil {
+	if err = conn.endpointUpdater.ConfigureWGEndpoint(ep, presharedKey, wgProxy); err != nil {
 		conn.handleConfigurationFailure(err, wgProxy)
 		return
 	}
-	wgConfigWorkaround()
 
 	if conn.wgProxyRelay != nil {
 		conn.Log.Debugf("redirect packets from relayed conn to WireGuard")
@@ -430,14 +425,10 @@ func (conn *Conn) onICEStateDisconnected() {
 	if conn.isReadyToUpgrade() {
 		conn.Log.Infof("ICE disconnected, set Relay to active connection")
 		conn.dumpState.SwitchToRelay()
-		conn.wgProxyRelay.Work()
-
 		presharedKey := conn.presharedKey(conn.rosenpassRemoteKey)
-		if err := conn.endpointUpdater.ConfigureWGEndpoint(conn.wgProxyRelay.EndpointAddr(), presharedKey); err != nil {
+		if err := conn.endpointUpdater.SwitchWGEndpoint(conn.wgProxyRelay.EndpointAddr(), presharedKey, conn.wgProxyRelay); err != nil {
 			conn.Log.Errorf("failed to switch to relay conn: %v", err)
 		}
-
-		conn.wgProxyRelay.Work()
 		conn.currentConnPriority = conntype.Relay
 	} else {
 		conn.Log.Infof("ICE disconnected, do not switch to Relay. Reset priority to: %s", conntype.None.String())
@@ -499,20 +490,15 @@ func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) {
 		return
 	}
 
-	wgProxy.Work()
-	presharedKey := conn.presharedKey(rci.rosenpassPubKey)
-
 	conn.enableWgWatcherIfNeeded()
 
-	if err := conn.endpointUpdater.ConfigureWGEndpoint(wgProxy.EndpointAddr(), presharedKey); err != nil {
+	if err := conn.endpointUpdater.ConfigureWGEndpoint(wgProxy.EndpointAddr(), conn.presharedKey(rci.rosenpassPubKey), wgProxy); err != nil {
 		if err := wgProxy.CloseConn(); err != nil {
 			conn.Log.Warnf("Failed to close relay connection: %v", err)
 		}
 		conn.Log.Errorf("Failed to update WireGuard peer configuration: %v", err)
 		return
 	}
-
-	wgConfigWorkaround()
 	conn.rosenpassRemoteKey = rci.rosenpassPubKey
 	conn.currentConnPriority = conntype.Relay
 	conn.statusRelay.SetConnected()
@@ -862,9 +848,3 @@ func isController(config ConnConfig) bool {
 func isRosenpassEnabled(remoteRosenpassPubKey []byte) bool {
 	return remoteRosenpassPubKey != nil
 }
-
-// wgConfigWorkaround is a workaround for the issue with WireGuard configuration update
-// When update a peer configuration in near to each other time, the second update can be ignored by WireGuard
-func wgConfigWorkaround() {
-	time.Sleep(100 * time.Millisecond)
-}

client/internal/peer/endpoint.go

@@ -8,6 +8,8 @@ import (
 
 	"github.com/sirupsen/logrus"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+
+	"github.com/netbirdio/netbird/client/iface/wgproxy"
 )
 
 const (
@@ -37,25 +39,36 @@ func NewEndpointUpdater(log *logrus.Entry, wgConfig WgConfig, initiator bool) *E
 // ConfigureWGEndpoint sets up the WireGuard endpoint configuration.
 // The initiator immediately configures the endpoint, while the non-initiator
 // waits for a fallback period before configuring to avoid handshake congestion.
-func (e *EndpointUpdater) ConfigureWGEndpoint(addr *net.UDPAddr, presharedKey *wgtypes.Key) error {
+//
+// If wgProxy is non-nil, Work() is called at the correct time relative to the
+// WireGuard peer update: before for the initiator (so the proxy is ready to
+// process the immediate handshake) and after for the non-initiator.
+func (e *EndpointUpdater) ConfigureWGEndpoint(addr *net.UDPAddr, presharedKey *wgtypes.Key, wgProxy wgproxy.Proxy) error {
 	e.mu.Lock()
 	defer e.mu.Unlock()
 
 	if e.initiator {
-		e.log.Debugf("configure up WireGuard as initiatr")
-		return e.updateWireGuardPeer(addr, presharedKey)
+		e.log.Debugf("configure up WireGuard as initiator")
+		return e.configureAsInitiator(addr, presharedKey, wgProxy)
 	}
 
+	e.log.Debugf("configure up WireGuard as responder")
+	return e.configureAsResponder(addr, presharedKey, wgProxy)
+}
+
+func (e *EndpointUpdater) SwitchWGEndpoint(addr *net.UDPAddr, presharedKey *wgtypes.Key, wgProxy wgproxy.Proxy) error {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
 	// prevent to run new update while cancel the previous update
 	e.waitForCloseTheDelayedUpdate()
 
-	var ctx context.Context
-	ctx, e.cancelFunc = context.WithCancel(context.Background())
-	e.updateWg.Add(1)
-	go e.scheduleDelayedUpdate(ctx, addr, presharedKey)
+	// Makes no sense to distinguish initiator and responder because the roaming feature in WG will overwrite the endpoint
+	if wgProxy != nil {
+		wgProxy.Work()
+	}
 
-	e.log.Debugf("configure up WireGuard and wait for handshake")
-	return e.updateWireGuardPeer(nil, presharedKey)
+	return e.updateWireGuardPeer(addr, presharedKey)
 }
 
 func (e *EndpointUpdater) RemoveWgPeer() error {
@@ -66,6 +79,41 @@ func (e *EndpointUpdater) RemoveWgPeer() error {
 	return e.wgConfig.WgInterface.RemovePeer(e.wgConfig.RemoteKey)
 }
 
+func (e *EndpointUpdater) configureAsInitiator(addr *net.UDPAddr, presharedKey *wgtypes.Key, wgProxy wgproxy.Proxy) error {
+	if wgProxy != nil {
+		wgProxy.Work()
+	}
+
+	if err := e.updateWireGuardPeer(addr, presharedKey); err != nil {
+		return err
+	}
+
+	wgConfigWorkaround()
+	return nil
+}
+
+func (e *EndpointUpdater) configureAsResponder(addr *net.UDPAddr, presharedKey *wgtypes.Key, wgProxy wgproxy.Proxy) error {
+	// prevent to run new update while cancel the previous update
+	e.waitForCloseTheDelayedUpdate()
+
+	var ctx context.Context
+	ctx, e.cancelFunc = context.WithCancel(context.Background())
+	e.updateWg.Add(1)
+	go e.scheduleDelayedUpdate(ctx, addr, presharedKey)
+
+	e.log.Debugf("configure up WireGuard and wait for handshake")
+	if err := e.updateWireGuardPeer(nil, presharedKey); err != nil {
+		e.waitForCloseTheDelayedUpdate()
+		return err
+	}
+	wgConfigWorkaround()
+
+	if wgProxy != nil {
+		wgProxy.Work()
+	}
+	return nil
+}
+
 func (e *EndpointUpdater) waitForCloseTheDelayedUpdate() {
 	if e.cancelFunc == nil {
 		return
@@ -101,3 +149,9 @@ func (e *EndpointUpdater) updateWireGuardPeer(endpoint *net.UDPAddr, presharedKe
 		presharedKey,
 	)
 }
+
+// wgConfigWorkaround is a workaround for the issue with WireGuard configuration update
+// When update a peer configuration in near to each other time, the second update can be ignored by WireGuard
+func wgConfigWorkaround() {
+	time.Sleep(100 * time.Millisecond)
+}

go.mod

@@ -40,7 +40,7 @@ require (
 	github.com/c-robinson/iplib v1.0.3
 	github.com/caddyserver/certmagic v0.21.3
 	github.com/cilium/ebpf v0.15.0
-	github.com/coder/websocket v1.8.13
+	github.com/coder/websocket v1.8.14
 	github.com/coreos/go-iptables v0.7.0
 	github.com/creack/pty v1.1.24
 	github.com/dexidp/dex v0.0.0-00010101000000-000000000000

go.sum

@@ -107,8 +107,8 @@ github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UF
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cilium/ebpf v0.15.0 h1:7NxJhNiBT3NG8pZJ3c+yfrVdHY8ScgKD27sScgjLMMk=
 github.com/cilium/ebpf v0.15.0/go.mod h1:DHp1WyrLeiBh19Cf/tfiSMhqheEiK8fXFZ4No0P1Hso=
-github.com/coder/websocket v1.8.13 h1:f3QZdXy7uGVz+4uCJy2nTZyM0yTBj8yANEHhqlXZ9FE=
-github.com/coder/websocket v1.8.13/go.mod h1:LNVeNrXQZfe5qhS9ALED3uA+l5pPqvwXg3CKoDBB2gs=
+github.com/coder/websocket v1.8.14 h1:9L0p0iKiNOibykf283eHkKUHHrpG7f65OE3BhhO7v9g=
+github.com/coder/websocket v1.8.14/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg=
 github.com/containerd/containerd v1.7.29 h1:90fWABQsaN9mJhGkoVnuzEY+o1XDPbg9BTC9QTAHnuE=
 github.com/containerd/containerd v1.7.29/go.mod h1:azUkWcOvHrWvaiUjSQH0fjzuHIwSPg1WL5PshGP4Szs=
 github.com/containerd/log v0.1.0 h1:TCJt7ioM2cr/tfR8GPbGf9/VRAX8D2B4PjzCpfX540I=