[management] Keep an established claim when its re-read is inconclusive

SaveProxy upserts on the proxy ID, so on a reconnect the row Connect just wrote is the claim the account has held since its first connect, and the session guard on DeleteProxy matches because the upsert wrote the new session. Withdrawing that row whenever the post-write re-read errored surrendered an established claim on a transient store error — a window in which any other account could take the address — where the pre-existing code left the row untouched. An inconclusive re-read still refuses the connect, but marks the session disconnected instead of deleting the row; only a conclusive answer that the address is claimed withdraws it. The write-then-re-read argument moves to the doc of the exported ErrClusterAddressUnavailable, where the API needs it, and both helpers point there instead of carrying it twice. Store-backed tests drive the re-read through the real queries — a reconnect keeps its row, the account's own pin is not a competing claim, another account's is — since the whole path now depends on the store excluding the account's own claims. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Sa3DsBDP3VciAi4PPG17L6
2026-09-12 17:59:06 +02:00 · 2026-09-12 13:36:36 +00:00
parent e6c69f674d
commit d80f0ff031
5 changed files with 161 additions and 39 deletions
@@ -92,28 +92,23 @@ 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.
+// confirmClusterAddressClaim re-reads availability once the proxy's row is
+// committed and withdraws the row if the claim is lost; see
+// proxy.ErrClusterAddressUnavailable for why the re-read is what closes the
+// race with a concurrent claim. An inconclusive re-read refuses the connect
+// but only marks the row disconnected: SaveProxy upserts on the proxy ID, so
+// on a reconnect the row is a claim the account already held, and a transient
+// store error must not surrender it.
 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 {
+	if err != nil {
+		if discErr := m.store.DisconnectProxy(ctx, p.ID, p.SessionID); discErr != nil {
+			log.WithContext(ctx).Errorf("failed to mark proxy %s session %s disconnected after an inconclusive claim check on %s: %v",
+				p.ID, p.SessionID, p.ClusterAddress, discErr)
+		}
+		return fmt.Errorf("confirm claim on cluster address %s: %w", p.ClusterAddress, err)
+	}
+	if available {
 		return nil
 	}

@@ -121,9 +116,6 @@ func (m *Manager) confirmClusterAddressClaim(ctx context.Context, p *proxy.Proxy
 		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)
@@ -0,0 +1,91 @@
+package manager
+
+import (
+	"context"
+	"runtime"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/metric/noop"
+
+	agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
+	nbstore "github.com/netbirdio/netbird/management/server/store"
+	nbtypes "github.com/netbirdio/netbird/management/server/types"
+)
+
+// newStoreBackedManager wires the manager to a real sqlite store, for the
+// cases where what matters is how the store's own queries answer the
+// post-write re-read — which the function-field mock cannot say.
+func newStoreBackedManager(t *testing.T) (*Manager, nbstore.Store) {
+	t.Helper()
+	if runtime.GOOS == "windows" {
+		t.Skip("sqlite store not properly supported on Windows yet")
+	}
+	t.Setenv("NETBIRD_STORE_ENGINE", string(nbtypes.SqliteStoreEngine))
+
+	st, cleanUp, err := nbstore.NewTestStoreFromSQL(context.Background(), "", t.TempDir())
+	require.NoError(t, err, "test store setup must succeed")
+	t.Cleanup(cleanUp)
+
+	mgr, err := NewManager(st, noop.NewMeterProvider().Meter("test"))
+	require.NoError(t, err)
+	return mgr, st
+}
+
+// TestConnect_RealStore_ConfirmsOwnClaims drives the post-write re-read
+// through the real queries. Every account-scoped connect now reads its own
+// just-written row back, so the whole path depends on the store excluding the
+// account's own claims: its own proxy row on a reconnect, and its own gateway
+// pin when the account deploys a proxy at the address it pinned first.
+func TestConnect_RealStore_ConfirmsOwnClaims(t *testing.T) {
+	ctx := context.Background()
+	accountID := "account1"
+	const host = "byop.account1.example.com"
+
+	t.Run("a reconnect keeps the proxy's own row", func(t *testing.T) {
+		mgr, st := newStoreBackedManager(t)
+
+		_, err := mgr.Connect(ctx, "proxy-1", "session-1", host, "10.0.0.1", &accountID, nil)
+		require.NoError(t, err, "first connect must succeed")
+		_, err = mgr.Connect(ctx, "proxy-1", "session-2", host, "10.0.0.1", &accountID, nil)
+		require.NoError(t, err, "a reconnect must not be refused by the row it is replacing")
+
+		rows, err := st.GetAllProxies(ctx)
+		require.NoError(t, err)
+		require.Len(t, rows, 1, "a reconnect upserts the same row")
+		assert.Equal(t, "session-2", rows[0].SessionID, "the row must carry the new session")
+		assert.Equal(t, proxy.StatusConnected, rows[0].Status)
+	})
+
+	t.Run("the account's own gateway pin is not a competing claim", func(t *testing.T) {
+		mgr, st := newStoreBackedManager(t)
+		settings := agentNetworkTypes.DefaultSettings(accountID)
+		settings.Domain = host
+		settings.ProxyAddress = host
+		require.NoError(t, st.CreateAgentNetworkSettings(ctx, settings), "seeding the account's own pin must succeed")
+
+		_, err := mgr.Connect(ctx, "proxy-1", "session-1", host, "10.0.0.1", &accountID, nil)
+		require.NoError(t, err, "pin first, deploy the proxy after is the documented order")
+
+		rows, err := st.GetAllProxies(ctx)
+		require.NoError(t, err)
+		assert.Len(t, rows, 1, "the proxy's row must stand next to the account's own pin")
+	})
+
+	t.Run("another account's gateway pin withdraws the row", func(t *testing.T) {
+		mgr, st := newStoreBackedManager(t)
+		settings := agentNetworkTypes.DefaultSettings("account2")
+		settings.Domain = host
+		settings.ProxyAddress = host
+		require.NoError(t, st.CreateAgentNetworkSettings(ctx, settings), "seeding the other account's pin must succeed")
+
+		_, err := mgr.Connect(ctx, "proxy-1", "session-1", host, "10.0.0.1", &accountID, nil)
+		require.ErrorIs(t, err, proxy.ErrClusterAddressUnavailable)
+
+		rows, err := st.GetAllProxies(ctx)
+		require.NoError(t, err)
+		assert.Empty(t, rows, "a withdrawn registration must leave no row behind")
+	})
+}
@@ -458,19 +458,28 @@ func TestConnect_WithdrawsClaimLostDuringRegistration(t *testing.T) {
 	}
 }

-// 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) {
+// TestConnect_KeepsClaimWhenRecheckFails pins what an inconclusive re-read
+// does: the connect is refused with the store's error, not
+// ErrClusterAddressUnavailable, since nothing established that the address is
+// taken — and the row is marked disconnected rather than deleted. SaveProxy
+// upserts on the proxy ID, so on a reconnect that row is the claim the account
+// has held since its first connect; a transient store error must not hand the
+// address to whoever asks next.
+func TestConnect_KeepsClaimWhenRecheckFails(t *testing.T) {
 	accountID := "acc-1"
-	var withdrawn int
+	var disconnected []string
 	s := &mockStore{
 		hasGatewayPinnedByOtherAccountFunc: func(_ context.Context, _, _ string) (bool, error) {
 			return false, errors.New("db unavailable")
 		},
-		deleteProxyFunc: func(_ context.Context, _, _ string) error { withdrawn++; return nil },
+		disconnectProxyFunc: func(_ context.Context, proxyID, sessionID string) error {
+			disconnected = append(disconnected, proxyID+"/"+sessionID)
+			return nil
+		},
+		deleteProxyFunc: func(_ context.Context, proxyID, _ string) error {
+			t.Fatalf("an inconclusive re-read must not withdraw the row, but proxy %s was deleted", proxyID)
+			return nil
+		},
 	}

 	mgr := newTestManager(s)
@@ -478,7 +487,26 @@ func TestConnect_WithdrawsClaimWhenRecheckFails(t *testing.T) {
 	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")
+	assert.Equal(t, []string{"proxy-1/session-1"}, disconnected, "the refused session must not stay marked connected")
+}
+
+// TestConnect_RefusesEvenWhenWithdrawalFails pins that a lost claim is
+// reported as lost whatever happens to the compensating delete: the caller
+// must never be told it holds an address another claim already has, and the
+// stale row is the reaper's problem, not a reason to lie.
+func TestConnect_RefusesEvenWhenWithdrawalFails(t *testing.T) {
+	accountID := "acc-1"
+	s := &mockStore{
+		hasGatewayPinnedByOtherAccountFunc: func(_ context.Context, _, _ string) (bool, error) { return true, nil },
+		deleteProxyFunc: func(_ context.Context, _, _ string) error {
+			return errors.New("delete failed")
+		},
+	}
+
+	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 failed withdrawal must not turn a lost claim into a held one")
+	assert.Nil(t, p)
 }

 // TestConnect_ConfirmedClaimKeepsRow is the common case: nothing landed in the
@@ -10,12 +10,23 @@ 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.
+// ErrClusterAddressUnavailable is returned by Manager.Connect when the cluster
+// address turns out to be claimed by someone else once the proxy's own row is
+// written: a conflicting proxy row, or another account's agent network gateway
+// pinned to the address. The row has been withdrawn by then, and the caller
+// reports the address as taken exactly as if the pre-write check had caught it.
+//
+// Both kinds of claim are made the same way, write then re-read then withdraw,
+// and the re-read is the whole mechanism. Each side's availability check and
+// its write are separate autocommit statements, so two concurrent claimants
+// can each pass their check with neither row committed yet. Because both write
+// before they re-read, of two concurrent claims at least one re-reads after
+// the other has committed and backs off; each statement sees every commit
+// before it on sqlite, postgres and mysql alike. 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. The gateway side of the same protocol is
+// agentnetwork's confirmGatewayClusterOwnership.
 var ErrClusterAddressUnavailable = errors.New("cluster address is not available")

 // Capabilities describes what a proxy can handle, as reported via gRPC.
@@ -6306,7 +6306,7 @@ func (s *SqlStore) DeleteProxy(ctx context.Context, proxyID, sessionID string) e
 		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)
+		log.WithContext(ctx).Debugf("proxy %s session %s: no row deleted (already gone or superseded by a newer session)", proxyID, sessionID)
 	}
 	return nil
 }