From 20eb1f50e3985bf65c39033566018d0219f388d2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= <zoltan.pmail@gmail.com>
Date: Tue, 23 Jul 2024 23:04:38 +0200
Subject: [PATCH] Fix loop close

---
 client/internal/peer/conn.go          |  2 +-
 client/internal/peer/worker_relay.go  | 19 ++++++++++++++-----
 client/internal/wgproxy/proxy_ebpf.go |  2 +-
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/client/internal/peer/conn.go b/client/internal/peer/conn.go
index 561437f04..4610afba2 100644
--- a/client/internal/peer/conn.go
+++ b/client/internal/peer/conn.go
@@ -337,7 +337,7 @@ func (conn *Conn) reconnectLoopWithRetry() {
 			}
 
 			// checks if there is peer connection is established via relay or ice
-			conn.log.Infof("ticker timedout, relay state: %s, ice state: %s", conn.statusRelay, conn.statusICE)
+			conn.log.Tracef("ticker timedout, relay state: %s, ice state: %s", conn.statusRelay, conn.statusICE)
 			if conn.isConnected() {
 				continue
 			}
diff --git a/client/internal/peer/worker_relay.go b/client/internal/peer/worker_relay.go
index f4acdba6a..2b90430f0 100644
--- a/client/internal/peer/worker_relay.go
+++ b/client/internal/peer/worker_relay.go
@@ -29,17 +29,20 @@ type WorkerRelayCallbacks struct {
 }
 
 type WorkerRelay struct {
-	ctx          context.Context
+	parentCtx    context.Context
 	log          *log.Entry
 	config       ConnConfig
 	wgInterface  iface.IWGIface
 	relayManager relayClient.ManagerService
 	conn         WorkerRelayCallbacks
+
+	ctx       context.Context
+	ctxCancel context.CancelFunc
 }
 
 func NewWorkerRelay(ctx context.Context, log *log.Entry, config ConnConfig, relayManager relayClient.ManagerService, callbacks WorkerRelayCallbacks) *WorkerRelay {
 	return &WorkerRelay{
-		ctx:          ctx,
+		parentCtx:    ctx,
 		log:          log,
 		config:       config,
 		relayManager: relayManager,
@@ -62,8 +65,10 @@ func (w *WorkerRelay) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 
 	srv := w.preferredRelayServer(currentRelayAddress, remoteOfferAnswer.RelaySrvAddress)
 
-	relayedConn, err := w.relayManager.OpenConn(srv, w.config.Key, w.conn.OnDisconnected)
+	w.ctx, w.ctxCancel = context.WithCancel(w.parentCtx)
+	relayedConn, err := w.relayManager.OpenConn(srv, w.config.Key, w.disconnected)
 	if err != nil {
+		w.ctxCancel()
 		// todo handle all type errors
 		if errors.Is(err, relayClient.ErrConnAlreadyExists) {
 			w.log.Infof("do not need to reopen relay connection")
@@ -99,7 +104,6 @@ func (w *WorkerRelay) RelayIsSupportedLocally() bool {
 func (w *WorkerRelay) wgStateCheck(conn net.Conn) {
 	timer := time.NewTimer(wgHandshakeOvertime)
 	defer timer.Stop()
-
 	for {
 		select {
 		case <-timer.C:
@@ -108,7 +112,7 @@ func (w *WorkerRelay) wgStateCheck(conn net.Conn) {
 				w.log.Errorf("failed to read wg stats: %v", err)
 				continue
 			}
-			log.Infof("last handshake: %v", lastHandshake)
+			w.log.Tracef("last handshake: %v", lastHandshake)
 
 			if time.Since(lastHandshake) > wgHandshakePeriod {
 				w.log.Infof("Wireguard handshake timed out, closing relay connection")
@@ -145,3 +149,8 @@ func (w *WorkerRelay) wgState() (time.Time, error) {
 	}
 	return wgState.LastHandshake, nil
 }
+
+func (w *WorkerRelay) disconnected() {
+	w.ctxCancel()
+	w.conn.OnDisconnected()
+}
diff --git a/client/internal/wgproxy/proxy_ebpf.go b/client/internal/wgproxy/proxy_ebpf.go
index bbaafa756..969a836ff 100644
--- a/client/internal/wgproxy/proxy_ebpf.go
+++ b/client/internal/wgproxy/proxy_ebpf.go
@@ -206,7 +206,7 @@ func (p *WGEBPFProxy) storeTurnConn(turnConn net.Conn) (uint16, error) {
 }
 
 func (p *WGEBPFProxy) removeTurnConn(turnConnID uint16) {
-	log.Tracef("remove turn conn from store by port: %d", turnConnID)
+	log.Debugf("remove turn conn from store by port: %d", turnConnID)
 	p.turnConnMutex.Lock()
 	defer p.turnConnMutex.Unlock()
 	delete(p.turnConnStore, turnConnID)