netbird/management/internals/modules/reverseproxy/proxy/manager/manager_test.go

package manager

import (
	"context"
	"errors"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.opentelemetry.io/otel/metric/noop"

	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
)

type mockStore struct {
	saveProxyFunc                            func(ctx context.Context, p *proxy.Proxy) error
	disconnectProxyFunc                      func(ctx context.Context, proxyID, sessionID string) error
	updateProxyHeartbeatFunc                 func(ctx context.Context, p *proxy.Proxy) error
	getActiveProxyClusterAddressesFunc       func(ctx context.Context) ([]string, error)
	getActiveProxyClusterAddressesForAccFunc func(ctx context.Context, accountID string) ([]string, error)
	cleanupStaleProxiesFunc                  func(ctx context.Context, d time.Duration) error
	getProxyByAccountIDFunc                  func(ctx context.Context, accountID string) (*proxy.Proxy, error)
	countProxiesByAccountIDFunc              func(ctx context.Context, accountID string) (int64, error)
	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 {
	if m.saveProxyFunc != nil {
		return m.saveProxyFunc(ctx, p)
	}
	return nil
}
func (m *mockStore) DisconnectProxy(ctx context.Context, proxyID, sessionID string) error {
	if m.disconnectProxyFunc != nil {
		return m.disconnectProxyFunc(ctx, proxyID, sessionID)
	}
	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)
	}
	return nil
}
func (m *mockStore) GetActiveProxyClusterAddresses(ctx context.Context) ([]string, error) {
	if m.getActiveProxyClusterAddressesFunc != nil {
		return m.getActiveProxyClusterAddressesFunc(ctx)
	}
	return nil, nil
}
func (m *mockStore) GetActiveProxyClusterAddressesForAccount(ctx context.Context, accountID string) ([]string, error) {
	if m.getActiveProxyClusterAddressesForAccFunc != nil {
		return m.getActiveProxyClusterAddressesForAccFunc(ctx, accountID)
	}
	return nil, nil
}
func (m *mockStore) GetProxyClusters(_ context.Context, _ string) ([]proxy.Cluster, error) {
	return nil, nil
}
func (m *mockStore) CleanupStaleProxies(ctx context.Context, d time.Duration) error {
	if m.cleanupStaleProxiesFunc != nil {
		return m.cleanupStaleProxiesFunc(ctx, d)
	}
	return nil
}
func (m *mockStore) GetProxyByAccountID(ctx context.Context, accountID string) (*proxy.Proxy, error) {
	if m.getProxyByAccountIDFunc != nil {
		return m.getProxyByAccountIDFunc(ctx, accountID)
	}
	return nil, fmt.Errorf("proxy not found for account %s", accountID)
}
func (m *mockStore) CountProxiesByAccountID(ctx context.Context, accountID string) (int64, error) {
	if m.countProxiesByAccountIDFunc != nil {
		return m.countProxiesByAccountIDFunc(ctx, accountID)
	}
	return 0, nil
}
func (m *mockStore) IsClusterAddressConflicting(ctx context.Context, clusterAddress, accountID string) (bool, error) {
	if m.isClusterAddressConflictingFunc != nil {
		return m.isClusterAddressConflictingFunc(ctx, clusterAddress, accountID)
	}
	return false, nil
}
func (m *mockStore) HasGatewayPinnedByOtherAccount(ctx context.Context, host, accountID string) (bool, error) {
	if m.hasGatewayPinnedByOtherAccountFunc != nil {
		return m.hasGatewayPinnedByOtherAccountFunc(ctx, host, accountID)
	}
	return false, nil
}
func (m *mockStore) DeleteAccountCluster(ctx context.Context, clusterAddress, accountID string) error {
	if m.deleteAccountClusterFunc != nil {
		return m.deleteAccountClusterFunc(ctx, clusterAddress, accountID)
	}
	return nil
}
func (m *mockStore) GetClusterSupportsCustomPorts(_ context.Context, _ string) *bool {
	return nil
}
func (m *mockStore) GetClusterRequireSubdomain(_ context.Context, _ string) *bool {
	return nil
}
func (m *mockStore) GetClusterSupportsCrowdSec(_ context.Context, _ string) *bool {
	return nil
}
func (m *mockStore) GetClusterSupportsPrivate(_ context.Context, _ string) *bool {
	return nil
}

func newTestManager(s store) *Manager {
	meter := noop.NewMeterProvider().Meter("test")
	m, err := NewManager(s, meter)
	if err != nil {
		panic(err)
	}
	return m
}

func TestConnect_WithAccountID(t *testing.T) {
	accountID := "acc-123"

	var savedProxy *proxy.Proxy
	s := &mockStore{
		saveProxyFunc: func(_ context.Context, p *proxy.Proxy) error {
			savedProxy = p
			return nil
		},
	}

	mgr := newTestManager(s)
	_, err := mgr.Connect(context.Background(), "proxy-1", "session-1", "cluster.example.com", "10.0.0.1", &accountID, nil)
	require.NoError(t, err)

	require.NotNil(t, savedProxy)
	assert.Equal(t, "proxy-1", savedProxy.ID)
	assert.Equal(t, "session-1", savedProxy.SessionID)
	assert.Equal(t, "cluster.example.com", savedProxy.ClusterAddress)
	assert.Equal(t, "10.0.0.1", savedProxy.IPAddress)
	assert.Equal(t, &accountID, savedProxy.AccountID)
	assert.Equal(t, proxy.StatusConnected, savedProxy.Status)
	assert.NotNil(t, savedProxy.ConnectedAt)
}

func TestConnect_WithoutAccountID(t *testing.T) {
	var savedProxy *proxy.Proxy
	s := &mockStore{
		saveProxyFunc: func(_ context.Context, p *proxy.Proxy) error {
			savedProxy = p
			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)

	require.NotNil(t, savedProxy)
	assert.Nil(t, savedProxy.AccountID)
	assert.Equal(t, proxy.StatusConnected, savedProxy.Status)
}

func TestConnect_StoreError(t *testing.T) {
	s := &mockStore{
		saveProxyFunc: func(_ context.Context, _ *proxy.Proxy) error {
			return errors.New("db error")
		},
	}

	mgr := newTestManager(s)
	_, err := mgr.Connect(context.Background(), "proxy-1", "session-1", "cluster.example.com", "10.0.0.1", nil, nil)
	assert.Error(t, err)
}

func TestIsClusterAddressAvailable(t *testing.T) {
	tests := []struct {
		name        string
		conflicting bool
		storeErr    error
		wantResult  bool
		wantErr     bool
	}{
		{
			name:        "available - no conflict",
			conflicting: false,
			wantResult:  true,
		},
		{
			name:        "not available - conflict exists",
			conflicting: true,
			wantResult:  false,
		},
		{
			name:     "store error",
			storeErr: errors.New("db error"),
			wantErr:  true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &mockStore{
				isClusterAddressConflictingFunc: func(_ context.Context, _, _ string) (bool, error) {
					return tt.conflicting, tt.storeErr
				},
			}

			mgr := newTestManager(s)
			result, err := mgr.IsClusterAddressAvailable(context.Background(), "cluster.example.com", "acc-123")
			if tt.wantErr {
				assert.Error(t, err)
				return
			}
			require.NoError(t, err)
			assert.Equal(t, tt.wantResult, result)
		})
	}
}

func TestCountAccountProxies(t *testing.T) {
	tests := []struct {
		name      string
		count     int64
		storeErr  error
		wantCount int64
		wantErr   bool
	}{
		{
			name:      "no proxies",
			count:     0,
			wantCount: 0,
		},
		{
			name:      "one proxy",
			count:     1,
			wantCount: 1,
		},
		{
			name:     "store error",
			storeErr: errors.New("db error"),
			wantErr:  true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &mockStore{
				countProxiesByAccountIDFunc: func(_ context.Context, _ string) (int64, error) {
					return tt.count, tt.storeErr
				},
			}

			mgr := newTestManager(s)
			count, err := mgr.CountAccountProxies(context.Background(), "acc-123")
			if tt.wantErr {
				assert.Error(t, err)
				return
			}
			require.NoError(t, err)
			assert.Equal(t, tt.wantCount, count)
		})
	}
}

func TestGetAccountProxy(t *testing.T) {
	accountID := "acc-123"

	t.Run("found", func(t *testing.T) {
		expected := &proxy.Proxy{
			ID:             "proxy-1",
			ClusterAddress: "byop.example.com",
			AccountID:      &accountID,
			Status:         proxy.StatusConnected,
		}
		s := &mockStore{
			getProxyByAccountIDFunc: func(_ context.Context, accID string) (*proxy.Proxy, error) {
				assert.Equal(t, accountID, accID)
				return expected, nil
			},
		}

		mgr := newTestManager(s)
		p, err := mgr.GetAccountProxy(context.Background(), accountID)
		require.NoError(t, err)
		assert.Equal(t, expected, p)
	})

	t.Run("not found", func(t *testing.T) {
		s := &mockStore{
			getProxyByAccountIDFunc: func(_ context.Context, _ string) (*proxy.Proxy, error) {
				return nil, errors.New("not found")
			},
		}

		mgr := newTestManager(s)
		_, err := mgr.GetAccountProxy(context.Background(), accountID)
		assert.Error(t, err)
	})
}

func TestDeleteAccountCluster(t *testing.T) {
	t.Run("success", func(t *testing.T) {
		var deletedCluster, deletedAccount string
		s := &mockStore{
			deleteAccountClusterFunc: func(_ context.Context, clusterAddress, accountID string) error {
				deletedCluster = clusterAddress
				deletedAccount = accountID
				return nil
			},
		}

		mgr := newTestManager(s)
		err := mgr.DeleteAccountCluster(context.Background(), "cluster.example.com", "acc-123")
		require.NoError(t, err)
		assert.Equal(t, "cluster.example.com", deletedCluster)
		assert.Equal(t, "acc-123", deletedAccount)
	})

	t.Run("store error", func(t *testing.T) {
		s := &mockStore{
			deleteAccountClusterFunc: func(_ context.Context, _, _ string) error {
				return errors.New("db error")
			},
		}

		mgr := newTestManager(s)
		err := mgr.DeleteAccountCluster(context.Background(), "cluster.example.com", "acc-123")
		assert.Error(t, err)
	})
}

func TestGetActiveClusterAddressesForAccount(t *testing.T) {
	expected := []string{"byop.example.com"}
	s := &mockStore{
		getActiveProxyClusterAddressesForAccFunc: func(_ context.Context, accID string) ([]string, error) {
			assert.Equal(t, "acc-123", accID)
			return expected, nil
		},
	}

	mgr := newTestManager(s)
	result, err := mgr.GetActiveClusterAddressesForAccount(context.Background(), "acc-123")
	require.NoError(t, err)
	assert.Equal(t, expected, result)
}

// TestIsClusterAddressAvailableConsidersGatewayPins pins that a proxy row is
// not the only claim on an address.
//
// An agent network gateway pinned to the address by another account is
// immutable and is served by whichever proxy declares that address, so a proxy
// from a different account taking it strands the pin — the mapping paths never
// hand an account-scoped proxy another account's mappings. Refusing the later
// claimant is what makes the bootstrap-time ownership check hold over time
// rather than only at the instant it runs: without this, an address a gateway
// pinned while no proxy served it could be taken a moment, or a week, later.
func TestIsClusterAddressAvailableConsidersGatewayPins(t *testing.T) {
	ctx := context.Background()

	tests := []struct {
		name        string
		conflicting bool
		pinned      bool
		available   bool
	}{
		{name: "free address", available: true},
		{name: "claimed by a proxy", conflicting: true},
		{name: "pinned by another account's gateway", pinned: true},
		{name: "claimed both ways", conflicting: true, pinned: true},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			st := &mockStore{
				isClusterAddressConflictingFunc: func(context.Context, string, string) (bool, error) {
					return tt.conflicting, nil
				},
				hasGatewayPinnedByOtherAccountFunc: func(context.Context, string, string) (bool, error) {
					return tt.pinned, nil
				},
			}
			m, err := NewManager(st, noop.NewMeterProvider().Meter(""))
			require.NoError(t, err)

			available, err := m.IsClusterAddressAvailable(ctx, "gw.example.com", "account1")
			require.NoError(t, err)
			assert.Equal(t, tt.available, available)
		})
	}
}

// TestIsClusterAddressAvailableSurfacesGatewayPinError pins that a failed pin
// lookup refuses the claim rather than falling through to available: this runs
// on the proxy-connect path, where "could not tell" must not read as "yes".
func TestIsClusterAddressAvailableSurfacesGatewayPinError(t *testing.T) {
	st := &mockStore{
		hasGatewayPinnedByOtherAccountFunc: func(context.Context, string, string) (bool, error) {
			return false, errors.New("db down")
		},
	}
	m, err := NewManager(st, noop.NewMeterProvider().Meter(""))
	require.NoError(t, err)

	available, err := m.IsClusterAddressAvailable(context.Background(), "gw.example.com", "account1")
	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)
}