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] [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 }