From 36339993bc4aacd22448910ed39ec0c5fceef18d Mon Sep 17 00:00:00 2001
From: Riccardo Manfrin <3090891+riccardomanfrin@users.noreply.github.com>
Date: Tue, 27 Oct 2026 14:01:16 +0100
Subject: [PATCH 1/5] [client] Mark Rosenpass as deprecated (#8372)
* [client] Mark Rosenpass flags as deprecated
Rosenpass support is being deprecated. Users enabling it from the CLI
should learn that before they depend on it, while the feature itself
keeps working.
* [client] Show Rosenpass deprecation notice in the UI settings
Desktop users enable Rosenpass from the settings toggle, not the CLI,
so the deprecation notice needs to reach them there too. A new string
key falls back to English until Crowdin translates it, instead of
showing a stale translation without the notice.
* [client] Keep switchOrCreateProfile doc comment attached
The Rosenpass warning helper landed between switchOrCreateProfile and
its doc comment, leaving the function undocumented.
* [client] Use MarkDeprecated for the Rosenpass flags
Follows the convention already used for --enable-lazy-connection, as
asked in review. The flags drop out of --help and pflag warns on any
use, so the custom helper is no longer needed.
* [client] Isolate the Rosenpass flag test from upCmd's flag set
Setting the shared upCmd flags left them marked Changed for any later
test in the package.
---
client/cmd/root.go | 6 ++--
client/cmd/up_rosenpass_test.go | 34 +++++++++++++++++++
.../src/modules/settings/SettingsSecurity.tsx | 9 ++++-
client/ui/i18n/locales/en/common.json | 4 +++
4 files changed, 50 insertions(+), 3 deletions(-)
create mode 100644 client/cmd/up_rosenpass_test.go
diff --git a/client/cmd/root.go b/client/cmd/root.go
index f976ca895..3c49f8cca 100644
--- a/client/cmd/root.go
+++ b/client/cmd/root.go
@@ -220,13 +220,15 @@ func init() {
`An empty string "" clears the previous configuration. `+
`E.g. --dns-resolver-address 127.0.0.1:5053 or --dns-resolver-address ""`,
)
- upCmd.PersistentFlags().BoolVar(&rosenpassEnabled, enableRosenpassFlag, false, "[Experimental] Enable Rosenpass feature. If enabled, the connection will be post-quantum secured via Rosenpass.")
- upCmd.PersistentFlags().BoolVar(&rosenpassPermissive, rosenpassPermissiveFlag, false, "[Experimental] Enable Rosenpass in permissive mode to allow this peer to accept WireGuard connections without requiring Rosenpass functionality from peers that do not have Rosenpass enabled.")
+ upCmd.PersistentFlags().BoolVar(&rosenpassEnabled, enableRosenpassFlag, false, "[Deprecated] Enable Rosenpass feature. If enabled, the connection will be post-quantum secured via Rosenpass.")
+ upCmd.PersistentFlags().BoolVar(&rosenpassPermissive, rosenpassPermissiveFlag, false, "[Deprecated] Enable Rosenpass in permissive mode to allow this peer to accept WireGuard connections without requiring Rosenpass functionality from peers that do not have Rosenpass enabled.")
upCmd.PersistentFlags().BoolVar(&autoConnectDisabled, disableAutoConnectFlag, false, "Disables auto-connect feature. If enabled, then the client won't connect automatically when the service starts.")
upCmd.PersistentFlags().BoolVar(&localMetricsEnabled, enableLocalMetricsFlag, false, "Enables a local Prometheus /metrics endpoint exposing connection state (peers, latency, P2P vs relay).")
upCmd.PersistentFlags().StringVar(&localMetricsAddr, localMetricsAddressFlag, localmetrics.DefaultListenAddress, "Listen address of the local Prometheus /metrics endpoint.")
upCmd.PersistentFlags().BoolVar(&lazyConnEnabled, enableLazyConnectionFlag, false, "Deprecated: no longer used. Lazy connections are controlled by the server and the NB_LAZY_CONN environment variable.")
_ = upCmd.PersistentFlags().MarkDeprecated(enableLazyConnectionFlag, "no longer used; lazy connections are controlled by the server and the NB_LAZY_CONN environment variable")
+ _ = upCmd.PersistentFlags().MarkDeprecated(enableRosenpassFlag, "Rosenpass support may be removed in a future release")
+ _ = upCmd.PersistentFlags().MarkDeprecated(rosenpassPermissiveFlag, "Rosenpass support may be removed in a future release")
}
diff --git a/client/cmd/up_rosenpass_test.go b/client/cmd/up_rosenpass_test.go
new file mode 100644
index 000000000..fd36df20f
--- /dev/null
+++ b/client/cmd/up_rosenpass_test.go
@@ -0,0 +1,34 @@
+package cmd
+
+import (
+ "bytes"
+ "testing"
+
+ "github.com/spf13/pflag"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func TestRosenpassFlagsWarnDeprecated(t *testing.T) {
+ origEnabled, origPermissive := rosenpassEnabled, rosenpassPermissive
+ t.Cleanup(func() {
+ rosenpassEnabled, rosenpassPermissive = origEnabled, origPermissive
+ })
+
+ for _, name := range []string{enableRosenpassFlag, rosenpassPermissiveFlag} {
+ // Copy the production flag into a private set so Set marks only the copy as
+ // Changed, leaving upCmd's shared flag set untouched for other tests.
+ prod := upCmd.PersistentFlags().Lookup(name)
+ require.NotNil(t, prod, "flag --%s must be registered on up", name)
+ flagCopy := *prod
+
+ var out bytes.Buffer
+ flags := pflag.NewFlagSet("rosenpass-test", pflag.ContinueOnError)
+ flags.SetOutput(&out)
+ flags.AddFlag(&flagCopy)
+
+ require.NoError(t, flags.Set(name, "true"))
+ assert.Contains(t, out.String(), "Flag --"+name+" has been deprecated", "setting --%s should print the deprecation notice", name)
+ assert.False(t, prod.Changed, "the shared --%s flag must stay unchanged", name)
+ }
+}
diff --git a/client/ui/frontend/src/modules/settings/SettingsSecurity.tsx b/client/ui/frontend/src/modules/settings/SettingsSecurity.tsx
index db1ad9c8f..eacf0ebfe 100644
--- a/client/ui/frontend/src/modules/settings/SettingsSecurity.tsx
+++ b/client/ui/frontend/src/modules/settings/SettingsSecurity.tsx
@@ -42,7 +42,14 @@ export function SettingsSecurity() {
if (!v) setField("rosenpassPermissive", false);
}}
label={t("settings.security.rosenpass.label")}
- helpText={t("settings.security.rosenpass.help")}
+ helpText={
+ <>
+ {t("settings.security.rosenpass.help")}
+
+ {t("settings.security.rosenpass.deprecated")}
+
+
+ }
/>
)}
{!hideRosenpassPermissive && (
diff --git a/client/ui/i18n/locales/en/common.json b/client/ui/i18n/locales/en/common.json
index 94d741b3e..fb3af16ea 100644
--- a/client/ui/i18n/locales/en/common.json
+++ b/client/ui/i18n/locales/en/common.json
@@ -899,6 +899,10 @@
"message": "Add a post-quantum key exchange via Rosenpass on top of WireGuard®.",
"description": "Helper text: adds a post-quantum key exchange via Rosenpass on top of WireGuard®. 'Rosenpass' and 'WireGuard®' are product names — do not translate; keep the ® symbol."
},
+ "settings.security.rosenpass.deprecated": {
+ "message": "Deprecated: Rosenpass support may be removed in a future release.",
+ "description": "Notice shown under the quantum-resistance toggle: the Rosenpass feature is deprecated and may be removed later. 'Rosenpass' is a product name — do not translate."
+ },
"settings.security.rosenpassPermissive.label": {
"message": "Enable Permissive Mode",
"description": "Toggle label: enable permissive mode (for quantum-resistance)."
From 09f1de12e81f53d2046407e8bf54d67f4e031c13 Mon Sep 17 00:00:00 2001
From: Pascal Fischer <32096965+pascal-fischer@users.noreply.github.com>
Date: Tue, 27 Oct 2026 16:24:46 +0100
Subject: [PATCH 2/5] [management] batch disconnect and lock free reconnect
(#8373)
---
management/internals/server/boot.go | 28 +-
management/internals/server/container.go | 37 ++-
management/internals/server/server.go | 53 +++-
management/internals/shared/grpc/server.go | 133 +++++++--
.../shared/grpc/server_disconnect_test.go | 145 ++++++++++
management/server/account.go | 17 +-
management/server/account/manager.go | 4 +-
management/server/account/manager_mock.go | 36 ++-
management/server/account_test.go | 2 +-
management/server/management_proto_test.go | 129 +++++++++
management/server/mock_server/account_mock.go | 22 +-
management/server/peer.go | 272 ++++++++++++------
management/server/peer_test.go | 134 +++++++++
management/server/store/sql_store_peer.go | 46 ++-
.../server/store/sql_store_peer_test.go | 147 ++++++++++
management/server/store/store.go | 5 +
management/server/store/store_mock.go | 15 +
management/server/types/peer.go | 13 +
18 files changed, 1094 insertions(+), 144 deletions(-)
diff --git a/management/internals/server/boot.go b/management/internals/server/boot.go
index 64d5ee4bc..d37784fae 100644
--- a/management/internals/server/boot.go
+++ b/management/internals/server/boot.go
@@ -269,6 +269,9 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
grpc.KeepaliveParams(kasp),
grpc.ChainUnaryInterceptor(realip.UnaryServerInterceptorOpts(realipOpts...), unaryInterceptor, proxyUnary),
grpc.ChainStreamInterceptor(realip.StreamServerInterceptorOpts(realipOpts...), streamInterceptor, proxyStream),
+ // Stop waits for the sync handlers, so their teardown has queued every
+ // disconnect before FlushDisconnects writes them.
+ grpc.WaitForHandlers(true),
}
// Append interceptors contributed by registered gRPC extensions. These
@@ -292,6 +295,21 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
}
gRPCAPIHandler := grpc.NewServer(gRPCOpts...)
+ mgmtProto.RegisterManagementServiceServer(gRPCAPIHandler, s.ManagementGRPCServer())
+
+ mgmtProto.RegisterProxyServiceServer(gRPCAPIHandler, s.ReverseProxyGRPCServer())
+ log.Info("ProxyService registered on gRPC server")
+
+ // Register services contributed by external modules via the extension seam.
+ registerExtensions(gRPCAPIHandler, s.grpcExtensions)
+
+ return gRPCAPIHandler
+ })
+}
+
+// ManagementGRPCServer returns the management gRPC service, created once and served by GRPCServer.
+func (s *BaseServer) ManagementGRPCServer() *nbgrpc.Server {
+ return Create(s, func() *nbgrpc.Server {
srv, err := nbgrpc.NewServer(s.Config, s.AccountManager(), s.SettingsManager(), s.JobManager(), s.SecretsManager(), s.Metrics(), s.AuthManager(), s.IntegratedValidator(), s.NetworkMapController(), s.OAuthConfigProvider(), s.SessionStore())
if err != nil {
log.Fatalf("failed to create management server: %v", err)
@@ -301,15 +319,7 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
if serviceMgr != nil {
serviceMgr.StartExposeReaper(context.Background())
}
- mgmtProto.RegisterManagementServiceServer(gRPCAPIHandler, srv)
-
- mgmtProto.RegisterProxyServiceServer(gRPCAPIHandler, s.ReverseProxyGRPCServer())
- log.Info("ProxyService registered on gRPC server")
-
- // Register services contributed by external modules via the extension seam.
- registerExtensions(gRPCAPIHandler, s.grpcExtensions)
-
- return gRPCAPIHandler
+ return srv
})
}
diff --git a/management/internals/server/container.go b/management/internals/server/container.go
index ee9e65016..760f528a8 100644
--- a/management/internals/server/container.go
+++ b/management/internals/server/container.go
@@ -1,6 +1,9 @@
package server
-import "fmt"
+import (
+ "fmt"
+ "reflect"
+)
// Create a dependency and add it to the BaseServer's container. A string key identifier will be based on its type definition.
func Create[T any](s Server, createFunc func() T) T {
@@ -33,15 +36,41 @@ func InjectNamed[T any](c Server, name string, thing T) {
}
func maybeCreate[T any](s Server, createFunc func() T) (result T, isNew bool) {
- key := fmt.Sprintf("%T", (*T)(nil))[1:]
- return maybeCreateKeyed(s, key, createFunc)
+ return maybeCreateKeyed(s, typeKey(reflect.TypeFor[T]()), createFunc)
}
func maybeCreateNamed[T any](s Server, name string, createFunc func() T) (result T, isNew bool) {
- key := fmt.Sprintf("%T:%s", (*T)(nil), name)[1:]
+ key := fmt.Sprintf("%s:%s", typeKey(reflect.TypeFor[T]()), name)
return maybeCreateKeyed(s, key, createFunc)
}
+// typeKey names t by the full package path of every named type it is built from, so
+// same-named types of different packages, such as google.golang.org/grpc.Server and a
+// local grpc.Server, never share a container key.
+func typeKey(t reflect.Type) string {
+ if t.Name() != "" {
+ if t.PkgPath() == "" {
+ return t.Name()
+ }
+ return t.PkgPath() + "." + t.Name()
+ }
+
+ switch t.Kind() {
+ case reflect.Pointer:
+ return "*" + typeKey(t.Elem())
+ case reflect.Slice:
+ return "[]" + typeKey(t.Elem())
+ case reflect.Array:
+ return fmt.Sprintf("[%d]%s", t.Len(), typeKey(t.Elem()))
+ case reflect.Map:
+ return "map[" + typeKey(t.Key()) + "]" + typeKey(t.Elem())
+ case reflect.Chan:
+ return t.ChanDir().String() + " " + typeKey(t.Elem())
+ default:
+ return t.String()
+ }
+}
+
func maybeCreateKeyed[T any](s Server, key string, createFunc func() T) (result T, isNew bool) {
if t, ok := s.GetContainer(key); ok {
if t == nil {
diff --git a/management/internals/server/server.go b/management/internals/server/server.go
index c95092d70..fa92be978 100644
--- a/management/internals/server/server.go
+++ b/management/internals/server/server.go
@@ -3,6 +3,7 @@ package server
import (
"context"
"crypto/tls"
+ "errors"
"fmt"
"net"
"net/http"
@@ -40,6 +41,15 @@ const (
ContainerKeyBaseServer = "baseServer"
applicationName = "management"
+
+ // shutdownTimeout stays below the 10s that systemd (TimeoutStopSec) and Docker grant
+ // before killing the process.
+ shutdownTimeout = 9 * time.Second
+ grpcStopTimeout = 3 * time.Second
+ // shutdownCloseReserve is kept from the flush for closing the stores.
+ shutdownCloseReserve = 2 * time.Second
+ // shutdownIdPReserve is kept from the extension shutdown hooks for stopping the embedded IdP.
+ shutdownIdPReserve = 1 * time.Second
)
type Server interface {
@@ -297,9 +307,13 @@ func (s *BaseServer) setupTLS(ctx context.Context) (bool, error) {
}
}
-// Stop attempts a graceful shutdown, waiting up to 5 seconds for active connections to finish
+// Stop attempts a graceful shutdown within shutdownTimeout. Writing the disconnects of the
+// drained streams may use the budget left after stopping the services and gRPC streams,
+// except for shutdownCloseReserve, which stays for closing the stores; the extension hooks
+// leave shutdownIdPReserve for stopping the embedded IdP. It returns an error when the
+// disconnects could not all be written or the shutdown ran out of time.
func (s *BaseServer) Stop() error {
- ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+ ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
defer cancel()
defer s.RunStopHandlers()
if s.domainCleanupStop != nil {
@@ -318,12 +332,20 @@ func (s *BaseServer) Stop() error {
if s.certListener != nil {
_ = s.certListener.Close()
}
- s.GRPCServer().Stop()
+ s.ManagementGRPCServer().StartDraining()
+ s.stopGRPCServer()
+ deadline, _ := ctx.Deadline()
+ flushCtx, cancelFlush := context.WithDeadline(ctx, deadline.Add(-shutdownCloseReserve))
+ flushErr := s.ManagementGRPCServer().FlushDisconnects(flushCtx)
+ cancelFlush()
+
if s.proxyAuthClose != nil {
s.proxyAuthClose()
s.proxyAuthClose = nil
}
- runExtensionShutdownHooks(ctx, s.grpcExtensions)
+ hooksCtx, cancelHooks := context.WithDeadline(ctx, deadline.Add(-shutdownIdPReserve))
+ runExtensionShutdownHooks(hooksCtx, s.grpcExtensions)
+ cancelHooks()
_ = s.Store().Close(ctx)
_ = s.EventStore().Close(ctx)
if s.update != nil {
@@ -338,12 +360,29 @@ func (s *BaseServer) Stop() error {
case err := <-s.Errors():
if err != nil {
log.WithContext(ctx).Infof("stopped Management Service: %v", err)
- return nil
+ return flushErr
}
log.WithContext(ctx).Infof("stopped Management Service")
- return nil
+ return flushErr
case <-ctx.Done():
- return ctx.Err()
+ return errors.Join(flushErr, ctx.Err())
+ }
+}
+
+// stopGRPCServer stops the gRPC server, which waits for the stream handlers to finish
+// their teardown, but gives up waiting after grpcStopTimeout so a stuck handler cannot
+// block the shutdown.
+func (s *BaseServer) stopGRPCServer() {
+ stopped := make(chan struct{})
+ go func() {
+ s.GRPCServer().Stop()
+ close(stopped)
+ }()
+
+ select {
+ case <-stopped:
+ case <-time.After(grpcStopTimeout):
+ log.Warnf("gRPC server did not stop within %s, continuing shutdown", grpcStopTimeout)
}
}
diff --git a/management/internals/shared/grpc/server.go b/management/internals/shared/grpc/server.go
index b48a86d55..a30f4c505 100644
--- a/management/internals/shared/grpc/server.go
+++ b/management/internals/shared/grpc/server.go
@@ -89,6 +89,12 @@ type Server struct {
reverseProxyManager rpservice.Manager
reverseProxyMu sync.RWMutex
+
+ // drainMu guards draining and pendingDisconnects. While draining, ending sync streams
+ // queue their disconnect for FlushDisconnects instead of writing it.
+ drainMu sync.Mutex
+ draining bool
+ pendingDisconnects []types.PeerSession
}
// NewServer creates a new Management server
@@ -311,7 +317,14 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
metahash := metaHash(peerMeta)
s.loginFilter.addLogin(peerKey.String(), metahash)
- peer, netMap, postureChecks, dnsFwdPort, err := s.accountManager.SyncAndMarkPeer(ctx, accountID, peerKey.String(), peerMeta, realIP, syncStart)
+ settings, err := s.settingsManager.GetSettings(ctx, accountID, activity.SystemInitiator)
+ if err != nil {
+ log.WithContext(ctx).Debugf("error while getting settings for peer %s: %v", peerKey.String(), err)
+ s.syncSem.Add(-1)
+ return status.Errorf(codes.Internal, "error handling request")
+ }
+
+ peer, netMap, postureChecks, dnsFwdPort, err := s.accountManager.SyncAndMarkPeer(ctx, accountID, peerKey.String(), peerMeta, realIP, syncStart, settings)
if err != nil {
log.WithContext(ctx).Debugf("error while syncing peer %s: %v", peerKey.String(), err)
s.syncSem.Add(-1)
@@ -319,11 +332,11 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
}
trackChallenge := func() { s.accountManager.TrackCertificateChallenges(ctx, accountID, peer.ID, syncStart) }
- err = s.sendInitialSync(ctx, peerKey, peer, netMap, postureChecks, srv, dnsFwdPort, trackChallenge)
+ err = s.sendInitialSync(ctx, peerKey, peer, netMap, postureChecks, srv, dnsFwdPort, settings, trackChallenge)
if err != nil {
log.WithContext(ctx).Debugf("error while sending initial sync for %s: %v", peerKey.String(), err)
s.syncSem.Add(-1)
- s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil)
+ s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil, false)
return err
}
@@ -331,10 +344,17 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
if err != nil {
log.WithContext(ctx).Debugf("error while notify peer connected for %s: %v", peerKey.String(), err)
s.syncSem.Add(-1)
- s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil)
+ s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil, false)
return err
}
+ if err = s.accountManager.CheckPeerMaySync(ctx, accountID, peerKey.String(), settings); err != nil {
+ log.WithContext(ctx).Debugf("peer %s may no longer sync after connecting: %v", peerKey.String(), err)
+ s.syncSem.Add(-1)
+ s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, updates, false)
+ return mapError(ctx, err)
+ }
+
s.secretsManager.SetupRefresh(ctx, accountID, peer.ID)
unlock()
@@ -442,21 +462,28 @@ func (s *Server) sendJob(ctx context.Context, peerKey wgtypes.Key, job *job.Even
return nil
}
+// cancelPeerRoutines tears down the stream under the peer lock. While draining, the
+// disconnect is queued before waiting for the lock, so the shutdown flush covers the
+// stream even when the lock is held past the gRPC stop timeout.
func (s *Server) cancelPeerRoutines(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time, session chan *network_map.UpdateMessage) {
uncanceledCTX := context.WithoutCancel(ctx)
+ disconnectQueued := s.queueDisconnect(accountID, peer, streamStartTime)
unlock := s.acquirePeerLockByUID(uncanceledCTX, peer.Key)
defer unlock()
- s.cancelPeerRoutinesWithoutLock(uncanceledCTX, accountID, peer, streamStartTime, session)
+ s.cancelPeerRoutinesWithoutLock(uncanceledCTX, accountID, peer, streamStartTime, session, disconnectQueued)
}
// cancelPeerRoutinesWithoutLock tears down the stream of the session identified by streamStartTime
// and its updates channel. A nil session means the stream failed before it registered a channel;
-// the controller then closes any channel still registered.
-func (s *Server) cancelPeerRoutinesWithoutLock(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time, session chan *network_map.UpdateMessage) {
- err := s.accountManager.OnPeerDisconnected(ctx, accountID, peer.Key, streamStartTime)
- if err != nil {
- log.WithContext(ctx).Errorf("failed to disconnect peer %s properly: %v", peer.Key, err)
+// the controller then closes any channel still registered. disconnectQueued tells that the
+// caller already queued the disconnect for the shutdown flush.
+func (s *Server) cancelPeerRoutinesWithoutLock(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time, session chan *network_map.UpdateMessage, disconnectQueued bool) {
+ if !disconnectQueued && !s.queueDisconnect(accountID, peer, streamStartTime) {
+ err := s.accountManager.OnPeerDisconnected(ctx, accountID, peer.Key, streamStartTime)
+ if err != nil {
+ log.WithContext(ctx).Errorf("failed to disconnect peer %s properly: %v", peer.Key, err)
+ }
}
if !s.networkMapController.OnPeerDisconnected(ctx, accountID, peer.ID, session) {
log.WithContext(ctx).Debugf("skipped peer routines teardown for %s: a newer session owns the peer", peer.Key)
@@ -468,6 +495,60 @@ func (s *Server) cancelPeerRoutinesWithoutLock(ctx context.Context, accountID st
log.WithContext(ctx).Debugf("peer %s has been disconnected", peer.Key)
}
+// queueDisconnect queues the disconnect of the stream for FlushDisconnects while the
+// server drains and reports whether it did. Embedded proxy peers are never queued
+// because their disconnect refreshes the network maps of the peers using them.
+func (s *Server) queueDisconnect(accountID string, peer *nbpeer.Peer, streamStartTime time.Time) bool {
+ if peer.ProxyMeta.Embedded {
+ return false
+ }
+
+ s.drainMu.Lock()
+ defer s.drainMu.Unlock()
+ if !s.draining {
+ return false
+ }
+ s.pendingDisconnects = append(s.pendingDisconnects, types.PeerSession{
+ AccountID: accountID,
+ PeerID: peer.ID,
+ SessionStartedAt: streamStartTime.UnixNano(),
+ })
+ return true
+}
+
+// StartDraining makes ending sync streams queue their disconnect for FlushDisconnects
+// instead of writing one each.
+func (s *Server) StartDraining() {
+ s.drainMu.Lock()
+ defer s.drainMu.Unlock()
+ s.draining = true
+}
+
+// FlushDisconnects stops draining and marks the queued sessions disconnected in batched
+// writes. Streams ending afterwards write their own disconnect again. It returns an error
+// when not every queued session could be written.
+func (s *Server) FlushDisconnects(ctx context.Context) error {
+ s.drainMu.Lock()
+ sessions := s.pendingDisconnects
+ s.pendingDisconnects = nil
+ s.draining = false
+ s.drainMu.Unlock()
+
+ if len(sessions) == 0 {
+ return nil
+ }
+ err := s.accountManager.MarkPeersDisconnected(ctx, sessions)
+ if err != nil && ctx.Err() == nil {
+ // Retrying the whole batch is safe: rows already marked no longer match their session token.
+ log.WithContext(ctx).Warnf("retrying to mark %d peers disconnected on shutdown: %v", len(sessions), err)
+ err = s.accountManager.MarkPeersDisconnected(ctx, sessions)
+ }
+ if err != nil {
+ return fmt.Errorf("mark %d peers disconnected on shutdown: %w", len(sessions), err)
+ }
+ return nil
+}
+
func (s *Server) validateToken(ctx context.Context, peerKey, jwtToken string) (string, error) {
if s.authManager == nil {
return "", status.Errorf(codes.Internal, "missing auth manager")
@@ -688,6 +769,7 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
//nolint
ctx = context.WithValue(ctx, nbContext.PeerIDKey, peerKey.String())
accountID, err := s.accountManager.GetAccountIDForPeerKey(ctx, peerKey.String())
+ accountKnown := err == nil
if err != nil {
// this case should not happen and already indicates an issue but we don't want the system to fail due to being unable to log in detail
accountID = "UNKNOWN"
@@ -714,6 +796,15 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
sshKey = loginReq.GetPeerKeys().GetSshPubKey()
}
+ var settings *types.Settings
+ if accountKnown {
+ settings, err = s.settingsManager.GetSettings(ctx, accountID, activity.SystemInitiator)
+ if err != nil {
+ log.WithContext(ctx).Warnf("failed getting settings for peer %s: %s", peerKey, err)
+ return nil, status.Errorf(codes.Internal, "failed logging in peer")
+ }
+ }
+
peer, network, postureChecks, enableSSH, err := s.accountManager.LoginPeer(ctx, types.PeerLogin{
WireGuardPubKey: peerKey.String(),
SSHKey: cleanClientString(string(sshKey)),
@@ -722,6 +813,7 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
SetupKey: loginReq.GetSetupKey(),
ConnectionIP: realIP,
ExtraDNSLabels: loginReq.GetDnsLabels(),
+ Settings: settings,
})
if err != nil {
if errors.Is(err, internalStatus.ErrNoAuthMethodProvided) {
@@ -732,7 +824,7 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
return nil, mapError(ctx, err)
}
- loginResp, err := s.prepareLoginResponse(ctx, peer, network, postureChecks, enableSSH)
+ loginResp, err := s.prepareLoginResponse(ctx, peer, network, postureChecks, enableSSH, settings)
if err != nil {
log.WithContext(ctx).Warnf("failed preparing login response for peer %s: %s", peerKey, err)
return nil, status.Errorf(codes.Internal, "failed logging in peer")
@@ -837,7 +929,7 @@ func (s *Server) ExtendAuthSession(ctx context.Context, req *proto.EncryptedMess
}, nil
}
-func (s *Server) prepareLoginResponse(ctx context.Context, peer *nbpeer.Peer, network *types.Network, postureChecks []*nmdata.PostureChecks, enableSSH bool) (*proto.LoginResponse, error) {
+func (s *Server) prepareLoginResponse(ctx context.Context, peer *nbpeer.Peer, network *types.Network, postureChecks []*nmdata.PostureChecks, enableSSH bool, settings *types.Settings) (*proto.LoginResponse, error) {
var relayToken *Token
var err error
if s.config.Relay != nil && len(s.config.Relay.Addresses) > 0 {
@@ -847,10 +939,12 @@ func (s *Server) prepareLoginResponse(ctx context.Context, peer *nbpeer.Peer, ne
}
}
- settings, err := s.settingsManager.GetSettings(ctx, peer.AccountID, activity.SystemInitiator)
- if err != nil {
- log.WithContext(ctx).Warnf("failed getting settings for peer %s: %s", peer.Key, err)
- return nil, status.Errorf(codes.Internal, "failed getting settings")
+ if settings == nil {
+ settings, err = s.settingsManager.GetSettings(ctx, peer.AccountID, activity.SystemInitiator)
+ if err != nil {
+ log.WithContext(ctx).Warnf("failed getting settings for peer %s: %s", peer.Key, err)
+ return nil, status.Errorf(codes.Internal, "failed getting settings")
+ }
}
// if peer has reached this point then it has logged in
@@ -925,7 +1019,7 @@ func (s *Server) IsHealthy(ctx context.Context, req *proto.Empty) (*proto.Empty,
}
// sendInitialSync sends initial proto.SyncResponse to the peer requesting synchronization
-func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer *nbpeer.Peer, networkMap *types.NetworkMap, postureChecks []*nmdata.PostureChecks, srv proto.ManagementService_SyncServer, dnsFwdPort int64, onChallengeStamped func()) error {
+func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer *nbpeer.Peer, networkMap *types.NetworkMap, postureChecks []*nmdata.PostureChecks, srv proto.ManagementService_SyncServer, dnsFwdPort int64, settings *types.Settings, onChallengeStamped func()) error {
var err error
var turnToken *Token
@@ -944,11 +1038,6 @@ func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer
}
}
- settings, err := s.settingsManager.GetSettings(ctx, peer.AccountID, activity.SystemInitiator)
- if err != nil {
- return status.Errorf(codes.Internal, "error handling request")
- }
-
peerGroups, err := s.accountManager.GetStore().GetPeerGroupIDs(ctx, store.LockingStrengthNone, peer.AccountID, peer.ID)
if err != nil {
return status.Errorf(codes.Internal, "failed to get peer groups %s", err)
diff --git a/management/internals/shared/grpc/server_disconnect_test.go b/management/internals/shared/grpc/server_disconnect_test.go
index 7d186b715..c664502e7 100644
--- a/management/internals/shared/grpc/server_disconnect_test.go
+++ b/management/internals/shared/grpc/server_disconnect_test.go
@@ -2,14 +2,18 @@ package grpc
import (
"context"
+ "errors"
+ "sync"
"testing"
"time"
+ "github.com/stretchr/testify/require"
"go.uber.org/mock/gomock"
"github.com/netbirdio/netbird/management/internals/controllers/network_map"
"github.com/netbirdio/netbird/management/server/account"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
+ "github.com/netbirdio/netbird/management/server/types"
)
func TestCancelPeerRoutines_SessionOwnership(t *testing.T) {
@@ -49,3 +53,144 @@ func TestCancelPeerRoutines_SessionOwnership(t *testing.T) {
})
}
}
+
+func TestCancelPeerRoutines_DrainingQueuesDisconnect(t *testing.T) {
+ streamStart := time.Unix(1700000000, 0)
+
+ tests := []struct {
+ name string
+ embedded bool
+ draining bool
+ flushBefore bool
+ wantQueued bool
+ wantOwnWrite bool
+ }{
+ {name: "regular peer while draining is queued", draining: true, wantQueued: true},
+ {name: "embedded proxy peer while draining writes its own disconnect", embedded: true, draining: true, wantOwnWrite: true},
+ {name: "regular peer without draining writes its own disconnect", wantOwnWrite: true},
+ {name: "regular peer after the flush writes its own disconnect", draining: true, flushBefore: true, wantOwnWrite: true},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ ctrl := gomock.NewController(t)
+ accountManager := account.NewMockManager(ctrl)
+ controller := network_map.NewMockController(ctrl)
+ secretsManager := NewMockSecretsManager(ctrl)
+ s := &Server{accountManager: accountManager, networkMapController: controller, secretsManager: secretsManager}
+
+ peer := &nbpeer.Peer{ID: "peer-1", Key: "peer-key"}
+ peer.ProxyMeta.Embedded = tt.embedded
+
+ if tt.draining {
+ s.StartDraining()
+ }
+ if tt.flushBefore {
+ require.NoError(t, s.FlushDisconnects(context.Background()))
+ }
+
+ if tt.wantOwnWrite {
+ accountManager.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.Key, streamStart).Return(nil)
+ }
+ controller.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.ID, gomock.Any()).Return(true)
+ secretsManager.EXPECT().CancelRefresh(peer.ID)
+ accountManager.EXPECT().UntrackCertificateChallenges("account-1", peer.ID, streamStart)
+
+ s.cancelPeerRoutines(context.Background(), "account-1", peer, streamStart, nil)
+
+ if tt.wantQueued {
+ accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), []types.PeerSession{
+ {AccountID: "account-1", PeerID: peer.ID, SessionStartedAt: streamStart.UnixNano()},
+ }).Return(nil)
+ }
+ require.NoError(t, s.FlushDisconnects(context.Background()))
+ })
+ }
+}
+
+func TestFlushDisconnects_FailedBatchStopsDraining(t *testing.T) {
+ ctrl := gomock.NewController(t)
+ accountManager := account.NewMockManager(ctrl)
+ controller := network_map.NewMockController(ctrl)
+ secretsManager := NewMockSecretsManager(ctrl)
+ s := &Server{accountManager: accountManager, networkMapController: controller, secretsManager: secretsManager}
+
+ peer := &nbpeer.Peer{ID: "peer-1", Key: "peer-key"}
+ first := time.Unix(1700000000, 0)
+ second := first.Add(time.Second)
+ controller.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.ID, gomock.Any()).Return(true).Times(2)
+ secretsManager.EXPECT().CancelRefresh(peer.ID).Times(2)
+ accountManager.EXPECT().UntrackCertificateChallenges("account-1", peer.ID, gomock.Any()).Times(2)
+
+ s.StartDraining()
+ s.cancelPeerRoutines(context.Background(), "account-1", peer, first, nil)
+
+ accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), gomock.Any()).Return(errors.New("db down")).Times(2)
+ require.Error(t, s.FlushDisconnects(context.Background()))
+
+ accountManager.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.Key, second).Return(nil)
+ s.cancelPeerRoutines(context.Background(), "account-1", peer, second, nil)
+}
+
+func TestFlushDisconnects_RetriesFailedBatch(t *testing.T) {
+ ctrl := gomock.NewController(t)
+ accountManager := account.NewMockManager(ctrl)
+ controller := network_map.NewMockController(ctrl)
+ secretsManager := NewMockSecretsManager(ctrl)
+ s := &Server{accountManager: accountManager, networkMapController: controller, secretsManager: secretsManager}
+
+ peer := &nbpeer.Peer{ID: "peer-1", Key: "peer-key"}
+ streamStart := time.Unix(1700000000, 0)
+ controller.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.ID, gomock.Any()).Return(true)
+ secretsManager.EXPECT().CancelRefresh(peer.ID)
+ accountManager.EXPECT().UntrackCertificateChallenges("account-1", peer.ID, streamStart)
+
+ s.StartDraining()
+ s.cancelPeerRoutines(context.Background(), "account-1", peer, streamStart, nil)
+
+ want := []types.PeerSession{{AccountID: "account-1", PeerID: peer.ID, SessionStartedAt: streamStart.UnixNano()}}
+ gomock.InOrder(
+ accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), want).Return(errors.New("deadlock detected")),
+ accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), want).Return(nil),
+ )
+ require.NoError(t, s.FlushDisconnects(context.Background()))
+}
+
+func TestCancelPeerRoutines_QueuesDisconnectBeforeWaitingForPeerLock(t *testing.T) {
+ ctrl := gomock.NewController(t)
+ accountManager := account.NewMockManager(ctrl)
+ controller := network_map.NewMockController(ctrl)
+ secretsManager := NewMockSecretsManager(ctrl)
+ s := &Server{accountManager: accountManager, networkMapController: controller, secretsManager: secretsManager}
+
+ peer := &nbpeer.Peer{ID: "peer-1", Key: "peer-key"}
+ streamStart := time.Unix(1700000000, 0)
+ controller.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.ID, gomock.Any()).Return(true)
+ secretsManager.EXPECT().CancelRefresh(peer.ID)
+ accountManager.EXPECT().UntrackCertificateChallenges("account-1", peer.ID, streamStart)
+
+ unlock := s.acquirePeerLockByUID(context.Background(), peer.Key)
+ var unlockOnce sync.Once
+ t.Cleanup(func() { unlockOnce.Do(unlock) })
+ s.StartDraining()
+
+ done := make(chan struct{})
+ go func() {
+ defer close(done)
+ s.cancelPeerRoutines(context.Background(), "account-1", peer, streamStart, nil)
+ }()
+
+ require.Eventually(t, func() bool {
+ s.drainMu.Lock()
+ defer s.drainMu.Unlock()
+ return len(s.pendingDisconnects) == 1
+ }, time.Second, 5*time.Millisecond, "the disconnect must be queued while the teardown waits for the peer lock")
+
+ accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), []types.PeerSession{
+ {AccountID: "account-1", PeerID: peer.ID, SessionStartedAt: streamStart.UnixNano()},
+ }).Return(nil)
+ require.NoError(t, s.FlushDisconnects(context.Background()))
+
+ unlockOnce.Do(unlock)
+ <-done
+}
diff --git a/management/server/account.go b/management/server/account.go
index 21b9fa431..f62a83d40 100644
--- a/management/server/account.go
+++ b/management/server/account.go
@@ -2033,13 +2033,13 @@ func domainIsUpToDate(domain string, domainCategory string, userAuth auth.UserAu
// derived from syncTime (the moment the gRPC stream opened). Any
// concurrent stream that started earlier loses the optimistic-lock race
// in MarkPeerConnected and bails without writing.
-func (am *DefaultAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
- peer, netMap, postureChecks, dnsfwdPort, err := am.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peerPubKey, Meta: meta, RealIP: realIP}, accountID)
+func (am *DefaultAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
+ peer, netMap, postureChecks, dnsfwdPort, err := am.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peerPubKey, Meta: meta, RealIP: realIP, Settings: settings}, accountID)
if err != nil {
return nil, nil, nil, 0, fmt.Errorf("error syncing peer: %w", err)
}
- if err := am.MarkPeerConnected(ctx, peerPubKey, accountID, syncTime.UnixNano(), netMap); err != nil {
+ if err := am.markPeerConnected(ctx, peerPubKey, accountID, syncTime.UnixNano(), netMap, settings); err != nil {
log.WithContext(ctx).Warnf("failed marking peer as connected %s %v", peerPubKey, err)
}
@@ -2059,6 +2059,17 @@ func (am *DefaultAccountManager) OnPeerDisconnected(ctx context.Context, account
return nil
}
+// MarkPeersDisconnected marks the peers of the given sessions disconnected in batched writes,
+// skipping every peer whose stored session differs from the given one.
+func (am *DefaultAccountManager) MarkPeersDisconnected(ctx context.Context, sessions []types.PeerSession) error {
+ marked, err := am.Store.MarkPeersDisconnectedIfSameSession(ctx, sessions)
+ if err != nil {
+ return fmt.Errorf("mark peers disconnected: %w", err)
+ }
+ log.WithContext(ctx).Infof("marked %d of %d peers as disconnected", marked, len(sessions))
+ return nil
+}
+
func (am *DefaultAccountManager) SyncPeerMeta(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) error {
accountID, err := am.Store.GetAccountIDByPeerPubKey(ctx, peerPubKey)
if err != nil {
diff --git a/management/server/account/manager.go b/management/server/account/manager.go
index 15fe11241..2f4df8a48 100644
--- a/management/server/account/manager.go
+++ b/management/server/account/manager.go
@@ -124,8 +124,10 @@ type Manager interface {
UpdateIntegratedValidator(ctx context.Context, accountID, userID, validator string, groups []string) error
GroupValidation(ctx context.Context, accountId string, groups []string) (bool, error)
GetValidatedPeers(ctx context.Context, accountID string) (map[string]struct{}, map[string]string, error)
- SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error)
+ SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error)
OnPeerDisconnected(ctx context.Context, accountID string, peerPubKey string, streamStartTime time.Time) error
+ CheckPeerMaySync(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error
+ MarkPeersDisconnected(ctx context.Context, sessions []types.PeerSession) error
SyncPeerMeta(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) error
FindExistingPostureCheck(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error)
GetAccountIDForPeerKey(ctx context.Context, peerKey string) (string, error)
diff --git a/management/server/account/manager_mock.go b/management/server/account/manager_mock.go
index b3868f2e2..e7b044bee 100644
--- a/management/server/account/manager_mock.go
+++ b/management/server/account/manager_mock.go
@@ -146,6 +146,20 @@ func (mr *MockManagerMockRecorder) BuildUserInfosForAccount(ctx, accountID, init
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BuildUserInfosForAccount", reflect.TypeOf((*MockManager)(nil).BuildUserInfosForAccount), ctx, accountID, initiatorUserID, accountUsers)
}
+// CheckPeerMaySync mocks base method.
+func (m *MockManager) CheckPeerMaySync(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error {
+ m.ctrl.T.Helper()
+ ret := m.ctrl.Call(m, "CheckPeerMaySync", ctx, accountID, peerPubKey, settings)
+ ret0, _ := ret[0].(error)
+ return ret0
+}
+
+// CheckPeerMaySync indicates an expected call of CheckPeerMaySync.
+func (mr *MockManagerMockRecorder) CheckPeerMaySync(ctx, accountID, peerPubKey, settings any) *gomock.Call {
+ mr.mock.ctrl.T.Helper()
+ return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckPeerMaySync", reflect.TypeOf((*MockManager)(nil).CheckPeerMaySync), ctx, accountID, peerPubKey, settings)
+}
+
// CreateGroup mocks base method.
func (m *MockManager) CreateGroup(ctx context.Context, accountID, userID string, group *types.Group) error {
m.ctrl.T.Helper()
@@ -1369,6 +1383,20 @@ func (mr *MockManagerMockRecorder) MarkPeerDisconnected(ctx, peerKey, accountID,
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeerDisconnected", reflect.TypeOf((*MockManager)(nil).MarkPeerDisconnected), ctx, peerKey, accountID, sessionStartedAt)
}
+// MarkPeersDisconnected mocks base method.
+func (m *MockManager) MarkPeersDisconnected(ctx context.Context, sessions []types.PeerSession) error {
+ m.ctrl.T.Helper()
+ ret := m.ctrl.Call(m, "MarkPeersDisconnected", ctx, sessions)
+ ret0, _ := ret[0].(error)
+ return ret0
+}
+
+// MarkPeersDisconnected indicates an expected call of MarkPeersDisconnected.
+func (mr *MockManagerMockRecorder) MarkPeersDisconnected(ctx, sessions any) *gomock.Call {
+ mr.mock.ctrl.T.Helper()
+ return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeersDisconnected", reflect.TypeOf((*MockManager)(nil).MarkPeersDisconnected), ctx, sessions)
+}
+
// OnPeerDisconnected mocks base method.
func (m *MockManager) OnPeerDisconnected(ctx context.Context, accountID, peerPubKey string, streamStartTime time.Time) error {
m.ctrl.T.Helper()
@@ -1569,9 +1597,9 @@ func (mr *MockManagerMockRecorder) StoreEvent(ctx, initiatorID, targetID, accoun
}
// SyncAndMarkPeer mocks base method.
-func (m *MockManager) SyncAndMarkPeer(ctx context.Context, accountID, peerPubKey string, meta peer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*peer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
+func (m *MockManager) SyncAndMarkPeer(ctx context.Context, accountID, peerPubKey string, meta peer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*peer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
m.ctrl.T.Helper()
- ret := m.ctrl.Call(m, "SyncAndMarkPeer", ctx, accountID, peerPubKey, meta, realIP, syncTime)
+ ret := m.ctrl.Call(m, "SyncAndMarkPeer", ctx, accountID, peerPubKey, meta, realIP, syncTime, settings)
ret0, _ := ret[0].(*peer.Peer)
ret1, _ := ret[1].(*types.NetworkMap)
ret2, _ := ret[2].([]*nmdata.PostureChecks)
@@ -1581,9 +1609,9 @@ func (m *MockManager) SyncAndMarkPeer(ctx context.Context, accountID, peerPubKey
}
// SyncAndMarkPeer indicates an expected call of SyncAndMarkPeer.
-func (mr *MockManagerMockRecorder) SyncAndMarkPeer(ctx, accountID, peerPubKey, meta, realIP, syncTime any) *gomock.Call {
+func (mr *MockManagerMockRecorder) SyncAndMarkPeer(ctx, accountID, peerPubKey, meta, realIP, syncTime, settings any) *gomock.Call {
mr.mock.ctrl.T.Helper()
- return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SyncAndMarkPeer", reflect.TypeOf((*MockManager)(nil).SyncAndMarkPeer), ctx, accountID, peerPubKey, meta, realIP, syncTime)
+ return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SyncAndMarkPeer", reflect.TypeOf((*MockManager)(nil).SyncAndMarkPeer), ctx, accountID, peerPubKey, meta, realIP, syncTime, settings)
}
// SyncPeer mocks base method.
diff --git a/management/server/account_test.go b/management/server/account_test.go
index c343575b8..d6971be85 100644
--- a/management/server/account_test.go
+++ b/management/server/account_test.go
@@ -3977,7 +3977,7 @@ func BenchmarkSyncAndMarkPeer(b *testing.B) {
b.ResetTimer()
start := time.Now()
for i := 0; i < b.N; i++ {
- _, _, _, _, err := manager.SyncAndMarkPeer(context.Background(), account.Id, account.Peers["peer-1"].Key, nbpeer.PeerSystemMeta{Hostname: strconv.Itoa(i)}, net.IP{1, 1, 1, 1}, time.Now().UTC())
+ _, _, _, _, err := manager.SyncAndMarkPeer(context.Background(), account.Id, account.Peers["peer-1"].Key, nbpeer.PeerSystemMeta{Hostname: strconv.Itoa(i)}, net.IP{1, 1, 1, 1}, time.Now().UTC(), nil)
assert.NoError(b, err)
}
diff --git a/management/server/management_proto_test.go b/management/server/management_proto_test.go
index 9ca94d2e6..8e495fa7e 100644
--- a/management/server/management_proto_test.go
+++ b/management/server/management_proto_test.go
@@ -17,8 +17,10 @@ import (
"go.uber.org/mock/gomock"
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
"google.golang.org/grpc"
+ "google.golang.org/grpc/codes"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/keepalive"
+ "google.golang.org/grpc/status"
"github.com/netbirdio/netbird/encryption"
"github.com/netbirdio/netbird/formatter/hook"
@@ -824,3 +826,130 @@ func testLoginPerformance(t *testing.T, loginCalls []func() error) {
tAvg := tSum / time.Duration(len(durations))
t.Logf("Min: %v, Max: %v, Avg: %v", tMin, tMax, tAvg)
}
+
+type expireOnConnectStore struct {
+ store.Store
+ once sync.Once
+ err error
+}
+
+func (s *expireOnConnectStore) MarkPeerConnectedIfNewerSession(ctx context.Context, accountID, peerID string, newSessionStartedAt int64) (bool, error) {
+ updated, err := s.Store.MarkPeerConnectedIfNewerSession(ctx, accountID, peerID, newSessionStartedAt)
+ if err != nil {
+ return updated, err
+ }
+
+ s.once.Do(func() {
+ peer, getErr := s.Store.GetPeerByID(ctx, store.LockingStrengthNone, accountID, peerID)
+ if getErr != nil {
+ s.err = getErr
+ return
+ }
+ peer.Status.LoginExpired = true
+ s.err = s.Store.SavePeerStatus(ctx, accountID, peerID, *peer.Status)
+ })
+ return updated, err
+}
+
+func Test_SyncRejectsPeerExpiredWhileConnecting(t *testing.T) {
+ mgmtServer, am, mgmtAddr, cleanup, err := startManagementForTest(t, "testdata/store_with_expired_peers.sql", &config.Config{
+ Datadir: t.TempDir(),
+ })
+ defer cleanup()
+ require.NoError(t, err)
+ defer mgmtServer.GracefulStop()
+
+ client, clientConn, err := createRawClient(mgmtAddr)
+ require.NoError(t, err)
+ defer clientConn.Close()
+
+ key, err := wgtypes.GeneratePrivateKey()
+ require.NoError(t, err)
+ _, err = loginPeerWithValidSetupKey(key, client)
+ require.NoError(t, err)
+
+ serverKey, err := getServerKey(client)
+ require.NoError(t, err)
+
+ decorated := &expireOnConnectStore{Store: am.Store}
+ am.Store = decorated
+
+ message, err := encryption.EncryptMessage(*serverKey, key, &mgmtProto.SyncRequest{Meta: &mgmtProto.PeerSystemMeta{}})
+ require.NoError(t, err)
+
+ ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+ defer cancel()
+
+ sync, err := client.Sync(ctx, &mgmtProto.EncryptedMessage{
+ WgPubKey: key.PublicKey().String(),
+ Body: message,
+ })
+ require.NoError(t, err)
+
+ resp := &mgmtProto.EncryptedMessage{}
+ for err == nil {
+ err = sync.RecvMsg(resp)
+ }
+ require.Equal(t, codes.PermissionDenied, status.Code(err), "sync stream ended with: %v", err)
+ require.NoError(t, decorated.err)
+
+ peer, err := am.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, key.PublicKey().String())
+ require.NoError(t, err)
+ require.True(t, peer.Status.LoginExpired)
+ require.Eventually(t, func() bool {
+ peer, err = am.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, key.PublicKey().String())
+ return err == nil && !peer.Status.Connected
+ }, 5*time.Second, 50*time.Millisecond)
+}
+
+type settingsCountingStore struct {
+ store.Store
+ reads atomic.Int32
+}
+
+func (s *settingsCountingStore) GetAccountSettings(ctx context.Context, lockStrength store.LockingStrength, accountID string) (*types.Settings, error) {
+ s.reads.Add(1)
+ return s.Store.GetAccountSettings(ctx, lockStrength, accountID)
+}
+
+func Test_LoginAndSyncReuseRequestSettings(t *testing.T) {
+ mgmtServer, am, mgmtAddr, cleanup, err := startManagementForTest(t, "testdata/store_with_expired_peers.sql", &config.Config{
+ Datadir: t.TempDir(),
+ })
+ defer cleanup()
+ require.NoError(t, err)
+ defer mgmtServer.GracefulStop()
+
+ client, clientConn, err := createRawClient(mgmtAddr)
+ require.NoError(t, err)
+ defer clientConn.Close()
+
+ key, err := wgtypes.GeneratePrivateKey()
+ require.NoError(t, err)
+ _, err = loginPeerWithValidSetupKey(key, client)
+ require.NoError(t, err)
+
+ counting := &settingsCountingStore{Store: am.Store}
+ am.Store = counting
+
+ _, err = loginPeerWithValidSetupKey(key, client)
+ require.NoError(t, err)
+ require.Zero(t, counting.reads.Load(), "login must not read settings through the account manager")
+
+ serverKey, err := getServerKey(client)
+ require.NoError(t, err)
+ message, err := encryption.EncryptMessage(*serverKey, key, &mgmtProto.SyncRequest{Meta: &mgmtProto.PeerSystemMeta{}})
+ require.NoError(t, err)
+
+ ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+ defer cancel()
+ sync, err := client.Sync(ctx, &mgmtProto.EncryptedMessage{WgPubKey: key.PublicKey().String(), Body: message})
+ require.NoError(t, err)
+ require.NoError(t, sync.RecvMsg(&mgmtProto.EncryptedMessage{}))
+
+ require.Eventually(t, func() bool {
+ peer, err := counting.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, key.PublicKey().String())
+ return err == nil && peer.Status.Connected
+ }, 5*time.Second, 50*time.Millisecond)
+ require.Zero(t, counting.reads.Load(), "sync must not read settings through the account manager")
+}
diff --git a/management/server/mock_server/account_mock.go b/management/server/mock_server/account_mock.go
index 50d18ca9b..69a0de251 100644
--- a/management/server/mock_server/account_mock.go
+++ b/management/server/mock_server/account_mock.go
@@ -42,7 +42,9 @@ type MockAccountManager struct {
GetPeersFunc func(ctx context.Context, accountID, userID, nameFilter, ipFilter, macFilter string) ([]*nbpeer.Peer, int, error)
MarkPeerConnectedFunc func(ctx context.Context, peerKey string, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error
MarkPeerDisconnectedFunc func(ctx context.Context, peerKey string, accountID string, sessionStartedAt int64) error
- SyncAndMarkPeerFunc func(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error)
+ MarkPeersDisconnectedFunc func(ctx context.Context, sessions []types.PeerSession) error
+ CheckPeerMaySyncFunc func(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error
+ SyncAndMarkPeerFunc func(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error)
DeletePeerFunc func(ctx context.Context, accountID, peerKey, userID string) error
GetNetworkMapFunc func(ctx context.Context, peerKey string) (*types.NetworkMap, error)
GetPeerNetworkFunc func(ctx context.Context, peerKey string) (*types.Network, error)
@@ -242,13 +244,27 @@ func (am *MockAccountManager) DeleteSetupKey(ctx context.Context, accountID, use
return status.Errorf(codes.Unimplemented, "method DeleteSetupKey is not implemented")
}
-func (am *MockAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
+func (am *MockAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
if am.SyncAndMarkPeerFunc != nil {
- return am.SyncAndMarkPeerFunc(ctx, accountID, peerPubKey, meta, realIP, syncTime)
+ return am.SyncAndMarkPeerFunc(ctx, accountID, peerPubKey, meta, realIP, syncTime, settings)
}
return nil, nil, nil, 0, status.Errorf(codes.Unimplemented, "method MarkPeerConnected is not implemented")
}
+func (am *MockAccountManager) MarkPeersDisconnected(ctx context.Context, sessions []types.PeerSession) error {
+ if am.MarkPeersDisconnectedFunc != nil {
+ return am.MarkPeersDisconnectedFunc(ctx, sessions)
+ }
+ return nil
+}
+
+func (am *MockAccountManager) CheckPeerMaySync(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error {
+ if am.CheckPeerMaySyncFunc != nil {
+ return am.CheckPeerMaySyncFunc(ctx, accountID, peerPubKey, settings)
+ }
+ return status.Errorf(codes.Unimplemented, "method CheckPeerMaySync is not implemented")
+}
+
func (am *MockAccountManager) OnPeerDisconnected(ctx context.Context, accountID string, peerPubKey string, streamStartTime time.Time) error {
// Mirror DefaultAccountManager.OnPeerDisconnected: drive the fencing
// hook so tests that inject MarkPeerDisconnectedFunc actually observe
diff --git a/management/server/peer.go b/management/server/peer.go
index c3454179e..3a48c37fb 100644
--- a/management/server/peer.go
+++ b/management/server/peer.go
@@ -88,6 +88,12 @@ func (am *DefaultAccountManager) GetPeers(ctx context.Context, accountID, userID
// Disconnects use MarkPeerDisconnected and require the session to match
// exactly; see PeerStatus.SessionStartedAt for the protocol.
func (am *DefaultAccountManager) MarkPeerConnected(ctx context.Context, peerPubKey string, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error {
+ return am.markPeerConnected(ctx, peerPubKey, accountID, sessionStartedAt, nmap, nil)
+}
+
+// markPeerConnected is MarkPeerConnected with the account settings the caller already
+// read; nil settings are loaded when needed.
+func (am *DefaultAccountManager) markPeerConnected(ctx context.Context, peerPubKey string, accountID string, sessionStartedAt int64, nmap *types.NetworkMap, settings *types.Settings) error {
start := time.Now()
defer func() {
am.metrics.AccountManagerMetrics().RecordPeerStatusUpdateDuration(telemetry.PeerStatusConnect, time.Since(start))
@@ -119,19 +125,13 @@ func (am *DefaultAccountManager) MarkPeerConnected(ctx context.Context, peerPubK
log.WithContext(ctx).Debugf("mark peer %s connected", peer.ID)
- if err = am.schedulePeerExpirations(ctx, accountID, peer); err != nil {
+ if err = am.schedulePeerExpirations(ctx, accountID, peer, settings); err != nil {
return err
}
- // A login-expired peer reconnecting, or an embedded proxy peer flipping to
- // connected (which triggers SynthesizePrivateServiceZones), must refresh the
- // peers reachable from it. The embedded-proxy fan-out tolerates a dispatch error.
- if peer.Status != nil && peer.Status.LoginExpired {
- affectedPeerIDs := am.markConnectedAffectedPeers(ctx, accountID, peer.ID, nmap)
- if err = am.networkMapController.OnPeersUpdated(ctx, accountID, []string{peer.ID}, affectedPeerIDs); err != nil {
- return fmt.Errorf("notify network map controller of peer update: %w", err)
- }
- }
+ // An embedded proxy peer flipping to connected (which triggers
+ // SynthesizePrivateServiceZones) must refresh the peers reachable from it.
+ // The fan-out tolerates a dispatch error.
if peer.ProxyMeta.Embedded {
affectedPeerIDs := am.markConnectedAffectedPeers(ctx, accountID, peer.ID, nmap)
if err := am.networkMapController.OnPeersUpdated(ctx, accountID, []string{peer.ID}, affectedPeerIDs); err != nil {
@@ -166,11 +166,11 @@ func (am *DefaultAccountManager) writePeerStatus(ctx context.Context, accountID
// schedulePeerExpirations arms the account's login/inactivity expiration
// timers for an SSO peer that just connected.
-func (am *DefaultAccountManager) schedulePeerExpirations(ctx context.Context, accountID string, peer *nbpeer.Peer) error {
+func (am *DefaultAccountManager) schedulePeerExpirations(ctx context.Context, accountID string, peer *nbpeer.Peer, settings *types.Settings) error {
if !peer.AddedWithSSOLogin() {
return nil
}
- settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
+ settings, err := am.accountSettings(ctx, accountID, settings)
if err != nil {
return err
}
@@ -1087,37 +1087,22 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
var metaDiff nbpeer.MetaDiff
var err error
- settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
+ settings, err := am.accountSettings(ctx, accountID, sync.Settings)
if err != nil {
return nil, nil, nil, 0, err
}
- err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
- peer, err = transaction.GetPeerByPeerPubKey(ctx, store.LockingStrengthUpdate, sync.WireGuardPubKey)
- if err != nil {
- return status.NewPeerNotRegisteredError()
- }
+ peer, metaDiff, ipv6CapabilityChanged, err = am.prepareSyncPeer(ctx, am.Store, store.LockingStrengthNone, sync, settings)
+ if err != nil {
+ return nil, nil, nil, 0, err
+ }
- if peer.UserID != "" {
- user, err := transaction.GetUserByUserID(ctx, store.LockingStrengthNone, peer.UserID)
- if err != nil {
+ if metaDiff.Updated() {
+ err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
+ peer, metaDiff, ipv6CapabilityChanged, err = am.prepareSyncPeer(ctx, transaction, store.LockingStrengthUpdate, sync, settings)
+ if err != nil || !metaDiff.Updated() {
return err
}
-
- if err = checkIfPeerOwnerIsBlocked(peer, user); err != nil {
- return err
- }
- }
-
- if peerLoginExpired(ctx, peer, settings) {
- return status.NewPeerLoginExpiredError()
- }
-
- oldHasIPv6Cap := peer.HasCapability(nbpeer.PeerCapabilityIPv6Overlay)
- newLocation := am.resolvePeerLocation(ctx, peer, sync.RealIP)
- metaDiff = peer.UpdateMetaIfNew(ctx, sync.Meta, newLocation)
- ipv6CapabilityChanged = oldHasIPv6Cap != peer.HasCapability(nbpeer.PeerCapabilityIPv6Overlay)
- if metaDiff.Updated() {
am.metrics.AccountManagerMetrics().CountPeerMetUpdate()
log.WithContext(ctx).Tracef("peer %s metadata updated", peer.ID)
if err = transaction.SavePeer(ctx, accountID, peer); err != nil {
@@ -1126,11 +1111,11 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
if err = transaction.IncrementNetworkSerial(ctx, accountID); err != nil {
return fmt.Errorf("increment network serial: %w", err)
}
+ return nil
+ })
+ if err != nil {
+ return nil, nil, nil, 0, err
}
- return nil
- })
- if err != nil {
- return nil, nil, nil, 0, err
}
peerGroupIDs, err := getPeerGroupIDs(ctx, am.Store, accountID, peer.ID)
@@ -1160,6 +1145,82 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
return peer, nmap, resPostureChecks, dnsFwdPort, nil
}
+// prepareSyncPeer loads the syncing peer with lockStrength, checks that it may sync and
+// applies the reported meta to it in memory. The returned diff tells whether the peer
+// has to be saved.
+func (am *DefaultAccountManager) prepareSyncPeer(ctx context.Context, st store.Store, lockStrength store.LockingStrength, sync types.PeerSync, settings *types.Settings) (*nbpeer.Peer, nbpeer.MetaDiff, bool, error) {
+ peer, err := st.GetPeerByPeerPubKey(ctx, lockStrength, sync.WireGuardPubKey)
+ if err != nil {
+ return nil, nbpeer.MetaDiff{}, false, peerLookupError(err)
+ }
+
+ if err = checkPeerMaySync(ctx, st, peer, settings); err != nil {
+ return nil, nbpeer.MetaDiff{}, false, err
+ }
+
+ oldHasIPv6Cap := peer.HasCapability(nbpeer.PeerCapabilityIPv6Overlay)
+ newLocation := am.resolvePeerLocation(ctx, peer, sync.RealIP)
+ metaDiff := peer.UpdateMetaIfNew(ctx, sync.Meta, newLocation)
+ ipv6CapabilityChanged := oldHasIPv6Cap != peer.HasCapability(nbpeer.PeerCapabilityIPv6Overlay)
+
+ return peer, metaDiff, ipv6CapabilityChanged, nil
+}
+
+// CheckPeerMaySync reads the peer again and rejects it like SyncPeer does when its owner
+// is blocked or its login expired. Nil settings are loaded from the store. The gRPC
+// server calls it once the sync stream can receive updates, so an expiry committed
+// while the stream was being set up is not missed.
+func (am *DefaultAccountManager) CheckPeerMaySync(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error {
+ settings, err := am.accountSettings(ctx, accountID, settings)
+ if err != nil {
+ return err
+ }
+
+ peer, err := am.Store.GetPeerByPeerPubKey(ctx, store.LockingStrengthNone, peerPubKey)
+ if err != nil {
+ return peerLookupError(err)
+ }
+
+ return checkPeerMaySync(ctx, am.Store, peer, settings)
+}
+
+// peerLookupError maps a missing peer to the not-registered error that makes the client
+// register again, and keeps any other store error so a transient failure is retried.
+func peerLookupError(err error) error {
+ if s, ok := status.FromError(err); ok && s.Type() == status.NotFound {
+ return status.NewPeerNotRegisteredError()
+ }
+ return err
+}
+
+// accountSettings returns settings when the caller already read them for this request
+// and loads them from the store otherwise.
+func (am *DefaultAccountManager) accountSettings(ctx context.Context, accountID string, settings *types.Settings) (*types.Settings, error) {
+ if settings != nil {
+ return settings, nil
+ }
+ return am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
+}
+
+// checkPeerMaySync rejects a peer whose owner is blocked or whose login expired.
+func checkPeerMaySync(ctx context.Context, st store.Store, peer *nbpeer.Peer, settings *types.Settings) error {
+ if peer.UserID != "" {
+ user, err := st.GetUserByUserID(ctx, store.LockingStrengthNone, peer.UserID)
+ if err != nil {
+ return err
+ }
+
+ if err = checkIfPeerOwnerIsBlocked(peer, user); err != nil {
+ return err
+ }
+ }
+
+ if peerLoginExpired(ctx, peer, settings) {
+ return status.NewPeerLoginExpiredError()
+ }
+ return nil
+}
+
// metaDiffAffectsPosture reports whether the meta change flips the verdict of any of
// the peer's posture checks, replaying them against the old and new state.
func metaDiffAffectsPosture(diff *nbpeer.MetaDiff, checks []*nmdata.PostureChecks) bool {
@@ -1206,8 +1267,8 @@ func (am *DefaultAccountManager) syncPeerAffectedPeers(ctx context.Context, acco
return affectedPeerIDsFromNetworkMap(nmap, peerID)
}
-// markConnectedAffectedPeers resolves the peers affected when a peer connects
-// (login-expiry reconnect or embedded-proxy connect). The connecting peer's
+// markConnectedAffectedPeers resolves the peers affected when an embedded proxy
+// peer connects. The connecting peer's
// network map already lists them bidirectionally — the synthesized
// private-service policy puts proxy access-group members in the proxy peer's own
// map, and these edges carry no source-posture gate. An invalid peer has an
@@ -1249,8 +1310,13 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login types.Peer
// when the client sends a login request with a JWT which is used to get the user ID,
// it means that the client has already checked if it needs login and had been through the SSO flow
// so, we can skip this check and directly proceed with the login
+ settings, err := am.accountSettings(ctx, accountID, login.Settings)
+ if err != nil {
+ return nil, nil, nil, false, err
+ }
+
if login.UserID == "" {
- err = am.checkIFPeerNeedsLoginWithoutLock(ctx, accountID, login)
+ err = am.checkIFPeerNeedsLoginWithoutLock(ctx, login, settings)
if err != nil {
return nil, nil, nil, false, err
}
@@ -1260,56 +1326,55 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login types.Peer
var shouldStorePeer, shouldUpdatePeers bool
var peerGroupIDs []string
- settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
+ peer, err = am.Store.GetPeerByPeerPubKey(ctx, store.LockingStrengthNone, login.WireGuardPubKey)
if err != nil {
return nil, nil, nil, false, err
}
+ if err = am.checkPeerLogin(ctx, am.Store, peer, login); err != nil {
+ return nil, nil, nil, false, err
+ }
- err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
- peer, err = transaction.GetPeerByPeerPubKey(ctx, store.LockingStrengthUpdate, login.WireGuardPubKey)
- if err != nil {
- return err
- }
-
- if login.UserID != "" {
- if peer.UserID != login.UserID {
- log.Warnf("user mismatch when logging in peer %s: peer user %s, login user %s ", peer.ID, peer.UserID, login.UserID)
- return status.NewPeerLoginMismatchError()
- }
-
- changed, err := am.handleUserPeer(ctx, transaction, peer, settings)
+ if peerLoginChangesPeer(ctx, peer, login, settings) {
+ err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
+ peer, err = transaction.GetPeerByPeerPubKey(ctx, store.LockingStrengthUpdate, login.WireGuardPubKey)
if err != nil {
return err
}
-
- if changed {
- shouldStorePeer = true
- shouldUpdatePeers = true
- }
- }
-
- if peer.SSHKey != login.SSHKey {
- peer.SSHKey = login.SSHKey
- shouldStorePeer = true
- }
-
- if !peer.AllowExtraDNSLabels && len(login.ExtraDNSLabels) > 0 {
- return status.Errorf(status.PreconditionFailed, "couldn't login peer: setup key doesn't allow extra DNS labels")
- }
-
- if shouldStorePeer {
- if err = transaction.SavePeer(ctx, accountID, peer); err != nil {
+ if err = am.checkPeerLogin(ctx, transaction, peer, login); err != nil {
return err
}
- if err = transaction.IncrementNetworkSerial(ctx, accountID); err != nil {
- return fmt.Errorf("increment network serial: %w", err)
- }
- }
- return nil
- })
- if err != nil {
- return nil, nil, nil, false, err
+ if login.UserID != "" {
+ changed, err := am.handleUserPeer(ctx, transaction, peer, settings)
+ if err != nil {
+ return err
+ }
+
+ if changed {
+ shouldStorePeer = true
+ shouldUpdatePeers = true
+ }
+ }
+
+ if peer.SSHKey != login.SSHKey {
+ peer.SSHKey = login.SSHKey
+ shouldStorePeer = true
+ }
+
+ if shouldStorePeer {
+ if err = transaction.SavePeer(ctx, accountID, peer); err != nil {
+ return err
+ }
+ if err = transaction.IncrementNetworkSerial(ctx, accountID); err != nil {
+ return fmt.Errorf("increment network serial: %w", err)
+ }
+ }
+
+ return nil
+ })
+ if err != nil {
+ return nil, nil, nil, false, err
+ }
}
// This is needed to keep in memory for the peer config. Otherwise browser client will end in a retry loop
@@ -1341,6 +1406,40 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login types.Peer
return peer, network, postureChecks, enableSSH, nil
}
+// checkPeerLogin rejects a login whose user does not own the peer, whose owner is
+// blocked, or that requests extra DNS labels the peer is not allowed to use.
+func (am *DefaultAccountManager) checkPeerLogin(ctx context.Context, st store.Store, peer *nbpeer.Peer, login types.PeerLogin) error {
+ if login.UserID != "" {
+ if peer.UserID != login.UserID {
+ log.Warnf("user mismatch when logging in peer %s: peer user %s, login user %s ", peer.ID, peer.UserID, login.UserID)
+ return status.NewPeerLoginMismatchError()
+ }
+
+ user, err := st.GetUserByUserID(ctx, store.LockingStrengthNone, peer.UserID)
+ if err != nil {
+ return err
+ }
+ if err = checkIfPeerOwnerIsBlocked(peer, user); err != nil {
+ return err
+ }
+ }
+
+ if !peer.AllowExtraDNSLabels && len(login.ExtraDNSLabels) > 0 {
+ return status.Errorf(status.PreconditionFailed, "couldn't login peer: setup key doesn't allow extra DNS labels")
+ }
+
+ return nil
+}
+
+// peerLoginChangesPeer reports whether logging in writes the peer: an SSO
+// re-login of an expired peer or a new SSH key.
+func peerLoginChangesPeer(ctx context.Context, peer *nbpeer.Peer, login types.PeerLogin, settings *types.Settings) bool {
+ if peer.SSHKey != login.SSHKey {
+ return true
+ }
+ return login.UserID != "" && peerLoginExpired(ctx, peer, settings)
+}
+
// ExtendPeerSession refreshes the peer's SSO session deadline by updating
// LastLogin after a successful JWT validation. The tunnel is untouched: no
// network map sync, no peer reconnect.
@@ -1507,7 +1606,7 @@ func processPeerPostureChecks(policy *types.Policy, peerID string, peerGroupIDs
// The NetBird client doesn't have a way to check if the peer needs login besides sending a login request
// with no JWT token and usually no setup-key. As the client can send up to two login request to check if it is expired
// and before starting the engine, we do the checks without an account lock to avoid piling up requests.
-func (am *DefaultAccountManager) checkIFPeerNeedsLoginWithoutLock(ctx context.Context, accountID string, login types.PeerLogin) error {
+func (am *DefaultAccountManager) checkIFPeerNeedsLoginWithoutLock(ctx context.Context, login types.PeerLogin, settings *types.Settings) error {
peer, err := am.Store.GetPeerByPeerPubKey(ctx, store.LockingStrengthNone, login.WireGuardPubKey)
if err != nil {
return err
@@ -1519,11 +1618,6 @@ func (am *DefaultAccountManager) checkIFPeerNeedsLoginWithoutLock(ctx context.Co
return nil
}
- settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
- if err != nil {
- return err
- }
-
if peerLoginExpired(ctx, peer, settings) {
return status.NewPeerLoginExpiredError()
}
diff --git a/management/server/peer_test.go b/management/server/peer_test.go
index bf8093c1d..48b75edbd 100644
--- a/management/server/peer_test.go
+++ b/management/server/peer_test.go
@@ -17,6 +17,7 @@ import (
"strconv"
"strings"
"sync"
+ "sync/atomic"
"testing"
"time"
@@ -3323,3 +3324,136 @@ func TestPeerWrites_BumpNetworkSerial(t *testing.T) {
assert.Greater(t, networkSerial(t, manager, account.Id), afterDisconnect, "embedded proxy connect must bump")
})
}
+
+type txCountingStore struct {
+ store.Store
+ txs atomic.Int32
+}
+
+func (s *txCountingStore) ExecuteInTransaction(ctx context.Context, f func(store.Store) error) error {
+ s.txs.Add(1)
+ return s.Store.ExecuteInTransaction(ctx, f)
+}
+
+func TestPeerSyncAndLogin_LockOnlyOnChange(t *testing.T) {
+ manager, _, account, peer1, _, _ := setupNetworkMapTest(t)
+ ctx := context.Background()
+ counting := &txCountingStore{Store: manager.Store}
+ manager.Store = counting
+
+ current, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+ require.NoError(t, err)
+
+ t.Run("SyncPeer without meta change opens no transaction", func(t *testing.T) {
+ counting.txs.Store(0)
+ _, _, _, _, err := manager.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peer1.Key, Meta: current.Meta}, account.Id)
+ require.NoError(t, err)
+ assert.Zero(t, counting.txs.Load())
+ })
+
+ t.Run("SyncPeer with meta change saves in one transaction", func(t *testing.T) {
+ counting.txs.Store(0)
+ meta := current.Meta
+ meta.Hostname = "lock-on-change-host"
+ _, _, _, _, err := manager.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peer1.Key, Meta: meta}, account.Id)
+ require.NoError(t, err)
+ assert.Equal(t, int32(1), counting.txs.Load())
+
+ saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+ require.NoError(t, err)
+ assert.Equal(t, "lock-on-change-host", saved.Meta.Hostname)
+ })
+
+ t.Run("LoginPeer without change opens no transaction", func(t *testing.T) {
+ saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+ require.NoError(t, err)
+
+ counting.txs.Store(0)
+ _, _, _, _, err = manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, SSHKey: saved.SSHKey, Meta: saved.Meta})
+ require.NoError(t, err)
+ assert.Zero(t, counting.txs.Load())
+ })
+
+ t.Run("LoginPeer with new ssh key saves in one transaction", func(t *testing.T) {
+ counting.txs.Store(0)
+ _, _, _, _, err := manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, SSHKey: "ssh-ed25519 AAAAlockonchange", Meta: current.Meta})
+ require.NoError(t, err)
+ assert.Equal(t, int32(1), counting.txs.Load())
+
+ saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+ require.NoError(t, err)
+ assert.Equal(t, "ssh-ed25519 AAAAlockonchange", saved.SSHKey)
+ })
+
+ t.Run("LoginPeer rejects a user mismatch without a transaction", func(t *testing.T) {
+ counting.txs.Store(0)
+ _, _, _, _, err := manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, UserID: "someone-else", SSHKey: "ssh-ed25519 AAAAother", Meta: current.Meta})
+ assert.Equal(t, status.NewPeerLoginMismatchError(), err)
+ assert.Zero(t, counting.txs.Load())
+
+ saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+ require.NoError(t, err)
+ assert.Equal(t, "ssh-ed25519 AAAAlockonchange", saved.SSHKey, "a rejected login must not store its ssh key")
+ })
+
+ t.Run("LoginPeer rejects disallowed extra DNS labels without a transaction", func(t *testing.T) {
+ counting.txs.Store(0)
+ _, _, _, _, err := manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, SSHKey: "ssh-ed25519 AAAAlabels", ExtraDNSLabels: []string{"extra"}, Meta: current.Meta})
+ errStatus, ok := status.FromError(err)
+ require.True(t, ok, "expected a status error, got %v", err)
+ assert.Equal(t, status.PreconditionFailed, errStatus.Type())
+ assert.Zero(t, counting.txs.Load())
+ })
+}
+
+func TestCheckPeerMaySync(t *testing.T) {
+ manager, _, account, peer1, peer2, _ := setupNetworkMapTest(t)
+ ctx := context.Background()
+
+ require.NoError(t, manager.CheckPeerMaySync(ctx, account.Id, peer1.Key, nil))
+
+ expired, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+ require.NoError(t, err)
+ expired.Status.LoginExpired = true
+ require.NoError(t, manager.Store.SavePeer(ctx, account.Id, expired))
+ assert.Equal(t, status.NewPeerLoginExpiredError(), manager.CheckPeerMaySync(ctx, account.Id, peer1.Key, nil))
+
+ owner := types.NewRegularUser("blocked-owner", "", "")
+ owner.AccountID = account.Id
+ owner.Blocked = true
+ require.NoError(t, manager.Store.SaveUser(ctx, owner))
+ blocked, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer2.ID)
+ require.NoError(t, err)
+ blocked.UserID = owner.Id
+ require.NoError(t, manager.Store.SavePeer(ctx, account.Id, blocked))
+ err = manager.CheckPeerMaySync(ctx, account.Id, peer2.Key, nil)
+ errStatus, ok := status.FromError(err)
+ require.True(t, ok, "expected a status error, got %v", err)
+ assert.Equal(t, status.PermissionDenied, errStatus.Type())
+
+ assert.Equal(t, status.NewPeerNotRegisteredError(), manager.CheckPeerMaySync(ctx, account.Id, "unknown-key", nil))
+}
+
+type failingPeerLookupStore struct {
+ store.Store
+}
+
+func (s *failingPeerLookupStore) GetPeerByPeerPubKey(context.Context, store.LockingStrength, string) (*nbpeer.Peer, error) {
+ return nil, status.Errorf(status.Internal, "issue getting peer from store: connection reset")
+}
+
+func TestPeerLookupErrorKeepsTransientStoreErrors(t *testing.T) {
+ manager, _, account, peer1, _, _ := setupNetworkMapTest(t)
+ ctx := context.Background()
+ manager.Store = &failingPeerLookupStore{Store: manager.Store}
+
+ err := manager.CheckPeerMaySync(ctx, account.Id, peer1.Key, nil)
+ errStatus, ok := status.FromError(err)
+ require.True(t, ok, "expected a status error, got %v", err)
+ assert.Equal(t, status.Internal, errStatus.Type(), "a store failure must not look like an unregistered peer")
+
+ _, _, _, _, err = manager.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peer1.Key}, account.Id)
+ errStatus, ok = status.FromError(err)
+ require.True(t, ok, "expected a status error, got %v", err)
+ assert.Equal(t, status.Internal, errStatus.Type(), "a store failure must not look like an unregistered peer")
+}
diff --git a/management/server/store/sql_store_peer.go b/management/server/store/sql_store_peer.go
index 77642eea7..b50ee6960 100644
--- a/management/server/store/sql_store_peer.go
+++ b/management/server/store/sql_store_peer.go
@@ -106,7 +106,6 @@ func (s *SqlStore) MarkPeerConnectedIfNewerSession(ctx context.Context, accountI
"peer_status_connected": true,
"peer_status_last_seen": gorm.Expr("CURRENT_TIMESTAMP"),
"peer_status_session_started_at": newSessionStartedAt,
- "peer_status_login_expired": false,
})
if result.Error != nil {
return false, status.Errorf(status.Internal, "mark peer connected: %v", result.Error)
@@ -142,6 +141,51 @@ func (s *SqlStore) MarkPeerDisconnectedIfSameSession(ctx context.Context, accoun
return result.RowsAffected > 0, nil
}
+// markPeersDisconnectedChunkSize bounds the sessions per statement so the
+// bound parameter count stays below the database limits.
+const markPeersDisconnectedChunkSize = 1000
+
+// MarkPeersDisconnectedIfSameSession is the batched form of
+// MarkPeerDisconnectedIfSameSession: one UPDATE per chunk of sessions, each
+// peer guarded by its own (account, peer, session) tuple. A failed chunk does
+// not stop the others, only a done ctx does; the count then covers the chunks
+// that succeeded and the error joins the failures.
+func (s *SqlStore) MarkPeersDisconnectedIfSameSession(ctx context.Context, sessions []types.PeerSession) (int64, error) {
+ var total int64
+ var errs []error
+ for start := 0; start < len(sessions); start += markPeersDisconnectedChunkSize {
+ if err := ctx.Err(); err != nil {
+ errs = append(errs, status.Errorf(status.Internal, "mark peers disconnected: %v", err))
+ break
+ }
+ end := min(start+markPeersDisconnectedChunkSize, len(sessions))
+ tuples := make([][]any, 0, end-start)
+ for _, session := range sessions[start:end] {
+ if session.SessionStartedAt == 0 {
+ continue
+ }
+ tuples = append(tuples, []any{session.AccountID, session.PeerID, session.SessionStartedAt})
+ }
+ if len(tuples) == 0 {
+ continue
+ }
+ result := s.db.WithContext(ctx).
+ Model(&nbpeer.Peer{}).
+ Where("(account_id, id, peer_status_session_started_at) IN ?", tuples).
+ Updates(map[string]any{
+ "peer_status_connected": false,
+ "peer_status_last_seen": gorm.Expr("CURRENT_TIMESTAMP"),
+ "peer_status_session_started_at": int64(0),
+ })
+ if result.Error != nil {
+ errs = append(errs, status.Errorf(status.Internal, "mark peers disconnected: %v", result.Error))
+ continue
+ }
+ total += result.RowsAffected
+ }
+ return total, errors.Join(errs...)
+}
+
// ApproveAccountPeers marks all peers that currently require approval in the given account as approved.
func (s *SqlStore) ApproveAccountPeers(ctx context.Context, accountID string) (int, error) {
result := s.db.Model(&nbpeer.Peer{}).
diff --git a/management/server/store/sql_store_peer_test.go b/management/server/store/sql_store_peer_test.go
index 0b467e37b..56498c2f8 100644
--- a/management/server/store/sql_store_peer_test.go
+++ b/management/server/store/sql_store_peer_test.go
@@ -3,16 +3,20 @@ package store
import (
"context"
"encoding/binary"
+ "errors"
"fmt"
"net"
"net/netip"
"reflect"
"sort"
+ "strings"
+ "sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
+ "gorm.io/gorm"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
@@ -941,3 +945,146 @@ func TestSqlStore_ApproveAccountPeers(t *testing.T) {
})
})
}
+
+func TestSqlStore_MarkPeersDisconnectedIfSameSession(t *testing.T) {
+ ctx := context.Background()
+ store, cleanUp, err := NewTestStoreFromSQL(ctx, "../testdata/extended-store.sql", t.TempDir())
+ t.Cleanup(cleanUp)
+ require.NoError(t, err)
+
+ accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+ sqlStore := store.(*SqlStore)
+
+ sessionOf := map[string]int64{"peer-a": 101, "peer-b": 102, "peer-c": 103, "peer-d": 104}
+ for id, session := range sessionOf {
+ require.NoError(t, store.AddPeerToAccount(ctx, &nbpeer.Peer{
+ ID: id,
+ AccountID: accountID,
+ Key: "key-" + id,
+ DNSLabel: id,
+ IP: netip.AddrFrom4([4]byte{10, 9, 0, byte(session)}),
+ IPv6: netip.MustParseAddr(fmt.Sprintf("fd00::9:%d", session)),
+ }))
+ require.NoError(t, sqlStore.db.Model(&nbpeer.Peer{}).Where("id = ?", id).Updates(map[string]any{
+ "peer_status_connected": true,
+ "peer_status_session_started_at": session,
+ }).Error)
+ }
+
+ sessions := []types.PeerSession{{AccountID: accountID, PeerID: "peer-a", SessionStartedAt: 101}}
+ for i := 0; i < markPeersDisconnectedChunkSize+5; i++ {
+ sessions = append(sessions, types.PeerSession{AccountID: accountID, PeerID: fmt.Sprintf("missing-%d", i), SessionStartedAt: 1})
+ }
+ sessions = append(sessions,
+ types.PeerSession{AccountID: accountID, PeerID: "peer-b", SessionStartedAt: 102},
+ types.PeerSession{AccountID: accountID, PeerID: "peer-c", SessionStartedAt: 103},
+ types.PeerSession{AccountID: accountID, PeerID: "peer-d", SessionStartedAt: 999},
+ types.PeerSession{AccountID: accountID, PeerID: "peer-d", SessionStartedAt: 0},
+ )
+
+ marked, err := store.MarkPeersDisconnectedIfSameSession(ctx, sessions)
+ require.NoError(t, err)
+ assert.Equal(t, int64(3), marked)
+
+ for id, wantConnected := range map[string]bool{"peer-a": false, "peer-b": false, "peer-c": false, "peer-d": true} {
+ peer, err := store.GetPeerByID(ctx, LockingStrengthNone, accountID, id)
+ require.NoError(t, err)
+ assert.Equal(t, wantConnected, peer.Status.Connected, id)
+ wantSession := int64(0)
+ if wantConnected {
+ wantSession = sessionOf[id]
+ }
+ assert.Equal(t, wantSession, peer.Status.SessionStartedAt, id)
+ }
+}
+
+func TestSqlStore_MarkPeersDisconnectedIfSameSessionContinuesAfterFailedChunk(t *testing.T) {
+ ctx := context.Background()
+ store, cleanUp, err := NewTestStoreFromSQL(ctx, "../testdata/extended-store.sql", t.TempDir())
+ t.Cleanup(cleanUp)
+ require.NoError(t, err)
+
+ accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+ sqlStore := store.(*SqlStore)
+
+ sessionOf := map[string]int64{"peer-a": 101, "peer-b": 102, "peer-c": 103}
+ for id, session := range sessionOf {
+ require.NoError(t, store.AddPeerToAccount(ctx, &nbpeer.Peer{
+ ID: id,
+ AccountID: accountID,
+ Key: "key-" + id,
+ DNSLabel: id,
+ IP: netip.AddrFrom4([4]byte{10, 9, 0, byte(session)}),
+ IPv6: netip.MustParseAddr(fmt.Sprintf("fd00::9:%d", session)),
+ }))
+ require.NoError(t, sqlStore.db.Model(&nbpeer.Peer{}).Where("id = ?", id).Updates(map[string]any{
+ "peer_status_connected": true,
+ "peer_status_session_started_at": session,
+ }).Error)
+ }
+
+ sessions := []types.PeerSession{{AccountID: accountID, PeerID: "peer-a", SessionStartedAt: 101}}
+ for i := 0; i < markPeersDisconnectedChunkSize-1; i++ {
+ sessions = append(sessions, types.PeerSession{AccountID: accountID, PeerID: fmt.Sprintf("missing-%d", i), SessionStartedAt: 1})
+ }
+ sessions = append(sessions,
+ types.PeerSession{AccountID: accountID, PeerID: "peer-b", SessionStartedAt: 102},
+ types.PeerSession{AccountID: accountID, PeerID: "peer-c", SessionStartedAt: 103},
+ )
+
+ chunkErr := errors.New("deadlock detected")
+ var updates atomic.Int32
+ require.NoError(t, sqlStore.db.Callback().Update().Before("gorm:update").Register("test:fail_first_chunk", func(db *gorm.DB) {
+ if updates.Add(1) == 1 {
+ _ = db.AddError(chunkErr)
+ }
+ }))
+
+ marked, err := store.MarkPeersDisconnectedIfSameSession(ctx, sessions)
+ require.Error(t, err)
+ assert.ErrorContains(t, err, chunkErr.Error())
+ assert.Equal(t, int64(2), marked, "only the second chunk is counted")
+ assert.Equal(t, int32(2), updates.Load(), "the second chunk must still run")
+
+ for id, wantConnected := range map[string]bool{"peer-a": true, "peer-b": false, "peer-c": false} {
+ peer, err := store.GetPeerByID(ctx, LockingStrengthNone, accountID, id)
+ require.NoError(t, err)
+ assert.Equal(t, wantConnected, peer.Status.Connected, id)
+ }
+
+ cancelled, cancel := context.WithCancel(ctx)
+ cancel()
+ updates.Store(1)
+ marked, err = store.MarkPeersDisconnectedIfSameSession(cancelled, sessions)
+ assert.Zero(t, marked)
+ assert.ErrorContains(t, err, context.Canceled.Error())
+ assert.Equal(t, int32(1), updates.Load(), "no chunk runs once the context is done")
+ assert.Equal(t, 1, strings.Count(err.Error(), context.Canceled.Error()), "a done context is reported once")
+}
+
+func TestSqlStore_MarkPeerConnectedKeepsLoginExpired(t *testing.T) {
+ ctx := context.Background()
+ store, cleanUp, err := NewTestStoreFromSQL(ctx, "../testdata/extended-store.sql", t.TempDir())
+ t.Cleanup(cleanUp)
+ require.NoError(t, err)
+
+ accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+ require.NoError(t, store.AddPeerToAccount(ctx, &nbpeer.Peer{
+ ID: "expired-peer",
+ AccountID: accountID,
+ Key: "key-expired-peer",
+ DNSLabel: "expired-peer",
+ IP: netip.AddrFrom4([4]byte{10, 9, 1, 1}),
+ IPv6: netip.MustParseAddr("fd00::9:101"),
+ Status: &nbpeer.PeerStatus{LoginExpired: true},
+ }))
+
+ updated, err := store.MarkPeerConnectedIfNewerSession(ctx, accountID, "expired-peer", 1)
+ require.NoError(t, err)
+ require.True(t, updated)
+
+ peer, err := store.GetPeerByID(ctx, LockingStrengthNone, accountID, "expired-peer")
+ require.NoError(t, err)
+ assert.True(t, peer.Status.Connected)
+ assert.True(t, peer.Status.LoginExpired, "connecting must not clear a login expiry")
+}
diff --git a/management/server/store/store.go b/management/server/store/store.go
index 670bd3956..b7c68f3c8 100644
--- a/management/server/store/store.go
+++ b/management/server/store/store.go
@@ -196,6 +196,11 @@ type Store interface {
// recorded by the database. Returns true when the update happened,
// false when a newer session has taken over.
MarkPeerDisconnectedIfSameSession(ctx context.Context, accountID, peerID string, sessionStartedAt int64) (bool, error)
+ // MarkPeersDisconnectedIfSameSession applies the semantics of
+ // MarkPeerDisconnectedIfSameSession to many peers in batched statements.
+ // Sessions with a zero SessionStartedAt are skipped. Returns the number
+ // of peers that were marked disconnected.
+ MarkPeersDisconnectedIfSameSession(ctx context.Context, sessions []types.PeerSession) (int64, error)
ApproveAccountPeers(ctx context.Context, accountID string) (int, error)
DeletePeer(ctx context.Context, accountID string, peerID string) error
diff --git a/management/server/store/store_mock.go b/management/server/store/store_mock.go
index 8b89219e5..7638a8b0c 100644
--- a/management/server/store/store_mock.go
+++ b/management/server/store/store_mock.go
@@ -3419,6 +3419,21 @@ func (mr *MockStoreMockRecorder) MarkPeerDisconnectedIfSameSession(ctx, accountI
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeerDisconnectedIfSameSession", reflect.TypeOf((*MockStore)(nil).MarkPeerDisconnectedIfSameSession), ctx, accountID, peerID, sessionStartedAt)
}
+// MarkPeersDisconnectedIfSameSession mocks base method.
+func (m *MockStore) MarkPeersDisconnectedIfSameSession(ctx context.Context, sessions []types3.PeerSession) (int64, error) {
+ m.ctrl.T.Helper()
+ ret := m.ctrl.Call(m, "MarkPeersDisconnectedIfSameSession", ctx, sessions)
+ ret0, _ := ret[0].(int64)
+ ret1, _ := ret[1].(error)
+ return ret0, ret1
+}
+
+// MarkPeersDisconnectedIfSameSession indicates an expected call of MarkPeersDisconnectedIfSameSession.
+func (mr *MockStoreMockRecorder) MarkPeersDisconnectedIfSameSession(ctx, sessions any) *gomock.Call {
+ mr.mock.ctrl.T.Helper()
+ return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeersDisconnectedIfSameSession", reflect.TypeOf((*MockStore)(nil).MarkPeersDisconnectedIfSameSession), ctx, sessions)
+}
+
// MarkPendingJobsAsFailed mocks base method.
func (m *MockStore) MarkPendingJobsAsFailed(ctx context.Context, accountID, peerID, jobID, reason string) error {
m.ctrl.T.Helper()
diff --git a/management/server/types/peer.go b/management/server/types/peer.go
index 885d67bba..3793ccb83 100644
--- a/management/server/types/peer.go
+++ b/management/server/types/peer.go
@@ -18,6 +18,9 @@ type PeerSync struct {
// UpdateAccountPeers indicate updating account peers,
// which occurs when the peer's metadata is updated
UpdateAccountPeers bool
+ // Settings are the account settings the caller already read for this request.
+ // When nil they are loaded from the store.
+ Settings *Settings
}
// PeerLogin used as a data object between the gRPC API and Manager on Login request.
@@ -37,4 +40,14 @@ type PeerLogin struct {
// ExtraDNSLabels is a list of extra DNS labels that the peer wants to use
ExtraDNSLabels []string
+ // Settings are the account settings the caller already read for this request.
+ // When nil they are loaded from the store.
+ Settings *Settings
+}
+
+// PeerSession identifies the sync session a peer is currently owned by.
+type PeerSession struct {
+ AccountID string
+ PeerID string
+ SessionStartedAt int64
}
From 678aff07de9c0cfc1612be8c7fa66a2932c2bdcd Mon Sep 17 00:00:00 2001
From: Tjeerd Ritsma
Date: Tue, 27 Oct 2026 16:57:23 +0100
Subject: [PATCH 3/5] [client] Send the relay hostname as Host when dialing a
relay by IP (#8359)
---
.../client/dialer/ws/dialopts_generic.go | 3 +-
shared/relay/client/dialer/ws/dialopts_js.go | 2 +-
shared/relay/client/dialer/ws/ws.go | 23 ++++-
shared/relay/client/dialer/ws/ws_test.go | 97 +++++++++++++++++++
4 files changed, 122 insertions(+), 3 deletions(-)
diff --git a/shared/relay/client/dialer/ws/dialopts_generic.go b/shared/relay/client/dialer/ws/dialopts_generic.go
index 8008d89d3..5e1cf4c4f 100644
--- a/shared/relay/client/dialer/ws/dialopts_generic.go
+++ b/shared/relay/client/dialer/ws/dialopts_generic.go
@@ -8,8 +8,9 @@ import (
"github.com/coder/websocket"
)
-func createDialOptions(serverName string, underlyingOut *net.Conn) *websocket.DialOptions {
+func createDialOptions(serverName, host string, underlyingOut *net.Conn) *websocket.DialOptions {
return &websocket.DialOptions{
HTTPClient: httpClientNbDialer(serverName, underlyingOut),
+ Host: host,
}
}
diff --git a/shared/relay/client/dialer/ws/dialopts_js.go b/shared/relay/client/dialer/ws/dialopts_js.go
index 5b11fe765..25c9b372c 100644
--- a/shared/relay/client/dialer/ws/dialopts_js.go
+++ b/shared/relay/client/dialer/ws/dialopts_js.go
@@ -8,7 +8,7 @@ import (
"github.com/coder/websocket"
)
-func createDialOptions(_ string, _ *net.Conn) *websocket.DialOptions {
+func createDialOptions(_, _ string, _ *net.Conn) *websocket.DialOptions {
// WASM version doesn't support HTTPClient or custom TLS config.
return &websocket.DialOptions{}
}
diff --git a/shared/relay/client/dialer/ws/ws.go b/shared/relay/client/dialer/ws/ws.go
index 6b310b73d..946d267fb 100644
--- a/shared/relay/client/dialer/ws/ws.go
+++ b/shared/relay/client/dialer/ws/ws.go
@@ -31,8 +31,13 @@ func (d Dialer) Dial(ctx context.Context, address, serverName string) (net.Conn,
return nil, err
}
+ host, err := requestHost(wsURL, serverName)
+ if err != nil {
+ return nil, err
+ }
+
var underlying net.Conn
- opts := createDialOptions(serverName, &underlying)
+ opts := createDialOptions(serverName, host, &underlying)
wsConn, resp, err := websocket.Dial(ctx, wsURL, opts)
if err != nil {
@@ -78,6 +83,22 @@ func prepareURL(address string) (string, error) {
return parsed.String(), nil
}
+// requestHost keeps the server name in the Host header when dialing by IP, so
+// proxies that route by host still match.
+func requestHost(wsURL, serverName string) (string, error) {
+ if serverName == "" {
+ return "", nil
+ }
+ parsed, err := url.Parse(wsURL)
+ if err != nil {
+ return "", fmt.Errorf("parse relay URL %q: %w", wsURL, err)
+ }
+ if port := parsed.Port(); port != "" {
+ return net.JoinHostPort(serverName, port), nil
+ }
+ return serverName, nil
+}
+
// httpClientNbDialer builds the http client used by the websocket library.
// underlyingOut, when non-nil, is populated with the raw conn from the
// transport's DialContext so the caller can read its RemoteAddr.
diff --git a/shared/relay/client/dialer/ws/ws_test.go b/shared/relay/client/dialer/ws/ws_test.go
index 7357adbc0..5f04284f6 100644
--- a/shared/relay/client/dialer/ws/ws_test.go
+++ b/shared/relay/client/dialer/ws/ws_test.go
@@ -1,7 +1,12 @@
package ws
import (
+ "context"
+ "net"
+ "net/http"
+ "net/http/httptest"
"testing"
+ "time"
)
func TestPrepareURL(t *testing.T) {
@@ -74,3 +79,95 @@ func TestPrepareURL(t *testing.T) {
})
}
}
+
+func TestRequestHost(t *testing.T) {
+ tests := []struct {
+ name string
+ wsURL string
+ serverName string
+ want string
+ }{
+ {
+ name: "no server name keeps the URL host",
+ wsURL: "wss://1.2.3.4:443/relay",
+ want: "",
+ },
+ {
+ name: "IP with port",
+ wsURL: "wss://1.2.3.4:443/relay",
+ serverName: "relay.example.com",
+ want: "relay.example.com:443",
+ },
+ {
+ name: "IP with non-standard port",
+ wsURL: "ws://1.2.3.4:45678/relay",
+ serverName: "relay.example.com",
+ want: "relay.example.com:45678",
+ },
+ {
+ name: "IP without port",
+ wsURL: "wss://1.2.3.4/relay",
+ serverName: "relay.example.com",
+ want: "relay.example.com",
+ },
+ {
+ name: "IPv6 with port",
+ wsURL: "wss://[2001:db8::1]:443/relay",
+ serverName: "relay.example.com",
+ want: "relay.example.com:443",
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got, err := requestHost(tt.wsURL, tt.serverName)
+ if err != nil {
+ t.Fatalf("requestHost(%q, %q) err = %v", tt.wsURL, tt.serverName, err)
+ }
+ if got != tt.want {
+ t.Errorf("requestHost(%q, %q) = %q, want %q", tt.wsURL, tt.serverName, got, tt.want)
+ }
+ })
+ }
+}
+
+func TestDialByIPSendsServerNameAsHost(t *testing.T) {
+ hosts := make(chan string, 1)
+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ hosts <- r.Host
+ http.NotFound(w, r)
+ }))
+ defer srv.Close()
+
+ _, port, err := net.SplitHostPort(srv.Listener.Addr().String())
+ if err != nil {
+ t.Fatal(err)
+ }
+
+ tests := []struct {
+ name string
+ serverName string
+ want string
+ }{
+ {name: "dial by IP with server name", serverName: "relay.example.com", want: "relay.example.com:" + port},
+ {name: "dial without server name", want: "127.0.0.1:" + port},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+ defer cancel()
+
+ if _, err := (Dialer{}).Dial(ctx, "rel://127.0.0.1:"+port, tt.serverName); err == nil {
+ t.Fatal("expected the handshake to fail against a plain HTTP server")
+ }
+
+ select {
+ case got := <-hosts:
+ if got != tt.want {
+ t.Errorf("server saw Host %q, want %q", got, tt.want)
+ }
+ case <-ctx.Done():
+ t.Fatal("server never received the handshake request")
+ }
+ })
+ }
+}
From c9936ba702234de07092ab31bca6d9dd140e201f Mon Sep 17 00:00:00 2001
From: Theodor Midtlien
Date: Wed, 28 Oct 2026 12:59:07 +0100
Subject: [PATCH 4/5] [client] Bump wails fork to v3.0.0 (#8370)
* Bump wails fork to v3.0.0
https://github.com/netbirdio/wails/commit/2b40023df5df6b7ae598dee279a769ce8c7efe95
https://github.com/netbirdio/wails/commit/cfa103160796b7a016bd49420ce7268ccea63dd7
---
go.mod | 2 +-
go.sum | 4 ++--
2 files changed, 3 insertions(+), 3 deletions(-)
diff --git a/go.mod b/go.mod
index 699a553af..df3edbb38 100644
--- a/go.mod
+++ b/go.mod
@@ -369,7 +369,7 @@ replace github.com/cloudflare/circl => codeberg.org/cunicu/circl v0.0.0-20230801
replace github.com/mailru/easyjson => github.com/netbirdio/easyjson v0.9.0
-replace github.com/wailsapp/wails/v3 => github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20261009120932-4e087a8bfd6d
+replace github.com/wailsapp/wails/v3 => github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20261028091752-cfa103160796
tool (
github.com/goreleaser/chglog/cmd/chglog
diff --git a/go.sum b/go.sum
index de7b77649..ee70f2c09 100644
--- a/go.sum
+++ b/go.sum
@@ -538,8 +538,8 @@ github.com/netbirdio/management-integrations/integrations v0.0.0-20261021190039-
github.com/netbirdio/management-integrations/integrations v0.0.0-20261021190039-6c1585eb11de/go.mod h1:LLvpw/iWbXX9+7HF7ueqkyE7SoFaeqoCmbskOQL9gbs=
github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20250805121659-6b4ac470ca45 h1:ujgviVYmx243Ksy7NdSwrdGPSRNE3pb8kEDSpH0QuAQ=
github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20250805121659-6b4ac470ca45/go.mod h1:5/sjFmLb8O96B5737VCqhHyGRzNFIaN/Bu7ZodXc3qQ=
-github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20261009120932-4e087a8bfd6d h1:DV0FHFhNAaU69JHXKclYQWr/ZJoc2BA/ZpqlnX8l22U=
-github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20261009120932-4e087a8bfd6d/go.mod h1:62UsqQRqanuCbFi6xElIw3/2X+yGF+TKlyiWImA5taU=
+github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20261028091752-cfa103160796 h1:+21hEHg5uXZiYP7vpEH3Eo3jk5anmynYRUq5dPGErxw=
+github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20261028091752-cfa103160796/go.mod h1:1U8Kg7QyeO09HanpT48heL7yOCPaYfh9ErFVzmxfiLQ=
github.com/netbirdio/wireguard-go v0.0.0-20261021152105-0de46b309cdb h1:3X7pVeZeneq4MK7hvvqeSGi1ir6OElgIqHd6cI4KsXo=
github.com/netbirdio/wireguard-go v0.0.0-20261021152105-0de46b309cdb/go.mod h1:Vk/c438m8FV2hHcT/6Efp8xe/JhHTzf3eUcuaohKKkw=
github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
From 1acdebf395eda4d88b15a50de626bc5ce452d091 Mon Sep 17 00:00:00 2001
From: Bethuel Mmbaga
Date: Wed, 28 Oct 2026 17:02:11 +0300
Subject: [PATCH 5/5] [management] Add configurable statement timeout to
postgres store (#8381)
---
.../accesslogs/manager/repository.go | 37 +++--
.../accesslogs/manager/repository_test.go | 70 +++++++++
management/internals/shared/db/conn.go | 12 +-
management/internals/shared/db/open.go | 27 +++-
management/internals/shared/db/timeout.go | 52 +++++++
.../internals/shared/db/timeout_test.go | 146 ++++++++++++++++++
util/duration.go | 16 ++
util/duration_test.go | 25 +++
8 files changed, 366 insertions(+), 19 deletions(-)
create mode 100644 management/internals/shared/db/timeout.go
create mode 100644 management/internals/shared/db/timeout_test.go
create mode 100644 util/duration_test.go
diff --git a/management/internals/modules/reverseproxy/accesslogs/manager/repository.go b/management/internals/modules/reverseproxy/accesslogs/manager/repository.go
index 3967442d8..1b2cd4976 100644
--- a/management/internals/modules/reverseproxy/accesslogs/manager/repository.go
+++ b/management/internals/modules/reverseproxy/accesslogs/manager/repository.go
@@ -14,18 +14,21 @@ import (
"github.com/netbirdio/netbird/shared/management/status"
)
+const deleteBatchSize = 10000
+
type sqlRepository struct {
- conn *db.Conn
- db *gorm.DB
+ conn *db.Conn
+ db *gorm.DB
+ deleteBatchSize int
}
// NewRepository returns the access log repository backed by conn.
func NewRepository(conn *db.Conn) accesslogs.Repository {
- return &sqlRepository{conn: conn, db: conn.DB(nil)}
+ return &sqlRepository{conn: conn, db: conn.DB(nil), deleteBatchSize: deleteBatchSize}
}
func (r *sqlRepository) WithTx(tx *db.Tx) accesslogs.Repository {
- return &sqlRepository{conn: r.conn, db: r.conn.DB(tx)}
+ return &sqlRepository{conn: r.conn, db: r.conn.DB(tx), deleteBatchSize: r.deleteBatchSize}
}
func (r *sqlRepository) Create(ctx context.Context, entry *accesslogs.AccessLogEntry) error {
@@ -73,12 +76,28 @@ func (r *sqlRepository) ListByAccount(ctx context.Context, lockStrength db.Locki
}
func (r *sqlRepository) DeleteOlderThan(ctx context.Context, olderThan time.Time) (int64, error) {
- result := r.db.Where("timestamp < ?", olderThan).Delete(&accesslogs.AccessLogEntry{})
- if result.Error != nil {
- log.WithContext(ctx).Errorf("failed to delete old access logs: %v", result.Error)
- return 0, status.Errorf(status.Internal, "failed to delete old access logs")
+ var deleted int64
+ for {
+ if err := ctx.Err(); err != nil {
+ return deleted, err
+ }
+
+ batch := r.db.Model(&accesslogs.AccessLogEntry{}).
+ Select("id").
+ Where("timestamp < ?", olderThan).
+ Order("timestamp").
+ Limit(r.deleteBatchSize)
+ result := r.db.Where("id IN (?)", batch).Delete(&accesslogs.AccessLogEntry{})
+ if result.Error != nil {
+ log.WithContext(ctx).Errorf("failed to delete old access logs: %v", result.Error)
+ return deleted, status.Errorf(status.Internal, "failed to delete old access logs")
+ }
+
+ deleted += result.RowsAffected
+ if result.RowsAffected < int64(r.deleteBatchSize) {
+ return deleted, nil
+ }
}
- return result.RowsAffected, nil
}
func applyFilters(query *gorm.DB, filter accesslogs.AccessLogFilter) *gorm.DB {
diff --git a/management/internals/modules/reverseproxy/accesslogs/manager/repository_test.go b/management/internals/modules/reverseproxy/accesslogs/manager/repository_test.go
index 7fea12a85..6cbbd3c60 100644
--- a/management/internals/modules/reverseproxy/accesslogs/manager/repository_test.go
+++ b/management/internals/modules/reverseproxy/accesslogs/manager/repository_test.go
@@ -3,6 +3,8 @@ package manager
import (
"context"
"errors"
+ "fmt"
+ "os"
"testing"
"time"
@@ -12,6 +14,7 @@ import (
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs"
"github.com/netbirdio/netbird/management/internals/shared/db"
"github.com/netbirdio/netbird/management/internals/shared/db/dbtest"
+ "github.com/netbirdio/netbird/management/server/testutil"
)
func newTestRepository(t *testing.T) (accesslogs.Repository, *db.Conn) {
@@ -157,3 +160,70 @@ func TestSqlRepository_ListByAccount_Search(t *testing.T) {
assert.Equal(t, []string{"req-api", "req-web"}, search("example.com"), "search should list every match newest first")
assert.Empty(t, search("nothing-matches"))
}
+
+func TestSqlRepository_DeleteOlderThan_DeletesInBatches(t *testing.T) {
+ cases := map[string]struct{ old, recent int }{
+ "partial last batch": {old: 5, recent: 2},
+ "exact multiple of batch": {old: 4, recent: 1},
+ "nothing to delete": {old: 0, recent: 3},
+ }
+ for name, tc := range cases {
+ t.Run(name, func(t *testing.T) {
+ repo, _ := newTestRepository(t)
+ assertBatchedDelete(t, repo, tc.old, tc.recent)
+ })
+ }
+}
+
+func TestSqlRepository_DeleteOlderThan_DeletesInBatchesOnPostgres(t *testing.T) {
+ if os.Getenv("NETBIRD_STORE_ENGINE") != string(db.PostgresStoreEngine) {
+ t.Skip("requires NETBIRD_STORE_ENGINE=postgres")
+ }
+ cleanup, dsn, err := testutil.CreatePostgresTestContainer()
+ require.NoError(t, err)
+ t.Cleanup(cleanup)
+ conn, err := db.OpenPostgres(context.Background(), dsn, db.DefaultPoolConfig)
+ require.NoError(t, err)
+ t.Cleanup(func() { _ = conn.Close() })
+ require.NoError(t, conn.AutoMigrate(&accesslogs.AccessLogEntry{}))
+ require.NoError(t, conn.DB(nil).Exec("DELETE FROM access_log_entries").Error)
+
+ assertBatchedDelete(t, NewRepository(conn), 5, 2)
+}
+
+func TestSqlRepository_DeleteOlderThan_StopsWhenContextDone(t *testing.T) {
+ repo, _ := newTestRepository(t)
+ require.NoError(t, repo.Create(context.Background(), newEntry("old", "acc", "GET", 48*time.Hour)))
+ ctx, cancel := context.WithCancel(context.Background())
+ cancel()
+
+ deleted, err := repo.DeleteOlderThan(ctx, time.Now().Add(-24*time.Hour))
+
+ require.ErrorIs(t, err, context.Canceled)
+ assert.Zero(t, deleted)
+ assert.EqualValues(t, 1, countEntries(t, repo), "nothing should be deleted once the context is done")
+}
+
+func assertBatchedDelete(t *testing.T, repo accesslogs.Repository, old, recent int) {
+ t.Helper()
+ repo.(*sqlRepository).deleteBatchSize = 2
+ ctx := context.Background()
+ for i := 0; i < old; i++ {
+ require.NoError(t, repo.Create(ctx, newEntry(fmt.Sprintf("old-%d", i), "acc", "GET", 48*time.Hour)))
+ }
+ for i := 0; i < recent; i++ {
+ require.NoError(t, repo.Create(ctx, newEntry(fmt.Sprintf("new-%d", i), "acc", "GET", time.Hour)))
+ }
+
+ deleted, err := repo.DeleteOlderThan(ctx, time.Now().Add(-24*time.Hour))
+ require.NoError(t, err)
+ assert.EqualValues(t, old, deleted, "every old entry should be deleted across batches")
+ assert.EqualValues(t, recent, countEntries(t, repo), "recent entries should be kept")
+}
+
+func countEntries(t *testing.T, repo accesslogs.Repository) int64 {
+ t.Helper()
+ _, total, err := repo.ListByAccount(context.Background(), db.LockingStrengthNone, "acc", accesslogs.AccessLogFilter{Page: 1, PageSize: 10})
+ require.NoError(t, err)
+ return total
+}
diff --git a/management/internals/shared/db/conn.go b/management/internals/shared/db/conn.go
index b21b066e2..e12d8e8c8 100644
--- a/management/internals/shared/db/conn.go
+++ b/management/internals/shared/db/conn.go
@@ -11,6 +11,8 @@ import (
"github.com/jackc/pgx/v5/pgxpool"
log "github.com/sirupsen/logrus"
"gorm.io/gorm"
+
+ "github.com/netbirdio/netbird/util"
)
const (
@@ -64,13 +66,11 @@ func NewConn(ctx context.Context, gormDB *gorm.DB, engine Engine, pool *pgxpool.
opt(&options)
}
- txTimeout := defaultTransactionTimeout
- if v := os.Getenv("NB_STORE_TRANSACTION_TIMEOUT"); v != "" {
- if parsed, err := time.ParseDuration(v); err == nil {
- txTimeout = parsed
- }
- }
+ txTimeout := util.DurationFromEnv(transactionTimeoutEnv, defaultTransactionTimeout)
log.WithContext(ctx).Infof("Setting transaction timeout to %v", txTimeout)
+ if d := util.DurationFromEnv(statementTimeoutEnv, 0); engine == PostgresStoreEngine && d > 0 {
+ log.WithContext(ctx).Infof("Setting statement timeout to %v", d)
+ }
conns := runtime.NumCPU()
configuredConns, err := strconv.Atoi(os.Getenv("NB_SQL_MAX_OPEN_CONNS"))
diff --git a/management/internals/shared/db/open.go b/management/internals/shared/db/open.go
index f813b0ce8..a64a4003f 100644
--- a/management/internals/shared/db/open.go
+++ b/management/internals/shared/db/open.go
@@ -10,7 +10,9 @@ import (
"strings"
"time"
+ "github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
+ "github.com/jackc/pgx/v5/stdlib"
"gorm.io/driver/postgres"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
@@ -100,9 +102,10 @@ func OpenSqliteFile(ctx context.Context, dataDir, storeFile string) (*Conn, erro
return conn, nil
}
-// OpenPostgres opens a Postgres database through gorm and a pgx pool sized by pool.
+// OpenPostgres opens a Postgres database through gorm and a pgx pool sized by
+// pool, applying the configured timeouts to every session.
func OpenPostgres(ctx context.Context, dsn string, pool PoolConfig, opts ...Option) (*Conn, error) {
- gormDB, err := gorm.Open(postgres.Open(dsn), GormConfig())
+ gormDB, err := openGormPostgres(ctx, dsn)
if err != nil {
return nil, err
}
@@ -123,7 +126,7 @@ func OpenPostgres(ctx context.Context, dsn string, pool PoolConfig, opts ...Opti
// OpenPostgresWithoutPool opens a Postgres database through gorm alone, for
// repositories that never take the pgx read paths.
func OpenPostgresWithoutPool(ctx context.Context, dsn string, opts ...Option) (*Conn, error) {
- gormDB, err := gorm.Open(postgres.Open(dsn), GormConfig())
+ gormDB, err := openGormPostgres(ctx, dsn)
if err != nil {
return nil, err
}
@@ -135,13 +138,29 @@ func OpenPostgresWithoutPool(ctx context.Context, dsn string, opts ...Option) (*
return conn, nil
}
-// OpenPgxPool opens a pgx pool to dsn sized by cfg and verifies it with a ping.
+func openGormPostgres(ctx context.Context, dsn string) (*gorm.DB, error) {
+ config, err := pgx.ParseConfig(dsn)
+ if err != nil {
+ return nil, fmt.Errorf("parse database config: %w", err)
+ }
+ sqlDB := stdlib.OpenDB(*config, stdlib.OptionAfterConnect(sessionTimeouts()))
+ gormDB, err := gorm.Open(postgres.New(postgres.Config{Conn: sqlDB}), GormConfig())
+ if err != nil {
+ _ = sqlDB.Close()
+ return nil, err
+ }
+ return gormDB, nil
+}
+
+// OpenPgxPool opens a pgx pool to dsn sized by cfg, applying the configured
+// timeouts to every session, and verifies it with a ping.
func OpenPgxPool(ctx context.Context, dsn string, cfg PoolConfig) (*pgxpool.Pool, error) {
config, err := pgxpool.ParseConfig(dsn)
if err != nil {
return nil, fmt.Errorf("unable to parse database config: %w", err)
}
+ config.AfterConnect = sessionTimeouts()
config.MaxConns = cfg.MaxConns
config.MinConns = cfg.MinConns
config.MaxConnLifetime = cfg.MaxConnLifetime
diff --git a/management/internals/shared/db/timeout.go b/management/internals/shared/db/timeout.go
new file mode 100644
index 000000000..d6bc70296
--- /dev/null
+++ b/management/internals/shared/db/timeout.go
@@ -0,0 +1,52 @@
+package db
+
+import (
+ "context"
+ "fmt"
+ "math"
+ "strconv"
+ "strings"
+ "time"
+
+ "github.com/jackc/pgx/v5"
+
+ "github.com/netbirdio/netbird/util"
+)
+
+const (
+ transactionTimeoutEnv = "NB_STORE_TRANSACTION_TIMEOUT"
+ statementTimeoutEnv = "NB_STORE_STATEMENT_TIMEOUT"
+)
+
+// sessionTimeouts sets the configured timeouts on each new Postgres session with
+// SET, since poolers such as PgBouncer reject them as startup parameters.
+func sessionTimeouts() func(context.Context, *pgx.Conn) error {
+ var settings []string
+ if d := util.DurationFromEnv(statementTimeoutEnv, 0); d > 0 {
+ settings = append(settings, "SET statement_timeout = "+postgresMillis(d))
+ }
+ if d := util.DurationFromEnv(transactionTimeoutEnv, defaultTransactionTimeout); d > 0 {
+ settings = append(settings, "SET idle_in_transaction_session_timeout = "+postgresMillis(d))
+ }
+ query := strings.Join(settings, "; ")
+
+ return func(ctx context.Context, conn *pgx.Conn) error {
+ if query == "" {
+ return nil
+ }
+ if _, err := conn.Exec(ctx, query); err != nil {
+ return fmt.Errorf("set session timeouts: %w", err)
+ }
+ return nil
+ }
+}
+
+// postgresMillis converts d to milliseconds for a Postgres timeout, rounding up
+// so a positive d never becomes 0 (disabled) and capping at the server maximum.
+func postgresMillis(d time.Duration) string {
+ ms := d.Milliseconds()
+ if d%time.Millisecond != 0 {
+ ms++
+ }
+ return strconv.FormatInt(min(ms, math.MaxInt32), 10)
+}
diff --git a/management/internals/shared/db/timeout_test.go b/management/internals/shared/db/timeout_test.go
new file mode 100644
index 000000000..b9678082d
--- /dev/null
+++ b/management/internals/shared/db/timeout_test.go
@@ -0,0 +1,146 @@
+package db
+
+import (
+ "context"
+ "errors"
+ "os"
+ "testing"
+ "time"
+
+ "github.com/jackc/pgx/v5"
+ "github.com/jackc/pgx/v5/pgconn"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+ "gorm.io/gorm"
+
+ "github.com/netbirdio/netbird/management/server/testutil"
+)
+
+// queryCanceledCode is the SQLSTATE Postgres reports when statement_timeout fires.
+const queryCanceledCode = "57014"
+
+func TestPostgresMillis(t *testing.T) {
+ cases := map[time.Duration]string{
+ time.Minute: "60000",
+ 1500 * time.Microsecond: "2",
+ 500 * time.Microsecond: "1",
+ 720 * time.Hour: "2147483647",
+ }
+ for d, want := range cases {
+ assert.Equal(t, want, postgresMillis(d), "milliseconds for %v", d)
+ }
+}
+
+func TestSessionTimeouts_NothingToApply(t *testing.T) {
+ t.Setenv(statementTimeoutEnv, "0")
+ t.Setenv(transactionTimeoutEnv, "0")
+ ctx := context.Background()
+
+ assert.NoError(t, sessionTimeouts()(ctx, nil), "the hook should not touch the session when both timeouts are disabled")
+}
+
+func TestOpenPostgres_BoundsEverySession(t *testing.T) {
+ t.Setenv(statementTimeoutEnv, "500ms")
+ t.Setenv(transactionTimeoutEnv, "3m")
+ conn := openPostgresTestConn(t)
+ ctx := context.Background()
+
+ requireStatementTimeout(t, conn.DB(nil).Exec("SELECT pg_sleep(1)").Error)
+
+ _, err := conn.Pool(nil).Exec(ctx, "SELECT pg_sleep(1)")
+ requireStatementTimeout(t, err)
+
+ var inTx string
+ require.NoError(t, conn.RunInTx(ctx, func(tx *Tx) error {
+ return conn.DB(tx).Raw("SHOW statement_timeout").Scan(&inTx).Error
+ }))
+ assert.Equal(t, "1min", inTx, "transactions should keep their own 1 minute bound")
+
+ assert.Equal(t, "3min", showSetting(t, conn.DB(nil), "idle_in_transaction_session_timeout"),
+ "the idle transaction timeout should follow the transaction timeout")
+}
+
+func TestOpenPostgres_NoStatementTimeoutKeepsServerSetting(t *testing.T) {
+ for name, value := range map[string]string{"unset": "", "zero": "0"} {
+ t.Run(name, func(t *testing.T) {
+ t.Setenv(statementTimeoutEnv, value)
+ if value == "" {
+ require.NoError(t, os.Unsetenv(statementTimeoutEnv))
+ }
+ ctx := context.Background()
+ dsn := postgresTestDSN(t)
+ admin, err := pgx.Connect(ctx, dsn)
+ require.NoError(t, err)
+ t.Cleanup(func() { _ = admin.Close(ctx) })
+ _, err = admin.Exec(ctx, "ALTER ROLE CURRENT_USER SET statement_timeout = '7s'")
+ require.NoError(t, err)
+
+ conn, err := OpenPostgres(ctx, dsn, DefaultPoolConfig)
+ require.NoError(t, err)
+ t.Cleanup(func() { _ = conn.Close() })
+
+ assert.Equal(t, "7s", showSetting(t, conn.DB(nil), "statement_timeout"), "sessions should keep the server setting")
+ })
+ }
+}
+
+func TestOpenPostgres_CapsTimeoutsAtServerMaximum(t *testing.T) {
+ t.Setenv(transactionTimeoutEnv, "720h")
+ conn := openPostgresTestConn(t)
+
+ assert.Equal(t, "2147483647ms", showSetting(t, conn.DB(nil), "idle_in_transaction_session_timeout"),
+ "a timeout above the server maximum should be capped instead of failing the connection")
+}
+
+func TestOpenPostgresWithoutPool_BoundsSessions(t *testing.T) {
+ t.Setenv(statementTimeoutEnv, "500ms")
+ conn, err := OpenPostgresWithoutPool(context.Background(), postgresTestDSN(t))
+ require.NoError(t, err)
+ t.Cleanup(func() { _ = conn.Close() })
+
+ requireStatementTimeout(t, conn.DB(nil).Exec("SELECT pg_sleep(1)").Error)
+}
+
+func TestOpenPgxPool_BoundsSessions(t *testing.T) {
+ t.Setenv(statementTimeoutEnv, "500ms")
+ ctx := context.Background()
+ pool, err := OpenPgxPool(ctx, postgresTestDSN(t), DefaultPoolConfig)
+ require.NoError(t, err)
+ t.Cleanup(pool.Close)
+
+ _, err = pool.Exec(ctx, "SELECT pg_sleep(1)")
+ requireStatementTimeout(t, err)
+}
+
+func postgresTestDSN(t *testing.T) string {
+ t.Helper()
+ if os.Getenv("NETBIRD_STORE_ENGINE") != string(PostgresStoreEngine) {
+ t.Skip("requires NETBIRD_STORE_ENGINE=postgres")
+ }
+ cleanup, dsn, err := testutil.CreatePostgresTestContainer()
+ require.NoError(t, err)
+ t.Cleanup(cleanup)
+ return dsn
+}
+
+func openPostgresTestConn(t *testing.T) *Conn {
+ t.Helper()
+ conn, err := OpenPostgres(context.Background(), postgresTestDSN(t), DefaultPoolConfig)
+ require.NoError(t, err)
+ t.Cleanup(func() { _ = conn.Close() })
+ return conn
+}
+
+func requireStatementTimeout(t *testing.T, err error) {
+ t.Helper()
+ var pgErr *pgconn.PgError
+ require.True(t, errors.As(err, &pgErr), "expected a Postgres error, got %v", err)
+ assert.Equal(t, queryCanceledCode, pgErr.Code, "the statement should be cancelled by statement_timeout")
+}
+
+func showSetting(t *testing.T, db *gorm.DB, name string) string {
+ t.Helper()
+ var value string
+ require.NoError(t, db.Raw("SHOW "+name).Scan(&value).Error)
+ return value
+}
diff --git a/util/duration.go b/util/duration.go
index b657a582d..db4fb3eec 100644
--- a/util/duration.go
+++ b/util/duration.go
@@ -3,7 +3,10 @@ package util
import (
"encoding/json"
"errors"
+ "os"
"time"
+
+ log "github.com/sirupsen/logrus"
)
// Duration is used strictly for JSON requests/responses due to duration marshalling issues
@@ -35,3 +38,16 @@ func (d *Duration) UnmarshalJSON(b []byte) error {
return errors.New("invalid duration")
}
}
+
+func DurationFromEnv(key string, fallback time.Duration) time.Duration {
+ val := os.Getenv(key)
+ if val == "" {
+ return fallback
+ }
+ parsed, err := time.ParseDuration(val)
+ if err != nil || parsed < 0 {
+ log.Warnf("ignoring invalid %s value %q, using %v", key, val, fallback)
+ return fallback
+ }
+ return parsed
+}
diff --git a/util/duration_test.go b/util/duration_test.go
new file mode 100644
index 000000000..784bfcf2f
--- /dev/null
+++ b/util/duration_test.go
@@ -0,0 +1,25 @@
+package util
+
+import (
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/assert"
+)
+
+func TestDurationFromEnv(t *testing.T) {
+ const key = "NB_TEST_DURATION"
+ cases := map[string]time.Duration{
+ "": time.Minute,
+ "30s": 30 * time.Second,
+ "0": 0,
+ "soon": time.Minute,
+ "-1s": time.Minute,
+ }
+ for value, want := range cases {
+ t.Run(value, func(t *testing.T) {
+ t.Setenv(key, value)
+ assert.Equal(t, want, DurationFromEnv(key, time.Minute), "duration for %q", value)
+ })
+ }
+}