From 28cbb4b70f370e0d3d23818fc8246ceee0b358c8 Mon Sep 17 00:00:00 2001
From: Zoltan Papp <zoltan.pmail@gmail.com>
Date: Tue, 17 Sep 2024 12:10:17 +0200
Subject: [PATCH] [client] Cancel the context of wg watcher when the go routine
 exit (#2612)

---
 client/internal/peer/conn.go         | 4 ++--
 client/internal/peer/worker_relay.go | 9 +++++----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/client/internal/peer/conn.go b/client/internal/peer/conn.go
index 9eb881087..5327e31d2 100644
--- 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
index 3457faa46..6bb385d3e 100644
--- 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)