Support multiple server

This commit is contained in:
Zoltán Papp 2024-05-29 16:40:26 +02:00
parent 7cc3964a4d
commit 4ff069a102
3 changed files with 225 additions and 70 deletions

View File

@ -3,7 +3,6 @@ package client
import ( import (
"context" "context"
"fmt" "fmt"
"github.com/netbirdio/netbird/relay/client/dialer/udp"
"io" "io"
"net" "net"
"sync" "sync"
@ -11,6 +10,7 @@ import (
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/relay/client/dialer/udp"
"github.com/netbirdio/netbird/relay/messages" "github.com/netbirdio/netbird/relay/messages"
) )
@ -48,6 +48,8 @@ type Client struct {
serviceIsRunningMutex sync.Mutex serviceIsRunningMutex sync.Mutex
wgReadLoop sync.WaitGroup wgReadLoop sync.WaitGroup
onDisconnected chan struct{} onDisconnected chan struct{}
remoteAddr net.Addr
} }
func NewClient(ctx context.Context, serverAddress, peerID string) *Client { func NewClient(ctx context.Context, serverAddress, peerID string) *Client {
@ -97,31 +99,35 @@ func (c *Client) Connect() error {
return nil return nil
} }
func (c *Client) reconnectGuard() { func (c *Client) ConnectWithoutReconnect() error {
for { c.serviceIsRunningMutex.Lock()
c.wgReadLoop.Wait() if c.serviceIsRunning {
c.serviceIsRunningMutex.Lock()
if !c.serviceIsRunning {
c.serviceIsRunningMutex.Unlock()
return
}
log.Infof("reconnecting to relay server")
err := c.connect()
if err != nil {
log.Errorf("failed to reconnect to relay server: %s", err)
c.serviceIsRunningMutex.Unlock()
time.Sleep(reconnectingTimeout)
continue
}
log.Infof("reconnected to relay server")
c.wgReadLoop.Add(1)
go c.readLoop()
c.serviceIsRunningMutex.Unlock() c.serviceIsRunningMutex.Unlock()
return nil
} }
err := c.connect()
if err != nil {
c.serviceIsRunningMutex.Unlock()
return err
}
c.serviceIsRunning = true
c.wgReadLoop.Add(1)
go c.readLoop()
c.serviceIsRunningMutex.Unlock()
go func() {
<-c.ctx.Done()
cErr := c.close()
if cErr != nil {
log.Errorf("failed to close relay connection: %s", cErr)
}
}()
return nil
} }
func (c *Client) OpenConn(dstPeerID string) (net.Conn, error) { func (c *Client) OpenConn(dstPeerID string) (net.Conn, error) {
@ -144,6 +150,15 @@ func (c *Client) OpenConn(dstPeerID string) (net.Conn, error) {
return conn, nil return conn, nil
} }
func (c *Client) RelayRemoteAddress() (net.Addr, error) {
c.serviceIsRunningMutex.Lock()
defer c.serviceIsRunningMutex.Unlock()
if c.remoteAddr == nil {
return nil, fmt.Errorf("relay connection is not established")
}
return c.remoteAddr, nil
}
func (c *Client) Close() error { func (c *Client) Close() error {
c.serviceIsRunningMutex.Lock() c.serviceIsRunningMutex.Lock()
if !c.serviceIsRunning { if !c.serviceIsRunning {
@ -172,6 +187,8 @@ func (c *Client) connect() error {
return err return err
} }
c.remoteAddr = conn.RemoteAddr()
c.readyToOpenConns = true c.readyToOpenConns = true
return nil return nil
} }
@ -193,6 +210,33 @@ func (c *Client) close() error {
return err return err
} }
func (c *Client) reconnectGuard() {
for {
c.wgReadLoop.Wait()
c.serviceIsRunningMutex.Lock()
if !c.serviceIsRunning {
c.serviceIsRunningMutex.Unlock()
return
}
log.Infof("reconnecting to relay server")
err := c.connect()
if err != nil {
log.Errorf("failed to reconnect to relay server: %s", err)
c.serviceIsRunningMutex.Unlock()
time.Sleep(reconnectingTimeout)
continue
}
log.Infof("reconnected to relay server")
c.wgReadLoop.Add(1)
go c.readLoop()
c.serviceIsRunningMutex.Unlock()
}
}
func (c *Client) handShake() error { func (c *Client) handShake() error {
defer func() { defer func() {
err := c.relayConn.SetReadDeadline(time.Time{}) err := c.relayConn.SetReadDeadline(time.Time{})

View File

@ -2,74 +2,88 @@ package client
import ( import (
"context" "context"
"fmt"
"net" "net"
"sync"
"time"
log "github.com/sirupsen/logrus"
) )
// Manager todo: thread safe
type Manager struct { type Manager struct {
ctx context.Context ctx context.Context
srvAddress string srvAddress string
peerID string peerID string
reconnectTime time.Duration relayClient *Client
mu sync.Mutex relayClients map[string]*Client
client *Client
} }
func NewManager(ctx context.Context, serverAddress string, peerID string) *Manager { func NewManager(ctx context.Context, serverAddress string, peerID string) *Manager {
return &Manager{ return &Manager{
ctx: ctx, ctx: ctx,
srvAddress: serverAddress, srvAddress: serverAddress,
peerID: peerID, peerID: peerID,
reconnectTime: 5 * time.Second, relayClients: make(map[string]*Client),
} }
} }
func (m *Manager) Serve() { func (m *Manager) Serve() error {
ok := m.mu.TryLock() m.relayClient = NewClient(m.ctx, m.srvAddress, m.peerID)
if !ok { err := m.relayClient.Connect()
return if err != nil {
return err
} }
return nil
}
m.client = NewClient(m.ctx, m.srvAddress, m.peerID) func (m *Manager) RelayAddress() (net.Addr, error) {
if m.relayClient == nil {
go func() { return nil, fmt.Errorf("relay client not connected")
defer m.mu.Unlock() }
return m.relayClient.RelayRemoteAddress()
// todo this is not thread safe
for {
select {
case <-m.ctx.Done():
return
default:
m.connect()
}
select {
case <-m.ctx.Done():
return
case <-time.After(2 * time.Second): //timeout
}
}
}()
} }
func (m *Manager) OpenConn(peerKey string) (net.Conn, error) { func (m *Manager) OpenConn(peerKey string) (net.Conn, error) {
// todo m.client nil check if m.relayClient == nil {
return m.client.OpenConn(peerKey) return nil, fmt.Errorf("relay client not connected")
}
rAddr, err := m.relayClient.RelayRemoteAddress()
if err != nil {
return nil, fmt.Errorf("relay client not connected")
}
return m.OpenConnTo(rAddr.String(), peerKey)
} }
// connect is blocking func (m *Manager) OpenConnTo(serverAddress, peerKey string) (net.Conn, error) {
func (m *Manager) connect() { if m.relayClient == nil {
err := m.client.Connect() return nil, fmt.Errorf("relay client not connected")
if err != nil {
if m.ctx.Err() != nil {
return
}
log.Errorf("connection error with '%s': %s", m.srvAddress, err)
} }
rAddr, err := m.relayClient.RelayRemoteAddress()
if err != nil {
return nil, fmt.Errorf("relay client not connected")
}
if rAddr.String() == serverAddress {
return m.relayClient.OpenConn(peerKey)
}
relayClient, ok := m.relayClients[serverAddress]
if ok {
return relayClient.OpenConn(peerKey)
}
relayClient = NewClient(m.ctx, serverAddress, m.peerID)
err = relayClient.ConnectWithoutReconnect()
if err != nil {
return nil, err
}
conn, err := relayClient.OpenConn(peerKey)
if err != nil {
return nil, err
}
m.relayClients[serverAddress] = relayClient
return conn, nil
} }

View File

@ -0,0 +1,97 @@
package client
import (
"context"
"testing"
"github.com/netbirdio/netbird/relay/server"
)
func TestNewManager(t *testing.T) {
ctx := context.Background()
idAlice := "alice"
idBob := "bob"
addr1 := "localhost:1234"
srv1 := server.NewServer()
go func() {
err := srv1.Listen(addr1)
if err != nil {
t.Fatalf("failed to bind server: %s", err)
}
}()
defer func() {
err := srv1.Close()
if err != nil {
t.Errorf("failed to close server: %s", err)
}
}()
addr2 := "localhost:2234"
srv2 := server.NewServer()
go func() {
err := srv2.Listen(addr2)
if err != nil {
t.Fatalf("failed to bind server: %s", err)
}
}()
defer func() {
err := srv2.Close()
if err != nil {
t.Errorf("failed to close server: %s", err)
}
}()
clientAlice := NewManager(ctx, addr1, idAlice)
err := clientAlice.Serve()
if err != nil {
t.Fatalf("failed to connect to server: %s", err)
}
clientBob := NewManager(ctx, addr2, idBob)
err = clientBob.Serve()
if err != nil {
t.Fatalf("failed to connect to server: %s", err)
}
bobsSrvAddr, err := clientBob.RelayAddress()
if err != nil {
t.Fatalf("failed to get relay address: %s", err)
}
connAliceToBob, err := clientAlice.OpenConnTo(bobsSrvAddr.String(), idBob)
if err != nil {
t.Fatalf("failed to bind channel: %s", err)
}
connBobToAlice, err := clientBob.OpenConn(idAlice)
if err != nil {
t.Fatalf("failed to bind channel: %s", err)
}
payload := "hello bob, I am alice"
_, err = connAliceToBob.Write([]byte(payload))
if err != nil {
t.Fatalf("failed to write to channel: %s", err)
}
buf := make([]byte, 65535)
n, err := connBobToAlice.Read(buf)
if err != nil {
t.Fatalf("failed to read from channel: %s", err)
}
_, err = connBobToAlice.Write(buf[:n])
if err != nil {
t.Fatalf("failed to write to channel: %s", err)
}
n, err = connAliceToBob.Read(buf)
if err != nil {
t.Fatalf("failed to read from channel: %s", err)
}
if payload != string(buf[:n]) {
t.Fatalf("expected %s, got %s", payload, string(buf[:n]))
}
}