netbird/relay/client/manager.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")
	}
	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()

}