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

Merge remote-tracking branch 'origin/main' into dmitri-catch-disconnected-peer

Browse Source
This commit is contained in:
pascal
2026-10-28 15:19:14 +01:00
parent d8f8bd7271 1acdebf395
commit 264a0865d1
36 changed files with 1637 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files
management/internals/shared/grpc/server.go
+111 -22
View File
@@ -91,6 +91,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
@@ -313,7 +319,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)
@@ -321,11 +334,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
}
@@ -333,10 +346,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()
@@ -455,21 +475,28 @@ func (s *Server) heartbeatPeerSession(ctx context.Context, accountID, peerID str
}
}
// 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)
@@ -481,6 +508,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")
@@ -701,6 +782,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"
@@ -727,6 +809,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)),
@@ -735,6 +826,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) {
@@ -745,7 +837,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")
@@ -850,7 +942,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 {
@@ -860,10 +952,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
@@ -938,7 +1032,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
@@ -957,11 +1051,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
}