sendnrw/netbird
2
0
Fork 0
mirror of https://github.com/netbirdio/netbird.git synced 2026-10-28 07:19:08 +01:00
Code Issues Packages Projects Releases Wiki Activity

[management] batch disconnect and lock free reconnect (#8373)

Browse Source
This commit is contained in:
Pascal Fischer
2026-10-27 16:24:46 +01:00
committed by GitHub
parent 36339993bc
commit 09f1de12e8
18 changed files with 1094 additions and 144 deletions
Download Patch File Download Diff File Expand all files Collapse all files
management/internals/server/boot.go
+19 -9
View File
@@ -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
})
}
management/internals/server/container.go
+33 -4
View File
@@ -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 {
management/internals/server/server.go
+46 -7
View File
@@ -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)
}
}
management/internals/shared/grpc/server.go
+111 -22
View File
@@ -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)
management/internals/shared/grpc/server_disconnect_test.go
+145
View File
@@ -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
}
management/server/account.go
+14 -3
View File
@@ -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 {
management/server/account/manager.go
+3 -1
View File
@@ -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)
management/server/account/manager_mock.go
+32 -4
View File
@@ -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.
management/server/account_test.go
+1 -1
View File
@@ -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)
}
management/server/management_proto_test.go
+129
View File
@@ -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")
}
management/server/mock_server/account_mock.go
+19 -3
View File
@@ -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
management/server/peer.go
+183 -89
View File
@@ -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()
}
management/server/peer_test.go
+134
View File
@@ -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")
}
management/server/store/sql_store_peer.go
+45 -1
View File
@@ -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{}).
management/server/store/sql_store_peer_test.go
+147
View File
@@ -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")
}
management/server/store/store.go
+5
View File
@@ -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
management/server/store/store_mock.go
+15
View File
@@ -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()
management/server/types/peer.go
+13
View File
@@ -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
}