Fix connstate indication (#732)

Fix the status indication in the client service. The status of the
management server and the signal server was incorrect if the network
connection was broken. Basically the status update was not used by
the management and signal library.

client/internal/connect.go

@@ -58,8 +58,7 @@ func RunClient(ctx context.Context, config *Config, statusRecorder *peer.Status)
 		return err
 	}
 
-	managementURL := config.ManagementURL.String()
-	statusRecorder.MarkManagementDisconnected(managementURL)
+	statusRecorder.MarkManagementDisconnected()
 
 	operation := func() error {
 		// if context cancelled we not start new backoff cycle
@@ -73,13 +72,16 @@ func RunClient(ctx context.Context, config *Config, statusRecorder *peer.Status)
 
 		engineCtx, cancel := context.WithCancel(ctx)
 		defer func() {
-			statusRecorder.MarkManagementDisconnected(managementURL)
+			statusRecorder.MarkManagementDisconnected()
 			statusRecorder.CleanLocalPeerState()
 			cancel()
 		}()
 
 		log.Debugf("conecting to the Management service %s", config.ManagementURL.Host)
 		mgmClient, err := mgm.NewClient(engineCtx, config.ManagementURL.Host, myPrivateKey, mgmTlsEnabled)
+		mgmNotifier := statusRecorderToMgmConnStateNotifier(statusRecorder)
+		mgmClient.SetConnStateListener(mgmNotifier)
+
 		if err != nil {
 			return wrapErr(gstatus.Errorf(codes.FailedPrecondition, "failed connecting to Management Service : %s", err))
 		}
@@ -101,7 +103,7 @@ func RunClient(ctx context.Context, config *Config, statusRecorder *peer.Status)
 			}
 			return wrapErr(err)
 		}
-		statusRecorder.MarkManagementConnected(managementURL)
+		statusRecorder.MarkManagementConnected()
 
 		localPeerState := peer.LocalPeerState{
 			IP:              loginResp.GetPeerConfig().GetAddress(),
@@ -117,8 +119,10 @@ func RunClient(ctx context.Context, config *Config, statusRecorder *peer.Status)
 			loginResp.GetWiretrusteeConfig().GetSignal().GetUri(),
 		)
 
-		statusRecorder.MarkSignalDisconnected(signalURL)
-		defer statusRecorder.MarkSignalDisconnected(signalURL)
+		statusRecorder.UpdateSignalAddress(signalURL)
+
+		statusRecorder.MarkSignalDisconnected()
+		defer statusRecorder.MarkSignalDisconnected()
 
 		// with the global Wiretrustee config in hand connect (just a connection, no stream yet) Signal
 		signalClient, err := connectToSignal(engineCtx, loginResp.GetWiretrusteeConfig(), myPrivateKey)
@@ -133,7 +137,10 @@ func RunClient(ctx context.Context, config *Config, statusRecorder *peer.Status)
 			}
 		}()
 
-		statusRecorder.MarkSignalConnected(signalURL)
+		signalNotifier := statusRecorderToSignalConnStateNotifier(statusRecorder)
+		signalClient.SetConnStateListener(signalNotifier)
+
+		statusRecorder.MarkSignalConnected()
 
 		peerConfig := loginResp.GetPeerConfig()
 
@@ -320,3 +327,15 @@ func UpdateOldManagementPort(ctx context.Context, config *Config, configPath str
 
 	return config, nil
 }
+
+func statusRecorderToMgmConnStateNotifier(statusRecorder *peer.Status) mgm.ConnStateNotifier {
+	var sri interface{} = statusRecorder
+	mgmNotifier, _ := sri.(mgm.ConnStateNotifier)
+	return mgmNotifier
+}
+
+func statusRecorderToSignalConnStateNotifier(statusRecorder *peer.Status) signal.ConnStateNotifier {
+	var sri interface{} = statusRecorder
+	notifier, _ := sri.(signal.ConnStateNotifier)
+	return notifier
+}

client/internal/engine_test.go

@@ -72,7 +72,7 @@ func TestEngine_SSH(t *testing.T) {
 		WgAddr:       "100.64.0.1/24",
 		WgPrivateKey: key,
 		WgPort:       33100,
-	}, peer.NewRecorder())
+	}, peer.NewRecorder("https://mgm"))
 
 	engine.dnsServer = &dns.MockServer{
 		UpdateDNSServerFunc: func(serial uint64, update nbdns.Config) error { return nil },
@@ -206,7 +206,7 @@ func TestEngine_UpdateNetworkMap(t *testing.T) {
 		WgAddr:       "100.64.0.1/24",
 		WgPrivateKey: key,
 		WgPort:       33100,
-	}, peer.NewRecorder())
+	}, peer.NewRecorder("https://mgm"))
 	engine.wgInterface, err = iface.NewWGIFace("utun102", "100.64.0.1/24", iface.DefaultMTU)
 	engine.routeManager = routemanager.NewManager(ctx, key.PublicKey().String(), engine.wgInterface, engine.statusRecorder)
 	engine.dnsServer = &dns.MockServer{
@@ -390,7 +390,7 @@ func TestEngine_Sync(t *testing.T) {
 		WgAddr:       "100.64.0.1/24",
 		WgPrivateKey: key,
 		WgPort:       33100,
-	}, peer.NewRecorder())
+	}, peer.NewRecorder("https://mgm"))
 
 	engine.dnsServer = &dns.MockServer{
 		UpdateDNSServerFunc: func(serial uint64, update nbdns.Config) error { return nil },
@@ -548,7 +548,7 @@ func TestEngine_UpdateNetworkMapWithRoutes(t *testing.T) {
 				WgAddr:       wgAddr,
 				WgPrivateKey: key,
 				WgPort:       33100,
-			}, peer.NewRecorder())
+			}, peer.NewRecorder("https://mgm"))
 			engine.wgInterface, err = iface.NewWGIFace(wgIfaceName, wgAddr, iface.DefaultMTU)
 			assert.NoError(t, err, "shouldn't return error")
 			input := struct {
@@ -713,7 +713,7 @@ func TestEngine_UpdateNetworkMapWithDNSUpdate(t *testing.T) {
 				WgAddr:       wgAddr,
 				WgPrivateKey: key,
 				WgPort:       33100,
-			}, peer.NewRecorder())
+			}, peer.NewRecorder("https://mgm"))
 			engine.wgInterface, err = iface.NewWGIFace(wgIfaceName, wgAddr, iface.DefaultMTU)
 			assert.NoError(t, err, "shouldn't return error")
 
@@ -978,7 +978,7 @@ func createEngine(ctx context.Context, cancel context.CancelFunc, setupKey strin
 		WgPort:       wgPort,
 	}
 
-	return NewEngine(ctx, cancel, signalClient, mgmtClient, conf, peer.NewRecorder()), nil
+	return NewEngine(ctx, cancel, signalClient, mgmtClient, conf, peer.NewRecorder("https://mgm")), nil
 }
 
 func startSignal() (*grpc.Server, string, error) {

client/internal/peer/conn_test.go

@@ -49,7 +49,7 @@ func TestConn_GetKey(t *testing.T) {
 
 func TestConn_OnRemoteOffer(t *testing.T) {
 
-	conn, err := NewConn(connConf, NewRecorder())
+	conn, err := NewConn(connConf, NewRecorder("https://mgm"))
 	if err != nil {
 		return
 	}
@@ -83,7 +83,7 @@ func TestConn_OnRemoteOffer(t *testing.T) {
 
 func TestConn_OnRemoteAnswer(t *testing.T) {
 
-	conn, err := NewConn(connConf, NewRecorder())
+	conn, err := NewConn(connConf, NewRecorder("https://mgm"))
 	if err != nil {
 		return
 	}
@@ -116,7 +116,7 @@ func TestConn_OnRemoteAnswer(t *testing.T) {
 }
 func TestConn_Status(t *testing.T) {
 
-	conn, err := NewConn(connConf, NewRecorder())
+	conn, err := NewConn(connConf, NewRecorder("https://mgm"))
 	if err != nil {
 		return
 	}
@@ -143,7 +143,7 @@ func TestConn_Status(t *testing.T) {
 
 func TestConn_Close(t *testing.T) {
 
-	conn, err := NewConn(connConf, NewRecorder())
+	conn, err := NewConn(connConf, NewRecorder("https://mgm"))
 	if err != nil {
 		return
 	}

client/internal/peer/status.go

@@ -49,21 +49,24 @@ type FullStatus struct {
 
 // Status holds a state of peers, signal and management connections
 type Status struct {
-	mux          sync.Mutex
-	peers        map[string]State
-	changeNotify map[string]chan struct{}
-	signal       SignalState
-	management   ManagementState
-	localPeer    LocalPeerState
-	offlinePeers []State
+	mux             sync.Mutex
+	peers           map[string]State
+	changeNotify    map[string]chan struct{}
+	signalState     bool
+	managementState bool
+	localPeer       LocalPeerState
+	offlinePeers    []State
+	mgmAddress      string
+	signalAddress   string
 }
 
 // NewRecorder returns a new Status instance
-func NewRecorder() *Status {
+func NewRecorder(mgmAddress string) *Status {
 	return &Status{
 		peers:        make(map[string]State),
 		changeNotify: make(map[string]chan struct{}),
 		offlinePeers: make([]State, 0),
+		mgmAddress:   mgmAddress,
 	}
 }
 
@@ -193,43 +196,45 @@ func (d *Status) CleanLocalPeerState() {
 }
 
 // MarkManagementDisconnected sets ManagementState to disconnected
-func (d *Status) MarkManagementDisconnected(managementURL string) {
+func (d *Status) MarkManagementDisconnected() {
 	d.mux.Lock()
 	defer d.mux.Unlock()
-	d.management = ManagementState{
-		URL:       managementURL,
-		Connected: false,
-	}
+	d.managementState = false
 }
 
 // MarkManagementConnected sets ManagementState to connected
-func (d *Status) MarkManagementConnected(managementURL string) {
+func (d *Status) MarkManagementConnected() {
 	d.mux.Lock()
 	defer d.mux.Unlock()
-	d.management = ManagementState{
-		URL:       managementURL,
-		Connected: true,
-	}
+	d.managementState = true
+}
+
+// UpdateSignalAddress update the address of the signal server
+func (d *Status) UpdateSignalAddress(signalURL string) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+	d.signalAddress = signalURL
+}
+
+// UpdateManagementAddress update the address of the management server
+func (d *Status) UpdateManagementAddress(mgmAddress string) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+	d.mgmAddress = mgmAddress
 }
 
 // MarkSignalDisconnected sets SignalState to disconnected
-func (d *Status) MarkSignalDisconnected(signalURL string) {
+func (d *Status) MarkSignalDisconnected() {
 	d.mux.Lock()
 	defer d.mux.Unlock()
-	d.signal = SignalState{
-		signalURL,
-		false,
-	}
+	d.signalState = false
 }
 
 // MarkSignalConnected sets SignalState to connected
-func (d *Status) MarkSignalConnected(signalURL string) {
+func (d *Status) MarkSignalConnected() {
 	d.mux.Lock()
 	defer d.mux.Unlock()
-	d.signal = SignalState{
-		signalURL,
-		true,
-	}
+	d.signalState = true
 }
 
 // GetFullStatus gets full status
@@ -238,9 +243,15 @@ func (d *Status) GetFullStatus() FullStatus {
 	defer d.mux.Unlock()
 
 	fullStatus := FullStatus{
-		ManagementState: d.management,
-		SignalState:     d.signal,
-		LocalPeerState:  d.localPeer,
+		ManagementState: ManagementState{
+			d.mgmAddress,
+			d.managementState,
+		},
+		SignalState: SignalState{
+			d.signalAddress,
+			d.signalState,
+		},
+		LocalPeerState: d.localPeer,
 	}
 
 	for _, status := range d.peers {

client/internal/peer/status_test.go

@@ -8,7 +8,7 @@ import (
 
 func TestAddPeer(t *testing.T) {
 	key := "abc"
-	status := NewRecorder()
+	status := NewRecorder("https://mgm")
 	err := status.AddPeer(key)
 	assert.NoError(t, err, "shouldn't return error")
 
@@ -22,7 +22,7 @@ func TestAddPeer(t *testing.T) {
 
 func TestGetPeer(t *testing.T) {
 	key := "abc"
-	status := NewRecorder()
+	status := NewRecorder("https://mgm")
 	err := status.AddPeer(key)
 	assert.NoError(t, err, "shouldn't return error")
 
@@ -38,7 +38,7 @@ func TestGetPeer(t *testing.T) {
 func TestUpdatePeerState(t *testing.T) {
 	key := "abc"
 	ip := "10.10.10.10"
-	status := NewRecorder()
+	status := NewRecorder("https://mgm")
 	peerState := State{
 		PubKey: key,
 	}
@@ -58,7 +58,7 @@ func TestUpdatePeerState(t *testing.T) {
 func TestStatus_UpdatePeerFQDN(t *testing.T) {
 	key := "abc"
 	fqdn := "peer-a.netbird.local"
-	status := NewRecorder()
+	status := NewRecorder("https://mgm")
 	peerState := State{
 		PubKey: key,
 	}
@@ -76,7 +76,7 @@ func TestStatus_UpdatePeerFQDN(t *testing.T) {
 func TestGetPeerStateChangeNotifierLogic(t *testing.T) {
 	key := "abc"
 	ip := "10.10.10.10"
-	status := NewRecorder()
+	status := NewRecorder("https://mgm")
 	peerState := State{
 		PubKey: key,
 	}
@@ -100,7 +100,7 @@ func TestGetPeerStateChangeNotifierLogic(t *testing.T) {
 
 func TestRemovePeer(t *testing.T) {
 	key := "abc"
-	status := NewRecorder()
+	status := NewRecorder("https://mgm")
 	peerState := State{
 		PubKey: key,
 	}
@@ -123,7 +123,7 @@ func TestUpdateLocalPeerState(t *testing.T) {
 		PubKey:          "abc",
 		KernelInterface: false,
 	}
-	status := NewRecorder()
+	status := NewRecorder("https://mgm")
 
 	status.UpdateLocalPeerState(localPeerState)
 
@@ -137,7 +137,7 @@ func TestCleanLocalPeerState(t *testing.T) {
 		PubKey:          "abc",
 		KernelInterface: false,
 	}
-	status := NewRecorder()
+	status := NewRecorder("https://mgm")
 
 	status.localPeer = localPeerState
 
@@ -151,29 +151,23 @@ func TestUpdateSignalState(t *testing.T) {
 	var tests = []struct {
 		name      string
 		connected bool
-		want      SignalState
+		want      bool
 	}{
-		{"should mark as connected", true, SignalState{
-
-			URL:       url,
-			Connected: true,
-		}},
-		{"should mark as disconnected", false, SignalState{
-			URL:       url,
-			Connected: false,
-		}},
+		{"should mark as connected", true, true},
+		{"should mark as disconnected", false, false},
 	}
 
-	status := NewRecorder()
+	status := NewRecorder("https://mgm")
+	status.UpdateSignalAddress(url)
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			if test.connected {
-				status.MarkSignalConnected(url)
+				status.MarkSignalConnected()
 			} else {
-				status.MarkSignalDisconnected(url)
+				status.MarkSignalDisconnected()
 			}
-			assert.Equal(t, test.want, status.signal, "signal status should be equal")
+			assert.Equal(t, test.want, status.signalState, "signal status should be equal")
 		})
 	}
 }
@@ -183,29 +177,22 @@ func TestUpdateManagementState(t *testing.T) {
 	var tests = []struct {
 		name      string
 		connected bool
-		want      ManagementState
+		want      bool
 	}{
-		{"should mark as connected", true, ManagementState{
-
-			URL:       url,
-			Connected: true,
-		}},
-		{"should mark as disconnected", false, ManagementState{
-			URL:       url,
-			Connected: false,
-		}},
+		{"should mark as connected", true, true},
+		{"should mark as disconnected", false, false},
 	}
 
-	status := NewRecorder()
+	status := NewRecorder(url)
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			if test.connected {
-				status.MarkManagementConnected(url)
+				status.MarkManagementConnected()
 			} else {
-				status.MarkManagementDisconnected(url)
+				status.MarkManagementDisconnected()
 			}
-			assert.Equal(t, test.want, status.management, "signal status should be equal")
+			assert.Equal(t, test.want, status.managementState, "signalState status should be equal")
 		})
 	}
 }
@@ -213,12 +200,13 @@ func TestUpdateManagementState(t *testing.T) {
 func TestGetFullStatus(t *testing.T) {
 	key1 := "abc"
 	key2 := "def"
+	signalAddr := "https://signal"
 	managementState := ManagementState{
-		URL:       "https://signal",
+		URL:       "https://mgm",
 		Connected: true,
 	}
 	signalState := SignalState{
-		URL:       "https://signal",
+		URL:       signalAddr,
 		Connected: true,
 	}
 	peerState1 := State{
@@ -229,10 +217,11 @@ func TestGetFullStatus(t *testing.T) {
 		PubKey: key2,
 	}
 
-	status := NewRecorder()
+	status := NewRecorder("https://mgm")
+	status.UpdateSignalAddress(signalAddr)
 
-	status.management = managementState
-	status.signal = signalState
+	status.managementState = managementState.Connected
+	status.signalState = signalState.Connected
 	status.peers[key1] = peerState1
 	status.peers[key2] = peerState2
 

client/internal/routemanager/manager_test.go

@@ -398,7 +398,7 @@ func TestManagerUpdateRoutes(t *testing.T) {
 			err = wgInterface.Create()
 			require.NoError(t, err, "should create testing wireguard interface")
 
-			statusRecorder := peer.NewRecorder()
+			statusRecorder := peer.NewRecorder("https://mgm")
 			ctx := context.TODO()
 			routeManager := NewManager(ctx, localPeerKey, wgInterface, statusRecorder)
 			defer routeManager.Stop()