From 68da6bf3aaecb5880b6d05e15f045b7627bdb081 Mon Sep 17 00:00:00 2001 From: mlsmaycon Date: Sat, 12 Sep 2026 12:53:15 +0000 Subject: [PATCH] [management] Withdraw a cluster address claim that is lost after the write MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A cluster address is claimed two ways: an account-scoped proxy row, and an agent network gateway pin on the address. Each side checked the other before writing — IsClusterAddressAvailable before SaveProxy, HasForeignAccountProxyAtHost before the settings insert — but check and write are separate autocommit statements, so two concurrent claimants could each pass their check and both commit, leaving a pin no proxy will ever serve next to the proxy row that displaces it. Both sides now re-read after they write. Manager.Connect re-asks availability once the proxy row is committed and, if the address is no longer free or the answer is inconclusive, deletes its own row and returns ErrClusterAddressUnavailable, which the connect path reports as AlreadyExists exactly as the pre-write check would have. bootstrapLabeled re-asks ownership once the settings row is committed and withdraws the pin on the same terms. Because both write before they re-read, of two concurrent claimants at least one re-reads after the other has committed and backs off — on sqlite, postgres and mysql alike, since each statement sees every commit before it. Both may back off, which costs a retry; neither keeps a claim the other holds. No lock spans the proxies and settings tables portably, and a claims table would be more machinery than the property needs, so the re-read is the whole mechanism. DeleteProxy is session-guarded like DisconnectProxy, so a stale session withdrawing itself cannot take out a newer session's row. Reported by CodeRabbit on #7402 (CWE-362). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Sa3DsBDP3VciAi4PPG17L6 --- .../internals/modules/agentnetwork/manager.go | 48 ++++++- .../agentnetwork/settings_bootstrap_test.go | 74 ++++++++++- .../reverseproxy/proxy/manager/manager.go | 45 +++++++ .../proxy/manager/manager_test.go | 117 ++++++++++++++++++ .../modules/reverseproxy/proxy/proxy.go | 9 ++ management/internals/shared/grpc/proxy.go | 6 + .../grpc/proxy_connect_authorizer_test.go | 27 +++- management/server/store/sql_store.go | 19 +++ .../store/sql_store_proxy_disconnect_test.go | 44 +++++++ management/server/store/store.go | 4 + management/server/store/store_mock.go | 14 +++ 11 files changed, 401 insertions(+), 6 deletions(-) diff --git a/management/internals/modules/agentnetwork/manager.go b/management/internals/modules/agentnetwork/manager.go index 44c7ce59e..46582e91d 100644 --- a/management/internals/modules/agentnetwork/manager.go +++ b/management/internals/modules/agentnetwork/manager.go @@ -1097,8 +1097,7 @@ func (m *managerImpl) validateGatewayCluster(ctx context.Context, accountID, clu return fmt.Errorf("check proxy cluster ownership: %w", err) } if foreign { - return status.Errorf(status.InvalidArgument, - "proxy cluster %s is not available to this account", clusterAddr) + return errForeignCluster(clusterAddr) } declared, err := m.accountClusterSpellings(ctx, accountID, clusterAddr) @@ -1227,12 +1226,55 @@ func (m *managerImpl) bootstrapLabeled(ctx context.Context, settings *types.Sett } return fmt.Errorf("create agent network settings: %w", err) } - return nil + return m.confirmGatewayClusterOwnership(ctx, settings) } return fmt.Errorf("allocate agent network endpoint for account %s: %d attempts exhausted", settings.AccountID, maxDomainAllocationAttempts) } +// confirmGatewayClusterOwnership re-asks, once the settings row is committed, +// whether another account's proxy declares the pinned cluster, and withdraws +// the row if one does. +// +// validateGatewayCluster answered that before the insert, but the two are +// separate statements: a foreign proxy can register at the host in between, +// and its own availability check — run before its row is written — would not +// have seen this pin yet either. Re-reading after the write closes that +// window from this side, and Manager.Connect does the same from the proxy's: +// both claimants write before they re-read, so of two concurrent claims at +// least one re-reads after the other has committed and backs off. Each +// statement runs autocommit, so that re-read sees every commit before it on +// sqlite, postgres and mysql alike. Both may back off, which costs the caller +// a retry; neither keeps a claim the other holds, which is the invariant. +// No lock spans the proxies and settings tables portably, and a claims table +// would be more machinery than the property needs. +// +// Only ownership is re-asked. The capability check is about what the cluster +// can do, not who holds it, and does not race a claim. +func (m *managerImpl) confirmGatewayClusterOwnership(ctx context.Context, settings *types.Settings) error { + foreign, err := m.store.HasForeignAccountProxyAtHost(ctx, settings.ProxyAddress, settings.AccountID) + if err == nil && !foreign { + return nil + } + + if delErr := m.store.DeleteAgentNetworkSettings(ctx, settings.AccountID); delErr != nil { + log.WithContext(ctx).Errorf("failed to withdraw agent network settings for account %s after losing the claim on %s: %v", + settings.AccountID, settings.ProxyAddress, delErr) + } + if err != nil { + return fmt.Errorf("confirm proxy cluster ownership: %w", err) + } + log.WithContext(ctx).Warnf("proxy cluster %s was claimed by another account while account %s bootstrapped onto it, withdrawing the pin", + settings.ProxyAddress, settings.AccountID) + return errForeignCluster(settings.ProxyAddress) +} + +// errForeignCluster is the refusal for a cluster another account's proxy +// declares, worded the same whether it is caught before or after the insert. +func errForeignCluster(clusterAddr string) error { + return status.Errorf(status.InvalidArgument, "proxy cluster %s is not available to this account", clusterAddr) +} + // isUniqueConstraintError reports whether err is a database unique-constraint // violation, matched on the driver message because CreateAgentNetworkSettings // deliberately returns the driver error unwrapped. diff --git a/management/internals/modules/agentnetwork/settings_bootstrap_test.go b/management/internals/modules/agentnetwork/settings_bootstrap_test.go index 34f7ccb59..0e7a9f6a6 100644 --- a/management/internals/modules/agentnetwork/settings_bootstrap_test.go +++ b/management/internals/modules/agentnetwork/settings_bootstrap_test.go @@ -7,9 +7,9 @@ import ( "testing" "time" - "go.uber.org/mock/gomock" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types" "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy" @@ -35,6 +35,16 @@ type bootstrapFixture struct { } func newBootstrapFixture(t *testing.T) *bootstrapFixture { + t.Helper() + return newBootstrapFixtureWith(t, func(st store.Store) store.Store { return st }) +} + +// newBootstrapFixtureWith hands the manager the real store as seen through +// wrap, while the fixture keeps the unwrapped store for seeding and +// assertions. It exists for cases that need something to happen between two +// of the manager's store calls — a competing claim landing mid-bootstrap — +// which a real store cannot be made to do on cue. +func newBootstrapFixtureWith(t *testing.T, wrap func(store.Store) store.Store) *bootstrapFixture { t.Helper() if runtime.GOOS == "windows" { t.Skip("sqlite store not properly supported on Windows yet") @@ -55,7 +65,7 @@ func newBootstrapFixture(t *testing.T) *bootstrapFixture { vendor := &stubLister{} return &bootstrapFixture{ - manager: NewManager(st, perms, accounts, nil, WithModelLister(vendor)), + manager: NewManager(wrap(st), perms, accounts, nil, WithModelLister(vendor)), store: st, perms: perms, vendor: vendor, @@ -403,6 +413,66 @@ func TestCreateSettingsRejectsHostAnotherAccountClaims(t *testing.T) { assert.Error(t, err, "no row may be left behind by a rejected bootstrap") } +// claimingStore is a store.Store on which another account's proxy registers +// at the host being pinned in the moment the settings row is written — the +// interleaving a concurrent proxy connect produces when it passes its own +// availability check before this bootstrap's row exists, so neither side's +// pre-write check sees the other. +type claimingStore struct { + store.Store + t *testing.T + claim *proxy.Proxy + claimed bool +} + +func (s *claimingStore) CreateAgentNetworkSettings(ctx context.Context, settings *types.Settings) error { + if !s.claimed { + s.claimed = true + require.NoError(s.t, s.Store.SaveProxy(ctx, s.claim), "the competing claim must land") + } + return s.Store.CreateAgentNetworkSettings(ctx, settings) +} + +// TestCreateSettingsWithdrawsPinClaimedDuringBootstrap covers the window +// between validateGatewayCluster and the insert: a foreign proxy that claims +// the host in that window is seen by the ownership re-read after the write, +// and the pin is withdrawn rather than left standing on a cluster that will +// never serve it. The refusal reads exactly as it would have had the +// pre-write check caught the claim. +func TestCreateSettingsWithdrawsPinClaimedDuringBootstrap(t *testing.T) { + ctx := context.Background() + const host = "shared.example.com" + + f := newBootstrapFixtureWith(t, func(st store.Store) store.Store { + return &claimingStore{Store: st, t: t, claim: &proxy.Proxy{ + ID: "foreign", + ClusterAddress: host, + Status: proxy.StatusConnected, + LastSeen: time.Now().UTC(), + AccountID: ptrTo("account2"), + Capabilities: proxy.Capabilities{Private: ptrTo(true)}, + }} + }) + // A shared embedded cluster, so the pre-write validation passes on its + // own merits and only the claim landing mid-bootstrap can refuse it. + f.seedEmbeddedCluster(t, host) + f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true) + + _, err := f.createSettings(ctx, "account1", "user1", host, "") + require.Error(t, err, "a host claimed by another account mid-bootstrap must be refused") + var sErr *status.Error + require.ErrorAs(t, err, &sErr) + assert.Equal(t, status.InvalidArgument, sErr.Type()) + assert.Contains(t, err.Error(), "not available to this account") + + _, err = f.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "account1") + assert.Error(t, err, "the pin written before the claim was seen must be withdrawn") + + foreign, err := f.store.HasForeignAccountProxyAtHost(ctx, host, "account1") + require.NoError(t, err) + assert.True(t, foreign, "the competing claim, having landed first, keeps the host") +} + // TestCreateSettingsAcceptsSharedClusterAlongsideOwnProxy pins the other side // of that ordering: a shared (NetBird-operated) proxy is not foreign, so // asking the ownership question first must not refuse the cluster most diff --git a/management/internals/modules/reverseproxy/proxy/manager/manager.go b/management/internals/modules/reverseproxy/proxy/manager/manager.go index f7a70b60f..bde583762 100644 --- a/management/internals/modules/reverseproxy/proxy/manager/manager.go +++ b/management/internals/modules/reverseproxy/proxy/manager/manager.go @@ -2,6 +2,7 @@ package manager import ( "context" + "fmt" "time" log "github.com/sirupsen/logrus" @@ -14,6 +15,7 @@ import ( type store interface { SaveProxy(ctx context.Context, p *proxy.Proxy) error DisconnectProxy(ctx context.Context, proxyID, sessionID string) error + DeleteProxy(ctx context.Context, proxyID, sessionID string) error UpdateProxyHeartbeat(ctx context.Context, p *proxy.Proxy) error GetActiveProxyClusterAddresses(ctx context.Context) ([]string, error) GetActiveProxyClusterAddressesForAccount(ctx context.Context, accountID string) ([]string, error) @@ -74,6 +76,12 @@ func (m *Manager) Connect(ctx context.Context, proxyID, sessionID, clusterAddres return nil, err } + if accountID != nil { + if err := m.confirmClusterAddressClaim(ctx, p, *accountID); err != nil { + return nil, err + } + } + log.WithContext(ctx).WithFields(log.Fields{ "proxyID": proxyID, "sessionID": sessionID, @@ -84,6 +92,43 @@ func (m *Manager) Connect(ctx context.Context, proxyID, sessionID, clusterAddres return p, nil } +// confirmClusterAddressClaim re-asks, once the proxy's row is committed, +// whether the account may hold the address, and withdraws the row if not. +// +// The connect path checks IsClusterAddressAvailable before Connect, but that +// read and the write here are separate statements: another claim — a foreign +// proxy row, or another account's agent network gateway pin — can land in +// between, and its own check would not have seen this row yet either. +// Re-reading after the write closes that window from this side, and the +// gateway bootstrap does the same from its side: both claimants write before +// they re-read, so of two concurrent claims at least one re-reads after the +// other has committed and backs off. Each statement runs autocommit, so that +// re-read sees every commit before it on sqlite, postgres and mysql alike. +// Both may back off, which costs a reconnect; neither keeps a claim the other +// holds, which is the invariant. No lock spans the proxies and settings +// tables portably, and a claims table would be more machinery than the +// property needs. +// +// The row is withdrawn on an inconclusive re-read too: a claim that cannot be +// confirmed must not stand, and the proxy reconnects on its own. +func (m *Manager) confirmClusterAddressClaim(ctx context.Context, p *proxy.Proxy, accountID string) error { + available, err := m.IsClusterAddressAvailable(ctx, p.ClusterAddress, accountID) + if err == nil && available { + return nil + } + + if delErr := m.store.DeleteProxy(ctx, p.ID, p.SessionID); delErr != nil { + log.WithContext(ctx).Errorf("failed to withdraw proxy %s session %s after losing the claim on %s: %v", + p.ID, p.SessionID, p.ClusterAddress, delErr) + } + if err != nil { + return fmt.Errorf("confirm claim on cluster address %s: %w", p.ClusterAddress, err) + } + log.WithContext(ctx).Warnf("cluster address %s was claimed while proxy %s registered for account %s, withdrawing its row", + p.ClusterAddress, p.ID, accountID) + return fmt.Errorf("cluster address %s: %w", p.ClusterAddress, proxy.ErrClusterAddressUnavailable) +} + // Disconnect marks a proxy as disconnected in the database. func (m *Manager) Disconnect(ctx context.Context, proxyID, sessionID string) error { if err := m.store.DisconnectProxy(ctx, proxyID, sessionID); err != nil { diff --git a/management/internals/modules/reverseproxy/proxy/manager/manager_test.go b/management/internals/modules/reverseproxy/proxy/manager/manager_test.go index a338a2cc5..ea0f19ef8 100644 --- a/management/internals/modules/reverseproxy/proxy/manager/manager_test.go +++ b/management/internals/modules/reverseproxy/proxy/manager/manager_test.go @@ -26,6 +26,7 @@ type mockStore struct { isClusterAddressConflictingFunc func(ctx context.Context, clusterAddress, accountID string) (bool, error) hasGatewayPinnedByOtherAccountFunc func(ctx context.Context, host, accountID string) (bool, error) deleteAccountClusterFunc func(ctx context.Context, clusterAddress, accountID string) error + deleteProxyFunc func(ctx context.Context, proxyID, sessionID string) error } func (m *mockStore) SaveProxy(ctx context.Context, p *proxy.Proxy) error { @@ -40,6 +41,12 @@ func (m *mockStore) DisconnectProxy(ctx context.Context, proxyID, sessionID stri } return nil } +func (m *mockStore) DeleteProxy(ctx context.Context, proxyID, sessionID string) error { + if m.deleteProxyFunc != nil { + return m.deleteProxyFunc(ctx, proxyID, sessionID) + } + return nil +} func (m *mockStore) UpdateProxyHeartbeat(ctx context.Context, p *proxy.Proxy) error { if m.updateProxyHeartbeatFunc != nil { return m.updateProxyHeartbeatFunc(ctx, p) @@ -407,3 +414,113 @@ func TestIsClusterAddressAvailableSurfacesGatewayPinError(t *testing.T) { require.Error(t, err) assert.False(t, available) } + +// TestConnect_WithdrawsClaimLostDuringRegistration covers the window between +// the connect path's availability check and the row being written: a claim +// that lands there — another account's proxy row or gateway pin — is seen by +// the re-read after the write, and the proxy's own row is withdrawn rather +// than left standing next to it. The refusal carries +// ErrClusterAddressUnavailable so the connect path reports it exactly as it +// would have had the pre-write check caught it. +func TestConnect_WithdrawsClaimLostDuringRegistration(t *testing.T) { + accountID := "acc-1" + + cases := map[string]func(s *mockStore, landed *bool){ + "another account pinned its gateway to the address": func(s *mockStore, landed *bool) { + s.hasGatewayPinnedByOtherAccountFunc = func(_ context.Context, _, _ string) (bool, error) { return *landed, nil } + }, + "another account's proxy declared the address": func(s *mockStore, landed *bool) { + s.isClusterAddressConflictingFunc = func(_ context.Context, _, _ string) (bool, error) { return *landed, nil } + }, + } + for name, arm := range cases { + t.Run(name, func(t *testing.T) { + landed := false + var withdrawn []string + s := &mockStore{ + // The competing claim commits as this row is written: the + // pre-write check (not exercised here) saw nothing, the + // re-read must. + saveProxyFunc: func(_ context.Context, _ *proxy.Proxy) error { landed = true; return nil }, + deleteProxyFunc: func(_ context.Context, proxyID, sessionID string) error { + withdrawn = append(withdrawn, proxyID+"/"+sessionID) + return nil + }, + } + arm(s, &landed) + + mgr := newTestManager(s) + p, err := mgr.Connect(context.Background(), "proxy-1", "session-1", "gw.example.com", "10.0.0.1", &accountID, nil) + require.ErrorIs(t, err, proxy.ErrClusterAddressUnavailable, "a claim lost after the write must surface as the address being unavailable") + assert.Nil(t, p, "no record may be handed back for a withdrawn registration") + assert.Equal(t, []string{"proxy-1/session-1"}, withdrawn, "exactly this session's row must be withdrawn") + }) + } +} + +// TestConnect_WithdrawsClaimWhenRecheckFails pins fail-closed: a re-read that +// cannot answer leaves the row withdrawn and the connect refused, rather than +// letting a claim stand that was never confirmed. The error is the store's, +// not ErrClusterAddressUnavailable — nothing established that the address is +// taken. +func TestConnect_WithdrawsClaimWhenRecheckFails(t *testing.T) { + accountID := "acc-1" + var withdrawn int + s := &mockStore{ + hasGatewayPinnedByOtherAccountFunc: func(_ context.Context, _, _ string) (bool, error) { + return false, errors.New("db unavailable") + }, + deleteProxyFunc: func(_ context.Context, _, _ string) error { withdrawn++; return nil }, + } + + mgr := newTestManager(s) + _, err := mgr.Connect(context.Background(), "proxy-1", "session-1", "gw.example.com", "10.0.0.1", &accountID, nil) + require.Error(t, err) + assert.NotErrorIs(t, err, proxy.ErrClusterAddressUnavailable, "an inconclusive re-read is not a conflict") + assert.ErrorContains(t, err, "db unavailable", "the store's error must be the one surfaced") + assert.Equal(t, 1, withdrawn, "an unconfirmed claim must be withdrawn") +} + +// TestConnect_ConfirmedClaimKeepsRow is the common case: nothing landed in the +// window, the re-read confirms the claim, and the row stays. +func TestConnect_ConfirmedClaimKeepsRow(t *testing.T) { + accountID := "acc-1" + s := &mockStore{ + deleteProxyFunc: func(_ context.Context, proxyID, _ string) error { + t.Fatalf("a confirmed claim must not be withdrawn, but proxy %s was", proxyID) + return nil + }, + } + + mgr := newTestManager(s) + p, err := mgr.Connect(context.Background(), "proxy-1", "session-1", "gw.example.com", "10.0.0.1", &accountID, nil) + require.NoError(t, err) + require.NotNil(t, p) + assert.Equal(t, proxy.StatusConnected, p.Status) +} + +// TestConnect_SharedProxySkipsClaimRecheck pins that a shared, NetBird-operated +// proxy — no account on its token — is not subject to the claim re-read: the +// connect path never asks availability for it before the write either, and a +// shared cluster is what accounts pin their gateways to, not a claim against +// them. +func TestConnect_SharedProxySkipsClaimRecheck(t *testing.T) { + s := &mockStore{ + isClusterAddressConflictingFunc: func(_ context.Context, _, _ string) (bool, error) { + t.Fatal("a shared proxy must not be checked for address conflicts") + return false, nil + }, + hasGatewayPinnedByOtherAccountFunc: func(_ context.Context, _, _ string) (bool, error) { + t.Fatal("a shared proxy must not be checked against gateway pins") + return false, nil + }, + deleteProxyFunc: func(_ context.Context, _, _ string) error { + t.Fatal("a shared proxy's row must not be withdrawn") + return nil + }, + } + + mgr := newTestManager(s) + _, err := mgr.Connect(context.Background(), "proxy-1", "session-1", "eu.proxy.netbird.io", "10.0.0.1", nil, nil) + require.NoError(t, err) +} diff --git a/management/internals/modules/reverseproxy/proxy/proxy.go b/management/internals/modules/reverseproxy/proxy/proxy.go index 4404b0d24..a7494cd79 100644 --- a/management/internals/modules/reverseproxy/proxy/proxy.go +++ b/management/internals/modules/reverseproxy/proxy/proxy.go @@ -1,6 +1,7 @@ package proxy import ( + "errors" "time" ) @@ -9,6 +10,14 @@ const ( StatusDisconnected = "disconnected" ) +// ErrClusterAddressUnavailable is returned by Manager.Connect when the +// cluster address turned out to be claimed by someone else once the proxy's +// own row was written — a conflicting proxy row or another account's agent +// network gateway pin that landed between the availability check and the +// write. The proxy's row has been withdrawn by then; the caller reports the +// address as taken, exactly as if the pre-write check had caught it. +var ErrClusterAddressUnavailable = errors.New("cluster address is not available") + // Capabilities describes what a proxy can handle, as reported via gRPC. // Nil fields mean the proxy never reported this capability. type Capabilities struct { diff --git a/management/internals/shared/grpc/proxy.go b/management/internals/shared/grpc/proxy.go index 48e64e314..47cc4cf5b 100644 --- a/management/internals/shared/grpc/proxy.go +++ b/management/internals/shared/grpc/proxy.go @@ -571,6 +571,12 @@ func (s *ProxyServiceServer) registerProxyConnection(ctx context.Context, params proxyRecord, err := s.proxyManager.Connect(ctx, params.proxyID, sessionID, params.address, peerInfo, accountID, caps) if err != nil { cancel() + if errors.Is(err, proxy.ErrClusterAddressUnavailable) { + // The claim was lost to a concurrent one after validateProxyConnect + // saw the address free; the row has been withdrawn. Same answer + // as the pre-write check gives, so the proxy treats both alike. + return nil, nil, status.Errorf(codes.AlreadyExists, "cluster address %s is already in use", params.address) + } if accountID != nil { return nil, nil, status.Errorf(codes.Internal, "failed to register BYOP proxy: %v", err) } diff --git a/management/internals/shared/grpc/proxy_connect_authorizer_test.go b/management/internals/shared/grpc/proxy_connect_authorizer_test.go index d0d196d20..c8dae29fb 100644 --- a/management/internals/shared/grpc/proxy_connect_authorizer_test.go +++ b/management/internals/shared/grpc/proxy_connect_authorizer_test.go @@ -3,11 +3,12 @@ package grpc import ( "context" "errors" + "fmt" "testing" - "go.uber.org/mock/gomock" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" "google.golang.org/grpc/codes" grpcstatus "google.golang.org/grpc/status" @@ -166,3 +167,27 @@ func TestValidateProxyConnect_AuthorizerRunsLast(t *testing.T) { assert.Equal(t, codes.AlreadyExists, st.Code(), "an address conflict must keep its own status") assert.Zero(t, auth.called, "a conflicting address must be rejected before policy runs") } + +// TestRegisterProxyConnection_LostClaimIsAlreadyExists pins the status a proxy +// sees when its claim is lost after validateProxyConnect passed: the manager +// withdraws the row and reports ErrClusterAddressUnavailable, and the connect +// path must answer AlreadyExists — the same code the pre-write check gives — +// rather than the Internal it uses for a store failure, so the proxy handles +// the two paths to "address taken" identically. +func TestRegisterProxyConnection_LostClaimIsAlreadyExists(t *testing.T) { + ctrl := gomock.NewController(t) + mgr := proxy.NewMockManager(ctrl) + mgr.EXPECT(). + Connect(gomock.Any(), "proxy-1", gomock.Any(), "cluster.example.com", gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("cluster address cluster.example.com: %w", proxy.ErrClusterAddressUnavailable)) + s := &ProxyServiceServer{proxyManager: mgr} + + _, _, err := s.registerProxyConnection(scopedCtx("acc-1"), proxyConnectParams{proxyID: "proxy-1", address: "cluster.example.com"}, &proxyConnection{}) + require.Error(t, err) + st, ok := grpcstatus.FromError(err) + require.True(t, ok) + assert.Equal(t, codes.AlreadyExists, st.Code(), "a claim lost after the check must read as the address being taken") + assert.Contains(t, st.Message(), "already in use") + _, tracked := s.connectedProxies.Load("proxy-1") + assert.False(t, tracked, "a refused registration must not be tracked as connected") +} diff --git a/management/server/store/sql_store.go b/management/server/store/sql_store.go index 6cd3abd95..ef811966f 100644 --- a/management/server/store/sql_store.go +++ b/management/server/store/sql_store.go @@ -6292,6 +6292,25 @@ func (s *SqlStore) DisconnectProxy(ctx context.Context, proxyID, sessionID strin return nil } +// DeleteProxy removes the proxy's row, but only while it still carries the +// given session: a registration withdrawing its own claim must not take out a +// newer session's row for the same proxy. A row already superseded or gone is +// not an error — the claim it would have withdrawn is no longer this session's +// to withdraw. +func (s *SqlStore) DeleteProxy(ctx context.Context, proxyID, sessionID string) error { + result := s.db. + Where("id = ? AND session_id = ?", proxyID, sessionID). + Delete(&proxy.Proxy{}) + if result.Error != nil { + log.WithContext(ctx).Errorf("failed to delete proxy %s session %s: %v", proxyID, sessionID, result.Error) + return status.Errorf(status.Internal, "failed to delete proxy") + } + if result.RowsAffected == 0 { + log.WithContext(ctx).Debugf("proxy %s session %s: no row deleted (superseded by newer session)", proxyID, sessionID) + } + return nil +} + // GetAllProxies returns all reverse proxy instance rows. func (s *SqlStore) GetAllProxies(ctx context.Context) ([]*proxy.Proxy, error) { var proxies []*proxy.Proxy diff --git a/management/server/store/sql_store_proxy_disconnect_test.go b/management/server/store/sql_store_proxy_disconnect_test.go index 2d0f34680..e388bb761 100644 --- a/management/server/store/sql_store_proxy_disconnect_test.go +++ b/management/server/store/sql_store_proxy_disconnect_test.go @@ -154,3 +154,47 @@ func TestSqlStore_GetAllProxies_Empty(t *testing.T) { assert.Empty(t, all) }) } + +// TestSqlStore_DeleteProxy guards the withdrawal a registration makes when +// its claim on a cluster address is lost after the row was written: +// +// 1. The delete is session-guarded, like DisconnectProxy — a stale session +// withdrawing itself must not take out the row a newer session of the +// same proxy has since written. +// 2. A row that is already gone, or already superseded, is not an error; +// the claim it would have withdrawn is no longer this session's. +// 3. Other proxies at the same address are untouched: only the one row is +// withdrawn, not the cluster. +func TestSqlStore_DeleteProxy(t *testing.T) { + if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" { + t.Skip("skip CI tests on darwin and windows") + } + + runTestForAllEngines(t, "", func(t *testing.T, store Store) { + ctx := context.Background() + accountID := "acct-withdraw" + now := time.Now() + + for _, p := range []*rpproxy.Proxy{ + {ID: "p-withdrawn", SessionID: "sess-new", ClusterAddress: "byop.example.com", LastSeen: now, Status: rpproxy.StatusConnected, AccountID: &accountID}, + {ID: "p-neighbour", SessionID: "sess-1", ClusterAddress: "byop.example.com", LastSeen: now, Status: rpproxy.StatusConnected, AccountID: &accountID}, + } { + require.NoError(t, store.SaveProxy(ctx, p)) + } + + require.NoError(t, store.DeleteProxy(ctx, "p-withdrawn", "sess-old"), + "a delete under a superseded session must be a no-op, not an error") + remaining, err := store.GetAllProxies(ctx) + require.NoError(t, err) + assert.Len(t, remaining, 2, "a superseded session must not withdraw the newer session's row") + + require.NoError(t, store.DeleteProxy(ctx, "p-withdrawn", "sess-new")) + remaining, err = store.GetAllProxies(ctx) + require.NoError(t, err) + require.Len(t, remaining, 1, "the withdrawing session's own row must be gone") + assert.Equal(t, "p-neighbour", remaining[0].ID, "the other proxy at the address must be untouched") + + require.NoError(t, store.DeleteProxy(ctx, "p-withdrawn", "sess-new"), + "withdrawing a row that is already gone must be a no-op, not an error") + }) +} diff --git a/management/server/store/store.go b/management/server/store/store.go index 203ced596..fb8f01d55 100644 --- a/management/server/store/store.go +++ b/management/server/store/store.go @@ -325,6 +325,7 @@ type Store interface { SaveProxy(ctx context.Context, proxy *proxy.Proxy) error DisconnectProxy(ctx context.Context, proxyID, sessionID string) error + DeleteProxy(ctx context.Context, proxyID, sessionID string) error UpdateProxyHeartbeat(ctx context.Context, p *proxy.Proxy) error GetActiveProxyClusterAddresses(ctx context.Context) ([]string, error) GetActiveProxyClusterAddressesForAccount(ctx context.Context, accountID string) ([]string, error) @@ -628,6 +629,9 @@ func getMigrationsPreAuto(ctx context.Context) []migrationFunc { func(db *gorm.DB) error { return migration.MigrateAgentNetworkSettingsToDomain(ctx, db) }, + func(db *gorm.DB) error { + return migration.NormalizeAgentNetworkSettingsIdentity(ctx, db) + }, } } diff --git a/management/server/store/store_mock.go b/management/server/store/store_mock.go index f0299a32c..a31004e84 100644 --- a/management/server/store/store_mock.go +++ b/management/server/store/store_mock.go @@ -754,6 +754,20 @@ func (mr *MockStoreMockRecorder) DeletePostureChecks(ctx, accountID, postureChec return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeletePostureChecks", reflect.TypeOf((*MockStore)(nil).DeletePostureChecks), ctx, accountID, postureChecksID) } +// DeleteProxy mocks base method. +func (m *MockStore) DeleteProxy(ctx context.Context, proxyID, sessionID string) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteProxy", ctx, proxyID, sessionID) + ret0, _ := ret[0].(error) + return ret0 +} + +// DeleteProxy indicates an expected call of DeleteProxy. +func (mr *MockStoreMockRecorder) DeleteProxy(ctx, proxyID, sessionID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteProxy", reflect.TypeOf((*MockStore)(nil).DeleteProxy), ctx, proxyID, sessionID) +} + // DeleteRoute mocks base method. func (m *MockStore) DeleteRoute(ctx context.Context, accountID, routeID string) error { m.ctrl.T.Helper()