[misc] Move shared components to shared directory (#4286)

Moved the following directories:

```
  - management/client → shared/management/client
  - management/domain → shared/management/domain
  - management/proto → shared/management/proto
  - signal/client → shared/signal/client
  - signal/proto → shared/signal/proto
  - relay/client → shared/relay/client
  - relay/auth → shared/relay/auth
```

and adjusted import paths

shared/relay/auth/allow/allow_all.go

@@ -0,0 +1,14 @@
+package allow
+
+// Auth is a Validator that allows all connections.
+// Used this for testing purposes only.
+type Auth struct {
+}
+
+func (a *Auth) Validate(any) error {
+	return nil
+}
+
+func (a *Auth) ValidateHelloMsgType(any) error {
+	return nil
+}

shared/relay/auth/doc.go

@@ -0,0 +1,26 @@
+/*
+Package auth manages the authentication process with the relay server.
+
+Key Components:
+
+Validator: The Validator interface defines the Validate method. Any type that provides this method can be used as a
+Validator.
+
+Methods:
+
+Validate(func() hash.Hash, any): This method is defined in the Validator interface and is used to validate the authentication.
+
+Usage:
+
+To create a new AllowAllAuth validator, simply instantiate it:
+
+	validator := &allow.Auth{}
+
+To validate the authentication, use the Validate method:
+
+	err := validator.Validate(sha256.New, any)
+
+This package provides a simple and effective way to manage authentication with the relay server, ensuring that the
+peers are authenticated properly.
+*/
+package auth

shared/relay/auth/go.sum

@@ -0,0 +1 @@
+golang.org/x/crypto v0.37.0/go.mod h1:vg+k43peMZ0pUMhYmVAWysMK35e6ioLh3wB8ZCAfbVc=

shared/relay/auth/hmac/doc.go

@@ -0,0 +1,8 @@
+/*
+This package uses a similar HMAC method for authentication with the TURN server. The Management server provides the
+tokens for the peers. The peers manage these tokens in the token store. The token store is a simple thread safe store
+that keeps the tokens in memory. These tokens are used to authenticate the peers with the Relay server in the hello
+message.
+*/
+
+package hmac

shared/relay/auth/hmac/store.go

@@ -0,0 +1,44 @@
+package hmac
+
+import (
+	"encoding/base64"
+	"fmt"
+	"sync"
+
+	v2 "github.com/netbirdio/netbird/shared/relay/auth/hmac/v2"
+)
+
+// TokenStore is a simple in-memory store for token
+// With this can update the token in thread safe way
+type TokenStore struct {
+	mu    sync.Mutex
+	token []byte
+}
+
+func (a *TokenStore) UpdateToken(token *Token) error {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	if token == nil {
+		return nil
+	}
+
+	sig, err := base64.StdEncoding.DecodeString(token.Signature)
+	if err != nil {
+		return fmt.Errorf("decode signature: %w", err)
+	}
+
+	tok := v2.Token{
+		AuthAlgo:  v2.AuthAlgoHMACSHA256,
+		Signature: sig,
+		Payload:   []byte(token.Payload),
+	}
+
+	a.token = tok.Marshal()
+	return nil
+}
+
+func (a *TokenStore) TokenBinary() []byte {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	return a.token
+}

shared/relay/auth/hmac/token.go

@@ -0,0 +1,94 @@
+package hmac
+
+import (
+	"bytes"
+	"crypto/hmac"
+	"encoding/base64"
+	"encoding/gob"
+	"fmt"
+	"hash"
+	"strconv"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+)
+
+type Token struct {
+	Payload   string
+	Signature string
+}
+
+func unmarshalToken(payload []byte) (Token, error) {
+	var creds Token
+	buffer := bytes.NewBuffer(payload)
+	decoder := gob.NewDecoder(buffer)
+	err := decoder.Decode(&creds)
+	return creds, err
+}
+
+// TimedHMAC generates a token with TTL and uses a pre-shared secret known to the relay server
+type TimedHMAC struct {
+	secret     string
+	timeToLive time.Duration
+}
+
+// NewTimedHMAC creates a new TimedHMAC instance
+func NewTimedHMAC(secret string, timeToLive time.Duration) *TimedHMAC {
+	return &TimedHMAC{
+		secret:     secret,
+		timeToLive: timeToLive,
+	}
+}
+
+// GenerateToken generates new time-based secret token - basically Payload is a unix timestamp and Signature is a HMAC
+// hash of a timestamp with a preshared TURN secret
+func (m *TimedHMAC) GenerateToken(algo func() hash.Hash) (*Token, error) {
+	timeAuth := time.Now().Add(m.timeToLive).Unix()
+	timeStamp := strconv.FormatInt(timeAuth, 10)
+
+	checksum, err := m.generate(algo, timeStamp)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Token{
+		Payload:   timeStamp,
+		Signature: base64.StdEncoding.EncodeToString(checksum),
+	}, nil
+}
+
+// Validate checks if the token is valid
+func (m *TimedHMAC) Validate(algo func() hash.Hash, token Token) error {
+	expectedMAC, err := m.generate(algo, token.Payload)
+	if err != nil {
+		return err
+	}
+
+	expectedSignature := base64.StdEncoding.EncodeToString(expectedMAC)
+
+	if !hmac.Equal([]byte(expectedSignature), []byte(token.Signature)) {
+		return fmt.Errorf("signature mismatch")
+	}
+
+	timeAuthInt, err := strconv.ParseInt(token.Payload, 10, 64)
+	if err != nil {
+		return fmt.Errorf("invalid payload: %w", err)
+	}
+
+	if time.Now().Unix() > timeAuthInt {
+		return fmt.Errorf("expired token")
+	}
+
+	return nil
+}
+
+func (m *TimedHMAC) generate(algo func() hash.Hash, payload string) ([]byte, error) {
+	mac := hmac.New(algo, []byte(m.secret))
+	_, err := mac.Write([]byte(payload))
+	if err != nil {
+		log.Debugf("failed to generate token: %s", err)
+		return nil, fmt.Errorf("failed to generate token: %w", err)
+	}
+
+	return mac.Sum(nil), nil
+}

shared/relay/auth/hmac/token_test.go

@@ -0,0 +1,105 @@
+package hmac
+
+import (
+	"crypto/sha1"
+	"crypto/sha256"
+	"encoding/base64"
+	"strconv"
+	"testing"
+	"time"
+)
+
+func TestGenerateCredentials(t *testing.T) {
+	secret := "secret"
+	timeToLive := 1 * time.Hour
+	v := NewTimedHMAC(secret, timeToLive)
+
+	creds, err := v.GenerateToken(sha1.New)
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	if creds.Payload == "" {
+		t.Fatalf("expected non-empty payload")
+	}
+
+	_, err = strconv.ParseInt(creds.Payload, 10, 64)
+	if err != nil {
+		t.Fatalf("expected payload to be a valid unix timestamp, got %v", err)
+	}
+
+	_, err = base64.StdEncoding.DecodeString(creds.Signature)
+	if err != nil {
+		t.Fatalf("expected signature to be base64 encoded, got %v", err)
+	}
+}
+
+func TestValidateCredentials(t *testing.T) {
+	secret := "supersecret"
+	timeToLive := 1 * time.Hour
+	manager := NewTimedHMAC(secret, timeToLive)
+
+	// Test valid token
+	creds, err := manager.GenerateToken(sha1.New)
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	if err := manager.Validate(sha1.New, *creds); err != nil {
+		t.Fatalf("expected valid token: %s", err)
+	}
+}
+
+func TestInvalidSignature(t *testing.T) {
+	secret := "supersecret"
+	timeToLive := 1 * time.Hour
+	manager := NewTimedHMAC(secret, timeToLive)
+
+	creds, err := manager.GenerateToken(sha256.New)
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	invalidCreds := &Token{
+		Payload:   creds.Payload,
+		Signature: "invalidsignature",
+	}
+
+	if err = manager.Validate(sha1.New, *invalidCreds); err == nil {
+		t.Fatalf("expected invalid token due to signature mismatch")
+	}
+}
+
+func TestExpired(t *testing.T) {
+	secret := "supersecret"
+	v := NewTimedHMAC(secret, -1*time.Hour)
+	expiredCreds, err := v.GenerateToken(sha256.New)
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	if err = v.Validate(sha1.New, *expiredCreds); err == nil {
+		t.Fatalf("expected invalid token due to expiration")
+	}
+}
+
+func TestInvalidPayload(t *testing.T) {
+	secret := "supersecret"
+	timeToLive := 1 * time.Hour
+	v := NewTimedHMAC(secret, timeToLive)
+
+	creds, err := v.GenerateToken(sha256.New)
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	// Test invalid payload
+	invalidPayloadCreds := &Token{
+		Payload:   "invalidtimestamp",
+		Signature: creds.Signature,
+	}
+
+	if err = v.Validate(sha1.New, *invalidPayloadCreds); err == nil {
+		t.Fatalf("expected invalid token due to invalid payload")
+	}
+}

shared/relay/auth/hmac/v2/algo.go

@@ -0,0 +1,40 @@
+package v2
+
+import (
+	"crypto/sha256"
+	"hash"
+)
+
+const (
+	AuthAlgoUnknown AuthAlgo = iota
+	AuthAlgoHMACSHA256
+)
+
+type AuthAlgo uint8
+
+func (a AuthAlgo) String() string {
+	switch a {
+	case AuthAlgoHMACSHA256:
+		return "HMAC-SHA256"
+	default:
+		return "Unknown"
+	}
+}
+
+func (a AuthAlgo) New() func() hash.Hash {
+	switch a {
+	case AuthAlgoHMACSHA256:
+		return sha256.New
+	default:
+		return nil
+	}
+}
+
+func (a AuthAlgo) Size() int {
+	switch a {
+	case AuthAlgoHMACSHA256:
+		return sha256.Size
+	default:
+		return 0
+	}
+}

shared/relay/auth/hmac/v2/generator.go

@@ -0,0 +1,45 @@
+package v2
+
+import (
+	"crypto/hmac"
+	"fmt"
+	"hash"
+	"strconv"
+	"time"
+)
+
+type Generator struct {
+	algo       func() hash.Hash
+	algoType   AuthAlgo
+	secret     []byte
+	timeToLive time.Duration
+}
+
+func NewGenerator(algo AuthAlgo, secret []byte, timeToLive time.Duration) (*Generator, error) {
+	algoFunc := algo.New()
+	if algoFunc == nil {
+		return nil, fmt.Errorf("unsupported auth algorithm: %s", algo)
+	}
+	return &Generator{
+		algo:       algoFunc,
+		algoType:   algo,
+		secret:     secret,
+		timeToLive: timeToLive,
+	}, nil
+}
+
+func (g *Generator) GenerateToken() (*Token, error) {
+	expirationTime := time.Now().Add(g.timeToLive).Unix()
+
+	payload := []byte(strconv.FormatInt(expirationTime, 10))
+
+	h := hmac.New(g.algo, g.secret)
+	h.Write(payload)
+	signature := h.Sum(nil)
+
+	return &Token{
+		AuthAlgo:  g.algoType,
+		Signature: signature,
+		Payload:   payload,
+	}, nil
+}

shared/relay/auth/hmac/v2/hmac_test.go

@@ -0,0 +1,110 @@
+package v2
+
+import (
+	"strconv"
+	"testing"
+	"time"
+)
+
+func TestGenerateCredentials(t *testing.T) {
+	secret := "supersecret"
+	timeToLive := 1 * time.Hour
+	g, err := NewGenerator(AuthAlgoHMACSHA256, []byte(secret), timeToLive)
+	if err != nil {
+		t.Fatalf("failed to create generator: %v", err)
+	}
+
+	token, err := g.GenerateToken()
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	if len(token.Payload) == 0 {
+		t.Fatalf("expected non-empty payload")
+	}
+
+	_, err = strconv.ParseInt(string(token.Payload), 10, 64)
+	if err != nil {
+		t.Fatalf("expected payload to be a valid unix timestamp, got %v", err)
+	}
+}
+
+func TestValidateCredentials(t *testing.T) {
+	secret := "supersecret"
+	timeToLive := 1 * time.Hour
+	g, err := NewGenerator(AuthAlgoHMACSHA256, []byte(secret), timeToLive)
+	if err != nil {
+		t.Fatalf("failed to create generator: %v", err)
+	}
+
+	token, err := g.GenerateToken()
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	v := NewValidator([]byte(secret))
+	if err := v.Validate(token.Marshal()); err != nil {
+		t.Fatalf("expected valid token: %s", err)
+	}
+}
+
+func TestInvalidSignature(t *testing.T) {
+	secret := "supersecret"
+	timeToLive := 1 * time.Hour
+	g, err := NewGenerator(AuthAlgoHMACSHA256, []byte(secret), timeToLive)
+	if err != nil {
+		t.Fatalf("failed to create generator: %v", err)
+	}
+
+	token, err := g.GenerateToken()
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	token.Signature = []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
+
+	v := NewValidator([]byte(secret))
+	if err := v.Validate(token.Marshal()); err == nil {
+		t.Fatalf("expected valid token: %s", err)
+	}
+}
+
+func TestExpired(t *testing.T) {
+	secret := "supersecret"
+	timeToLive := -1 * time.Hour
+	g, err := NewGenerator(AuthAlgoHMACSHA256, []byte(secret), timeToLive)
+	if err != nil {
+		t.Fatalf("failed to create generator: %v", err)
+	}
+
+	token, err := g.GenerateToken()
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	v := NewValidator([]byte(secret))
+	if err := v.Validate(token.Marshal()); err == nil {
+		t.Fatalf("expected valid token: %s", err)
+	}
+}
+
+func TestInvalidPayload(t *testing.T) {
+	secret := "supersecret"
+	timeToLive := 1 * time.Hour
+	g, err := NewGenerator(AuthAlgoHMACSHA256, []byte(secret), timeToLive)
+	if err != nil {
+		t.Fatalf("failed to create generator: %v", err)
+	}
+
+	token, err := g.GenerateToken()
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	token.Payload = []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
+
+	v := NewValidator([]byte(secret))
+	if err := v.Validate(token.Marshal()); err == nil {
+		t.Fatalf("expected invalid token due to invalid payload")
+	}
+}

shared/relay/auth/hmac/v2/token.go

@@ -0,0 +1,39 @@
+package v2
+
+import "errors"
+
+type Token struct {
+	AuthAlgo  AuthAlgo
+	Signature []byte
+	Payload   []byte
+}
+
+func (t *Token) Marshal() []byte {
+	size := 1 + len(t.Signature) + len(t.Payload)
+
+	buf := make([]byte, size)
+
+	buf[0] = byte(t.AuthAlgo)
+	copy(buf[1:], t.Signature)
+	copy(buf[1+len(t.Signature):], t.Payload)
+
+	return buf
+}
+
+func UnmarshalToken(data []byte) (*Token, error) {
+	if len(data) == 0 {
+		return nil, errors.New("invalid token data")
+	}
+
+	algo := AuthAlgo(data[0])
+	sigSize := algo.Size()
+	if len(data) < 1+sigSize {
+		return nil, errors.New("invalid token data: insufficient length")
+	}
+
+	return &Token{
+		AuthAlgo:  algo,
+		Signature: data[1 : 1+sigSize],
+		Payload:   data[1+sigSize:],
+	}, nil
+}

shared/relay/auth/hmac/v2/validator.go

@@ -0,0 +1,59 @@
+package v2
+
+import (
+	"crypto/hmac"
+	"errors"
+	"fmt"
+	"strconv"
+	"time"
+)
+
+const minLengthUnixTimestamp = 10
+
+type Validator struct {
+	secret []byte
+}
+
+func NewValidator(secret []byte) *Validator {
+	return &Validator{secret: secret}
+}
+
+func (v *Validator) Validate(data any) error {
+	d, ok := data.([]byte)
+	if !ok {
+		return fmt.Errorf("invalid data type")
+	}
+
+	token, err := UnmarshalToken(d)
+	if err != nil {
+		return fmt.Errorf("unmarshal token: %w", err)
+	}
+
+	if len(token.Payload) < minLengthUnixTimestamp {
+		return errors.New("invalid payload: insufficient length")
+	}
+
+	hashFunc := token.AuthAlgo.New()
+	if hashFunc == nil {
+		return fmt.Errorf("unsupported auth algorithm: %s", token.AuthAlgo)
+	}
+
+	h := hmac.New(hashFunc, v.secret)
+	h.Write(token.Payload)
+	expectedMAC := h.Sum(nil)
+
+	if !hmac.Equal(token.Signature, expectedMAC) {
+		return errors.New("invalid signature")
+	}
+
+	timestamp, err := strconv.ParseInt(string(token.Payload), 10, 64)
+	if err != nil {
+		return fmt.Errorf("invalid payload: %w", err)
+	}
+
+	if time.Now().Unix() > timestamp {
+		return fmt.Errorf("expired token")
+	}
+
+	return nil
+}

shared/relay/auth/hmac/validator.go

@@ -0,0 +1,33 @@
+package hmac
+
+import (
+	"crypto/sha256"
+	"fmt"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+)
+
+type TimedHMACValidator struct {
+	*TimedHMAC
+}
+
+func NewTimedHMACValidator(secret string, duration time.Duration) *TimedHMACValidator {
+	ta := NewTimedHMAC(secret, duration)
+	return &TimedHMACValidator{
+		ta,
+	}
+}
+
+func (a *TimedHMACValidator) Validate(credentials any) error {
+	b, ok := credentials.([]byte)
+	if !ok {
+		return fmt.Errorf("invalid credentials type")
+	}
+	c, err := unmarshalToken(b)
+	if err != nil {
+		log.Debugf("failed to unmarshal token: %s", err)
+		return err
+	}
+	return a.TimedHMAC.Validate(sha256.New, c)
+}

shared/relay/auth/validator.go

@@ -0,0 +1,28 @@
+package auth
+
+import (
+	"time"
+
+	auth "github.com/netbirdio/netbird/shared/relay/auth/hmac"
+	authv2 "github.com/netbirdio/netbird/shared/relay/auth/hmac/v2"
+)
+
+type TimedHMACValidator struct {
+	authenticatorV2 *authv2.Validator
+	authenticator   *auth.TimedHMACValidator
+}
+
+func NewTimedHMACValidator(secret []byte, duration time.Duration) *TimedHMACValidator {
+	return &TimedHMACValidator{
+		authenticatorV2: authv2.NewValidator(secret),
+		authenticator:   auth.NewTimedHMACValidator(string(secret), duration),
+	}
+}
+
+func (a *TimedHMACValidator) Validate(credentials any) error {
+	return a.authenticatorV2.Validate(credentials)
+}
+
+func (a *TimedHMACValidator) ValidateHelloMsgType(credentials any) error {
+	return a.authenticator.Validate(credentials)
+}

shared/relay/client/addr.go

@@ -0,0 +1,13 @@
+package client
+
+type RelayAddr struct {
+	addr string
+}
+
+func (a RelayAddr) Network() string {
+	return "relay"
+}
+
+func (a RelayAddr) String() string {
+	return a.addr
+}

shared/relay/client/client.go

@@ -0,0 +1,665 @@
+package client
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	auth "github.com/netbirdio/netbird/shared/relay/auth/hmac"
+	"github.com/netbirdio/netbird/shared/relay/client/dialer"
+	"github.com/netbirdio/netbird/shared/relay/client/dialer/quic"
+	"github.com/netbirdio/netbird/shared/relay/client/dialer/ws"
+	"github.com/netbirdio/netbird/relay/healthcheck"
+	"github.com/netbirdio/netbird/relay/messages"
+)
+
+const (
+	bufferSize            = 8820
+	serverResponseTimeout = 8 * time.Second
+)
+
+var (
+	ErrConnAlreadyExists = fmt.Errorf("connection already exists")
+)
+
+type internalStopFlag struct {
+	sync.Mutex
+	stop bool
+}
+
+func newInternalStopFlag() *internalStopFlag {
+	return &internalStopFlag{}
+}
+
+func (isf *internalStopFlag) set() {
+	isf.Lock()
+	defer isf.Unlock()
+	isf.stop = true
+}
+
+func (isf *internalStopFlag) isSet() bool {
+	isf.Lock()
+	defer isf.Unlock()
+	return isf.stop
+}
+
+// Msg carry the payload from the server to the client. With this struct, the net.Conn can free the buffer.
+type Msg struct {
+	Payload []byte
+
+	bufPool *sync.Pool
+	bufPtr  *[]byte
+}
+
+func (m *Msg) Free() {
+	m.bufPool.Put(m.bufPtr)
+}
+
+// connContainer is a container for the connection to the peer. It is responsible for managing the messages from the
+// server and forwarding them to the upper layer content reader.
+type connContainer struct {
+	log         *log.Entry
+	conn        *Conn
+	messages    chan Msg
+	msgChanLock sync.Mutex
+	closed      bool // flag to check if channel is closed
+	ctx         context.Context
+	cancel      context.CancelFunc
+}
+
+func newConnContainer(log *log.Entry, conn *Conn, messages chan Msg) *connContainer {
+	ctx, cancel := context.WithCancel(context.Background())
+	return &connContainer{
+		log:      log,
+		conn:     conn,
+		messages: messages,
+		ctx:      ctx,
+		cancel:   cancel,
+	}
+}
+
+func (cc *connContainer) writeMsg(msg Msg) {
+	cc.msgChanLock.Lock()
+	defer cc.msgChanLock.Unlock()
+
+	if cc.closed {
+		msg.Free()
+		return
+	}
+
+	select {
+	case cc.messages <- msg:
+	case <-cc.ctx.Done():
+		msg.Free()
+	default:
+		msg.Free()
+	}
+}
+
+func (cc *connContainer) close() {
+	cc.cancel()
+
+	cc.msgChanLock.Lock()
+	defer cc.msgChanLock.Unlock()
+
+	if cc.closed {
+		return
+	}
+
+	cc.closed = true
+	close(cc.messages)
+
+	for msg := range cc.messages {
+		msg.Free()
+	}
+}
+
+// Client is a client for the relay server. It is responsible for establishing a connection to the relay server and
+// managing connections to other peers. All exported functions are safe to call concurrently. After close the connection,
+// the client can be reused by calling Connect again. When the client is closed, all connections are closed too.
+// While the Connect is in progress, the OpenConn function will block until the connection is established with relay server.
+type Client struct {
+	log            *log.Entry
+	connectionURL  string
+	authTokenStore *auth.TokenStore
+	hashedID       messages.PeerID
+
+	bufPool *sync.Pool
+
+	relayConn        net.Conn
+	conns            map[messages.PeerID]*connContainer
+	serviceIsRunning bool
+	mu               sync.Mutex // protect serviceIsRunning and conns
+	readLoopMutex    sync.Mutex
+	wgReadLoop       sync.WaitGroup
+	instanceURL      *RelayAddr
+	muInstanceURL    sync.Mutex
+
+	onDisconnectListener func(string)
+	listenerMutex        sync.Mutex
+
+	stateSubscription *PeersStateSubscription
+}
+
+// NewClient creates a new client for the relay server. The client is not connected to the server until the Connect
+func NewClient(serverURL string, authTokenStore *auth.TokenStore, peerID string) *Client {
+	hashedID := messages.HashID(peerID)
+	relayLog := log.WithFields(log.Fields{"relay": serverURL})
+
+	c := &Client{
+		log:            relayLog,
+		connectionURL:  serverURL,
+		authTokenStore: authTokenStore,
+		hashedID:       hashedID,
+		bufPool: &sync.Pool{
+			New: func() any {
+				buf := make([]byte, bufferSize)
+				return &buf
+			},
+		},
+		conns: make(map[messages.PeerID]*connContainer),
+	}
+
+	c.log.Infof("create new relay connection: local peerID: %s, local peer hashedID: %s", peerID, hashedID)
+	return c
+}
+
+// Connect establishes a connection to the relay server. It blocks until the connection is established or an error occurs.
+func (c *Client) Connect(ctx context.Context) error {
+	c.log.Infof("connecting to relay server")
+	c.readLoopMutex.Lock()
+	defer c.readLoopMutex.Unlock()
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if c.serviceIsRunning {
+		return nil
+	}
+
+	instanceURL, err := c.connect(ctx)
+	if err != nil {
+		return err
+	}
+	c.muInstanceURL.Lock()
+	c.instanceURL = instanceURL
+	c.muInstanceURL.Unlock()
+
+	c.stateSubscription = NewPeersStateSubscription(c.log, c.relayConn, c.closeConnsByPeerID)
+
+	c.log = c.log.WithField("relay", instanceURL.String())
+	c.log.Infof("relay connection established")
+
+	c.serviceIsRunning = true
+
+	internallyStoppedFlag := newInternalStopFlag()
+	hc := healthcheck.NewReceiver(c.log)
+	go c.listenForStopEvents(ctx, hc, c.relayConn, internallyStoppedFlag)
+
+	c.wgReadLoop.Add(1)
+	go c.readLoop(hc, c.relayConn, internallyStoppedFlag)
+
+	return nil
+}
+
+// OpenConn create a new net.Conn for the destination peer ID. In case if the connection is in progress
+// to the relay server, the function will block until the connection is established or timed out. Otherwise,
+// it will return immediately.
+// It block until the server confirm the peer is online.
+// todo: what should happen if call with the same peerID with multiple times?
+func (c *Client) OpenConn(ctx context.Context, dstPeerID string) (net.Conn, error) {
+	peerID := messages.HashID(dstPeerID)
+
+	c.mu.Lock()
+	if !c.serviceIsRunning {
+		c.mu.Unlock()
+		return nil, fmt.Errorf("relay connection is not established")
+	}
+	_, ok := c.conns[peerID]
+	if ok {
+		c.mu.Unlock()
+		return nil, ErrConnAlreadyExists
+	}
+	c.mu.Unlock()
+
+	if err := c.stateSubscription.WaitToBeOnlineAndSubscribe(ctx, peerID); err != nil {
+		c.log.Errorf("peer not available: %s, %s", peerID, err)
+		return nil, err
+	}
+
+	c.log.Infof("remote peer is available, prepare the relayed connection: %s", peerID)
+	msgChannel := make(chan Msg, 100)
+
+	c.mu.Lock()
+	if !c.serviceIsRunning {
+		c.mu.Unlock()
+		return nil, fmt.Errorf("relay connection is not established")
+	}
+
+	c.muInstanceURL.Lock()
+	instanceURL := c.instanceURL
+	c.muInstanceURL.Unlock()
+	conn := NewConn(c, peerID, msgChannel, instanceURL)
+
+	_, ok = c.conns[peerID]
+	if ok {
+		c.mu.Unlock()
+		_ = conn.Close()
+		return nil, ErrConnAlreadyExists
+	}
+	c.conns[peerID] = newConnContainer(c.log, conn, msgChannel)
+	c.mu.Unlock()
+	return conn, nil
+}
+
+// ServerInstanceURL returns the address of the relay server. It could change after the close and reopen the connection.
+func (c *Client) ServerInstanceURL() (string, error) {
+	c.muInstanceURL.Lock()
+	defer c.muInstanceURL.Unlock()
+	if c.instanceURL == nil {
+		return "", fmt.Errorf("relay connection is not established")
+	}
+	return c.instanceURL.String(), nil
+}
+
+// SetOnDisconnectListener sets a function that will be called when the connection to the relay server is closed.
+func (c *Client) SetOnDisconnectListener(fn func(string)) {
+	c.listenerMutex.Lock()
+	defer c.listenerMutex.Unlock()
+	c.onDisconnectListener = fn
+}
+
+// HasConns returns true if there are connections.
+func (c *Client) HasConns() bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return len(c.conns) > 0
+}
+
+func (c *Client) Ready() bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.serviceIsRunning
+}
+
+// Close closes the connection to the relay server and all connections to other peers.
+func (c *Client) Close() error {
+	return c.close(true)
+}
+
+func (c *Client) connect(ctx context.Context) (*RelayAddr, error) {
+	rd := dialer.NewRaceDial(c.log, dialer.DefaultConnectionTimeout, c.connectionURL, quic.Dialer{}, ws.Dialer{})
+	conn, err := rd.Dial()
+	if err != nil {
+		return nil, err
+	}
+	c.relayConn = conn
+
+	instanceURL, err := c.handShake(ctx)
+	if err != nil {
+		cErr := conn.Close()
+		if cErr != nil {
+			c.log.Errorf("failed to close connection: %s", cErr)
+		}
+		return nil, err
+	}
+
+	return instanceURL, nil
+}
+
+func (c *Client) handShake(ctx context.Context) (*RelayAddr, error) {
+	msg, err := messages.MarshalAuthMsg(c.hashedID, c.authTokenStore.TokenBinary())
+	if err != nil {
+		c.log.Errorf("failed to marshal auth message: %s", err)
+		return nil, err
+	}
+
+	_, err = c.relayConn.Write(msg)
+	if err != nil {
+		c.log.Errorf("failed to send auth message: %s", err)
+		return nil, err
+	}
+	buf := make([]byte, messages.MaxHandshakeRespSize)
+	n, err := c.readWithTimeout(ctx, buf)
+	if err != nil {
+		c.log.Errorf("failed to read auth response: %s", err)
+		return nil, err
+	}
+
+	_, err = messages.ValidateVersion(buf[:n])
+	if err != nil {
+		return nil, fmt.Errorf("validate version: %w", err)
+	}
+
+	msgType, err := messages.DetermineServerMessageType(buf[:n])
+	if err != nil {
+		c.log.Errorf("failed to determine message type: %s", err)
+		return nil, err
+	}
+
+	if msgType != messages.MsgTypeAuthResponse {
+		c.log.Errorf("unexpected message type: %s", msgType)
+		return nil, fmt.Errorf("unexpected message type")
+	}
+
+	addr, err := messages.UnmarshalAuthResponse(buf[:n])
+	if err != nil {
+		return nil, err
+	}
+
+	return &RelayAddr{addr: addr}, nil
+}
+
+func (c *Client) readLoop(hc *healthcheck.Receiver, relayConn net.Conn, internallyStoppedFlag *internalStopFlag) {
+	var (
+		errExit error
+		n       int
+	)
+	for {
+		bufPtr := c.bufPool.Get().(*[]byte)
+		buf := *bufPtr
+		n, errExit = relayConn.Read(buf)
+		if errExit != nil {
+			c.log.Infof("start to Relay read loop exit")
+			c.mu.Lock()
+			if c.serviceIsRunning && !internallyStoppedFlag.isSet() {
+				c.log.Errorf("failed to read message from relay server: %s", errExit)
+			}
+			c.mu.Unlock()
+			c.bufPool.Put(bufPtr)
+			break
+		}
+
+		buf = buf[:n]
+
+		_, err := messages.ValidateVersion(buf)
+		if err != nil {
+			c.log.Errorf("failed to validate protocol version: %s", err)
+			c.bufPool.Put(bufPtr)
+			continue
+		}
+
+		msgType, err := messages.DetermineServerMessageType(buf)
+		if err != nil {
+			c.log.Errorf("failed to determine message type: %s", err)
+			c.bufPool.Put(bufPtr)
+			continue
+		}
+
+		if !c.handleMsg(msgType, buf, bufPtr, hc, internallyStoppedFlag) {
+			break
+		}
+	}
+
+	hc.Stop()
+
+	c.stateSubscription.Cleanup()
+	c.wgReadLoop.Done()
+	_ = c.close(false)
+	c.notifyDisconnected()
+}
+
+func (c *Client) handleMsg(msgType messages.MsgType, buf []byte, bufPtr *[]byte, hc *healthcheck.Receiver, internallyStoppedFlag *internalStopFlag) (continueLoop bool) {
+	switch msgType {
+	case messages.MsgTypeHealthCheck:
+		c.handleHealthCheck(hc, internallyStoppedFlag)
+		c.bufPool.Put(bufPtr)
+	case messages.MsgTypeTransport:
+		return c.handleTransportMsg(buf, bufPtr, internallyStoppedFlag)
+	case messages.MsgTypePeersOnline:
+		c.handlePeersOnlineMsg(buf)
+		c.bufPool.Put(bufPtr)
+		return true
+	case messages.MsgTypePeersWentOffline:
+		c.handlePeersWentOfflineMsg(buf)
+		c.bufPool.Put(bufPtr)
+		return true
+	case messages.MsgTypeClose:
+		c.log.Debugf("relay connection close by server")
+		c.bufPool.Put(bufPtr)
+		return false
+	}
+
+	return true
+}
+
+func (c *Client) handleHealthCheck(hc *healthcheck.Receiver, internallyStoppedFlag *internalStopFlag) {
+	msg := messages.MarshalHealthcheck()
+	_, wErr := c.relayConn.Write(msg)
+	if wErr != nil {
+		if c.serviceIsRunning && !internallyStoppedFlag.isSet() {
+			c.log.Errorf("failed to send heartbeat: %s", wErr)
+		}
+	}
+	hc.Heartbeat()
+}
+
+func (c *Client) handleTransportMsg(buf []byte, bufPtr *[]byte, internallyStoppedFlag *internalStopFlag) bool {
+	peerID, payload, err := messages.UnmarshalTransportMsg(buf)
+	if err != nil {
+		if c.serviceIsRunning && !internallyStoppedFlag.isSet() {
+			c.log.Errorf("failed to parse transport message: %v", err)
+		}
+
+		c.bufPool.Put(bufPtr)
+		return true
+	}
+
+	c.mu.Lock()
+	if !c.serviceIsRunning {
+		c.mu.Unlock()
+		c.bufPool.Put(bufPtr)
+		return false
+	}
+	container, ok := c.conns[*peerID]
+	c.mu.Unlock()
+	if !ok {
+		c.log.Errorf("peer not found: %s", peerID.String())
+		c.bufPool.Put(bufPtr)
+		return true
+	}
+	msg := Msg{
+		bufPool: c.bufPool,
+		bufPtr:  bufPtr,
+		Payload: payload,
+	}
+	container.writeMsg(msg)
+	return true
+}
+
+func (c *Client) writeTo(connReference *Conn, dstID messages.PeerID, payload []byte) (int, error) {
+	c.mu.Lock()
+	conn, ok := c.conns[dstID]
+	c.mu.Unlock()
+	if !ok {
+		return 0, net.ErrClosed
+	}
+
+	if conn.conn != connReference {
+		return 0, net.ErrClosed
+	}
+
+	// todo: use buffer pool instead of create new transport msg.
+	msg, err := messages.MarshalTransportMsg(dstID, payload)
+	if err != nil {
+		c.log.Errorf("failed to marshal transport message: %s", err)
+		return 0, err
+	}
+
+	// the write always return with 0 length because the underling does not support the size feedback.
+	_, err = c.relayConn.Write(msg)
+	if err != nil {
+		c.log.Errorf("failed to write transport message: %s", err)
+	}
+	return len(payload), err
+}
+
+func (c *Client) listenForStopEvents(ctx context.Context, hc *healthcheck.Receiver, conn net.Conn, internalStopFlag *internalStopFlag) {
+	for {
+		select {
+		case _, ok := <-hc.OnTimeout:
+			if !ok {
+				return
+			}
+			c.log.Errorf("health check timeout")
+			internalStopFlag.set()
+			if err := conn.Close(); err != nil {
+				// ignore the err handling because the readLoop will handle it
+				c.log.Warnf("failed to close connection: %s", err)
+			}
+			return
+		case <-ctx.Done():
+			err := c.close(true)
+			if err != nil {
+				c.log.Errorf("failed to teardown connection: %s", err)
+			}
+			return
+		}
+	}
+}
+
+func (c *Client) closeAllConns() {
+	for _, container := range c.conns {
+		container.close()
+	}
+	c.conns = make(map[messages.PeerID]*connContainer)
+}
+
+func (c *Client) closeConnsByPeerID(peerIDs []messages.PeerID) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	for _, peerID := range peerIDs {
+		container, ok := c.conns[peerID]
+		if !ok {
+			c.log.Warnf("can not close connection, peer not found: %s", peerID)
+			continue
+		}
+
+		container.log.Infof("remote peer has been disconnected, free up connection: %s", peerID)
+		container.close()
+		delete(c.conns, peerID)
+	}
+
+	if err := c.stateSubscription.UnsubscribeStateChange(peerIDs); err != nil {
+		c.log.Errorf("failed to unsubscribe from peer state change: %s, %s", peerIDs, err)
+	}
+}
+
+func (c *Client) closeConn(connReference *Conn, id messages.PeerID) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	container, ok := c.conns[id]
+	if !ok {
+		return net.ErrClosed
+	}
+
+	if container.conn != connReference {
+		return fmt.Errorf("conn reference mismatch")
+	}
+
+	if err := c.stateSubscription.UnsubscribeStateChange([]messages.PeerID{id}); err != nil {
+		container.log.Errorf("failed to unsubscribe from peer state change: %s", err)
+	}
+
+	c.log.Infof("free up connection to peer: %s", id)
+	delete(c.conns, id)
+	container.close()
+
+	return nil
+}
+
+func (c *Client) close(gracefullyExit bool) error {
+	c.readLoopMutex.Lock()
+	defer c.readLoopMutex.Unlock()
+
+	c.mu.Lock()
+	var err error
+	if !c.serviceIsRunning {
+		c.mu.Unlock()
+		c.log.Warn("relay connection was already marked as not running")
+		return nil
+	}
+	c.serviceIsRunning = false
+
+	c.muInstanceURL.Lock()
+	c.instanceURL = nil
+	c.muInstanceURL.Unlock()
+
+	c.log.Infof("closing all peer connections")
+	c.closeAllConns()
+	if gracefullyExit {
+		c.writeCloseMsg()
+	}
+	err = c.relayConn.Close()
+	c.mu.Unlock()
+
+	c.log.Infof("waiting for read loop to close")
+	c.wgReadLoop.Wait()
+	c.log.Infof("relay connection closed")
+	return err
+}
+
+func (c *Client) notifyDisconnected() {
+	c.listenerMutex.Lock()
+	defer c.listenerMutex.Unlock()
+
+	if c.onDisconnectListener == nil {
+		return
+	}
+	go c.onDisconnectListener(c.connectionURL)
+}
+
+func (c *Client) writeCloseMsg() {
+	msg := messages.MarshalCloseMsg()
+	_, err := c.relayConn.Write(msg)
+	if err != nil {
+		c.log.Errorf("failed to send close message: %s", err)
+	}
+}
+
+func (c *Client) readWithTimeout(ctx context.Context, buf []byte) (int, error) {
+	ctx, cancel := context.WithTimeout(ctx, serverResponseTimeout)
+	defer cancel()
+
+	readDone := make(chan struct{})
+	var (
+		n   int
+		err error
+	)
+
+	go func() {
+		n, err = c.relayConn.Read(buf)
+		close(readDone)
+	}()
+
+	select {
+	case <-ctx.Done():
+		return 0, fmt.Errorf("read operation timed out")
+	case <-readDone:
+		return n, err
+	}
+}
+
+func (c *Client) handlePeersOnlineMsg(buf []byte) {
+	peersID, err := messages.UnmarshalPeersOnlineMsg(buf)
+	if err != nil {
+		c.log.Errorf("failed to unmarshal peers online msg: %s", err)
+		return
+	}
+	c.stateSubscription.OnPeersOnline(peersID)
+}
+
+func (c *Client) handlePeersWentOfflineMsg(buf []byte) {
+	peersID, err := messages.UnMarshalPeersWentOffline(buf)
+	if err != nil {
+		c.log.Errorf("failed to unmarshal peers went offline msg: %s", err)
+		return
+	}
+	c.stateSubscription.OnPeersWentOffline(peersID)
+}

shared/relay/client/client_test.go

@@ -0,0 +1,743 @@
+package client
+
+import (
+	"context"
+	"net"
+	"os"
+	"testing"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"go.opentelemetry.io/otel"
+
+	"github.com/netbirdio/netbird/shared/relay/auth/allow"
+	"github.com/netbirdio/netbird/shared/relay/auth/hmac"
+	"github.com/netbirdio/netbird/util"
+
+	"github.com/netbirdio/netbird/relay/server"
+)
+
+var (
+	hmacTokenStore   = &hmac.TokenStore{}
+	serverListenAddr = "127.0.0.1:1234"
+	serverURL        = "rel://127.0.0.1:1234"
+	serverCfg        = server.Config{
+		Meter:          otel.Meter(""),
+		ExposedAddress: serverURL,
+		TLSSupport:     false,
+		AuthValidator:  &allow.Auth{},
+	}
+)
+
+func TestMain(m *testing.M) {
+	_ = util.InitLog("debug", util.LogConsole)
+	code := m.Run()
+	os.Exit(code)
+}
+
+func TestClient(t *testing.T) {
+	ctx := context.Background()
+
+	srv, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+	go func() {
+		listenCfg := server.ListenerConfig{Address: serverListenAddr}
+		err := srv.Listen(listenCfg)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	defer func() {
+		err := srv.Shutdown(ctx)
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	// wait for server to start
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+	t.Log("alice connecting to server")
+	clientAlice := NewClient(serverURL, hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+	defer clientAlice.Close()
+
+	t.Log("placeholder connecting to server")
+	clientPlaceHolder := NewClient(serverURL, hmacTokenStore, "clientPlaceHolder")
+	err = clientPlaceHolder.Connect(ctx)
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+	defer clientPlaceHolder.Close()
+
+	t.Log("Bob connecting to server")
+	clientBob := NewClient(serverURL, hmacTokenStore, "bob")
+	err = clientBob.Connect(ctx)
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+	defer clientBob.Close()
+
+	t.Log("Alice open connection to Bob")
+	connAliceToBob, err := clientAlice.OpenConn(ctx, "bob")
+	if err != nil {
+		t.Fatalf("failed to bind channel: %s", err)
+	}
+
+	t.Log("Bob open connection to Alice")
+	connBobToAlice, err := clientBob.OpenConn(ctx, "alice")
+	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)
+	}
+	log.Debugf("alice sent message to bob")
+
+	buf := make([]byte, 65535)
+	n, err := connBobToAlice.Read(buf)
+	if err != nil {
+		t.Fatalf("failed to read from channel: %s", err)
+	}
+	log.Debugf("on new message from alice to bob")
+
+	if payload != string(buf[:n]) {
+		t.Fatalf("expected %s, got %s", payload, string(buf[:n]))
+	}
+}
+
+func TestRegistration(t *testing.T) {
+	ctx := context.Background()
+	srvCfg := server.ListenerConfig{Address: serverListenAddr}
+	srv, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+	go func() {
+		err := srv.Listen(srvCfg)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	// wait for server to start
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	clientAlice := NewClient(serverURL, hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
+	if err != nil {
+		_ = srv.Shutdown(ctx)
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+	err = clientAlice.Close()
+	if err != nil {
+		t.Errorf("failed to close conn: %s", err)
+	}
+	err = srv.Shutdown(ctx)
+	if err != nil {
+		t.Errorf("failed to close server: %s", err)
+	}
+}
+
+func TestRegistrationTimeout(t *testing.T) {
+	ctx := context.Background()
+	fakeUDPListener, err := net.ListenUDP("udp", &net.UDPAddr{
+		Port: 1234,
+		IP:   net.ParseIP("0.0.0.0"),
+	})
+	if err != nil {
+		t.Fatalf("failed to bind UDP server: %s", err)
+	}
+	defer func(fakeUDPListener *net.UDPConn) {
+		_ = fakeUDPListener.Close()
+	}(fakeUDPListener)
+
+	fakeTCPListener, err := net.ListenTCP("tcp", &net.TCPAddr{
+		Port: 1234,
+		IP:   net.ParseIP("0.0.0.0"),
+	})
+	if err != nil {
+		t.Fatalf("failed to bind TCP server: %s", err)
+	}
+	defer func(fakeTCPListener *net.TCPListener) {
+		_ = fakeTCPListener.Close()
+	}(fakeTCPListener)
+
+	clientAlice := NewClient("127.0.0.1:1234", hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
+	if err == nil {
+		t.Errorf("failed to connect to server: %s", err)
+	}
+	log.Debugf("%s", err)
+	err = clientAlice.Close()
+	if err != nil {
+		t.Errorf("failed to close conn: %s", err)
+	}
+}
+
+func TestEcho(t *testing.T) {
+	ctx := context.Background()
+	idAlice := "alice"
+	idBob := "bob"
+	srvCfg := server.ListenerConfig{Address: serverListenAddr}
+	srv, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+	go func() {
+		err := srv.Listen(srvCfg)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	defer func() {
+		err := srv.Shutdown(ctx)
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	// wait for servers to start
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	clientAlice := NewClient(serverURL, hmacTokenStore, idAlice)
+	err = clientAlice.Connect(ctx)
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+	defer func() {
+		err := clientAlice.Close()
+		if err != nil {
+			t.Errorf("failed to close Alice client: %s", err)
+		}
+	}()
+
+	clientBob := NewClient(serverURL, hmacTokenStore, idBob)
+	err = clientBob.Connect(ctx)
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+	defer func() {
+		err := clientBob.Close()
+		if err != nil {
+			t.Errorf("failed to close Bob client: %s", err)
+		}
+	}()
+
+	connAliceToBob, err := clientAlice.OpenConn(ctx, idBob)
+	if err != nil {
+		t.Fatalf("failed to bind channel: %s", err)
+	}
+
+	connBobToAlice, err := clientBob.OpenConn(ctx, 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]))
+	}
+}
+
+func TestBindToUnavailabePeer(t *testing.T) {
+	ctx := context.Background()
+
+	srvCfg := server.ListenerConfig{Address: serverListenAddr}
+	srv, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+	go func() {
+		err := srv.Listen(srvCfg)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	defer func() {
+		log.Infof("closing server")
+		err := srv.Shutdown(ctx)
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	// wait for servers to start
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	clientAlice := NewClient(serverURL, hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
+	if err != nil {
+		t.Errorf("failed to connect to server: %s", err)
+	}
+	_, err = clientAlice.OpenConn(ctx, "bob")
+	if err == nil {
+		t.Errorf("expected error when binding to unavailable peer, got nil")
+	}
+
+	log.Infof("closing client")
+	err = clientAlice.Close()
+	if err != nil {
+		t.Errorf("failed to close client: %s", err)
+	}
+}
+
+func TestBindReconnect(t *testing.T) {
+	ctx := context.Background()
+
+	srvCfg := server.ListenerConfig{Address: serverListenAddr}
+	srv, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+	go func() {
+		err := srv.Listen(srvCfg)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	defer func() {
+		log.Infof("closing server")
+		err := srv.Shutdown(ctx)
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	// wait for servers to start
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	clientAlice := NewClient(serverURL, hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+
+	clientBob := NewClient(serverURL, hmacTokenStore, "bob")
+	err = clientBob.Connect(ctx)
+	if err != nil {
+		t.Errorf("failed to connect to server: %s", err)
+	}
+
+	_, err = clientAlice.OpenConn(ctx, "bob")
+	if err != nil {
+		t.Fatalf("failed to bind channel: %s", err)
+	}
+
+	chBob, err := clientBob.OpenConn(ctx, "alice")
+	if err != nil {
+		t.Errorf("failed to bind channel: %s", err)
+	}
+
+	log.Infof("closing client Alice")
+	err = clientAlice.Close()
+	if err != nil {
+		t.Errorf("failed to close client: %s", err)
+	}
+
+	clientAlice = NewClient(serverURL, hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
+	if err != nil {
+		t.Errorf("failed to connect to server: %s", err)
+	}
+
+	chAlice, err := clientAlice.OpenConn(ctx, "bob")
+	if err != nil {
+		t.Errorf("failed to bind channel: %s", err)
+	}
+
+	testString := "hello alice, I am bob"
+	_, err = chBob.Write([]byte(testString))
+	if err == nil {
+		t.Errorf("expected error when writing to channel, got nil")
+	}
+
+	chBob, err = clientBob.OpenConn(ctx, "alice")
+	if err != nil {
+		t.Errorf("failed to bind channel: %s", err)
+	}
+
+	_, err = chBob.Write([]byte(testString))
+	if err != nil {
+		t.Errorf("failed to write to channel: %s", err)
+	}
+
+	buf := make([]byte, 65535)
+	n, err := chAlice.Read(buf)
+	if err != nil {
+		t.Errorf("failed to read from channel: %s", err)
+	}
+
+	if testString != string(buf[:n]) {
+		t.Errorf("expected %s, got %s", testString, string(buf[:n]))
+	}
+
+	log.Infof("closing client")
+	err = clientAlice.Close()
+	if err != nil {
+		t.Errorf("failed to close client: %s", err)
+	}
+}
+
+func TestCloseConn(t *testing.T) {
+	ctx := context.Background()
+
+	srvCfg := server.ListenerConfig{Address: serverListenAddr}
+	srv, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+	go func() {
+		err := srv.Listen(srvCfg)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	defer func() {
+		log.Infof("closing server")
+		err := srv.Shutdown(ctx)
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	// wait for servers to start
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	bob := NewClient(serverURL, hmacTokenStore, "bob")
+	err = bob.Connect(ctx)
+	if err != nil {
+		t.Errorf("failed to connect to server: %s", err)
+	}
+
+	clientAlice := NewClient(serverURL, hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
+	if err != nil {
+		t.Errorf("failed to connect to server: %s", err)
+	}
+
+	conn, err := clientAlice.OpenConn(ctx, "bob")
+	if err != nil {
+		t.Errorf("failed to bind channel: %s", err)
+	}
+
+	log.Infof("closing connection")
+	err = conn.Close()
+	if err != nil {
+		t.Errorf("failed to close connection: %s", err)
+	}
+
+	_, err = conn.Read(make([]byte, 1))
+	if err == nil {
+		t.Errorf("unexpected reading from closed connection")
+	}
+
+	_, err = conn.Write([]byte("hello"))
+	if err == nil {
+		t.Errorf("unexpected writing from closed connection")
+	}
+}
+
+func TestCloseRelayConn(t *testing.T) {
+	ctx := context.Background()
+
+	srvCfg := server.ListenerConfig{Address: serverListenAddr}
+	srv, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+	go func() {
+		err := srv.Listen(srvCfg)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	defer func() {
+		err := srv.Shutdown(ctx)
+		if err != nil {
+			log.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	// wait for servers to start
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	bob := NewClient(serverURL, hmacTokenStore, "bob")
+	err = bob.Connect(ctx)
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+
+	clientAlice := NewClient(serverURL, hmacTokenStore, "alice")
+	err = clientAlice.Connect(ctx)
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+
+	conn, err := clientAlice.OpenConn(ctx, "bob")
+	if err != nil {
+		t.Errorf("failed to bind channel: %s", err)
+	}
+
+	_ = clientAlice.relayConn.Close()
+
+	_, err = conn.Read(make([]byte, 1))
+	if err == nil {
+		t.Errorf("unexpected reading from closed connection")
+	}
+
+	_, err = clientAlice.OpenConn(ctx, "bob")
+	if err == nil {
+		t.Errorf("unexpected opening connection to closed server")
+	}
+}
+
+func TestCloseByServer(t *testing.T) {
+	ctx := context.Background()
+
+	srvCfg := server.ListenerConfig{Address: serverListenAddr}
+	srv1, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+
+	go func() {
+		err := srv1.Listen(srvCfg)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	// wait for servers to start
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	idAlice := "alice"
+	log.Debugf("connect by alice")
+	relayClient := NewClient(serverURL, hmacTokenStore, idAlice)
+	if err = relayClient.Connect(ctx); err != nil {
+		log.Fatalf("failed to connect to server: %s", err)
+	}
+	defer func() {
+		if err := relayClient.Close(); err != nil {
+			log.Errorf("failed to close client: %s", err)
+		}
+	}()
+
+	disconnected := make(chan struct{})
+	relayClient.SetOnDisconnectListener(func(_ string) {
+		log.Infof("client disconnected")
+		close(disconnected)
+	})
+
+	err = srv1.Shutdown(ctx)
+	if err != nil {
+		t.Fatalf("failed to close server: %s", err)
+	}
+
+	select {
+	case <-disconnected:
+	case <-time.After(3 * time.Second):
+		log.Errorf("timeout waiting for client to disconnect")
+	}
+
+	_, err = relayClient.OpenConn(ctx, "bob")
+	if err == nil {
+		t.Errorf("unexpected opening connection to closed server")
+	}
+}
+
+func TestCloseByClient(t *testing.T) {
+	ctx := context.Background()
+
+	srvCfg := server.ListenerConfig{Address: serverListenAddr}
+	srv, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+	go func() {
+		err := srv.Listen(srvCfg)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	// wait for servers to start
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	idAlice := "alice"
+	log.Debugf("connect by alice")
+	relayClient := NewClient(serverURL, hmacTokenStore, idAlice)
+	err = relayClient.Connect(ctx)
+	if err != nil {
+		log.Fatalf("failed to connect to server: %s", err)
+	}
+
+	err = relayClient.Close()
+	if err != nil {
+		t.Errorf("failed to close client: %s", err)
+	}
+
+	_, err = relayClient.OpenConn(ctx, "bob")
+	if err == nil {
+		t.Errorf("unexpected opening connection to closed server")
+	}
+
+	err = srv.Shutdown(ctx)
+	if err != nil {
+		t.Fatalf("failed to close server: %s", err)
+	}
+}
+
+func TestCloseNotDrainedChannel(t *testing.T) {
+	ctx := context.Background()
+	idAlice := "alice"
+	idBob := "bob"
+	srvCfg := server.ListenerConfig{Address: serverListenAddr}
+	srv, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+	go func() {
+		err := srv.Listen(srvCfg)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	defer func() {
+		err := srv.Shutdown(ctx)
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	// wait for servers to start
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	clientAlice := NewClient(serverURL, hmacTokenStore, idAlice)
+	err = clientAlice.Connect(ctx)
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+	defer func() {
+		err := clientAlice.Close()
+		if err != nil {
+			t.Errorf("failed to close Alice client: %s", err)
+		}
+	}()
+
+	clientBob := NewClient(serverURL, hmacTokenStore, idBob)
+	err = clientBob.Connect(ctx)
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+	defer func() {
+		err := clientBob.Close()
+		if err != nil {
+			t.Errorf("failed to close Bob client: %s", err)
+		}
+	}()
+
+	connAliceToBob, err := clientAlice.OpenConn(ctx, idBob)
+	if err != nil {
+		t.Fatalf("failed to bind channel: %s", err)
+	}
+
+	connBobToAlice, err := clientBob.OpenConn(ctx, idAlice)
+	if err != nil {
+		t.Fatalf("failed to bind channel: %s", err)
+	}
+
+	payload := "hello bob, I am alice"
+	// the internal channel buffer size is 2. So we should overflow it
+	for i := 0; i < 5; i++ {
+		_, err = connAliceToBob.Write([]byte(payload))
+		if err != nil {
+			t.Fatalf("failed to write to channel: %s", err)
+		}
+
+	}
+
+	// wait for delivery
+	time.Sleep(1 * time.Second)
+	err = connBobToAlice.Close()
+	if err != nil {
+		t.Errorf("failed to close channel: %s", err)
+	}
+}
+
+func waitForServerToStart(errChan chan error) error {
+	select {
+	case err := <-errChan:
+		if err != nil {
+			return err
+		}
+	case <-time.After(300 * time.Millisecond):
+		return nil
+	}
+	return nil
+}

shared/relay/client/conn.go

@@ -0,0 +1,74 @@
+package client
+
+import (
+	"net"
+	"time"
+
+	"github.com/netbirdio/netbird/relay/messages"
+)
+
+// Conn represent a connection to a relayed remote peer.
+type Conn struct {
+	client      *Client
+	dstID       messages.PeerID
+	messageChan chan Msg
+	instanceURL *RelayAddr
+}
+
+// NewConn creates a new connection to a relayed remote peer.
+// client: the client instance, it used to send messages to the destination peer
+// dstID: the destination peer ID
+// messageChan: the channel where the messages will be received
+// instanceURL: the relay instance URL, it used to get the proper server instance address for the remote peer
+func NewConn(client *Client, dstID messages.PeerID, messageChan chan Msg, instanceURL *RelayAddr) *Conn {
+	c := &Conn{
+		client:      client,
+		dstID:       dstID,
+		messageChan: messageChan,
+		instanceURL: instanceURL,
+	}
+
+	return c
+}
+
+func (c *Conn) Write(p []byte) (n int, err error) {
+	return c.client.writeTo(c, c.dstID, p)
+}
+
+func (c *Conn) Read(b []byte) (n int, err error) {
+	msg, ok := <-c.messageChan
+	if !ok {
+		return 0, net.ErrClosed
+	}
+
+	n = copy(b, msg.Payload)
+	msg.Free()
+	return n, nil
+}
+
+func (c *Conn) Close() error {
+	return c.client.closeConn(c, c.dstID)
+}
+
+func (c *Conn) LocalAddr() net.Addr {
+	return c.client.relayConn.LocalAddr()
+}
+
+func (c *Conn) RemoteAddr() net.Addr {
+	return c.instanceURL
+}
+
+func (c *Conn) SetDeadline(t time.Time) error {
+	//TODO implement me
+	panic("SetDeadline is not implemented")
+}
+
+func (c *Conn) SetReadDeadline(t time.Time) error {
+	//TODO implement me
+	panic("SetReadDeadline is not implemented")
+}
+
+func (c *Conn) SetWriteDeadline(t time.Time) error {
+	//TODO implement me
+	panic("SetReadDeadline is not implemented")
+}

shared/relay/client/dialer/net/err.go

@@ -0,0 +1,7 @@
+package net
+
+import "errors"
+
+var (
+	ErrClosedByServer = errors.New("closed by server")
+)

shared/relay/client/dialer/quic/conn.go

@@ -0,0 +1,97 @@
+package quic
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"time"
+
+	"github.com/quic-go/quic-go"
+	log "github.com/sirupsen/logrus"
+
+	netErr "github.com/netbirdio/netbird/shared/relay/client/dialer/net"
+)
+
+const (
+	Network = "quic"
+)
+
+type Addr struct {
+	addr string
+}
+
+func (a Addr) Network() string {
+	return Network
+}
+
+func (a Addr) String() string {
+	return a.addr
+}
+
+type Conn struct {
+	session quic.Connection
+	ctx     context.Context
+}
+
+func NewConn(session quic.Connection) net.Conn {
+	return &Conn{
+		session: session,
+		ctx:     context.Background(),
+	}
+}
+
+func (c *Conn) Read(b []byte) (n int, err error) {
+	dgram, err := c.session.ReceiveDatagram(c.ctx)
+	if err != nil {
+		return 0, c.remoteCloseErrHandling(err)
+	}
+
+	n = copy(b, dgram)
+	return n, nil
+}
+
+func (c *Conn) Write(b []byte) (int, error) {
+	err := c.session.SendDatagram(b)
+	if err != nil {
+		err = c.remoteCloseErrHandling(err)
+		log.Errorf("failed to write to QUIC stream: %v", err)
+		return 0, err
+	}
+	return len(b), nil
+}
+
+func (c *Conn) RemoteAddr() net.Addr {
+	return c.session.RemoteAddr()
+}
+
+func (c *Conn) LocalAddr() net.Addr {
+	if c.session != nil {
+		return c.session.LocalAddr()
+	}
+	return Addr{addr: "unknown"}
+}
+
+func (c *Conn) SetReadDeadline(t time.Time) error {
+	return fmt.Errorf("SetReadDeadline is not implemented")
+}
+
+func (c *Conn) SetWriteDeadline(t time.Time) error {
+	return fmt.Errorf("SetWriteDeadline is not implemented")
+}
+
+func (c *Conn) SetDeadline(t time.Time) error {
+	return nil
+}
+
+func (c *Conn) Close() error {
+	return c.session.CloseWithError(0, "normal closure")
+}
+
+func (c *Conn) remoteCloseErrHandling(err error) error {
+	var appErr *quic.ApplicationError
+	if errors.As(err, &appErr) && appErr.ErrorCode == 0x0 {
+		return netErr.ErrClosedByServer
+	}
+	return err
+}

shared/relay/client/dialer/quic/quic.go

@@ -0,0 +1,99 @@
+package quic
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"strings"
+	"time"
+
+	"github.com/quic-go/quic-go"
+	log "github.com/sirupsen/logrus"
+
+	quictls "github.com/netbirdio/netbird/relay/tls"
+	nbnet "github.com/netbirdio/netbird/util/net"
+)
+
+type Dialer struct {
+}
+
+func (d Dialer) Protocol() string {
+	return Network
+}
+
+func (d Dialer) Dial(ctx context.Context, address string) (net.Conn, error) {
+	quicURL, err := prepareURL(address)
+	if err != nil {
+		return nil, err
+	}
+
+	// Get the base TLS config
+	tlsClientConfig := quictls.ClientQUICTLSConfig()
+
+	// Set ServerName to hostname if not an IP address
+	host, _, splitErr := net.SplitHostPort(quicURL)
+	if splitErr == nil && net.ParseIP(host) == nil {
+		// It's a hostname, not an IP - modify directly
+		tlsClientConfig.ServerName = host
+	}
+
+	quicConfig := &quic.Config{
+		KeepAlivePeriod:   30 * time.Second,
+		MaxIdleTimeout:    4 * time.Minute,
+		EnableDatagrams:   true,
+		InitialPacketSize: 1452,
+	}
+
+	udpConn, err := nbnet.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
+	if err != nil {
+		log.Errorf("failed to listen on UDP: %s", err)
+		return nil, err
+	}
+
+	udpAddr, err := net.ResolveUDPAddr("udp", quicURL)
+	if err != nil {
+		log.Errorf("failed to resolve UDP address: %s", err)
+		return nil, err
+	}
+
+	session, err := quic.Dial(ctx, udpConn, udpAddr, tlsClientConfig, quicConfig)
+	if err != nil {
+		if errors.Is(err, context.Canceled) {
+			return nil, err
+		}
+		log.Errorf("failed to dial to Relay server via QUIC '%s': %s", quicURL, err)
+		return nil, err
+	}
+
+	conn := NewConn(session)
+	return conn, nil
+}
+
+func prepareURL(address string) (string, error) {
+	var host string
+	var defaultPort string
+
+	switch {
+	case strings.HasPrefix(address, "rels://"):
+		host = address[7:]
+		defaultPort = "443"
+	case strings.HasPrefix(address, "rel://"):
+		host = address[6:]
+		defaultPort = "80"
+	default:
+		return "", fmt.Errorf("unsupported scheme: %s", address)
+	}
+
+	finalHost, finalPort, err := net.SplitHostPort(host)
+	if err != nil {
+		if strings.Contains(err.Error(), "missing port") {
+			return host + ":" + defaultPort, nil
+		}
+
+		// return any other split error as is
+		return "", err
+	}
+
+	return finalHost + ":" + finalPort, nil
+}

shared/relay/client/dialer/race_dialer.go

@@ -0,0 +1,98 @@
+package dialer
+
+import (
+	"context"
+	"errors"
+	"net"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	DefaultConnectionTimeout = 30 * time.Second
+)
+
+type DialeFn interface {
+	Dial(ctx context.Context, address string) (net.Conn, error)
+	Protocol() string
+}
+
+type dialResult struct {
+	Conn     net.Conn
+	Protocol string
+	Err      error
+}
+
+type RaceDial struct {
+	log               *log.Entry
+	serverURL         string
+	dialerFns         []DialeFn
+	connectionTimeout time.Duration
+}
+
+func NewRaceDial(log *log.Entry, connectionTimeout time.Duration, serverURL string, dialerFns ...DialeFn) *RaceDial {
+	return &RaceDial{
+		log:               log,
+		serverURL:         serverURL,
+		dialerFns:         dialerFns,
+		connectionTimeout: connectionTimeout,
+	}
+}
+
+func (r *RaceDial) Dial() (net.Conn, error) {
+	connChan := make(chan dialResult, len(r.dialerFns))
+	winnerConn := make(chan net.Conn, 1)
+	abortCtx, abort := context.WithCancel(context.Background())
+	defer abort()
+
+	for _, dfn := range r.dialerFns {
+		go r.dial(dfn, abortCtx, connChan)
+	}
+
+	go r.processResults(connChan, winnerConn, abort)
+
+	conn, ok := <-winnerConn
+	if !ok {
+		return nil, errors.New("failed to dial to Relay server on any protocol")
+	}
+	return conn, nil
+}
+
+func (r *RaceDial) dial(dfn DialeFn, abortCtx context.Context, connChan chan dialResult) {
+	ctx, cancel := context.WithTimeout(abortCtx, r.connectionTimeout)
+	defer cancel()
+
+	r.log.Infof("dialing Relay server via %s", dfn.Protocol())
+	conn, err := dfn.Dial(ctx, r.serverURL)
+	connChan <- dialResult{Conn: conn, Protocol: dfn.Protocol(), Err: err}
+}
+
+func (r *RaceDial) processResults(connChan chan dialResult, winnerConn chan net.Conn, abort context.CancelFunc) {
+	var hasWinner bool
+	for i := 0; i < len(r.dialerFns); i++ {
+		dr := <-connChan
+		if dr.Err != nil {
+			if errors.Is(dr.Err, context.Canceled) {
+				r.log.Infof("connection attempt aborted via: %s", dr.Protocol)
+			} else {
+				r.log.Errorf("failed to dial via %s: %s", dr.Protocol, dr.Err)
+			}
+			continue
+		}
+
+		if hasWinner {
+			if cerr := dr.Conn.Close(); cerr != nil {
+				r.log.Warnf("failed to close connection via %s: %s", dr.Protocol, cerr)
+			}
+			continue
+		}
+
+		r.log.Infof("successfully dialed via: %s", dr.Protocol)
+
+		abort()
+		hasWinner = true
+		winnerConn <- dr.Conn
+	}
+	close(winnerConn)
+}

shared/relay/client/dialer/race_dialer_test.go

@@ -0,0 +1,252 @@
+package dialer
+
+import (
+	"context"
+	"errors"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/sirupsen/logrus"
+)
+
+type MockAddr struct {
+	network string
+}
+
+func (m *MockAddr) Network() string {
+	return m.network
+}
+
+func (m *MockAddr) String() string {
+	return "1.2.3.4"
+}
+
+// MockDialer is a mock implementation of DialeFn
+type MockDialer struct {
+	dialFunc    func(ctx context.Context, address string) (net.Conn, error)
+	protocolStr string
+}
+
+func (m *MockDialer) Dial(ctx context.Context, address string) (net.Conn, error) {
+	return m.dialFunc(ctx, address)
+}
+
+func (m *MockDialer) Protocol() string {
+	return m.protocolStr
+}
+
+// MockConn implements net.Conn for testing
+type MockConn struct {
+	remoteAddr net.Addr
+}
+
+func (m *MockConn) Read(b []byte) (n int, err error) {
+	return 0, nil
+}
+
+func (m *MockConn) Write(b []byte) (n int, err error) {
+	return 0, nil
+}
+
+func (m *MockConn) Close() error {
+	return nil
+}
+
+func (m *MockConn) LocalAddr() net.Addr {
+	return nil
+}
+
+func (m *MockConn) RemoteAddr() net.Addr {
+	return m.remoteAddr
+}
+
+func (m *MockConn) SetDeadline(t time.Time) error {
+	return nil
+}
+
+func (m *MockConn) SetReadDeadline(t time.Time) error {
+	return nil
+}
+
+func (m *MockConn) SetWriteDeadline(t time.Time) error {
+	return nil
+}
+
+func TestRaceDialEmptyDialers(t *testing.T) {
+	logger := logrus.NewEntry(logrus.New())
+	serverURL := "test.server.com"
+
+	rd := NewRaceDial(logger, DefaultConnectionTimeout, serverURL)
+	conn, err := rd.Dial()
+	if err == nil {
+		t.Errorf("Expected an error with empty dialers, got nil")
+	}
+	if conn != nil {
+		t.Errorf("Expected nil connection with empty dialers, got %v", conn)
+	}
+}
+
+func TestRaceDialSingleSuccessfulDialer(t *testing.T) {
+	logger := logrus.NewEntry(logrus.New())
+	serverURL := "test.server.com"
+	proto := "test-protocol"
+
+	mockConn := &MockConn{
+		remoteAddr: &MockAddr{network: proto},
+	}
+
+	mockDialer := &MockDialer{
+		dialFunc: func(ctx context.Context, address string) (net.Conn, error) {
+			return mockConn, nil
+		},
+		protocolStr: proto,
+	}
+
+	rd := NewRaceDial(logger, DefaultConnectionTimeout, serverURL, mockDialer)
+	conn, err := rd.Dial()
+	if err != nil {
+		t.Errorf("Expected no error, got %v", err)
+	}
+	if conn == nil {
+		t.Errorf("Expected non-nil connection")
+	}
+}
+
+func TestRaceDialMultipleDialersWithOneSuccess(t *testing.T) {
+	logger := logrus.NewEntry(logrus.New())
+	serverURL := "test.server.com"
+	proto2 := "protocol2"
+
+	mockConn2 := &MockConn{
+		remoteAddr: &MockAddr{network: proto2},
+	}
+
+	mockDialer1 := &MockDialer{
+		dialFunc: func(ctx context.Context, address string) (net.Conn, error) {
+			return nil, errors.New("first dialer failed")
+		},
+		protocolStr: "proto1",
+	}
+
+	mockDialer2 := &MockDialer{
+		dialFunc: func(ctx context.Context, address string) (net.Conn, error) {
+			return mockConn2, nil
+		},
+		protocolStr: "proto2",
+	}
+
+	rd := NewRaceDial(logger, DefaultConnectionTimeout, serverURL, mockDialer1, mockDialer2)
+	conn, err := rd.Dial()
+	if err != nil {
+		t.Errorf("Expected no error, got %v", err)
+	}
+	if conn.RemoteAddr().Network() != proto2 {
+		t.Errorf("Expected connection with protocol %s, got %s", proto2, conn.RemoteAddr().Network())
+	}
+	_ = conn.Close()
+}
+
+func TestRaceDialTimeout(t *testing.T) {
+	logger := logrus.NewEntry(logrus.New())
+	serverURL := "test.server.com"
+
+	mockDialer := &MockDialer{
+		dialFunc: func(ctx context.Context, address string) (net.Conn, error) {
+			<-ctx.Done()
+			return nil, ctx.Err()
+		},
+		protocolStr: "proto1",
+	}
+
+	rd := NewRaceDial(logger, 3*time.Second, serverURL, mockDialer)
+	conn, err := rd.Dial()
+	if err == nil {
+		t.Errorf("Expected an error, got nil")
+	}
+	if conn != nil {
+		t.Errorf("Expected nil connection, got %v", conn)
+	}
+}
+
+func TestRaceDialAllDialersFail(t *testing.T) {
+	logger := logrus.NewEntry(logrus.New())
+	serverURL := "test.server.com"
+
+	mockDialer1 := &MockDialer{
+		dialFunc: func(ctx context.Context, address string) (net.Conn, error) {
+			return nil, errors.New("first dialer failed")
+		},
+		protocolStr: "protocol1",
+	}
+
+	mockDialer2 := &MockDialer{
+		dialFunc: func(ctx context.Context, address string) (net.Conn, error) {
+			return nil, errors.New("second dialer failed")
+		},
+		protocolStr: "protocol2",
+	}
+
+	rd := NewRaceDial(logger, DefaultConnectionTimeout, serverURL, mockDialer1, mockDialer2)
+	conn, err := rd.Dial()
+	if err == nil {
+		t.Errorf("Expected an error, got nil")
+	}
+	if conn != nil {
+		t.Errorf("Expected nil connection, got %v", conn)
+	}
+}
+
+func TestRaceDialFirstSuccessfulDialerWins(t *testing.T) {
+	logger := logrus.NewEntry(logrus.New())
+	serverURL := "test.server.com"
+	proto1 := "protocol1"
+	proto2 := "protocol2"
+
+	mockConn1 := &MockConn{
+		remoteAddr: &MockAddr{network: proto1},
+	}
+
+	mockConn2 := &MockConn{
+		remoteAddr: &MockAddr{network: proto2},
+	}
+
+	mockDialer1 := &MockDialer{
+		dialFunc: func(ctx context.Context, address string) (net.Conn, error) {
+			time.Sleep(1 * time.Second)
+			return mockConn1, nil
+		},
+		protocolStr: proto1,
+	}
+
+	mock2err := make(chan error)
+	mockDialer2 := &MockDialer{
+		dialFunc: func(ctx context.Context, address string) (net.Conn, error) {
+			<-ctx.Done()
+			mock2err <- ctx.Err()
+			return mockConn2, ctx.Err()
+		},
+		protocolStr: proto2,
+	}
+
+	rd := NewRaceDial(logger, DefaultConnectionTimeout, serverURL, mockDialer1, mockDialer2)
+	conn, err := rd.Dial()
+	if err != nil {
+		t.Errorf("Expected no error, got %v", err)
+	}
+	if conn == nil {
+		t.Errorf("Expected non-nil connection")
+	}
+	if conn != mockConn1 {
+		t.Errorf("Expected first connection, got %v", conn)
+	}
+
+	select {
+	case <-time.After(3 * time.Second):
+		t.Errorf("Timed out waiting for second dialer to finish")
+	case err := <-mock2err:
+		if !errors.Is(err, context.Canceled) {
+			t.Errorf("Expected context.Canceled error, got %v", err)
+		}
+	}
+}

shared/relay/client/dialer/ws/addr.go

@@ -0,0 +1,17 @@
+package ws
+
+const (
+	Network = "ws"
+)
+
+type WebsocketAddr struct {
+	addr string
+}
+
+func (a WebsocketAddr) Network() string {
+	return Network
+}
+
+func (a WebsocketAddr) String() string {
+	return a.addr
+}

shared/relay/client/dialer/ws/conn.go

@@ -0,0 +1,67 @@
+package ws
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"time"
+
+	"github.com/coder/websocket"
+)
+
+type Conn struct {
+	ctx context.Context
+	*websocket.Conn
+	remoteAddr WebsocketAddr
+}
+
+func NewConn(wsConn *websocket.Conn, serverAddress string) net.Conn {
+	return &Conn{
+		ctx:        context.Background(),
+		Conn:       wsConn,
+		remoteAddr: WebsocketAddr{serverAddress},
+	}
+}
+
+func (c *Conn) Read(b []byte) (n int, err error) {
+	t, ioReader, err := c.Conn.Reader(c.ctx)
+	if err != nil {
+		// todo use ErrClosedByServer
+		return 0, err
+	}
+
+	if t != websocket.MessageBinary {
+		return 0, fmt.Errorf("unexpected message type")
+	}
+
+	return ioReader.Read(b)
+}
+
+func (c *Conn) Write(b []byte) (n int, err error) {
+	err = c.Conn.Write(c.ctx, websocket.MessageBinary, b)
+	return 0, err
+}
+
+func (c *Conn) RemoteAddr() net.Addr {
+	return c.remoteAddr
+}
+
+func (c *Conn) LocalAddr() net.Addr {
+	return WebsocketAddr{addr: "unknown"}
+}
+
+func (c *Conn) SetReadDeadline(t time.Time) error {
+	return fmt.Errorf("SetReadDeadline is not implemented")
+}
+
+func (c *Conn) SetWriteDeadline(t time.Time) error {
+	return fmt.Errorf("SetWriteDeadline is not implemented")
+}
+
+func (c *Conn) SetDeadline(t time.Time) error {
+	return fmt.Errorf("SetDeadline is not implemented")
+}
+
+func (c *Conn) Close() error {
+	return c.Conn.CloseNow()
+}

shared/relay/client/dialer/ws/ws.go

@@ -0,0 +1,90 @@
+package ws
+
+import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"errors"
+	"fmt"
+	"net"
+	"net/http"
+	"net/url"
+	"strings"
+
+	"github.com/coder/websocket"
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/relay/server/listener/ws"
+	"github.com/netbirdio/netbird/util/embeddedroots"
+	nbnet "github.com/netbirdio/netbird/util/net"
+)
+
+type Dialer struct {
+}
+
+func (d Dialer) Protocol() string {
+	return "WS"
+}
+
+func (d Dialer) Dial(ctx context.Context, address string) (net.Conn, error) {
+	wsURL, err := prepareURL(address)
+	if err != nil {
+		return nil, err
+	}
+
+	opts := &websocket.DialOptions{
+		HTTPClient: httpClientNbDialer(),
+	}
+
+	parsedURL, err := url.Parse(wsURL)
+	if err != nil {
+		return nil, err
+	}
+	parsedURL.Path = ws.URLPath
+
+	wsConn, resp, err := websocket.Dial(ctx, parsedURL.String(), opts)
+	if err != nil {
+		if errors.Is(err, context.Canceled) {
+			return nil, err
+		}
+		log.Errorf("failed to dial to Relay server '%s': %s", wsURL, err)
+		return nil, err
+	}
+	if resp.Body != nil {
+		_ = resp.Body.Close()
+	}
+
+	conn := NewConn(wsConn, address)
+	return conn, nil
+}
+
+func prepareURL(address string) (string, error) {
+	if !strings.HasPrefix(address, "rel:") && !strings.HasPrefix(address, "rels:") {
+		return "", fmt.Errorf("unsupported scheme: %s", address)
+	}
+
+	return strings.Replace(address, "rel", "ws", 1), nil
+}
+
+func httpClientNbDialer() *http.Client {
+	customDialer := nbnet.NewDialer()
+
+	certPool, err := x509.SystemCertPool()
+	if err != nil || certPool == nil {
+		log.Debugf("System cert pool not available; falling back to embedded cert, error: %v", err)
+		certPool = embeddedroots.Get()
+	}
+
+	customTransport := &http.Transport{
+		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
+			return customDialer.DialContext(ctx, network, addr)
+		},
+		TLSClientConfig: &tls.Config{
+			RootCAs: certPool,
+		},
+	}
+
+	return &http.Client{
+		Transport: customTransport,
+	}
+}

shared/relay/client/doc.go

@@ -0,0 +1,12 @@
+/*
+Package client contains the implementation of the Relay client.
+
+The Relay client is responsible for establishing a connection with the Relay server and sending and receiving messages,
+Keep persistent connection with the Relay server and handle the connection issues.
+It uses the WebSocket protocol for communication and optionally supports TLS (Transport Layer Security).
+
+If a peer wants to communicate with a peer 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. The peers
+negotiate the common relay instance via signaling service.
+*/
+package client

shared/relay/client/go.sum

@@ -0,0 +1,10 @@
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

shared/relay/client/guard.go

@@ -0,0 +1,149 @@
+package client
+
+import (
+	"context"
+	"time"
+
+	"github.com/cenkalti/backoff/v4"
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	// TODO: make it configurable, the manager should validate all configurable parameters
+	reconnectingTimeout = 60 * time.Second
+)
+
+// Guard manage the reconnection tries to the Relay server in case of disconnection event.
+type Guard struct {
+	// OnNewRelayClient is a channel that is used to notify the relay manager about a new relay client instance.
+	OnNewRelayClient chan *Client
+	OnReconnected    chan struct{}
+	serverPicker     *ServerPicker
+}
+
+// NewGuard creates a new guard for the relay client.
+func NewGuard(sp *ServerPicker) *Guard {
+	g := &Guard{
+		OnNewRelayClient: make(chan *Client, 1),
+		OnReconnected:    make(chan struct{}, 1),
+		serverPicker:     sp,
+	}
+	return g
+}
+
+// StartReconnectTrys is called when the relay client is disconnected from the relay server.
+// It attempts to reconnect to the relay server. The function first tries a quick reconnect
+// to the same server that was used before, if the server URL is still valid. If the quick
+// reconnect fails, it starts a ticker to periodically attempt server picking until it
+// succeeds or the context is done.
+//
+// Parameters:
+// - ctx: The context to control the lifecycle of the reconnection attempts.
+// - relayClient: The relay client instance that was disconnected.
+// todo prevent multiple reconnection instances. In the current usage it should not happen, but it is better to prevent
+func (g *Guard) StartReconnectTrys(ctx context.Context, relayClient *Client) {
+	// try to reconnect to the same server
+	if ok := g.tryToQuickReconnect(ctx, relayClient); ok {
+		g.notifyReconnected()
+		return
+	}
+
+	// start a ticker to pick a new server
+	ticker := exponentTicker(ctx)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			if err := g.retry(ctx); err != nil {
+				log.Errorf("failed to pick new Relay server: %s", err)
+				continue
+			}
+			return
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
+func (g *Guard) tryToQuickReconnect(parentCtx context.Context, rc *Client) bool {
+	if rc == nil {
+		return false
+	}
+
+	if !g.isServerURLStillValid(rc) {
+		return false
+	}
+
+	if cancelled := waiteBeforeRetry(parentCtx); !cancelled {
+		return false
+	}
+
+	log.Infof("try to reconnect to Relay server: %s", rc.connectionURL)
+
+	if err := rc.Connect(parentCtx); err != nil {
+		log.Errorf("failed to reconnect to relay server: %s", err)
+		return false
+	}
+	return true
+}
+
+func (g *Guard) retry(ctx context.Context) error {
+	log.Infof("try to pick up a new Relay server")
+	relayClient, err := g.serverPicker.PickServer(ctx)
+	if err != nil {
+		return err
+	}
+
+	// prevent to work with a deprecated Relay client instance
+	g.drainRelayClientChan()
+
+	g.OnNewRelayClient <- relayClient
+	return nil
+}
+
+func (g *Guard) drainRelayClientChan() {
+	select {
+	case <-g.OnNewRelayClient:
+	default:
+	}
+}
+
+func (g *Guard) isServerURLStillValid(rc *Client) bool {
+	for _, url := range g.serverPicker.ServerURLs.Load().([]string) {
+		if url == rc.connectionURL {
+			return true
+		}
+	}
+	return false
+}
+
+func (g *Guard) notifyReconnected() {
+	select {
+	case g.OnReconnected <- struct{}{}:
+	default:
+	}
+}
+
+func exponentTicker(ctx context.Context) *backoff.Ticker {
+	bo := backoff.WithContext(&backoff.ExponentialBackOff{
+		InitialInterval: 2 * time.Second,
+		Multiplier:      2,
+		MaxInterval:     reconnectingTimeout,
+		Clock:           backoff.SystemClock,
+	}, ctx)
+
+	return backoff.NewTicker(bo)
+}
+
+func waiteBeforeRetry(ctx context.Context) bool {
+	timer := time.NewTimer(1500 * time.Millisecond)
+	defer timer.Stop()
+
+	select {
+	case <-timer.C:
+		return true
+	case <-ctx.Done():
+		return false
+	}
+}

shared/relay/client/manager.go

@@ -0,0 +1,404 @@
+package client
+
+import (
+	"container/list"
+	"context"
+	"fmt"
+	"net"
+	"reflect"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	relayAuth "github.com/netbirdio/netbird/shared/relay/auth/hmac"
+)
+
+var (
+	relayCleanupInterval = 60 * time.Second
+	keepUnusedServerTime = 5 * time.Second
+
+	ErrRelayClientNotConnected = fmt.Errorf("relay client not connected")
+)
+
+// RelayTrack hold the relay clients for the foreign 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
+	err         error
+	created     time.Time
+}
+
+func NewRelayTrack() *RelayTrack {
+	return &RelayTrack{
+		created: time.Now(),
+	}
+}
+
+type OnServerCloseListener func()
+
+// Manager is a manager for the relay client instances. It establishes one persistent connection to the given relay URL
+// and automatically reconnect to them in case disconnection.
+// The manager also manage temporary relay connection. If a client wants to communicate with a 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
+	peerID       string
+	running      bool
+	tokenStore   *relayAuth.TokenStore
+	serverPicker *ServerPicker
+
+	relayClient *Client
+	// the guard logic can overwrite the relayClient variable, this mutex protect the usage of the variable
+	relayClientMu  sync.RWMutex
+	reconnectGuard *Guard
+
+	relayClients      map[string]*RelayTrack
+	relayClientsMutex sync.RWMutex
+
+	onDisconnectedListeners map[string]*list.List
+	onReconnectedListenerFn func()
+	listenerLock            sync.Mutex
+}
+
+// NewManager creates a new manager instance.
+// The serverURL address can be empty. In this case, the manager will not serve.
+func NewManager(ctx context.Context, serverURLs []string, peerID string) *Manager {
+	tokenStore := &relayAuth.TokenStore{}
+
+	m := &Manager{
+		ctx:        ctx,
+		peerID:     peerID,
+		tokenStore: tokenStore,
+		serverPicker: &ServerPicker{
+			TokenStore: tokenStore,
+			PeerID:     peerID,
+		},
+		relayClients:            make(map[string]*RelayTrack),
+		onDisconnectedListeners: make(map[string]*list.List),
+	}
+	m.serverPicker.ServerURLs.Store(serverURLs)
+	m.reconnectGuard = NewGuard(m.serverPicker)
+	return m
+}
+
+// Serve starts the manager, attempting to establish a connection with the relay server.
+// If the connection fails, it will keep trying to reconnect in the background.
+// Additionally, it starts a cleanup loop to remove unused relay connections.
+// The manager will automatically reconnect to the relay server in case of disconnection.
+func (m *Manager) Serve() error {
+	if m.running {
+		return fmt.Errorf("manager already serving")
+	}
+	m.running = true
+	log.Debugf("starting relay client manager with %v relay servers", m.serverPicker.ServerURLs.Load())
+
+	client, err := m.serverPicker.PickServer(m.ctx)
+	if err != nil {
+		go m.reconnectGuard.StartReconnectTrys(m.ctx, nil)
+	} else {
+		m.storeClient(client)
+	}
+
+	go m.listenGuardEvent(m.ctx)
+	go m.startCleanupLoop()
+	return err
+}
+
+// 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. It returns back with a net.Conn what represent the remote peer connection.
+func (m *Manager) OpenConn(ctx context.Context, serverAddress, peerKey string) (net.Conn, error) {
+	m.relayClientMu.RLock()
+	defer m.relayClientMu.RUnlock()
+
+	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(ctx, peerKey)
+	} else {
+		log.Debugf("open peer connection via foreign server: %s", serverAddress)
+		netConn, err = m.openConnVia(ctx, serverAddress, peerKey)
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	return netConn, err
+}
+
+// Ready returns true if the home Relay client is connected to the relay server.
+func (m *Manager) Ready() bool {
+	m.relayClientMu.RLock()
+	defer m.relayClientMu.RUnlock()
+
+	if m.relayClient == nil {
+		return false
+	}
+	return m.relayClient.Ready()
+}
+
+func (m *Manager) SetOnReconnectedListener(f func()) {
+	m.listenerLock.Lock()
+	defer m.listenerLock.Unlock()
+
+	m.onReconnectedListenerFn = f
+}
+
+// AddCloseListener adds a listener to the given server instance address. The listener will be called if the connection
+// closed.
+func (m *Manager) AddCloseListener(serverAddress string, onClosedListener OnServerCloseListener) error {
+	m.relayClientMu.RLock()
+	defer m.relayClientMu.RUnlock()
+
+	if m.relayClient == nil {
+		return ErrRelayClientNotConnected
+	}
+
+	foreign, err := m.isForeignServer(serverAddress)
+	if err != nil {
+		return err
+	}
+
+	var listenerAddr string
+	if foreign {
+		listenerAddr = serverAddress
+	} else {
+		listenerAddr = m.relayClient.connectionURL
+	}
+	m.addListener(listenerAddr, onClosedListener)
+	return nil
+}
+
+// RelayInstanceAddress 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) RelayInstanceAddress() (string, error) {
+	m.relayClientMu.RLock()
+	defer m.relayClientMu.RUnlock()
+
+	if m.relayClient == nil {
+		return "", ErrRelayClientNotConnected
+	}
+	return m.relayClient.ServerInstanceURL()
+}
+
+// ServerURLs returns the addresses of the relay servers.
+func (m *Manager) ServerURLs() []string {
+	return m.serverPicker.ServerURLs.Load().([]string)
+}
+
+// HasRelayAddress returns true if the manager is serving. With this method can check if the peer can communicate with
+// Relay service.
+func (m *Manager) HasRelayAddress() bool {
+	return len(m.serverPicker.ServerURLs.Load().([]string)) > 0
+}
+
+func (m *Manager) UpdateServerURLs(serverURLs []string) {
+	log.Infof("update relay server URLs: %v", serverURLs)
+	m.serverPicker.ServerURLs.Store(serverURLs)
+}
+
+// UpdateToken updates the token in the token store.
+func (m *Manager) UpdateToken(token *relayAuth.Token) error {
+	return m.tokenStore.UpdateToken(token)
+}
+
+func (m *Manager) openConnVia(ctx context.Context, 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()
+		if rt.err != nil {
+			return nil, rt.err
+		}
+		return rt.relayClient.OpenConn(ctx, 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()
+		if rt.err != nil {
+			return nil, rt.err
+		}
+		return rt.relayClient.OpenConn(ctx, 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(serverAddress, m.tokenStore, m.peerID)
+	err := relayClient.Connect(m.ctx)
+	if err != nil {
+		rt.err = err
+		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(m.onServerDisconnected)
+	rt.relayClient = relayClient
+	rt.Unlock()
+
+	conn, err := relayClient.OpenConn(ctx, peerKey)
+	if err != nil {
+		return nil, err
+	}
+	return conn, nil
+}
+
+func (m *Manager) onServerConnected() {
+	m.listenerLock.Lock()
+	defer m.listenerLock.Unlock()
+
+	if m.onReconnectedListenerFn == nil {
+		return
+	}
+	go m.onReconnectedListenerFn()
+}
+
+// onServerDisconnected start to reconnection for home server only
+func (m *Manager) onServerDisconnected(serverAddress string) {
+	m.relayClientMu.Lock()
+	if serverAddress == m.relayClient.connectionURL {
+		go func(client *Client) {
+			m.reconnectGuard.StartReconnectTrys(m.ctx, client)
+		}(m.relayClient)
+	}
+	m.relayClientMu.Unlock()
+
+	m.notifyOnDisconnectListeners(serverAddress)
+}
+
+func (m *Manager) listenGuardEvent(ctx context.Context) {
+	for {
+		select {
+		case <-m.reconnectGuard.OnReconnected:
+			m.onServerConnected()
+		case rc := <-m.reconnectGuard.OnNewRelayClient:
+			m.storeClient(rc)
+			m.onServerConnected()
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
+func (m *Manager) storeClient(client *Client) {
+	m.relayClientMu.Lock()
+	defer m.relayClientMu.Unlock()
+
+	m.relayClient = client
+	m.relayClient.SetOnDisconnectListener(m.onServerDisconnected)
+}
+
+func (m *Manager) isForeignServer(address string) (bool, error) {
+	rAddr, err := m.relayClient.ServerInstanceURL()
+	if err != nil {
+		return false, fmt.Errorf("relay client not connected")
+	}
+	return rAddr != address, nil
+}
+
+func (m *Manager) startCleanupLoop() {
+	ticker := time.NewTicker(relayCleanupInterval)
+	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 the connection failed to the server the relay client will be nil
+		// but the instance will be kept in the relayClients until the next locking
+		if rt.err != nil {
+			rt.Unlock()
+			continue
+		}
+
+		if time.Since(rt.created) <= keepUnusedServerTime {
+			rt.Unlock()
+			continue
+		}
+
+		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 OnServerCloseListener) {
+	m.listenerLock.Lock()
+	defer m.listenerLock.Unlock()
+	l, ok := m.onDisconnectedListeners[serverAddress]
+	if !ok {
+		l = list.New()
+	}
+	for e := l.Front(); e != nil; e = e.Next() {
+		if reflect.ValueOf(e.Value).Pointer() == reflect.ValueOf(onClosedListener).Pointer() {
+			return
+		}
+	}
+	l.PushBack(onClosedListener)
+	m.onDisconnectedListeners[serverAddress] = l
+}
+
+func (m *Manager) notifyOnDisconnectListeners(serverAddress string) {
+	m.listenerLock.Lock()
+	defer m.listenerLock.Unlock()
+
+	l, ok := m.onDisconnectedListeners[serverAddress]
+	if !ok {
+		return
+	}
+	for e := l.Front(); e != nil; e = e.Next() {
+		go e.Value.(OnServerCloseListener)()
+	}
+	delete(m.onDisconnectedListeners, serverAddress)
+}

shared/relay/client/manager_test.go

@@ -0,0 +1,469 @@
+package client
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"go.opentelemetry.io/otel"
+
+	"github.com/netbirdio/netbird/shared/relay/auth/allow"
+	"github.com/netbirdio/netbird/relay/server"
+)
+
+func TestEmptyURL(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	mgr := NewManager(ctx, nil, "alice")
+	err := mgr.Serve()
+	if err == nil {
+		t.Errorf("expected error, got nil")
+	}
+}
+
+func TestForeignConn(t *testing.T) {
+	ctx := context.Background()
+
+	lstCfg1 := server.ListenerConfig{
+		Address: "localhost:1234",
+	}
+
+	srv1, err := server.NewServer(server.Config{
+		Meter:          otel.Meter(""),
+		ExposedAddress: lstCfg1.Address,
+		TLSSupport:     false,
+		AuthValidator:  &allow.Auth{},
+	})
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+	go func() {
+		err := srv1.Listen(lstCfg1)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	defer func() {
+		err := srv1.Shutdown(ctx)
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	srvCfg2 := server.ListenerConfig{
+		Address: "localhost:2234",
+	}
+	srv2, err := server.NewServer(server.Config{
+		Meter:          otel.Meter(""),
+		ExposedAddress: srvCfg2.Address,
+		TLSSupport:     false,
+		AuthValidator:  &allow.Auth{},
+	})
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan2 := make(chan error, 1)
+	go func() {
+		err := srv2.Listen(srvCfg2)
+		if err != nil {
+			errChan2 <- err
+		}
+	}()
+
+	defer func() {
+		err := srv2.Shutdown(ctx)
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	if err := waitForServerToStart(errChan2); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	mCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	clientAlice := NewManager(mCtx, toURL(lstCfg1), "alice")
+	if err := clientAlice.Serve(); err != nil {
+		t.Fatalf("failed to serve manager: %s", err)
+	}
+
+	clientBob := NewManager(mCtx, toURL(srvCfg2), "bob")
+	if err := clientBob.Serve(); err != nil {
+		t.Fatalf("failed to serve manager: %s", err)
+	}
+	bobsSrvAddr, err := clientBob.RelayInstanceAddress()
+	if err != nil {
+		t.Fatalf("failed to get relay address: %s", err)
+	}
+	connAliceToBob, err := clientAlice.OpenConn(ctx, bobsSrvAddr, "bob")
+	if err != nil {
+		t.Fatalf("failed to bind channel: %s", err)
+	}
+	connBobToAlice, err := clientBob.OpenConn(ctx, bobsSrvAddr, "alice")
+	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]))
+	}
+}
+
+func TestForeginConnClose(t *testing.T) {
+	ctx := context.Background()
+
+	srvCfg1 := server.ListenerConfig{
+		Address: "localhost:1234",
+	}
+	srv1, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+	go func() {
+		err := srv1.Listen(srvCfg1)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	defer func() {
+		err := srv1.Shutdown(ctx)
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	srvCfg2 := server.ListenerConfig{
+		Address: "localhost:2234",
+	}
+	srv2, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan2 := make(chan error, 1)
+	go func() {
+		err := srv2.Listen(srvCfg2)
+		if err != nil {
+			errChan2 <- err
+		}
+	}()
+
+	defer func() {
+		err := srv2.Shutdown(ctx)
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	if err := waitForServerToStart(errChan2); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	mCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	mgrBob := NewManager(mCtx, toURL(srvCfg2), "bob")
+	if err := mgrBob.Serve(); err != nil {
+		t.Fatalf("failed to serve manager: %s", err)
+	}
+
+	mgr := NewManager(mCtx, toURL(srvCfg1), "alice")
+	err = mgr.Serve()
+	if err != nil {
+		t.Fatalf("failed to serve manager: %s", err)
+	}
+	conn, err := mgr.OpenConn(ctx, toURL(srvCfg2)[0], "bob")
+	if err != nil {
+		t.Fatalf("failed to bind channel: %s", err)
+	}
+
+	err = conn.Close()
+	if err != nil {
+		t.Fatalf("failed to close connection: %s", err)
+	}
+}
+
+func TestForeignAutoClose(t *testing.T) {
+	ctx := context.Background()
+	relayCleanupInterval = 1 * time.Second
+	keepUnusedServerTime = 2 * time.Second
+
+	srvCfg1 := server.ListenerConfig{
+		Address: "localhost:1234",
+	}
+	srv1, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+	go func() {
+		t.Log("binding server 1.")
+		if err := srv1.Listen(srvCfg1); err != nil {
+			errChan <- err
+		}
+	}()
+
+	defer func() {
+		t.Logf("closing server 1.")
+		if err := srv1.Shutdown(ctx); err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+		t.Logf("server 1. closed")
+	}()
+
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	srvCfg2 := server.ListenerConfig{
+		Address: "localhost:2234",
+	}
+	srv2, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan2 := make(chan error, 1)
+	go func() {
+		t.Log("binding server 2.")
+		err := srv2.Listen(srvCfg2)
+		if err != nil {
+			errChan2 <- err
+		}
+	}()
+	defer func() {
+		t.Logf("closing server 2.")
+		err := srv2.Shutdown(ctx)
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+		t.Logf("server 2 closed.")
+	}()
+
+	if err := waitForServerToStart(errChan2); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	idAlice := "alice"
+	t.Log("connect to server 1.")
+	mCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	mgr := NewManager(mCtx, toURL(srvCfg1), idAlice)
+	err = mgr.Serve()
+	if err != nil {
+		t.Fatalf("failed to serve manager: %s", err)
+	}
+
+	// Set up a disconnect listener to track when foreign server disconnects
+	foreignServerURL := toURL(srvCfg2)[0]
+	disconnected := make(chan struct{})
+	onDisconnect := func() {
+		select {
+		case disconnected <- struct{}{}:
+		default:
+		}
+	}
+
+	t.Log("open connection to another peer")
+	if _, err = mgr.OpenConn(ctx, foreignServerURL, "anotherpeer"); err == nil {
+		t.Fatalf("should have failed to open connection to another peer")
+	}
+
+	// Add the disconnect listener after the connection attempt
+	if err := mgr.AddCloseListener(foreignServerURL, onDisconnect); err != nil {
+		t.Logf("failed to add close listener (expected if connection failed): %s", err)
+	}
+
+	// Wait for cleanup to happen
+	timeout := relayCleanupInterval + keepUnusedServerTime + 2*time.Second
+	t.Logf("waiting for relay cleanup: %s", timeout)
+
+	select {
+	case <-disconnected:
+		t.Log("foreign relay connection cleaned up successfully")
+	case <-time.After(timeout):
+		t.Log("timeout waiting for cleanup - this might be expected if connection never established")
+	}
+
+	t.Logf("closing manager")
+}
+
+func TestAutoReconnect(t *testing.T) {
+	ctx := context.Background()
+
+	srvCfg := server.ListenerConfig{
+		Address: "localhost:1234",
+	}
+	srv, err := server.NewServer(serverCfg)
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+	go func() {
+		if err := srv.Listen(srvCfg); err != nil {
+			errChan <- err
+		}
+	}()
+
+	defer func() {
+		err := srv.Shutdown(ctx)
+		if err != nil {
+			log.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	mCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	clientBob := NewManager(mCtx, toURL(srvCfg), "bob")
+	err = clientBob.Serve()
+	if err != nil {
+		t.Fatalf("failed to serve manager: %s", err)
+	}
+
+	clientAlice := NewManager(mCtx, toURL(srvCfg), "alice")
+	err = clientAlice.Serve()
+	if err != nil {
+		t.Fatalf("failed to serve manager: %s", err)
+	}
+	ra, err := clientAlice.RelayInstanceAddress()
+	if err != nil {
+		t.Errorf("failed to get relay address: %s", err)
+	}
+	conn, err := clientAlice.OpenConn(ctx, ra, "bob")
+	if err != nil {
+		t.Errorf("failed to bind channel: %s", err)
+	}
+
+	t.Log("closing client relay connection")
+	// todo figure out moc server
+	_ = clientAlice.relayClient.relayConn.Close()
+	t.Log("start test reading")
+	_, err = conn.Read(make([]byte, 1))
+	if err == nil {
+		t.Errorf("unexpected reading from closed connection")
+	}
+
+	log.Infof("waiting for reconnection")
+	time.Sleep(reconnectingTimeout + 1*time.Second)
+
+	log.Infof("reopent the connection")
+	_, err = clientAlice.OpenConn(ctx, ra, "bob")
+	if err != nil {
+		t.Errorf("failed to open channel: %s", err)
+	}
+}
+
+func TestNotifierDoubleAdd(t *testing.T) {
+	ctx := context.Background()
+
+	listenerCfg1 := server.ListenerConfig{
+		Address: "localhost:1234",
+	}
+	srv, err := server.NewServer(server.Config{
+		Meter:          otel.Meter(""),
+		ExposedAddress: listenerCfg1.Address,
+		TLSSupport:     false,
+		AuthValidator:  &allow.Auth{},
+	})
+	if err != nil {
+		t.Fatalf("failed to create server: %s", err)
+	}
+	errChan := make(chan error, 1)
+	go func() {
+		if err := srv.Listen(listenerCfg1); err != nil {
+			errChan <- err
+		}
+	}()
+
+	defer func() {
+		if err := srv.Shutdown(ctx); err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	if err := waitForServerToStart(errChan); err != nil {
+		t.Fatalf("failed to start server: %s", err)
+	}
+
+	log.Debugf("connect by alice")
+	mCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	clientBob := NewManager(mCtx, toURL(listenerCfg1), "bob")
+	if err = clientBob.Serve(); err != nil {
+		t.Fatalf("failed to serve manager: %s", err)
+	}
+
+	clientAlice := NewManager(mCtx, toURL(listenerCfg1), "alice")
+	if err = clientAlice.Serve(); err != nil {
+		t.Fatalf("failed to serve manager: %s", err)
+	}
+
+	conn1, err := clientAlice.OpenConn(ctx, clientAlice.ServerURLs()[0], "bob")
+	if err != nil {
+		t.Fatalf("failed to bind channel: %s", err)
+	}
+
+	fnCloseListener := OnServerCloseListener(func() {
+		log.Infof("close listener")
+	})
+
+	err = clientAlice.AddCloseListener(clientAlice.ServerURLs()[0], fnCloseListener)
+	if err != nil {
+		t.Fatalf("failed to add close listener: %s", err)
+	}
+
+	err = clientAlice.AddCloseListener(clientAlice.ServerURLs()[0], fnCloseListener)
+	if err != nil {
+		t.Fatalf("failed to add close listener: %s", err)
+	}
+
+	err = conn1.Close()
+	if err != nil {
+		t.Errorf("failed to close connection: %s", err)
+	}
+
+}
+
+func toURL(address server.ListenerConfig) []string {
+	return []string{"rel://" + address.Address}
+}

shared/relay/client/peer_subscription.go

@@ -0,0 +1,191 @@
+package client
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/relay/messages"
+)
+
+const (
+	OpenConnectionTimeout = 30 * time.Second
+)
+
+type relayedConnWriter interface {
+	Write(p []byte) (n int, err error)
+}
+
+// PeersStateSubscription manages subscriptions to peer state changes (online/offline)
+// over a relay connection. It allows tracking peers' availability and handling offline
+// events via a callback. We get online notification from the server only once.
+type PeersStateSubscription struct {
+	log             *log.Entry
+	relayConn       relayedConnWriter
+	offlineCallback func(peerIDs []messages.PeerID)
+
+	listenForOfflinePeers map[messages.PeerID]struct{}
+	waitingPeers          map[messages.PeerID]chan struct{}
+	mu                    sync.Mutex // Mutex to protect access to waitingPeers and listenForOfflinePeers
+}
+
+func NewPeersStateSubscription(log *log.Entry, relayConn relayedConnWriter, offlineCallback func(peerIDs []messages.PeerID)) *PeersStateSubscription {
+	return &PeersStateSubscription{
+		log:                   log,
+		relayConn:             relayConn,
+		offlineCallback:       offlineCallback,
+		listenForOfflinePeers: make(map[messages.PeerID]struct{}),
+		waitingPeers:          make(map[messages.PeerID]chan struct{}),
+	}
+}
+
+// OnPeersOnline should be called when a notification is received that certain peers have come online.
+// It checks if any of the peers are being waited on and signals their availability.
+func (s *PeersStateSubscription) OnPeersOnline(peersID []messages.PeerID) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	for _, peerID := range peersID {
+		waitCh, ok := s.waitingPeers[peerID]
+		if !ok {
+			// If meanwhile the peer was unsubscribed, we don't need to signal it
+			continue
+		}
+
+		waitCh <- struct{}{}
+		delete(s.waitingPeers, peerID)
+		close(waitCh)
+	}
+}
+
+func (s *PeersStateSubscription) OnPeersWentOffline(peersID []messages.PeerID) {
+	s.mu.Lock()
+	relevantPeers := make([]messages.PeerID, 0, len(peersID))
+	for _, peerID := range peersID {
+		if _, ok := s.listenForOfflinePeers[peerID]; ok {
+			relevantPeers = append(relevantPeers, peerID)
+		}
+	}
+	s.mu.Unlock()
+
+	if len(relevantPeers) > 0 {
+		s.offlineCallback(relevantPeers)
+	}
+}
+
+// WaitToBeOnlineAndSubscribe waits for a specific peer to come online and subscribes to its state changes.
+func (s *PeersStateSubscription) WaitToBeOnlineAndSubscribe(ctx context.Context, peerID messages.PeerID) error {
+	// Check if already waiting for this peer
+	s.mu.Lock()
+	if _, exists := s.waitingPeers[peerID]; exists {
+		s.mu.Unlock()
+		return errors.New("already waiting for peer to come online")
+	}
+
+	// Create a channel to wait for the peer to come online
+	waitCh := make(chan struct{}, 1)
+	s.waitingPeers[peerID] = waitCh
+	s.listenForOfflinePeers[peerID] = struct{}{}
+	s.mu.Unlock()
+
+	if err := s.subscribeStateChange(peerID); err != nil {
+		s.log.Errorf("failed to subscribe to peer state: %s", err)
+		s.mu.Lock()
+		if ch, exists := s.waitingPeers[peerID]; exists && ch == waitCh {
+			close(waitCh)
+			delete(s.waitingPeers, peerID)
+			delete(s.listenForOfflinePeers, peerID)
+		}
+		s.mu.Unlock()
+		return err
+	}
+
+	// Wait for peer to come online or context to be cancelled
+	timeoutCtx, cancel := context.WithTimeout(ctx, OpenConnectionTimeout)
+	defer cancel()
+	select {
+	case _, ok := <-waitCh:
+		if !ok {
+			return fmt.Errorf("wait for peer to come online has been cancelled")
+		}
+
+		s.log.Debugf("peer %s is now online", peerID)
+		return nil
+	case <-timeoutCtx.Done():
+		s.log.Debugf("context timed out while waiting for peer %s to come online", peerID)
+		if err := s.unsubscribeStateChange([]messages.PeerID{peerID}); err != nil {
+			s.log.Errorf("failed to unsubscribe from peer state: %s", err)
+		}
+		s.mu.Lock()
+		if ch, exists := s.waitingPeers[peerID]; exists && ch == waitCh {
+			close(waitCh)
+			delete(s.waitingPeers, peerID)
+			delete(s.listenForOfflinePeers, peerID)
+		}
+		s.mu.Unlock()
+		return timeoutCtx.Err()
+	}
+}
+
+func (s *PeersStateSubscription) UnsubscribeStateChange(peerIDs []messages.PeerID) error {
+	msgErr := s.unsubscribeStateChange(peerIDs)
+
+	s.mu.Lock()
+	for _, peerID := range peerIDs {
+		if wch, ok := s.waitingPeers[peerID]; ok {
+			close(wch)
+			delete(s.waitingPeers, peerID)
+		}
+
+		delete(s.listenForOfflinePeers, peerID)
+	}
+	s.mu.Unlock()
+
+	return msgErr
+}
+
+func (s *PeersStateSubscription) Cleanup() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	for _, waitCh := range s.waitingPeers {
+		close(waitCh)
+	}
+
+	s.waitingPeers = make(map[messages.PeerID]chan struct{})
+	s.listenForOfflinePeers = make(map[messages.PeerID]struct{})
+}
+
+func (s *PeersStateSubscription) subscribeStateChange(peerID messages.PeerID) error {
+	msgs, err := messages.MarshalSubPeerStateMsg([]messages.PeerID{peerID})
+	if err != nil {
+		return err
+	}
+
+	for _, msg := range msgs {
+		if _, err := s.relayConn.Write(msg); err != nil {
+			return err
+		}
+
+	}
+	return nil
+}
+
+func (s *PeersStateSubscription) unsubscribeStateChange(peerIDs []messages.PeerID) error {
+	msgs, err := messages.MarshalUnsubPeerStateMsg(peerIDs)
+	if err != nil {
+		return err
+	}
+
+	var connWriteErr error
+	for _, msg := range msgs {
+		if _, err := s.relayConn.Write(msg); err != nil {
+			connWriteErr = err
+		}
+	}
+	return connWriteErr
+}

shared/relay/client/peer_subscription_test.go

@@ -0,0 +1,99 @@
+package client
+
+import (
+	"bytes"
+	"context"
+	"testing"
+	"time"
+
+	"github.com/netbirdio/netbird/relay/messages"
+
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type mockRelayedConn struct {
+}
+
+func (m *mockRelayedConn) Write(p []byte) (n int, err error) {
+	return len(p), nil
+}
+
+func TestWaitToBeOnlineAndSubscribe_Success(t *testing.T) {
+	peerID := messages.HashID("peer1")
+	mockConn := &mockRelayedConn{}
+	logger := logrus.New()
+	logger.SetOutput(&bytes.Buffer{}) // discard log output
+	sub := NewPeersStateSubscription(logrus.NewEntry(logger), mockConn, nil)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// Launch wait in background
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		sub.OnPeersOnline([]messages.PeerID{peerID})
+	}()
+
+	err := sub.WaitToBeOnlineAndSubscribe(ctx, peerID)
+	assert.NoError(t, err)
+}
+
+func TestWaitToBeOnlineAndSubscribe_Timeout(t *testing.T) {
+	peerID := messages.HashID("peer2")
+	mockConn := &mockRelayedConn{}
+	logger := logrus.New()
+	logger.SetOutput(&bytes.Buffer{})
+	sub := NewPeersStateSubscription(logrus.NewEntry(logger), mockConn, nil)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+
+	err := sub.WaitToBeOnlineAndSubscribe(ctx, peerID)
+	assert.Error(t, err)
+	assert.Equal(t, context.DeadlineExceeded, err)
+}
+
+func TestWaitToBeOnlineAndSubscribe_Duplicate(t *testing.T) {
+	peerID := messages.HashID("peer3")
+	mockConn := &mockRelayedConn{}
+	logger := logrus.New()
+	logger.SetOutput(&bytes.Buffer{})
+	sub := NewPeersStateSubscription(logrus.NewEntry(logger), mockConn, nil)
+
+	ctx := context.Background()
+	go func() {
+		_ = sub.WaitToBeOnlineAndSubscribe(ctx, peerID)
+
+	}()
+	time.Sleep(100 * time.Millisecond)
+	err := sub.WaitToBeOnlineAndSubscribe(ctx, peerID)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "already waiting")
+}
+
+func TestUnsubscribeStateChange(t *testing.T) {
+	peerID := messages.HashID("peer4")
+	mockConn := &mockRelayedConn{}
+	logger := logrus.New()
+	logger.SetOutput(&bytes.Buffer{})
+	sub := NewPeersStateSubscription(logrus.NewEntry(logger), mockConn, nil)
+
+	doneChan := make(chan struct{})
+	go func() {
+		_ = sub.WaitToBeOnlineAndSubscribe(context.Background(), peerID)
+		close(doneChan)
+	}()
+	time.Sleep(100 * time.Millisecond)
+
+	err := sub.UnsubscribeStateChange([]messages.PeerID{peerID})
+	assert.NoError(t, err)
+
+	select {
+	case <-doneChan:
+	case <-time.After(200 * time.Millisecond):
+		// Expected timeout, meaning the subscription was successfully unsubscribed
+		t.Errorf("timeout")
+	}
+}

shared/relay/client/picker.go

@@ -0,0 +1,104 @@
+package client
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync/atomic"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	auth "github.com/netbirdio/netbird/shared/relay/auth/hmac"
+)
+
+const (
+	maxConcurrentServers = 7
+)
+
+var (
+	connectionTimeout = 30 * time.Second
+)
+
+type connResult struct {
+	RelayClient *Client
+	Url         string
+	Err         error
+}
+
+type ServerPicker struct {
+	TokenStore *auth.TokenStore
+	ServerURLs atomic.Value
+	PeerID     string
+}
+
+func (sp *ServerPicker) PickServer(parentCtx context.Context) (*Client, error) {
+	ctx, cancel := context.WithTimeout(parentCtx, connectionTimeout)
+	defer cancel()
+
+	totalServers := len(sp.ServerURLs.Load().([]string))
+
+	connResultChan := make(chan connResult, totalServers)
+	successChan := make(chan connResult, 1)
+	concurrentLimiter := make(chan struct{}, maxConcurrentServers)
+
+	log.Debugf("pick server from list: %v", sp.ServerURLs.Load().([]string))
+	for _, url := range sp.ServerURLs.Load().([]string) {
+		// todo check if we have a successful connection so we do not need to connect to other servers
+		concurrentLimiter <- struct{}{}
+		go func(url string) {
+			defer func() {
+				<-concurrentLimiter
+			}()
+			sp.startConnection(parentCtx, connResultChan, url)
+		}(url)
+	}
+
+	go sp.processConnResults(connResultChan, successChan)
+
+	select {
+	case cr, ok := <-successChan:
+		if !ok {
+			return nil, errors.New("failed to connect to any relay server: all attempts failed")
+		}
+		log.Infof("chosen home Relay server: %s", cr.Url)
+		return cr.RelayClient, nil
+	case <-ctx.Done():
+		return nil, fmt.Errorf("failed to connect to any relay server: %w", ctx.Err())
+	}
+}
+
+func (sp *ServerPicker) startConnection(ctx context.Context, resultChan chan connResult, url string) {
+	log.Infof("try to connecting to relay server: %s", url)
+	relayClient := NewClient(url, sp.TokenStore, sp.PeerID)
+	err := relayClient.Connect(ctx)
+	resultChan <- connResult{
+		RelayClient: relayClient,
+		Url:         url,
+		Err:         err,
+	}
+}
+
+func (sp *ServerPicker) processConnResults(resultChan chan connResult, successChan chan connResult) {
+	var hasSuccess bool
+	for numOfResults := 0; numOfResults < cap(resultChan); numOfResults++ {
+		cr := <-resultChan
+		if cr.Err != nil {
+			log.Tracef("failed to connect to Relay server: %s: %v", cr.Url, cr.Err)
+			continue
+		}
+		log.Infof("connected to Relay server: %s", cr.Url)
+
+		if hasSuccess {
+			log.Infof("closing unnecessary Relay connection to: %s", cr.Url)
+			if err := cr.RelayClient.Close(); err != nil {
+				log.Errorf("failed to close connection to %s: %v", cr.Url, err)
+			}
+			continue
+		}
+
+		hasSuccess = true
+		successChan <- cr
+	}
+	close(successChan)
+}

shared/relay/client/picker_test.go

@@ -0,0 +1,34 @@
+package client
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+)
+
+func TestServerPicker_UnavailableServers(t *testing.T) {
+	connectionTimeout = 5 * time.Second
+
+	sp := ServerPicker{
+		TokenStore: nil,
+		PeerID:     "test",
+	}
+	sp.ServerURLs.Store([]string{"rel://dummy1", "rel://dummy2"})
+
+	ctx, cancel := context.WithTimeout(context.Background(), connectionTimeout+1)
+	defer cancel()
+
+	go func() {
+		_, err := sp.PickServer(ctx)
+		if err == nil {
+			t.Error(err)
+		}
+		cancel()
+	}()
+
+	<-ctx.Done()
+	if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+		t.Errorf("PickServer() took too long to complete")
+	}
+}