Merge branch 'main' into poc/certificate-posture

This commit is contained in:
pascal
2026-09-16 23:31:07 +02:00
29 changed files with 3033 additions and 79 deletions
@@ -9,6 +9,7 @@ import (
"runtime"
"strings"
"testing"
"time"
"go.uber.org/mock/gomock"
"github.com/gorilla/mux"
@@ -17,6 +18,7 @@ import (
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
"github.com/netbirdio/netbird/management/server/account"
nbcontext "github.com/netbirdio/netbird/management/server/context"
"github.com/netbirdio/netbird/management/server/permissions"
@@ -29,6 +31,9 @@ import (
const (
testAccountID = "acc-1"
testUserID = "user-bob"
// testClusterAddress is the shared proxy cluster the settings tests pin
// their gateway to; the fixture seeds a connected private-capable proxy for it.
testClusterAddress = "eu.proxy.netbird.io"
)
// agentNetworkHandlerFixture builds a real agentnetwork.Manager with
@@ -75,6 +80,12 @@ func newAgentNetworkHandlerFixture(t *testing.T) *agentNetworkHandlerFixture {
manager := agentnetwork.NewManager(st, perms, accounts, nil)
h := &handler{manager: manager}
// The labeled bootstrap validates its proxy_address against the live
// clusters, so seed the shared cluster these tests pin to as a real,
// private-capable one — the wire-shape assertions then run through the
// validated path rather than the "nothing connected yet" carve-out.
seedSharedPrivateCluster(t, st, testClusterAddress)
router := mux.NewRouter()
router.HandleFunc("/agent-network/providers", h.createProvider).Methods("POST")
router.HandleFunc("/agent-network/providers/{providerId}", h.getProvider).Methods("GET")
@@ -268,3 +279,21 @@ func TestConsumptionHandler_PopulatedAccountListsRows(t *testing.T) {
assert.Equal(t, groupRow.WindowStartUtc, userRow.WindowStartUtc,
"rows recorded in the same window must share the aligned window_start_utc")
}
// seedSharedPrivateCluster registers a connected, NetBird-operated proxy
// with private capabilities (the `private` capability) so
// clusterAddr is a cluster any account may pin its agent-network gateway to.
func seedSharedPrivateCluster(t *testing.T, st store.Store, clusterAddr string) {
t.Helper()
private := true
now := time.Now().UTC()
require.NoError(t, st.SaveProxy(context.Background(), &rpproxy.Proxy{
ID: "shared-proxy-" + clusterAddr,
SessionID: "shared-session",
ClusterAddress: clusterAddr,
LastSeen: now,
ConnectedAt: &now,
Status: rpproxy.StatusConnected,
Capabilities: rpproxy.Capabilities{Private: &private},
}), "seeding the shared proxy cluster must succeed")
}
@@ -1036,6 +1036,18 @@ func (m *managerImpl) bootstrapSelfAddressed(ctx context.Context, settings *type
if err != nil {
return status.Errorf(status.InvalidArgument, "invalid endpoint: %s", err)
}
if err := m.requireHostNotForeign(ctx, settings.AccountID, hostname); err != nil {
return err
}
// Another account's labeled pin beneath this hostname makes it their
// cluster: a proxy serving them there would never serve this endpoint.
// The domain unique index already arbitrates two endpoints on one name.
if err := m.requireNotClaimedByOtherAccount(ctx, settings.AccountID, hostname, m.store.HasGatewayClusterPinnedByOtherAccount); err != nil {
return err
}
if err := m.validateGatewayCluster(ctx, settings.AccountID, hostname); err != nil {
return err
}
settings.Domain = hostname
settings.ProxyAddress = hostname
@@ -1054,6 +1066,99 @@ func (m *managerImpl) bootstrapSelfAddressed(ctx context.Context, settings *type
return nil
}
// validateGatewayCluster rejects a bootstrap pinned to a cluster that cannot
// serve the account's gateway — a labeled endpoint beneath the cluster and a
// self-addressed one on the very address a proxy declares alike, since the
// service behind either is the same private one.
//
// The synthesised gateway service is unconditionally private
// (buildAccountService): agents reach it over the WireGuard tunnel and are
// authorised by ValidateTunnelPeer against the policies' source groups, and
// its single target is the cluster itself with DirectUpstream. Only a cluster
// with private capabilities can serve that. Management reports it per cluster
// as the `private` capability, the same flag the dashboard renders as
// supports_private when it gates NetBird-only services.
//
// Without this check the bootstrap happily pins to any cluster the caller
// names, including one without private capabilities — and the endpoint it
// allocates is immutable, so the account is left with a dead gateway that only
// a DeleteSettings/re-bootstrap can undo.
//
// Whether management knows the cluster is decided on the proxy rows
// themselves, never on how fresh their heartbeats are: a cluster's rows
// outlive its proxies' liveness (only the stale-proxy reaper removes them), so
// a cluster that exists stays judged as one. Judging on liveness instead would
// make the same centralised cluster pass or fail depending on whether its
// proxies happened to have heartbeated in the last couple of minutes.
//
// The single opening left is a cluster management holds no proxy row for at
// all: pinning ahead of a proxy's first connection is a legitimate order — the
// dedicated path claims an address the same way, before any proxy declares it.
func (m *managerImpl) validateGatewayCluster(ctx context.Context, accountID, clusterAddr string) error {
declared, err := m.accountClusterSpellings(ctx, accountID, clusterAddr)
if err != nil {
return err
}
if len(declared) == 0 {
// No proxy has ever declared this address: an address-first pin.
return nil
}
// A cluster management knows has to prove it can serve the gateway, and
// only a live proxy reporting the capability proves that. Both an explicit false and an
// unreported capability (nothing live in the cluster, or proxies predating
// capability reporting) fail here: unusable and unproven are the same
// answer for a decision that cannot be revisited later.
//
// The capability is read per declared spelling and taken as any-true, the
// same way it aggregates over a cluster's proxies: the store matches
// cluster_address exactly, so a host two proxies spelled differently must
// not come back unproven just because it was asked about under one of them.
for _, address := range declared {
if private := m.store.GetClusterSupportsPrivate(ctx, address); private != nil && *private {
return nil
}
}
return status.Errorf(status.InvalidArgument,
"proxy cluster %s has no private capabilities: the agent network gateway requires a reverse proxy cluster "+
"with private capabilities", clusterAddr)
}
// accountClusterSpellings returns every proxy cluster address in the account's
// view — its own (BYOP) clusters plus the shared ones — that names the same
// host as clusterAddr. Empty means management holds no proxy row for that host
// in this account's view.
//
// A proxy declares its cluster address as the operator spelled it, so identity
// is compared on the normalised form rather than byte-equal — an in-memory pass
// over the account's clusters, not a query. What comes back is the stored
// spelling, because the capability lookup matches cluster_address exactly and
// would silently find nothing under a spelling the store never held. The
// cluster listing is not gated on heartbeats, so this answer does not change
// while a cluster's proxies are merely offline.
func (m *managerImpl) accountClusterSpellings(ctx context.Context, accountID, clusterAddr string) ([]string, error) {
clusters, err := m.store.GetProxyClusters(ctx, accountID)
if err != nil {
return nil, fmt.Errorf("list proxy clusters: %w", err)
}
var spellings []string
for _, cluster := range clusters {
normalized, err := types.NormalizeHostname(cluster.Address)
if err != nil {
// An address declared in a shape we cannot normalise is not one an
// endpoint can be allocated beneath.
log.WithContext(ctx).Debugf("skipping unusable proxy cluster address %q: %s", cluster.Address, err)
continue
}
if normalized == clusterAddr {
spellings = append(spellings, cluster.Address)
}
}
return spellings, nil
}
// bootstrapLabeled allocates a labeled endpoint one label beneath the given
// cluster address: Domain = <label>.<proxyAddress>, served by whichever proxy
// declares the parent. Labels are adjective-noun tuples; a candidate is
@@ -1065,6 +1170,20 @@ func (m *managerImpl) bootstrapLabeled(ctx context.Context, settings *types.Sett
if err != nil {
return status.Errorf(status.InvalidArgument, "invalid proxy_address: %s", err)
}
if err := m.requireHostNotForeign(ctx, settings.AccountID, parent); err != nil {
return err
}
// Another account's endpoint at this exact hostname means the proxy that
// declares it is theirs, so nothing would serve a label beneath it. Other
// accounts' labeled pins under the same cluster are not asked about: a
// shared cluster carries many of them by design.
if err := m.requireNotClaimedByOtherAccount(ctx, settings.AccountID, parent, m.store.HasGatewayEndpointByOtherAccount); err != nil {
return err
}
if err := m.validateGatewayCluster(ctx, settings.AccountID, parent); err != nil {
return err
}
for attempt := 1; attempt <= maxDomainAllocationAttempts; attempt++ {
label := labelgen.PickTuple()
@@ -1111,6 +1230,41 @@ func (m *managerImpl) bootstrapLabeled(ctx context.Context, settings *types.Sett
return fmt.Errorf("allocate agent network endpoint for account %s: %d attempts exhausted", settings.AccountID, maxDomainAllocationAttempts)
}
// requireHostNotForeign refuses to pin the account's gateway onto a host that
// another account's proxy declares. The pin's proxy_address is what selects
// the proxy that serves the endpoint, and an account-scoped proxy only ever
// receives its own account's mappings, so such a pin could never be served —
// and the endpoint it assigns is immutable. Shared proxies are not foreign, and
// a host no proxy has declared stays pinnable: claiming the address before the
// proxy's first connection is the documented order.
func (m *managerImpl) requireHostNotForeign(ctx context.Context, accountID, host string) error {
foreign, err := m.store.HasForeignAccountProxyAtHost(ctx, host, accountID)
if err != nil {
return fmt.Errorf("check proxy host ownership: %w", err)
}
if foreign {
return errHostNotAvailable(host)
}
return nil
}
// requireNotClaimedByOtherAccount refuses the pin when another account's
// gateway settings already claim the host in the shape claimed answers for.
func (m *managerImpl) requireNotClaimedByOtherAccount(ctx context.Context, accountID, host string, claimed func(context.Context, string, string) (bool, error)) error {
taken, err := claimed(ctx, host, accountID)
if err != nil {
return fmt.Errorf("check agent network gateway claims at host: %w", err)
}
if taken {
return errHostNotAvailable(host)
}
return nil
}
func errHostNotAvailable(host string) error {
return status.Errorf(status.InvalidArgument, "proxy cluster %s is not available to this account", host)
}
// isUniqueConstraintError reports whether err is a database unique-constraint
// violation, matched on the driver message because CreateAgentNetworkSettings
// deliberately returns the driver error unwrapped.
@@ -5,12 +5,14 @@ import (
"runtime"
"strings"
"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"
"github.com/netbirdio/netbird/management/server/account"
"github.com/netbirdio/netbird/management/server/permissions"
"github.com/netbirdio/netbird/management/server/permissions/modules"
@@ -70,6 +72,57 @@ func (f *bootstrapFixture) createSettings(ctx context.Context, accountID, userID
return f.manager.CreateSettings(ctx, userID, types.DefaultSettings(accountID), proxyAddress, endpoint)
}
func ptrTo[T any](v T) *T { return &v }
// seedProxy registers a proxy in clusterAddr, heartbeating now, so the labeled
// bootstrap path has a real cluster to validate against. accountID empty makes
// it a shared (NetBird-operated) cluster; private mirrors the capability an
// proxy with private capabilities reports, nil an unreported one.
func (f *bootstrapFixture) seedProxy(t *testing.T, proxyID, accountID, clusterAddr string, private *bool) {
t.Helper()
f.seedProxyAt(t, proxyID, accountID, clusterAddr, private, time.Now().UTC())
}
// seedProxyAt is seedProxy with an explicit last-seen, for cases that need a
// proxy whose heartbeat has aged past the active window while its row (and so
// its cluster) is still on record.
func (f *bootstrapFixture) seedProxyAt(t *testing.T, proxyID, accountID, clusterAddr string, private *bool, lastSeen time.Time) {
t.Helper()
p := &proxy.Proxy{
ID: proxyID,
ClusterAddress: clusterAddr,
Status: proxy.StatusConnected,
LastSeen: lastSeen,
Capabilities: proxy.Capabilities{Private: private},
}
if accountID != "" {
p.AccountID = &accountID
}
require.NoError(t, f.store.SaveProxy(context.Background(), p), "seeding a proxy must succeed")
}
// seedPrivateCluster is the common case: a shared cluster with a connected
// proxy that has private capabilities, which is what a bootstrap requires.
func (f *bootstrapFixture) seedPrivateCluster(t *testing.T, clusterAddr string) {
t.Helper()
f.seedProxy(t, "proxy-"+clusterAddr, "", clusterAddr, ptrTo(true))
}
// requireForeignClusterRefusal asserts the refusal a pin onto another
// account's host gets, and that it left no row behind.
func (f *bootstrapFixture) requireForeignClusterRefusal(t *testing.T, err error, accountID string) {
t.Helper()
require.Error(t, err, "another account's host must be refused")
var sErr *status.Error
require.ErrorAs(t, err, &sErr)
assert.Equal(t, status.InvalidArgument, sErr.Type(), "rejection must be a validation error")
assert.Contains(t, err.Error(), "not available to this account",
"the error must say the host is not the account's to use")
_, err = f.store.GetAgentNetworkSettings(context.Background(), store.LockingStrengthNone, accountID)
assert.Error(t, err, "no row may be left behind by a rejected bootstrap")
}
// TestCreateSettingsRequiresPermission pins the gate: bootstrap assigns the
// account's immutable endpoint, a settings write requiring the settings
// Create permission — and a denial leaves no row behind.
@@ -94,6 +147,7 @@ func TestCreateSettingsRequiresPermission(t *testing.T) {
func TestCreateSettingsLabeled(t *testing.T) {
ctx := context.Background()
f := newBootstrapFixture(t)
f.seedPrivateCluster(t, "cluster1.example.com")
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
created, err := f.createSettings(ctx, "account1", "user1", "Cluster1.Example.com", "")
@@ -167,6 +221,7 @@ func TestCreateSettingsIdentityFieldValidation(t *testing.T) {
func TestCreateSettingsConflictsOnSecondBootstrap(t *testing.T) {
ctx := context.Background()
f := newBootstrapFixture(t)
f.seedPrivateCluster(t, "cluster1.example.com")
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
first, err := f.createSettings(ctx, "account1", "user1", "cluster1.example.com", "")
@@ -230,3 +285,279 @@ func TestCreateProviderHasNoSettingsSideEffects(t *testing.T) {
_, err = f.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "account1")
assert.Error(t, err, "provider create must not conjure a settings row")
}
// TestCreateSettingsRejectsOfflineCluster is the guard against deciding on
// heartbeat freshness. A centralised cluster is refused while its proxies are
// live; the same cluster must stay refused once they stop heartbeating, which
// takes only a couple of minutes (proxyActiveThreshold). Judging on liveness
// would turn "wait for the proxy to go quiet" into a way to pin the account's
// immutable endpoint to a cluster that can never serve it.
func TestCreateSettingsRejectsOfflineCluster(t *testing.T) {
ctx := context.Background()
notPrivate := false
cases := map[string]*bool{
"centralised proxy gone quiet": &notPrivate,
// A cluster that could serve the gateway still has to have something
// live in it to prove so at bootstrap: refusing is the safe direction
// (reconnect the proxy and retry) where accepting is permanent.
"private proxy gone quiet": ptrTo(true),
}
for name, private := range cases {
t.Run(name, func(t *testing.T) {
f := newBootstrapFixture(t)
f.seedProxyAt(t, "proxy1", "", "offline.example.com", private,
time.Now().UTC().Add(-time.Hour))
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
_, err := f.createSettings(ctx, "account1", "user1", "offline.example.com", "")
require.Error(t, err, "a known cluster with nothing live in it must be rejected")
var sErr *status.Error
require.ErrorAs(t, err, &sErr)
assert.Equal(t, status.InvalidArgument, sErr.Type(), "rejection must be a validation error")
assert.Contains(t, err.Error(), "private capabilities",
"the error must say private capabilities are what is missing")
_, err = f.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "account1")
assert.Error(t, err, "no row may be left behind by a rejected bootstrap")
})
}
}
// TestCreateSettingsRequiresPrivateCluster pins the capability gate: the
// synthesised gateway service is always private, so a live cluster whose
// proxies lack private capabilities cannot serve it and must not
// become the account's immutable endpoint.
func TestCreateSettingsRequiresPrivateCluster(t *testing.T) {
ctx := context.Background()
f := newBootstrapFixture(t)
notPrivate := false
f.seedProxy(t, "proxy1", "", "central.example.com", &notPrivate)
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
_, err := f.createSettings(ctx, "account1", "user1", "central.example.com", "")
require.Error(t, err, "a cluster without private capabilities must be rejected")
var sErr *status.Error
require.ErrorAs(t, err, &sErr)
assert.Equal(t, status.InvalidArgument, sErr.Type(), "rejection must be a validation error")
assert.Contains(t, err.Error(), "private capabilities", "the error must name what the cluster is missing")
_, err = f.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "account1")
assert.Error(t, err, "no row may be left behind by a rejected bootstrap")
}
// TestCreateSettingsAcceptsOwnPrivateCluster pins the BYOP happy path: the
// account's own cluster with a connected private-capable proxy is a valid pin.
func TestCreateSettingsAcceptsOwnPrivateCluster(t *testing.T) {
ctx := context.Background()
f := newBootstrapFixture(t)
f.seedProxy(t, "proxy1", "account1", "byop.account1.example.com", ptrTo(true))
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
created, err := f.createSettings(ctx, "account1", "user1", "byop.account1.example.com", "")
require.NoError(t, err, "the account's own private cluster must be accepted")
assert.Equal(t, "byop.account1.example.com", created.ProxyAddress)
}
// TestCreateSettingsMatchesClusterCasing pins that a cluster spelled with
// capitals in the store is still recognised as the same cluster the normalised
// proxy_address names, in both directions: a private cluster is accepted and a
// centralised one is refused, whatever the casing. The comparison is in memory
// over the account's cluster list; the capability lookup is still asked under
// the spelling the store actually holds, which is what an exact match needs.
func TestCreateSettingsMatchesClusterCasing(t *testing.T) {
ctx := context.Background()
t.Run("own private cluster is found", func(t *testing.T) {
f := newBootstrapFixture(t)
f.seedProxy(t, "proxy1", "", "EU.Proxy.Example.com", ptrTo(true))
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
created, err := f.createSettings(ctx, "account1", "user1", "eu.proxy.example.com", "")
require.NoError(t, err, "a private cluster declared with capitals must still be accepted")
assert.Equal(t, "eu.proxy.example.com", created.ProxyAddress)
})
t.Run("non-private cluster is still refused", func(t *testing.T) {
f := newBootstrapFixture(t)
f.seedProxy(t, "proxy1", "", "Central.Example.com", ptrTo(false))
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
_, err := f.createSettings(ctx, "account1", "user1", "central.example.com", "")
require.Error(t, err, "casing must not become a way past the capability check")
assert.Contains(t, err.Error(), "private capabilities")
})
}
// TestCreateSettingsRejectsForeignCluster pins tenant consistency on the pin:
// an account may not pin its gateway onto a host another account's proxy
// declares. That proxy only ever receives its own account's mappings, so the
// pin could never be served, and the endpoint it assigns is immutable.
// Ownership is decided on the proxy rows, not on heartbeat freshness — a
// cluster whose proxies are merely offline is still somebody's — and on the
// normalised host, since proxies declare their address as the operator
// spelled it.
func TestCreateSettingsRejectsForeignCluster(t *testing.T) {
ctx := context.Background()
cases := map[string]struct {
spelling string
lastSeen time.Time
}{
"live": {"byop.account2.example.com", time.Now().UTC()},
"offline": {"byop.account2.example.com", time.Now().UTC().Add(-time.Hour)},
"spelled in caps": {"BYOP.Account2.Example.com", time.Now().UTC()},
}
for name, tc := range cases {
t.Run("labeled "+name, func(t *testing.T) {
f := newBootstrapFixture(t)
f.seedProxyAt(t, "proxy1", "account2", tc.spelling, ptrTo(true), tc.lastSeen)
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
_, err := f.createSettings(ctx, "account1", "user1", "byop.account2.example.com", "")
f.requireForeignClusterRefusal(t, err, "account1")
})
t.Run("self-addressed "+name, func(t *testing.T) {
f := newBootstrapFixture(t)
f.seedProxyAt(t, "proxy1", "account2", tc.spelling, ptrTo(true), tc.lastSeen)
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
_, err := f.createSettings(ctx, "account1", "user1", "", "byop.account2.example.com")
f.requireForeignClusterRefusal(t, err, "account1")
})
}
}
// TestCreateSettingsSharedClusterStaysPinnable pins the constraint the
// ownership check must respect: a shared (NetBird-operated) cluster is not
// anybody's, so any number of accounts pin their gateways to it — including
// an account that also runs a proxy of its own elsewhere.
func TestCreateSettingsSharedClusterStaysPinnable(t *testing.T) {
ctx := context.Background()
f := newBootstrapFixture(t)
f.seedProxy(t, "shared", "", "eu.proxy.netbird.io", ptrTo(true))
f.seedProxy(t, "own", "account1", "byop.account1.example.com", ptrTo(true))
for _, account := range []string{"account1", "account2"} {
f.expectPermission(account, "user", modules.AgentNetworkSettings, operations.Create, true)
created, err := f.createSettings(ctx, account, "user", "eu.proxy.netbird.io", "")
require.NoError(t, err, "a shared cluster must stay pinnable by %s", account)
assert.Equal(t, "eu.proxy.netbird.io", created.ProxyAddress)
}
}
// TestCreateSettingsOwnClusterIsPinnable is the BYOP order in both directions:
// the account's own proxy is not a competing claim, whether the pin is labeled
// beneath its cluster or self-addressed onto the very host it declares.
func TestCreateSettingsOwnClusterIsPinnable(t *testing.T) {
ctx := context.Background()
t.Run("labeled", func(t *testing.T) {
f := newBootstrapFixture(t)
f.seedProxy(t, "own", "account1", "byop.account1.example.com", ptrTo(true))
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
created, err := f.createSettings(ctx, "account1", "user1", "byop.account1.example.com", "")
require.NoError(t, err, "the account's own cluster must be pinnable")
assert.True(t, strings.HasSuffix(created.Domain, ".byop.account1.example.com"))
})
t.Run("self-addressed", func(t *testing.T) {
f := newBootstrapFixture(t)
f.seedProxy(t, "own", "account1", "gw.account1.example.com", ptrTo(true))
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
created, err := f.createSettings(ctx, "account1", "user1", "", "gw.account1.example.com")
require.NoError(t, err, "the host the account's own proxy declares must be pinnable")
assert.Equal(t, "gw.account1.example.com", created.ProxyAddress)
})
}
// TestCreateSettingsUnknownHostIsPinnable pins the address-first order: a host
// no proxy has ever declared is nobody's, so the pin goes through and the
// proxy is deployed after.
func TestCreateSettingsUnknownHostIsPinnable(t *testing.T) {
ctx := context.Background()
f := newBootstrapFixture(t)
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
created, err := f.createSettings(ctx, "account1", "user1", "future.example.com", "")
require.NoError(t, err, "a host no proxy has declared must stay pinnable")
assert.Equal(t, "future.example.com", created.ProxyAddress)
}
// TestCreateSettingsRejectsHostAnotherAccountPinned covers claims made by pins
// rather than proxies, which the proxy-row check cannot see. A labeled pin
// beneath a host makes that host the other account's cluster, so a
// self-addressed endpoint on it would never be served; a self-addressed
// endpoint on a host makes the proxy declaring it theirs, so a label beneath
// it would never be served either. Neither is a shared-cluster shape: many
// labeled pins under one cluster are asked about in neither direction.
func TestCreateSettingsRejectsHostAnotherAccountPinned(t *testing.T) {
ctx := context.Background()
t.Run("self-addressed onto another account's cluster", func(t *testing.T) {
f := newBootstrapFixture(t)
f.expectPermission("account2", "user2", modules.AgentNetworkSettings, operations.Create, true)
_, err := f.createSettings(ctx, "account2", "user2", "gw.example.com", "")
require.NoError(t, err, "account2's labeled pin beneath the host must go through first")
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
_, err = f.createSettings(ctx, "account1", "user1", "", "gw.example.com")
f.requireForeignClusterRefusal(t, err, "account1")
})
t.Run("labeled beneath another account's endpoint", func(t *testing.T) {
f := newBootstrapFixture(t)
f.expectPermission("account2", "user2", modules.AgentNetworkSettings, operations.Create, true)
_, err := f.createSettings(ctx, "account2", "user2", "", "gw.example.com")
require.NoError(t, err, "account2's self-addressed endpoint must go through first")
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
_, err = f.createSettings(ctx, "account1", "user1", "gw.example.com", "")
f.requireForeignClusterRefusal(t, err, "account1")
})
t.Run("labeled beside another account's labeled pin stays allowed", func(t *testing.T) {
f := newBootstrapFixture(t)
for _, account := range []string{"account1", "account2"} {
f.expectPermission(account, "user", modules.AgentNetworkSettings, operations.Create, true)
_, err := f.createSettings(ctx, account, "user", "eu.proxy.netbird.io", "")
require.NoError(t, err, "labeled pins under one cluster are the shared-cluster shape and must not refuse each other")
}
})
}
// TestCreateSettingsSelfAddressedRequiresPrivateCluster pins that the
// capability gate applies to a self-addressed endpoint too: the service behind
// it is the same private one, so a proxy that already declares the hostname
// must have private capabilities, whether the account's own or a shared cluster's. A
// hostname no proxy declares yet stays claimable (TestCreateSettingsSelfAddressed).
func TestCreateSettingsSelfAddressedRequiresPrivateCluster(t *testing.T) {
ctx := context.Background()
t.Run("centralised proxy at the hostname is refused", func(t *testing.T) {
f := newBootstrapFixture(t)
f.seedProxy(t, "central", "", "gw.example.com", ptrTo(false))
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
_, err := f.createSettings(ctx, "account1", "user1", "", "gw.example.com")
require.Error(t, err, "a self-addressed endpoint on a centralised proxy can never be served")
var sErr *status.Error
require.ErrorAs(t, err, &sErr)
assert.Equal(t, status.InvalidArgument, sErr.Type())
assert.Contains(t, err.Error(), "private capabilities")
_, err = f.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "account1")
assert.Error(t, err, "no row may be left behind by a rejected bootstrap")
})
t.Run("private proxy at the hostname is accepted", func(t *testing.T) {
f := newBootstrapFixture(t)
f.seedProxy(t, "private", "", "gw.example.com", ptrTo(true))
f.expectPermission("account1", "user1", modules.AgentNetworkSettings, operations.Create, true)
created, err := f.createSettings(ctx, "account1", "user1", "", "gw.example.com")
require.NoError(t, err)
assert.Equal(t, "gw.example.com", created.ProxyAddress)
})
}
+7 -3
View File
@@ -3516,7 +3516,13 @@ func buildTestManager(t testing.TB, store store.Store, nmdataStore *networkmapdb
eventStore := &activity.InMemoryEventStore{}
metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
// Everything built here watches this context; cancelling it on cleanup stops
// the metrics flushers, caches and controllers instead of leaking them for
// the rest of the package run.
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
metrics, err := telemetry.NewDefaultAppMetrics(ctx)
if err != nil {
return nil, nil, err
}
@@ -3535,8 +3541,6 @@ func buildTestManager(t testing.TB, store store.Store, nmdataStore *networkmapdb
Return(nil).
AnyTimes()
ctx := context.Background()
cacheStore, err := cache.NewStore(ctx, 100*time.Millisecond, 300*time.Millisecond, 100)
if err != nil {
return nil, nil, err
@@ -89,6 +89,7 @@ func TestAgentNetwork_UpdateSettings_PreservesImmutableAndTogglesCollection(t *t
// Bootstrap is an explicit settings create; providers have no settings
// side effects anymore.
seedPrivateProxyCluster(t, am.Store, clusterAddr)
before, err := mgr.CreateSettings(ctx, adminUserID, agenttypes.DefaultSettings(accountID), clusterAddr, "")
require.NoError(t, err, "CreateSettings must bootstrap the row")
require.Equal(t, clusterAddr, before.ProxyAddress, "proxy address pinned at bootstrap")
@@ -13,6 +13,7 @@ import (
networkmap "github.com/netbirdio/netbird/management/internals/controllers/network_map"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
agenttypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/permissions"
"github.com/netbirdio/netbird/management/server/store"
@@ -92,6 +93,7 @@ func TestAgentNetwork_ProviderCRUD_FansOutToProxyAndClientPeers(t *testing.T) {
// UpdateAccountPeers, which is the path under test.
agentMgr := agentnetwork.NewManager(am.Store, permissions.NewManager(am.Store), am, nil)
seedPrivateProxyCluster(t, am.Store, clusterAddr)
_, err = agentMgr.CreateSettings(ctx, adminUserID, agenttypes.DefaultSettings(accountID), clusterAddr, "")
require.NoError(t, err, "CreateSettings must bootstrap the endpoint")
// The bootstrap itself reconciles and queues updates on both channels;
@@ -222,3 +224,22 @@ func synthZoneRData(sync *nbproto.SyncResponse, clusterAddr, fqdn string) string
}
return ""
}
// seedPrivateProxyCluster registers a connected proxy with private capabilities in a
// netbird client for clusterAddr, matching what a real deployment looks like
// when the account bootstraps: the agent-network gateway service is always
// private, so its cluster has to be one that can serve private services.
func seedPrivateProxyCluster(t *testing.T, st store.Store, clusterAddr string) {
t.Helper()
private := true
now := time.Now().UTC()
require.NoError(t, st.SaveProxy(context.Background(), &rpproxy.Proxy{
ID: "agent-net-proxy-" + clusterAddr,
SessionID: "agent-net-session",
ClusterAddress: clusterAddr,
LastSeen: now,
ConnectedAt: &now,
Status: rpproxy.StatusConnected,
Capabilities: rpproxy.Capabilities{Private: &private},
}), "seeding the proxy cluster must succeed")
}
+66 -22
View File
@@ -3157,7 +3157,12 @@ func NewMysqlStore(ctx context.Context, dsn string, metrics telemetry.AppMetrics
return nil, err
}
return NewSqlStore(ctx, db, types.MysqlStoreEngine, metrics, skipMigration)
store, err := NewSqlStore(ctx, db, types.MysqlStoreEngine, metrics, skipMigration)
if err != nil {
closeGormDB(db)
return nil, err
}
return store, nil
}
func getGormConfig() *gorm.Config {
@@ -3216,23 +3221,20 @@ func NewSqliteStoreFromFileStore(ctx context.Context, fileStore *FileStore, data
// NewPostgresqlStoreFromSqlStore restores a store from SqlStore and stores Postgres DB.
func NewPostgresqlStoreFromSqlStore(ctx context.Context, sqliteStore *SqlStore, dsn string, metrics telemetry.AppMetrics) (*SqlStore, error) {
store, err := NewPostgresqlStoreForTests(ctx, dsn, metrics, false)
return newPostgresqlStoreFromSqlStore(ctx, sqliteStore, dsn, metrics, false)
}
func newPostgresqlStoreFromSqlStore(ctx context.Context, sqliteStore *SqlStore, dsn string, metrics telemetry.AppMetrics, skipMigration bool) (*SqlStore, error) {
store, err := NewPostgresqlStoreForTests(ctx, dsn, metrics, skipMigration)
if err != nil {
return nil, err
}
err = store.SaveInstallationID(ctx, sqliteStore.GetInstallationID())
if err != nil {
if err := seedFromSqliteStore(ctx, store, sqliteStore); err != nil {
closeStore(ctx, store)
return nil, err
}
for _, account := range sqliteStore.GetAllAccounts(ctx) {
err := store.SaveAccount(ctx, account)
if err != nil {
return nil, err
}
}
return store, nil
}
@@ -3244,11 +3246,14 @@ func NewPostgresqlStoreForTests(ctx context.Context, dsn string, metrics telemet
}
pool, err := connectToPgDbForTests(context.Background(), dsn)
if err != nil {
closeGormDB(db)
return nil, err
}
store, err := NewSqlStore(ctx, db, types.PostgresStoreEngine, metrics, skipMigration)
if err != nil {
// Release the sessions, or the caller cannot drop the database.
pool.Close()
closeGormDB(db)
return nil, err
}
store.pool = pool
@@ -3282,22 +3287,42 @@ func connectToPgDbForTests(ctx context.Context, dsn string) (*pgxpool.Pool, erro
// NewMysqlStoreFromSqlStore restores a store from SqlStore and stores MySQL DB.
func NewMysqlStoreFromSqlStore(ctx context.Context, sqliteStore *SqlStore, dsn string, metrics telemetry.AppMetrics) (*SqlStore, error) {
store, err := NewMysqlStore(ctx, dsn, metrics, false)
if err != nil {
return nil, err
}
return newMysqlStoreFromSqlStore(ctx, sqliteStore, dsn, metrics, false)
}
err = store.SaveInstallationID(ctx, sqliteStore.GetInstallationID())
if err != nil {
return nil, err
// seedFromSqliteStore copies the installation ID and the accounts of the
// sqlite seed store into a freshly created engine store.
func seedFromSqliteStore(ctx context.Context, store, sqliteStore *SqlStore) error {
if err := store.SaveInstallationID(ctx, sqliteStore.GetInstallationID()); err != nil {
return err
}
for _, account := range sqliteStore.GetAllAccounts(ctx) {
err := store.SaveAccount(ctx, account)
if err != nil {
return nil, err
if err := store.SaveAccount(ctx, account); err != nil {
return err
}
}
return nil
}
// closeStore releases a store that is not handed to the caller, so a failed
// seed does not leak its connection and pool.
func closeStore(ctx context.Context, store *SqlStore) {
store.Close(ctx)
if store.pool != nil {
store.pool.Close()
}
}
func newMysqlStoreFromSqlStore(ctx context.Context, sqliteStore *SqlStore, dsn string, metrics telemetry.AppMetrics, skipMigration bool) (*SqlStore, error) {
store, err := NewMysqlStore(ctx, dsn, metrics, skipMigration)
if err != nil {
return nil, err
}
if err := seedFromSqliteStore(ctx, store, sqliteStore); err != nil {
closeStore(ctx, store)
return nil, err
}
return store, nil
}
@@ -6449,6 +6474,25 @@ func (s *SqlStore) IsClusterAddressConflicting(ctx context.Context, clusterAddre
return count > 0, nil
}
// HasForeignAccountProxyAtHost reports whether a proxy owned by a different
// account declares this host. Shared proxies (account_id IS NULL) are not
// foreign: a shared cluster is what most accounts pin their agent network
// gateway to. The match folds case because proxies declare their address as
// the operator spelled it while the caller's host is normalised; that costs a
// scan of the proxies table, taken once per account when its gateway is
// bootstrapped, not on the per-connect path IsClusterAddressConflicting serves.
func (s *SqlStore) HasForeignAccountProxyAtHost(ctx context.Context, host, accountID string) (bool, error) {
var count int64
result := s.db.
Model(&proxy.Proxy{}).
Where("LOWER(cluster_address) = LOWER(?) AND account_id IS NOT NULL AND account_id != ?", host, accountID).
Count(&count)
if result.Error != nil {
return false, status.Errorf(status.Internal, "check proxy host ownership: %v", result.Error)
}
return count > 0, nil
}
func (s *SqlStore) DeleteAccountCluster(ctx context.Context, clusterAddress, accountID string) error {
result := s.db.
Where("cluster_address = ? AND account_id = ?", clusterAddress, accountID).
@@ -315,6 +315,36 @@ func (s *SqlStore) GetAllAgentNetworkSettings(ctx context.Context, lockStrength
return settings, nil
}
// HasGatewayClusterPinnedByOtherAccount reports whether another account has a
// labeled agent network gateway pinned beneath host, making host its cluster.
// A self-addressed endpoint on the very same hostname is not counted: that
// collision is the domain unique index's to refuse, as a conflict. Case-folded,
// since a settings row written before hostnames were normalised may carry
// capitals; one row per account, so the scan is cheap.
func (s *SqlStore) HasGatewayClusterPinnedByOtherAccount(ctx context.Context, host, accountID string) (bool, error) {
return s.countGatewayRowsByOtherAccount(ctx, "LOWER(proxy_address) = LOWER(?) AND LOWER(domain) <> LOWER(proxy_address)", host, accountID)
}
// HasGatewayEndpointByOtherAccount reports whether host is another account's
// agent network endpoint hostname (domain). Case-folded for the same reason as
// HasGatewayClusterPinnedByOtherAccount.
func (s *SqlStore) HasGatewayEndpointByOtherAccount(ctx context.Context, host, accountID string) (bool, error) {
return s.countGatewayRowsByOtherAccount(ctx, "LOWER(domain) = LOWER(?)", host, accountID)
}
func (s *SqlStore) countGatewayRowsByOtherAccount(ctx context.Context, predicate, host, accountID string) (bool, error) {
var count int64
result := s.db.
Model(&agentNetworkTypes.Settings{}).
Where(predicate+" AND account_id != ?", host, accountID).
Count(&count)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to check agent network gateway claims at host: %v", result.Error)
return false, status.Errorf(status.Internal, "check agent network gateway claims at host")
}
return count > 0, nil
}
// GetAgentNetworkSettingsByProxyAddress returns every Settings row whose
// gateway is served by the proxy declaring the given cluster address. Used by
// cluster-scoped synthesis to find the accounts a shared proxy serves.
+242 -14
View File
@@ -4,6 +4,7 @@ package store
import (
"context"
"database/sql"
"errors"
"fmt"
"net"
@@ -15,6 +16,7 @@ import (
"runtime"
"slices"
"strings"
"sync"
"time"
"github.com/google/uuid"
@@ -340,6 +342,9 @@ type Store interface {
CountProxiesByAccountID(ctx context.Context, accountID string) (int64, error)
IsClusterAddressConflicting(ctx context.Context, clusterAddress, accountID string) (bool, error)
HasActiveProxyAtClusterAddress(ctx context.Context, clusterAddress string) (bool, error)
HasForeignAccountProxyAtHost(ctx context.Context, host, accountID string) (bool, error)
HasGatewayClusterPinnedByOtherAccount(ctx context.Context, host, accountID string) (bool, error)
HasGatewayEndpointByOtherAccount(ctx context.Context, host, accountID string) (bool, error)
DeleteAccountCluster(ctx context.Context, clusterAddress, accountID string) error
GetCustomDomainsCounts(ctx context.Context) (total int64, validated int64, err error)
@@ -732,6 +737,7 @@ func NewTestStoreFromSQL(ctx context.Context, filename string, dataDir string) (
time.Sleep(100 * time.Millisecond)
}
}
store.Close(ctx)
return nil, nil, fmt.Errorf("failed to create test store after %d attempts: %v", maxRetries, err)
}
@@ -758,14 +764,15 @@ func addAllGroupToAccount(ctx context.Context, store Store) error {
return nil
}
func getSqlStoreEngine(ctx context.Context, store *SqlStore, kind types.Engine) (Store, func(), error) {
func getSqlStoreEngine(ctx context.Context, sqliteStore *SqlStore, kind types.Engine) (Store, func(), error) {
store := sqliteStore
var cleanup func()
var err error
switch kind {
case types.PostgresStoreEngine:
store, cleanup, err = newReusedPostgresStore(ctx, store, kind)
store, cleanup, err = newReusedPostgresStore(ctx, sqliteStore, kind)
case types.MysqlStoreEngine:
store, cleanup, err = newReusedMysqlStore(ctx, store, kind)
store, cleanup, err = newReusedMysqlStore(ctx, sqliteStore, kind)
default:
cleanup = func() {
// sqlite doesn't need to be cleaned up
@@ -781,6 +788,11 @@ func getSqlStoreEngine(ctx context.Context, store *SqlStore, kind types.Engine)
if store.pool != nil {
store.pool.Close()
}
if store != sqliteStore {
// The sqlite store only seeded the engine under test; without this
// every test leaks its connection and the opener goroutines.
sqliteStore.Close(ctx)
}
}
return store, closeConnection, nil
@@ -805,19 +817,23 @@ func newReusedPostgresStore(ctx context.Context, store *SqlStore, kind types.Eng
return nil, nil, fmt.Errorf("failed to open postgres connection: %v", err)
}
dsn, cleanup, err := createRandomDB(dsn, db, kind)
sqlDB, _ := db.DB()
if sqlDB != nil {
sqlDB.Close()
template, err := postgresSchemaTemplate(ctx, dsn, db)
if err != nil {
closeGormDB(db)
return nil, nil, err
}
dsn, cleanup, err := createRandomDB(dsn, db, kind, template)
closeGormDB(db)
if err != nil {
return nil, nil, err
}
store, err = NewPostgresqlStoreFromSqlStore(ctx, store, dsn, nil)
store, err = newPostgresqlStoreFromSqlStore(ctx, store, dsn, nil, true)
if err != nil {
cleanup()
return nil, nil, err
}
@@ -850,7 +866,13 @@ func newReusedMysqlStore(ctx context.Context, store *SqlStore, kind types.Engine
sqlDB.SetMaxOpenConns(1)
sqlDB.SetMaxIdleConns(1)
dsn, cleanup, err := createRandomDB(dsn, db, kind)
tableDDL, err := mysqlSchemaTemplate(ctx, dsn, db)
if err != nil {
sqlDB.Close()
return nil, nil, err
}
dsn, cleanup, err := createRandomDB(dsn, db, kind, "")
sqlDB.Close()
@@ -858,14 +880,200 @@ func newReusedMysqlStore(ctx context.Context, store *SqlStore, kind types.Engine
return nil, nil, err
}
store, err = NewMysqlStoreFromSqlStore(ctx, store, dsn, nil)
if err := cloneMysqlSchema(ctx, dsn, tableDDL); err != nil {
cleanup()
return nil, nil, err
}
store, err = newMysqlStoreFromSqlStore(ctx, store, dsn, nil, true)
if err != nil {
cleanup()
return nil, nil, err
}
return store, cleanup, nil
}
// schemaTemplates remembers, per engine and server, a database that went
// through the full migration once in this process. Every later test database
// is cloned from it, so a test pays for CREATE DATABASE and a schema copy
// instead of the 40-table AutoMigrate plus every pre and post migration, which
// is what made each MySQL test store cost well over a second in CI.
var (
schemaTemplatesMu sync.Mutex
schemaTemplates = map[string]*schemaTemplate{}
)
type schemaTemplate struct {
dbName string
// tableDDL holds the CREATE TABLE statements of the template. MySQL has no
// server-side database template, so the schema is replayed statement by
// statement into each test database.
tableDDL []string
}
func schemaTemplateKey(engine types.Engine, dsn string) string {
return string(engine) + "|" + dsn
}
func newTestDBName(prefix string) string {
return fmt.Sprintf("%s_%s", prefix, strings.ReplaceAll(uuid.New().String(), "-", "_"))
}
// postgresSchemaTemplate returns the name of a fully migrated database that
// CREATE DATABASE ... TEMPLATE can copy, creating it on first use.
func postgresSchemaTemplate(ctx context.Context, baseDSN string, admin *gorm.DB) (string, error) {
schemaTemplatesMu.Lock()
defer schemaTemplatesMu.Unlock()
key := schemaTemplateKey(types.PostgresStoreEngine, baseDSN)
if tpl, ok := schemaTemplates[key]; ok {
return tpl.dbName, nil
}
name := newTestDBName("test_template")
if err := admin.Exec(fmt.Sprintf("CREATE DATABASE %s", name)).Error; err != nil {
return "", fmt.Errorf("create postgres template database: %w", err)
}
tplStore, err := NewPostgresqlStoreForTests(ctx, replaceDBName(baseDSN, name), nil, false)
if err != nil {
dropDatabase(admin, name)
return "", fmt.Errorf("migrate postgres template database: %w", err)
}
// TEMPLATE refuses a source that still has sessions, so release both handles
// before the first clone.
tplStore.Close(ctx)
if tplStore.pool != nil {
tplStore.pool.Close()
}
schemaTemplates[key] = &schemaTemplate{dbName: name}
return name, nil
}
// mysqlSchemaTemplate returns the CREATE TABLE statements of a fully migrated
// database, migrating one on first use.
func mysqlSchemaTemplate(ctx context.Context, baseDSN string, admin *gorm.DB) ([]string, error) {
schemaTemplatesMu.Lock()
defer schemaTemplatesMu.Unlock()
key := schemaTemplateKey(types.MysqlStoreEngine, baseDSN)
if tpl, ok := schemaTemplates[key]; ok {
return tpl.tableDDL, nil
}
name := newTestDBName("test_template")
if err := admin.Exec(fmt.Sprintf("CREATE DATABASE %s", name)).Error; err != nil {
return nil, fmt.Errorf("create mysql template database: %w", err)
}
tplStore, err := NewMysqlStore(ctx, replaceDBName(baseDSN, name), nil, false)
if err != nil {
dropDatabase(admin, name)
return nil, fmt.Errorf("migrate mysql template database: %w", err)
}
tableDDL, err := mysqlTableDDL(ctx, tplStore.db, name)
tplStore.Close(ctx)
if err != nil {
dropDatabase(admin, name)
return nil, err
}
schemaTemplates[key] = &schemaTemplate{dbName: name, tableDDL: tableDDL}
return tableDDL, nil
}
func mysqlTableDDL(ctx context.Context, db *gorm.DB, dbName string) ([]string, error) {
sqlDB, err := db.DB()
if err != nil {
return nil, err
}
tables, err := mysqlTableNames(ctx, sqlDB, dbName)
if err != nil {
return nil, err
}
tableDDL := make([]string, 0, len(tables))
for _, table := range tables {
var name, createStmt string
row := sqlDB.QueryRowContext(ctx, fmt.Sprintf("SHOW CREATE TABLE %s.%s", dbName, table))
if err := row.Scan(&name, &createStmt); err != nil {
return nil, fmt.Errorf("read create statement of %s: %w", table, err)
}
tableDDL = append(tableDDL, createStmt)
}
return tableDDL, nil
}
func mysqlTableNames(ctx context.Context, sqlDB *sql.DB, dbName string) ([]string, error) {
rows, err := sqlDB.QueryContext(ctx, fmt.Sprintf("SHOW TABLES FROM %s", dbName))
if err != nil {
return nil, fmt.Errorf("list template tables: %w", err)
}
defer rows.Close()
var tables []string
for rows.Next() {
var table string
if err := rows.Scan(&table); err != nil {
return nil, fmt.Errorf("scan template table name: %w", err)
}
tables = append(tables, table)
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("list template tables: %w", err)
}
return tables, nil
}
// cloneMysqlSchema replays the template's CREATE TABLE statements into the
// database the DSN points at.
func cloneMysqlSchema(ctx context.Context, dsn string, tableDDL []string) error {
db, err := gorm.Open(mysql.Open(dsn+"?charset=utf8&parseTime=True&loc=Local"), getGormConfig())
if err != nil {
return fmt.Errorf("connect to test database: %w", err)
}
sqlDB, err := db.DB()
if err != nil {
return err
}
defer sqlDB.Close()
// The statements come out of SHOW TABLES in name order, not dependency
// order, and their foreign keys reference tables of the session's default
// database. Pin a single connection so the session setting below covers
// every statement, and connect straight to the new database so unqualified
// references land there.
sqlDB.SetMaxOpenConns(1)
if _, err := sqlDB.ExecContext(ctx, "SET FOREIGN_KEY_CHECKS = 0"); err != nil {
return fmt.Errorf("disable foreign key checks: %w", err)
}
for _, stmt := range tableDDL {
if _, err := sqlDB.ExecContext(ctx, stmt); err != nil {
return fmt.Errorf("replay table definition: %w", err)
}
}
return nil
}
// dropDatabase removes a template that never became usable, so a failed setup
// does not leave it behind on a shared server. The server may still be tearing
// down the sessions the failed migration held, so the drop retries while
// Postgres reports the database as in use.
func dropDatabase(admin *gorm.DB, name string) {
if err := execWithTemplateRetry(admin, fmt.Sprintf("DROP DATABASE IF EXISTS %s", name)); err != nil {
log.Warnf("failed to drop template database %s: %v", name, err)
}
}
func closeGormDB(db *gorm.DB) {
if sqlDB, _ := db.DB(); sqlDB != nil {
sqlDB.Close()
}
}
func openDBWithRetry(dsn string, engine types.Engine, maxRetries int) (*gorm.DB, error) {
var db *gorm.DB
var err error
@@ -891,10 +1099,16 @@ func openDBWithRetry(dsn string, engine types.Engine, maxRetries int) (*gorm.DB,
return nil, err
}
func createRandomDB(dsn string, db *gorm.DB, engine types.Engine) (string, func(), error) {
dbName := fmt.Sprintf("test_db_%s", strings.ReplaceAll(uuid.New().String(), "-", "_"))
// createRandomDB creates a uniquely named database for one test. On postgres a
// non-empty template is copied server-side with CREATE DATABASE ... TEMPLATE.
func createRandomDB(dsn string, db *gorm.DB, engine types.Engine, template string) (string, func(), error) {
dbName := newTestDBName("test_db")
if err := db.Exec(fmt.Sprintf("CREATE DATABASE %s", dbName)).Error; err != nil {
createStmt := fmt.Sprintf("CREATE DATABASE %s", dbName)
if template != "" && engine == types.PostgresStoreEngine {
createStmt = fmt.Sprintf("CREATE DATABASE %s TEMPLATE %s", dbName, template)
}
if err := execWithTemplateRetry(db, createStmt); err != nil {
return "", nil, fmt.Errorf("failed to create database: %v", err)
}
@@ -960,6 +1174,20 @@ func createRandomDB(dsn string, db *gorm.DB, engine types.Engine) (string, func(
return replaceDBName(dsn, dbName), cleanup, nil
}
// execWithTemplateRetry runs a statement, retrying briefly when postgres still
// sees the template's just-closed sessions and refuses to copy it.
func execWithTemplateRetry(db *gorm.DB, stmt string) error {
var err error
for attempt := 0; attempt < 20; attempt++ {
err = db.Exec(stmt).Error
if err == nil || !strings.Contains(err.Error(), "is being accessed by other users") {
return err
}
time.Sleep(100 * time.Millisecond)
}
return err
}
func replaceDBName(dsn, newDBName string) string {
re := regexp.MustCompile(`(?P<pre>[:/@])(?P<dbname>[^/?]+)(?P<post>\?|$)`)
return re.ReplaceAllString(dsn, `${pre}`+newDBName+`${post}`)
+45
View File
@@ -3065,6 +3065,51 @@ func (mr *MockStoreMockRecorder) HasActiveProxyAtClusterAddress(ctx, clusterAddr
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HasActiveProxyAtClusterAddress", reflect.TypeOf((*MockStore)(nil).HasActiveProxyAtClusterAddress), ctx, clusterAddress)
}
// HasForeignAccountProxyAtHost mocks base method.
func (m *MockStore) HasForeignAccountProxyAtHost(ctx context.Context, host, accountID string) (bool, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "HasForeignAccountProxyAtHost", ctx, host, accountID)
ret0, _ := ret[0].(bool)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// HasForeignAccountProxyAtHost indicates an expected call of HasForeignAccountProxyAtHost.
func (mr *MockStoreMockRecorder) HasForeignAccountProxyAtHost(ctx, host, accountID any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HasForeignAccountProxyAtHost", reflect.TypeOf((*MockStore)(nil).HasForeignAccountProxyAtHost), ctx, host, accountID)
}
// HasGatewayClusterPinnedByOtherAccount mocks base method.
func (m *MockStore) HasGatewayClusterPinnedByOtherAccount(ctx context.Context, host, accountID string) (bool, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "HasGatewayClusterPinnedByOtherAccount", ctx, host, accountID)
ret0, _ := ret[0].(bool)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// HasGatewayClusterPinnedByOtherAccount indicates an expected call of HasGatewayClusterPinnedByOtherAccount.
func (mr *MockStoreMockRecorder) HasGatewayClusterPinnedByOtherAccount(ctx, host, accountID any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HasGatewayClusterPinnedByOtherAccount", reflect.TypeOf((*MockStore)(nil).HasGatewayClusterPinnedByOtherAccount), ctx, host, accountID)
}
// HasGatewayEndpointByOtherAccount mocks base method.
func (m *MockStore) HasGatewayEndpointByOtherAccount(ctx context.Context, host, accountID string) (bool, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "HasGatewayEndpointByOtherAccount", ctx, host, accountID)
ret0, _ := ret[0].(bool)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// HasGatewayEndpointByOtherAccount indicates an expected call of HasGatewayEndpointByOtherAccount.
func (mr *MockStoreMockRecorder) HasGatewayEndpointByOtherAccount(ctx, host, accountID any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HasGatewayEndpointByOtherAccount", reflect.TypeOf((*MockStore)(nil).HasGatewayEndpointByOtherAccount), ctx, host, accountID)
}
// IncrementAgentNetworkConsumption mocks base method.
func (m *MockStore) IncrementAgentNetworkConsumption(ctx context.Context, accountID string, kind types.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time, tokensIn, tokensOut int64, costUSD float64) error {
m.ctrl.T.Helper()
+12
View File
@@ -37,6 +37,18 @@ func CreateMysqlTestContainer() (func(), string, error) {
mysql.WithDatabase("testing"),
mysql.WithUsername("root"),
mysql.WithPassword("testing"),
// Every test creates and drops a database with about 40 tables, so with
// the server defaults the run is dominated by durability work: each
// CREATE TABLE fsyncs the redo log, the binary log and the doublewrite
// buffer. None of it protects anything in a container that is discarded
// after the run. Tables stay in per-table files on purpose: in the shared
// system tablespace the cost of every CREATE and DROP grew with the number
// of databases the run had already created.
testcontainers.WithCmd("mysqld",
"--innodb-flush-log-at-trx-commit=0",
"--innodb-doublewrite=OFF",
"--skip-log-bin",
),
testcontainers.WithWaitStrategy(
wait.ForLog("/usr/sbin/mysqld: ready for connections").
WithOccurrence(1).WithStartupTimeout(15*time.Second).WithPollInterval(100*time.Millisecond),