From 8425660074d5947f6e4bd0028031749333f34650 Mon Sep 17 00:00:00 2001
From: Grzegorz Baranski <grzegorz.baranski@nordsec.com>
Date: Wed, 15 Jan 2025 12:01:52 +0100
Subject: [PATCH] fix: reorganize the job actual ip related functions

---
 daemon/job_actual_ip.go      | 70 ++++++++++++++++++++++++++++--------
 daemon/job_actual_ip_test.go |  9 ++---
 daemon/jobs.go               | 42 +---------------------
 3 files changed, 59 insertions(+), 62 deletions(-)

diff --git a/daemon/job_actual_ip.go b/daemon/job_actual_ip.go
index 6621aece..c6c3f0a9 100644
--- a/daemon/job_actual_ip.go
+++ b/daemon/job_actual_ip.go
@@ -7,6 +7,8 @@ import (
 	"time"
 
 	"github.com/NordSecurity/nordvpn-linux/core"
+	"github.com/NordSecurity/nordvpn-linux/daemon/state"
+	"github.com/NordSecurity/nordvpn-linux/events"
 	"github.com/NordSecurity/nordvpn-linux/internal"
 	"github.com/NordSecurity/nordvpn-linux/network"
 )
@@ -64,25 +66,65 @@ func insightsIPUntilSuccess(ctx context.Context, api core.InsightsAPI, backoff f
 	}
 }
 
-func JobActualIP(dm *DataManager, api core.InsightsAPI) func(context.Context, bool) error {
-	return func(ctx context.Context, isConnected bool) error {
-		var newIP netip.Addr
-		defer func() {
-			dm.SetActualIP(newIP)
-		}()
+func updateActualIP(dm *DataManager, api core.InsightsAPI, ctx context.Context, isConnected bool) error {
+	var newIP netip.Addr
+	defer func() {
+		dm.SetActualIP(newIP)
+	}()
 
-		if !isConnected {
-			return nil
-		}
+	if !isConnected {
+		return nil
+	}
 
-		insightsIP, err := insightsIPUntilSuccess(ctx, api, network.ExponentialBackoff)
+	insightsIP, err := insightsIPUntilSuccess(ctx, api, network.ExponentialBackoff)
+	if err != nil {
+		return err
+	}
+	if insightsIP.IsValid() {
+		newIP = insightsIP
+	}
+
+	return nil
+}
+
+// JobActualIP is a long-running job that will update the actual IP address indefinitely
+// it reacts to state updates from the statePublisher
+func JobActualIP(statePublisher *state.StatePublisher, dm *DataManager, api core.InsightsAPI) {
+	call := func(ctx context.Context, isConnected bool) {
+		err := updateActualIP(dm, api, ctx, isConnected)
 		if err != nil {
-			return err
+			if err == context.Canceled {
+				return
+			}
+			log.Println(internal.ErrorPrefix, "actual ip job error: ", err)
 		}
-		if insightsIP.IsValid() {
-			newIP = insightsIP
+	}
+
+	stateChan, _ := statePublisher.AddSubscriber()
+	var cancel context.CancelFunc
+
+	for ev := range stateChan {
+		_, isConnect := ev.(events.DataConnect)
+		_, isDisconnect := ev.(events.DataDisconnect)
+
+		if isConnect || isDisconnect {
+			if cancel != nil {
+				cancel()
+			}
+
+			var ctx context.Context
+			ctx, cancel = context.WithCancel(context.Background())
+
+			if isConnect {
+				go call(ctx, true)
+			} else {
+				call(ctx, false) // should finish immediately, that's why it's not a separate goroutine
+			}
 		}
+	}
 
-		return nil
+	// Ensure the context is canceled when the loop exits
+	if cancel != nil {
+		cancel()
 	}
 }
diff --git a/daemon/job_actual_ip_test.go b/daemon/job_actual_ip_test.go
index 3579c21e..db7fbdbf 100644
--- a/daemon/job_actual_ip_test.go
+++ b/daemon/job_actual_ip_test.go
@@ -13,10 +13,6 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-// func TestJobActualIP(t *testing.T) {
-
-// }
-
 type mockInsights struct {
 	insightsFunc insightFunc
 }
@@ -130,7 +126,7 @@ func TestInsightsIPUntilSuccess(t *testing.T) {
 	}
 }
 
-func TestJobActualIP(t *testing.T) {
+func TestUpdateActualIP(t *testing.T) {
 	tests := []struct {
 		name        string
 		isConnected bool
@@ -177,8 +173,7 @@ func TestJobActualIP(t *testing.T) {
 				insightsFunc: tt.insightsAPI,
 			}
 
-			job := JobActualIP(dm, api)
-			err := job(ctx, tt.isConnected)
+			err := updateActualIP(dm, api, ctx, tt.isConnected)
 
 			assert.ErrorIs(t, err, tt.expectedErr)
 
diff --git a/daemon/jobs.go b/daemon/jobs.go
index 5a17e757..6d3446fb 100644
--- a/daemon/jobs.go
+++ b/daemon/jobs.go
@@ -75,47 +75,7 @@ func (r *RPC) StartJobs(
 		}
 	}
 
-	actualIP := JobActualIP(r.dm, r.api)
-
-	go func() {
-		call := func(ctx context.Context, isConnected bool) {
-			err := actualIP(ctx, isConnected)
-			if err != nil {
-				if err == context.Canceled {
-					return
-				}
-				log.Println(internal.ErrorPrefix, "actual ip job error: ", err)
-			}
-		}
-
-		stateChan, _ := statePublisher.AddSubscriber()
-		var cancel context.CancelFunc
-
-		for ev := range stateChan {
-			_, isConnect := ev.(events.DataConnect)
-			_, isDisconnect := ev.(events.DataDisconnect)
-
-			if isConnect || isDisconnect {
-				if cancel != nil {
-					cancel()
-				}
-
-				var ctx context.Context
-				ctx, cancel = context.WithCancel(context.Background())
-
-				if isConnect {
-					go call(ctx, true)
-				} else {
-					call(ctx, false) // should finish immediately, that's why it's not a separate goroutine
-				}
-			}
-		}
-
-		// Ensure the context is canceled when the loop exits
-		if cancel != nil {
-			cancel()
-		}
-	}()
+	go JobActualIP(statePublisher, r.dm, r.api)
 
 	go func() {
 		stateChan, _ := statePublisher.AddSubscriber()