[client] Extend WG watcher for ICE connection too (#5133)

Extend WG watcher for ICE connection too

client/internal/peer/wg_watcher.go

@@ -30,10 +30,8 @@ type WGWatcher struct {
 	peerKey       string
 	stateDump     *stateDump
 
-	ctx         context.Context
-	ctxCancel   context.CancelFunc
-	ctxLock     sync.Mutex
-	enabledTime time.Time
+	enabled   bool
+	muEnabled sync.RWMutex
 }
 
 func NewWGWatcher(log *log.Entry, wgIfaceStater WGInterfaceStater, peerKey string, stateDump *stateDump) *WGWatcher {
@@ -46,52 +44,44 @@ func NewWGWatcher(log *log.Entry, wgIfaceStater WGInterfaceStater, peerKey strin
 }
 
 // EnableWgWatcher starts the WireGuard watcher. If it is already enabled, it will return immediately and do nothing.
-func (w *WGWatcher) EnableWgWatcher(parentCtx context.Context, onDisconnectedFn func()) {
-	w.log.Debugf("enable WireGuard watcher")
-	w.ctxLock.Lock()
-	w.enabledTime = time.Now()
-
-	if w.ctx != nil && w.ctx.Err() == nil {
-		w.log.Errorf("WireGuard watcher already enabled")
-		w.ctxLock.Unlock()
+// The watcher runs until ctx is cancelled. Caller is responsible for context lifecycle management.
+func (w *WGWatcher) EnableWgWatcher(ctx context.Context, onDisconnectedFn func()) {
+	w.muEnabled.Lock()
+	if w.enabled {
+		w.muEnabled.Unlock()
 		return
 	}
 
-	ctx, ctxCancel := context.WithCancel(parentCtx)
-	w.ctx = ctx
-	w.ctxCancel = ctxCancel
-	w.ctxLock.Unlock()
+	w.log.Debugf("enable WireGuard watcher")
+	enabledTime := time.Now()
+	w.enabled = true
+	w.muEnabled.Unlock()
 
 	initialHandshake, err := w.wgState()
 	if err != nil {
 		w.log.Warnf("failed to read initial wg stats: %v", err)
 	}
 
-	w.periodicHandshakeCheck(ctx, ctxCancel, onDisconnectedFn, initialHandshake)
+	w.periodicHandshakeCheck(ctx, onDisconnectedFn, enabledTime, initialHandshake)
+
+	w.muEnabled.Lock()
+	w.enabled = false
+	w.muEnabled.Unlock()
 }
 
-// DisableWgWatcher stops the WireGuard watcher and wait for the watcher to exit
-func (w *WGWatcher) DisableWgWatcher() {
-	w.ctxLock.Lock()
-	defer w.ctxLock.Unlock()
-
-	if w.ctxCancel == nil {
-		return
-	}
-
-	w.log.Debugf("disable WireGuard watcher")
-
-	w.ctxCancel()
-	w.ctxCancel = nil
+// IsEnabled returns true if the WireGuard watcher is currently enabled
+func (w *WGWatcher) IsEnabled() bool {
+	w.muEnabled.RLock()
+	defer w.muEnabled.RUnlock()
+	return w.enabled
 }
 
 // wgStateCheck help to check the state of the WireGuard handshake and relay connection
-func (w *WGWatcher) periodicHandshakeCheck(ctx context.Context, ctxCancel context.CancelFunc, onDisconnectedFn func(), initialHandshake time.Time) {
+func (w *WGWatcher) periodicHandshakeCheck(ctx context.Context, onDisconnectedFn func(), enabledTime time.Time, initialHandshake time.Time) {
 	w.log.Infof("WireGuard watcher started")
 
 	timer := time.NewTimer(wgHandshakeOvertime)
 	defer timer.Stop()
-	defer ctxCancel()
 
 	lastHandshake := initialHandshake
 
@@ -104,7 +94,7 @@ func (w *WGWatcher) periodicHandshakeCheck(ctx context.Context, ctxCancel contex
 				return
 			}
 			if lastHandshake.IsZero() {
-				elapsed := handshake.Sub(w.enabledTime).Seconds()
+				elapsed := calcElapsed(enabledTime, *handshake)
 				w.log.Infof("first wg handshake detected within: %.2fsec, (%s)", elapsed, handshake)
 			}
 
@@ -134,19 +124,19 @@ func (w *WGWatcher) handshakeCheck(lastHandshake time.Time) (*time.Time, bool) {
 
 	// the current know handshake did not change
 	if handshake.Equal(lastHandshake) {
-		w.log.Warnf("WireGuard handshake timed out, closing relay connection: %v", handshake)
+		w.log.Warnf("WireGuard handshake timed out: %v", handshake)
 		return nil, false
 	}
 
 	// in case if the machine is suspended, the handshake time will be in the past
 	if handshake.Add(checkPeriod).Before(time.Now()) {
-		w.log.Warnf("WireGuard handshake timed out, closing relay connection: %v", handshake)
+		w.log.Warnf("WireGuard handshake timed out: %v", handshake)
 		return nil, false
 	}
 
 	// error handling for handshake time in the future
 	if handshake.After(time.Now()) {
-		w.log.Warnf("WireGuard handshake is in the future, closing relay connection: %v", handshake)
+		w.log.Warnf("WireGuard handshake is in the future: %v", handshake)
 		return nil, false
 	}
 
@@ -164,3 +154,13 @@ func (w *WGWatcher) wgState() (time.Time, error) {
 	}
 	return wgState.LastHandshake, nil
 }
+
+// calcElapsed calculates elapsed time since watcher was enabled.
+// The watcher started after the wg configuration happens, because of this need to normalise the negative value
+func calcElapsed(enabledTime, handshake time.Time) float64 {
+	elapsed := handshake.Sub(enabledTime).Seconds()
+	if elapsed < 0 {
+		elapsed = 0
+	}
+	return elapsed
+}