mirror of
https://github.com/netbirdio/netbird.git
synced 2026-03-31 06:24:18 -04:00
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
This commit is contained in:
Zoltan Papp
GitHub
@@ -1,241 +0,0 @@
|
||||
package status
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// PeerState contains the latest state of a peer
|
||||
type PeerState struct {
|
||||
IP string
|
||||
PubKey string
|
||||
FQDN string
|
||||
ConnStatus string
|
||||
ConnStatusUpdate time.Time
|
||||
Relayed bool
|
||||
Direct bool
|
||||
LocalIceCandidateType string
|
||||
RemoteIceCandidateType string
|
||||
}
|
||||
|
||||
// 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 []PeerState
|
||||
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]PeerState
|
||||
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]PeerState),
|
||||
changeNotify: make(map[string]chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// 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] = PeerState{PubKey: peerPubKey}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetPeer adds peer to Daemon status map
|
||||
func (d *Status) GetPeer(peerPubKey string) (PeerState, error) {
|
||||
d.mux.Lock()
|
||||
defer d.mux.Unlock()
|
||||
|
||||
state, ok := d.peers[peerPubKey]
|
||||
if !ok {
|
||||
return PeerState{}, 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 PeerState) 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
|
||||
}
|
||||
@@ -1,243 +0,0 @@
|
||||
package status
|
||||
|
||||
import (
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAddPeer(t *testing.T) {
|
||||
key := "abc"
|
||||
status := NewRecorder()
|
||||
err := status.AddPeer(key)
|
||||
assert.NoError(t, err, "shouldn't return error")
|
||||
|
||||
_, 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 := PeerState{
|
||||
PubKey: key,
|
||||
}
|
||||
|
||||
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 := PeerState{
|
||||
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 := PeerState{
|
||||
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 := PeerState{
|
||||
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 := PeerState{
|
||||
PubKey: key1,
|
||||
}
|
||||
|
||||
peerState2 := PeerState{
|
||||
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, []PeerState{peerState1, peerState2}, fullStatus.Peers, "peers states should match")
|
||||
}
|
||||
Reference in New Issue
Block a user