[client] Cancel the context of wg watcher when the go routine exit (#…

…2612)
netbirdio · Sep 17, 2024 · 28cbb4b · 28cbb4b
1 parent 1104c9c
commit 28cbb4b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/client/internal/peer/conn.go b/client/internal/peer/conn.go
@@ -484,11 +484,11 @@ func (conn *Conn) onWorkerICEStateDisconnected(newState ConnStatus) {
 	// switch back to relay connection
 	if conn.endpointRelay != nil && conn.currentConnPriority != connPriorityRelay {
 		conn.log.Debugf("ICE disconnected, set Relay to active connection")
-		conn.workerRelay.EnableWgWatcher(conn.ctx)
 		err := conn.configureWGEndpoint(conn.endpointRelay)
 		if err != nil {
 			conn.log.Errorf("failed to switch to relay conn: %v", err)
 		}
+		conn.workerRelay.EnableWgWatcher(conn.ctx)
 		conn.currentConnPriority = connPriorityRelay
 	}
 
@@ -551,7 +551,6 @@ func (conn *Conn) relayConnectionIsReady(rci RelayConnInfo) {
 		}
 	}
 
-	conn.workerRelay.EnableWgWatcher(conn.ctx)
 	err = conn.configureWGEndpoint(endpointUdpAddr)
 	if err != nil {
 		if err := wgProxy.CloseConn(); err != nil {
@@ -560,6 +559,7 @@ func (conn *Conn) relayConnectionIsReady(rci RelayConnInfo) {
 		conn.log.Errorf("Failed to update wg peer configuration: %v", err)
 		return
 	}
+	conn.workerRelay.EnableWgWatcher(conn.ctx)
 	wgConfigWorkaround()
 
 	if conn.wgProxyRelay != nil {

diff --git a/client/internal/peer/worker_relay.go b/client/internal/peer/worker_relay.go
@@ -109,10 +109,10 @@ func (w *WorkerRelay) EnableWgWatcher(ctx context.Context) {
 	}
 
 	ctx, ctxCancel := context.WithCancel(ctx)
-	w.wgStateCheck(ctx)
 	w.ctxWgWatch = ctx
 	w.ctxCancelWgWatch = ctxCancel
 
+	w.wgStateCheck(ctx, ctxCancel)
 }
 
 func (w *WorkerRelay) DisableWgWatcher() {
@@ -158,21 +158,22 @@ func (w *WorkerRelay) CloseConn() {
 }
 
 // wgStateCheck help to check the state of the WireGuard handshake and relay connection
-func (w *WorkerRelay) wgStateCheck(ctx context.Context) {
+func (w *WorkerRelay) wgStateCheck(ctx context.Context, ctxCancel context.CancelFunc) {
+	w.log.Debugf("WireGuard watcher started")
 	lastHandshake, err := w.wgState()
 	if err != nil {
-		w.log.Errorf("failed to read wg stats: %v", err)
+		w.log.Warnf("failed to read wg stats: %v", err)
 		lastHandshake = time.Time{}
 	}
 
 	go func(lastHandshake time.Time) {
 		timer := time.NewTimer(wgHandshakeOvertime)
 		defer timer.Stop()
+		defer ctxCancel()
 
 		for {
 			select {
 			case <-timer.C:
-
 				handshake, err := w.wgState()
 				if err != nil {
 					w.log.Errorf("failed to read wg stats: %v", err)