Add comment

- eliminate reader function generation
- fix write to closed channel panic

client/android/client.go

@@ -57,15 +57,17 @@ type Client struct {
 	ctxCancel             context.CancelFunc
 	ctxCancelLock         *sync.Mutex
 	deviceName            string
+	uiVersion             string
 	networkChangeListener listener.NetworkChangeListener
 }
 
 // NewClient instantiate a new Client
-func NewClient(cfgFile, deviceName string, tunAdapter TunAdapter, iFaceDiscover IFaceDiscover, networkChangeListener NetworkChangeListener) *Client {
+func NewClient(cfgFile, deviceName string, uiVersion string, tunAdapter TunAdapter, iFaceDiscover IFaceDiscover, networkChangeListener NetworkChangeListener) *Client {
 	net.SetAndroidProtectSocketFn(tunAdapter.ProtectSocket)
 	return &Client{
 		cfgFile:               cfgFile,
 		deviceName:            deviceName,
+		uiVersion:             uiVersion,
 		tunAdapter:            tunAdapter,
 		iFaceDiscover:         iFaceDiscover,
 		recorder:              peer.NewRecorder(""),
@@ -88,6 +90,9 @@ func (c *Client) Run(urlOpener URLOpener, dns *DNSList, dnsReadyListener DnsRead
 	var ctx context.Context
 	//nolint
 	ctxWithValues := context.WithValue(context.Background(), system.DeviceNameCtxKey, c.deviceName)
+	//nolint
+	ctxWithValues = context.WithValue(ctxWithValues, system.UiVersionCtxKey, c.uiVersion)
+
 	c.ctxCancelLock.Lock()
 	ctx, c.ctxCancel = context.WithCancel(ctxWithValues)
 	defer c.ctxCancel()

client/cmd/debug.go

@@ -59,7 +59,7 @@ var forCmd = &cobra.Command{
 }
 
 func debugBundle(cmd *cobra.Command, _ []string) error {
-	conn, err := getClient(cmd.Context())
+	conn, err := getClient(cmd)
 	if err != nil {
 		return err
 	}
@@ -80,7 +80,7 @@ func debugBundle(cmd *cobra.Command, _ []string) error {
 }
 
 func setLogLevel(cmd *cobra.Command, args []string) error {
-	conn, err := getClient(cmd.Context())
+	conn, err := getClient(cmd)
 	if err != nil {
 		return err
 	}
@@ -109,7 +109,7 @@ func runForDuration(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("invalid duration format: %v", err)
 	}
 
-	conn, err := getClient(cmd.Context())
+	conn, err := getClient(cmd)
 	if err != nil {
 		return err
 	}

client/cmd/root.go

@@ -353,8 +353,11 @@ func migrateToNetbird(oldPath, newPath string) bool {
 	return true
 }
 
-func getClient(ctx context.Context) (*grpc.ClientConn, error) {
-	conn, err := DialClientGRPCServer(ctx, daemonAddr)
+func getClient(cmd *cobra.Command) (*grpc.ClientConn, error) {
+	SetFlagsFromEnvVars(rootCmd)
+	cmd.SetOut(cmd.OutOrStdout())
+
+	conn, err := DialClientGRPCServer(cmd.Context(), daemonAddr)
 	if err != nil {
 		return nil, fmt.Errorf("failed to connect to daemon error: %v\n"+
 			"If the daemon is not running please run: "+

client/cmd/route.go

@@ -49,7 +49,7 @@ func init() {
 }
 
 func routesList(cmd *cobra.Command, _ []string) error {
-	conn, err := getClient(cmd.Context())
+	conn, err := getClient(cmd)
 	if err != nil {
 		return err
 	}
@@ -79,7 +79,7 @@ func routesList(cmd *cobra.Command, _ []string) error {
 }
 
 func routesSelect(cmd *cobra.Command, args []string) error {
-	conn, err := getClient(cmd.Context())
+	conn, err := getClient(cmd)
 	if err != nil {
 		return err
 	}
@@ -106,7 +106,7 @@ func routesSelect(cmd *cobra.Command, args []string) error {
 }
 
 func routesDeselect(cmd *cobra.Command, args []string) error {
-	conn, err := getClient(cmd.Context())
+	conn, err := getClient(cmd)
 	if err != nil {
 		return err
 	}

client/internal/engine.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"maps"
 	"math/rand"
 	"net"
 	"net/netip"
@@ -117,7 +118,8 @@ type Engine struct {
 	TURNs []*stun.URI
 
 	// clientRoutes is the most recent list of clientRoutes received from the Management Service
-	clientRoutes route.HAMap
+	clientRoutes   route.HAMap
+	clientRoutesMu sync.RWMutex
 
 	clientCtx    context.Context
 	clientCancel context.CancelFunc
@@ -240,7 +242,9 @@ func (e *Engine) Stop() error {
 		return err
 	}
 
+	e.clientRoutesMu.Lock()
 	e.clientRoutes = nil
+	e.clientRoutesMu.Unlock()
 
 	// very ugly but we want to remove peers from the WireGuard interface first before removing interface.
 	// Removing peers happens in the conn.Close() asynchronously
@@ -738,7 +742,9 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
 		log.Errorf("failed to update clientRoutes, err: %v", err)
 	}
 
+	e.clientRoutesMu.Lock()
 	e.clientRoutes = clientRoutes
+	e.clientRoutesMu.Unlock()
 
 	protoDNSConfig := networkMap.GetDNSConfig()
 	if protoDNSConfig == nil {
@@ -977,7 +983,6 @@ func (e *Engine) createPeerConn(pubKey string, allowedIPs string) (*peer.Conn, e
 		WgConfig:             wgConfig,
 		LocalWgPort:          e.config.WgPort,
 		NATExternalIPs:       e.parseNATExternalIPMappings(),
-		UserspaceBind:        e.wgInterface.IsUserspaceBind(),
 		RosenpassPubKey:      e.getRosenpassPubKey(),
 		RosenpassAddr:        e.getRosenpassAddr(),
 	}
@@ -1040,8 +1045,6 @@ func (e *Engine) receiveSignalEvents() {
 					return err
 				}
 
-				conn.RegisterProtoSupportMeta(msg.Body.GetFeaturesSupported())
-
 				var rosenpassPubKey []byte
 				rosenpassAddr := ""
 				if msg.GetBody().GetRosenpassConfig() != nil {
@@ -1064,8 +1067,6 @@ func (e *Engine) receiveSignalEvents() {
 					return err
 				}
 
-				conn.RegisterProtoSupportMeta(msg.GetBody().GetFeaturesSupported())
-
 				var rosenpassPubKey []byte
 				rosenpassAddr := ""
 				if msg.GetBody().GetRosenpassConfig() != nil {
@@ -1088,7 +1089,8 @@ func (e *Engine) receiveSignalEvents() {
 					log.Errorf("failed on parsing remote candidate %s -> %s", candidate, err)
 					return err
 				}
-				conn.OnRemoteCandidate(candidate)
+
+				conn.OnRemoteCandidate(candidate, e.GetClientRoutes())
 			case sProto.Body_MODE:
 			}
 
@@ -1282,11 +1284,17 @@ func (e *Engine) newDnsServer() ([]*route.Route, dns.Server, error) {
 
 // GetClientRoutes returns the current routes from the route map
 func (e *Engine) GetClientRoutes() route.HAMap {
-	return e.clientRoutes
+	e.clientRoutesMu.RLock()
+	defer e.clientRoutesMu.RUnlock()
+
+	return maps.Clone(e.clientRoutes)
 }
 
 // GetClientRoutesWithNetID returns the current routes from the route map, but the keys consist of the network ID only
 func (e *Engine) GetClientRoutesWithNetID() map[route.NetID][]*route.Route {
+	e.clientRoutesMu.RLock()
+	defer e.clientRoutesMu.RUnlock()
+
 	routes := make(map[route.NetID][]*route.Route, len(e.clientRoutes))
 	for id, v := range e.clientRoutes {
 		routes[id.NetID()] = v

client/internal/peer/conn.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"net/netip"
 	"runtime"
 	"strings"
 	"sync"
@@ -18,7 +19,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/wgproxy"
 	"github.com/netbirdio/netbird/iface"
 	"github.com/netbirdio/netbird/iface/bind"
-	signal "github.com/netbirdio/netbird/signal/client"
+	"github.com/netbirdio/netbird/route"
 	sProto "github.com/netbirdio/netbird/signal/proto"
 	nbnet "github.com/netbirdio/netbird/util/net"
 	"github.com/netbirdio/netbird/version"
@@ -68,9 +69,6 @@ type ConnConfig struct {
 
 	NATExternalIPs []string
 
-	// UsesBind indicates whether the WireGuard interface is userspace and uses bind.ICEBind
-	UserspaceBind bool
-
 	// RosenpassPubKey is this peer's Rosenpass public key
 	RosenpassPubKey []byte
 	// RosenpassPubKey is this peer's RosenpassAddr server address (IP:port)
@@ -133,32 +131,15 @@ type Conn struct {
 	wgProxyFactory *wgproxy.Factory
 	wgProxy        wgproxy.Proxy
 
-	remoteModeCh chan ModeMessage
-	meta         meta
-
 	adapter        iface.TunAdapter
 	iFaceDiscover  stdnet.ExternalIFaceDiscover
 	sentExtraSrflx bool
 
-	remoteEndpoint *net.UDPAddr
-	remoteConn     *ice.Conn
-
 	connID               nbnet.ConnectionID
 	beforeAddPeerHooks   []BeforeAddPeerHookFunc
 	afterRemovePeerHooks []AfterRemovePeerHookFunc
 }
 
-// meta holds meta information about a connection
-type meta struct {
-	protoSupport signal.FeaturesSupport
-}
-
-// ModeMessage represents a connection mode chosen by the peer
-type ModeMessage struct {
-	// Direct indicates that it decided to use a direct connection
-	Direct bool
-}
-
 // GetConf returns the connection config
 func (conn *Conn) GetConf() ConnConfig {
 	return conn.config
@@ -185,7 +166,6 @@ func NewConn(config ConnConfig, statusRecorder *Status, wgProxyFactory *wgproxy.
 		remoteOffersCh: make(chan OfferAnswer),
 		remoteAnswerCh: make(chan OfferAnswer),
 		statusRecorder: statusRecorder,
-		remoteModeCh:   make(chan ModeMessage, 1),
 		wgProxyFactory: wgProxyFactory,
 		adapter:        adapter,
 		iFaceDiscover:  iFaceDiscover,
@@ -353,7 +333,7 @@ func (conn *Conn) Open(ctx context.Context) error {
 
 	err = conn.agent.GatherCandidates()
 	if err != nil {
-		return err
+		return fmt.Errorf("gather candidates: %v", err)
 	}
 
 	// will block until connection succeeded
@@ -370,14 +350,12 @@ func (conn *Conn) Open(ctx context.Context) error {
 		return err
 	}
 
-	// dynamically set remote WireGuard port is other side specified a different one from the default one
+	// dynamically set remote WireGuard port if other side specified a different one from the default one
 	remoteWgPort := iface.DefaultWgPort
 	if remoteOfferAnswer.WgListenPort != 0 {
 		remoteWgPort = remoteOfferAnswer.WgListenPort
 	}
 
-	conn.remoteConn = remoteConn
-
 	// the ice connection has been established successfully so we are ready to start the proxy
 	remoteAddr, err := conn.configureConnection(remoteConn, remoteWgPort, remoteOfferAnswer.RosenpassPubKey,
 		remoteOfferAnswer.RosenpassAddr)
@@ -435,7 +413,6 @@ func (conn *Conn) configureConnection(remoteConn net.Conn, remoteWgPort int, rem
 	}
 
 	endpointUdpAddr, _ := net.ResolveUDPAddr(endpoint.Network(), endpoint.String())
-	conn.remoteEndpoint = endpointUdpAddr
 	log.Debugf("Conn resolved IP for %s: %s", endpoint, endpointUdpAddr.IP)
 
 	conn.connID = nbnet.GenerateConnID()
@@ -621,40 +598,39 @@ func (conn *Conn) SetSendSignalMessage(handler func(message *sProto.Message) err
 // onICECandidate is a callback attached to an ICE Agent to receive new local connection candidates
 // and then signals them to the remote peer
 func (conn *Conn) onICECandidate(candidate ice.Candidate) {
-	if candidate != nil {
-		// TODO: reported port is incorrect for CandidateTypeHost, makes understanding ICE use via logs confusing as port is ignored
-		log.Debugf("discovered local candidate %s", candidate.String())
-		go func() {
-			err := conn.signalCandidate(candidate)
-			if err != nil {
-				log.Errorf("failed signaling candidate to the remote peer %s %s", conn.config.Key, err)
-			}
-
-			// sends an extra server reflexive candidate to the remote peer with our related port (usually the wireguard port)
-			// this is useful when network has an existing port forwarding rule for the wireguard port and this peer
-			if !conn.sentExtraSrflx && candidate.Type() == ice.CandidateTypeServerReflexive && candidate.Port() != candidate.RelatedAddress().Port {
-				relatedAdd := candidate.RelatedAddress()
-				extraSrflx, err := ice.NewCandidateServerReflexive(&ice.CandidateServerReflexiveConfig{
-					Network:   candidate.NetworkType().String(),
-					Address:   candidate.Address(),
-					Port:      relatedAdd.Port,
-					Component: candidate.Component(),
-					RelAddr:   relatedAdd.Address,
-					RelPort:   relatedAdd.Port,
-				})
-				if err != nil {
-					log.Errorf("failed creating extra server reflexive candidate %s", err)
-					return
-				}
-				err = conn.signalCandidate(extraSrflx)
-				if err != nil {
-					log.Errorf("failed signaling the extra server reflexive candidate to the remote peer %s: %s", conn.config.Key, err)
-					return
-				}
-				conn.sentExtraSrflx = true
-			}
-		}()
+	// nil means candidate gathering has been ended
+	if candidate == nil {
+		return
 	}
+
+	// TODO: reported port is incorrect for CandidateTypeHost, makes understanding ICE use via logs confusing as port is ignored
+	log.Debugf("discovered local candidate %s", candidate.String())
+	go func() {
+		err := conn.signalCandidate(candidate)
+		if err != nil {
+			log.Errorf("failed signaling candidate to the remote peer %s %s", conn.config.Key, err)
+		}
+	}()
+
+	if !conn.shouldSendExtraSrflxCandidate(candidate) {
+		return
+	}
+
+	// sends an extra server reflexive candidate to the remote peer with our related port (usually the wireguard port)
+	// this is useful when network has an existing port forwarding rule for the wireguard port and this peer
+	extraSrflx, err := extraSrflxCandidate(candidate)
+	if err != nil {
+		log.Errorf("failed creating extra server reflexive candidate %s", err)
+		return
+	}
+	conn.sentExtraSrflx = true
+
+	go func() {
+		err = conn.signalCandidate(extraSrflx)
+		if err != nil {
+			log.Errorf("failed signaling the extra server reflexive candidate to the remote peer %s: %s", conn.config.Key, err)
+		}
+	}()
 }
 
 func (conn *Conn) onICESelectedCandidatePair(c1 ice.Candidate, c2 ice.Candidate) {
@@ -779,7 +755,7 @@ func (conn *Conn) OnRemoteAnswer(answer OfferAnswer) bool {
 }
 
 // OnRemoteCandidate Handles ICE connection Candidate provided by the remote peer.
-func (conn *Conn) OnRemoteCandidate(candidate ice.Candidate) {
+func (conn *Conn) OnRemoteCandidate(candidate ice.Candidate, haRoutes route.HAMap) {
 	log.Debugf("OnRemoteCandidate from peer %s -> %s", conn.config.Key, candidate.String())
 	go func() {
 		conn.mu.Lock()
@@ -789,6 +765,10 @@ func (conn *Conn) OnRemoteCandidate(candidate ice.Candidate) {
 			return
 		}
 
+		if candidateViaRoutes(candidate, haRoutes) {
+			return
+		}
+
 		err := conn.agent.AddRemoteCandidate(candidate)
 		if err != nil {
 			log.Errorf("error while handling remote candidate from peer %s", conn.config.Key)
@@ -801,8 +781,49 @@ func (conn *Conn) GetKey() string {
 	return conn.config.Key
 }
 
-// RegisterProtoSupportMeta register supported proto message in the connection metadata
-func (conn *Conn) RegisterProtoSupportMeta(support []uint32) {
-	protoSupport := signal.ParseFeaturesSupported(support)
-	conn.meta.protoSupport = protoSupport
+func (conn *Conn) shouldSendExtraSrflxCandidate(candidate ice.Candidate) bool {
+	if !conn.sentExtraSrflx && candidate.Type() == ice.CandidateTypeServerReflexive && candidate.Port() != candidate.RelatedAddress().Port {
+		return true
+	}
+	return false
+}
+
+func extraSrflxCandidate(candidate ice.Candidate) (*ice.CandidateServerReflexive, error) {
+	relatedAdd := candidate.RelatedAddress()
+	return ice.NewCandidateServerReflexive(&ice.CandidateServerReflexiveConfig{
+		Network:   candidate.NetworkType().String(),
+		Address:   candidate.Address(),
+		Port:      relatedAdd.Port,
+		Component: candidate.Component(),
+		RelAddr:   relatedAdd.Address,
+		RelPort:   relatedAdd.Port,
+	})
+}
+
+func candidateViaRoutes(candidate ice.Candidate, clientRoutes route.HAMap) bool {
+	var routePrefixes []netip.Prefix
+	for _, routes := range clientRoutes {
+		if len(routes) > 0 && routes[0] != nil {
+			routePrefixes = append(routePrefixes, routes[0].Network)
+		}
+	}
+
+	addr, err := netip.ParseAddr(candidate.Address())
+	if err != nil {
+		log.Errorf("Failed to parse IP address %s: %v", candidate.Address(), err)
+		return false
+	}
+
+	for _, prefix := range routePrefixes {
+		// default route is
+		if prefix.Bits() == 0 {
+			continue
+		}
+
+		if prefix.Contains(addr) {
+			log.Debugf("Ignoring candidate [%s], its address is part of routed network %s", candidate.String(), prefix)
+			return true
+		}
+	}
+	return false
 }

client/internal/templates/pkce-auth-msg.html

@@ -1,7 +1,7 @@
 <!DOCTYPE html>
-<html>
+<html lang="en">
 <head>
-    <meta name="viewport" content="width=device-width, initial-scale=1"/>
+    <meta content="width=device-width, initial-scale=1" name="viewport"/>
     <style>
         body {
             display: flex;
@@ -50,16 +50,17 @@
             color: black;
         }
     </style>
+    <title>NetBird Login Successful</title>
 </head>
 <body>
 <div class="container">
     <div class="logo">
-        <img src="https://img.mailinblue.com/6211297/images/content_library/original/64bd4ce82e1ea753e439b6a2.png">
+        <img alt="netbird_logo" src="https://img.mailinblue.com/6211297/images/content_library/original/64bd4ce82e1ea753e439b6a2.png">
     </div>
     <br>
     {{ if .Error }}
-    <svg xmlns="http://www.w3.org/2000/svg" height="50" viewBox="0 0 100 100">
-        <circle cx="50" cy="50" r="45" fill="none" stroke="red" stroke-width="3"/>
+    <svg height="50" viewBox="0 0 100 100" xmlns="http://www.w3.org/2000/svg">
+        <circle cx="50" cy="50" fill="none" r="45" stroke="red" stroke-width="3"/>
         <path d="M30 30 L70 70 M30 70 L70 30" fill="none" stroke="red" stroke-width="3"/>
     </svg>
     <div class="content">
@@ -69,8 +70,8 @@
         {{ .Error }}.
     </div>
     {{ else }}
-    <svg xmlns="http://www.w3.org/2000/svg" height="50" viewBox="0 0 100 100">
-        <circle cx="50" cy="50" r="45" fill="none" stroke="#5cb85c" stroke-width="3"/>
+    <svg height="50" viewBox="0 0 100 100" xmlns="http://www.w3.org/2000/svg">
+        <circle cx="50" cy="50" fill="none" r="45" stroke="#5cb85c" stroke-width="3"/>
         <path d="M30 50 L45 65 L70 35" fill="none" stroke="#5cb85c" stroke-width="5"/>
     </svg>
     <div class="content">

client/system/info.go

@@ -20,6 +20,9 @@ const OsVersionCtxKey = "OsVersion"
 // OsNameCtxKey context key for operating system name
 const OsNameCtxKey = "OsName"
 
+// UiVersionCtxKey context key for user UI version
+const UiVersionCtxKey = "user-agent"
+
 type NetworkAddress struct {
 	NetIP netip.Prefix
 	Mac   string

client/system/info_android.go

@@ -28,10 +28,18 @@ func GetInfo(ctx context.Context) *Info {
 		kernelVersion = osInfo[2]
 	}
 
-	gio := &Info{Kernel: kernel, Platform: "unknown", OS: "android", OSVersion: osVersion(), GoOS: runtime.GOOS, CPUs: runtime.NumCPU(), KernelVersion: kernelVersion}
-	gio.Hostname = extractDeviceName(ctx, "android")
-	gio.WiretrusteeVersion = version.NetbirdVersion()
-	gio.UIVersion = extractUserAgent(ctx)
+	gio := &Info{
+		GoOS:               runtime.GOOS,
+		Kernel:             kernel,
+		Platform:           "unknown",
+		OS:                 "android",
+		OSVersion:          osVersion(),
+		Hostname:           extractDeviceName(ctx, "android"),
+		CPUs:               runtime.NumCPU(),
+		WiretrusteeVersion: version.NetbirdVersion(),
+		UIVersion:          extractUIVersion(ctx),
+		KernelVersion:      kernelVersion,
+	}
 
 	return gio
 }
@@ -45,6 +53,14 @@ func osVersion() string {
 	return run("/system/bin/getprop", "ro.build.version.release")
 }
 
+func extractUIVersion(ctx context.Context) string {
+	v, ok := ctx.Value(UiVersionCtxKey).(string)
+	if !ok {
+		return ""
+	}
+	return v
+}
+
 func run(name string, arg ...string) string {
 	cmd := exec.Command(name, arg...)
 	cmd.Stdin = strings.NewReader("some")

go.mod

@@ -12,7 +12,7 @@ require (
 	github.com/gorilla/mux v1.8.0
 	github.com/kardianos/service v1.2.1-0.20210728001519-a323c3813bc7
 	github.com/onsi/ginkgo v1.16.5
-	github.com/onsi/gomega v1.23.0
+	github.com/onsi/gomega v1.27.6
 	github.com/pion/ice/v3 v3.0.2
 	github.com/rs/cors v1.8.0
 	github.com/sirupsen/logrus v1.9.3
@@ -47,6 +47,7 @@ require (
 	github.com/google/martian/v3 v3.0.0
 	github.com/google/nftables v0.0.0-20220808154552-2eca00135732
 	github.com/gopacket/gopacket v1.1.1
+	github.com/gorilla/websocket v1.5.1
 	github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.0.2-0.20240212192251-757544f21357
 	github.com/hashicorp/go-multierror v1.1.1
 	github.com/hashicorp/go-secure-stdlib/base62 v0.1.2
@@ -67,6 +68,7 @@ require (
 	github.com/pion/transport/v3 v3.0.1
 	github.com/pion/turn/v3 v3.0.1
 	github.com/prometheus/client_golang v1.19.1
+	github.com/quic-go/quic-go v0.45.0
 	github.com/rs/xid v1.3.0
 	github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966
 	github.com/stretchr/testify v1.9.0
@@ -91,6 +93,7 @@ require (
 	gorm.io/driver/postgres v1.5.7
 	gorm.io/driver/sqlite v1.5.3
 	gorm.io/gorm v1.25.7-0.20240204074919-46816ad31dde
+	nhooyr.io/websocket v1.8.11
 )
 
 require (
@@ -132,10 +135,12 @@ require (
 	github.com/go-ole/go-ole v1.3.0 // indirect
 	github.com/go-redis/redis/v8 v8.11.5 // indirect
 	github.com/go-stack/stack v1.8.0 // indirect
+	github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/goki/freetype v0.0.0-20181231101311-fa8a33aabaff // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
 	github.com/google/btree v1.0.1 // indirect
+	github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38 // indirect
 	github.com/google/s2a-go v0.1.7 // indirect
 	github.com/googleapis/enterprise-certificate-proxy v0.3.2 // indirect
 	github.com/googleapis/gax-go/v2 v2.12.3 // indirect
@@ -161,6 +166,7 @@ require (
 	github.com/moby/term v0.5.0 // indirect
 	github.com/morikuni/aec v1.0.0 // indirect
 	github.com/nxadm/tail v1.4.8 // indirect
+	github.com/onsi/ginkgo/v2 v2.9.5 // indirect
 	github.com/opencontainers/go-digest v1.0.0 // indirect
 	github.com/opencontainers/image-spec v1.1.0 // indirect
 	github.com/oxtoacart/bpool v0.0.0-20190530202638-03653db5a59c // indirect
@@ -188,9 +194,12 @@ require (
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.51.0 // indirect
 	go.opentelemetry.io/otel/sdk v1.26.0 // indirect
 	go.opentelemetry.io/otel/trace v1.26.0 // indirect
+	go.uber.org/mock v0.4.0 // indirect
 	golang.org/x/image v0.10.0 // indirect
+	golang.org/x/mod v0.17.0 // indirect
 	golang.org/x/text v0.15.0 // indirect
 	golang.org/x/time v0.5.0 // indirect
+	golang.org/x/tools v0.21.0 // indirect
 	golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20240509183442-62759503f434 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240515191416-fc5f0ca64291 // indirect

go.sum

@@ -50,6 +50,9 @@ github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA
 github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI=
+github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI=
+github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU=
 github.com/cilium/ebpf v0.15.0 h1:7NxJhNiBT3NG8pZJ3c+yfrVdHY8ScgKD27sScgjLMMk=
 github.com/cilium/ebpf v0.15.0/go.mod h1:DHp1WyrLeiBh19Cf/tfiSMhqheEiK8fXFZ4No0P1Hso=
 github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
@@ -154,6 +157,8 @@ github.com/go-redis/redis/v8 v8.11.5/go.mod h1:gREzHqY1hg6oD9ngVRbLStwAWKhA0FEgq
 github.com/go-stack/stack v1.8.0 h1:5SgMzNM5HxrEjV0ww2lTmX6E2Izsfxas4+YHWRs3Lsk=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/go-task/slim-sprig v0.0.0-20210107165309-348f09dbbbc0/go.mod h1:fyg7847qk6SyHyPtNmDHnmrv/HOrqktSC+C9fM+CJOE=
+github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEevZMzYi5KSi8KkcZtzBcTgAUUtapy0OI=
+github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4B2jHnOSGXyyzV8ROjYa2ojvAY6HCGYYfMoC3Ls=
 github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/godbus/dbus/v5 v5.1.0 h1:4KLkAxT3aOY8Li4FRJe/KvhoNFFxo0m6fNuFUO8QJUk=
 github.com/godbus/dbus/v5 v5.1.0/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
@@ -213,6 +218,8 @@ github.com/google/martian/v3 v3.0.0 h1:pMen7vLs8nvgEYhywH3KDWJIJTeEr2ULsVWHWYHQy
 github.com/google/martian/v3 v3.0.0/go.mod h1:y5Zk1BBys9G+gd6Jrk0W3cC1+ELVxBWuIGO+w/tUAp0=
 github.com/google/nftables v0.0.0-20220808154552-2eca00135732 h1:csc7dT82JiSLvq4aMyQMIQDL7986NH6Wxf/QrvOj55A=
 github.com/google/nftables v0.0.0-20220808154552-2eca00135732/go.mod h1:b97ulCCFipUC+kSin+zygkvUVpx0vyIAwxXFdY3PlNc=
+github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38 h1:yAJXTCF9TqKcTiHJAE8dj7HMvPfh66eeA2JYW7eFpSE=
+github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
 github.com/google/s2a-go v0.1.7 h1:60BLSyTrOV4/haCDW4zb1guZItoSq8foHCXrAnjBo/o=
 github.com/google/s2a-go v0.1.7/go.mod h1:50CgR4k1jNlWBu4UfS4AcfhVe1r6pdZPygJ3R8F0Qdw=
 github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
@@ -230,6 +237,8 @@ github.com/gopherjs/gopherjs v0.0.0-20220410123724-9e86199038b0 h1:fWY+zXdWhvWnd
 github.com/gopherjs/gopherjs v0.0.0-20220410123724-9e86199038b0/go.mod h1:pRRIvn/QzFLrKfvEz3qUuEhtE/zLCWfreZ6J5gM2i+k=
 github.com/gorilla/mux v1.8.0 h1:i40aqfkR1h2SlN9hojwV5ZA91wcXFOvkdNIeFDP5koI=
 github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So=
+github.com/gorilla/websocket v1.5.1 h1:gmztn0JnHVt9JZquRuzLw3g4wouNVzKL15iLr/zn/QY=
+github.com/gorilla/websocket v1.5.1/go.mod h1:x3kM2JMyaluk02fnUJpQuwD2dCS5NDG2ZHL0uE0tcaY=
 github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.0.2-0.20240212192251-757544f21357 h1:Fkzd8ktnpOR9h47SXHe2AYPwelXLH2GjGsjlAloiWfo=
 github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.0.2-0.20240212192251-757544f21357/go.mod h1:w9Y7gY31krpLmrVU5ZPG9H7l9fZuRu5/3R3S3FMtVQ4=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.16.0 h1:YBftPWNWd4WwGqtY2yeZL2ef8rHAxPBD8KFhJpmcqms=
@@ -247,6 +256,7 @@ github.com/hashicorp/go-version v1.6.0 h1:feTTfFNnjP967rlCxM/I9g701jU+RN74YKx2mO
 github.com/hashicorp/go-version v1.6.0/go.mod h1:fltr4n8CU8Ke44wwGCBoEymUuxUHl09ZGVZPK5anwXA=
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
+github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
 github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
@@ -356,14 +366,14 @@ github.com/onsi/ginkgo v1.10.1/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+
 github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk=
 github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE=
 github.com/onsi/ginkgo v1.16.5/go.mod h1:+E8gABHa3K6zRBolWtd+ROzc/U5bkGt0FwiG042wbpU=
-github.com/onsi/ginkgo/v2 v2.4.0 h1:+Ig9nvqgS5OBSACXNk15PLdp0U9XPYROt9CFzVdFGIs=
-github.com/onsi/ginkgo/v2 v2.4.0/go.mod h1:iHkDK1fKGcBoEHT5W7YBq4RFWaQulw+caOMkAt4OrFo=
+github.com/onsi/ginkgo/v2 v2.9.5 h1:+6Hr4uxzP4XIUyAkg61dWBw8lb/gc4/X5luuxN/EC+Q=
+github.com/onsi/ginkgo/v2 v2.9.5/go.mod h1:tvAoo1QUJwNEU2ITftXTpR7R1RbCzoZUOs3RonqW57k=
 github.com/onsi/gomega v0.0.0-20170829124025-dcabb60a477c/go.mod h1:C1qb7wdrVGGVU+Z6iS04AVkA3Q65CEZX59MT0QO5uiA=
 github.com/onsi/gomega v1.7.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY=
 github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY=
 github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo=
-github.com/onsi/gomega v1.23.0 h1:/oxKu9c2HVap+F3PfKort2Hw5DEU+HGlW8n+tguWsys=
-github.com/onsi/gomega v1.23.0/go.mod h1:Z/NWtiqwBrwUt4/2loMmHL63EDLnYHmVbuBpDr2vQAg=
+github.com/onsi/gomega v1.27.6 h1:ENqfyGeS5AX/rlXDd/ETokDz93u0YufY1Pgxuy/PvWE=
+github.com/onsi/gomega v1.27.6/go.mod h1:PIQNjfQwkP3aQAH7lf7j87O/5FiNr+ZR8+ipb+qQlhg=
 github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
 github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
 github.com/opencontainers/image-spec v1.1.0 h1:8SG7/vwALn54lVB/0yZ/MMwhFrPYtpEHQb2IpWsCzug=
@@ -415,6 +425,8 @@ github.com/prometheus/common v0.53.0 h1:U2pL9w9nmJwJDa4qqLQ3ZaePJ6ZTwt7cMD3AG3+a
 github.com/prometheus/common v0.53.0/go.mod h1:BrxBKv3FWBIGXw89Mg1AeBq7FSyRzXWI3l3e7W3RN5U=
 github.com/prometheus/procfs v0.15.0 h1:A82kmvXJq2jTu5YUhSGNlYoxh85zLnKgPz4bMZgI5Ek=
 github.com/prometheus/procfs v0.15.0/go.mod h1:Y0RJ/Y5g5wJpkTisOtqwDSo4HwhGmLB4VQSw2sQJLHk=
+github.com/quic-go/quic-go v0.45.0 h1:OHmkQGM37luZITyTSu6ff03HP/2IrwDX1ZFiNEhSFUE=
+github.com/quic-go/quic-go v0.45.0/go.mod h1:1dLehS7TIR64+vxGR70GDcatWTOtMX2PUtnKsjbTurI=
 github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDNXVBh4M=
 github.com/rogpeppe/go-internal v1.11.0/go.mod h1:ddIwULY96R17DhadqLgMfk9H9tvdUzkipdSkR5nkCZA=
 github.com/rs/cors v1.8.0 h1:P2KMzcFwrPoSjkF1WLRPsp3UMLyql8L4v9hQpVeK5so=
@@ -462,6 +474,7 @@ github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXf
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
+github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
@@ -522,6 +535,8 @@ go.opentelemetry.io/otel/trace v1.26.0 h1:1ieeAUb4y0TE26jUFrCIXKpTuVK7uJGN9/Z/2L
 go.opentelemetry.io/otel/trace v1.26.0/go.mod h1:4iDxvGDQuUkHve82hJJ8UqrwswHYsZuWCBllGV2U2y0=
 go.opentelemetry.io/proto/otlp v1.0.0 h1:T0TX0tmXU8a3CbNXzEKGeU5mIVOdf0oykP+u2lIVU/I=
 go.opentelemetry.io/proto/otlp v1.0.0/go.mod h1:Sy6pihPLfYHkr3NkUbEhGHFhINUSI/v80hjKIs5JXpM=
+go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU=
+go.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc=
 goauthentik.io/api/v3 v3.2023051.3 h1:NebAhD/TeTWNo/9X3/Uj+rM5fG1HaiLOlKTNLQv9Qq4=
 goauthentik.io/api/v3 v3.2023051.3/go.mod h1:nYECml4jGbp/541hj8GcylKQG1gVBsKppHy4+7G8u4U=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
@@ -554,6 +569,8 @@ golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
 golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
+golang.org/x/mod v0.17.0 h1:zY54UmvipHiNd+pm+m0x9KhZ9hl1/7QNMyxXbc6ICqA=
+golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
 golang.org/x/net v0.0.0-20170114055629-f2499483f923/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
@@ -608,6 +625,7 @@ golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20191010194322-b09406accb47/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20191120155948-bd437916bb0e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20191204072324-ce4227a45e2e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200217220822-9197077df867/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200515095857-1151b9dac4a9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
@@ -676,6 +694,8 @@ golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/tools v0.1.5/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
 golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
+golang.org/x/tools v0.21.0 h1:qc0xYgIbsSDt9EyWz05J5wfa7LOVW0YTLOXrqdLAWIw=
+golang.org/x/tools v0.21.0/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
@@ -770,5 +790,7 @@ k8s.io/gengo v0.0.0-20190128074634-0689ccc1d7d6/go.mod h1:ezvh/TsK7cY6rbqRK0oQQ8
 k8s.io/klog v0.0.0-20181102134211-b9b56d5dfc92/go.mod h1:Gq+BEi5rUBO/HRz0bTSXDUcqjScdoY3a9IHpCEIOOfk=
 k8s.io/klog v1.0.0/go.mod h1:4Bi6QPql/J/LkTDqv7R/cd3hPo4k2DG6Ptcz060Ez5I=
 k8s.io/kube-openapi v0.0.0-20191107075043-30be4d16710a/go.mod h1:1TqjTSzOxsLGIKfj0lK8EeCP7K1iUG65v09OM0/WG5E=
+nhooyr.io/websocket v1.8.11 h1:f/qXNc2/3DpoSZkHt1DQu6rj4zGC8JmkkLkWss0MgN0=
+nhooyr.io/websocket v1.8.11/go.mod h1:rN9OFWIUwuxg4fR5tELlYC04bXYowCP9GX47ivo2l+c=
 sigs.k8s.io/structured-merge-diff v0.0.0-20190525122527-15d366b2352e/go.mod h1:wWxsB5ozmmv/SG7nM11ayaAW51xMvak/t1r0CSlcokI=
 sigs.k8s.io/yaml v1.1.0/go.mod h1:UJmg0vDUVViEyp3mgSv9WPwZCDxu4rQW1olrI1uml+o=

management/server/account.go

@@ -132,6 +132,7 @@ type AccountManager interface {
 	GetValidatedPeers(account *Account) (map[string]struct{}, error)
 	SyncAndMarkPeer(peerPubKey string, realIP net.IP) (*nbpeer.Peer, *NetworkMap, error)
 	CancelPeerRoutines(peer *nbpeer.Peer) error
+	FindExistingPostureCheck(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error)
 }
 
 type DefaultAccountManager struct {
@@ -241,6 +242,11 @@ type Account struct {
 	Settings *Settings `gorm:"embedded;embeddedPrefix:settings_"`
 }
 
+// Subclass used in gorm to only load settings and not whole account
+type AccountSettings struct {
+	Settings *Settings `gorm:"embedded;embeddedPrefix:settings_"`
+}
+
 type UserPermissions struct {
 	DashboardView string `json:"dashboard_view"`
 }
@@ -1768,6 +1774,8 @@ func (am *DefaultAccountManager) GetAccountFromToken(claims jwtclaims.Authorizat
 //
 // Existing user + Existing account + Existing domain reclassified Domain as private -> Nothing changes (index domain)
 func (am *DefaultAccountManager) getAccountWithAuthorizationClaims(claims jwtclaims.AuthorizationClaims) (*Account, error) {
+	log.Tracef("getting account with authorization claims. User ID: \"%s\", Account ID: \"%s\", Domain: \"%s\", Domain Category: \"%s\"",
+		claims.UserId, claims.AccountId, claims.Domain, claims.DomainCategory)
 	if claims.UserId == "" {
 		return nil, fmt.Errorf("user ID is empty")
 	}
@@ -1959,6 +1967,10 @@ func (am *DefaultAccountManager) onPeersInvalidated(accountID string) {
 	am.updateAccountPeers(updatedAccount)
 }
 
+func (am *DefaultAccountManager) FindExistingPostureCheck(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error) {
+	return am.Store.GetPostureCheckByChecksDefinition(accountID, checks)
+}
+
 // addAllGroup to account object if it doesn't exist
 func addAllGroup(account *Account) error {
 	if len(account.Groups) == 0 {

management/server/file_store.go

@@ -12,6 +12,7 @@ import (
 
 	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/netbirdio/netbird/management/server/status"
 	"github.com/netbirdio/netbird/management/server/telemetry"
 
@@ -508,7 +509,7 @@ func (s *FileStore) GetAccountByUser(userID string) (*Account, error) {
 
 	accountID, ok := s.UserID2AccountID[userID]
 	if !ok {
-		return nil, status.Errorf(status.NotFound, "account not found")
+		return nil, status.NewUserNotFoundError(userID)
 	}
 
 	account, err := s.getAccount(accountID)
@@ -539,7 +540,7 @@ func (s *FileStore) GetAccountByPeerID(peerID string) (*Account, error) {
 	if _, ok := account.Peers[peerID]; !ok {
 		delete(s.PeerID2AccountID, peerID)
 		log.Warnf("removed stale peerID %s to accountID %s index", peerID, accountID)
-		return nil, status.Errorf(status.NotFound, "provided peer doesn't exists %s", peerID)
+		return nil, status.NewPeerNotFoundError(peerID)
 	}
 
 	return account.Copy(), nil
@@ -552,7 +553,7 @@ func (s *FileStore) GetAccountByPeerPubKey(peerKey string) (*Account, error) {
 
 	accountID, ok := s.PeerKeyID2AccountID[peerKey]
 	if !ok {
-		return nil, status.Errorf(status.NotFound, "provided peer key doesn't exists %s", peerKey)
+		return nil, status.NewPeerNotFoundError(peerKey)
 	}
 
 	account, err := s.getAccount(accountID)
@@ -572,7 +573,7 @@ func (s *FileStore) GetAccountByPeerPubKey(peerKey string) (*Account, error) {
 	if stale {
 		delete(s.PeerKeyID2AccountID, peerKey)
 		log.Warnf("removed stale peerKey %s to accountID %s index", peerKey, accountID)
-		return nil, status.Errorf(status.NotFound, "provided peer doesn't exists %s", peerKey)
+		return nil, status.NewPeerNotFoundError(peerKey)
 	}
 
 	return account.Copy(), nil
@@ -584,12 +585,71 @@ func (s *FileStore) GetAccountIDByPeerPubKey(peerKey string) (string, error) {
 
 	accountID, ok := s.PeerKeyID2AccountID[peerKey]
 	if !ok {
-		return "", status.Errorf(status.NotFound, "provided peer key doesn't exists %s", peerKey)
+		return "", status.NewPeerNotFoundError(peerKey)
 	}
 
 	return accountID, nil
 }
 
+func (s *FileStore) GetAccountIDByUserID(userID string) (string, error) {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
+	accountID, ok := s.UserID2AccountID[userID]
+	if !ok {
+		return "", status.NewUserNotFoundError(userID)
+	}
+
+	return accountID, nil
+}
+
+func (s *FileStore) GetAccountIDBySetupKey(setupKey string) (string, error) {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
+	accountID, ok := s.SetupKeyID2AccountID[strings.ToUpper(setupKey)]
+	if !ok {
+		return "", status.Errorf(status.NotFound, "account not found: provided setup key doesn't exists")
+	}
+
+	return accountID, nil
+}
+
+func (s *FileStore) GetPeerByPeerPubKey(peerKey string) (*nbpeer.Peer, error) {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
+	accountID, ok := s.PeerKeyID2AccountID[peerKey]
+	if !ok {
+		return nil, status.NewPeerNotFoundError(peerKey)
+	}
+
+	account, err := s.getAccount(accountID)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, peer := range account.Peers {
+		if peer.Key == peerKey {
+			return peer.Copy(), nil
+		}
+	}
+
+	return nil, status.NewPeerNotFoundError(peerKey)
+}
+
+func (s *FileStore) GetAccountSettings(accountID string) (*Settings, error) {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
+	account, err := s.getAccount(accountID)
+	if err != nil {
+		return nil, err
+	}
+
+	return account.Settings.Copy(), nil
+}
+
 // GetInstallationID returns the installation ID from the store
 func (s *FileStore) GetInstallationID() string {
 	return s.InstallationID
@@ -667,6 +727,10 @@ func (s *FileStore) SaveUserLastLogin(accountID, userID string, lastLogin time.T
 	return nil
 }
 
+func (s *FileStore) GetPostureCheckByChecksDefinition(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error) {
+	return nil, status.Errorf(status.Internal, "GetPostureCheckByChecksDefinition is not implemented")
+}
+
 // Close the FileStore persisting data to disk
 func (s *FileStore) Close() error {
 	s.mux.Lock()

management/server/group.go

@@ -245,6 +245,11 @@ func (am *DefaultAccountManager) DeleteGroup(accountId, userId, groupID string)
 				return &GroupLinkError{"route", string(r.NetID)}
 			}
 		}
+		for _, g := range r.PeerGroups {
+			if g == groupID {
+				return &GroupLinkError{"route", string(r.NetID)}
+			}
+		}
 	}
 
 	// check DNS links

management/server/group_test.go

@@ -70,6 +70,11 @@ func TestDefaultAccountManager_DeleteGroup(t *testing.T) {
 			"grp-for-route",
 			"route",
 		},
+		{
+			"route with peer groups",
+			"grp-for-route2",
+			"route",
+		},
 		{
 			"name server groups",
 			"grp-for-name-server-grp",
@@ -138,6 +143,14 @@ func initTestGroupAccount(am *DefaultAccountManager) (*Account, error) {
 		Peers:     make([]string, 0),
 	}
 
+	groupForRoute2 := &nbgroup.Group{
+		ID:        "grp-for-route2",
+		AccountID: "account-id",
+		Name:      "Group for route",
+		Issued:    nbgroup.GroupIssuedAPI,
+		Peers:     make([]string, 0),
+	}
+
 	groupForNameServerGroups := &nbgroup.Group{
 		ID:        "grp-for-name-server-grp",
 		AccountID: "account-id",
@@ -183,6 +196,11 @@ func initTestGroupAccount(am *DefaultAccountManager) (*Account, error) {
 		Groups: []string{groupForRoute.ID},
 	}
 
+	routePeerGroupResource := &route.Route{
+		ID:         "example route with peer groups",
+		PeerGroups: []string{groupForRoute2.ID},
+	}
+
 	nameServerGroup := &nbdns.NameServerGroup{
 		ID:     "example name server group",
 		Groups: []string{groupForNameServerGroups.ID},
@@ -209,6 +227,7 @@ func initTestGroupAccount(am *DefaultAccountManager) (*Account, error) {
 	}
 	account := newAccountWithId(accountID, groupAdminUserID, domain)
 	account.Routes[routeResource.ID] = routeResource
+	account.Routes[routePeerGroupResource.ID] = routePeerGroupResource
 	account.NameServerGroups[nameServerGroup.ID] = nameServerGroup
 	account.Policies = append(account.Policies, policy)
 	account.SetupKeys[setupKey.Id] = setupKey
@@ -220,6 +239,7 @@ func initTestGroupAccount(am *DefaultAccountManager) (*Account, error) {
 	}
 
 	_ = am.SaveGroup(accountID, groupAdminUserID, groupForRoute)
+	_ = am.SaveGroup(accountID, groupAdminUserID, groupForRoute2)
 	_ = am.SaveGroup(accountID, groupAdminUserID, groupForNameServerGroups)
 	_ = am.SaveGroup(accountID, groupAdminUserID, groupForPolicies)
 	_ = am.SaveGroup(accountID, groupAdminUserID, groupForSetupKeys)

management/server/grpcserver.go

@@ -368,7 +368,7 @@ func (s *GRPCServer) Login(ctx context.Context, req *proto.EncryptedMessage) (*p
 	})
 
 	if err != nil {
-		log.Warnf("failed logging in peer %s", peerKey)
+		log.Warnf("failed logging in peer %s: %s", peerKey, err)
 		return nil, mapError(err)
 	}
 

management/server/http/posture_checks_handler.go

@@ -3,12 +3,10 @@ package http
 import (
 	"encoding/json"
 	"net/http"
-	"net/netip"
 	"regexp"
 	"slices"
 
 	"github.com/gorilla/mux"
-	"github.com/rs/xid"
 
 	"github.com/netbirdio/netbird/management/server"
 	"github.com/netbirdio/netbird/management/server/geolocation"
@@ -59,7 +57,7 @@ func (p *PostureChecksHandler) GetAllPostureChecks(w http.ResponseWriter, r *htt
 
 	postureChecks := []*api.PostureCheck{}
 	for _, postureCheck := range accountPostureChecks {
-		postureChecks = append(postureChecks, toPostureChecksResponse(postureCheck))
+		postureChecks = append(postureChecks, postureCheck.ToAPIResponse())
 	}
 
 	util.WriteJSONObject(w, postureChecks)
@@ -130,7 +128,7 @@ func (p *PostureChecksHandler) GetPostureCheck(w http.ResponseWriter, r *http.Re
 		return
 	}
 
-	util.WriteJSONObject(w, toPostureChecksResponse(postureChecks))
+	util.WriteJSONObject(w, postureChecks.ToAPIResponse())
 }
 
 // DeletePostureCheck handles posture check deletion request
@@ -178,55 +176,26 @@ func (p *PostureChecksHandler) savePostureChecks(
 		return
 	}
 
-	if postureChecksID == "" {
-		postureChecksID = xid.New().String()
-	}
-
-	postureChecks := posture.Checks{
-		ID:          postureChecksID,
-		Name:        req.Name,
-		Description: req.Description,
-	}
-
-	if nbVersionCheck := req.Checks.NbVersionCheck; nbVersionCheck != nil {
-		postureChecks.Checks.NBVersionCheck = &posture.NBVersionCheck{
-			MinVersion: nbVersionCheck.MinVersion,
-		}
-	}
-
-	if osVersionCheck := req.Checks.OsVersionCheck; osVersionCheck != nil {
-		postureChecks.Checks.OSVersionCheck = &posture.OSVersionCheck{
-			Android: (*posture.MinVersionCheck)(osVersionCheck.Android),
-			Darwin:  (*posture.MinVersionCheck)(osVersionCheck.Darwin),
-			Ios:     (*posture.MinVersionCheck)(osVersionCheck.Ios),
-			Linux:   (*posture.MinKernelVersionCheck)(osVersionCheck.Linux),
-			Windows: (*posture.MinKernelVersionCheck)(osVersionCheck.Windows),
-		}
-	}
-
 	if geoLocationCheck := req.Checks.GeoLocationCheck; geoLocationCheck != nil {
 		if p.geolocationManager == nil {
 			// TODO: update error message to include geo db self hosted doc link when ready
 			util.WriteError(status.Errorf(status.PreconditionFailed, "Geo location database is not initialized"), w)
 			return
 		}
-		postureChecks.Checks.GeoLocationCheck = toPostureGeoLocationCheck(geoLocationCheck)
 	}
 
-	if peerNetworkRangeCheck := req.Checks.PeerNetworkRangeCheck; peerNetworkRangeCheck != nil {
-		postureChecks.Checks.PeerNetworkRangeCheck, err = toPeerNetworkRangeCheck(peerNetworkRangeCheck)
-		if err != nil {
-			util.WriteError(status.Errorf(status.InvalidArgument, "invalid network prefix"), w)
-			return
-		}
-	}
-
-	if err := p.accountManager.SavePostureChecks(account.Id, user.Id, &postureChecks); err != nil {
+	postureChecks, err := posture.NewChecksFromAPIPostureCheckUpdate(req, postureChecksID)
+	if err != nil {
 		util.WriteError(err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, toPostureChecksResponse(&postureChecks))
+	if err := p.accountManager.SavePostureChecks(account.Id, user.Id, postureChecks); err != nil {
+		util.WriteError(err, w)
+		return
+	}
+
+	util.WriteJSONObject(w, postureChecks.ToAPIResponse())
 }
 
 func validatePostureChecksUpdate(req api.PostureCheckUpdate) error {
@@ -294,105 +263,3 @@ func validatePostureChecksUpdate(req api.PostureCheckUpdate) error {
 
 	return nil
 }
-
-func toPostureChecksResponse(postureChecks *posture.Checks) *api.PostureCheck {
-	var checks api.Checks
-
-	if postureChecks.Checks.NBVersionCheck != nil {
-		checks.NbVersionCheck = &api.NBVersionCheck{
-			MinVersion: postureChecks.Checks.NBVersionCheck.MinVersion,
-		}
-	}
-
-	if postureChecks.Checks.OSVersionCheck != nil {
-		checks.OsVersionCheck = &api.OSVersionCheck{
-			Android: (*api.MinVersionCheck)(postureChecks.Checks.OSVersionCheck.Android),
-			Darwin:  (*api.MinVersionCheck)(postureChecks.Checks.OSVersionCheck.Darwin),
-			Ios:     (*api.MinVersionCheck)(postureChecks.Checks.OSVersionCheck.Ios),
-			Linux:   (*api.MinKernelVersionCheck)(postureChecks.Checks.OSVersionCheck.Linux),
-			Windows: (*api.MinKernelVersionCheck)(postureChecks.Checks.OSVersionCheck.Windows),
-		}
-	}
-
-	if postureChecks.Checks.GeoLocationCheck != nil {
-		checks.GeoLocationCheck = toGeoLocationCheckResponse(postureChecks.Checks.GeoLocationCheck)
-	}
-
-	if postureChecks.Checks.PeerNetworkRangeCheck != nil {
-		checks.PeerNetworkRangeCheck = toPeerNetworkRangeCheckResponse(postureChecks.Checks.PeerNetworkRangeCheck)
-	}
-
-	return &api.PostureCheck{
-		Id:          postureChecks.ID,
-		Name:        postureChecks.Name,
-		Description: &postureChecks.Description,
-		Checks:      checks,
-	}
-}
-
-func toGeoLocationCheckResponse(geoLocationCheck *posture.GeoLocationCheck) *api.GeoLocationCheck {
-	locations := make([]api.Location, 0, len(geoLocationCheck.Locations))
-	for _, loc := range geoLocationCheck.Locations {
-		l := loc // make G601 happy
-		var cityName *string
-		if loc.CityName != "" {
-			cityName = &l.CityName
-		}
-		locations = append(locations, api.Location{
-			CityName:    cityName,
-			CountryCode: loc.CountryCode,
-		})
-	}
-
-	return &api.GeoLocationCheck{
-		Action:    api.GeoLocationCheckAction(geoLocationCheck.Action),
-		Locations: locations,
-	}
-}
-
-func toPostureGeoLocationCheck(apiGeoLocationCheck *api.GeoLocationCheck) *posture.GeoLocationCheck {
-	locations := make([]posture.Location, 0, len(apiGeoLocationCheck.Locations))
-	for _, loc := range apiGeoLocationCheck.Locations {
-		cityName := ""
-		if loc.CityName != nil {
-			cityName = *loc.CityName
-		}
-		locations = append(locations, posture.Location{
-			CountryCode: loc.CountryCode,
-			CityName:    cityName,
-		})
-	}
-
-	return &posture.GeoLocationCheck{
-		Action:    string(apiGeoLocationCheck.Action),
-		Locations: locations,
-	}
-}
-
-func toPeerNetworkRangeCheckResponse(check *posture.PeerNetworkRangeCheck) *api.PeerNetworkRangeCheck {
-	netPrefixes := make([]string, 0, len(check.Ranges))
-	for _, netPrefix := range check.Ranges {
-		netPrefixes = append(netPrefixes, netPrefix.String())
-	}
-
-	return &api.PeerNetworkRangeCheck{
-		Ranges: netPrefixes,
-		Action: api.PeerNetworkRangeCheckAction(check.Action),
-	}
-}
-
-func toPeerNetworkRangeCheck(check *api.PeerNetworkRangeCheck) (*posture.PeerNetworkRangeCheck, error) {
-	prefixes := make([]netip.Prefix, 0)
-	for _, prefix := range check.Ranges {
-		parsedPrefix, err := netip.ParsePrefix(prefix)
-		if err != nil {
-			return nil, err
-		}
-		prefixes = append(prefixes, parsedPrefix)
-	}
-
-	return &posture.PeerNetworkRangeCheck{
-		Ranges: prefixes,
-		Action: string(check.Action),
-	}, nil
-}

management/server/idp/zitadel.go

@@ -154,7 +154,7 @@ func (zc *ZitadelCredentials) requestJWTToken() (*http.Response, error) {
 	data.Set("client_id", zc.clientConfig.ClientID)
 	data.Set("client_secret", zc.clientConfig.ClientSecret)
 	data.Set("grant_type", zc.clientConfig.GrantType)
-	data.Set("scope", "urn:zitadel:iam:org:project:id:zitadel:aud")
+	data.Set("scope", "openid urn:zitadel:iam:org:project:id:zitadel:aud")
 
 	payload := strings.NewReader(data.Encode())
 	req, err := http.NewRequest(http.MethodPost, zc.clientConfig.TokenEndpoint, payload)

management/server/metrics/selfhosted.go

@@ -26,7 +26,7 @@ const (
 	// defaultPushInterval default interval to push metrics
 	defaultPushInterval = 24 * time.Hour
 	// requestTimeout http request timeout
-	requestTimeout = 30 * time.Second
+	requestTimeout = 45 * time.Second
 )
 
 type getTokenResponse struct {
@@ -98,10 +98,7 @@ func (w *Worker) Run() {
 }
 
 func (w *Worker) sendMetrics() error {
-	ctx, cancel := context.WithTimeout(w.ctx, requestTimeout)
-	defer cancel()
-
-	apiKey, err := getAPIKey(ctx)
+	apiKey, err := getAPIKey(w.ctx)
 	if err != nil {
 		return err
 	}
@@ -115,7 +112,7 @@ func (w *Worker) sendMetrics() error {
 
 	httpClient := http.Client{}
 
-	exportJobReq, err := createPostRequest(ctx, payloadEndpoint+"/capture/", payloadString)
+	exportJobReq, err := createPostRequest(w.ctx, payloadEndpoint+"/capture/", payloadString)
 	if err != nil {
 		return fmt.Errorf("unable to create metrics post request %v", err)
 	}
@@ -328,6 +325,8 @@ func (w *Worker) generateProperties() properties {
 }
 
 func getAPIKey(ctx context.Context) (string, error) {
+	ctx, cancel := context.WithTimeout(ctx, requestTimeout)
+	defer cancel()
 
 	httpClient := http.Client{}
 
@@ -375,6 +374,8 @@ func buildMetricsPayload(payload pushPayload) (string, error) {
 }
 
 func createPostRequest(ctx context.Context, endpoint string, payloadStr string) (*http.Request, error) {
+	ctx, cancel := context.WithTimeout(ctx, requestTimeout)
+	defer cancel()
 	reqURL := endpoint
 
 	payload := strings.NewReader(payloadStr)

management/server/mock_server/account_mock.go

@@ -95,6 +95,7 @@ type MockAccountManager struct {
 	GetIdpManagerFunc                   func() idp.Manager
 	UpdateIntegratedValidatorGroupsFunc func(accountID string, userID string, groups []string) error
 	GroupValidationFunc                 func(accountId string, groups []string) (bool, error)
+	FindExistingPostureCheckFunc        func(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error)
 }
 
 func (am *MockAccountManager) SyncAndMarkPeer(peerPubKey string, realIP net.IP) (*nbpeer.Peer, *server.NetworkMap, error) {
@@ -734,3 +735,11 @@ func (am *MockAccountManager) GroupValidation(accountId string, groups []string)
 	}
 	return false, status.Errorf(codes.Unimplemented, "method GroupValidation is not implemented")
 }
+
+// FindExistingPostureCheck mocks FindExistingPostureCheck of the AccountManager interface
+func (am *MockAccountManager) FindExistingPostureCheck(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error) {
+	if am.FindExistingPostureCheckFunc != nil {
+		return am.FindExistingPostureCheckFunc(accountID, checks)
+	}
+	return nil, status.Errorf(codes.Unimplemented, "method FindExistingPostureCheck is not implemented")
+}

management/server/peer.go

@@ -335,24 +335,29 @@ func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *nbpeer.P
 	}
 
 	upperKey := strings.ToUpper(setupKey)
-	var account *Account
+	var accountID string
 	var err error
 	addedByUser := false
 	if len(userID) > 0 {
 		addedByUser = true
-		account, err = am.Store.GetAccountByUser(userID)
+		accountID, err = am.Store.GetAccountIDByUserID(userID)
 	} else {
-		account, err = am.Store.GetAccountBySetupKey(setupKey)
+		accountID, err = am.Store.GetAccountIDBySetupKey(setupKey)
 	}
 	if err != nil {
 		return nil, nil, status.Errorf(status.NotFound, "failed adding new peer: account not found")
 	}
 
-	unlock := am.Store.AcquireAccountWriteLock(account.Id)
-	defer unlock()
+	unlock := am.Store.AcquireAccountWriteLock(accountID)
+	defer func() {
+		if unlock != nil {
+			unlock()
+		}
+	}()
 
+	var account *Account
 	// ensure that we consider modification happened meanwhile (because we were outside the account lock when we fetched the account)
-	account, err = am.Store.GetAccount(account.Id)
+	account, err = am.Store.GetAccount(accountID)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -485,6 +490,10 @@ func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *nbpeer.P
 		return nil, nil, err
 	}
 
+	// Account is saved, we can release the lock
+	unlock()
+	unlock = nil
+
 	opEvent.TargetID = newPeer.ID
 	opEvent.Meta = newPeer.EventMeta(am.GetDNSDomain())
 	if !addedByUser {
@@ -507,7 +516,7 @@ func (am *DefaultAccountManager) AddPeer(setupKey, userID string, peer *nbpeer.P
 func (am *DefaultAccountManager) SyncPeer(sync PeerSync, account *Account) (*nbpeer.Peer, *NetworkMap, error) {
 	peer, err := account.FindPeerByPubKey(sync.WireGuardPubKey)
 	if err != nil {
-		return nil, nil, status.Errorf(status.Unauthenticated, "peer is not registered")
+		return nil, nil, status.NewPeerNotRegisteredError()
 	}
 
 	err = checkIfPeerOwnerIsBlocked(peer, account)
@@ -515,7 +524,7 @@ func (am *DefaultAccountManager) SyncPeer(sync PeerSync, account *Account) (*nbp
 		return nil, nil, err
 	}
 
-	if peerLoginExpired(peer, account) {
+	if peerLoginExpired(peer, account.Settings) {
 		return nil, nil, status.Errorf(status.PermissionDenied, "peer login has expired, please log in once more")
 	}
 
@@ -545,7 +554,7 @@ func (am *DefaultAccountManager) SyncPeer(sync PeerSync, account *Account) (*nbp
 // LoginPeer logs in or registers a peer.
 // If peer doesn't exist the function checks whether a setup key or a user is present and registers a new peer if so.
 func (am *DefaultAccountManager) LoginPeer(login PeerLogin) (*nbpeer.Peer, *NetworkMap, error) {
-	account, err := am.Store.GetAccountByPeerPubKey(login.WireGuardPubKey)
+	accountID, err := am.Store.GetAccountIDByPeerPubKey(login.WireGuardPubKey)
 	if err != nil {
 		if errStatus, ok := status.FromError(err); ok && errStatus.Type() == status.NotFound {
 			// we couldn't find this peer by its public key which can mean that peer hasn't been registered yet.
@@ -574,19 +583,59 @@ func (am *DefaultAccountManager) LoginPeer(login PeerLogin) (*nbpeer.Peer, *Netw
 		return nil, nil, status.Errorf(status.Internal, "failed while logging in peer")
 	}
 
-	// we found the peer, and we follow a normal login flow
-	unlock := am.Store.AcquireAccountWriteLock(account.Id)
-	defer unlock()
+	peer, err := am.Store.GetPeerByPeerPubKey(login.WireGuardPubKey)
+	if err != nil {
+		return nil, nil, status.NewPeerNotRegisteredError()
+	}
+
+	accSettings, err := am.Store.GetAccountSettings(accountID)
+	if err != nil {
+		return nil, nil, status.Errorf(status.Internal, "failed to get account settings: %s", err)
+	}
+
+	var isWriteLock bool
+
+	// duplicated logic from after the lock to have an early exit
+	expired := peerLoginExpired(peer, accSettings)
+	switch {
+	case expired:
+		if err := checkAuth(login.UserID, peer); err != nil {
+			return nil, nil, err
+		}
+		isWriteLock = true
+		log.Debugf("peer login expired, acquiring write lock")
+
+	case peer.UpdateMetaIfNew(login.Meta):
+		isWriteLock = true
+		log.Debugf("peer changed meta, acquiring write lock")
+
+	default:
+		isWriteLock = false
+		log.Debugf("peer meta is the same, acquiring read lock")
+	}
+
+	var unlock func()
+
+	if isWriteLock {
+		unlock = am.Store.AcquireAccountWriteLock(accountID)
+	} else {
+		unlock = am.Store.AcquireAccountReadLock(accountID)
+	}
+	defer func() {
+		if unlock != nil {
+			unlock()
+		}
+	}()
 
 	// fetch the account from the store once more after acquiring lock to avoid concurrent updates inconsistencies
-	account, err = am.Store.GetAccount(account.Id)
+	account, err := am.Store.GetAccount(accountID)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	peer, err := account.FindPeerByPubKey(login.WireGuardPubKey)
+	peer, err = account.FindPeerByPubKey(login.WireGuardPubKey)
 	if err != nil {
-		return nil, nil, status.Errorf(status.Unauthenticated, "peer is not registered")
+		return nil, nil, status.NewPeerNotRegisteredError()
 	}
 
 	err = checkIfPeerOwnerIsBlocked(peer, account)
@@ -597,7 +646,7 @@ func (am *DefaultAccountManager) LoginPeer(login PeerLogin) (*nbpeer.Peer, *Netw
 	// this flag prevents unnecessary calls to the persistent store.
 	shouldStoreAccount := false
 	updateRemotePeers := false
-	if peerLoginExpired(peer, account) {
+	if peerLoginExpired(peer, account.Settings) {
 		err = checkAuth(login.UserID, peer)
 		if err != nil {
 			return nil, nil, err
@@ -633,11 +682,17 @@ func (am *DefaultAccountManager) LoginPeer(login PeerLogin) (*nbpeer.Peer, *Netw
 	}
 
 	if shouldStoreAccount {
+		if !isWriteLock {
+			log.Errorf("account %s should be stored but is not write locked", accountID)
+			return nil, nil, status.Errorf(status.Internal, "account should be stored but is not write locked")
+		}
 		err = am.Store.SaveAccount(account)
 		if err != nil {
 			return nil, nil, err
 		}
 	}
+	unlock()
+	unlock = nil
 
 	if updateRemotePeers || isStatusChanged {
 		am.updateAccountPeers(account)
@@ -683,9 +738,9 @@ func checkAuth(loginUserID string, peer *nbpeer.Peer) error {
 	return nil
 }
 
-func peerLoginExpired(peer *nbpeer.Peer, account *Account) bool {
-	expired, expiresIn := peer.LoginExpired(account.Settings.PeerLoginExpiration)
-	expired = account.Settings.PeerLoginExpirationEnabled && expired
+func peerLoginExpired(peer *nbpeer.Peer, settings *Settings) bool {
+	expired, expiresIn := peer.LoginExpired(settings.PeerLoginExpiration)
+	expired = settings.PeerLoginExpirationEnabled && expired
 	if expired || peer.Status.LoginExpired {
 		log.Debugf("peer's %s login expired %v ago", peer.ID, expiresIn)
 		return true

management/server/posture/checks.go

@@ -5,8 +5,11 @@ import (
 	"net/netip"
 
 	"github.com/hashicorp/go-version"
+	"github.com/rs/xid"
 
+	"github.com/netbirdio/netbird/management/server/http/api"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/status"
 )
 
 const (
@@ -136,6 +139,96 @@ func (pc *Checks) GetChecks() []Check {
 	return checks
 }
 
+func NewChecksFromAPIPostureCheck(source api.PostureCheck) (*Checks, error) {
+	description := ""
+	if source.Description != nil {
+		description = *source.Description
+	}
+
+	return buildPostureCheck(source.Id, source.Name, description, source.Checks)
+}
+
+func NewChecksFromAPIPostureCheckUpdate(source api.PostureCheckUpdate, postureChecksID string) (*Checks, error) {
+	return buildPostureCheck(postureChecksID, source.Name, source.Description, *source.Checks)
+}
+
+func buildPostureCheck(postureChecksID string, name string, description string, checks api.Checks) (*Checks, error) {
+	if postureChecksID == "" {
+		postureChecksID = xid.New().String()
+	}
+
+	postureChecks := Checks{
+		ID:          postureChecksID,
+		Name:        name,
+		Description: description,
+	}
+
+	if nbVersionCheck := checks.NbVersionCheck; nbVersionCheck != nil {
+		postureChecks.Checks.NBVersionCheck = &NBVersionCheck{
+			MinVersion: nbVersionCheck.MinVersion,
+		}
+	}
+
+	if osVersionCheck := checks.OsVersionCheck; osVersionCheck != nil {
+		postureChecks.Checks.OSVersionCheck = &OSVersionCheck{
+			Android: (*MinVersionCheck)(osVersionCheck.Android),
+			Darwin:  (*MinVersionCheck)(osVersionCheck.Darwin),
+			Ios:     (*MinVersionCheck)(osVersionCheck.Ios),
+			Linux:   (*MinKernelVersionCheck)(osVersionCheck.Linux),
+			Windows: (*MinKernelVersionCheck)(osVersionCheck.Windows),
+		}
+	}
+
+	if geoLocationCheck := checks.GeoLocationCheck; geoLocationCheck != nil {
+		postureChecks.Checks.GeoLocationCheck = toPostureGeoLocationCheck(geoLocationCheck)
+	}
+
+	var err error
+	if peerNetworkRangeCheck := checks.PeerNetworkRangeCheck; peerNetworkRangeCheck != nil {
+		postureChecks.Checks.PeerNetworkRangeCheck, err = toPeerNetworkRangeCheck(peerNetworkRangeCheck)
+		if err != nil {
+			return nil, status.Errorf(status.InvalidArgument, "invalid network prefix")
+		}
+	}
+
+	return &postureChecks, nil
+}
+
+func (pc *Checks) ToAPIResponse() *api.PostureCheck {
+	var checks api.Checks
+
+	if pc.Checks.NBVersionCheck != nil {
+		checks.NbVersionCheck = &api.NBVersionCheck{
+			MinVersion: pc.Checks.NBVersionCheck.MinVersion,
+		}
+	}
+
+	if pc.Checks.OSVersionCheck != nil {
+		checks.OsVersionCheck = &api.OSVersionCheck{
+			Android: (*api.MinVersionCheck)(pc.Checks.OSVersionCheck.Android),
+			Darwin:  (*api.MinVersionCheck)(pc.Checks.OSVersionCheck.Darwin),
+			Ios:     (*api.MinVersionCheck)(pc.Checks.OSVersionCheck.Ios),
+			Linux:   (*api.MinKernelVersionCheck)(pc.Checks.OSVersionCheck.Linux),
+			Windows: (*api.MinKernelVersionCheck)(pc.Checks.OSVersionCheck.Windows),
+		}
+	}
+
+	if pc.Checks.GeoLocationCheck != nil {
+		checks.GeoLocationCheck = toGeoLocationCheckResponse(pc.Checks.GeoLocationCheck)
+	}
+
+	if pc.Checks.PeerNetworkRangeCheck != nil {
+		checks.PeerNetworkRangeCheck = toPeerNetworkRangeCheckResponse(pc.Checks.PeerNetworkRangeCheck)
+	}
+
+	return &api.PostureCheck{
+		Id:          pc.ID,
+		Name:        pc.Name,
+		Description: &pc.Description,
+		Checks:      checks,
+	}
+}
+
 func (pc *Checks) Validate() error {
 	if check := pc.Checks.NBVersionCheck; check != nil {
 		if !isVersionValid(check.MinVersion) {
@@ -192,3 +285,70 @@ func isVersionValid(ver string) bool {
 
 	return false
 }
+
+func toGeoLocationCheckResponse(geoLocationCheck *GeoLocationCheck) *api.GeoLocationCheck {
+	locations := make([]api.Location, 0, len(geoLocationCheck.Locations))
+	for _, loc := range geoLocationCheck.Locations {
+		l := loc // make G601 happy
+		var cityName *string
+		if loc.CityName != "" {
+			cityName = &l.CityName
+		}
+		locations = append(locations, api.Location{
+			CityName:    cityName,
+			CountryCode: loc.CountryCode,
+		})
+	}
+
+	return &api.GeoLocationCheck{
+		Action:    api.GeoLocationCheckAction(geoLocationCheck.Action),
+		Locations: locations,
+	}
+}
+
+func toPostureGeoLocationCheck(apiGeoLocationCheck *api.GeoLocationCheck) *GeoLocationCheck {
+	locations := make([]Location, 0, len(apiGeoLocationCheck.Locations))
+	for _, loc := range apiGeoLocationCheck.Locations {
+		cityName := ""
+		if loc.CityName != nil {
+			cityName = *loc.CityName
+		}
+		locations = append(locations, Location{
+			CountryCode: loc.CountryCode,
+			CityName:    cityName,
+		})
+	}
+
+	return &GeoLocationCheck{
+		Action:    string(apiGeoLocationCheck.Action),
+		Locations: locations,
+	}
+}
+
+func toPeerNetworkRangeCheckResponse(check *PeerNetworkRangeCheck) *api.PeerNetworkRangeCheck {
+	netPrefixes := make([]string, 0, len(check.Ranges))
+	for _, netPrefix := range check.Ranges {
+		netPrefixes = append(netPrefixes, netPrefix.String())
+	}
+
+	return &api.PeerNetworkRangeCheck{
+		Ranges: netPrefixes,
+		Action: api.PeerNetworkRangeCheckAction(check.Action),
+	}
+}
+
+func toPeerNetworkRangeCheck(check *api.PeerNetworkRangeCheck) (*PeerNetworkRangeCheck, error) {
+	prefixes := make([]netip.Prefix, 0)
+	for _, prefix := range check.Ranges {
+		parsedPrefix, err := netip.ParsePrefix(prefix)
+		if err != nil {
+			return nil, err
+		}
+		prefixes = append(prefixes, parsedPrefix)
+	}
+
+	return &PeerNetworkRangeCheck{
+		Ranges: prefixes,
+		Action: string(check.Action),
+	}, nil
+}

management/server/sql_store.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
 	"path/filepath"
@@ -390,7 +391,7 @@ func (s *SqlStore) GetAccount(accountID string) (*Account, error) {
 		Preload(clause.Associations).
 		First(&account, "id = ?", accountID)
 	if result.Error != nil {
-		log.Errorf("error when getting account from the store: %s", result.Error)
+		log.Errorf("error when getting account %s from the store: %s", accountID, result.Error)
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "account not found")
 		}
@@ -457,7 +458,6 @@ func (s *SqlStore) GetAccountByUser(userID string) (*Account, error) {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
 		}
-		log.Errorf("error when getting user from the store: %s", result.Error)
 		return nil, status.Errorf(status.Internal, "issue getting account from store")
 	}
 
@@ -520,6 +520,61 @@ func (s *SqlStore) GetAccountIDByPeerPubKey(peerKey string) (string, error) {
 	return accountID, nil
 }
 
+func (s *SqlStore) GetAccountIDByUserID(userID string) (string, error) {
+	var user User
+	var accountID string
+	result := s.db.Model(&user).Select("account_id").Where("id = ?", userID).First(&accountID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
+		}
+		return "", status.Errorf(status.Internal, "issue getting account from store")
+	}
+
+	return accountID, nil
+}
+
+func (s *SqlStore) GetAccountIDBySetupKey(setupKey string) (string, error) {
+	var key SetupKey
+	var accountID string
+	result := s.db.Model(&key).Select("account_id").Where("key = ?", strings.ToUpper(setupKey)).First(&accountID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
+		}
+		log.Errorf("error when getting setup key from the store: %s", result.Error)
+		return "", status.Errorf(status.Internal, "issue getting setup key from store")
+	}
+
+	return accountID, nil
+}
+
+func (s *SqlStore) GetPeerByPeerPubKey(peerKey string) (*nbpeer.Peer, error) {
+	var peer nbpeer.Peer
+	result := s.db.First(&peer, "key = ?", peerKey)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.Errorf(status.NotFound, "peer not found")
+		}
+		log.Errorf("error when getting peer from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "issue getting peer from store")
+	}
+
+	return &peer, nil
+}
+
+func (s *SqlStore) GetAccountSettings(accountID string) (*Settings, error) {
+	var accountSettings AccountSettings
+	if err := s.db.Model(&Account{}).Where("id = ?", accountID).First(&accountSettings).Error; err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, status.Errorf(status.NotFound, "settings not found")
+		}
+		log.Errorf("error when getting settings from the store: %s", err)
+		return nil, status.Errorf(status.Internal, "issue getting settings from store")
+	}
+	return accountSettings.Settings, nil
+}
+
 // SaveUserLastLogin stores the last login time for a user in DB.
 func (s *SqlStore) SaveUserLastLogin(accountID, userID string, lastLogin time.Time) error {
 	var user User
@@ -529,7 +584,6 @@ func (s *SqlStore) SaveUserLastLogin(accountID, userID string, lastLogin time.Ti
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return status.Errorf(status.NotFound, "user %s not found", userID)
 		}
-		log.Errorf("error when getting user from the store: %s", result.Error)
 		return status.Errorf(status.Internal, "issue getting user from store")
 	}
 
@@ -538,6 +592,21 @@ func (s *SqlStore) SaveUserLastLogin(accountID, userID string, lastLogin time.Ti
 	return s.db.Save(user).Error
 }
 
+func (s *SqlStore) GetPostureCheckByChecksDefinition(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error) {
+	definitionJSON, err := json.Marshal(checks)
+	if err != nil {
+		return nil, err
+	}
+
+	var postureCheck posture.Checks
+	err = s.db.Where("account_id = ? AND checks = ?", accountID, string(definitionJSON)).First(&postureCheck).Error
+	if err != nil {
+		return nil, err
+	}
+
+	return &postureCheck, nil
+}
+
 // Close closes the underlying DB connection
 func (s *SqlStore) Close() error {
 	sql, err := s.db.DB()

management/server/status/error.go

@@ -75,3 +75,23 @@ func FromError(err error) (s *Error, ok bool) {
 	}
 	return nil, false
 }
+
+// NewPeerNotFoundError creates a new Error with NotFound type for a missing peer
+func NewPeerNotFoundError(peerKey string) error {
+	return Errorf(NotFound, "peer not found: %s", peerKey)
+}
+
+// NewAccountNotFoundError creates a new Error with NotFound type for a missing account
+func NewAccountNotFoundError(accountKey string) error {
+	return Errorf(NotFound, "account not found: %s", accountKey)
+}
+
+// NewUserNotFoundError creates a new Error with NotFound type for a missing user
+func NewUserNotFoundError(userKey string) error {
+	return Errorf(NotFound, "user not found: %s", userKey)
+}
+
+// NewPeerNotRegisteredError creates a new Error with NotFound type for a missing peer
+func NewPeerNotRegisteredError() error {
+	return Errorf(Unauthenticated, "peer is not registered")
+}

management/server/store.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/netbirdio/netbird/management/server/migration"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/netbirdio/netbird/management/server/telemetry"
 	"github.com/netbirdio/netbird/management/server/testutil"
 	"github.com/netbirdio/netbird/route"
@@ -26,11 +27,14 @@ type Store interface {
 	GetAccountByUser(userID string) (*Account, error)
 	GetAccountByPeerPubKey(peerKey string) (*Account, error)
 	GetAccountIDByPeerPubKey(peerKey string) (string, error)
+	GetAccountIDByUserID(peerKey string) (string, error)
+	GetAccountIDBySetupKey(peerKey string) (string, error)
 	GetAccountByPeerID(peerID string) (*Account, error)
 	GetAccountBySetupKey(setupKey string) (*Account, error) // todo use key hash later
 	GetAccountByPrivateDomain(domain string) (*Account, error)
 	GetTokenIDByHashedToken(secret string) (string, error)
 	GetUserByTokenID(tokenID string) (*User, error)
+	GetPostureCheckByChecksDefinition(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error)
 	SaveAccount(account *Account) error
 	DeleteHashedPAT2TokenIDIndex(hashedToken string) error
 	DeleteTokenID2UserIDIndex(tokenID string) error
@@ -50,6 +54,8 @@ type Store interface {
 	// GetStoreEngine should return StoreEngine of the current store implementation.
 	// This is also a method of metrics.DataSource interface.
 	GetStoreEngine() StoreEngine
+	GetPeerByPeerPubKey(peerKey string) (*nbpeer.Peer, error)
+	GetAccountSettings(accountID string) (*Settings, error)
 }
 
 type StoreEngine string

relay/client/client.go

@@ -0,0 +1,419 @@
+package client
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	ws "github.com/netbirdio/netbird/relay/client/dialer/wsnhooyr"
+	"github.com/netbirdio/netbird/relay/messages"
+)
+
+const (
+	bufferSize            = 8820
+	serverResponseTimeout = 8 * time.Second
+)
+
+// Msg carry the payload from the server to the client. With this sturct, 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)
+}
+
+type connContainer struct {
+	conn        *Conn
+	messages    chan Msg
+	msgChanLock sync.Mutex
+	closed      bool // flag to check if channel is closed
+}
+
+func newConnContainer(conn *Conn, messages chan Msg) *connContainer {
+	return &connContainer{
+		conn:     conn,
+		messages: messages,
+	}
+}
+
+func (cc *connContainer) writeMsg(msg Msg) {
+	cc.msgChanLock.Lock()
+	defer cc.msgChanLock.Unlock()
+	if cc.closed {
+		return
+	}
+	cc.messages <- msg
+}
+
+func (cc *connContainer) close() {
+	cc.msgChanLock.Lock()
+	defer cc.msgChanLock.Unlock()
+	if cc.closed {
+		return
+	}
+	close(cc.messages)
+	cc.closed = true
+}
+
+// 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.
+type Client struct {
+	log           *log.Entry
+	parentCtx     context.Context
+	ctxCancel     context.CancelFunc
+	serverAddress string
+	hashedID      []byte
+
+	bufPool *sync.Pool
+
+	relayConn        net.Conn
+	conns            map[string]*connContainer
+	serviceIsRunning bool
+	mu               sync.Mutex
+	readLoopMutex    sync.Mutex
+	wgReadLoop       sync.WaitGroup
+
+	remoteAddr net.Addr
+
+	onDisconnectListener func()
+	listenerMutex        sync.Mutex
+}
+
+// NewClient creates a new client for the relay server. The client is not connected to the server until the Connect
+func NewClient(ctx context.Context, serverAddress, peerID string) *Client {
+	hashedID, hashedStringId := messages.HashID(peerID)
+	return &Client{
+		log:           log.WithField("client_id", hashedStringId),
+		parentCtx:     ctx,
+		ctxCancel:     func() {},
+		serverAddress: serverAddress,
+		hashedID:      hashedID,
+		bufPool: &sync.Pool{
+			New: func() any {
+				buf := make([]byte, bufferSize)
+				return &buf
+			},
+		},
+		conns: make(map[string]*connContainer),
+	}
+}
+
+// Connect establishes a connection to the relay server. It blocks until the connection is established or an error occurs.
+func (c *Client) Connect() error {
+	c.readLoopMutex.Lock()
+	defer c.readLoopMutex.Unlock()
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if c.serviceIsRunning {
+		return nil
+	}
+
+	err := c.connect()
+	if err != nil {
+		return err
+	}
+
+	c.serviceIsRunning = true
+
+	var ctx context.Context
+	ctx, c.ctxCancel = context.WithCancel(c.parentCtx)
+	context.AfterFunc(ctx, func() {
+		cErr := c.close(false)
+		if cErr != nil {
+			log.Errorf("failed to close relay connection: %s", cErr)
+		}
+	})
+	c.wgReadLoop.Add(1)
+	go c.readLoop(c.relayConn)
+
+	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.
+// todo: what should happen if call with the same peerID with multiple times?
+func (c *Client) OpenConn(dstPeerID string) (net.Conn, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if !c.serviceIsRunning {
+		return nil, fmt.Errorf("relay connection is not established")
+	}
+
+	hashedID, hashedStringID := messages.HashID(dstPeerID)
+	log.Infof("open connection to peer: %s", hashedStringID)
+	msgChannel := make(chan Msg, 2)
+	conn := NewConn(c, hashedID, hashedStringID, msgChannel)
+
+	c.conns[hashedStringID] = newConnContainer(conn, msgChannel)
+	return conn, nil
+}
+
+// RelayRemoteAddress returns the IP address of the relay server. It could change after the close and reopen the connection.
+func (c *Client) RelayRemoteAddress() (net.Addr, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.remoteAddr == nil {
+		return nil, fmt.Errorf("relay connection is not established")
+	}
+	return c.remoteAddr, nil
+}
+
+// SetOnDisconnectListener sets a function that will be called when the connection to the relay server is closed.
+func (c *Client) SetOnDisconnectListener(fn func()) {
+	c.listenerMutex.Lock()
+	defer c.listenerMutex.Unlock()
+	c.onDisconnectListener = fn
+}
+
+func (c *Client) HasConns() bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return len(c.conns) > 0
+}
+
+// Close closes the connection to the relay server and all connections to other peers.
+func (c *Client) Close() error {
+	return c.close(false)
+}
+
+func (c *Client) close(byServer bool) error {
+	c.readLoopMutex.Lock()
+	defer c.readLoopMutex.Unlock()
+
+	c.mu.Lock()
+	var err error
+	if !c.serviceIsRunning {
+		c.mu.Unlock()
+		return nil
+	}
+
+	c.serviceIsRunning = false
+	c.closeAllConns()
+	if !byServer {
+		c.writeCloseMsg()
+		err = c.relayConn.Close()
+	}
+	c.mu.Unlock()
+
+	c.wgReadLoop.Wait()
+	c.log.Infof("relay connection closed with: %s", c.serverAddress)
+	c.ctxCancel()
+	return err
+}
+
+func (c *Client) connect() error {
+	conn, err := ws.Dial(c.serverAddress)
+	if err != nil {
+		return err
+	}
+	c.relayConn = conn
+
+	err = c.handShake()
+	if err != nil {
+		cErr := conn.Close()
+		if cErr != nil {
+			log.Errorf("failed to close connection: %s", cErr)
+		}
+		c.relayConn = nil
+		return err
+	}
+
+	c.remoteAddr = conn.RemoteAddr()
+
+	return nil
+}
+
+func (c *Client) handShake() error {
+	defer func() {
+		err := c.relayConn.SetReadDeadline(time.Time{})
+		if err != nil {
+			log.Errorf("failed to reset read deadline: %s", err)
+		}
+	}()
+
+	msg, err := messages.MarshalHelloMsg(c.hashedID)
+	if err != nil {
+		log.Errorf("failed to marshal hello message: %s", err)
+		return err
+	}
+	_, err = c.relayConn.Write(msg)
+	if err != nil {
+		log.Errorf("failed to send hello message: %s", err)
+		return err
+	}
+
+	err = c.relayConn.SetReadDeadline(time.Now().Add(serverResponseTimeout))
+	if err != nil {
+		log.Errorf("failed to set read deadline: %s", err)
+		return err
+	}
+
+	buf := make([]byte, 1500) // todo: optimise buffer size
+	n, err := c.relayConn.Read(buf)
+	if err != nil {
+		log.Errorf("failed to read hello response: %s", err)
+		return err
+	}
+
+	msgType, err := messages.DetermineServerMsgType(buf[:n])
+	if err != nil {
+		log.Errorf("failed to determine message type: %s", err)
+		return err
+	}
+
+	if msgType != messages.MsgTypeHelloResponse {
+		log.Errorf("unexpected message type: %s", msgType)
+		return fmt.Errorf("unexpected message type")
+	}
+	return nil
+}
+
+func (c *Client) readLoop(relayConn net.Conn) {
+	var (
+		errExit        error
+		n              int
+		closedByServer bool
+	)
+	for {
+		bufPtr := c.bufPool.Get().(*[]byte)
+		buf := *bufPtr
+		n, errExit = relayConn.Read(buf)
+		if errExit != nil {
+			c.mu.Lock()
+			if c.serviceIsRunning {
+				c.log.Debugf("failed to read message from relay server: %s", errExit)
+			}
+			c.mu.Unlock()
+			goto Exit
+		}
+
+		msgType, err := messages.DetermineServerMsgType(buf[:n])
+		if err != nil {
+			c.log.Errorf("failed to determine message type: %s", err)
+			continue
+		}
+
+		switch msgType {
+		case messages.MsgTypeTransport:
+			peerID, payload, err := messages.UnmarshalTransportMsg(buf[:n])
+			if err != nil {
+				c.log.Errorf("failed to parse transport message: %v", err)
+				continue
+			}
+			stringID := messages.HashIDToString(peerID)
+
+			c.mu.Lock()
+			if !c.serviceIsRunning {
+				c.mu.Unlock()
+				goto Exit
+			}
+			container, ok := c.conns[stringID]
+			c.mu.Unlock()
+			if !ok {
+				c.log.Errorf("peer not found: %s", stringID)
+				continue
+			}
+
+			container.writeMsg(Msg{
+				bufPool: c.bufPool,
+				bufPtr:  bufPtr,
+				Payload: payload})
+		case messages.MsgClose:
+			closedByServer = true
+			log.Debugf("relay connection close by server")
+			goto Exit
+		}
+	}
+
+Exit:
+	c.notifyDisconnected()
+	c.wgReadLoop.Done()
+	_ = c.close(closedByServer)
+}
+
+// todo check by reference too, the id is not enought because the id come from the outer conn
+func (c *Client) writeTo(id string, dstID []byte, payload []byte) (int, error) {
+	c.mu.Lock()
+	//	conn, ok := c.conns[id]
+	_, ok := c.conns[id]
+	c.mu.Unlock()
+	if !ok {
+		return 0, io.EOF
+	}
+	/*
+		if conn != clientRef {
+			return 0, io.EOF
+		}
+	*/
+	msg := messages.MarshalTransportMsg(dstID, payload)
+	n, err := c.relayConn.Write(msg)
+	if err != nil {
+		log.Errorf("failed to write transport message: %s", err)
+	}
+	return n, err
+}
+
+func (c *Client) closeAllConns() {
+	for _, container := range c.conns {
+		container.close()
+	}
+	c.conns = make(map[string]*connContainer)
+}
+
+// todo check by reference too, the id is not enought because the id come from the outer conn
+func (c *Client) closeConn(id string) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	container, ok := c.conns[id]
+	if !ok {
+		return fmt.Errorf("connection already closed")
+	}
+	container.close()
+	delete(c.conns, id)
+
+	return nil
+}
+
+func (c *Client) onDisconnect() {
+	c.listenerMutex.Lock()
+	defer c.listenerMutex.Unlock()
+
+	if c.onDisconnectListener == nil {
+		return
+	}
+	c.onDisconnectListener()
+}
+
+func (c *Client) notifyDisconnected() {
+	c.listenerMutex.Lock()
+	defer c.listenerMutex.Unlock()
+
+	if c.onDisconnectListener == nil {
+		return
+	}
+	go c.onDisconnectListener()
+}
+
+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)
+	}
+}

relay/client/client_test.go

@@ -0,0 +1,523 @@
+package client
+
+import (
+	"context"
+	"net"
+	"os"
+	"testing"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/util"
+
+	"github.com/netbirdio/netbird/relay/server"
+)
+
+func TestMain(m *testing.M) {
+	_ = util.InitLog("trace", "console")
+	code := m.Run()
+	os.Exit(code)
+}
+
+func TestClient(t *testing.T) {
+	ctx := context.Background()
+
+	addr := "localhost:1234"
+	srv := server.NewServer()
+	go func() {
+		err := srv.Listen(addr)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		err := srv.Close()
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	clientAlice := NewClient(ctx, addr, "alice")
+	err := clientAlice.Connect()
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+	defer clientAlice.Close()
+
+	clientPlaceHolder := NewClient(ctx, addr, "clientPlaceHolder")
+	err = clientPlaceHolder.Connect()
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+	defer clientPlaceHolder.Close()
+
+	clientBob := NewClient(ctx, addr, "bob")
+	err = clientBob.Connect()
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+	defer clientBob.Close()
+
+	connAliceToBob, err := clientAlice.OpenConn("bob")
+	if err != nil {
+		t.Fatalf("failed to bind channel: %s", err)
+	}
+
+	connBobToAlice, err := clientBob.OpenConn("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()
+	addr := "localhost:1234"
+	srv := server.NewServer()
+	go func() {
+		err := srv.Listen(addr)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+
+	clientAlice := NewClient(ctx, addr, "alice")
+	err := clientAlice.Connect()
+	if err != nil {
+		_ = srv.Close()
+		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.Close()
+	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(ctx, "127.0.0.1:1234", "alice")
+	err = clientAlice.Connect()
+	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"
+	addr := "localhost:1234"
+	srv := server.NewServer()
+	go func() {
+		err := srv.Listen(addr)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		err := srv.Close()
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	clientAlice := NewClient(ctx, addr, idAlice)
+	err := clientAlice.Connect()
+	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(ctx, addr, idBob)
+	err = clientBob.Connect()
+	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(idBob)
+	if err != nil {
+		t.Fatalf("failed to bind channel: %s", err)
+	}
+
+	connBobToAlice, err := clientBob.OpenConn(idAlice)
+	if err != nil {
+		t.Fatalf("failed to bind channel: %s", err)
+	}
+
+	payload := "hello bob, I am alice"
+	_, err = connAliceToBob.Write([]byte(payload))
+	if err != nil {
+		t.Fatalf("failed to write to channel: %s", err)
+	}
+
+	buf := make([]byte, 65535)
+	n, err := connBobToAlice.Read(buf)
+	if err != nil {
+		t.Fatalf("failed to read from channel: %s", err)
+	}
+
+	_, err = connBobToAlice.Write(buf[:n])
+	if err != nil {
+		t.Fatalf("failed to write to channel: %s", err)
+	}
+
+	n, err = connAliceToBob.Read(buf)
+	if err != nil {
+		t.Fatalf("failed to read from channel: %s", err)
+	}
+
+	if payload != string(buf[:n]) {
+		t.Fatalf("expected %s, got %s", payload, string(buf[:n]))
+	}
+}
+
+func TestBindToUnavailabePeer(t *testing.T) {
+	ctx := context.Background()
+
+	addr := "localhost:1234"
+	srv := server.NewServer()
+	go func() {
+		err := srv.Listen(addr)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		log.Infof("closing server")
+		err := srv.Close()
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	clientAlice := NewClient(ctx, addr, "alice")
+	err := clientAlice.Connect()
+	if err != nil {
+		t.Errorf("failed to connect to server: %s", err)
+	}
+	_, err = clientAlice.OpenConn("bob")
+	if err != nil {
+		t.Errorf("failed to bind channel: %s", err)
+	}
+
+	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()
+
+	addr := "localhost:1234"
+	srv := server.NewServer()
+	go func() {
+		err := srv.Listen(addr)
+		if err != nil {
+			t.Errorf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		log.Infof("closing server")
+		err := srv.Close()
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	clientAlice := NewClient(ctx, addr, "alice")
+	err := clientAlice.Connect()
+	if err != nil {
+		t.Errorf("failed to connect to server: %s", err)
+	}
+
+	_, err = clientAlice.OpenConn("bob")
+	if err != nil {
+		t.Errorf("failed to bind channel: %s", err)
+	}
+
+	clientBob := NewClient(ctx, addr, "bob")
+	err = clientBob.Connect()
+	if err != nil {
+		t.Errorf("failed to connect to server: %s", err)
+	}
+
+	chBob, err := clientBob.OpenConn("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(ctx, addr, "alice")
+	err = clientAlice.Connect()
+	if err != nil {
+		t.Errorf("failed to connect to server: %s", err)
+	}
+
+	chAlice, err := clientAlice.OpenConn("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("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()
+
+	addr := "localhost:1234"
+	srv := server.NewServer()
+	go func() {
+		err := srv.Listen(addr)
+		if err != nil {
+			t.Errorf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		log.Infof("closing server")
+		err := srv.Close()
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	clientAlice := NewClient(ctx, addr, "alice")
+	err := clientAlice.Connect()
+	if err != nil {
+		t.Errorf("failed to connect to server: %s", err)
+	}
+
+	conn, err := clientAlice.OpenConn("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()
+
+	addr := "localhost:1234"
+	srv := server.NewServer()
+	go func() {
+		err := srv.Listen(addr)
+		if err != nil {
+			t.Errorf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		err := srv.Close()
+		if err != nil {
+			log.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	clientAlice := NewClient(ctx, addr, "alice")
+	err := clientAlice.Connect()
+	if err != nil {
+		t.Fatalf("failed to connect to server: %s", err)
+	}
+
+	conn, err := clientAlice.OpenConn("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("bob")
+	if err == nil {
+		t.Errorf("unexpected opening connection to closed server")
+	}
+}
+
+func TestCloseByServer(t *testing.T) {
+	ctx := context.Background()
+
+	addr1 := "localhost:1234"
+	srv1 := server.NewServer()
+	go func() {
+		err := srv1.Listen(addr1)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+
+	idAlice := "alice"
+	log.Debugf("connect by alice")
+	relayClient := NewClient(ctx, addr1, idAlice)
+	err := relayClient.Connect()
+	if err != nil {
+		log.Fatalf("failed to connect to server: %s", err)
+	}
+
+	disconnected := make(chan struct{})
+	relayClient.SetOnDisconnectListener(func() {
+		log.Infof("client disconnected")
+		close(disconnected)
+	})
+
+	err = srv1.Close()
+	if err != nil {
+		t.Fatalf("failed to close server: %s", err)
+	}
+
+	select {
+	case <-disconnected:
+	case <-time.After(3 * time.Second):
+		log.Fatalf("timeout waiting for client to disconnect")
+	}
+
+	_, err = relayClient.OpenConn("bob")
+	if err == nil {
+		t.Errorf("unexpected opening connection to closed server")
+	}
+}
+
+func TestCloseByClient(t *testing.T) {
+	ctx := context.Background()
+
+	addr1 := "localhost:1234"
+	srv := server.NewServer()
+	go func() {
+		err := srv.Listen(addr1)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+
+	idAlice := "alice"
+	log.Debugf("connect by alice")
+	relayClient := NewClient(ctx, addr1, idAlice)
+	err := relayClient.Connect()
+	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("bob")
+	if err == nil {
+		t.Errorf("unexpected opening connection to closed server")
+	}
+
+	err = srv.Close()
+	if err != nil {
+		t.Fatalf("failed to close server: %s", err)
+	}
+}

relay/client/conn.go

@@ -0,0 +1,67 @@
+package client
+
+import (
+	"io"
+	"net"
+	"time"
+)
+
+type Conn struct {
+	client      *Client
+	dstID       []byte
+	dstStringID string
+	messageChan chan Msg
+}
+
+func NewConn(client *Client, dstID []byte, dstStringID string, messageChan chan Msg) *Conn {
+	c := &Conn{
+		client:      client,
+		dstID:       dstID,
+		dstStringID: dstStringID,
+		messageChan: messageChan,
+	}
+
+	return c
+}
+
+func (c *Conn) Write(p []byte) (n int, err error) {
+	return c.client.writeTo(c.dstStringID, c.dstID, p)
+}
+
+func (c *Conn) Read(b []byte) (n int, err error) {
+	msg, ok := <-c.messageChan
+	if !ok {
+		return 0, io.EOF
+	}
+
+	n = copy(b, msg.Payload)
+	msg.Free()
+	return n, nil
+}
+
+func (c *Conn) Close() error {
+	return c.client.closeConn(c.dstStringID)
+}
+
+func (c *Conn) LocalAddr() net.Addr {
+	return c.client.relayConn.LocalAddr()
+}
+
+func (c *Conn) RemoteAddr() net.Addr {
+	return c.client.relayConn.RemoteAddr()
+}
+
+func (c *Conn) SetDeadline(t time.Time) error {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (c *Conn) SetReadDeadline(t time.Time) error {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (c *Conn) SetWriteDeadline(t time.Time) error {
+	//TODO implement me
+	panic("implement me")
+}

relay/client/dialer/quic/conn.go

@@ -0,0 +1,52 @@
+package quic
+
+import (
+	"net"
+
+	"github.com/quic-go/quic-go"
+	log "github.com/sirupsen/logrus"
+)
+
+type Conn struct {
+	quic.Stream
+	qConn quic.Connection
+}
+
+func NewConn(stream quic.Stream, qConn quic.Connection) net.Conn {
+	return &Conn{
+		Stream: stream,
+		qConn:  qConn,
+	}
+}
+
+func (q *Conn) Write(b []byte) (n int, err error) {
+	log.Debugf("writing: %d, %x\n", len(b), b)
+	n, err = q.Stream.Write(b)
+	if n != len(b) {
+		log.Errorf("failed to write out the full message")
+	}
+	return
+}
+
+func (q *Conn) Close() error {
+	err := q.Stream.Close()
+	if err != nil {
+		log.Errorf("failed to close stream: %s", err)
+		return err
+	}
+	err = q.qConn.CloseWithError(0, "")
+	if err != nil {
+		log.Errorf("failed to close connection: %s", err)
+		return err
+
+	}
+	return err
+}
+
+func (c *Conn) LocalAddr() net.Addr {
+	return c.qConn.LocalAddr()
+}
+
+func (c *Conn) RemoteAddr() net.Addr {
+	return c.qConn.RemoteAddr()
+}

relay/client/dialer/quic/quic.go

@@ -0,0 +1,32 @@
+package quic
+
+import (
+	"context"
+	"crypto/tls"
+	"net"
+
+	"github.com/quic-go/quic-go"
+	log "github.com/sirupsen/logrus"
+)
+
+func Dial(address string) (net.Conn, error) {
+	tlsConf := &tls.Config{
+		InsecureSkipVerify: true,
+		NextProtos:         []string{"quic-echo-example"},
+	}
+	qConn, err := quic.DialAddr(context.Background(), address, tlsConf, &quic.Config{
+		EnableDatagrams: true,
+	})
+	if err != nil {
+		log.Errorf("dial quic address %s failed: %s", address, err)
+		return nil, err
+	}
+
+	stream, err := qConn.OpenStreamSync(context.Background())
+	if err != nil {
+		return nil, err
+	}
+
+	conn := NewConn(stream, qConn)
+	return conn, nil
+}

relay/client/dialer/tcp/tcp.go

@@ -0,0 +1,7 @@
+package tcp
+
+import "net"
+
+func Dial(address string) (net.Conn, error) {
+	return net.Dial("tcp", address)
+}

relay/client/dialer/udp/udp.go

@@ -0,0 +1,14 @@
+package udp
+
+import (
+	"net"
+)
+
+func Dial(address string) (net.Conn, error) {
+	udpAddr, err := net.ResolveUDPAddr("udp", address)
+	if err != nil {
+		return nil, err
+	}
+
+	return net.DialUDP("udp", nil, udpAddr)
+}

relay/client/dialer/ws/client_conn.go

@@ -0,0 +1,59 @@
+package ws
+
+import (
+	"fmt"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+type Conn struct {
+	*websocket.Conn
+	mu sync.Mutex
+}
+
+func NewConn(wsConn *websocket.Conn) net.Conn {
+	return &Conn{
+		Conn: wsConn,
+	}
+}
+
+func (c *Conn) Read(b []byte) (n int, err error) {
+	t, r, err := c.NextReader()
+	if err != nil {
+		return 0, err
+	}
+
+	if t != websocket.BinaryMessage {
+		return 0, fmt.Errorf("unexpected message type")
+	}
+
+	return r.Read(b)
+}
+
+func (c *Conn) Write(b []byte) (int, error) {
+	c.mu.Lock()
+	err := c.WriteMessage(websocket.BinaryMessage, b)
+	c.mu.Unlock()
+	return len(b), err
+}
+
+func (c *Conn) SetDeadline(t time.Time) error {
+	errR := c.SetReadDeadline(t)
+	errW := c.SetWriteDeadline(t)
+
+	if errR != nil {
+		return errR
+	}
+
+	if errW != nil {
+		return errW
+	}
+	return nil
+}
+
+func (c *Conn) Close() error {
+	return c.Conn.Close()
+}

relay/client/dialer/ws/ws.go

@@ -0,0 +1,22 @@
+package ws
+
+import (
+	"fmt"
+	"net"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+func Dial(address string) (net.Conn, error) {
+	addr := fmt.Sprintf("ws://" + address)
+	wsDialer := websocket.Dialer{
+		HandshakeTimeout: 3 * time.Second,
+	}
+	wsConn, _, err := wsDialer.Dial(addr, nil)
+	if err != nil {
+		return nil, err
+	}
+	conn := NewConn(wsConn)
+	return conn, nil
+}

relay/client/dialer/wsnhooyr/client_conn.go

@@ -0,0 +1,79 @@
+package wsnhooyr
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"time"
+
+	"nhooyr.io/websocket"
+)
+
+type Conn struct {
+	*websocket.Conn
+	ctx context.Context
+}
+
+func NewConn(wsConn *websocket.Conn) net.Conn {
+	return &Conn{
+		Conn: wsConn,
+		ctx:  context.Background(),
+	}
+}
+
+func (c *Conn) Read(b []byte) (n int, err error) {
+	t, ioReader, err := c.Conn.Reader(c.ctx)
+	if err != nil {
+		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 len(b), err
+}
+
+func (c *Conn) RemoteAddr() net.Addr {
+	// todo: implement me
+	return nil
+}
+
+func (c *Conn) LocalAddr() net.Addr {
+	// todo: implement me
+	return nil
+}
+
+func (c *Conn) SetReadDeadline(t time.Time) error {
+	// todo: implement me
+	return nil
+}
+
+func (c *Conn) SetWriteDeadline(t time.Time) error {
+	// todo: implement me
+	return nil
+}
+
+func (c *Conn) SetDeadline(t time.Time) error {
+	// todo: implement me
+	errR := c.SetReadDeadline(t)
+	errW := c.SetWriteDeadline(t)
+
+	if errR != nil {
+		return errR
+	}
+
+	if errW != nil {
+		return errW
+	}
+	return nil
+}
+
+func (c *Conn) Close() error {
+	return c.Conn.CloseNow()
+}

relay/client/dialer/wsnhooyr/ws.go

@@ -0,0 +1,22 @@
+package wsnhooyr
+
+import (
+	"context"
+	"fmt"
+	"net"
+
+	"nhooyr.io/websocket"
+)
+
+func Dial(address string) (net.Conn, error) {
+
+	addr := fmt.Sprintf("ws://" + address)
+	wsConn, _, err := websocket.Dial(context.Background(), addr, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	conn := NewConn(wsConn)
+
+	return conn, nil
+}

relay/client/guard.go

@@ -0,0 +1,33 @@
+package client
+
+import (
+	"context"
+	"time"
+)
+
+var (
+	reconnectingTimeout = 5 * time.Second
+)
+
+type Guard struct {
+	ctx         context.Context
+	relayClient *Client
+}
+
+func NewGuard(context context.Context, relayClient *Client) *Guard {
+	g := &Guard{
+		ctx:         context,
+		relayClient: relayClient,
+	}
+
+	return g
+}
+
+func (g *Guard) OnDisconnected() {
+	select {
+	case <-time.After(time.Second):
+		_ = g.relayClient.Connect()
+	case <-g.ctx.Done():
+		return
+	}
+}

relay/client/manager.go

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

relay/client/manager_test.go

@@ -0,0 +1,271 @@
+package client
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/relay/server"
+)
+
+func TestForeignConn(t *testing.T) {
+	ctx := context.Background()
+
+	addr1 := "localhost:1234"
+	srv1 := server.NewServer()
+	go func() {
+		err := srv1.Listen(addr1)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		err := srv1.Close()
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	addr2 := "localhost:2234"
+	srv2 := server.NewServer()
+	go func() {
+		err := srv2.Listen(addr2)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		err := srv2.Close()
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	idAlice := "alice"
+	log.Debugf("connect by alice")
+	mCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	clientAlice := NewManager(mCtx, addr1, idAlice)
+	clientAlice.Serve()
+
+	idBob := "bob"
+	log.Debugf("connect by bob")
+	clientBob := NewManager(mCtx, addr2, idBob)
+	clientBob.Serve()
+
+	bobsSrvAddr, err := clientBob.RelayAddress()
+	if err != nil {
+		t.Fatalf("failed to get relay address: %s", err)
+	}
+	connAliceToBob, err := clientAlice.OpenConn(bobsSrvAddr.String(), idBob)
+	if err != nil {
+		t.Fatalf("failed to bind channel: %s", err)
+	}
+	connBobToAlice, err := clientBob.OpenConn(bobsSrvAddr.String(), 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 TestForeginConnClose(t *testing.T) {
+	ctx := context.Background()
+
+	addr1 := "localhost:1234"
+	srv1 := server.NewServer()
+	go func() {
+		err := srv1.Listen(addr1)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		err := srv1.Close()
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	addr2 := "localhost:2234"
+	srv2 := server.NewServer()
+	go func() {
+		err := srv2.Listen(addr2)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		err := srv2.Close()
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	idAlice := "alice"
+	log.Debugf("connect by alice")
+	mCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	mgr := NewManager(mCtx, addr1, idAlice)
+	mgr.Serve()
+
+	conn, err := mgr.OpenConn(addr2, "anotherpeer")
+	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 TestForeginAutoClose(t *testing.T) {
+	ctx := context.Background()
+	relayCleanupInterval = 1 * time.Second
+	addr1 := "localhost:1234"
+	srv1 := server.NewServer()
+	go func() {
+		t.Log("binding server 1.")
+		err := srv1.Listen(addr1)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		t.Logf("closing server 1.")
+		err := srv1.Close()
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+		t.Logf("server 1. closed")
+	}()
+
+	addr2 := "localhost:2234"
+	srv2 := server.NewServer()
+	go func() {
+		t.Log("binding server 2.")
+		err := srv2.Listen(addr2)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+	defer func() {
+		t.Logf("closing server 2.")
+		err := srv2.Close()
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+		t.Logf("server 2 closed.")
+	}()
+
+	idAlice := "alice"
+	t.Log("connect to server 1.")
+	mCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	mgr := NewManager(mCtx, addr1, idAlice)
+	mgr.Serve()
+
+	t.Log("open connection to another peer")
+	conn, err := mgr.OpenConn(addr2, "anotherpeer")
+	if err != nil {
+		t.Fatalf("failed to bind channel: %s", err)
+	}
+
+	t.Log("close conn")
+	err = conn.Close()
+	if err != nil {
+		t.Fatalf("failed to close connection: %s", err)
+	}
+
+	t.Logf("waiting for relay cleanup: %s", relayCleanupInterval+1*time.Second)
+	time.Sleep(relayCleanupInterval + 1*time.Second)
+	if len(mgr.relayClients) != 0 {
+		t.Errorf("expected 0, got %d", len(mgr.relayClients))
+	}
+
+	t.Logf("closing manager")
+}
+
+func TestAutoReconnect(t *testing.T) {
+	ctx := context.Background()
+	reconnectingTimeout = 2 * time.Second
+
+	addr := "localhost:1234"
+	srv := server.NewServer()
+	go func() {
+		err := srv.Listen(addr)
+		if err != nil {
+			t.Errorf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		err := srv.Close()
+		if err != nil {
+			log.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	mCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	clientAlice := NewManager(mCtx, addr, "alice")
+	clientAlice.Serve()
+	ra, err := clientAlice.RelayAddress()
+	if err != nil {
+		t.Errorf("failed to get relay address: %s", err)
+	}
+	conn, err := clientAlice.OpenConn(ra.String(), "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(ra.String(), "bob")
+	if err != nil {
+		t.Errorf("failed to open channel: %s", err)
+	}
+}

relay/cmd/main.go

@@ -0,0 +1,40 @@
+package main
+
+import (
+	"os"
+	"os/signal"
+	"syscall"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/relay/server"
+	"github.com/netbirdio/netbird/util"
+)
+
+func init() {
+	util.InitLog("trace", "console")
+}
+
+func waitForExitSignal() {
+	osSigs := make(chan os.Signal, 1)
+	signal.Notify(osSigs, syscall.SIGINT, syscall.SIGTERM)
+	_ = <-osSigs
+}
+
+func main() {
+	address := "10.145.236.1:1235"
+	srv := server.NewServer()
+	err := srv.Listen(address)
+	if err != nil {
+		log.Errorf("failed to bind server: %s", err)
+		os.Exit(1)
+	}
+
+	waitForExitSignal()
+
+	err = srv.Close()
+	if err != nil {
+		log.Errorf("failed to close server: %s", err)
+		os.Exit(1)
+	}
+}

relay/messages/id.go

@@ -0,0 +1,20 @@
+package messages
+
+import (
+	"crypto/sha256"
+	"encoding/base64"
+)
+
+const (
+	IDSize = sha256.Size
+)
+
+func HashID(peerID string) ([]byte, string) {
+	idHash := sha256.Sum256([]byte(peerID))
+	idHashString := base64.StdEncoding.EncodeToString(idHash[:])
+	return idHash[:], idHashString
+}
+
+func HashIDToString(idHash []byte) string {
+	return base64.StdEncoding.EncodeToString(idHash[:])
+}

relay/messages/message.go

@@ -0,0 +1,137 @@
+package messages
+
+import (
+	"fmt"
+
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	MsgTypeHello         MsgType = 0
+	MsgTypeHelloResponse MsgType = 1
+	MsgTypeTransport     MsgType = 2
+	MsgClose             MsgType = 3
+)
+
+var (
+	ErrInvalidMessageLength = fmt.Errorf("invalid message length")
+)
+
+type MsgType byte
+
+func (m MsgType) String() string {
+	switch m {
+	case MsgTypeHello:
+		return "hello"
+	case MsgTypeHelloResponse:
+		return "hello response"
+	case MsgTypeTransport:
+		return "transport"
+	case MsgClose:
+		return "close"
+	default:
+		return "unknown"
+	}
+}
+
+func DetermineClientMsgType(msg []byte) (MsgType, error) {
+	// todo: validate magic byte
+	msgType := MsgType(msg[0])
+	switch msgType {
+	case MsgTypeHello:
+		return msgType, nil
+	case MsgTypeTransport:
+		return msgType, nil
+	case MsgClose:
+		return msgType, nil
+	default:
+		return 0, fmt.Errorf("invalid msg type, len: %d", len(msg))
+	}
+}
+
+func DetermineServerMsgType(msg []byte) (MsgType, error) {
+	// todo: validate magic byte
+	msgType := MsgType(msg[0])
+	switch msgType {
+	case MsgTypeHelloResponse:
+		return msgType, nil
+	case MsgTypeTransport:
+		return msgType, nil
+	case MsgClose:
+		return msgType, nil
+	default:
+		return 0, fmt.Errorf("invalid msg type (len: %d)", len(msg))
+	}
+}
+
+// MarshalHelloMsg initial hello message
+func MarshalHelloMsg(peerID []byte) ([]byte, error) {
+	if len(peerID) != IDSize {
+		return nil, fmt.Errorf("invalid peerID length")
+	}
+	msg := make([]byte, 1, 1+len(peerID))
+	msg[0] = byte(MsgTypeHello)
+	msg = append(msg, peerID...)
+	return msg, nil
+}
+
+func UnmarshalHelloMsg(msg []byte) ([]byte, error) {
+	if len(msg) < 2 {
+		return nil, fmt.Errorf("invalid 'hello' messge")
+	}
+	return msg[1:], nil
+}
+
+func MarshalHelloResponse() []byte {
+	msg := make([]byte, 1)
+	msg[0] = byte(MsgTypeHelloResponse)
+	return msg
+}
+
+// Close message
+
+func MarshalCloseMsg() []byte {
+	msg := make([]byte, 1)
+	msg[0] = byte(MsgClose)
+	return msg
+}
+
+// Transport message
+
+func MarshalTransportMsg(peerID []byte, payload []byte) []byte {
+	if len(peerID) != IDSize {
+		return nil
+	}
+
+	msg := make([]byte, 1+IDSize, 1+IDSize+len(payload))
+	msg[0] = byte(MsgTypeTransport)
+	copy(msg[1:], peerID)
+	msg = append(msg, payload...)
+	return msg
+}
+
+func UnmarshalTransportMsg(buf []byte) ([]byte, []byte, error) {
+	headerSize := 1 + IDSize
+	if len(buf) < headerSize {
+		return nil, nil, ErrInvalidMessageLength
+	}
+
+	return buf[1:headerSize], buf[headerSize:], nil
+}
+
+func UnmarshalTransportID(buf []byte) ([]byte, error) {
+	headerSize := 1 + IDSize
+	if len(buf) < headerSize {
+		log.Debugf("invalid message length: %d, expected: %d, %x", len(buf), headerSize, buf)
+		return nil, ErrInvalidMessageLength
+	}
+	return buf[1:headerSize], nil
+}
+
+func UpdateTransportMsg(msg []byte, peerID []byte) error {
+	if len(msg) < 1+len(peerID) {
+		return ErrInvalidMessageLength
+	}
+	copy(msg[1:], peerID)
+	return nil
+}

relay/server/listener/listener.go

@@ -0,0 +1,9 @@
+package listener
+
+import "net"
+
+type Listener interface {
+	Listen(func(conn net.Conn)) error
+	Close() error
+	WaitForExitAcceptedConns()
+}

relay/server/listener/quic/conn.go

@@ -0,0 +1,36 @@
+package quic
+
+import (
+	"net"
+
+	"github.com/quic-go/quic-go"
+	log "github.com/sirupsen/logrus"
+)
+
+type QuicConn struct {
+	quic.Stream
+	qConn quic.Connection
+}
+
+func NewConn(stream quic.Stream, qConn quic.Connection) net.Conn {
+	return &QuicConn{
+		Stream: stream,
+		qConn:  qConn,
+	}
+}
+
+func (q QuicConn) Write(b []byte) (n int, err error) {
+	n, err = q.Stream.Write(b)
+	if n != len(b) {
+		log.Errorf("failed to write out the full message")
+	}
+	return
+}
+
+func (q QuicConn) LocalAddr() net.Addr {
+	return q.qConn.LocalAddr()
+}
+
+func (q QuicConn) RemoteAddr() net.Addr {
+	return q.qConn.RemoteAddr()
+}

relay/server/listener/quic/listener.go

@@ -0,0 +1,111 @@
+package quic
+
+import (
+	"context"
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/tls"
+	"crypto/x509"
+	"encoding/pem"
+	"math/big"
+	"net"
+	"sync"
+
+	"github.com/quic-go/quic-go"
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/relay/server/listener"
+)
+
+type Listener struct {
+	address    string
+	onAcceptFn func(conn net.Conn)
+
+	listener *quic.Listener
+	quit     chan struct{}
+	wg       sync.WaitGroup
+}
+
+func NewListener(address string) listener.Listener {
+	return &Listener{
+		address: address,
+	}
+}
+
+func (l *Listener) Listen(onAcceptFn func(conn net.Conn)) error {
+	ql, err := quic.ListenAddr(l.address, generateTLSConfig(), &quic.Config{
+		EnableDatagrams: true,
+	})
+	if err != nil {
+		return err
+	}
+	l.listener = ql
+	l.quit = make(chan struct{})
+
+	log.Infof("quic server is listening on address: %s", l.address)
+	l.wg.Add(1)
+	go l.acceptLoop(onAcceptFn)
+
+	<-l.quit
+	return nil
+}
+
+func (l *Listener) Close() error {
+	close(l.quit)
+	err := l.listener.Close()
+	l.wg.Wait()
+	return err
+}
+
+func (l *Listener) acceptLoop(acceptFn func(conn net.Conn)) {
+	defer l.wg.Done()
+
+	for {
+		qConn, err := l.listener.Accept(context.Background())
+		if err != nil {
+			select {
+			case <-l.quit:
+				return
+			default:
+				log.Errorf("failed to accept connection: %s", err)
+				continue
+			}
+		}
+
+		log.Infof("new connection from: %s", qConn.RemoteAddr())
+
+		stream, err := qConn.AcceptStream(context.Background())
+		if err != nil {
+			log.Errorf("failed to open stream: %s", err)
+			continue
+		}
+
+		conn := NewConn(stream, qConn)
+
+		go acceptFn(conn)
+	}
+}
+
+// Setup a bare-bones TLS config for the server
+func generateTLSConfig() *tls.Config {
+	key, err := rsa.GenerateKey(rand.Reader, 1024)
+	if err != nil {
+		panic(err)
+	}
+	template := x509.Certificate{SerialNumber: big.NewInt(1)}
+	certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, &key.PublicKey, key)
+	if err != nil {
+		panic(err)
+	}
+	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})
+	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
+
+	tlsCert, err := tls.X509KeyPair(certPEM, keyPEM)
+	if err != nil {
+		panic(err)
+	}
+	return &tls.Config{
+		Certificates: []tls.Certificate{tlsCert},
+		NextProtos:   []string{"quic-echo-example"},
+	}
+}

relay/server/listener/tcp/listener.go

@@ -0,0 +1,80 @@
+package tcp
+
+import (
+	"net"
+	"sync"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/relay/server/listener"
+)
+
+// Listener
+// Is it just demo code. It does not work in real life environment because the TCP is a streaming protocol, adn
+// it does not handle framing.
+type Listener struct {
+	address string
+
+	onAcceptFn func(conn net.Conn)
+	wg         sync.WaitGroup
+	quit       chan struct{}
+	listener   net.Listener
+	lock       sync.Mutex
+}
+
+func NewListener(address string) listener.Listener {
+	return &Listener{
+		address: address,
+	}
+}
+
+func (l *Listener) Listen(onAcceptFn func(conn net.Conn)) error {
+	l.lock.Lock()
+
+	l.onAcceptFn = onAcceptFn
+	l.quit = make(chan struct{})
+
+	li, err := net.Listen("tcp", l.address)
+	if err != nil {
+		log.Errorf("failed to listen on address: %s, %s", l.address, err)
+		l.lock.Unlock()
+		return err
+	}
+	log.Debugf("TCP server is listening on address: %s", l.address)
+	l.listener = li
+	l.wg.Add(1)
+	go l.acceptLoop()
+
+	l.lock.Unlock()
+	<-l.quit
+	return nil
+}
+
+// Close todo: prevent multiple call (do not close two times the channel)
+func (l *Listener) Close() error {
+	l.lock.Lock()
+	defer l.lock.Unlock()
+
+	close(l.quit)
+	err := l.listener.Close()
+	l.wg.Wait()
+	return err
+}
+
+func (l *Listener) acceptLoop() {
+	defer l.wg.Done()
+
+	for {
+		conn, err := l.listener.Accept()
+		if err != nil {
+			select {
+			case <-l.quit:
+				return
+			default:
+				log.Errorf("failed to accept connection: %s", err)
+				continue
+			}
+		}
+		go l.onAcceptFn(conn)
+	}
+}

relay/server/listener/udp/conn.go

@@ -0,0 +1,68 @@
+package udp
+
+import (
+	"io"
+	"net"
+	"time"
+)
+
+type UDPConn struct {
+	*net.UDPConn
+	addr       *net.UDPAddr
+	msgChannel chan []byte
+}
+
+func NewConn(conn *net.UDPConn, addr *net.UDPAddr) *UDPConn {
+	return &UDPConn{
+		UDPConn:    conn,
+		addr:       addr,
+		msgChannel: make(chan []byte),
+	}
+}
+
+func (u *UDPConn) Read(b []byte) (n int, err error) {
+	msg, ok := <-u.msgChannel
+	if !ok {
+		return 0, io.EOF
+	}
+
+	n = copy(b, msg)
+	return n, nil
+}
+
+func (u *UDPConn) Write(b []byte) (n int, err error) {
+	return u.UDPConn.WriteTo(b, u.addr)
+}
+
+func (u *UDPConn) Close() error {
+	//TODO implement me
+	//panic("implement me")
+	return nil
+}
+
+func (u *UDPConn) LocalAddr() net.Addr {
+	return u.UDPConn.LocalAddr()
+}
+
+func (u *UDPConn) RemoteAddr() net.Addr {
+	return u.addr
+}
+
+func (u *UDPConn) SetDeadline(t time.Time) error {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (u *UDPConn) SetReadDeadline(t time.Time) error {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (u *UDPConn) SetWriteDeadline(t time.Time) error {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (u *UDPConn) onNewMsg(b []byte) {
+	u.msgChannel <- b
+}

relay/server/listener/udp/listener.go

@@ -0,0 +1,109 @@
+package udp
+
+import (
+	"net"
+	"sync"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/relay/server/listener"
+)
+
+type Listener struct {
+	address    string
+	conns      map[string]*UDPConn
+	onAcceptFn func(conn net.Conn)
+
+	listener *net.UDPConn
+
+	wg   sync.WaitGroup
+	quit chan struct{}
+	lock sync.Mutex
+}
+
+func (l *Listener) WaitForExitAcceptedConns() {
+	l.wg.Wait()
+	return
+}
+
+func NewListener(address string) listener.Listener {
+	return &Listener{
+		address: address,
+		conns:   make(map[string]*UDPConn),
+	}
+}
+
+func (l *Listener) Listen(onAcceptFn func(conn net.Conn)) error {
+	l.lock.Lock()
+
+	l.onAcceptFn = onAcceptFn
+	l.quit = make(chan struct{})
+
+	addr, err := net.ResolveUDPAddr("udp", l.address)
+	if err != nil {
+		log.Errorf("invalid listen address '%s': %s", l.address, err)
+		l.lock.Unlock()
+		return err
+	}
+
+	li, err := net.ListenUDP("udp", addr)
+	if err != nil {
+		log.Fatalf("%s", err)
+		l.lock.Unlock()
+		return err
+	}
+	log.Debugf("udp server is listening on address: %s", addr.String())
+	l.listener = li
+	l.wg.Add(1)
+	go l.readLoop()
+
+	l.lock.Unlock()
+	<-l.quit
+	return nil
+}
+
+func (l *Listener) Close() error {
+	l.lock.Lock()
+	defer l.lock.Unlock()
+
+	if l.listener == nil {
+		return nil
+	}
+
+	log.Infof("closing UDP listener")
+	close(l.quit)
+	err := l.listener.Close()
+	l.wg.Wait()
+	l.listener = nil
+	return err
+}
+
+func (l *Listener) readLoop() {
+	defer l.wg.Done()
+
+	for {
+		buf := make([]byte, 1500)
+		n, addr, err := l.listener.ReadFromUDP(buf)
+		if err != nil {
+			select {
+			case <-l.quit:
+				return
+			default:
+				log.Errorf("failed to accept connection: %s", err)
+				continue
+			}
+		}
+
+		pConn, ok := l.conns[addr.String()]
+		if ok {
+			pConn.onNewMsg(buf[:n])
+			continue
+		}
+
+		pConn = NewConn(l.listener, addr)
+		log.Infof("new connection from: %s", pConn.RemoteAddr())
+		l.conns[addr.String()] = pConn
+		go l.onAcceptFn(pConn)
+		pConn.onNewMsg(buf[:n])
+	}
+}

relay/server/listener/ws/conn.go

@@ -0,0 +1,74 @@
+package ws
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+	log "github.com/sirupsen/logrus"
+)
+
+type Conn struct {
+	*websocket.Conn
+
+	mu sync.Mutex
+}
+
+func NewConn(wsConn *websocket.Conn) *Conn {
+	return &Conn{
+		Conn: wsConn,
+	}
+}
+
+func (c *Conn) Read(b []byte) (n int, err error) {
+	t, r, err := c.NextReader()
+	if err != nil {
+		return 0, ioErrHandling(err)
+	}
+
+	if t != websocket.BinaryMessage {
+		log.Errorf("unexpected message type: %d", t)
+		return 0, fmt.Errorf("unexpected message type")
+	}
+
+	n, err = r.Read(b)
+	if err != nil {
+		return 0, ioErrHandling(err)
+	}
+	return n, err
+}
+
+func (c *Conn) Write(b []byte) (int, error) {
+	c.mu.Lock()
+	err := c.WriteMessage(websocket.BinaryMessage, b)
+	c.mu.Unlock()
+	return len(b), err
+}
+
+func (c *Conn) SetDeadline(t time.Time) error {
+	errR := c.SetReadDeadline(t)
+	errW := c.SetWriteDeadline(t)
+
+	if errR != nil {
+		return errR
+	}
+
+	if errW != nil {
+		return errW
+	}
+	return nil
+}
+
+func ioErrHandling(err error) error {
+	var wErr *websocket.CloseError
+	if !errors.As(err, &wErr) {
+		return err
+	}
+	if wErr.Code == websocket.CloseNormalClosure {
+		return io.EOF
+	}
+	return err
+}

relay/server/listener/ws/listener.go

@@ -0,0 +1,92 @@
+package ws
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/relay/server/listener"
+)
+
+var (
+	upgrader = websocket.Upgrader{} // use default options
+)
+
+type Listener struct {
+	address string
+
+	wg       sync.WaitGroup
+	server   *http.Server
+	acceptFn func(conn net.Conn)
+}
+
+func NewListener(address string) listener.Listener {
+	return &Listener{
+		address: address,
+	}
+}
+
+func (l *Listener) Listen(acceptFn func(conn net.Conn)) error {
+	if l.server != nil {
+		return errors.New("server is already running")
+	}
+
+	l.acceptFn = acceptFn
+	mux := http.NewServeMux()
+	mux.HandleFunc("/", l.onAccept)
+
+	l.server = &http.Server{
+		Addr:    l.address,
+		Handler: mux,
+	}
+
+	log.Infof("WS server is listening on address: %s", l.address)
+	err := l.server.ListenAndServe()
+	if errors.Is(err, http.ErrServerClosed) {
+		return nil
+	}
+	return err
+}
+
+func (l *Listener) Close() error {
+	if l.server == nil {
+		return nil
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	log.Debugf("closing WS server")
+	if err := l.server.Shutdown(ctx); err != nil {
+		return fmt.Errorf("server shutdown failed: %v", err)
+	}
+
+	l.wg.Wait()
+	return nil
+}
+
+func (l *Listener) WaitForExitAcceptedConns() {
+	l.wg.Wait()
+}
+
+func (l *Listener) onAccept(writer http.ResponseWriter, request *http.Request) {
+	l.wg.Add(1)
+	defer l.wg.Done()
+
+	wsConn, err := upgrader.Upgrade(writer, request, nil)
+	if err != nil {
+		log.Errorf("failed to upgrade connection: %s", err)
+		return
+	}
+	conn := NewConn(wsConn)
+	log.Infof("new connection from: %s", conn.RemoteAddr())
+	l.acceptFn(conn)
+	return
+}

relay/server/listener/wsnhooyr/conn.go

@@ -0,0 +1,107 @@
+package ws
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"nhooyr.io/websocket"
+)
+
+type Conn struct {
+	*websocket.Conn
+	lAddr *net.TCPAddr
+	rAddr *net.TCPAddr
+
+	closed   bool
+	closedMu sync.Mutex
+	ctx      context.Context
+}
+
+func NewConn(wsConn *websocket.Conn, lAddr, rAddr *net.TCPAddr) *Conn {
+	return &Conn{
+		Conn:  wsConn,
+		lAddr: lAddr,
+		rAddr: rAddr,
+		ctx:   context.Background(),
+	}
+}
+
+func (c *Conn) Read(b []byte) (n int, err error) {
+	t, r, err := c.Reader(c.ctx)
+	if err != nil {
+		return 0, c.ioErrHandling(err)
+	}
+
+	if t != websocket.MessageBinary {
+		log.Errorf("unexpected message type: %d", t)
+		return 0, fmt.Errorf("unexpected message type")
+	}
+
+	n, err = r.Read(b)
+	if err != nil {
+		return 0, c.ioErrHandling(err)
+	}
+	return n, err
+}
+
+func (c *Conn) Write(b []byte) (int, error) {
+	err := c.Conn.Write(c.ctx, websocket.MessageBinary, b)
+	return len(b), err
+}
+
+func (c *Conn) LocalAddr() net.Addr {
+	return c.lAddr
+}
+
+func (c *Conn) RemoteAddr() net.Addr {
+	return c.rAddr
+}
+
+func (c *Conn) SetReadDeadline(t time.Time) error {
+	// todo: implement me
+	return nil
+}
+
+func (c *Conn) SetWriteDeadline(t time.Time) error {
+	// todo: implement me
+	return nil
+}
+
+func (c *Conn) SetDeadline(t time.Time) error {
+	// todo: implement me
+	return nil
+}
+
+func (c *Conn) Close() error {
+	c.closedMu.Lock()
+	c.closed = true
+	c.closedMu.Unlock()
+	return c.Conn.CloseNow()
+}
+
+func (c *Conn) isClosed() bool {
+	c.closedMu.Lock()
+	defer c.closedMu.Unlock()
+	return c.closed
+}
+
+func (c *Conn) ioErrHandling(err error) error {
+	if c.isClosed() {
+		return io.EOF
+	}
+
+	var wErr *websocket.CloseError
+	if !errors.As(err, &wErr) {
+		return err
+	}
+	if wErr.Code == websocket.StatusNormalClosure {
+		return io.EOF
+	}
+	return err
+}

relay/server/listener/wsnhooyr/listener.go

@@ -0,0 +1,96 @@
+package ws
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"net/http"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"nhooyr.io/websocket"
+
+	"github.com/netbirdio/netbird/relay/server/listener"
+)
+
+type Listener struct {
+	address string
+
+	wg       sync.WaitGroup
+	server   *http.Server
+	acceptFn func(conn net.Conn)
+}
+
+func NewListener(address string) listener.Listener {
+	return &Listener{
+		address: address,
+	}
+}
+
+// Listen todo: prevent multiple call
+func (l *Listener) Listen(acceptFn func(conn net.Conn)) error {
+	l.acceptFn = acceptFn
+	mux := http.NewServeMux()
+	mux.HandleFunc("/", l.onAccept)
+
+	l.server = &http.Server{
+		Addr:    l.address,
+		Handler: mux,
+	}
+
+	log.Infof("WS server is listening on address: %s", l.address)
+	err := l.server.ListenAndServe()
+	if errors.Is(err, http.ErrServerClosed) {
+		return nil
+	}
+	return err
+}
+
+func (l *Listener) Close() error {
+	if l.server == nil {
+		return nil
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	log.Infof("stop WS listener")
+	if err := l.server.Shutdown(ctx); err != nil {
+		return fmt.Errorf("server shutdown failed: %v", err)
+	}
+	log.Infof("WS listener stopped")
+	return nil
+}
+
+func (l *Listener) WaitForExitAcceptedConns() {
+	l.wg.Wait()
+}
+
+func (l *Listener) onAccept(w http.ResponseWriter, r *http.Request) {
+	l.wg.Add(1)
+	defer l.wg.Done()
+
+	wsConn, err := websocket.Accept(w, r, nil)
+	if err != nil {
+		log.Errorf("failed to accept ws connection: %s", err)
+		return
+	}
+
+	rAddr, err := net.ResolveTCPAddr("tcp", r.RemoteAddr)
+	if err != nil {
+		_ = wsConn.Close(websocket.StatusInternalError, "internal error")
+		return
+	}
+
+	lAddr, err := net.ResolveTCPAddr("tcp", l.server.Addr)
+	if err != nil {
+		_ = wsConn.Close(websocket.StatusInternalError, "internal error")
+		return
+	}
+
+	conn := NewConn(wsConn, lAddr, rAddr)
+	l.acceptFn(conn)
+	return
+}

relay/server/peer.go

@@ -0,0 +1,33 @@
+package server
+
+import (
+	"net"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/relay/messages"
+)
+
+type Peer struct {
+	Log  *log.Entry
+	idS  string
+	idB  []byte
+	conn net.Conn
+}
+
+func NewPeer(id []byte, conn net.Conn) *Peer {
+	stringID := messages.HashIDToString(id)
+	return &Peer{
+		Log:  log.WithField("peer_id", stringID),
+		idB:  id,
+		idS:  stringID,
+		conn: conn,
+	}
+}
+func (p *Peer) ID() []byte {
+	return p.idB
+}
+
+func (p *Peer) String() string {
+	return p.idS
+}

relay/server/server.go

@@ -0,0 +1,192 @@
+package server
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"sync"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/relay/messages"
+	"github.com/netbirdio/netbird/relay/server/listener"
+	"github.com/netbirdio/netbird/relay/server/listener/udp"
+	ws "github.com/netbirdio/netbird/relay/server/listener/wsnhooyr"
+)
+
+type Server struct {
+	store   *Store
+	storeMu sync.RWMutex
+
+	UDPListener listener.Listener
+	WSListener  listener.Listener
+}
+
+func NewServer() *Server {
+	return &Server{
+		store:   NewStore(),
+		storeMu: sync.RWMutex{},
+	}
+}
+
+func (r *Server) Listen(address string) error {
+	wg := sync.WaitGroup{}
+	wg.Add(2)
+
+	r.WSListener = ws.NewListener(address)
+	var wslErr error
+	go func() {
+		defer wg.Done()
+		wslErr = r.WSListener.Listen(r.accept)
+		if wslErr != nil {
+			log.Errorf("failed to bind ws server: %s", wslErr)
+		}
+	}()
+
+	r.UDPListener = udp.NewListener(address)
+	var udpLErr error
+	go func() {
+		defer wg.Done()
+		udpLErr = r.UDPListener.Listen(r.accept)
+		if udpLErr != nil {
+			log.Errorf("failed to bind ws server: %s", udpLErr)
+		}
+	}()
+
+	err := errors.Join(wslErr, udpLErr)
+	return err
+}
+
+func (r *Server) Close() error {
+	var wErr error
+	if r.WSListener != nil {
+		wErr = r.WSListener.Close()
+	}
+
+	var uErr error
+	if r.UDPListener != nil {
+		uErr = r.UDPListener.Close()
+	}
+
+	r.sendCloseMsgs()
+
+	r.WSListener.WaitForExitAcceptedConns()
+
+	err := errors.Join(wErr, uErr)
+	return err
+}
+
+func (r *Server) accept(conn net.Conn) {
+	peer, err := handShake(conn)
+	if err != nil {
+		log.Errorf("failed to handshake wiht %s: %s", conn.RemoteAddr(), err)
+		cErr := conn.Close()
+		if cErr != nil {
+			log.Errorf("failed to close connection, %s: %s", conn.RemoteAddr(), cErr)
+		}
+		return
+	}
+	peer.Log.Infof("peer connected from: %s", conn.RemoteAddr())
+
+	r.store.AddPeer(peer)
+	defer func() {
+		r.store.DeletePeer(peer)
+		peer.Log.Infof("relay connection closed")
+	}()
+
+	for {
+		buf := make([]byte, 1500) // todo: optimize buffer size
+		n, err := conn.Read(buf)
+		if err != nil {
+			if err != io.EOF {
+				peer.Log.Errorf("failed to read message: %s", err)
+			}
+			return
+		}
+
+		msgType, err := messages.DetermineClientMsgType(buf[:n])
+		if err != nil {
+			peer.Log.Errorf("failed to determine message type: %s", err)
+			return
+		}
+		switch msgType {
+		case messages.MsgTypeTransport:
+			msg := buf[:n]
+			peerID, err := messages.UnmarshalTransportID(msg)
+			if err != nil {
+				peer.Log.Errorf("failed to unmarshal transport message: %s", err)
+				continue
+			}
+			go func() {
+				stringPeerID := messages.HashIDToString(peerID)
+				dp, ok := r.store.Peer(stringPeerID)
+				if !ok {
+					peer.Log.Errorf("peer not found: %s", stringPeerID)
+					return
+				}
+				err := messages.UpdateTransportMsg(msg, peer.ID())
+				if err != nil {
+					peer.Log.Errorf("failed to update transport message: %s", err)
+					return
+				}
+				_, err = dp.conn.Write(msg)
+				if err != nil {
+					peer.Log.Errorf("failed to write transport message to: %s", dp.String())
+				}
+				return
+			}()
+		case messages.MsgClose:
+			peer.Log.Infof("peer disconnected gracefully")
+			_ = conn.Close()
+			return
+		}
+	}
+}
+
+func (r *Server) sendCloseMsgs() {
+	msg := messages.MarshalCloseMsg()
+
+	r.storeMu.Lock()
+	log.Debugf("sending close messages to %d peers", len(r.store.peers))
+	for _, p := range r.store.peers {
+		_, err := p.conn.Write(msg)
+		if err != nil {
+			log.Errorf("failed to send close message to peer: %s", p.String())
+		}
+
+		err = p.conn.Close()
+		if err != nil {
+			log.Errorf("failed to close connection to peer: %s", err)
+		}
+	}
+	r.storeMu.Unlock()
+}
+
+func handShake(conn net.Conn) (*Peer, error) {
+	buf := make([]byte, 1500)
+	n, err := conn.Read(buf)
+	if err != nil {
+		log.Errorf("failed to read message: %s", err)
+		return nil, err
+	}
+	msgType, err := messages.DetermineClientMsgType(buf[:n])
+	if err != nil {
+		return nil, err
+	}
+	if msgType != messages.MsgTypeHello {
+		tErr := fmt.Errorf("invalid message type")
+		log.Errorf("failed to handshake: %s", tErr)
+		return nil, tErr
+	}
+	peerId, err := messages.UnmarshalHelloMsg(buf[:n])
+	if err != nil {
+		log.Errorf("failed to handshake: %s", err)
+		return nil, err
+	}
+	p := NewPeer(peerId, conn)
+
+	msg := messages.MarshalHelloResponse()
+	_, err = conn.Write(msg)
+	return p, err
+}

relay/server/store.go

@@ -0,0 +1,47 @@
+package server
+
+import (
+	"sync"
+)
+
+type Store struct {
+	peers     map[string]*Peer // consider to use [32]byte as key. The Peer(id string) would be faster
+	peersLock sync.RWMutex
+}
+
+func NewStore() *Store {
+	return &Store{
+		peers: make(map[string]*Peer),
+	}
+}
+
+func (s *Store) AddPeer(peer *Peer) {
+	s.peersLock.Lock()
+	defer s.peersLock.Unlock()
+	s.peers[peer.String()] = peer
+}
+
+func (s *Store) DeletePeer(peer *Peer) {
+	s.peersLock.Lock()
+	defer s.peersLock.Unlock()
+	delete(s.peers, peer.String())
+}
+
+func (s *Store) Peer(id string) (*Peer, bool) {
+	s.peersLock.RLock()
+	defer s.peersLock.RUnlock()
+
+	p, ok := s.peers[id]
+	return p, ok
+}
+
+func (s *Store) Peers() []*Peer {
+	s.peersLock.RLock()
+	defer s.peersLock.RUnlock()
+
+	peers := make([]*Peer, 0, len(s.peers))
+	for _, p := range s.peers {
+		peers = append(peers, p)
+	}
+	return peers
+}

relay/test/manager_test.go

@@ -0,0 +1,57 @@
+package test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/netbirdio/netbird/relay/client"
+	"github.com/netbirdio/netbird/relay/server"
+)
+
+func TestManager(t *testing.T) {
+	addr := "localhost:1239"
+
+	srv := server.NewServer()
+	go func() {
+		err := srv.Listen(addr)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		err := srv.Close()
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	cm := client.NewManager(ctx, addr, "me")
+	cm.Serve()
+
+	// wait for the relay handshake to complete
+	time.Sleep(1 * time.Second)
+	conn, err := cm.OpenConn("remotepeer")
+	if err != nil {
+		t.Errorf("failed to open connection: %s", err)
+	}
+
+	readCtx, readCancel := context.WithCancel(context.Background())
+	defer readCancel()
+	go func() {
+		_, _ = conn.Read(make([]byte, 1))
+		readCancel()
+	}()
+
+	cancel()
+
+	select {
+	case <-time.After(2 * time.Second):
+		t.Errorf("client peer conn did not close automatically")
+	case <-readCtx.Done():
+		// conn exited well
+	}
+}

signal/client/client.go

@@ -25,11 +25,6 @@ const (
 	DirectCheck uint32 = 1
 )
 
-// FeaturesSupport register protocol supported features
-type FeaturesSupport struct {
-	DirectCheck bool
-}
-
 type Client interface {
 	io.Closer
 	StreamConnected() bool
@@ -79,15 +74,3 @@ type Credential struct {
 	UFrag string
 	Pwd   string
 }
-
-// ParseFeaturesSupported parses a slice of supported features into FeaturesSupport
-func ParseFeaturesSupported(featuresMessage []uint32) FeaturesSupport {
-	var protoSupport FeaturesSupport
-	for _, feature := range featuresMessage {
-		if feature == DirectCheck {
-			protoSupport.DirectCheck = true
-			return protoSupport
-		}
-	}
-	return protoSupport
-}

signal/client/client_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"net"
 	"sync"
-	"testing"
 	"time"
 
 	. "github.com/onsi/ginkgo"
@@ -168,41 +167,6 @@ var _ = Describe("GrpcClient", func() {
 
 })
 
-func TestParseFeaturesSupported(t *testing.T) {
-	expectedOnEmptyOrUnsupported := FeaturesSupport{DirectCheck: false}
-	expectedWithDirectCheck := FeaturesSupport{DirectCheck: true}
-	testCases := []struct {
-		name     string
-		input    []uint32
-		expected FeaturesSupport
-	}{
-		{
-			name:     "Should Return DirectCheck Supported",
-			input:    []uint32{DirectCheck},
-			expected: expectedWithDirectCheck,
-		},
-		{
-			name:     "Should Return DirectCheck Unsupported When Nil",
-			input:    nil,
-			expected: expectedOnEmptyOrUnsupported,
-		},
-		{
-			name:     "Should Return DirectCheck Unsupported When Not Known Feature",
-			input:    []uint32{9999},
-			expected: expectedOnEmptyOrUnsupported,
-		},
-	}
-
-	for _, testCase := range testCases {
-		t.Run(testCase.name, func(t *testing.T) {
-			result := ParseFeaturesSupported(testCase.input)
-			if result.DirectCheck != testCase.expected.DirectCheck {
-				t.Errorf("Direct check feature should match: Expected: %t, Got: %t", testCase.expected.DirectCheck, result.DirectCheck)
-			}
-		})
-	}
-}
-
 func createSignalClient(addr string, key wgtypes.Key) *GrpcClient {
 	var sigTLSEnabled = false
 	client, err := NewClient(context.Background(), addr, key, sigTLSEnabled)