Merge branch 'main' into embedded-vnc

# Conflicts:
#	client/ui/frontend/src/app.tsx
#	client/ui/frontend/src/modules/main/MainConnectionStatusSwitch.tsx
#	client/ui/i18n/locales/uk/common.json
#	go.sum
This commit is contained in:
Viktor Liu
2026-09-22 14:53:34 +02:00
205 changed files with 12118 additions and 1280 deletions
+114 -28
View File
@@ -3160,7 +3160,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 {
@@ -3219,23 +3224,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
}
@@ -3247,11 +3249,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
@@ -3285,22 +3290,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
}
@@ -3479,7 +3504,7 @@ func (s *SqlStore) GetPeerGroups(ctx context.Context, lockStrength LockingStreng
var groups []*types.Group
query := tx.
Joins("JOIN group_peers ON group_peers.group_id = groups.id").
Where("group_peers.peer_id = ?", peerId).
Where("groups.account_id = ? AND group_peers.peer_id = ?", accountId, peerId).
Preload(clause.Associations).
Find(&groups)
@@ -5059,7 +5084,7 @@ func (s *SqlStore) GetPeersByGroupIDs(ctx context.Context, accountID string, gro
Select("DISTINCT peer_id").
Where("account_id = ? AND group_id IN ?", accountID, groupIDs)
result := s.db.Where("id IN (?)", peerIDsSubquery).Find(&peers)
result := s.db.Where("account_id = ? AND id IN (?)", accountID, peerIDsSubquery).Find(&peers)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to get peers by group IDs: %s", result.Error)
return nil, status.Errorf(status.Internal, "failed to get peers by group IDs")
@@ -5692,6 +5717,23 @@ func (s *SqlStore) ListCustomDomains(ctx context.Context, accountID string) ([]*
return domains, nil
}
// GetCustomDomainByName returns the custom domain row holding the given name,
// regardless of which account owns it.
func (s *SqlStore) GetCustomDomainByName(ctx context.Context, domainName string) (*domain.Domain, error) {
customDomain := &domain.Domain{}
result := s.db.Take(customDomain, "domain = ?", domainName)
if result.Error != nil {
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
return nil, status.Errorf(status.NotFound, "custom domain %s not found", domainName)
}
log.WithContext(ctx).Errorf("failed to get custom domain by name from store: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get custom domain from store")
}
return customDomain, nil
}
func (s *SqlStore) CreateCustomDomain(ctx context.Context, accountID string, domainName string, targetCluster string, validated bool) (*domain.Domain, error) {
newDomain := &domain.Domain{
ID: xid.New().String(), // Generate our own ID because gorm doesn't always configure the database to handle this for us.
@@ -5701,8 +5743,24 @@ func (s *SqlStore) CreateCustomDomain(ctx context.Context, accountID string, dom
Type: domain.TypeCustom,
Validated: validated,
}
if !validated {
expiresAt := time.Now().UTC().Add(domain.ValidationTTL)
newDomain.ValidationExpiresAt = &expiresAt
}
result := s.db.Create(newDomain)
if result.Error != nil {
// The unique index is the last guard when two requests clear the
// manager's availability check at the same time. The one that loses the
// insert is a conflict, not an internal failure.
var count int64
if err := s.db.Model(&domain.Domain{}).Where("domain = ?", domainName).Count(&count).Error; err == nil && count > 0 {
// The insert error is logged even on this path: the name being taken
// is what the caller has to act on, but if the insert also failed for
// an unrelated reason the operator still needs to see it.
log.WithContext(ctx).Warnf("create reverse proxy custom domain %s rejected, name already registered: %v", domainName, result.Error)
return nil, status.Errorf(status.AlreadyExists, "domain %s is already registered", domainName)
}
log.WithContext(ctx).Errorf("failed to create reverse proxy custom domain to store: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to create reverse proxy custom domain to store")
}
@@ -5710,12 +5768,21 @@ func (s *SqlStore) CreateCustomDomain(ctx context.Context, accountID string, dom
return newDomain, nil
}
// UpdateCustomDomain completes validation only while the original registration is pending.
func (s *SqlStore) UpdateCustomDomain(ctx context.Context, accountID string, d *domain.Domain) (*domain.Domain, error) {
d.AccountID = accountID
result := s.db.Select("*").Save(d)
if !d.Validated {
return nil, status.Errorf(status.InvalidArgument, "custom domain update must complete validation")
}
result := s.db.WithContext(ctx).Model(&domain.Domain{}).
Where(accountAndIDQueryCondition, accountID, d.ID).
Where("domain = ? AND target_cluster = ?", d.Domain, d.TargetCluster).
Where("validated = ? AND validation_expires_at > ?", false, time.Now().UTC()).
Update("validated", true)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to update reverse proxy custom domain to store: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to update reverse proxy custom domain to store")
return nil, fmt.Errorf("validate custom domain in store: %w", result.Error)
}
if result.RowsAffected == 0 {
return nil, status.Errorf(status.PreconditionFailed, "custom domain registration is no longer pending validation")
}
return d, nil
@@ -6410,6 +6477,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.
@@ -0,0 +1,60 @@
package store
import (
"context"
"fmt"
"strings"
"time"
"gorm.io/gorm"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
"github.com/netbirdio/netbird/shared/management/status"
)
// GetExpiredCustomDomains lists pending registrations in stable batches across accounts.
func (s *SqlStore) GetExpiredCustomDomains(ctx context.Context, now time.Time, afterID domain.ID, limit int) ([]*domain.Domain, error) {
var domains []*domain.Domain
result := s.db.WithContext(ctx).
Where("validated = ? AND validation_expires_at <= ? AND id > ?", false, now, string(afterID)).
Order("id").Limit(limit).Find(&domains)
if result.Error != nil {
return nil, fmt.Errorf("list expired custom domains: %w", result.Error)
}
return domains, nil
}
// DeleteExpiredCustomDomain deletes an expired registration only if no service uses its namespace.
func (s *SqlStore) DeleteExpiredCustomDomain(ctx context.Context, d *domain.Domain, now time.Time) (bool, error) {
db := s.db.WithContext(ctx)
services := customDomainServices(db, d)
result := db.Where(accountAndIDQueryCondition, d.AccountID, d.ID).
Where("domain = ? AND validated = ? AND validation_expires_at <= ?", d.Domain, false, now).
Where("NOT EXISTS (?)", services.Select("1")).Delete(&domain.Domain{})
if result.Error != nil {
return false, fmt.Errorf("delete expired custom domain: %w", result.Error)
}
if result.RowsAffected > 0 {
return true, nil
}
var count int64
if err := customDomainServices(db, d).Count(&count).Error; err != nil {
return false, fmt.Errorf("check expired custom domain services: %w", err)
}
if count > 0 {
return false, status.Errorf(status.PreconditionFailed, "expired custom domain still has dependent services")
}
return false, nil
}
func customDomainServices(db *gorm.DB, d *domain.Domain) *gorm.DB {
name := strings.ToLower(strings.TrimSuffix(d.Domain, "."))
// Shared domain validation permits underscores, and older rows may contain
// other LIKE metacharacters.
escaped := strings.NewReplacer("!", "!!", "%", "!%", "_", "!_").Replace(name)
return db.Model(&rpservice.Service{}).Where(
"LOWER(domain) IN ? OR LOWER(domain) LIKE ? ESCAPE '!' OR LOWER(domain) LIKE ? ESCAPE '!'",
[]string{name, name + "."}, "%."+escaped, "%."+escaped+".",
)
}
@@ -0,0 +1,78 @@
package store
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
)
func TestDeleteExpiredCustomDomain_ServiceDependencies(t *testing.T) {
runTestForAllEngines(t, "", func(t *testing.T, store Store) {
ctx := context.Background()
now := time.Now().UTC()
db := store.(*SqlStore).db
require.NoError(t, store.SaveAccount(ctx, newAccountWithId(ctx, "owner", "admin", "")))
for _, tt := range []struct {
name string
domainName string
serviceHost string
protected bool
}{
{"exact", "example.com", "example.com", true},
{"subdomain", "example.com", "deep.app.example.com", true},
{"case", "example.com", "APP.EXAMPLE.COM.", true},
{"suffix-boundary", "example.com", "notexample.com", false},
{"literal underscore", "a_b.example.com", "app.a_b.example.com", true},
{"underscore wildcard", "a_b.example.com", "app.axb.example.com", false},
{"legacy percent wildcard", "a%b.example.com", "app.axxb.example.com", false},
{"legacy escape character", "a!b.example.com", "app.ab.example.com", false},
} {
t.Run(tt.name, func(t *testing.T) {
d, err := store.CreateCustomDomain(ctx, "owner", tt.domainName, "cluster", false)
require.NoError(t, err)
require.NoError(t, db.Model(d).Update("validation_expires_at", now.Add(-time.Hour)).Error)
svc := &rpservice.Service{ID: "legacy", AccountID: "owner", Domain: tt.serviceHost}
require.NoError(t, store.CreateService(ctx, svc))
deleted, err := store.DeleteExpiredCustomDomain(ctx, d, now)
if tt.protected {
require.Error(t, err)
assert.False(t, deleted, "service namespaces must remain reserved")
} else {
require.NoError(t, err)
assert.True(t, deleted, "a hostname outside the namespace must not prevent cleanup")
}
require.NoError(t, db.Delete(svc).Error)
require.NoError(t, db.Delete(d).Error)
})
}
})
}
func TestDeleteExpiredCustomDomain_RechecksValidation(t *testing.T) {
runTestForAllEngines(t, "", func(t *testing.T, store Store) {
ctx := context.Background()
require.NoError(t, store.SaveAccount(ctx, newAccountWithId(ctx, "owner", "admin", "")))
d, err := store.CreateCustomDomain(ctx, "owner", "validated.example.com", "cluster", false)
require.NoError(t, err)
d, err = store.GetCustomDomain(ctx, "owner", d.ID)
require.NoError(t, err)
stale := d.Copy()
d.Validated = true
_, err = store.UpdateCustomDomain(ctx, "owner", d)
require.NoError(t, err)
deleted, err := store.DeleteExpiredCustomDomain(ctx, stale, time.Now().Add(domain.ValidationTTL))
require.NoError(t, err)
assert.False(t, deleted, "a stale cleanup candidate must not delete a validated registration")
stored, err := store.GetCustomDomain(ctx, "owner", d.ID)
require.NoError(t, err)
assert.True(t, stored.Validated, "the validated registration must remain usable")
require.NotNil(t, stored.ValidationExpiresAt)
assert.Equal(t, stale.ValidationExpiresAt, stored.ValidationExpiresAt, "validation must preserve the original deadline")
})
}
@@ -64,7 +64,7 @@ func assertGetAccountLoadsCustomDomains(t *testing.T, store Store) {
_, err := store.CreateCustomDomain(ctx, accountID, "example.com", "eu.proxy.netbird.io", true)
require.NoError(t, err, "creating the first custom domain must succeed")
_, err = store.CreateCustomDomain(ctx, accountID, "apps.acme.io", "us.proxy.netbird.io", false)
pending, err := store.CreateCustomDomain(ctx, accountID, "apps.acme.io", "us.proxy.netbird.io", false)
require.NoError(t, err, "creating the second custom domain must succeed")
account, err := store.GetAccount(ctx, accountID)
@@ -75,6 +75,10 @@ func assertGetAccountLoadsCustomDomains(t *testing.T, store Store) {
for _, d := range account.Domains {
require.NotNil(t, d)
byDomain[d.Domain] = d.TargetCluster
if d.ID == pending.ID {
require.NotNil(t, d.ValidationExpiresAt)
assert.WithinDuration(t, *pending.ValidationExpiresAt, *d.ValidationExpiresAt, time.Millisecond, "both account loaders must preserve the validation deadline")
}
}
assert.Equal(t, "eu.proxy.netbird.io", byDomain["example.com"], "custom domain must carry its target cluster")
assert.Equal(t, "us.proxy.netbird.io", byDomain["apps.acme.io"], "custom domain must carry its target cluster")
+14
View File
@@ -2847,6 +2847,14 @@ func TestSqlStore_GetPeerGroups(t *testing.T) {
groups, err = store.GetPeerGroups(context.Background(), LockingStrengthNone, accountID, peerID)
require.NoError(t, err)
assert.Len(t, groups, 2)
foreignPeerID := "foreign-peer"
err = store.AddPeerToGroup(context.Background(), accountID, foreignPeerID, "cfefqs706sqkneg59g4h")
require.NoError(t, err)
groups, err = store.GetPeerGroups(context.Background(), LockingStrengthNone, "other-account", foreignPeerID)
require.NoError(t, err)
assert.Empty(t, groups, "groups of another account must not be returned")
}
func TestSqlStore_GetAccountPeers(t *testing.T) {
@@ -4042,9 +4050,15 @@ func TestSqlStore_GetPeersByGroupIDs(t *testing.T) {
}
require.NoError(t, store.CreateGroups(ctx, accountID, groups))
otherAccount := newAccountWithId(ctx, "other-account", "other-user", "")
require.NoError(t, store.SaveAccount(ctx, otherAccount))
foreignPeer := &nbpeer.Peer{ID: "foreign-peer", AccountID: otherAccount.Id}
require.NoError(t, store.AddPeerToAccount(ctx, foreignPeer))
require.NoError(t, store.AddPeerToGroup(ctx, accountID, peer1, group1ID))
require.NoError(t, store.AddPeerToGroup(ctx, accountID, peer2, group1ID))
require.NoError(t, store.AddPeerToGroup(ctx, accountID, peer1, group2ID))
require.NoError(t, store.AddPeerToGroup(ctx, accountID, foreignPeer.ID, group1ID))
peers, err := store.GetPeersByGroupIDs(ctx, accountID, tt.groupIDs)
require.NoError(t, err)
+248 -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"
@@ -302,8 +304,11 @@ type Store interface {
GetCustomDomain(ctx context.Context, accountID string, domainID string) (*domain.Domain, error)
ListFreeDomains(ctx context.Context, accountID string) ([]string, error)
ListCustomDomains(ctx context.Context, accountID string) ([]*domain.Domain, error)
GetCustomDomainByName(ctx context.Context, domainName string) (*domain.Domain, error)
CreateCustomDomain(ctx context.Context, accountID string, domainName string, targetCluster string, validated bool) (*domain.Domain, error)
UpdateCustomDomain(ctx context.Context, accountID string, d *domain.Domain) (*domain.Domain, error)
GetExpiredCustomDomains(ctx context.Context, now time.Time, afterID domain.ID, limit int) ([]*domain.Domain, error)
DeleteExpiredCustomDomain(ctx context.Context, d *domain.Domain, now time.Time) (bool, error)
DeleteCustomDomain(ctx context.Context, accountID string, domainID string) error
CreateAccessLog(ctx context.Context, log *accesslogs.AccessLogEntry) error
@@ -337,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)
@@ -641,6 +649,9 @@ func migratePostAuto(ctx context.Context, db *gorm.DB) error {
func getMigrationsPostAuto(ctx context.Context) []migrationFunc {
return []migrationFunc{
func(db *gorm.DB) error {
return migration.MigrateCustomDomainValidationExpiry(ctx, db)
},
func(db *gorm.DB) error {
return migration.CreateIndexIfNotExists[nbpeer.Peer](ctx, db, "idx_account_ip", "account_id", "ip")
},
@@ -726,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)
}
@@ -752,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
@@ -775,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
@@ -799,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
}
@@ -844,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()
@@ -852,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
@@ -885,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)
}
@@ -954,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}`)
+90
View File
@@ -555,6 +555,21 @@ func (mr *MockStoreMockRecorder) DeleteDNSRecord(ctx, accountID, zoneID, recordI
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteDNSRecord", reflect.TypeOf((*MockStore)(nil).DeleteDNSRecord), ctx, accountID, zoneID, recordID)
}
// DeleteExpiredCustomDomain mocks base method.
func (m *MockStore) DeleteExpiredCustomDomain(ctx context.Context, d *domain.Domain, now time.Time) (bool, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "DeleteExpiredCustomDomain", ctx, d, now)
ret0, _ := ret[0].(bool)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// DeleteExpiredCustomDomain indicates an expected call of DeleteExpiredCustomDomain.
func (mr *MockStoreMockRecorder) DeleteExpiredCustomDomain(ctx, d, now any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteExpiredCustomDomain", reflect.TypeOf((*MockStore)(nil).DeleteExpiredCustomDomain), ctx, d, now)
}
// DeleteGroup mocks base method.
func (m *MockStore) DeleteGroup(ctx context.Context, accountID, groupID string) error {
m.ctrl.T.Helper()
@@ -1941,6 +1956,21 @@ func (mr *MockStoreMockRecorder) GetCustomDomain(ctx, accountID, domainID any) *
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetCustomDomain", reflect.TypeOf((*MockStore)(nil).GetCustomDomain), ctx, accountID, domainID)
}
// GetCustomDomainByName mocks base method.
func (m *MockStore) GetCustomDomainByName(ctx context.Context, domainName string) (*domain.Domain, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetCustomDomainByName", ctx, domainName)
ret0, _ := ret[0].(*domain.Domain)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetCustomDomainByName indicates an expected call of GetCustomDomainByName.
func (mr *MockStoreMockRecorder) GetCustomDomainByName(ctx, domainName any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetCustomDomainByName", reflect.TypeOf((*MockStore)(nil).GetCustomDomainByName), ctx, domainName)
}
// GetCustomDomainsCounts mocks base method.
func (m *MockStore) GetCustomDomainsCounts(ctx context.Context) (int64, int64, error) {
m.ctrl.T.Helper()
@@ -1987,6 +2017,21 @@ func (mr *MockStoreMockRecorder) GetEmbeddedProxyPeerIDsByCluster(ctx, accountID
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEmbeddedProxyPeerIDsByCluster", reflect.TypeOf((*MockStore)(nil).GetEmbeddedProxyPeerIDsByCluster), ctx, accountID)
}
// GetExpiredCustomDomains mocks base method.
func (m *MockStore) GetExpiredCustomDomains(ctx context.Context, now time.Time, afterID domain.ID, limit int) ([]*domain.Domain, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetExpiredCustomDomains", ctx, now, afterID, limit)
ret0, _ := ret[0].([]*domain.Domain)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetExpiredCustomDomains indicates an expected call of GetExpiredCustomDomains.
func (mr *MockStoreMockRecorder) GetExpiredCustomDomains(ctx, now, afterID, limit any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetExpiredCustomDomains", reflect.TypeOf((*MockStore)(nil).GetExpiredCustomDomains), ctx, now, afterID, limit)
}
// GetExpiredEphemeralServices mocks base method.
func (m *MockStore) GetExpiredEphemeralServices(ctx context.Context, ttl time.Duration, limit int) ([]*service.Service, error) {
m.ctrl.T.Helper()
@@ -3020,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()