Add random routes for the test

client/firewall/iptables/manager_linux.go

@@ -83,11 +83,9 @@ func (m *Manager) Init(stateManager *statemanager.Manager) error {
 	}
 
 	// persist early to ensure cleanup of chains
-	go func() {
-		if err := stateManager.PersistState(context.Background()); err != nil {
-			log.Errorf("failed to persist state: %v", err)
-		}
-	}()
+	if err := stateManager.PersistState(context.Background()); err != nil {
+		log.Errorf("failed to persist state: %v", err)
+	}
 
 	return nil
 }

client/firewall/nftables/manager_linux.go

@@ -99,11 +99,9 @@ func (m *Manager) Init(stateManager *statemanager.Manager) error {
 	}
 
 	// persist early
-	go func() {
-		if err := stateManager.PersistState(context.Background()); err != nil {
-			log.Errorf("failed to persist state: %v", err)
-		}
-	}()
+	if err := stateManager.PersistState(context.Background()); err != nil {
+		log.Errorf("failed to persist state: %v", err)
+	}
 
 	return nil
 }

client/internal/config.go

@@ -164,7 +164,7 @@ func UpdateOrCreateConfig(input ConfigInput) (*Config, error) {
 		if err != nil {
 			return nil, err
 		}
-		err = util.WriteJsonWithRestrictedPermission(context.Background(), input.ConfigPath, cfg)
+		err = util.WriteJsonWithRestrictedPermission(input.ConfigPath, cfg)
 		return cfg, err
 	}
 
@@ -185,7 +185,7 @@ func CreateInMemoryConfig(input ConfigInput) (*Config, error) {
 
 // WriteOutConfig write put the prepared config to the given path
 func WriteOutConfig(path string, config *Config) error {
-	return util.WriteJson(context.Background(), path, config)
+	return util.WriteJson(path, config)
 }
 
 // createNewConfig creates a new config generating a new Wireguard key and saving to file
@@ -215,7 +215,7 @@ func update(input ConfigInput) (*Config, error) {
 	}
 
 	if updated {
-		if err := util.WriteJson(context.Background(), input.ConfigPath, config); err != nil {
+		if err := util.WriteJson(input.ConfigPath, config); err != nil {
 			return nil, err
 		}
 	}

client/internal/connect.go

@@ -230,8 +230,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, probes *ProbeHold
 		c.statusRecorder.MarkSignalConnected()
 
 		relayURLs, token := parseRelayInfo(loginResp)
-		relayManager := relayClient.NewManager(engineCtx, []string{"rels://temp-relay-quic.relay.netbird.io:443"}, myPrivateKey.PublicKey().String())
-		c.statusRecorder.SetRelayMgr(relayManager)
+		relayManager := relayClient.NewManager(engineCtx, relayURLs, myPrivateKey.PublicKey().String())
 		if len(relayURLs) > 0 {
 			if token != nil {
 				if err := relayManager.UpdateToken(token); err != nil {
@@ -242,7 +241,9 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, probes *ProbeHold
 			log.Infof("connecting to the Relay service(s): %s", strings.Join(relayURLs, ", "))
 			if err = relayManager.Serve(); err != nil {
 				log.Error(err)
+				return wrapErr(err)
 			}
+			c.statusRecorder.SetRelayMgr(relayManager)
 		}
 
 		peerConfig := loginResp.GetPeerConfig()

client/internal/dns/server.go

@@ -326,13 +326,9 @@ func (s *DefaultServer) applyConfiguration(update nbdns.Config) error {
 	// persist dns state right away
 	ctx, cancel := context.WithTimeout(s.ctx, 3*time.Second)
 	defer cancel()
-
-	// don't block
-	go func() {
-		if err := s.stateManager.PersistState(ctx); err != nil {
-			log.Errorf("Failed to persist dns state: %v", err)
-		}
-	}()
+	if err := s.stateManager.PersistState(ctx); err != nil {
+		log.Errorf("Failed to persist dns state: %v", err)
+	}
 
 	if s.searchDomainNotifier != nil {
 		s.searchDomainNotifier.onNewSearchDomains(s.SearchDomains())

client/internal/engine.go

@@ -38,6 +38,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/routemanager/systemops"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 
+
 	nbssh "github.com/netbirdio/netbird/client/ssh"
 	"github.com/netbirdio/netbird/client/system"
 	nbdns "github.com/netbirdio/netbird/dns"
@@ -170,7 +171,7 @@ type Engine struct {
 
 	relayManager *relayClient.Manager
 	stateManager *statemanager.Manager
-	srWatcher    *guard.SRWatcher
+	srWatcher *guard.SRWatcher
 }
 
 // Peer is an instance of the Connection Peer
@@ -537,7 +538,6 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 
 		relayMsg := wCfg.GetRelay()
 		if relayMsg != nil {
-			// when we receive token we expect valid address list too
 			c := &auth.Token{
 				Payload:   relayMsg.GetTokenPayload(),
 				Signature: relayMsg.GetTokenSignature(),
@@ -546,16 +546,9 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 				log.Errorf("failed to update relay token: %v", err)
 				return fmt.Errorf("update relay token: %w", err)
 			}
-
-			e.relayManager.UpdateServerURLs([]string{"rels://temp-relay-quic.relay.netbird.io:443"})
-
-			// Just in case the agent started with an MGM server where the relay was disabled but was later enabled.
-			// We can ignore all errors because the guard will manage the reconnection retries.
-			_ = e.relayManager.Serve()
-		} else {
-			e.relayManager.UpdateServerURLs(nil)
 		}
 
+		// todo update relay address in the relay manager
 		// todo update signal
 	}
 
@@ -648,10 +641,6 @@ func (e *Engine) updateSSH(sshConf *mgmProto.SSHConfig) error {
 }
 
 func (e *Engine) updateConfig(conf *mgmProto.PeerConfig) error {
-	if e.wgInterface == nil {
-		return errors.New("wireguard interface is not initialized")
-	}
-
 	if e.wgInterface.Address().String() != conf.Address {
 		oldAddr := e.wgInterface.Address().String()
 		log.Debugf("updating peer address from %s to %s", oldAddr, conf.Address)

client/internal/peer/status.go

@@ -676,23 +676,25 @@ func (d *Status) GetRelayStates() []relay.ProbeResult {
 	// extend the list of stun, turn servers with relay address
 	relayStates := slices.Clone(d.relayStates)
 
+	var relayState relay.ProbeResult
+
 	// if the server connection is not established then we will use the general address
 	// in case of connection we will use the instance specific address
 	instanceAddr, err := d.relayMgr.RelayInstanceAddress()
 	if err != nil {
 		// TODO add their status
-		for _, r := range d.relayMgr.ServerURLs() {
-			relayStates = append(relayStates, relay.ProbeResult{
-				URI: r,
-				Err: err,
-			})
+		if errors.Is(err, relayClient.ErrRelayClientNotConnected) {
+			for _, r := range d.relayMgr.ServerURLs() {
+				relayStates = append(relayStates, relay.ProbeResult{
+					URI: r,
+				})
+			}
+			return relayStates
 		}
-		return relayStates
+		relayState.Err = err
 	}
 
-	relayState := relay.ProbeResult{
-		URI: instanceAddr,
-	}
+	relayState.URI = instanceAddr
 	return append(relayStates, relayState)
 }
 

client/internal/peer/worker_ice.go

@@ -46,6 +46,8 @@ type WorkerICE struct {
 	hasRelayOnLocally bool
 	conn              WorkerICECallbacks
 
+	selectedPriority ConnPriority
+
 	agent    *ice.Agent
 	muxAgent sync.Mutex
 
@@ -93,8 +95,10 @@ func (w *WorkerICE) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 
 	var preferredCandidateTypes []ice.CandidateType
 	if w.hasRelayOnLocally && remoteOfferAnswer.RelaySrvAddress != "" {
+		w.selectedPriority = connPriorityICEP2P
 		preferredCandidateTypes = icemaker.CandidateTypesP2P()
 	} else {
+		w.selectedPriority = connPriorityICETurn
 		preferredCandidateTypes = icemaker.CandidateTypes()
 	}
 
@@ -155,7 +159,7 @@ func (w *WorkerICE) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 		RelayedOnLocal:             isRelayCandidate(pair.Local),
 	}
 	w.log.Debugf("on ICE conn read to use ready")
-	go w.conn.OnConnReady(selectedPriority(pair), ci)
+	go w.conn.OnConnReady(w.selectedPriority, ci)
 }
 
 // OnRemoteCandidate Handles ICE connection Candidate provided by the remote peer.
@@ -390,11 +394,3 @@ func isRelayed(pair *ice.CandidatePair) bool {
 	}
 	return false
 }
-
-func selectedPriority(pair *ice.CandidatePair) ConnPriority {
-	if isRelayed(pair) {
-		return connPriorityICETurn
-	} else {
-		return connPriorityICEP2P
-	}
-}

client/internal/routemanager/client_test.go

@@ -228,64 +228,6 @@ func TestGetBestrouteFromStatuses(t *testing.T) {
 			currentRoute:    "route1",
 			expectedRouteID: "route1",
 		},
-		{
-			name: "relayed routes with latency 0 should maintain previous choice",
-			statuses: map[route.ID]routerPeerStatus{
-				"route1": {
-					connected: true,
-					relayed:   true,
-					latency:   0 * time.Millisecond,
-				},
-				"route2": {
-					connected: true,
-					relayed:   true,
-					latency:   0 * time.Millisecond,
-				},
-			},
-			existingRoutes: map[route.ID]*route.Route{
-				"route1": {
-					ID:     "route1",
-					Metric: route.MaxMetric,
-					Peer:   "peer1",
-				},
-				"route2": {
-					ID:     "route2",
-					Metric: route.MaxMetric,
-					Peer:   "peer2",
-				},
-			},
-			currentRoute:    "route1",
-			expectedRouteID: "route1",
-		},
-		{
-			name: "p2p routes with latency 0 should maintain previous choice",
-			statuses: map[route.ID]routerPeerStatus{
-				"route1": {
-					connected: true,
-					relayed:   false,
-					latency:   0 * time.Millisecond,
-				},
-				"route2": {
-					connected: true,
-					relayed:   false,
-					latency:   0 * time.Millisecond,
-				},
-			},
-			existingRoutes: map[route.ID]*route.Route{
-				"route1": {
-					ID:     "route1",
-					Metric: route.MaxMetric,
-					Peer:   "peer1",
-				},
-				"route2": {
-					ID:     "route2",
-					Metric: route.MaxMetric,
-					Peer:   "peer2",
-				},
-			},
-			currentRoute:    "route1",
-			expectedRouteID: "route1",
-		},
 		{
 			name: "current route with bad score should be changed to route with better score",
 			statuses: map[route.ID]routerPeerStatus{
@@ -352,7 +294,7 @@ func TestGetBestrouteFromStatuses(t *testing.T) {
 			dummyRoute := &route.Route{
 				ID:     route.ID(fmt.Sprintf("dummy_p1_%d", i)),
 				Metric: route.MinMetric,
-				Peer:   fmt.Sprintf("dummy_p1_%d", i),
+				Peer:   "peer1",
 			}
 			tc.existingRoutes[dummyRoute.ID] = dummyRoute
 		}
@@ -360,29 +302,10 @@ func TestGetBestrouteFromStatuses(t *testing.T) {
 			dummyRoute := &route.Route{
 				ID:     route.ID(fmt.Sprintf("dummy_p2_%d", i)),
 				Metric: route.MinMetric,
-				Peer:   fmt.Sprintf("dummy_p1_%d", i),
+				Peer:   "peer2",
 			}
 			tc.existingRoutes[dummyRoute.ID] = dummyRoute
 		}
-
-		for i := 0; i < 50; i++ {
-			id := route.ID(fmt.Sprintf("dummy_p1_%d", i))
-			dummyStatus := routerPeerStatus{
-				connected: false,
-				relayed:   true,
-				latency:   0,
-			}
-			tc.statuses[id] = dummyStatus
-		}
-		for i := 0; i < 50; i++ {
-			id := route.ID(fmt.Sprintf("dummy_p2_%d", i))
-			dummyStatus := routerPeerStatus{
-				connected: false,
-				relayed:   true,
-				latency:   0,
-			}
-			tc.statuses[id] = dummyStatus
-		}
 	}
 
 	for _, tc := range testCases {

client/internal/statemanager/manager.go

@@ -16,7 +16,6 @@ import (
 	"golang.org/x/exp/maps"
 
 	nberrors "github.com/netbirdio/netbird/client/errors"
-	"github.com/netbirdio/netbird/util"
 )
 
 // State interface defines the methods that all state types must implement
@@ -179,14 +178,25 @@ func (m *Manager) PersistState(ctx context.Context) error {
 		return nil
 	}
 
-	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
 	defer cancel()
 
 	done := make(chan error, 1)
 
-	start := time.Now()
 	go func() {
-		done <- util.WriteJsonWithRestrictedPermission(ctx, m.filePath, m.states)
+		data, err := json.MarshalIndent(m.states, "", "  ")
+		if err != nil {
+			done <- fmt.Errorf("marshal states: %w", err)
+			return
+		}
+
+		// nolint:gosec
+		if err := os.WriteFile(m.filePath, data, 0640); err != nil {
+			done <- fmt.Errorf("write state file: %w", err)
+			return
+		}
+
+		done <- nil
 	}()
 
 	select {
@@ -198,7 +208,7 @@ func (m *Manager) PersistState(ctx context.Context) error {
 		}
 	}
 
-	log.Debugf("persisted shutdown states: %v, took %v", maps.Keys(m.dirty), time.Since(start))
+	log.Debugf("persisted shutdown states: %v", maps.Keys(m.dirty))
 
 	clear(m.dirty)
 

client/internal/statemanager/path.go

@@ -4,20 +4,32 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
+
+	log "github.com/sirupsen/logrus"
 )
 
 // GetDefaultStatePath returns the path to the state file based on the operating system
-// It returns an empty string if the path cannot be determined.
+// It returns an empty string if the path cannot be determined. It also creates the directory if it does not exist.
 func GetDefaultStatePath() string {
+	var path string
+
 	switch runtime.GOOS {
 	case "windows":
-		return filepath.Join(os.Getenv("PROGRAMDATA"), "Netbird", "state.json")
+		path = filepath.Join(os.Getenv("PROGRAMDATA"), "Netbird", "state.json")
 	case "darwin", "linux":
-		return "/var/lib/netbird/state.json"
+		path = "/var/lib/netbird/state.json"
 	case "freebsd", "openbsd", "netbsd", "dragonfly":
-		return "/var/db/netbird/state.json"
+		path = "/var/db/netbird/state.json"
+	// ios/android don't need state
+	default:
+		return ""
 	}
 
-	return ""
+	dir := filepath.Dir(path)
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		log.Errorf("Error creating directory %s: %v. Continuing without state support.", dir, err)
+		return ""
+	}
 
+	return path
 }

go.mod

@@ -71,7 +71,6 @@ 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.48.1
 	github.com/rs/xid v1.3.0
 	github.com/shirou/gopsutil/v3 v3.24.4
 	github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966
@@ -152,13 +151,11 @@ require (
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/go-ole/go-ole v1.3.0 // indirect
 	github.com/go-redis/redis/v8 v8.11.5 // indirect
-	github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 // indirect
 	github.com/go-text/render v0.1.0 // indirect
 	github.com/go-text/typesetting v0.1.0 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
 	github.com/google/btree v1.1.2 // indirect
-	github.com/google/pprof v0.0.0-20211214055906-6f57359322fd // 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
@@ -219,7 +216,6 @@ 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
 	go.uber.org/multierr v1.11.0 // indirect
 	golang.org/x/image v0.18.0 // indirect
 	golang.org/x/mod v0.17.0 // indirect

go.sum

@@ -400,7 +400,6 @@ github.com/hashicorp/serf v0.8.2/go.mod h1:6hOLApaqBFA1NXqRQAsxw9QxuDEvNxSQRwA/J
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
 github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
 github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
-github.com/ianlancetaylor/demangle v0.0.0-20210905161508-09a460cdf81d/go.mod h1:aYm2/VgdVmcIU8iMfdMvDMsRAQjcfZSKFby6HOFvi/w=
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
@@ -606,8 +605,6 @@ 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.48.1 h1:y/8xmfWI9qmGTc+lBr4jKRUWLGSlSigv847ULJ4hYXA=
-github.com/quic-go/quic-go v0.48.1/go.mod h1:yBgs3rWBOADpga7F+jJsb6Ybg1LSYiQvwWlLX+/6HMs=
 github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
 github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDNXVBh4M=
@@ -756,8 +753,6 @@ go.opentelemetry.io/proto/otlp v1.0.0/go.mod h1:Sy6pihPLfYHkr3NkUbEhGHFhINUSI/v8
 go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
 go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
-go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU=
-go.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc=
 go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
 go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
 go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
@@ -968,7 +963,6 @@ golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20210616094352-59db8d763f22/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20211007075335-d3039528d8ac/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

management/server/account.go

@@ -2029,7 +2029,7 @@ func (am *DefaultAccountManager) syncJWTGroups(ctx context.Context, accountID st
 			return fmt.Errorf("error getting user: %w", err)
 		}
 
-		groups, err := transaction.GetAccountGroups(ctx, LockingStrengthShare, accountID)
+		groups, err := transaction.GetAccountGroups(ctx, accountID)
 		if err != nil {
 			return fmt.Errorf("error getting account groups: %w", err)
 		}
@@ -2059,7 +2059,7 @@ func (am *DefaultAccountManager) syncJWTGroups(ctx context.Context, accountID st
 
 		// Propagate changes to peers if group propagation is enabled
 		if settings.GroupsPropagationEnabled {
-			groups, err = transaction.GetAccountGroups(ctx, LockingStrengthShare, accountID)
+			groups, err = transaction.GetAccountGroups(ctx, accountID)
 			if err != nil {
 				return fmt.Errorf("error getting account groups: %w", err)
 			}
@@ -2129,7 +2129,7 @@ func (am *DefaultAccountManager) syncJWTGroups(ctx context.Context, accountID st
 	if settings.GroupsPropagationEnabled {
 		account, err := am.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
 		if err != nil {
-			return status.NewGetAccountError(err)
+			return fmt.Errorf("error getting account: %w", err)
 		}
 
 		if areGroupChangesAffectPeers(account, addNewGroups) || areGroupChangesAffectPeers(account, removeOldGroups) {
@@ -2290,12 +2290,12 @@ func (am *DefaultAccountManager) SyncAndMarkPeer(ctx context.Context, accountID
 
 	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
-		return nil, nil, nil, status.NewGetAccountError(err)
+		return nil, nil, nil, err
 	}
 
 	peer, netMap, postureChecks, err := am.SyncPeer(ctx, PeerSync{WireGuardPubKey: peerPubKey, Meta: meta}, account)
 	if err != nil {
-		return nil, nil, nil, fmt.Errorf("error syncing peer: %w", err)
+		return nil, nil, nil, err
 	}
 
 	err = am.MarkPeerConnected(ctx, peerPubKey, true, realIP, account)
@@ -2314,7 +2314,7 @@ func (am *DefaultAccountManager) OnPeerDisconnected(ctx context.Context, account
 
 	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
-		return status.NewGetAccountError(err)
+		return err
 	}
 
 	err = am.MarkPeerConnected(ctx, peerPubKey, false, nil, account)

management/server/account_test.go

@@ -2773,7 +2773,7 @@ func TestAccount_SetJWTGroups(t *testing.T) {
 		err = manager.syncJWTGroups(context.Background(), "accountID", claims)
 		assert.NoError(t, err, "unable to sync jwt groups")
 
-		groups, err := manager.Store.GetAccountGroups(context.Background(), LockingStrengthShare, "accountID")
+		groups, err := manager.Store.GetAccountGroups(context.Background(), "accountID")
 		assert.NoError(t, err)
 		assert.Len(t, groups, 3, "new group3 should be added")
 

management/server/file_store.go

@@ -223,7 +223,7 @@ func restore(ctx context.Context, file string) (*FileStore, error) {
 // It is recommended to call it with locking FileStore.mux
 func (s *FileStore) persist(ctx context.Context, file string) error {
 	start := time.Now()
-	err := util.WriteJson(context.Background(), file, s)
+	err := util.WriteJson(file, s)
 	if err != nil {
 		return err
 	}

management/server/group.go

@@ -6,11 +6,10 @@ import (
 	"fmt"
 	"slices"
 
-	"github.com/rs/xid"
-	log "github.com/sirupsen/logrus"
-
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/route"
+	"github.com/rs/xid"
+	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/management/server/activity"
 	nbgroup "github.com/netbirdio/netbird/management/server/group"
@@ -28,12 +27,17 @@ func (e *GroupLinkError) Error() string {
 
 // CheckGroupPermissions validates if a user has the necessary permissions to view groups
 func (am *DefaultAccountManager) CheckGroupPermissions(ctx context.Context, accountID, userID string) error {
+	settings, err := am.Store.GetAccountSettings(ctx, LockingStrengthShare, accountID)
+	if err != nil {
+		return err
+	}
+
 	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
 		return err
 	}
 
-	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
+	if (!user.IsAdminOrServiceUser() && settings.RegularUsersViewBlocked) || user.AccountID != accountID {
 		return status.Errorf(status.PermissionDenied, "groups are blocked for users")
 	}
 
@@ -55,7 +59,7 @@ func (am *DefaultAccountManager) GetAllGroups(ctx context.Context, accountID, us
 		return nil, err
 	}
 
-	return am.Store.GetAccountGroups(ctx, LockingStrengthShare, accountID)
+	return am.Store.GetAccountGroups(ctx, accountID)
 }
 
 // GetGroupByName filters all groups in an account by name and returns the one with the most peers

management/server/group/group.go

@@ -49,8 +49,3 @@ func (g *Group) Copy() *Group {
 func (g *Group) HasPeers() bool {
 	return len(g.Peers) > 0
 }
-
-// IsGroupAll checks if the group is a default "All" group.
-func (g *Group) IsGroupAll() bool {
-	return g.Name == "All"
-}

management/server/grpcserver.go

@@ -180,7 +180,6 @@ func (s *GRPCServer) Sync(req *proto.EncryptedMessage, srv proto.ManagementServi
 
 	peer, netMap, postureChecks, err := s.accountManager.SyncAndMarkPeer(ctx, accountID, peerKey.String(), extractPeerMeta(ctx, syncReq.GetMeta()), realIP)
 	if err != nil {
-		log.WithContext(ctx).Debugf("error while syncing peer %s: %v", peerKey.String(), err)
 		return mapError(ctx, err)
 	}
 
@@ -208,7 +207,6 @@ func (s *GRPCServer) Sync(req *proto.EncryptedMessage, srv proto.ManagementServi
 
 // handleUpdates sends updates to the connected peer until the updates channel is closed.
 func (s *GRPCServer) handleUpdates(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, updates chan *UpdateMessage, srv proto.ManagementService_SyncServer) error {
-	log.WithContext(ctx).Tracef("starting to handle updates for peer %s", peerKey.String())
 	for {
 		select {
 		// condition when there are some updates
@@ -262,15 +260,10 @@ func (s *GRPCServer) cancelPeerRoutines(ctx context.Context, accountID string, p
 	unlock := s.acquirePeerLockByUID(ctx, peer.Key)
 	defer unlock()
 
-	err := s.accountManager.OnPeerDisconnected(ctx, accountID, peer.Key)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to disconnect peer %s properly: %v", peer.Key, err)
-	}
+	_ = s.accountManager.OnPeerDisconnected(ctx, accountID, peer.Key)
 	s.peersUpdateManager.CloseChannel(ctx, peer.ID)
 	s.secretsManager.CancelRefresh(peer.ID)
 	s.ephemeralManager.OnPeerDisconnected(ctx, peer)
-
-	log.WithContext(ctx).Tracef("peer %s has been disconnected", peer.Key)
 }
 
 func (s *GRPCServer) validateToken(ctx context.Context, jwtToken string) (string, error) {

management/server/http/peers_handler.go

@@ -184,26 +184,14 @@ func (h *PeersHandler) GetAllPeers(w http.ResponseWriter, r *http.Request) {
 
 	dnsDomain := h.accountManager.GetDNSDomain()
 
-	peers, err := h.accountManager.GetPeers(r.Context(), accountID, userID)
-	if err != nil {
-		util.WriteError(r.Context(), err, w)
-		return
-	}
-
-	groupsMap := map[string]*nbgroup.Group{}
-	groups, _ := h.accountManager.GetAllGroups(r.Context(), accountID, userID)
-	for _, group := range groups {
-		groupsMap[group.ID] = group
-	}
-
-	respBody := make([]*api.PeerBatch, 0, len(peers))
-	for _, peer := range peers {
+	respBody := make([]*api.PeerBatch, 0, len(account.Peers))
+	for _, peer := range account.Peers {
 		peerToReturn, err := h.checkPeerStatus(peer)
 		if err != nil {
 			util.WriteError(r.Context(), err, w)
 			return
 		}
-		groupMinimumInfo := toGroupsInfo(groupsMap, peer.ID)
+		groupMinimumInfo := toGroupsInfo(account.Groups, peer.ID)
 
 		respBody = append(respBody, toPeerListItemResponse(peerToReturn, groupMinimumInfo, dnsDomain, 0))
 	}
@@ -316,7 +304,7 @@ func peerToAccessiblePeer(peer *nbpeer.Peer, dnsDomain string) api.AccessiblePee
 }
 
 func toGroupsInfo(groups map[string]*nbgroup.Group, peerID string) []api.GroupMinimum {
-	groupsInfo := []api.GroupMinimum{}
+	var groupsInfo []api.GroupMinimum
 	groupsChecked := make(map[string]struct{})
 	for _, group := range groups {
 		_, ok := groupsChecked[group.ID]

management/server/peer.go

@@ -110,16 +110,14 @@ func (am *DefaultAccountManager) GetPeers(ctx context.Context, accountID, userID
 func (am *DefaultAccountManager) MarkPeerConnected(ctx context.Context, peerPubKey string, connected bool, realIP net.IP, account *Account) error {
 	peer, err := account.FindPeerByPubKey(peerPubKey)
 	if err != nil {
-		return fmt.Errorf("failed to find peer by pub key: %w", err)
+		return err
 	}
 
 	expired, err := am.updatePeerStatusAndLocation(ctx, peer, connected, realIP, account)
 	if err != nil {
-		return fmt.Errorf("failed to update peer status and location: %w", err)
+		return err
 	}
 
-	log.WithContext(ctx).Debugf("mark peer %s connected: %t", peer.ID, connected)
-
 	if peer.AddedWithSSOLogin() {
 		if peer.LoginExpirationEnabled && account.Settings.PeerLoginExpirationEnabled {
 			am.checkAndSchedulePeerLoginExpiration(ctx, account)
@@ -170,7 +168,7 @@ func (am *DefaultAccountManager) updatePeerStatusAndLocation(ctx context.Context
 
 	err := am.Store.SavePeerStatus(account.Id, peer.ID, *newStatus)
 	if err != nil {
-		return false, fmt.Errorf("failed to save peer status: %w", err)
+		return false, err
 	}
 
 	return oldStatus.LoginExpired, nil
@@ -589,7 +587,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 
 	account, err := am.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
 	if err != nil {
-		return nil, nil, nil, status.NewGetAccountError(err)
+		return nil, nil, nil, fmt.Errorf("error getting account: %w", err)
 	}
 
 	allGroup, err := account.GetGroupAll()
@@ -642,7 +640,7 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
 	if peer.UserID != "" {
 		user, err := account.FindUser(peer.UserID)
 		if err != nil {
-			return nil, nil, nil, fmt.Errorf("failed to get user: %w", err)
+			return nil, nil, nil, err
 		}
 
 		err = checkIfPeerOwnerIsBlocked(peer, user)
@@ -659,7 +657,7 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
 	if updated {
 		err = am.Store.SavePeer(ctx, account.Id, peer)
 		if err != nil {
-			return nil, nil, nil, fmt.Errorf("failed to save peer: %w", err)
+			return nil, nil, nil, err
 		}
 
 		if sync.UpdateAccountPeers {
@@ -669,7 +667,7 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
 
 	peerNotValid, isStatusChanged, err := am.integratedPeerValidator.IsNotValidPeer(ctx, account.Id, peer, account.GetPeerGroupsList(peer.ID), account.Settings.Extra)
 	if err != nil {
-		return nil, nil, nil, fmt.Errorf("failed to validate peer: %w", err)
+		return nil, nil, nil, err
 	}
 
 	var postureChecks []*posture.Checks
@@ -687,7 +685,7 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
 
 	validPeersMap, err := am.GetValidatedPeers(account)
 	if err != nil {
-		return nil, nil, nil, fmt.Errorf("failed to get validated peers: %w", err)
+		return nil, nil, nil, err
 	}
 	postureChecks = am.getPeerPostureChecks(account, peer)
 
@@ -767,7 +765,7 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 		}
 	}
 
-	groups, err := am.Store.GetAccountGroups(ctx, LockingStrengthShare, accountID)
+	groups, err := am.Store.GetAccountGroups(ctx, accountID)
 	if err != nil {
 		return nil, nil, nil, err
 	}

management/server/setupkey.go

@@ -4,17 +4,18 @@ import (
 	"context"
 	"crypto/sha256"
 	b64 "encoding/base64"
+	"fmt"
 	"hash/fnv"
-	"slices"
 	"strconv"
 	"strings"
 	"time"
 	"unicode/utf8"
 
 	"github.com/google/uuid"
+	log "github.com/sirupsen/logrus"
+
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/status"
-	log "github.com/sirupsen/logrus"
 )
 
 const (
@@ -228,43 +229,32 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
-	}
-
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
-	}
-
-	var setupKey *SetupKey
-	var plainKey string
-	var eventsToStore []func()
-
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		if err = validateSetupKeyAutoGroups(ctx, transaction, accountID, autoGroups); err != nil {
-			return err
-		}
-
-		setupKey, plainKey = GenerateSetupKey(keyName, keyType, expiresIn, autoGroups, usageLimit, ephemeral)
-		setupKey.AccountID = accountID
-
-		events := am.prepareSetupKeyEvents(ctx, transaction, accountID, userID, autoGroups, nil, setupKey)
-		eventsToStore = append(eventsToStore, events...)
-
-		return transaction.SaveSetupKey(ctx, LockingStrengthUpdate, setupKey)
-	})
-	if err != nil {
+	if err := validateSetupKeyAutoGroups(account, autoGroups); err != nil {
 		return nil, err
 	}
 
+	setupKey, plainKey := GenerateSetupKey(keyName, keyType, expiresIn, autoGroups, usageLimit, ephemeral)
+	account.SetupKeys[setupKey.Key] = setupKey
+	err = am.Store.SaveAccount(ctx, account)
+	if err != nil {
+		return nil, status.Errorf(status.Internal, "failed adding account key")
+	}
+
 	am.StoreEvent(ctx, userID, setupKey.Id, accountID, activity.SetupKeyCreated, setupKey.EventMeta())
-	for _, storeEvent := range eventsToStore {
-		storeEvent()
+
+	for _, g := range setupKey.AutoGroups {
+		group := account.GetGroup(g)
+		if group != nil {
+			am.StoreEvent(ctx, userID, setupKey.Id, accountID, activity.GroupAddedToSetupKey,
+				map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": setupKey.Name})
+		} else {
+			log.WithContext(ctx).Errorf("group %s not found while saving setup key activity event of account %s", g, account.Id)
+		}
 	}
 
 	// for the creation return the plain key to the caller
@@ -278,56 +268,43 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 // (e.g. the key itself, creation date, ID, etc).
 // These properties are overwritten: Name, AutoGroups, Revoked. The rest is copied from the existing key.
 func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID string, keyToSave *SetupKey, userID string) (*SetupKey, error) {
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
+	defer unlock()
+
 	if keyToSave == nil {
 		return nil, status.Errorf(status.InvalidArgument, "provided setup key to update is nil")
 	}
 
-	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
-	defer unlock()
-
-	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
+	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
-	}
-
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
-	}
-
 	var oldKey *SetupKey
-	var newKey *SetupKey
-	var eventsToStore []func()
-
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		if err = validateSetupKeyAutoGroups(ctx, transaction, accountID, keyToSave.AutoGroups); err != nil {
-			return err
+	for _, key := range account.SetupKeys {
+		if key.Id == keyToSave.Id {
+			oldKey = key.Copy()
+			break
 		}
+	}
+	if oldKey == nil {
+		return nil, status.Errorf(status.NotFound, "setup key not found")
+	}
 
-		oldKey, err = transaction.GetSetupKeyByID(ctx, LockingStrengthShare, accountID, keyToSave.Id)
-		if err != nil {
-			return err
-		}
+	if err := validateSetupKeyAutoGroups(account, keyToSave.AutoGroups); err != nil {
+		return nil, err
+	}
 
-		// only auto groups, revoked status, and name can be updated for now
-		newKey = oldKey.Copy()
-		newKey.Name = keyToSave.Name
-		newKey.AutoGroups = keyToSave.AutoGroups
-		newKey.Revoked = keyToSave.Revoked
-		newKey.UpdatedAt = time.Now().UTC()
+	// only auto groups, revoked status, and name can be updated for now
+	newKey := oldKey.Copy()
+	newKey.Name = keyToSave.Name
+	newKey.AutoGroups = keyToSave.AutoGroups
+	newKey.Revoked = keyToSave.Revoked
+	newKey.UpdatedAt = time.Now().UTC()
 
-		addedGroups := difference(newKey.AutoGroups, oldKey.AutoGroups)
-		removedGroups := difference(oldKey.AutoGroups, newKey.AutoGroups)
+	account.SetupKeys[newKey.Key] = newKey
 
-		events := am.prepareSetupKeyEvents(ctx, transaction, accountID, userID, addedGroups, removedGroups, oldKey)
-		eventsToStore = append(eventsToStore, events...)
-
-		return transaction.SaveSetupKey(ctx, LockingStrengthUpdate, newKey)
-	})
-	if err != nil {
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
 		return nil, err
 	}
 
@@ -335,9 +312,30 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 		am.StoreEvent(ctx, userID, newKey.Id, accountID, activity.SetupKeyRevoked, newKey.EventMeta())
 	}
 
-	for _, storeEvent := range eventsToStore {
-		storeEvent()
-	}
+	defer func() {
+		addedGroups := difference(newKey.AutoGroups, oldKey.AutoGroups)
+		removedGroups := difference(oldKey.AutoGroups, newKey.AutoGroups)
+		for _, g := range removedGroups {
+			group := account.GetGroup(g)
+			if group != nil {
+				am.StoreEvent(ctx, userID, oldKey.Id, accountID, activity.GroupRemovedFromSetupKey,
+					map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": newKey.Name})
+			} else {
+				log.WithContext(ctx).Errorf("group %s not found while saving setup key activity event of account %s", g, account.Id)
+			}
+
+		}
+
+		for _, g := range addedGroups {
+			group := account.GetGroup(g)
+			if group != nil {
+				am.StoreEvent(ctx, userID, oldKey.Id, accountID, activity.GroupAddedToSetupKey,
+					map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": newKey.Name})
+			} else {
+				log.WithContext(ctx).Errorf("group %s not found while saving setup key activity event of account %s", g, account.Id)
+			}
+		}
+	}()
 
 	return newKey, nil
 }
@@ -349,15 +347,16 @@ func (am *DefaultAccountManager) ListSetupKeys(ctx context.Context, accountID, u
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
+	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
+		return nil, status.NewUnauthorizedToViewSetupKeysError()
 	}
 
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
+	setupKeys, err := am.Store.GetAccountSetupKeys(ctx, LockingStrengthShare, accountID)
+	if err != nil {
+		return nil, err
 	}
 
-	return am.Store.GetAccountSetupKeys(ctx, LockingStrengthShare, accountID)
+	return setupKeys, nil
 }
 
 // GetSetupKey looks up a SetupKey by KeyID, returns NotFound error if not found.
@@ -367,12 +366,8 @@ func (am *DefaultAccountManager) GetSetupKey(ctx context.Context, accountID, use
 		return nil, err
 	}
 
-	if user.AccountID != accountID {
-		return nil, status.NewUserNotPartOfAccountError()
-	}
-
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
+	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
+		return nil, status.NewUnauthorizedToViewSetupKeysError()
 	}
 
 	setupKey, err := am.Store.GetSetupKeyByID(ctx, LockingStrengthShare, keyID, accountID)
@@ -392,29 +387,21 @@ func (am *DefaultAccountManager) GetSetupKey(ctx context.Context, accountID, use
 func (am *DefaultAccountManager) DeleteSetupKey(ctx context.Context, accountID, userID, keyID string) error {
 	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to get user: %w", err)
 	}
 
-	if user.AccountID != accountID {
-		return status.NewUserNotPartOfAccountError()
+	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
+		return status.NewUnauthorizedToViewSetupKeysError()
 	}
 
-	if user.IsRegularUser() {
-		return status.NewAdminPermissionError()
-	}
-
-	var deletedSetupKey *SetupKey
-
-	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		deletedSetupKey, err = transaction.GetSetupKeyByID(ctx, LockingStrengthShare, accountID, keyID)
-		if err != nil {
-			return err
-		}
-
-		return transaction.DeleteSetupKey(ctx, LockingStrengthUpdate, accountID, keyID)
-	})
+	deletedSetupKey, err := am.Store.GetSetupKeyByID(ctx, LockingStrengthShare, keyID, accountID)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to get setup key: %w", err)
+	}
+
+	err = am.Store.DeleteSetupKey(ctx, accountID, keyID)
+	if err != nil {
+		return fmt.Errorf("failed to delete setup key: %w", err)
 	}
 
 	am.StoreEvent(ctx, userID, keyID, accountID, activity.SetupKeyDeleted, deletedSetupKey.EventMeta())
@@ -422,62 +409,15 @@ func (am *DefaultAccountManager) DeleteSetupKey(ctx context.Context, accountID,
 	return nil
 }
 
-func validateSetupKeyAutoGroups(ctx context.Context, transaction Store, accountID string, autoGroupIDs []string) error {
-	groups, err := transaction.GetGroupsByIDs(ctx, LockingStrengthShare, accountID, autoGroupIDs)
-	if err != nil {
-		return err
-	}
-
-	for _, groupID := range autoGroupIDs {
-		group, ok := groups[groupID]
+func validateSetupKeyAutoGroups(account *Account, autoGroups []string) error {
+	for _, group := range autoGroups {
+		g, ok := account.Groups[group]
 		if !ok {
-			return status.Errorf(status.NotFound, "group not found: %s", groupID)
+			return status.Errorf(status.NotFound, "group %s doesn't exist", group)
 		}
-
-		if group.IsGroupAll() {
-			return status.Errorf(status.InvalidArgument, "can't add 'All' group to the setup key")
+		if g.Name == "All" {
+			return status.Errorf(status.InvalidArgument, "can't add All group to the setup key")
 		}
 	}
-
 	return nil
 }
-
-// prepareSetupKeyEvents prepares a list of event functions to be stored.
-func (am *DefaultAccountManager) prepareSetupKeyEvents(ctx context.Context, transaction Store, accountID, userID string, addedGroups, removedGroups []string, key *SetupKey) []func() {
-	var eventsToStore []func()
-
-	modifiedGroups := slices.Concat(addedGroups, removedGroups)
-	groups, err := transaction.GetGroupsByIDs(ctx, LockingStrengthShare, accountID, modifiedGroups)
-	if err != nil {
-		log.WithContext(ctx).Errorf("issue getting groups for setup key events: %v", err)
-		return nil
-	}
-
-	for _, g := range removedGroups {
-		group, ok := groups[g]
-		if !ok {
-			log.WithContext(ctx).Debugf("skipped adding group: %s GroupRemovedFromSetupKey activity: %v", g, err)
-			continue
-		}
-
-		eventsToStore = append(eventsToStore, func() {
-			meta := map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": key.Name}
-			am.StoreEvent(ctx, userID, key.Id, accountID, activity.GroupRemovedFromSetupKey, meta)
-		})
-	}
-
-	for _, g := range addedGroups {
-		group, ok := groups[g]
-		if !ok {
-			log.WithContext(ctx).Debugf("skipped adding group: %s GroupAddedToSetupKey activity: %v", g, err)
-			continue
-		}
-
-		eventsToStore = append(eventsToStore, func() {
-			meta := map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": key.Name}
-			am.StoreEvent(ctx, userID, key.Id, accountID, activity.GroupAddedToSetupKey, meta)
-		})
-	}
-
-	return eventsToStore
-}

management/server/sql_store.go

@@ -485,10 +485,9 @@ func (s *SqlStore) GetAccountBySetupKey(ctx context.Context, setupKey string) (*
 	result := s.db.Select("account_id").First(&key, keyQueryCondition, setupKey)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.NewSetupKeyNotFoundError(setupKey)
+			return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
 		}
-		log.WithContext(ctx).Errorf("failed to get account by setup key from store: %v", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get account by setup key from store")
+		return nil, status.NewSetupKeyNotFoundError(result.Error)
 	}
 
 	if key.AccountID == "" {
@@ -569,15 +568,15 @@ func (s *SqlStore) GetAccountUsers(ctx context.Context, accountID string) ([]*Us
 	return users, nil
 }
 
-func (s *SqlStore) GetAccountGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbgroup.Group, error) {
+func (s *SqlStore) GetAccountGroups(ctx context.Context, accountID string) ([]*nbgroup.Group, error) {
 	var groups []*nbgroup.Group
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&groups, accountIDCondition, accountID)
+	result := s.db.Find(&groups, accountIDCondition, accountID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "accountID not found: index lookup failed")
 		}
-		log.WithContext(ctx).Errorf("failed to get account groups from the store: %s", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get account groups from the store")
+		log.WithContext(ctx).Errorf("error when getting groups from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "issue getting groups from store")
 	}
 
 	return groups, nil
@@ -757,10 +756,9 @@ func (s *SqlStore) GetAccountIDBySetupKey(ctx context.Context, setupKey string)
 	result := s.db.Model(&SetupKey{}).Select("account_id").Where(keyQueryCondition, setupKey).First(&accountID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return "", status.NewSetupKeyNotFoundError(setupKey)
+			return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
 		}
-		log.WithContext(ctx).Errorf("failed to get account ID by setup key from store: %v", result.Error)
-		return "", status.Errorf(status.Internal, "failed to get account ID by setup key from store")
+		return "", status.NewSetupKeyNotFoundError(result.Error)
 	}
 
 	if accountID == "" {
@@ -988,10 +986,9 @@ func (s *SqlStore) GetSetupKeyBySecret(ctx context.Context, lockStrength Locking
 		First(&setupKey, keyQueryCondition, key)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.NewSetupKeyNotFoundError(key)
+			return nil, status.Errorf(status.NotFound, "setup key not found")
 		}
-		log.WithContext(ctx).Errorf("failed to get setup key by secret from store: %v", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get setup key by secret from store")
+		return nil, status.NewSetupKeyNotFoundError(result.Error)
 	}
 	return &setupKey, nil
 }
@@ -1009,7 +1006,7 @@ func (s *SqlStore) IncrementSetupKeyUsage(ctx context.Context, setupKeyID string
 	}
 
 	if result.RowsAffected == 0 {
-		return status.NewSetupKeyNotFoundError(setupKeyID)
+		return status.Errorf(status.NotFound, "setup key not found")
 	}
 
 	return nil
@@ -1182,23 +1179,6 @@ func (s *SqlStore) GetGroupByName(ctx context.Context, lockStrength LockingStren
 	return &group, nil
 }
 
-// GetGroupsByIDs retrieves groups by their IDs and account ID.
-func (s *SqlStore) GetGroupsByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, groupIDs []string) (map[string]*nbgroup.Group, error) {
-	var groups []*nbgroup.Group
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&groups, "account_id = ? AND id in ?", accountID, groupIDs)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to get groups by ID's from the store: %s", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to get groups by ID's from the store")
-	}
-
-	groupsMap := make(map[string]*nbgroup.Group)
-	for _, group := range groups {
-		groupsMap[group.ID] = group
-	}
-
-	return groupsMap, nil
-}
-
 // SaveGroup saves a group to the store.
 func (s *SqlStore) SaveGroup(ctx context.Context, lockStrength LockingStrength, group *nbgroup.Group) error {
 	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(group)
@@ -1240,57 +1220,12 @@ func (s *SqlStore) GetRouteByID(ctx context.Context, lockStrength LockingStrengt
 
 // GetAccountSetupKeys retrieves setup keys for an account.
 func (s *SqlStore) GetAccountSetupKeys(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*SetupKey, error) {
-	var setupKeys []*SetupKey
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		Find(&setupKeys, accountIDCondition, accountID)
-	if err := result.Error; err != nil {
-		log.WithContext(ctx).Errorf("failed to get setup keys from the store: %s", err)
-		return nil, status.Errorf(status.Internal, "failed to get setup keys from store")
-	}
-
-	return setupKeys, nil
+	return getRecords[*SetupKey](s.db, lockStrength, accountID)
 }
 
 // GetSetupKeyByID retrieves a setup key by its ID and account ID.
-func (s *SqlStore) GetSetupKeyByID(ctx context.Context, lockStrength LockingStrength, accountID, setupKeyID string) (*SetupKey, error) {
-	var setupKey *SetupKey
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		First(&setupKey, accountAndIDQueryCondition, accountID, setupKeyID)
-	if err := result.Error; err != nil {
-		if errors.Is(err, gorm.ErrRecordNotFound) {
-			return nil, status.NewSetupKeyNotFoundError(setupKeyID)
-		}
-		log.WithContext(ctx).Errorf("failed to get setup key from the store: %s", err)
-		return nil, status.Errorf(status.Internal, "failed to get setup key from store")
-	}
-
-	return setupKey, nil
-}
-
-// SaveSetupKey saves a setup key to the database.
-func (s *SqlStore) SaveSetupKey(ctx context.Context, lockStrength LockingStrength, setupKey *SetupKey) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(setupKey)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to save setup key to store: %s", result.Error)
-		return status.Errorf(status.Internal, "failed to save setup key to store")
-	}
-
-	return nil
-}
-
-// DeleteSetupKey deletes a setup key from the database.
-func (s *SqlStore) DeleteSetupKey(ctx context.Context, lockStrength LockingStrength, accountID, keyID string) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Delete(&SetupKey{}, accountAndIDQueryCondition, accountID, keyID)
-	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to delete setup key from store: %s", result.Error)
-		return status.Errorf(status.Internal, "failed to delete setup key from store")
-	}
-
-	if result.RowsAffected == 0 {
-		return status.NewSetupKeyNotFoundError(keyID)
-	}
-
-	return nil
+func (s *SqlStore) GetSetupKeyByID(ctx context.Context, lockStrength LockingStrength, setupKeyID string, accountID string) (*SetupKey, error) {
+	return getRecordByID[SetupKey](s.db, lockStrength, setupKeyID, accountID)
 }
 
 // GetAccountNameServerGroups retrieves name server groups for an account.
@@ -1303,6 +1238,10 @@ func (s *SqlStore) GetNameServerGroupByID(ctx context.Context, lockStrength Lock
 	return getRecordByID[nbdns.NameServerGroup](s.db, lockStrength, nsGroupID, accountID)
 }
 
+func (s *SqlStore) DeleteSetupKey(ctx context.Context, accountID, keyID string) error {
+	return deleteRecordByID[SetupKey](s.db, LockingStrengthUpdate, keyID, accountID)
+}
+
 // getRecords retrieves records from the database based on the account ID.
 func getRecords[T any](db *gorm.DB, lockStrength LockingStrength, accountID string) ([]T, error) {
 	var record []T
@@ -1335,3 +1274,21 @@ func getRecordByID[T any](db *gorm.DB, lockStrength LockingStrength, recordID, a
 	}
 	return &record, nil
 }
+
+// deleteRecordByID deletes a record by its ID and account ID from the database.
+func deleteRecordByID[T any](db *gorm.DB, lockStrength LockingStrength, recordID, accountID string) error {
+	var record T
+	result := db.Clauses(clause.Locking{Strength: string(lockStrength)}).Delete(record, accountAndIDQueryCondition, accountID, recordID)
+	if err := result.Error; err != nil {
+		parts := strings.Split(fmt.Sprintf("%T", record), ".")
+		recordType := parts[len(parts)-1]
+
+		return status.Errorf(status.Internal, "failed to delete %s from store: %v", recordType, err)
+	}
+
+	if result.RowsAffected == 0 {
+		return status.Errorf(status.NotFound, "record not found")
+	}
+
+	return nil
+}

management/server/sql_store_test.go

@@ -1274,7 +1274,7 @@ func Test_DeleteSetupKeySuccessfully(t *testing.T) {
 	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
 	setupKeyID := "A2C8E62B-38F5-4553-B31E-DD66C696CEBB"
 
-	err = store.DeleteSetupKey(context.Background(), LockingStrengthUpdate, accountID, setupKeyID)
+	err = store.DeleteSetupKey(context.Background(), accountID, setupKeyID)
 	require.NoError(t, err)
 
 	_, err = store.GetSetupKeyByID(context.Background(), LockingStrengthShare, setupKeyID, accountID)
@@ -1290,6 +1290,6 @@ func Test_DeleteSetupKeyFailsForNonExistingKey(t *testing.T) {
 	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
 	nonExistingKeyID := "non-existing-key-id"
 
-	err = store.DeleteSetupKey(context.Background(), LockingStrengthUpdate, accountID, nonExistingKeyID)
+	err = store.DeleteSetupKey(context.Background(), accountID, nonExistingKeyID)
 	require.Error(t, err)
 }

management/server/status/error.go

@@ -103,29 +103,19 @@ func NewPeerLoginExpiredError() error {
 }
 
 // NewSetupKeyNotFoundError creates a new Error with NotFound type for a missing setup key
-func NewSetupKeyNotFoundError(setupKeyID string) error {
-	return Errorf(NotFound, "setup key: %s not found", setupKeyID)
+func NewSetupKeyNotFoundError(err error) error {
+	return Errorf(NotFound, "setup key not found: %s", err)
 }
 
 func NewGetAccountFromStoreError(err error) error {
 	return Errorf(Internal, "issue getting account from store: %s", err)
 }
 
-// NewUserNotPartOfAccountError creates a new Error with PermissionDenied type for a user not being part of an account
-func NewUserNotPartOfAccountError() error {
-	return Errorf(PermissionDenied, "user is not part of this account")
-}
-
 // NewGetUserFromStoreError creates a new Error with Internal type for an issue getting user from store
 func NewGetUserFromStoreError() error {
 	return Errorf(Internal, "issue getting user from store")
 }
 
-// NewAdminPermissionError creates a new Error with PermissionDenied type for actions requiring admin role.
-func NewAdminPermissionError() error {
-	return Errorf(PermissionDenied, "admin role required to perform this action")
-}
-
 // NewStoreContextCanceledError creates a new Error with Internal type for a canceled store context
 func NewStoreContextCanceledError(duration time.Duration) error {
 	return Errorf(Internal, "store access: context canceled after %v", duration)
@@ -136,7 +126,7 @@ func NewInvalidKeyIDError() error {
 	return Errorf(InvalidArgument, "invalid key ID")
 }
 
-// NewGetAccountError creates a new Error with Internal type for an issue getting account
-func NewGetAccountError(err error) error {
-	return Errorf(Internal, "error getting account: %s", err)
+// NewUnauthorizedToViewSetupKeysError creates a new Error with Unauthorized type for an issue getting a setup key
+func NewUnauthorizedToViewSetupKeysError() error {
+	return Errorf(Unauthorized, "only users with admin power can view setup keys")
 }

management/server/store.go

@@ -70,10 +70,9 @@ type Store interface {
 	DeleteHashedPAT2TokenIDIndex(hashedToken string) error
 	DeleteTokenID2UserIDIndex(tokenID string) error
 
-	GetAccountGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbgroup.Group, error)
+	GetAccountGroups(ctx context.Context, accountID string) ([]*nbgroup.Group, error)
 	GetGroupByID(ctx context.Context, lockStrength LockingStrength, groupID, accountID string) (*nbgroup.Group, error)
 	GetGroupByName(ctx context.Context, lockStrength LockingStrength, groupName, accountID string) (*nbgroup.Group, error)
-	GetGroupsByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, groupIDs []string) (map[string]*nbgroup.Group, error)
 	SaveGroups(ctx context.Context, lockStrength LockingStrength, groups []*nbgroup.Group) error
 	SaveGroup(ctx context.Context, lockStrength LockingStrength, group *nbgroup.Group) error
 
@@ -97,9 +96,7 @@ type Store interface {
 	GetSetupKeyBySecret(ctx context.Context, lockStrength LockingStrength, key string) (*SetupKey, error)
 	IncrementSetupKeyUsage(ctx context.Context, setupKeyID string) error
 	GetAccountSetupKeys(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*SetupKey, error)
-	GetSetupKeyByID(ctx context.Context, lockStrength LockingStrength, accountID, setupKeyID string) (*SetupKey, error)
-	SaveSetupKey(ctx context.Context, lockStrength LockingStrength, setupKey *SetupKey) error
-	DeleteSetupKey(ctx context.Context, lockStrength LockingStrength, accountID, keyID string) error
+	GetSetupKeyByID(ctx context.Context, lockStrength LockingStrength, setupKeyID string, accountID string) (*SetupKey, error)
 
 	GetAccountRoutes(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*route.Route, error)
 	GetRouteByID(ctx context.Context, lockStrength LockingStrength, routeID string, accountID string) (*route.Route, error)
@@ -127,6 +124,7 @@ type Store interface {
 	// This is also a method of metrics.DataSource interface.
 	GetStoreEngine() StoreEngine
 	ExecuteInTransaction(ctx context.Context, f func(store Store) error) error
+	DeleteSetupKey(ctx context.Context, accountID, keyID string) error
 }
 
 type StoreEngine string

management/server/updatechannel.go

@@ -96,12 +96,9 @@ func (p *PeersUpdateManager) closeChannel(ctx context.Context, peerID string) {
 	if channel, ok := p.peerChannels[peerID]; ok {
 		delete(p.peerChannels, peerID)
 		close(channel)
-
-		log.WithContext(ctx).Debugf("closed updates channel of a peer %s", peerID)
-		return
 	}
 
-	log.WithContext(ctx).Debugf("closing updates channel: peer %s has no channel", peerID)
+	log.WithContext(ctx).Debugf("closed updates channel of a peer %s", peerID)
 }
 
 // CloseChannels closes updates channel for each given peer

management/server/user.go

@@ -9,16 +9,14 @@ import (
 	"time"
 
 	"github.com/google/uuid"
-	log "github.com/sirupsen/logrus"
-
 	"github.com/netbirdio/netbird/management/server/activity"
-	nbContext "github.com/netbirdio/netbird/management/server/context"
 	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	"github.com/netbirdio/netbird/management/server/idp"
 	"github.com/netbirdio/netbird/management/server/integration_reference"
 	"github.com/netbirdio/netbird/management/server/jwtclaims"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/status"
+	log "github.com/sirupsen/logrus"
 )
 
 const (
@@ -105,11 +103,6 @@ func (u *User) IsAdminOrServiceUser() bool {
 	return u.HasAdminPower() || u.IsServiceUser
 }
 
-// IsRegularUser checks if the user is a regular user.
-func (u *User) IsRegularUser() bool {
-	return !u.HasAdminPower() && !u.IsServiceUser
-}
-
 // ToUserInfo converts a User object to a UserInfo object.
 func (u *User) ToUserInfo(userData *idp.UserData, settings *Settings) (*UserInfo, error) {
 	autoGroups := u.AutoGroups
@@ -1107,9 +1100,6 @@ func (am *DefaultAccountManager) GetUsersFromAccount(ctx context.Context, accoun
 func (am *DefaultAccountManager) expireAndUpdatePeers(ctx context.Context, account *Account, peers []*nbpeer.Peer) error {
 	var peerIDs []string
 	for _, peer := range peers {
-		// nolint:staticcheck
-		ctx = context.WithValue(ctx, nbContext.PeerIDKey, peer.Key)
-
 		if peer.Status.LoginExpired {
 			continue
 		}
@@ -1117,11 +1107,8 @@ func (am *DefaultAccountManager) expireAndUpdatePeers(ctx context.Context, accou
 		peer.MarkLoginExpired(true)
 		account.UpdatePeer(peer)
 		if err := am.Store.SavePeerStatus(account.Id, peer.ID, *peer.Status); err != nil {
-			return fmt.Errorf("failed saving peer status for peer %s: %s", peer.ID, err)
+			return err
 		}
-
-		log.WithContext(ctx).Tracef("mark peer %s login expired", peer.ID)
-
 		am.StoreEvent(
 			ctx,
 			peer.UserID, peer.ID, account.Id,

relay/client/client.go

@@ -10,7 +10,6 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	auth "github.com/netbirdio/netbird/relay/auth/hmac"
-	"github.com/netbirdio/netbird/relay/client/dialer/quic"
 	"github.com/netbirdio/netbird/relay/client/dialer/ws"
 	"github.com/netbirdio/netbird/relay/healthcheck"
 	"github.com/netbirdio/netbird/relay/messages"
@@ -94,6 +93,10 @@ func (cc *connContainer) writeMsg(msg Msg) {
 	case cc.messages <- msg:
 	case <-cc.ctx.Done():
 		msg.Free()
+	default:
+		msg.Free()
+		cc.log.Infof("message queue is full")
+		// todo consider to close the connection
 	}
 }
 
@@ -137,7 +140,7 @@ type Client struct {
 	instanceURL      *RelayAddr
 	muInstanceURL    sync.Mutex
 
-	onDisconnectListener func(string)
+	onDisconnectListener func()
 	onConnectedListener  func()
 	listenerMutex        sync.Mutex
 }
@@ -176,7 +179,8 @@ func (c *Client) Connect() error {
 		return nil
 	}
 
-	if err := c.connect(); err != nil {
+	err := c.connect()
+	if err != nil {
 		return err
 	}
 
@@ -229,7 +233,7 @@ func (c *Client) ServerInstanceURL() (string, error) {
 }
 
 // SetOnDisconnectListener sets a function that will be called when the connection to the relay server is closed.
-func (c *Client) SetOnDisconnectListener(fn func(string)) {
+func (c *Client) SetOnDisconnectListener(fn func()) {
 	c.listenerMutex.Lock()
 	defer c.listenerMutex.Unlock()
 	c.onDisconnectListener = fn
@@ -260,20 +264,14 @@ func (c *Client) Close() error {
 }
 
 func (c *Client) connect() error {
-	var conn net.Conn
-	var err error
-	if c.connectionURL == "rels://temp-relay-quic.relay.netbird.io:443" {
-		log.Infof("connecting to relay server %s using quic protocol", c.connectionURL)
-		conn, err = quic.Dial(c.connectionURL)
-	} else {
-		conn, err = ws.Dial(c.connectionURL)
-	}
+	conn, err := ws.Dial(c.connectionURL)
 	if err != nil {
 		return err
 	}
 	c.relayConn = conn
 
-	if err = c.handShake(); err != nil {
+	err = c.handShake()
+	if err != nil {
 		cErr := conn.Close()
 		if cErr != nil {
 			c.log.Errorf("failed to close connection: %s", cErr)
@@ -347,7 +345,7 @@ func (c *Client) readLoop(relayConn net.Conn) {
 			c.log.Infof("start to Relay read loop exit")
 			c.mu.Lock()
 			if c.serviceIsRunning && !internallyStoppedFlag.isSet() {
-				c.log.Errorf("failed to read message from relay server: %s", errExit)
+				c.log.Debugf("failed to read message from relay server: %s", errExit)
 			}
 			c.mu.Unlock()
 			break
@@ -556,7 +554,7 @@ func (c *Client) notifyDisconnected() {
 	if c.onDisconnectListener == nil {
 		return
 	}
-	go c.onDisconnectListener(c.connectionURL)
+	go c.onDisconnectListener()
 }
 
 func (c *Client) notifyConnected() {

relay/client/client_test.go

@@ -551,7 +551,7 @@ func TestCloseByServer(t *testing.T) {
 	}
 
 	disconnected := make(chan struct{})
-	relayClient.SetOnDisconnectListener(func(_ string) {
+	relayClient.SetOnDisconnectListener(func() {
 		log.Infof("client disconnected")
 		close(disconnected)
 	})

relay/client/dialer/quic/conn.go

@@ -1,87 +0,0 @@
-package quic
-
-import (
-	"context"
-	"net"
-	"time"
-
-	"github.com/quic-go/quic-go"
-	log "github.com/sirupsen/logrus"
-)
-
-type QuicAddr struct {
-	addr string
-}
-
-func (a QuicAddr) Network() string {
-	return "quic"
-}
-
-func (a QuicAddr) String() string {
-	return a.addr
-}
-
-type Conn struct {
-	session    quic.Connection
-	remoteAddr QuicAddr
-	ctx        context.Context
-}
-
-func NewConn(session quic.Connection, serverAddress string) net.Conn {
-	return &Conn{
-		session:    session,
-		remoteAddr: QuicAddr{addr: serverAddress},
-		ctx:        context.Background(),
-	}
-}
-
-func (c *Conn) Read(b []byte) (n int, err error) {
-	dgram, err := c.session.ReceiveDatagram(c.ctx)
-	if err != nil {
-		log.Errorf("failed to read from QUIC session: %v", err)
-		return 0, err
-	}
-
-	n = copy(b, dgram)
-	return n, nil
-}
-
-func (c *Conn) Write(b []byte) (int, error) {
-	err := c.session.SendDatagram(b)
-	if err != nil {
-		log.Errorf("failed to write to QUIC stream: %v", err)
-		return 0, err
-	}
-	return len(b), nil
-}
-
-func (c *Conn) RemoteAddr() net.Addr {
-	if c.session != nil {
-		return c.session.RemoteAddr()
-	}
-	return c.remoteAddr
-}
-
-func (c *Conn) LocalAddr() net.Addr {
-	if c.session != nil {
-		return c.session.LocalAddr()
-	}
-	return QuicAddr{addr: "unknown"}
-}
-
-func (c *Conn) SetReadDeadline(t time.Time) error {
-	return nil
-}
-
-func (c *Conn) SetWriteDeadline(t time.Time) error {
-	return nil
-}
-
-func (c *Conn) SetDeadline(t time.Time) error {
-
-	return nil
-}
-
-func (c *Conn) Close() error {
-	return c.session.CloseWithError(0, "normal closure")
-}

relay/client/dialer/quic/quic.go

@@ -1,60 +0,0 @@
-package quic
-
-import (
-	"context"
-	"crypto/tls"
-	"fmt"
-	"net"
-	"strings"
-	"time"
-
-	"github.com/quic-go/quic-go"
-	log "github.com/sirupsen/logrus"
-)
-
-const (
-	dialTimeout = 30 * time.Second
-)
-
-func Dial(address string) (net.Conn, error) {
-	quicURL, err := prepareURL(address)
-	if err != nil {
-		return nil, err
-	}
-
-	ctx, cancel := context.WithTimeout(context.Background(), dialTimeout)
-	defer cancel()
-
-	tlsConf := &tls.Config{
-		InsecureSkipVerify: true,           // Set to true only for testing
-		NextProtos:         []string{"h2"}, // Ensure this matches the server's ALPN
-	}
-
-	quicConfig := &quic.Config{
-		KeepAlivePeriod: 15 * time.Second,
-		MaxIdleTimeout:  60 * time.Second,
-		EnableDatagrams: true,
-	}
-
-	// todo add support for custom dialer
-
-	session, err := quic.DialAddr(ctx, quicURL, tlsConf, quicConfig)
-	if err != nil {
-		log.Errorf("failed to dial to Relay server via QUIC '%s': %s", quicURL, err)
-		return nil, err
-	}
-
-	conn := NewConn(session, address)
-	return conn, nil
-}
-
-func prepareURL(address string) (string, error) {
-	if !strings.HasPrefix(address, "rel://") && !strings.HasPrefix(address, "rels://") {
-		return "", fmt.Errorf("unsupported scheme: %s", address)
-	}
-
-	if strings.HasPrefix(address, "rels://") {
-		return address[7:], nil
-	}
-	return address[6:], nil
-}

relay/client/dialer/ws/ws.go

@@ -2,7 +2,6 @@ package ws
 
 import (
 	"context"
-	"crypto/tls"
 	"fmt"
 	"net"
 	"net/http"
@@ -60,10 +59,6 @@ func httpClientNbDialer() *http.Client {
 		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
 			return customDialer.DialContext(ctx, network, addr)
 		},
-		// Set up a TLS configuration that skips certificate verification
-		TLSClientConfig: &tls.Config{
-			InsecureSkipVerify: true, // This accepts invalid TLS certificates
-		},
 	}
 
 	return &http.Client{

relay/client/guard.go

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

relay/client/manager.go

@@ -57,15 +57,12 @@ type ManagerService interface {
 // relay servers will be closed if there is no active connection. Periodically the manager will check if there is any
 // unused relay connection and close it.
 type Manager struct {
-	ctx          context.Context
-	peerID       string
-	running      bool
-	tokenStore   *relayAuth.TokenStore
-	serverPicker *ServerPicker
+	ctx        context.Context
+	serverURLs []string
+	peerID     string
+	tokenStore *relayAuth.TokenStore
 
-	relayClient *Client
-	// the guard logic can overwrite the relayClient variable, this mutex protect the usage of the variable
-	relayClientMu  sync.Mutex
+	relayClient    *Client
 	reconnectGuard *Guard
 
 	relayClients      map[string]*RelayTrack
@@ -79,54 +76,48 @@ type Manager struct {
 // NewManager creates a new manager instance.
 // The serverURL address can be empty. In this case, the manager will not serve.
 func NewManager(ctx context.Context, serverURLs []string, peerID string) *Manager {
-	tokenStore := &relayAuth.TokenStore{}
-
-	m := &Manager{
-		ctx:        ctx,
-		peerID:     peerID,
-		tokenStore: tokenStore,
-		serverPicker: &ServerPicker{
-			TokenStore: tokenStore,
-			PeerID:     peerID,
-		},
+	return &Manager{
+		ctx:                     ctx,
+		serverURLs:              serverURLs,
+		peerID:                  peerID,
+		tokenStore:              &relayAuth.TokenStore{},
 		relayClients:            make(map[string]*RelayTrack),
 		onDisconnectedListeners: make(map[string]*list.List),
 	}
-	m.serverPicker.ServerURLs.Store(serverURLs)
-	m.reconnectGuard = NewGuard(m.serverPicker)
-	return m
 }
 
-// Serve starts the manager, attempting to establish a connection with the relay server.
-// If the connection fails, it will keep trying to reconnect in the background.
-// Additionally, it starts a cleanup loop to remove unused relay connections.
-// The manager will automatically reconnect to the relay server in case of disconnection.
+// Serve starts the manager. It will establish a connection to the relay server and start the relay cleanup loop for
+// the unused relay connections. The manager will automatically reconnect to the relay server in case of disconnection.
 func (m *Manager) Serve() error {
-	if m.running {
+	if m.relayClient != nil {
 		return fmt.Errorf("manager already serving")
 	}
-	m.running = true
-	log.Debugf("starting relay client manager with %v relay servers", m.serverPicker.ServerURLs.Load())
+	log.Debugf("starting relay client manager with %v relay servers", m.serverURLs)
 
-	client, err := m.serverPicker.PickServer(m.ctx)
-	if err != nil {
-		go m.reconnectGuard.StartReconnectTrys(m.ctx, nil)
-	} else {
-		m.storeClient(client)
+	sp := ServerPicker{
+		TokenStore: m.tokenStore,
+		PeerID:     m.peerID,
 	}
 
-	go m.listenGuardEvent(m.ctx)
-	go m.startCleanupLoop()
-	return err
+	client, err := sp.PickServer(m.ctx, m.serverURLs)
+	if err != nil {
+		return err
+	}
+	m.relayClient = client
+
+	m.reconnectGuard = NewGuard(m.ctx, m.relayClient)
+	m.relayClient.SetOnConnectedListener(m.onServerConnected)
+	m.relayClient.SetOnDisconnectListener(func() {
+		m.onServerDisconnected(client.connectionURL)
+	})
+	m.startCleanupLoop()
+	return nil
 }
 
 // OpenConn opens a connection to the given peer key. If the peer is on the same relay server, the connection will be
 // established via the relay server. If the peer is on a different relay server, the manager will establish a new
 // connection to the relay server. It returns back with a net.Conn what represent the remote peer connection.
 func (m *Manager) OpenConn(serverAddress, peerKey string) (net.Conn, error) {
-	m.relayClientMu.Lock()
-	defer m.relayClientMu.Unlock()
-
 	if m.relayClient == nil {
 		return nil, ErrRelayClientNotConnected
 	}
@@ -155,9 +146,6 @@ func (m *Manager) OpenConn(serverAddress, peerKey string) (net.Conn, error) {
 
 // Ready returns true if the home Relay client is connected to the relay server.
 func (m *Manager) Ready() bool {
-	m.relayClientMu.Lock()
-	defer m.relayClientMu.Unlock()
-
 	if m.relayClient == nil {
 		return false
 	}
@@ -171,13 +159,6 @@ func (m *Manager) SetOnReconnectedListener(f func()) {
 // AddCloseListener adds a listener to the given server instance address. The listener will be called if the connection
 // closed.
 func (m *Manager) AddCloseListener(serverAddress string, onClosedListener OnServerCloseListener) error {
-	m.relayClientMu.Lock()
-	defer m.relayClientMu.Unlock()
-
-	if m.relayClient == nil {
-		return ErrRelayClientNotConnected
-	}
-
 	foreign, err := m.isForeignServer(serverAddress)
 	if err != nil {
 		return err
@@ -196,9 +177,6 @@ func (m *Manager) AddCloseListener(serverAddress string, onClosedListener OnServ
 // RelayInstanceAddress returns the address of the permanent relay server. It could change if the network connection is
 // lost. This address will be sent to the target peer to choose the common relay server for the communication.
 func (m *Manager) RelayInstanceAddress() (string, error) {
-	m.relayClientMu.Lock()
-	defer m.relayClientMu.Unlock()
-
 	if m.relayClient == nil {
 		return "", ErrRelayClientNotConnected
 	}
@@ -207,18 +185,13 @@ func (m *Manager) RelayInstanceAddress() (string, error) {
 
 // ServerURLs returns the addresses of the relay servers.
 func (m *Manager) ServerURLs() []string {
-	return m.serverPicker.ServerURLs.Load().([]string)
+	return m.serverURLs
 }
 
 // HasRelayAddress returns true if the manager is serving. With this method can check if the peer can communicate with
 // Relay service.
 func (m *Manager) HasRelayAddress() bool {
-	return len(m.serverPicker.ServerURLs.Load().([]string)) > 0
-}
-
-func (m *Manager) UpdateServerURLs(serverURLs []string) {
-	log.Infof("update relay server URLs: %v", serverURLs)
-	m.serverPicker.ServerURLs.Store(serverURLs)
+	return len(m.serverURLs) > 0
 }
 
 // UpdateToken updates the token in the token store.
@@ -272,7 +245,9 @@ func (m *Manager) openConnVia(serverAddress, peerKey string) (net.Conn, error) {
 		return nil, err
 	}
 	// if connection closed then delete the relay client from the list
-	relayClient.SetOnDisconnectListener(m.onServerDisconnected)
+	relayClient.SetOnDisconnectListener(func() {
+		m.onServerDisconnected(serverAddress)
+	})
 	rt.relayClient = relayClient
 	rt.Unlock()
 
@@ -290,37 +265,14 @@ func (m *Manager) onServerConnected() {
 	go m.onReconnectedListenerFn()
 }
 
-// onServerDisconnected start to reconnection for home server only
 func (m *Manager) onServerDisconnected(serverAddress string) {
-	m.relayClientMu.Lock()
 	if serverAddress == m.relayClient.connectionURL {
-		go m.reconnectGuard.StartReconnectTrys(m.ctx, m.relayClient)
+		go m.reconnectGuard.OnDisconnected()
 	}
-	m.relayClientMu.Unlock()
 
 	m.notifyOnDisconnectListeners(serverAddress)
 }
 
-func (m *Manager) listenGuardEvent(ctx context.Context) {
-	for {
-		select {
-		case rc := <-m.reconnectGuard.OnNewRelayClient:
-			m.storeClient(rc)
-		case <-ctx.Done():
-			return
-		}
-	}
-}
-
-func (m *Manager) storeClient(client *Client) {
-	m.relayClientMu.Lock()
-	defer m.relayClientMu.Unlock()
-
-	m.relayClient = client
-	m.relayClient.SetOnConnectedListener(m.onServerConnected)
-	m.relayClient.SetOnDisconnectListener(m.onServerDisconnected)
-}
-
 func (m *Manager) isForeignServer(address string) (bool, error) {
 	rAddr, err := m.relayClient.ServerInstanceURL()
 	if err != nil {
@@ -330,16 +282,22 @@ func (m *Manager) isForeignServer(address string) (bool, error) {
 }
 
 func (m *Manager) startCleanupLoop() {
-	ticker := time.NewTicker(relayCleanupInterval)
-	defer ticker.Stop()
-	for {
-		select {
-		case <-m.ctx.Done():
-			return
-		case <-ticker.C:
-			m.cleanUpUnusedRelays()
-		}
+	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() {

relay/client/picker.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"sync/atomic"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -13,13 +12,10 @@ import (
 )
 
 const (
+	connectionTimeout    = 30 * time.Second
 	maxConcurrentServers = 7
 )
 
-var (
-	connectionTimeout = 30 * time.Second
-)
-
 type connResult struct {
 	RelayClient *Client
 	Url         string
@@ -28,22 +24,20 @@ type connResult struct {
 
 type ServerPicker struct {
 	TokenStore *auth.TokenStore
-	ServerURLs atomic.Value
 	PeerID     string
 }
 
-func (sp *ServerPicker) PickServer(parentCtx context.Context) (*Client, error) {
+func (sp *ServerPicker) PickServer(parentCtx context.Context, urls []string) (*Client, error) {
 	ctx, cancel := context.WithTimeout(parentCtx, connectionTimeout)
 	defer cancel()
 
-	totalServers := len(sp.ServerURLs.Load().([]string))
+	totalServers := len(urls)
 
 	connResultChan := make(chan connResult, totalServers)
 	successChan := make(chan connResult, 1)
 	concurrentLimiter := make(chan struct{}, maxConcurrentServers)
 
-	log.Debugf("pick server from list: %v", sp.ServerURLs.Load().([]string))
-	for _, url := range sp.ServerURLs.Load().([]string) {
+	for _, url := range urls {
 		// todo check if we have a successful connection so we do not need to connect to other servers
 		concurrentLimiter <- struct{}{}
 		go func(url string) {
@@ -84,7 +78,7 @@ func (sp *ServerPicker) processConnResults(resultChan chan connResult, successCh
 	for numOfResults := 0; numOfResults < cap(resultChan); numOfResults++ {
 		cr := <-resultChan
 		if cr.Err != nil {
-			log.Tracef("failed to connect to Relay server: %s: %v", cr.Url, cr.Err)
+			log.Debugf("failed to connect to Relay server: %s: %v", cr.Url, cr.Err)
 			continue
 		}
 		log.Infof("connected to Relay server: %s", cr.Url)

relay/client/picker_test.go

@@ -7,19 +7,16 @@ import (
 )
 
 func TestServerPicker_UnavailableServers(t *testing.T) {
-	connectionTimeout = 5 * time.Second
-
 	sp := ServerPicker{
 		TokenStore: nil,
 		PeerID:     "test",
 	}
-	sp.ServerURLs.Store([]string{"rel://dummy1", "rel://dummy2"})
 
 	ctx, cancel := context.WithTimeout(context.Background(), connectionTimeout+1)
 	defer cancel()
 
 	go func() {
-		_, err := sp.PickServer(ctx)
+		_, err := sp.PickServer(ctx, []string{"rel://dummy1", "rel://dummy2"})
 		if err == nil {
 			t.Error(err)
 		}

relay/server/listener/quic/conn.go

@@ -1,112 +0,0 @@
-package quic
-
-import (
-	"context"
-	"errors"
-	"io"
-	"net"
-	"sync"
-	"time"
-
-	"github.com/quic-go/quic-go"
-)
-
-const (
-	writeTimeout = 10 * time.Second
-)
-
-type Conn struct {
-	session   quic.Connection
-	closed    bool
-	closedMu  sync.Mutex
-	ctx       context.Context
-	ctxCancel context.CancelFunc
-}
-
-func NewConn(session quic.Connection) *Conn {
-	ctx, cancel := context.WithCancel(context.Background())
-	return &Conn{
-		session:   session,
-		ctx:       ctx,
-		ctxCancel: cancel,
-	}
-}
-
-func (c *Conn) Read(b []byte) (n int, err error) {
-	if c.isClosed() {
-		return 0, io.EOF
-	}
-
-	dgram, err := c.session.ReceiveDatagram(c.ctx)
-	if err != nil {
-		return 0, c.ioErrHandling(err)
-	}
-	// Copy data to b, ensuring we don’t exceed the size of b
-	n = copy(b, dgram)
-	return n, nil
-}
-
-func (c *Conn) Write(b []byte) (int, error) {
-	err := c.session.SendDatagram(b)
-	return len(b), err
-}
-
-func (c *Conn) LocalAddr() net.Addr {
-	return c.session.LocalAddr()
-}
-
-func (c *Conn) RemoteAddr() net.Addr {
-	return c.session.RemoteAddr()
-}
-
-func (c *Conn) SetReadDeadline(t time.Time) error {
-	return nil
-}
-
-func (c *Conn) SetWriteDeadline(t time.Time) error {
-	return nil
-}
-
-func (c *Conn) SetDeadline(t time.Time) error {
-	return nil
-}
-
-func (c *Conn) Close() error {
-	c.closedMu.Lock()
-	if c.closed {
-		c.closedMu.Unlock()
-		return nil
-	}
-	c.closed = true
-	c.closedMu.Unlock()
-
-	c.ctxCancel() // Cancel the context
-
-	sessionErr := c.session.CloseWithError(0, "normal closure")
-	return sessionErr
-}
-
-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
-	}
-
-	// Handle QUIC-specific errors
-	if err == nil {
-		return nil
-	}
-
-	// Check if the connection was closed remotely
-	var appErr *quic.ApplicationError
-	if errors.As(err, &appErr) && appErr.ErrorCode == 0 { // 0 is normal closure
-		return io.EOF
-	}
-
-	return err
-}

relay/server/listener/quic/listener.go

@@ -1,66 +0,0 @@
-package quic
-
-import (
-	"context"
-	"crypto/tls"
-	"errors"
-	"fmt"
-	"net"
-
-	"github.com/quic-go/quic-go"
-	log "github.com/sirupsen/logrus"
-)
-
-type Listener struct {
-	// Address is the address to listen on
-	Address string
-	// TLSConfig is the TLS configuration for the server
-	TLSConfig *tls.Config
-
-	listener *quic.Listener
-	acceptFn func(conn net.Conn)
-}
-
-func (l *Listener) Listen(acceptFn func(conn net.Conn)) error {
-	l.acceptFn = acceptFn
-
-	quicCfg := &quic.Config{
-		EnableDatagrams: true,
-	}
-	listener, err := quic.ListenAddr(l.Address, l.TLSConfig, quicCfg)
-	if err != nil {
-		return fmt.Errorf("failed to create QUIC listener: %v", err)
-	}
-
-	l.listener = listener
-	log.Infof("QUIC server listening on address: %s", l.Address)
-
-	for {
-		session, err := listener.Accept(context.Background())
-		if err != nil {
-			if errors.Is(err, quic.ErrServerClosed) {
-				return nil
-			}
-
-			log.Errorf("Failed to accept QUIC session: %v", err)
-			continue
-		}
-
-		log.Infof("QUIC client connected from: %s", session.RemoteAddr())
-		conn := NewConn(session)
-		l.acceptFn(conn)
-	}
-}
-
-func (l *Listener) Shutdown(ctx context.Context) error {
-	if l.listener == nil {
-		return nil
-	}
-
-	log.Infof("stopping QUIC listener")
-	if err := l.listener.Close(); err != nil {
-		return fmt.Errorf("listener shutdown failed: %v", err)
-	}
-	log.Infof("QUIC listener stopped")
-	return nil
-}

relay/server/listener/ws/listener.go

@@ -23,7 +23,6 @@ type Listener struct {
 
 	server   *http.Server
 	acceptFn func(conn net.Conn)
-	log      *log.Entry
 }
 
 func (l *Listener) Listen(acceptFn func(conn net.Conn)) error {
@@ -89,8 +88,6 @@ func (l *Listener) onAccept(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	log.Infof("WS client connected from: %s", rAddr)
-
 	conn := NewConn(wsConn, lAddr, rAddr)
 	l.acceptFn(conn)
 }

relay/server/peer.go

@@ -16,8 +16,6 @@ import (
 
 const (
 	bufferSize = 8820
-
-	errCloseConn = "failed to close connection to peer: %s"
 )
 
 // Peer represents a peer connection
@@ -48,12 +46,6 @@ func NewPeer(metrics *metrics.Metrics, id []byte, conn net.Conn, store *Store) *
 // It manages the protocol (healthcheck, transport, close). Read the message and determine the message type and handle
 // the message accordingly.
 func (p *Peer) Work() {
-	defer func() {
-		if err := p.conn.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
-			p.log.Errorf(errCloseConn, err)
-		}
-	}()
-
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
@@ -105,7 +97,7 @@ func (p *Peer) handleMsgType(ctx context.Context, msgType messages.MsgType, hc *
 	case messages.MsgTypeClose:
 		p.log.Infof("peer exited gracefully")
 		if err := p.conn.Close(); err != nil {
-			log.Errorf(errCloseConn, err)
+			log.Errorf("failed to close connection to peer: %s", err)
 		}
 	default:
 		p.log.Warnf("received unexpected message type: %s", msgType)
@@ -129,8 +121,9 @@ func (p *Peer) CloseGracefully(ctx context.Context) {
 		p.log.Errorf("failed to send close message to peer: %s", p.String())
 	}
 
-	if err := p.conn.Close(); err != nil {
-		p.log.Errorf(errCloseConn, err)
+	err = p.conn.Close()
+	if err != nil {
+		p.log.Errorf("failed to close connection to peer: %s", err)
 	}
 }
 
@@ -139,7 +132,7 @@ func (p *Peer) Close() {
 	defer p.connMu.Unlock()
 
 	if err := p.conn.Close(); err != nil {
-		p.log.Errorf(errCloseConn, err)
+		p.log.Errorf("failed to close connection to peer: %s", err)
 	}
 }
 

relay/server/server.go

@@ -2,25 +2,13 @@ package server
 
 import (
 	"context"
-	"crypto/rand"
-	"crypto/rsa"
 	"crypto/tls"
-	"crypto/x509"
-	"crypto/x509/pkix"
-	"encoding/pem"
-	"math/big"
-	"net"
-	"sync"
-	"time"
 
-	"github.com/hashicorp/go-multierror"
 	log "github.com/sirupsen/logrus"
 	"go.opentelemetry.io/otel/metric"
 
-	nberrors "github.com/netbirdio/netbird/client/errors"
 	"github.com/netbirdio/netbird/relay/auth"
 	"github.com/netbirdio/netbird/relay/server/listener"
-	"github.com/netbirdio/netbird/relay/server/listener/quic"
 	"github.com/netbirdio/netbird/relay/server/listener/ws"
 )
 
@@ -36,8 +24,8 @@ type ListenerConfig struct {
 // It is the gate between the WebSocket listener and the Relay server logic.
 // In a new HTTP connection, the server will accept the connection and pass it to the Relay server via the Accept method.
 type Server struct {
-	relay     *Relay
-	listeners []listener.Listener
+	relay      *Relay
+	wSListener listener.Listener
 }
 
 // NewServer creates a new relay server instance.
@@ -51,120 +39,38 @@ func NewServer(meter metric.Meter, exposedAddress string, tlsSupport bool, authV
 		return nil, err
 	}
 	return &Server{
-		relay:     relay,
-		listeners: make([]listener.Listener, 0, 2),
+		relay: relay,
 	}, nil
 }
 
 // Listen starts the relay server.
 func (r *Server) Listen(cfg ListenerConfig) error {
-	wSListener := &ws.Listener{
+	r.wSListener = &ws.Listener{
 		Address:   cfg.Address,
 		TLSConfig: cfg.TLSConfig,
 	}
-	r.listeners = append(r.listeners, wSListener)
 
-	quicListener := &quic.Listener{
-		Address: cfg.Address,
+	wslErr := r.wSListener.Listen(r.relay.Accept)
+	if wslErr != nil {
+		log.Errorf("failed to bind ws server: %s", wslErr)
 	}
 
-	if cfg.TLSConfig != nil {
-		quicListener.TLSConfig = cfg.TLSConfig
-	} else {
-		tlsConfig, err := generateTestTLSConfig()
-		if err != nil {
-			return err
-		}
-		quicListener.TLSConfig = tlsConfig
-	}
-	r.listeners = append(r.listeners, quicListener)
-
-	errChan := make(chan error, len(r.listeners))
-	wg := sync.WaitGroup{}
-	for _, l := range r.listeners {
-		wg.Add(1)
-		go func(listener listener.Listener) {
-			defer wg.Done()
-			errChan <- listener.Listen(r.relay.Accept)
-		}(l)
-	}
-
-	wg.Wait()
-	close(errChan)
-	var multiErr *multierror.Error
-	for err := range errChan {
-		multiErr = multierror.Append(multiErr, err)
-	}
-
-	return nberrors.FormatErrorOrNil(multiErr)
+	return wslErr
 }
 
 // Shutdown stops the relay server. If there are active connections, they will be closed gracefully. In case of a context,
 // the connections will be forcefully closed.
-func (r *Server) Shutdown(ctx context.Context) error {
-	var multiErr *multierror.Error
-	for _, l := range r.listeners {
-		if err := l.Shutdown(ctx); err != nil {
-			multiErr = multierror.Append(multiErr, err)
-		}
+func (r *Server) Shutdown(ctx context.Context) (err error) {
+	// stop service new connections
+	if r.wSListener != nil {
+		err = r.wSListener.Shutdown(ctx)
 	}
 
 	r.relay.Shutdown(ctx)
-	return nberrors.FormatErrorOrNil(multiErr)
+	return
 }
 
 // InstanceURL returns the instance URL of the relay server.
 func (r *Server) InstanceURL() string {
 	return r.relay.instanceURL
 }
-
-// GenerateTestTLSConfig creates a self-signed certificate for testing
-func generateTestTLSConfig() (*tls.Config, error) {
-	log.Infof("generating test TLS config")
-	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
-	if err != nil {
-		return nil, err
-	}
-
-	template := x509.Certificate{
-		SerialNumber: big.NewInt(1),
-		Subject: pkix.Name{
-			Organization: []string{"Test Organization"},
-		},
-		NotBefore: time.Now(),
-		NotAfter:  time.Now().Add(time.Hour * 24 * 180), // Valid for 180 days
-		KeyUsage:  x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
-		ExtKeyUsage: []x509.ExtKeyUsage{
-			x509.ExtKeyUsageServerAuth,
-		},
-		BasicConstraintsValid: true,
-		DNSNames:              []string{"localhost"},
-		IPAddresses:           []net.IP{net.ParseIP("192.168.0.10")},
-	}
-
-	// Create certificate
-	certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, &privateKey.PublicKey, privateKey)
-	if err != nil {
-		return nil, err
-	}
-
-	certPEM := pem.EncodeToMemory(&pem.Block{
-		Type:  "CERTIFICATE",
-		Bytes: certDER,
-	})
-
-	privateKeyPEM := pem.EncodeToMemory(&pem.Block{
-		Type:  "RSA PRIVATE KEY",
-		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
-	})
-
-	tlsCert, err := tls.X509KeyPair(certPEM, privateKeyPEM)
-	if err != nil {
-		return nil, err
-	}
-
-	return &tls.Config{
-		Certificates: []tls.Certificate{tlsCert},
-		NextProtos:   []string{"netbird-relay"},
-	}, nil
-}

util/file.go

@@ -15,7 +15,7 @@ import (
 )
 
 // WriteJsonWithRestrictedPermission writes JSON config object to a file. Enforces permission on the parent directory
-func WriteJsonWithRestrictedPermission(ctx context.Context, file string, obj interface{}) error {
+func WriteJsonWithRestrictedPermission(file string, obj interface{}) error {
 	configDir, configFileName, err := prepareConfigFileDir(file)
 	if err != nil {
 		return err
@@ -26,18 +26,18 @@ func WriteJsonWithRestrictedPermission(ctx context.Context, file string, obj int
 		return err
 	}
 
-	return writeJson(ctx, file, obj, configDir, configFileName)
+	return writeJson(file, obj, configDir, configFileName)
 }
 
 // WriteJson writes JSON config object to a file creating parent directories if required
 // The output JSON is pretty-formatted
-func WriteJson(ctx context.Context, file string, obj interface{}) error {
+func WriteJson(file string, obj interface{}) error {
 	configDir, configFileName, err := prepareConfigFileDir(file)
 	if err != nil {
 		return err
 	}
 
-	return writeJson(ctx, file, obj, configDir, configFileName)
+	return writeJson(file, obj, configDir, configFileName)
 }
 
 // DirectWriteJson writes JSON config object to a file creating parent directories if required without creating a temporary file
@@ -79,11 +79,7 @@ func DirectWriteJson(ctx context.Context, file string, obj interface{}) error {
 	return nil
 }
 
-func writeJson(ctx context.Context, file string, obj interface{}, configDir string, configFileName string) error {
-	// Check context before expensive operations
-	if ctx.Err() != nil {
-		return ctx.Err()
-	}
+func writeJson(file string, obj interface{}, configDir string, configFileName string) error {
 
 	// make it pretty
 	bs, err := json.MarshalIndent(obj, "", "    ")
@@ -91,10 +87,6 @@ func writeJson(ctx context.Context, file string, obj interface{}, configDir stri
 		return err
 	}
 
-	if ctx.Err() != nil {
-		return ctx.Err()
-	}
-
 	tempFile, err := os.CreateTemp(configDir, ".*"+configFileName)
 	if err != nil {
 		return err
@@ -119,11 +111,6 @@ func writeJson(ctx context.Context, file string, obj interface{}, configDir stri
 		return err
 	}
 
-	// Check context again
-	if ctx.Err() != nil {
-		return ctx.Err()
-	}
-
 	err = os.Rename(tempFileName, file)
 	if err != nil {
 		return err

util/file_test.go

@@ -1,7 +1,6 @@
 package util
 
 import (
-	"context"
 	"crypto/md5"
 	"encoding/hex"
 	"io"
@@ -40,7 +39,7 @@ func TestConfigJSON(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			tmpDir := t.TempDir()
 
-			err := WriteJson(context.Background(), tmpDir+"/testconfig.json", tt.config)
+			err := WriteJson(tmpDir+"/testconfig.json", tt.config)
 			require.NoError(t, err)
 
 			read, err := ReadJson(tmpDir+"/testconfig.json", &TestConfig{})
@@ -74,7 +73,7 @@ func TestCopyFileContents(t *testing.T) {
 			src := tmpDir + "/copytest_src"
 			dst := tmpDir + "/copytest_dst"
 
-			err := WriteJson(context.Background(), src, tt.srcContent)
+			err := WriteJson(src, tt.srcContent)
 			require.NoError(t, err)
 
 			err = CopyFileContents(src, dst)
@@ -128,7 +127,7 @@ func TestHandleConfigFileWithoutFullPath(t *testing.T) {
 				_ = os.Remove(cfgFile)
 			}()
 
-			err := WriteJson(context.Background(), cfgFile, tt.config)
+			err := WriteJson(cfgFile, tt.config)
 			require.NoError(t, err)
 
 			read, err := ReadJson(cfgFile, &TestConfig{})