mirror of
https://github.com/netbirdio/netbird.git
synced 2024-12-15 03:11:02 +01:00
085d072b17
- Add magic cookie in hello msg - Add tests
264 lines
6.8 KiB
Go
264 lines
6.8 KiB
Go
package client
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net"
|
|
"sync"
|
|
"time"
|
|
|
|
log "github.com/sirupsen/logrus"
|
|
)
|
|
|
|
var (
|
|
relayCleanupInterval = 60 * time.Second
|
|
|
|
errRelayClientNotConnected = fmt.Errorf("relay client not connected")
|
|
)
|
|
|
|
// RelayTrack hold the relay clients for the foregin relay servers.
|
|
// With the mutex can ensure we can open new connection in case the relay connection has been established with
|
|
// the relay server.
|
|
type RelayTrack struct {
|
|
sync.RWMutex
|
|
relayClient *Client
|
|
}
|
|
|
|
func NewRelayTrack() *RelayTrack {
|
|
return &RelayTrack{}
|
|
}
|
|
|
|
// Manager is a manager for the relay client. It establish one persistent connection to the given relay server. In case
|
|
// of network error the manager will try to reconnect to the server.
|
|
// The manager also manage temproary relay connection. If a client wants to communicate with an another client on a
|
|
// different relay server, the manager will establish a new connection to the relay server. The connection with these
|
|
// relay servers will be closed if there is no active connection. Periodically the manager will check if there is any
|
|
// unused relay connection and close it.
|
|
type Manager struct {
|
|
ctx context.Context
|
|
srvAddress string
|
|
peerID string
|
|
|
|
relayClient *Client
|
|
reconnectGuard *Guard
|
|
|
|
relayClients map[string]*RelayTrack
|
|
relayClientsMutex sync.RWMutex
|
|
|
|
onDisconnectedListeners map[string]map[*func()]struct{}
|
|
listenerLock sync.Mutex
|
|
}
|
|
|
|
func NewManager(ctx context.Context, serverAddress string, peerID string) *Manager {
|
|
return &Manager{
|
|
ctx: ctx,
|
|
srvAddress: serverAddress,
|
|
peerID: peerID,
|
|
relayClients: make(map[string]*RelayTrack),
|
|
onDisconnectedListeners: make(map[string]map[*func()]struct{}),
|
|
}
|
|
}
|
|
|
|
// Serve starts the manager. It will establish a connection to the relay server and start the relay cleanup loop.
|
|
func (m *Manager) Serve() error {
|
|
if m.relayClient != nil {
|
|
return fmt.Errorf("manager already serving")
|
|
}
|
|
|
|
m.relayClient = NewClient(m.ctx, m.srvAddress, m.peerID)
|
|
err := m.relayClient.Connect()
|
|
if err != nil {
|
|
log.Errorf("failed to connect to relay server: %s", err)
|
|
return err
|
|
}
|
|
|
|
m.reconnectGuard = NewGuard(m.ctx, m.relayClient)
|
|
m.relayClient.SetOnDisconnectListener(func() {
|
|
m.onServerDisconnected(m.srvAddress)
|
|
})
|
|
m.startCleanupLoop()
|
|
|
|
return nil
|
|
}
|
|
|
|
// OpenConn opens a connection to the given peer key. If the peer is on the same relay server, the connection will be
|
|
// established via the relay server. If the peer is on a different relay server, the manager will establish a new
|
|
// connection to the relay server.
|
|
func (m *Manager) OpenConn(serverAddress, peerKey string, onClosedListener func()) (net.Conn, error) {
|
|
if m.relayClient == nil {
|
|
return nil, errRelayClientNotConnected
|
|
}
|
|
|
|
foreign, err := m.isForeignServer(serverAddress)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var (
|
|
netConn net.Conn
|
|
)
|
|
if !foreign {
|
|
log.Debugf("open peer connection via permanent server: %s", peerKey)
|
|
netConn, err = m.relayClient.OpenConn(peerKey)
|
|
} else {
|
|
log.Debugf("open peer connection via foreign server: %s", serverAddress)
|
|
netConn, err = m.openConnVia(serverAddress, peerKey)
|
|
}
|
|
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if onClosedListener != nil {
|
|
m.addListener(serverAddress, onClosedListener)
|
|
}
|
|
|
|
return netConn, err
|
|
}
|
|
|
|
// RelayAddress returns the address of the permanent relay server. It could change if the network connection is lost.
|
|
// This address will be sent to the target peer to choose the common relay server for the communication.
|
|
func (m *Manager) RelayAddress() (net.Addr, error) {
|
|
if m.relayClient == nil {
|
|
return nil, errRelayClientNotConnected
|
|
}
|
|
return m.relayClient.RelayRemoteAddress()
|
|
}
|
|
|
|
func (m *Manager) HasRelayAddress() bool {
|
|
return m.srvAddress != ""
|
|
}
|
|
|
|
func (m *Manager) openConnVia(serverAddress, peerKey string) (net.Conn, error) {
|
|
// check if already has a connection to the desired relay server
|
|
m.relayClientsMutex.RLock()
|
|
rt, ok := m.relayClients[serverAddress]
|
|
if ok {
|
|
rt.RLock()
|
|
m.relayClientsMutex.RUnlock()
|
|
defer rt.RUnlock()
|
|
return rt.relayClient.OpenConn(peerKey)
|
|
}
|
|
m.relayClientsMutex.RUnlock()
|
|
|
|
// if not, establish a new connection but check it again (because changed the lock type) before starting the
|
|
// connection
|
|
m.relayClientsMutex.Lock()
|
|
rt, ok = m.relayClients[serverAddress]
|
|
if ok {
|
|
rt.RLock()
|
|
m.relayClientsMutex.Unlock()
|
|
defer rt.RUnlock()
|
|
return rt.relayClient.OpenConn(peerKey)
|
|
}
|
|
|
|
// create a new relay client and store it in the relayClients map
|
|
rt = NewRelayTrack()
|
|
rt.Lock()
|
|
m.relayClients[serverAddress] = rt
|
|
m.relayClientsMutex.Unlock()
|
|
|
|
relayClient := NewClient(m.ctx, serverAddress, m.peerID)
|
|
err := relayClient.Connect()
|
|
if err != nil {
|
|
rt.Unlock()
|
|
m.relayClientsMutex.Lock()
|
|
delete(m.relayClients, serverAddress)
|
|
m.relayClientsMutex.Unlock()
|
|
return nil, err
|
|
}
|
|
// if connection closed then delete the relay client from the list
|
|
relayClient.SetOnDisconnectListener(func() {
|
|
m.onServerDisconnected(serverAddress)
|
|
})
|
|
rt.relayClient = relayClient
|
|
rt.Unlock()
|
|
|
|
conn, err := relayClient.OpenConn(peerKey)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return conn, nil
|
|
}
|
|
|
|
func (m *Manager) onServerDisconnected(serverAddress string) {
|
|
if serverAddress == m.srvAddress {
|
|
m.reconnectGuard.OnDisconnected()
|
|
}
|
|
|
|
m.notifyOnDisconnectListeners(serverAddress)
|
|
}
|
|
|
|
func (m *Manager) isForeignServer(address string) (bool, error) {
|
|
rAddr, err := m.relayClient.RelayRemoteAddress()
|
|
if err != nil {
|
|
return false, fmt.Errorf("relay client not connected")
|
|
}
|
|
log.Debugf("check if foreign server: %s != %s", rAddr.String(), address)
|
|
return rAddr.String() != address, nil
|
|
}
|
|
|
|
func (m *Manager) startCleanupLoop() {
|
|
if m.ctx.Err() != nil {
|
|
return
|
|
}
|
|
|
|
ticker := time.NewTicker(relayCleanupInterval)
|
|
go func() {
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-m.ctx.Done():
|
|
return
|
|
case <-ticker.C:
|
|
m.cleanUpUnusedRelays()
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
|
|
func (m *Manager) cleanUpUnusedRelays() {
|
|
m.relayClientsMutex.Lock()
|
|
defer m.relayClientsMutex.Unlock()
|
|
|
|
for addr, rt := range m.relayClients {
|
|
rt.Lock()
|
|
if rt.relayClient.HasConns() {
|
|
rt.Unlock()
|
|
continue
|
|
}
|
|
rt.relayClient.SetOnDisconnectListener(nil)
|
|
go func() {
|
|
_ = rt.relayClient.Close()
|
|
}()
|
|
log.Debugf("clean up unused relay server connection: %s", addr)
|
|
delete(m.relayClients, addr)
|
|
rt.Unlock()
|
|
}
|
|
}
|
|
|
|
func (m *Manager) addListener(serverAddress string, onClosedListener func()) {
|
|
m.listenerLock.Lock()
|
|
l, ok := m.onDisconnectedListeners[serverAddress]
|
|
if !ok {
|
|
l = make(map[*func()]struct{})
|
|
}
|
|
l[&onClosedListener] = struct{}{}
|
|
m.onDisconnectedListeners[serverAddress] = l
|
|
m.listenerLock.Unlock()
|
|
}
|
|
|
|
func (m *Manager) notifyOnDisconnectListeners(serverAddress string) {
|
|
m.listenerLock.Lock()
|
|
l, ok := m.onDisconnectedListeners[serverAddress]
|
|
if !ok {
|
|
return
|
|
}
|
|
for f := range l {
|
|
go (*f)()
|
|
}
|
|
delete(m.onDisconnectedListeners, serverAddress)
|
|
m.listenerLock.Unlock()
|
|
|
|
}
|