Use the conn state of peer on proper way (#717)

The ConnStatus is a custom type based on iota
like an enum. The problem was nowhere used to the
benefits of this implementation. All ConnStatus
instances has been compared with strings. I
suppose the reason to do it to avoid a circle
dependency. In this commit the separated status
package has been moved to peer package.

Remove unused, exported functions from engine

client/internal/peer/conn.go

@@ -7,13 +7,13 @@ import (
 	"sync"
 	"time"
 
-	"github.com/netbirdio/netbird/client/internal/proxy"
-	nbStatus "github.com/netbirdio/netbird/client/status"
-	"github.com/netbirdio/netbird/client/system"
-	"github.com/netbirdio/netbird/iface"
 	"github.com/pion/ice/v2"
 	log "github.com/sirupsen/logrus"
 	"golang.zx2c4.com/wireguard/wgctrl"
+
+	"github.com/netbirdio/netbird/client/internal/proxy"
+	"github.com/netbirdio/netbird/client/system"
+	"github.com/netbirdio/netbird/iface"
 )
 
 // ConnConfig is a peer Connection configuration
@@ -83,7 +83,7 @@ type Conn struct {
 	agent  *ice.Agent
 	status ConnStatus
 
-	statusRecorder *nbStatus.Status
+	statusRecorder *Status
 
 	proxy proxy.Proxy
 }
@@ -100,7 +100,7 @@ func (conn *Conn) UpdateConf(conf ConnConfig) {
 
 // NewConn creates a new not opened Conn to the remote peer.
 // To establish a connection run Conn.Open
-func NewConn(config ConnConfig, statusRecorder *nbStatus.Status) (*Conn, error) {
+func NewConn(config ConnConfig, statusRecorder *Status) (*Conn, error) {
 	return &Conn{
 		config:         config,
 		mu:             sync.Mutex{},
@@ -190,11 +190,11 @@ func (conn *Conn) reCreateAgent() error {
 func (conn *Conn) Open() error {
 	log.Debugf("trying to connect to peer %s", conn.config.Key)
 
-	peerState := nbStatus.PeerState{PubKey: conn.config.Key}
+	peerState := State{PubKey: conn.config.Key}
 
 	peerState.IP = strings.Split(conn.config.ProxyConfig.AllowedIps, "/")[0]
 	peerState.ConnStatusUpdate = time.Now()
-	peerState.ConnStatus = conn.status.String()
+	peerState.ConnStatus = conn.status
 
 	err := conn.statusRecorder.UpdatePeerState(peerState)
 	if err != nil {
@@ -250,9 +250,9 @@ func (conn *Conn) Open() error {
 	defer conn.notifyDisconnected()
 	conn.mu.Unlock()
 
-	peerState = nbStatus.PeerState{PubKey: conn.config.Key}
+	peerState = State{PubKey: conn.config.Key}
 
-	peerState.ConnStatus = conn.status.String()
+	peerState.ConnStatus = conn.status
 	peerState.ConnStatusUpdate = time.Now()
 	err = conn.statusRecorder.UpdatePeerState(peerState)
 	if err != nil {
@@ -359,7 +359,7 @@ func (conn *Conn) startProxy(remoteConn net.Conn, remoteWgPort int) error {
 		return err
 	}
 
-	peerState := nbStatus.PeerState{PubKey: conn.config.Key}
+	peerState := State{PubKey: conn.config.Key}
 	useProxy := shouldUseProxy(pair)
 	var p proxy.Proxy
 	if useProxy {
@@ -377,7 +377,7 @@ func (conn *Conn) startProxy(remoteConn net.Conn, remoteWgPort int) error {
 
 	conn.status = StatusConnected
 
-	peerState.ConnStatus = conn.status.String()
+	peerState.ConnStatus = conn.status
 	peerState.ConnStatusUpdate = time.Now()
 	peerState.LocalIceCandidateType = pair.Local.Type().String()
 	peerState.RemoteIceCandidateType = pair.Remote.Type().String()
@@ -422,8 +422,8 @@ func (conn *Conn) cleanup() error {
 
 	conn.status = StatusDisconnected
 
-	peerState := nbStatus.PeerState{PubKey: conn.config.Key}
-	peerState.ConnStatus = conn.status.String()
+	peerState := State{PubKey: conn.config.Key}
+	peerState.ConnStatus = conn.status
 	peerState.ConnStatusUpdate = time.Now()
 
 	err := conn.statusRecorder.UpdatePeerState(peerState)

client/internal/peer/conn_status.go

@@ -0,0 +1,29 @@
+package peer
+
+import log "github.com/sirupsen/logrus"
+
+const (
+	// StatusConnected indicate the peer is in connected state
+	StatusConnected ConnStatus = iota
+	// StatusConnecting indicate the peer is in connecting state
+	StatusConnecting
+	// StatusDisconnected indicate the peer is in disconnected state
+	StatusDisconnected
+)
+
+// ConnStatus describe the status of a peer's connection
+type ConnStatus int
+
+func (s ConnStatus) String() string {
+	switch s {
+	case StatusConnecting:
+		return "Connecting"
+	case StatusConnected:
+		return "Connected"
+	case StatusDisconnected:
+		return "Disconnected"
+	default:
+		log.Errorf("unknown status: %d", s)
+		return "INVALID_PEER_CONNECTION_STATUS"
+	}
+}

client/internal/peer/conn_status_test.go

@@ -0,0 +1,27 @@
+package peer
+
+import (
+	"github.com/magiconair/properties/assert"
+	"testing"
+)
+
+func TestConnStatus_String(t *testing.T) {
+
+	tables := []struct {
+		name   string
+		status ConnStatus
+		want   string
+	}{
+		{"StatusConnected", StatusConnected, "Connected"},
+		{"StatusDisconnected", StatusDisconnected, "Disconnected"},
+		{"StatusConnecting", StatusConnecting, "Connecting"},
+	}
+
+	for _, table := range tables {
+		t.Run(table.name, func(t *testing.T) {
+			got := table.status.String()
+			assert.Equal(t, got, table.want, "they should be equal")
+		})
+	}
+
+}

client/internal/peer/conn_test.go

@@ -1,14 +1,15 @@
 package peer
 
 import (
-	"github.com/magiconair/properties/assert"
-	"github.com/netbirdio/netbird/client/internal/proxy"
-	nbstatus "github.com/netbirdio/netbird/client/status"
-	"github.com/netbirdio/netbird/iface"
-	"github.com/pion/ice/v2"
 	"sync"
 	"testing"
 	"time"
+
+	"github.com/magiconair/properties/assert"
+	"github.com/pion/ice/v2"
+
+	"github.com/netbirdio/netbird/client/internal/proxy"
+	"github.com/netbirdio/netbird/iface"
 )
 
 var connConf = ConnConfig{
@@ -46,7 +47,7 @@ func TestConn_GetKey(t *testing.T) {
 
 func TestConn_OnRemoteOffer(t *testing.T) {
 
-	conn, err := NewConn(connConf, nbstatus.NewRecorder())
+	conn, err := NewConn(connConf, NewRecorder())
 	if err != nil {
 		return
 	}
@@ -80,7 +81,7 @@ func TestConn_OnRemoteOffer(t *testing.T) {
 
 func TestConn_OnRemoteAnswer(t *testing.T) {
 
-	conn, err := NewConn(connConf, nbstatus.NewRecorder())
+	conn, err := NewConn(connConf, NewRecorder())
 	if err != nil {
 		return
 	}
@@ -113,7 +114,7 @@ func TestConn_OnRemoteAnswer(t *testing.T) {
 }
 func TestConn_Status(t *testing.T) {
 
-	conn, err := NewConn(connConf, nbstatus.NewRecorder())
+	conn, err := NewConn(connConf, NewRecorder())
 	if err != nil {
 		return
 	}
@@ -140,7 +141,7 @@ func TestConn_Status(t *testing.T) {
 
 func TestConn_Close(t *testing.T) {
 
-	conn, err := NewConn(connConf, nbstatus.NewRecorder())
+	conn, err := NewConn(connConf, NewRecorder())
 	if err != nil {
 		return
 	}

client/internal/peer/status.go

@@ -1,25 +1,241 @@
 package peer
 
-import log "github.com/sirupsen/logrus"
+import (
+	"errors"
+	"sync"
+	"time"
+)
 
-type ConnStatus int
+// State contains the latest state of a peer
+type State struct {
+	IP                     string
+	PubKey                 string
+	FQDN                   string
+	ConnStatus             ConnStatus
+	ConnStatusUpdate       time.Time
+	Relayed                bool
+	Direct                 bool
+	LocalIceCandidateType  string
+	RemoteIceCandidateType string
+}
 
-func (s ConnStatus) String() string {
-	switch s {
-	case StatusConnecting:
-		return "Connecting"
-	case StatusConnected:
-		return "Connected"
-	case StatusDisconnected:
-		return "Disconnected"
-	default:
-		log.Errorf("unknown status: %d", s)
-		return "INVALID_PEER_CONNECTION_STATUS"
+// LocalPeerState contains the latest state of the local peer
+type LocalPeerState struct {
+	IP              string
+	PubKey          string
+	KernelInterface bool
+	FQDN            string
+}
+
+// SignalState contains the latest state of a signal connection
+type SignalState struct {
+	URL       string
+	Connected bool
+}
+
+// ManagementState contains the latest state of a management connection
+type ManagementState struct {
+	URL       string
+	Connected bool
+}
+
+// FullStatus contains the full state held by the Status instance
+type FullStatus struct {
+	Peers           []State
+	ManagementState ManagementState
+	SignalState     SignalState
+	LocalPeerState  LocalPeerState
+}
+
+// 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
+}
+
+// NewRecorder returns a new Status instance
+func NewRecorder() *Status {
+	return &Status{
+		peers:        make(map[string]State),
+		changeNotify: make(map[string]chan struct{}),
 	}
 }
 
-const (
-	StatusConnected ConnStatus = iota
-	StatusConnecting
-	StatusDisconnected
-)
+// AddPeer adds peer to Daemon status map
+func (d *Status) AddPeer(peerPubKey string) error {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	_, ok := d.peers[peerPubKey]
+	if ok {
+		return errors.New("peer already exist")
+	}
+	d.peers[peerPubKey] = State{PubKey: peerPubKey}
+	return nil
+}
+
+// GetPeer adds peer to Daemon status map
+func (d *Status) GetPeer(peerPubKey string) (State, error) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	state, ok := d.peers[peerPubKey]
+	if !ok {
+		return State{}, errors.New("peer not found")
+	}
+	return state, nil
+}
+
+// RemovePeer removes peer from Daemon status map
+func (d *Status) RemovePeer(peerPubKey string) error {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	_, ok := d.peers[peerPubKey]
+	if ok {
+		delete(d.peers, peerPubKey)
+		return nil
+	}
+
+	return errors.New("no peer with to remove")
+}
+
+// UpdatePeerState updates peer status
+func (d *Status) UpdatePeerState(receivedState State) error {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	peerState, ok := d.peers[receivedState.PubKey]
+	if !ok {
+		return errors.New("peer doesn't exist")
+	}
+
+	if receivedState.IP != "" {
+		peerState.IP = receivedState.IP
+	}
+
+	if receivedState.ConnStatus != peerState.ConnStatus {
+		peerState.ConnStatus = receivedState.ConnStatus
+		peerState.ConnStatusUpdate = receivedState.ConnStatusUpdate
+		peerState.Direct = receivedState.Direct
+		peerState.Relayed = receivedState.Relayed
+		peerState.LocalIceCandidateType = receivedState.LocalIceCandidateType
+		peerState.RemoteIceCandidateType = receivedState.RemoteIceCandidateType
+	}
+
+	d.peers[receivedState.PubKey] = peerState
+
+	ch, found := d.changeNotify[receivedState.PubKey]
+	if found && ch != nil {
+		close(ch)
+		d.changeNotify[receivedState.PubKey] = nil
+	}
+
+	return nil
+}
+
+// UpdatePeerFQDN update peer's state fqdn only
+func (d *Status) UpdatePeerFQDN(peerPubKey, fqdn string) error {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	peerState, ok := d.peers[peerPubKey]
+	if !ok {
+		return errors.New("peer doesn't exist")
+	}
+
+	peerState.FQDN = fqdn
+	d.peers[peerPubKey] = peerState
+
+	return nil
+}
+
+// GetPeerStateChangeNotifier returns a change notifier channel for a peer
+func (d *Status) GetPeerStateChangeNotifier(peer string) <-chan struct{} {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+	ch, found := d.changeNotify[peer]
+	if !found || ch == nil {
+		ch = make(chan struct{})
+		d.changeNotify[peer] = ch
+	}
+	return ch
+}
+
+// UpdateLocalPeerState updates local peer status
+func (d *Status) UpdateLocalPeerState(localPeerState LocalPeerState) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	d.localPeer = localPeerState
+}
+
+// CleanLocalPeerState cleans local peer status
+func (d *Status) CleanLocalPeerState() {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	d.localPeer = LocalPeerState{}
+}
+
+// MarkManagementDisconnected sets ManagementState to disconnected
+func (d *Status) MarkManagementDisconnected(managementURL string) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+	d.management = ManagementState{
+		URL:       managementURL,
+		Connected: false,
+	}
+}
+
+// MarkManagementConnected sets ManagementState to connected
+func (d *Status) MarkManagementConnected(managementURL string) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+	d.management = ManagementState{
+		URL:       managementURL,
+		Connected: true,
+	}
+}
+
+// MarkSignalDisconnected sets SignalState to disconnected
+func (d *Status) MarkSignalDisconnected(signalURL string) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+	d.signal = SignalState{
+		signalURL,
+		false,
+	}
+}
+
+// MarkSignalConnected sets SignalState to connected
+func (d *Status) MarkSignalConnected(signalURL string) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+	d.signal = SignalState{
+		signalURL,
+		true,
+	}
+}
+
+// GetFullStatus gets full status
+func (d *Status) GetFullStatus() FullStatus {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	fullStatus := FullStatus{
+		ManagementState: d.management,
+		SignalState:     d.signal,
+		LocalPeerState:  d.localPeer,
+	}
+
+	for _, status := range d.peers {
+		fullStatus.Peers = append(fullStatus.Peers, status)
+	}
+
+	return fullStatus
+}

client/internal/peer/status_test.go

@@ -1,27 +1,244 @@
 package peer
 
 import (
-	"github.com/magiconair/properties/assert"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
-func TestConnStatus_String(t *testing.T) {
+func TestAddPeer(t *testing.T) {
+	key := "abc"
+	status := NewRecorder()
+	err := status.AddPeer(key)
+	assert.NoError(t, err, "shouldn't return error")
 
-	tables := []struct {
-		name   string
-		status ConnStatus
-		want   string
-	}{
-		{"StatusConnected", StatusConnected, "Connected"},
-		{"StatusDisconnected", StatusDisconnected, "Disconnected"},
-		{"StatusConnecting", StatusConnecting, "Connecting"},
+	_, exists := status.peers[key]
+	assert.True(t, exists, "value was found")
+
+	err = status.AddPeer(key)
+
+	assert.Error(t, err, "should return error on duplicate")
+}
+
+func TestGetPeer(t *testing.T) {
+	key := "abc"
+	status := NewRecorder()
+	err := status.AddPeer(key)
+	assert.NoError(t, err, "shouldn't return error")
+
+	peerStatus, err := status.GetPeer(key)
+	assert.NoError(t, err, "shouldn't return error on getting peer")
+
+	assert.Equal(t, key, peerStatus.PubKey, "retrieved public key should match")
+
+	_, err = status.GetPeer("non_existing_key")
+	assert.Error(t, err, "should return error when peer doesn't exist")
+}
+
+func TestUpdatePeerState(t *testing.T) {
+	key := "abc"
+	ip := "10.10.10.10"
+	status := NewRecorder()
+	peerState := State{
+		PubKey: key,
 	}
 
-	for _, table := range tables {
-		t.Run(table.name, func(t *testing.T) {
-			got := table.status.String()
-			assert.Equal(t, got, table.want, "they should be equal")
+	status.peers[key] = peerState
+
+	peerState.IP = ip
+
+	err := status.UpdatePeerState(peerState)
+	assert.NoError(t, err, "shouldn't return error")
+
+	state, exists := status.peers[key]
+	assert.True(t, exists, "state should be found")
+	assert.Equal(t, ip, state.IP, "ip should be equal")
+}
+
+func TestStatus_UpdatePeerFQDN(t *testing.T) {
+	key := "abc"
+	fqdn := "peer-a.netbird.local"
+	status := NewRecorder()
+	peerState := State{
+		PubKey: key,
+	}
+
+	status.peers[key] = peerState
+
+	err := status.UpdatePeerFQDN(key, fqdn)
+	assert.NoError(t, err, "shouldn't return error")
+
+	state, exists := status.peers[key]
+	assert.True(t, exists, "state should be found")
+	assert.Equal(t, fqdn, state.FQDN, "fqdn should be equal")
+}
+
+func TestGetPeerStateChangeNotifierLogic(t *testing.T) {
+	key := "abc"
+	ip := "10.10.10.10"
+	status := NewRecorder()
+	peerState := State{
+		PubKey: key,
+	}
+
+	status.peers[key] = peerState
+
+	ch := status.GetPeerStateChangeNotifier(key)
+	assert.NotNil(t, ch, "channel shouldn't be nil")
+
+	peerState.IP = ip
+
+	err := status.UpdatePeerState(peerState)
+	assert.NoError(t, err, "shouldn't return error")
+
+	select {
+	case <-ch:
+	default:
+		t.Errorf("channel wasn't closed after update")
+	}
+}
+
+func TestRemovePeer(t *testing.T) {
+	key := "abc"
+	status := NewRecorder()
+	peerState := State{
+		PubKey: key,
+	}
+
+	status.peers[key] = peerState
+
+	err := status.RemovePeer(key)
+	assert.NoError(t, err, "shouldn't return error")
+
+	_, exists := status.peers[key]
+	assert.False(t, exists, "state value shouldn't be found")
+
+	err = status.RemovePeer("not existing")
+	assert.Error(t, err, "should return error when peer doesn't exist")
+}
+
+func TestUpdateLocalPeerState(t *testing.T) {
+	localPeerState := LocalPeerState{
+		IP:              "10.10.10.10",
+		PubKey:          "abc",
+		KernelInterface: false,
+	}
+	status := NewRecorder()
+
+	status.UpdateLocalPeerState(localPeerState)
+
+	assert.Equal(t, localPeerState, status.localPeer, "local peer status should be equal")
+}
+
+func TestCleanLocalPeerState(t *testing.T) {
+	emptyLocalPeerState := LocalPeerState{}
+	localPeerState := LocalPeerState{
+		IP:              "10.10.10.10",
+		PubKey:          "abc",
+		KernelInterface: false,
+	}
+	status := NewRecorder()
+
+	status.localPeer = localPeerState
+
+	status.CleanLocalPeerState()
+
+	assert.Equal(t, emptyLocalPeerState, status.localPeer, "local peer status should be empty")
+}
+
+func TestUpdateSignalState(t *testing.T) {
+	url := "https://signal"
+	var tests = []struct {
+		name      string
+		connected bool
+		want      SignalState
+	}{
+		{"should mark as connected", true, SignalState{
+
+			URL:       url,
+			Connected: true,
+		}},
+		{"should mark as disconnected", false, SignalState{
+			URL:       url,
+			Connected: false,
+		}},
+	}
+
+	status := NewRecorder()
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			if test.connected {
+				status.MarkSignalConnected(url)
+			} else {
+				status.MarkSignalDisconnected(url)
+			}
+			assert.Equal(t, test.want, status.signal, "signal status should be equal")
 		})
 	}
-
+}
+
+func TestUpdateManagementState(t *testing.T) {
+	url := "https://management"
+	var tests = []struct {
+		name      string
+		connected bool
+		want      ManagementState
+	}{
+		{"should mark as connected", true, ManagementState{
+
+			URL:       url,
+			Connected: true,
+		}},
+		{"should mark as disconnected", false, ManagementState{
+			URL:       url,
+			Connected: false,
+		}},
+	}
+
+	status := NewRecorder()
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			if test.connected {
+				status.MarkManagementConnected(url)
+			} else {
+				status.MarkManagementDisconnected(url)
+			}
+			assert.Equal(t, test.want, status.management, "signal status should be equal")
+		})
+	}
+}
+
+func TestGetFullStatus(t *testing.T) {
+	key1 := "abc"
+	key2 := "def"
+	managementState := ManagementState{
+		URL:       "https://signal",
+		Connected: true,
+	}
+	signalState := SignalState{
+		URL:       "https://signal",
+		Connected: true,
+	}
+	peerState1 := State{
+		PubKey: key1,
+	}
+
+	peerState2 := State{
+		PubKey: key2,
+	}
+
+	status := NewRecorder()
+
+	status.management = managementState
+	status.signal = signalState
+	status.peers[key1] = peerState1
+	status.peers[key2] = peerState2
+
+	fullStatus := status.GetFullStatus()
+
+	assert.Equal(t, managementState, fullStatus.ManagementState, "management status should be equal")
+	assert.Equal(t, signalState, fullStatus.SignalState, "signal status should be equal")
+	assert.ElementsMatch(t, []State{peerState1, peerState2}, fullStatus.Peers, "peers states should match")
 }