netbird/proxy/internal/roundtrip/netbird_test.go

package roundtrip

import (
	"context"
	"net/http"
	"net/netip"
	"sync"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"google.golang.org/grpc"

	"github.com/netbirdio/netbird/client/embed"
	"github.com/netbirdio/netbird/proxy/internal/types"
	"github.com/netbirdio/netbird/shared/management/proto"
)

type mockMgmtClient struct{}

func (m *mockMgmtClient) CreateProxyPeer(_ context.Context, _ *proto.CreateProxyPeerRequest, _ ...grpc.CallOption) (*proto.CreateProxyPeerResponse, error) {
	return &proto.CreateProxyPeerResponse{Success: true}, nil
}

type mockStatusNotifier struct {
	mu       sync.Mutex
	statuses []statusCall
}

type statusCall struct {
	accountID types.AccountID
	serviceID types.ServiceID
	connected bool
	// ctx is captured so tests can assert the notifier received a
	// fresh background context rather than an inherited request ctx.
	ctx context.Context
}

func (m *mockStatusNotifier) NotifyStatus(ctx context.Context, accountID types.AccountID, serviceID types.ServiceID, connected bool) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.statuses = append(m.statuses, statusCall{accountID, serviceID, connected, ctx})
	return nil
}

func (m *mockStatusNotifier) calls() []statusCall {
	m.mu.Lock()
	defer m.mu.Unlock()
	return append([]statusCall{}, m.statuses...)
}

// mockNetBird creates a NetBird instance for testing without actually connecting.
// It uses an invalid management URL to prevent real connections.
func mockNetBird() *NetBird {
	return NewNetBird(context.Background(), "test-proxy", "invalid.test", ClientConfig{
		MgmtAddr:     "http://invalid.test:9999",
		WGPort:       0,
		PreSharedKey: "",
	}, nil, nil, &mockMgmtClient{})
}

func TestNetBird_AddPeer_CreatesClientForNewAccount(t *testing.T) {
	nb := mockNetBird()
	accountID := types.AccountID("account-1")

	// Initially no client exists.
	assert.False(t, nb.HasClient(accountID), "should not have client before AddPeer")
	assert.Equal(t, 0, nb.ServiceCount(accountID), "service count should be 0")

	// Add first service - this should create a new client.
	err := nb.AddPeer(context.Background(), accountID, "domain1.test", "setup-key-1", types.ServiceID("proxy-1"))
	require.NoError(t, err)

	assert.True(t, nb.HasClient(accountID), "should have client after AddPeer")
	assert.Equal(t, 1, nb.ServiceCount(accountID), "service count should be 1")
}

func TestNetBird_AddPeer_ReuseClientForSameAccount(t *testing.T) {
	nb := mockNetBird()
	accountID := types.AccountID("account-1")

	// Add first service.
	err := nb.AddPeer(context.Background(), accountID, "domain1.test", "setup-key-1", types.ServiceID("proxy-1"))
	require.NoError(t, err)
	assert.Equal(t, 1, nb.ServiceCount(accountID))

	// Add second service for the same account - should reuse existing client.
	err = nb.AddPeer(context.Background(), accountID, "domain2.test", "setup-key-1", types.ServiceID("proxy-2"))
	require.NoError(t, err)
	assert.Equal(t, 2, nb.ServiceCount(accountID), "service count should be 2 after adding second service")

	// Add third service.
	err = nb.AddPeer(context.Background(), accountID, "domain3.test", "setup-key-1", types.ServiceID("proxy-3"))
	require.NoError(t, err)
	assert.Equal(t, 3, nb.ServiceCount(accountID), "service count should be 3 after adding third service")

	// Still only one client.
	assert.True(t, nb.HasClient(accountID))
}

func TestNetBird_AddPeer_SeparateClientsForDifferentAccounts(t *testing.T) {
	nb := mockNetBird()
	account1 := types.AccountID("account-1")
	account2 := types.AccountID("account-2")

	// Add service for account 1.
	err := nb.AddPeer(context.Background(), account1, "domain1.test", "setup-key-1", types.ServiceID("proxy-1"))
	require.NoError(t, err)

	// Add service for account 2.
	err = nb.AddPeer(context.Background(), account2, "domain2.test", "setup-key-2", types.ServiceID("proxy-2"))
	require.NoError(t, err)

	// Both accounts should have their own clients.
	assert.True(t, nb.HasClient(account1), "account1 should have client")
	assert.True(t, nb.HasClient(account2), "account2 should have client")
	assert.Equal(t, 1, nb.ServiceCount(account1), "account1 service count should be 1")
	assert.Equal(t, 1, nb.ServiceCount(account2), "account2 service count should be 1")
}

func TestNetBird_RemovePeer_KeepsClientWhenServicesRemain(t *testing.T) {
	nb := mockNetBird()
	accountID := types.AccountID("account-1")

	// Add multiple services.
	err := nb.AddPeer(context.Background(), accountID, "domain1.test", "setup-key-1", types.ServiceID("proxy-1"))
	require.NoError(t, err)
	err = nb.AddPeer(context.Background(), accountID, "domain2.test", "setup-key-1", types.ServiceID("proxy-2"))
	require.NoError(t, err)
	err = nb.AddPeer(context.Background(), accountID, "domain3.test", "setup-key-1", types.ServiceID("proxy-3"))
	require.NoError(t, err)
	assert.Equal(t, 3, nb.ServiceCount(accountID))

	// Remove one service - client should remain.
	err = nb.RemovePeer(context.Background(), accountID, "domain1.test")
	require.NoError(t, err)
	assert.True(t, nb.HasClient(accountID), "client should remain after removing one service")
	assert.Equal(t, 2, nb.ServiceCount(accountID), "service count should be 2")

	// Remove another service - client should still remain.
	err = nb.RemovePeer(context.Background(), accountID, "domain2.test")
	require.NoError(t, err)
	assert.True(t, nb.HasClient(accountID), "client should remain after removing second service")
	assert.Equal(t, 1, nb.ServiceCount(accountID), "service count should be 1")
}

func TestNetBird_RemovePeer_RemovesClientWhenLastServiceRemoved(t *testing.T) {
	nb := mockNetBird()
	accountID := types.AccountID("account-1")

	// Add single service.
	err := nb.AddPeer(context.Background(), accountID, "domain1.test", "setup-key-1", types.ServiceID("proxy-1"))
	require.NoError(t, err)
	assert.True(t, nb.HasClient(accountID))

	// Remove the only service - client should be removed.
	_ = nb.RemovePeer(context.Background(), accountID, "domain1.test")

	// After removing all services, client should be gone.
	assert.False(t, nb.HasClient(accountID), "client should be removed after removing last service")
	assert.Equal(t, 0, nb.ServiceCount(accountID), "service count should be 0")
}

func TestNetBird_RemovePeer_NonExistentAccountIsNoop(t *testing.T) {
	nb := mockNetBird()
	accountID := types.AccountID("nonexistent-account")

	// Removing from non-existent account should not error.
	err := nb.RemovePeer(context.Background(), accountID, "domain1.test")
	assert.NoError(t, err, "removing from non-existent account should not error")
}

func TestNetBird_RemovePeer_NonExistentServiceIsNoop(t *testing.T) {
	nb := mockNetBird()
	accountID := types.AccountID("account-1")

	// Add one service.
	err := nb.AddPeer(context.Background(), accountID, "domain1.test", "setup-key-1", types.ServiceID("proxy-1"))
	require.NoError(t, err)

	// Remove non-existent service - should not affect existing service.
	err = nb.RemovePeer(context.Background(), accountID, "nonexistent.test")
	require.NoError(t, err)

	// Original service should still be registered.
	assert.True(t, nb.HasClient(accountID))
	assert.Equal(t, 1, nb.ServiceCount(accountID), "original service should remain")
}

func TestWithAccountID_AndAccountIDFromContext(t *testing.T) {
	ctx := context.Background()
	accountID := types.AccountID("test-account")

	// Initially no account ID in context.
	retrieved := AccountIDFromContext(ctx)
	assert.True(t, retrieved == "", "should be empty when not set")

	// Add account ID to context.
	ctx = WithAccountID(ctx, accountID)
	retrieved = AccountIDFromContext(ctx)
	assert.Equal(t, accountID, retrieved, "should retrieve the same account ID")
}

func TestAccountIDFromContext_ReturnsEmptyForWrongType(t *testing.T) {
	// Create context with wrong type for account ID key.
	ctx := context.WithValue(context.Background(), accountIDContextKey{}, "wrong-type-string")

	retrieved := AccountIDFromContext(ctx)
	assert.True(t, retrieved == "", "should return empty for wrong type")
}

func TestNetBird_StopAll_StopsAllClients(t *testing.T) {
	nb := mockNetBird()
	account1 := types.AccountID("account-1")
	account2 := types.AccountID("account-2")
	account3 := types.AccountID("account-3")

	// Add services for multiple accounts.
	err := nb.AddPeer(context.Background(), account1, "domain1.test", "key-1", types.ServiceID("proxy-1"))
	require.NoError(t, err)
	err = nb.AddPeer(context.Background(), account2, "domain2.test", "key-2", types.ServiceID("proxy-2"))
	require.NoError(t, err)
	err = nb.AddPeer(context.Background(), account3, "domain3.test", "key-3", types.ServiceID("proxy-3"))
	require.NoError(t, err)

	assert.Equal(t, 3, nb.ClientCount(), "should have 3 clients")

	// Stop all clients.
	_ = nb.StopAll(context.Background())

	assert.Equal(t, 0, nb.ClientCount(), "should have 0 clients after StopAll")
	assert.False(t, nb.HasClient(account1), "account1 should not have client")
	assert.False(t, nb.HasClient(account2), "account2 should not have client")
	assert.False(t, nb.HasClient(account3), "account3 should not have client")
}

func TestNetBird_ClientCount(t *testing.T) {
	nb := mockNetBird()

	assert.Equal(t, 0, nb.ClientCount(), "should start with 0 clients")

	// Add clients for different accounts.
	err := nb.AddPeer(context.Background(), types.AccountID("account-1"), "domain1.test", "key-1", types.ServiceID("proxy-1"))
	require.NoError(t, err)
	assert.Equal(t, 1, nb.ClientCount())

	err = nb.AddPeer(context.Background(), types.AccountID("account-2"), "domain2.test", "key-2", types.ServiceID("proxy-2"))
	require.NoError(t, err)
	assert.Equal(t, 2, nb.ClientCount())

	// Adding service to existing account should not increase count.
	err = nb.AddPeer(context.Background(), types.AccountID("account-1"), "domain1b.test", "key-1", types.ServiceID("proxy-1b"))
	require.NoError(t, err)
	assert.Equal(t, 2, nb.ClientCount(), "adding service to existing account should not increase client count")
}

func TestNetBird_RoundTrip_RequiresAccountIDInContext(t *testing.T) {
	nb := mockNetBird()

	// Create a request without account ID in context.
	req, err := http.NewRequest("GET", "http://example.com/", nil)
	require.NoError(t, err)

	// RoundTrip should fail because no account ID in context.
	_, err = nb.RoundTrip(req) //nolint:bodyclose
	require.ErrorIs(t, err, ErrNoAccountID)
}

func TestNetBird_RoundTrip_RequiresExistingClient(t *testing.T) {
	nb := mockNetBird()
	accountID := types.AccountID("nonexistent-account")

	// Create a request with account ID but no client exists.
	req, err := http.NewRequest("GET", "http://example.com/", nil)
	require.NoError(t, err)
	req = req.WithContext(WithAccountID(req.Context(), accountID))

	// RoundTrip should fail because no client for this account.
	_, err = nb.RoundTrip(req) //nolint:bodyclose // Error case, no response body
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "no peer connection found for account")
}

func TestNetBird_AddPeer_ExistingStartedClient_NotifiesStatus(t *testing.T) {
	notifier := &mockStatusNotifier{}
	nb := NewNetBird(context.Background(), "test-proxy", "invalid.test", ClientConfig{
		MgmtAddr:     "http://invalid.test:9999",
		WGPort:       0,
		PreSharedKey: "",
	}, nil, notifier, &mockMgmtClient{})
	accountID := types.AccountID("account-1")

	// Add first service — creates a new client entry.
	err := nb.AddPeer(context.Background(), accountID, "domain1.test", "key-1", types.ServiceID("svc-1"))
	require.NoError(t, err)

	// Manually mark client as started to simulate background startup completing.
	nb.clientsMux.Lock()
	nb.clients[accountID].started = true
	nb.clientsMux.Unlock()

	// Add second service with an already-cancelled caller context —
	// should notify immediately (client is started) AND the notification
	// must not inherit the cancelled ctx.
	cancelledCtx, cancel := context.WithCancel(context.Background())
	cancel()
	err = nb.AddPeer(cancelledCtx, accountID, "domain2.test", "key-1", types.ServiceID("svc-2"))
	require.NoError(t, err)

	calls := notifier.calls()
	require.Len(t, calls, 1)
	assert.Equal(t, accountID, calls[0].accountID)
	assert.Equal(t, types.ServiceID("svc-2"), calls[0].serviceID)
	assert.True(t, calls[0].connected)
	require.NotNil(t, calls[0].ctx, "NotifyStatus must receive a context")
	require.NoError(t, calls[0].ctx.Err(),
		"already-started NotifyStatus must use a background ctx, not the cancelled caller ctx")
}

// TestNetBird_IdentityForIP_UnknownAccountReturnsFalse confirms that the
// public lookup short-circuits when no client has been registered for
// the queried account. The auth middleware uses ok=false as a fast deny.
func TestNetBird_IdentityForIP_UnknownAccountReturnsFalse(t *testing.T) {
	nb := mockNetBird()
	_, _, ok := nb.IdentityForIP("acct-missing", netip.MustParseAddr("100.64.0.10"))
	assert.False(t, ok, "unknown account must yield ok=false")
}

// TestClientEntry_IdentityForIP_NilClientGuard ensures the receiver
// methods stay safe when called on partially-initialized state, which
// can happen briefly during AddPeer setup or test fixtures.
func TestClientEntry_IdentityForIP_NilClientGuard(t *testing.T) {
	var e *clientEntry
	_, _, ok := e.IdentityForIP(netip.MustParseAddr("100.64.0.10"))
	assert.False(t, ok, "nil clientEntry must yield ok=false")

	e = &clientEntry{}
	_, _, ok = e.IdentityForIP(netip.MustParseAddr("100.64.0.10"))
	assert.False(t, ok, "clientEntry with nil embed.Client must yield ok=false")
}

// TestClientEntry_IdentityForIP_InvalidIPReturnsFalse covers the input
// guard so callers don't have to repeat the check.
func TestClientEntry_IdentityForIP_InvalidIPReturnsFalse(t *testing.T) {
	e := &clientEntry{}
	_, _, ok := e.IdentityForIP(netip.Addr{})
	assert.False(t, ok, "invalid IP must yield ok=false")
}

func TestNetBird_RemovePeer_NotifiesDisconnection(t *testing.T) {
	notifier := &mockStatusNotifier{}
	nb := NewNetBird(context.Background(), "test-proxy", "invalid.test", ClientConfig{
		MgmtAddr:     "http://invalid.test:9999",
		WGPort:       0,
		PreSharedKey: "",
	}, nil, notifier, &mockMgmtClient{})
	accountID := types.AccountID("account-1")

	err := nb.AddPeer(context.Background(), accountID, "domain1.test", "key-1", types.ServiceID("svc-1"))
	require.NoError(t, err)
	err = nb.AddPeer(context.Background(), accountID, "domain2.test", "key-1", types.ServiceID("svc-2"))
	require.NoError(t, err)

	// Remove one service — client stays, but disconnection notification fires.
	err = nb.RemovePeer(context.Background(), accountID, "domain1.test")
	require.NoError(t, err)
	assert.True(t, nb.HasClient(accountID))

	calls := notifier.calls()
	require.Len(t, calls, 1)
	assert.Equal(t, types.ServiceID("svc-1"), calls[0].serviceID)
	assert.False(t, calls[0].connected)
}

// TestNotifyClientReady_UsesBackgroundCtx pins the contract that the
// post-Start hooks (readyHandler + statusNotifier.NotifyStatus) run on
// a fresh context.Background() rather than inheriting the AddPeer
// caller's request- or stream-scoped ctx. Without this, a cancelled
// caller ctx could abort the inbound listener bring-up or cause the
// management status notification to fail spuriously and leave the
// account in a half-connected state.
func TestNotifyClientReady_UsesBackgroundCtx(t *testing.T) {
	notifier := &mockStatusNotifier{}
	nb := NewNetBird(context.Background(), "test-proxy", "invalid.test", ClientConfig{
		MgmtAddr: "http://invalid.test:9999",
	}, nil, notifier, &mockMgmtClient{})

	accountID := types.AccountID("acct-async")
	// Pre-populate a client entry so notifyClientReady has something
	// to mark started + something to enumerate for NotifyStatus.
	nb.clientsMux.Lock()
	nb.clients[accountID] = &clientEntry{
		services: map[ServiceKey]serviceInfo{
			DomainServiceKey("svc.example"): {serviceID: types.ServiceID("svc-1")},
		},
	}
	nb.clientsMux.Unlock()

	var capturedReadyCtx context.Context
	nb.SetClientLifecycle(
		func(ctx context.Context, _ types.AccountID, _ *embed.Client) any {
			capturedReadyCtx = ctx
			return nil
		},
		nil,
	)

	// Drive the post-Start path directly; a real client.Start would
	// need a working management URL.
	nb.notifyClientReady(accountID, nil)

	require.NotNil(t, capturedReadyCtx, "readyHandler must have been invoked")
	require.NoError(t, capturedReadyCtx.Err(),
		"readyHandler must receive a background context, not an inherited cancelled one")
	deadline, ok := capturedReadyCtx.Deadline()
	assert.False(t, ok, "readyHandler ctx must have no deadline (background); got %v", deadline)

	calls := notifier.calls()
	require.Len(t, calls, 1, "NotifyStatus must be invoked once per registered service")
	require.NotNil(t, calls[0].ctx, "NotifyStatus must receive a context")
	require.NoError(t, calls[0].ctx.Err(),
		"NotifyStatus must receive a background context, not an inherited cancelled one")
}