[client] Refactor WG endpoint setup with role-based proxy activation

client/internal/peer/conn.go

@@ -434,14 +434,14 @@ func (conn *Conn) onICEStateDisconnected(sessionChanged bool) {
 			conn.resetEndpoint()
 		}
 
+		// todo consider to move after the ConfigureWGEndpoint
 		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); 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())
@@ -503,20 +503,22 @@ func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) {
 		return
 	}
 
-	wgProxy.Work()
-	presharedKey := conn.presharedKey(rci.rosenpassPubKey)
+	controller := isController(conn.config)
 
+	if controller {
+		wgProxy.Work()
+	}
 	conn.enableWgWatcherIfNeeded()
-
-	if err := conn.endpointUpdater.ConfigureWGEndpoint(wgProxy.EndpointAddr(), presharedKey); err != nil {
+	if err := conn.endpointUpdater.ConfigureWGEndpoint(wgProxy.EndpointAddr(), conn.presharedKey(rci.rosenpassPubKey)); 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()
+	if !controller {
+		wgProxy.Work()
+	}
 	conn.rosenpassRemoteKey = rci.rosenpassPubKey
 	conn.currentConnPriority = conntype.Relay
 	conn.statusRelay.SetConnected()
@@ -877,9 +879,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

@@ -34,28 +34,27 @@ 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 {
 	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)
 	}
 
+	e.log.Debugf("configure up WireGuard as responder")
+	return e.configureAsResponder(addr, presharedKey)
+}
+
+func (e *EndpointUpdater) SwitchWGEndpoint(addr *net.UDPAddr, presharedKey *wgtypes.Key) 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)
-
-	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 {
@@ -67,9 +66,37 @@ func (e *EndpointUpdater) RemoveWgPeer() error {
 }
 
 func (e *EndpointUpdater) RemoveEndpointAddress() error {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	e.waitForCloseTheDelayedUpdate()
 	return e.wgConfig.WgInterface.RemoveEndpointAddress(e.wgConfig.RemoteKey)
 }
 
+func (e *EndpointUpdater) configureAsInitiator(addr *net.UDPAddr, presharedKey *wgtypes.Key) error {
+	if err := e.updateWireGuardPeer(addr, presharedKey); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (e *EndpointUpdater) configureAsResponder(addr *net.UDPAddr, presharedKey *wgtypes.Key) error {
+	// prevent to run new update while cancel the previous update
+	e.waitForCloseTheDelayedUpdate()
+
+	e.log.Debugf("configure up WireGuard and wait for handshake")
+	var ctx context.Context
+	ctx, e.cancelFunc = context.WithCancel(context.Background())
+	e.updateWg.Add(1)
+	go e.scheduleDelayedUpdate(ctx, addr, presharedKey)
+
+	if err := e.updateWireGuardPeer(nil, presharedKey); err != nil {
+		e.waitForCloseTheDelayedUpdate()
+		return err
+	}
+	return nil
+}
+
 func (e *EndpointUpdater) waitForCloseTheDelayedUpdate() {
 	if e.cancelFunc == nil {
 		return
@@ -105,3 +132,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/coreos/go-oidc/v3 v3.14.1
 	github.com/creack/pty v1.1.24

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=