Merge remote-tracking branch 'origin/main' into refactor/permissions-manager

This commit is contained in:
pascal
2026-09-14 13:33:20 +02:00
130 changed files with 4707 additions and 560 deletions
@@ -1,5 +1,13 @@
package domain
import "time"
// ValidationTTL is the time available to validate a custom domain registration.
const ValidationTTL = 48 * time.Hour
// ID identifies a custom domain registration.
type ID string
type Type string
const (
@@ -8,12 +16,13 @@ const (
)
type Domain struct {
ID string `gorm:"unique;primaryKey;autoIncrement"`
Domain string `gorm:"unique"` // Domain records must be unique, this avoids domain reuse across accounts.
AccountID string `gorm:"index"`
TargetCluster string // The proxy cluster this domain should be validated against
Type Type `gorm:"-"`
Validated bool
ID string `gorm:"unique;primaryKey;autoIncrement"`
Domain string `gorm:"unique"` // Domain records must be unique, this avoids domain reuse across accounts.
AccountID string `gorm:"index"`
TargetCluster string // The proxy cluster this domain should be validated against
Type Type `gorm:"-"`
Validated bool
ValidationExpiresAt *time.Time `gorm:"index"`
// SupportsCustomPorts is populated at query time for free domains from the
// proxy cluster capabilities. Not persisted.
SupportsCustomPorts *bool `gorm:"-"`
@@ -36,7 +45,12 @@ func (d *Domain) EventMeta() map[string]any {
}
}
// Copy returns a copy with an independent validation deadline.
func (d *Domain) Copy() *Domain {
dCopy := *d
if d.ValidationExpiresAt != nil {
expiresAt := *d.ValidationExpiresAt
dCopy.ValidationExpiresAt = &expiresAt
}
return &dCopy
}
@@ -0,0 +1,73 @@
package manager
import (
"context"
"time"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
"github.com/netbirdio/netbird/management/server/activity"
)
const (
validationCleanupInterval = 60 * time.Minute
validationCleanupBatch = 100
)
// RunValidationCleanup removes expired registrations on startup and hourly until cancellation.
func (m Manager) RunValidationCleanup(ctx context.Context) {
ticker := time.NewTicker(validationCleanupInterval)
defer ticker.Stop()
for {
m.cleanupExpiredDomains(ctx, time.Now().UTC())
select {
case <-ctx.Done():
return
case <-ticker.C:
}
}
}
func (m Manager) cleanupExpiredDomains(ctx context.Context, now time.Time) {
var afterID domain.ID
for ctx.Err() == nil {
domains, err := m.store.GetExpiredCustomDomains(ctx, now, afterID, validationCleanupBatch)
if err != nil {
if ctx.Err() == nil {
log.WithContext(ctx).WithError(err).Error("list expired custom domain registrations")
}
return
}
for _, d := range domains {
if ctx.Err() != nil {
return
}
m.deleteExpiredDomain(ctx, d, now)
afterID = domain.ID(d.ID)
}
if len(domains) < validationCleanupBatch {
return
}
}
}
func (m Manager) deleteExpiredDomain(ctx context.Context, d *domain.Domain, now time.Time) {
deleted, err := m.store.DeleteExpiredCustomDomain(ctx, d, now)
if err != nil {
if ctx.Err() == nil {
log.WithContext(ctx).WithFields(log.Fields{"accountID": d.AccountID, "domainID": d.ID}).
WithError(err).Warn("could not expire custom domain registration")
}
return
}
if !deleted {
return
}
meta := d.EventMeta()
if d.ValidationExpiresAt != nil {
meta["validation_expires_at"] = d.ValidationExpiresAt.UTC().Format(time.RFC3339)
}
m.accountManager.StoreEvent(ctx, activity.SystemInitiator, d.ID, d.AccountID,
activity.CustomDomainValidationExpired, meta)
}
@@ -0,0 +1,274 @@
package manager
import (
"context"
"fmt"
"sync"
"testing"
"testing/synctest"
"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"
"github.com/netbirdio/netbird/management/server/activity"
"github.com/netbirdio/netbird/management/server/mock_server"
nbstore "github.com/netbirdio/netbird/management/server/store"
)
func TestValidateDomain_ExpiredRegistration(t *testing.T) {
env := setupDomainTest(t)
ctx := context.Background()
d, err := env.manager.CreateDomain(ctx, accountA, accountAUser, "expired.example.com", testCluster)
require.NoError(t, err)
expiresAt := time.Now().Add(-time.Second)
db := env.store.(*nbstore.SqlStore).GetDB()
require.NoError(t, db.Model(&domain.Domain{}).Where("id = ?", d.ID).
Update("validation_expires_at", expiresAt).Error)
env.resolver.set("validation.expired.example.com", testCluster)
env.manager.ValidateDomain(ctx, accountA, accountAUser, d.ID)
stored := storedDomain(t, env.store, accountA, d.Domain)
require.NotNil(t, stored)
assert.False(t, stored.Validated, "an expired registration must not become usable before cleanup runs")
}
func TestCreateDomain_ValidationDeadline(t *testing.T) {
env := setupClockDomainTest(t)
synctest.Test(t, func(t *testing.T) {
ctx := context.Background()
createdAt := time.Now().UTC()
d, err := env.manager.CreateDomain(ctx, accountA, accountAUser, "pending.example.com", testCluster)
require.NoError(t, err)
require.NotNil(t, d.ValidationExpiresAt)
assert.Equal(t, createdAt.Add(48*time.Hour), *d.ValidationExpiresAt, "new registrations get 48 hours")
time.Sleep(time.Hour)
env.manager.ValidateDomain(ctx, accountA, accountAUser, d.ID)
stored := storedDomain(t, env.store, accountA, d.Domain)
require.NotNil(t, stored)
require.NotNil(t, stored.ValidationExpiresAt)
assert.WithinDuration(t, *d.ValidationExpiresAt, *stored.ValidationExpiresAt, 0, "failed validation must not extend the deadline")
})
}
func TestCleanupExpiredDomains_Boundaries(t *testing.T) {
env := setupDomainTest(t)
events := captureDomainEvents(env)
ctx := context.Background()
now := time.Now().UTC().Truncate(time.Second)
tests := []struct {
name string
expiresAt time.Time
validated bool
deleted bool
}{
{"expired", now.Add(-time.Second), false, true},
{"deadline", now, false, true},
{"pending", now.Add(time.Second), false, false},
{"validated", now.Add(-time.Hour), true, false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
d := createExpiringDomain(t, env, tt.name+".example.com", tt.expiresAt)
if tt.validated {
require.NoError(t, env.store.(*nbstore.SqlStore).GetDB().Model(d).Update("validated", true).Error)
}
env.manager.cleanupExpiredDomains(ctx, now)
stored := storedDomain(t, env.store, accountA, d.Domain)
if !tt.deleted {
assert.NotNil(t, stored, "pending and validated registrations must survive cleanup")
return
}
assert.Nil(t, stored, "expired unused registrations must be removed")
replacement, err := env.manager.CreateDomain(ctx, accountB, accountBUser, d.Domain, testCluster)
require.NoError(t, err)
assert.NotEqual(t, d.ID, replacement.ID, "the released name must receive a fresh registration")
assert.False(t, replacement.Validated, "the new account must validate its own registration")
})
}
got := events.get()
require.Len(t, got, 2, "only successful expiration deletions emit events")
for _, event := range got {
assert.Equal(t, activity.CustomDomainValidationExpired, event.Activity, "use the requested expiration event")
assert.Equal(t, activity.SystemInitiator, event.InitiatorID, "cleanup is attributed to the system")
assert.Equal(t, accountA, event.AccountID, "expiration belongs to the original account")
assert.NotEmpty(t, event.TargetID, "retain the deleted domain ID")
assert.NotEmpty(t, event.Meta["domain"], "retain the deleted domain name")
assert.NotEmpty(t, event.Meta["validation_expires_at"], "include the validation deadline")
}
}
func TestCleanupExpiredDomains_ContinuesPastProtectedBatch(t *testing.T) {
env := setupDomainTest(t)
ctx := context.Background()
now := time.Now().UTC()
for i := range validationCleanupBatch {
d := createExpiringDomain(t, env, fmt.Sprintf("protected-%d.example.com", i), now.Add(-time.Hour))
require.NoError(t, env.store.CreateService(ctx, &rpservice.Service{
ID: fmt.Sprintf("service-%d", i), AccountID: accountA, Domain: "app." + d.Domain,
}))
}
unprotected := createExpiringDomain(t, env, "unused.example.com", now.Add(-time.Hour))
env.manager.cleanupExpiredDomains(ctx, now)
assert.Nil(t, storedDomain(t, env.store, accountA, unprotected.Domain), "protected registrations must not starve later batches")
remaining, err := env.store.ListCustomDomains(ctx, accountA)
require.NoError(t, err)
assert.Len(t, remaining, validationCleanupBatch, "all registrations with dependent services must survive")
}
func TestCleanupExpiredDomains_ConcurrentWorkers(t *testing.T) {
env := setupDomainTest(t)
events := captureDomainEvents(env)
now := time.Now().UTC()
d := createExpiringDomain(t, env, "concurrent.example.com", now.Add(-time.Hour))
var workers sync.WaitGroup
for range 2 {
workers.Go(func() { env.manager.cleanupExpiredDomains(context.Background(), now) })
}
workers.Wait()
assert.Nil(t, storedDomain(t, env.store, accountA, d.Domain), "one worker must remove the expired registration")
assert.Len(t, events.get(), 1, "only the worker that deletes the row may emit the event")
}
func TestRunValidationCleanup_HourlyAndRestart(t *testing.T) {
env := setupClockDomainTest(t)
synctest.Test(t, func(t *testing.T) {
events := captureDomainEvents(env)
now := time.Now().UTC()
startup := createExpiringDomain(t, env, "startup.example.com", now.Add(-time.Hour))
hourly := createExpiringDomain(t, env, "hourly.example.com", now.Add(time.Minute))
ctx, cancel := context.WithCancel(context.Background())
done := make(chan struct{})
go func() {
defer close(done)
env.manager.RunValidationCleanup(ctx)
}()
synctest.Wait()
assert.Nil(t, storedDomain(t, env.store, accountA, startup.Domain), "startup must collect overdue registrations")
time.Sleep(59 * time.Minute)
synctest.Wait()
assert.NotNil(t, storedDomain(t, env.store, accountA, hourly.Domain), "cleanup must wait for the 60-minute interval")
time.Sleep(time.Minute)
synctest.Wait()
assert.Nil(t, storedDomain(t, env.store, accountA, hourly.Domain), "the hourly scan must collect expired registrations")
cancel()
<-done
offline := createExpiringDomain(t, env, "offline.example.com", time.Now().UTC().Add(time.Minute))
time.Sleep(2 * time.Hour)
assert.NotNil(t, storedDomain(t, env.store, accountA, offline.Domain), "a stopped worker must not continue deleting")
ctx, cancel = context.WithCancel(context.Background())
done = make(chan struct{})
go func() {
defer close(done)
env.manager.RunValidationCleanup(ctx)
}()
synctest.Wait()
assert.Nil(t, storedDomain(t, env.store, accountA, offline.Domain), "restart must use the persisted deadline")
cancel()
<-done
assert.Len(t, events.get(), 3, "each deletion should emit an expiration event")
})
}
type blockingDomainResolver struct {
started chan struct{}
release chan struct{}
}
func (r blockingDomainResolver) LookupCNAME(context.Context, string) (string, error) {
close(r.started)
<-r.release
return testCluster + ".", nil
}
func TestValidateDomain_DeadlinePassesDuringLookup(t *testing.T) {
for _, cleanup := range []bool{false, true} {
t.Run(fmt.Sprintf("cleanup=%t", cleanup), func(t *testing.T) {
env := setupClockDomainTest(t)
synctest.Test(t, func(t *testing.T) {
events := captureDomainEvents(env)
ctx := context.Background()
d, err := env.manager.CreateDomain(ctx, accountA, accountAUser, "late.example.com", testCluster)
require.NoError(t, err)
resolver := blockingDomainResolver{started: make(chan struct{}), release: make(chan struct{})}
env.manager.validator.Resolver = resolver
done := make(chan struct{})
go func() {
defer close(done)
env.manager.ValidateDomain(ctx, accountA, accountAUser, d.ID)
}()
<-resolver.started
time.Sleep(48 * time.Hour)
if cleanup {
env.manager.cleanupExpiredDomains(ctx, time.Now().UTC())
_, err = env.store.CreateCustomDomain(ctx, accountB, d.Domain, testCluster, false)
require.NoError(t, err)
}
close(resolver.release)
<-done
owner := accountA
if cleanup {
assert.Nil(t, storedDomain(t, env.store, accountA, d.Domain), "late validation must not restore the old claim")
owner = accountB
}
stored := storedDomain(t, env.store, owner, d.Domain)
require.NotNil(t, stored)
assert.False(t, stored.Validated, "late validation must not validate either claim")
for _, event := range events.get() {
assert.NotEqual(t, activity.DomainValidated, event.Activity, "a rejected write must not emit a validation event")
}
})
})
}
}
func setupClockDomainTest(t *testing.T) *domainTestEnv {
t.Helper()
// Network driver watchers cannot share cancellation channels across synctest bubbles.
// Store boundary and concurrency tests still exercise the selected database engine.
t.Setenv("NETBIRD_STORE_ENGINE", "sqlite")
return setupDomainTest(t)
}
func createExpiringDomain(t *testing.T, env *domainTestEnv, name string, expiresAt time.Time) *domain.Domain {
t.Helper()
d, err := env.store.CreateCustomDomain(context.Background(), accountA, name, testCluster, false)
require.NoError(t, err)
require.NoError(t, env.store.(*nbstore.SqlStore).GetDB().Model(d).Update("validation_expires_at", expiresAt).Error)
d.ValidationExpiresAt = &expiresAt
return d
}
type domainEvents struct {
mu sync.Mutex
events []*activity.Event
}
func captureDomainEvents(env *domainTestEnv) *domainEvents {
events := &domainEvents{}
env.manager.accountManager = &mock_server.MockAccountManager{
StoreEventFunc: func(_ context.Context, initiator, target, account string, code activity.ActivityDescriber, meta map[string]any) {
if code == activity.DomainAdded {
return
}
events.mu.Lock()
defer events.mu.Unlock()
events.events = append(events.events, &activity.Event{
InitiatorID: initiator, TargetID: target, AccountID: account,
Activity: code.(activity.Activity), Meta: meta,
})
},
}
return events
}
func (e *domainEvents) get() []*activity.Event {
e.mu.Lock()
defer e.mu.Unlock()
return append([]*activity.Event(nil), e.events...)
}
@@ -6,6 +6,7 @@ import (
"fmt"
"net"
"strings"
"time"
log "github.com/sirupsen/logrus"
@@ -15,6 +16,7 @@ import (
"github.com/netbirdio/netbird/management/server/activity"
nbstore "github.com/netbirdio/netbird/management/server/store"
"github.com/netbirdio/netbird/management/server/types"
nbdomain "github.com/netbirdio/netbird/shared/management/domain"
"github.com/netbirdio/netbird/shared/management/status"
)
@@ -29,6 +31,8 @@ type store interface {
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)
DeleteCustomDomain(ctx context.Context, accountID string, domainID string) 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)
}
type proxyManager interface {
@@ -93,12 +97,13 @@ func (m Manager) GetDomains(ctx context.Context, accountID, userID string) ([]*d
// Add custom domains.
for _, d := range domains {
cd := &domain.Domain{
ID: d.ID,
Domain: d.Domain,
AccountID: accountID,
TargetCluster: d.TargetCluster,
Type: domain.TypeCustom,
Validated: d.Validated,
ID: d.ID,
Domain: d.Domain,
AccountID: accountID,
TargetCluster: d.TargetCluster,
Type: domain.TypeCustom,
Validated: d.Validated,
ValidationExpiresAt: d.ValidationExpiresAt,
}
if d.TargetCluster != "" {
cd.SupportsCustomPorts = m.proxyManager.ClusterSupportsCustomPorts(ctx, d.TargetCluster)
@@ -113,7 +118,17 @@ func (m Manager) GetDomains(ctx context.Context, accountID, userID string) ([]*d
return ret, nil
}
// CreateDomain registers a normalized custom domain and attempts DNS validation.
func (m Manager) CreateDomain(ctx context.Context, accountID, userID, domainName, targetCluster string) (*domain.Domain, error) {
parsed, err := nbdomain.FromString(strings.TrimSuffix(domainName, "."))
if err != nil {
return nil, status.Errorf(status.InvalidArgument, "invalid domain: %v", err)
}
domainName = parsed.PunycodeString()
if !nbdomain.IsValidDomainNoWildcard(domainName) {
return nil, status.Errorf(status.InvalidArgument, "invalid domain format")
}
// Verify the target cluster is in the available clusters for this account
allowList, err := m.getClusterAllowList(ctx, accountID)
if err != nil {
@@ -197,6 +212,14 @@ func (m Manager) ValidateDomain(ctx context.Context, accountID, userID, domainID
}).WithError(err).Error("get custom domain from store")
return
}
if d.Validated {
return
}
if d.ValidationExpiresAt == nil || !time.Now().Before(*d.ValidationExpiresAt) {
log.WithFields(log.Fields{"accountID": accountID, "domainID": domainID}).
Debug("custom domain validation window has expired")
return
}
// Validate only against the domain's target cluster
targetCluster := d.TargetCluster
@@ -217,20 +240,21 @@ func (m Manager) ValidateDomain(ctx context.Context, accountID, userID, domainID
}).Info("validating domain against target cluster")
if m.validator.IsValid(context.Background(), d.Domain, []string{targetCluster}) {
log.WithFields(log.Fields{
"accountID": accountID,
"domainID": domainID,
"domain": d.Domain,
}).Info("domain validated successfully")
d.Validated = true
if _, err := m.store.UpdateCustomDomain(context.Background(), accountID, d); err != nil {
log.WithFields(log.Fields{
entry := log.WithFields(log.Fields{
"accountID": accountID,
"domainID": domainID,
"domain": d.Domain,
}).WithError(err).Error("update custom domain in store")
}).WithError(err)
if sErr, ok := status.FromError(err); ok && sErr.Type() == status.PreconditionFailed {
entry.Debug("custom domain registration is no longer pending validation")
return
}
entry.Error("update custom domain in store")
return
}
log.WithFields(log.Fields{"accountID": accountID, "domainID": domainID}).
Info("custom domain validated successfully")
m.accountManager.StoreEvent(context.Background(), userID, domainID, accountID, activity.DomainValidated, d.EventMeta())
} else {
@@ -268,11 +268,8 @@ func TestStore_DuplicateDomainRejectedByIndexAsConflict(t *testing.T) {
assert.Equal(t, status.AlreadyExists, sErr.Type(), "a lost race is a 409, not a 500")
}
// Validation runs asynchronously, so it can finish after the domain was
// deleted and then write a stale row back. gorm's Save falls back to an insert
// when an update affects no rows, which would resurrect the domain as
// validated; UpdateCustomDomain avoids that by selecting explicit columns.
// This pins that behaviour, since dropping the Select would reintroduce it.
// A validation finishing after deletion must reject the stale write, without
// restoring the registration or reporting successful validation.
func TestUpdateCustomDomain_DoesNotResurrectDeletedDomain(t *testing.T) {
ctx := context.Background()
env := setupDomainTest(t)
@@ -287,11 +284,9 @@ func TestUpdateCustomDomain_DoesNotResurrectDeletedDomain(t *testing.T) {
require.Nil(t, storedDomain(t, env.store, accountA, "racy.example.com"), "the domain should be gone")
// What an in-flight validation would write once its CNAME check succeeded.
// The write has to succeed for the assertion below to mean anything: a
// rejected write would leave the domain absent for the wrong reason.
stale.Validated = true
_, err = env.store.UpdateCustomDomain(ctx, accountA, stale)
require.NoError(t, err, "the update itself must succeed, so absence is not just a failed write")
require.Error(t, err, "a deleted registration must reject a late validation")
assert.Nil(t, storedDomain(t, env.store, accountA, "racy.example.com"),
"a late validation write must not recreate a deleted domain")
@@ -4,6 +4,7 @@ import (
"context"
"errors"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -208,6 +209,14 @@ func (s *stubStore) DeleteCustomDomain(context.Context, string, string) error {
panic("not used in allow-list tests")
}
func (s *stubStore) GetExpiredCustomDomains(context.Context, time.Time, domain.ID, int) ([]*domain.Domain, error) {
panic("not used in allow-list tests")
}
func (s *stubStore) DeleteExpiredCustomDomain(context.Context, *domain.Domain, time.Time) (bool, error) {
panic("not used in allow-list tests")
}
// TestGetClusterAllowList_DedicatedGatewayAddressExcluded pins invariant (B)'s
// chokepoint: a self-addressed settings pin reserves the account's gateway
// address, so it is dropped from the allow list — which, because the
@@ -0,0 +1,83 @@
package manager
import (
"context"
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/shared/management/status"
)
func TestCreateDomain_NormalizesName(t *testing.T) {
for _, tt := range []struct {
name string
input string
canonical string
}{
{"mixed case", "Apps.Example.COM", "apps.example.com"},
{"unicode", "münchen.example.com", "xn--mnchen-3ya.example.com"},
{"trailing dot", "apps.example.com.", "apps.example.com"},
{"underscore", "My_App.example.com", "my_app.example.com"},
} {
t.Run(tt.name, func(t *testing.T) {
ctx := context.Background()
env := setupDomainTest(t)
env.resolver.set("validation."+tt.canonical, testCluster)
created, err := env.manager.CreateDomain(ctx, accountA, accountAUser, tt.input, testCluster)
require.NoError(t, err)
assert.Equal(t, tt.canonical, created.Domain, "the response must use the normalized name")
assert.True(t, created.Validated, "the CNAME lookup must use the normalized name")
stored, err := env.store.GetCustomDomain(ctx, accountA, created.ID)
require.NoError(t, err)
assert.Equal(t, tt.canonical, stored.Domain, "the database must retain the normalized name")
_, err = env.manager.CreateDomain(ctx, accountB, accountBUser, tt.canonical, testCluster)
require.Error(t, err)
sErr, ok := status.FromError(err)
require.True(t, ok, "an equivalent name must return a typed conflict")
assert.Equal(t, status.AlreadyExists, sErr.Type(), "normalization must precede the availability check")
})
}
}
func TestCreateDomain_NormalizedNameCanValidateLater(t *testing.T) {
ctx := context.Background()
env := setupDomainTest(t)
created, err := env.manager.CreateDomain(ctx, accountA, accountAUser, "Apps.Example.COM.", testCluster)
require.NoError(t, err)
require.False(t, created.Validated, "a missing CNAME must leave the normalized registration pending")
env.resolver.set("validation.apps.example.com", testCluster)
env.manager.ValidateDomain(ctx, accountA, accountAUser, created.ID)
stored, err := env.store.GetCustomDomain(ctx, accountA, created.ID)
require.NoError(t, err)
assert.Equal(t, "apps.example.com", stored.Domain, "retrying validation must retain the normalized name")
assert.True(t, stored.Validated, "later validation must look up the normalized name")
}
func TestCreateDomain_RejectsInvalidName(t *testing.T) {
ctx := context.Background()
env := setupDomainTest(t)
for _, name := range []string{
"", ".", "app..example.com", "app.example.com..", "-app.example.com",
"app%.example.com", "app!.example.com", "*.example.com", "app example.com",
"https://example.com", strings.Repeat("a", 64) + ".example.com",
} {
t.Run(name, func(t *testing.T) {
// A matching DNS response must not make a malformed name acceptable.
env.resolver.set("validation."+name, testCluster)
_, err := env.manager.CreateDomain(ctx, accountA, accountAUser, name, testCluster)
require.Error(t, err)
sErr, ok := status.FromError(err)
require.True(t, ok, "invalid names must return a typed client error")
assert.Equal(t, status.InvalidArgument, sErr.Type(), "malformed names must be rejected before storage")
})
}
stored, err := env.store.ListCustomDomains(ctx, accountA)
require.NoError(t, err)
assert.Empty(t, stored, "invalid registration attempts must not reserve any names")
}
+36 -20
View File
@@ -31,6 +31,7 @@ import (
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
networkmapdbfactory "github.com/netbirdio/netbird/management/internals/network_map_db/factory"
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
"github.com/netbirdio/netbird/management/server/activity"
activitystore "github.com/netbirdio/netbird/management/server/activity/store"
@@ -111,7 +112,8 @@ func (s *BaseServer) NetworkMapStore() *networkmapdb.NetworkMapDBStoreImpl {
s.Config.StoreConfig.Engine,
s.Config.Datadir,
s.IntegratedValidator(),
s.SettingsManager())
s.SettingsManager(),
)
// networkmap db store supports postgres and sqlite backends only
// for other backends a fallback is used, so NotSupportedStoreEngineError
// is not a fatal error
@@ -180,24 +182,7 @@ func (s *BaseServer) RateLimiter() *middleware.APIRateLimiter {
func (s *BaseServer) GRPCServer() *grpc.Server {
return Create(s, func() *grpc.Server {
trustedPeers := s.Config.ReverseProxy.TrustedPeers
defaultTrustedPeers := []netip.Prefix{netip.MustParsePrefix("0.0.0.0/0"), netip.MustParsePrefix("::/0")}
if len(trustedPeers) == 0 || slices.Equal[[]netip.Prefix](trustedPeers, defaultTrustedPeers) {
log.WithContext(context.Background()).Warn("TrustedPeers are configured to default value '0.0.0.0/0', '::/0'. This allows connection IP spoofing.")
trustedPeers = defaultTrustedPeers
}
trustedHTTPProxies := s.Config.ReverseProxy.TrustedHTTPProxies
trustedProxiesCount := s.Config.ReverseProxy.TrustedHTTPProxiesCount
if len(trustedHTTPProxies) > 0 && trustedProxiesCount > 0 {
log.WithContext(context.Background()).Warn("TrustedHTTPProxies and TrustedHTTPProxiesCount both are configured. " +
"This is not recommended way to extract X-Forwarded-For. Consider using one of these options.")
}
realipOpts := []realip.Option{
realip.WithTrustedPeers(trustedPeers),
realip.WithTrustedProxies(trustedHTTPProxies),
realip.WithTrustedProxiesCount(trustedProxiesCount),
realip.WithHeaders([]string{realip.XForwardedFor, realip.XRealIp}),
}
realipOpts := realIPOptions(s.Config.ReverseProxy)
proxyUnary, proxyStream, proxyAuthClose := nbgrpc.NewProxyAuthInterceptors(s.Store())
s.proxyAuthClose = proxyAuthClose
gRPCOpts := []grpc.ServerOption{
@@ -333,7 +318,7 @@ func (s *BaseServer) AccessLogsManager() accesslogs.Manager {
})
}
func loadTLSConfig(certFile string, certKey string) (*tls.Config, error) {
func loadTLSConfig(certFile, certKey string) (*tls.Config, error) {
// Load server's certificate and private key
serverCert, err := tls.LoadX509KeyPair(certFile, certKey)
if err != nil {
@@ -380,3 +365,34 @@ func streamInterceptor(
wrapped.WrappedContext = context.WithValue(ctx, nbContext.RequestIDKey, reqID)
return handler(srv, wrapped)
}
// realIPOptions builds the real-IP middleware options from the reverse proxy config.
//
// TrustedPeers controls which transport peers are allowed to supply forwarded-IP
// headers. If empty, forwarded headers are ignored and the transport peer address
// is used directly. Operators terminating connections at a reverse proxy should
// configure TrustedPeers with that proxy's address or network.
//
// Only X-Forwarded-For is trusted. X-Real-IP contains a single client-supplied
// address with no proxy chain to validate, and none of the reverse proxies we ship
// use it on the gRPC path.
func realIPOptions(cfg nbconfig.ReverseProxy) []realip.Option {
if idx := slices.IndexFunc(cfg.TrustedPeers, func(p netip.Prefix) bool { return p.Bits() == 0 }); idx >= 0 {
log.WithContext(context.Background()).Warnf("TrustedPeers contains the default route %s, which trusts "+
"X-Forwarded-For from every client and allows connection IP spoofing. Set TrustedPeers to the address "+
"of your reverse proxy, or leave it empty to use the connection's source address.", cfg.TrustedPeers[idx])
}
if cfg.TrustedHTTPProxiesCount > 0 {
log.WithContext(context.Background()).Warn(
"TrustedHTTPProxiesCount skips X-Forwarded-For entries by position before TrustedHTTPProxies filters by address. " +
"An incorrect count may skip the real client IP and produce an incorrect source address.",
)
}
return []realip.Option{
realip.WithTrustedPeers(cfg.TrustedPeers),
realip.WithTrustedProxies(cfg.TrustedHTTPProxies),
realip.WithTrustedProxiesCount(cfg.TrustedHTTPProxiesCount),
realip.WithHeaders([]string{realip.XForwardedFor}),
}
}
+171
View File
@@ -0,0 +1,171 @@
package server
import (
"context"
"io"
"net"
"net/netip"
"testing"
"time"
"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/realip"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/metadata"
"google.golang.org/protobuf/types/known/emptypb"
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
)
const (
realIPProbeMethod = "/netbird.test.RealIPProbe/Probe"
realIPProbeStreamMethod = "/netbird.test.RealIPProbe/ProbeStream"
)
// realIPProbe records the real IP the middleware derived for each call.
type realIPProbe struct {
got chan string
}
func (p *realIPProbe) record(ctx context.Context) {
addr, _ := realip.FromContext(ctx)
p.got <- addr.String()
}
func (p *realIPProbe) wait(t *testing.T) string {
t.Helper()
select {
case got := <-p.got:
return got
case <-time.After(5 * time.Second):
t.Fatal("timed out waiting for probe")
return ""
}
}
func startProbeServer(t *testing.T, cfg nbconfig.ReverseProxy) (*grpc.ClientConn, *realIPProbe) {
t.Helper()
listener, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
probe := &realIPProbe{got: make(chan string, 1)}
opts := realIPOptions(cfg)
srv := grpc.NewServer(
grpc.ChainUnaryInterceptor(realip.UnaryServerInterceptorOpts(opts...)),
grpc.ChainStreamInterceptor(realip.StreamServerInterceptorOpts(opts...)),
)
srv.RegisterService(&grpc.ServiceDesc{
ServiceName: "netbird.test.RealIPProbe",
HandlerType: (*any)(nil),
Methods: []grpc.MethodDesc{{
MethodName: "Probe",
Handler: func(_ any, ctx context.Context, dec func(any) error, interceptor grpc.UnaryServerInterceptor) (any, error) {
req := new(emptypb.Empty)
if err := dec(req); err != nil {
return nil, err
}
handler := func(ctx context.Context, _ any) (any, error) {
probe.record(ctx)
return &emptypb.Empty{}, nil
}
if interceptor == nil {
return handler(ctx, req)
}
return interceptor(ctx, req, &grpc.UnaryServerInfo{FullMethod: realIPProbeMethod}, handler)
},
}},
Streams: []grpc.StreamDesc{{
StreamName: "ProbeStream",
ServerStreams: true,
Handler: func(_ any, stream grpc.ServerStream) error {
probe.record(stream.Context())
return nil
},
}},
}, probe)
go func() { _ = srv.Serve(listener) }()
t.Cleanup(srv.Stop)
conn, err := grpc.NewClient(listener.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
require.NoError(t, err)
t.Cleanup(func() { _ = conn.Close() })
return conn, probe
}
func callUnary(t *testing.T, conn *grpc.ClientConn, probe *realIPProbe, kv ...string) string {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
ctx = metadata.AppendToOutgoingContext(ctx, kv...)
require.NoError(t, conn.Invoke(ctx, realIPProbeMethod, &emptypb.Empty{}, &emptypb.Empty{}))
return probe.wait(t)
}
func callStream(t *testing.T, conn *grpc.ClientConn, probe *realIPProbe, kv ...string) string {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
ctx = metadata.AppendToOutgoingContext(ctx, kv...)
desc := &grpc.StreamDesc{StreamName: "ProbeStream", ServerStreams: true}
stream, err := conn.NewStream(ctx, desc, realIPProbeStreamMethod)
require.NoError(t, err)
require.NoError(t, stream.CloseSend())
require.ErrorIs(t, stream.RecvMsg(&emptypb.Empty{}), io.EOF)
return probe.wait(t)
}
func assertRealIP(t *testing.T, cfg nbconfig.ReverseProxy, want string, kv ...string) {
t.Helper()
conn, probe := startProbeServer(t, cfg)
t.Run("unary", func(t *testing.T) {
assert.Equal(t, want, callUnary(t, conn, probe, kv...))
})
t.Run("stream", func(t *testing.T) {
assert.Equal(t, want, callStream(t, conn, probe, kv...))
})
}
func TestRealIPDefaultIgnoresClientForwardedHeaders(t *testing.T) {
assertRealIP(t, nbconfig.ReverseProxy{}, "127.0.0.1",
realip.XForwardedFor, "203.0.113.44",
realip.XRealIp, "203.0.113.44",
)
}
func TestRealIPUntrustedPeerIgnoresForwardedHeaders(t *testing.T) {
cfg := nbconfig.ReverseProxy{TrustedPeers: []netip.Prefix{netip.MustParsePrefix("10.9.8.7/32")}}
assertRealIP(t, cfg, "127.0.0.1",
realip.XForwardedFor, "203.0.113.44",
realip.XRealIp, "203.0.113.44",
)
}
func TestRealIPTrustedPeerHonoursForwardedHeaders(t *testing.T) {
cfg := nbconfig.ReverseProxy{TrustedPeers: []netip.Prefix{netip.MustParsePrefix("127.0.0.1/32")}}
assertRealIP(t, cfg, "203.0.113.44",
realip.XForwardedFor, "203.0.113.44",
realip.XRealIp, "203.0.113.44",
)
}
func TestRealIPIgnoresXRealIPWhenProxyCountIsSet(t *testing.T) {
cfg := nbconfig.ReverseProxy{
TrustedPeers: []netip.Prefix{netip.MustParsePrefix("127.0.0.1/32")},
TrustedHTTPProxiesCount: 1,
}
assertRealIP(t, cfg, "127.0.0.1", realip.XRealIp, "203.0.113.44")
}
+55 -18
View File
@@ -66,7 +66,8 @@ type BaseServer struct {
disableLegacyManagementPort bool
autoResolveDomains bool
proxyAuthClose func()
proxyAuthClose func()
domainCleanupStop func()
// grpcExtensions holds additional gRPC services, interceptors, and shutdown
// hooks registered by external modules via RegisterGRPCExtension. Populated
@@ -74,6 +75,7 @@ type BaseServer struct {
grpcExtensions []GRPCExtension
listener net.Listener
tlsConfig *tls.Config
certManager *autocert.Manager
update *version.Update
@@ -94,6 +96,7 @@ type Config struct {
DisableGeoliteUpdate bool
UserDeleteFromIDPEnabled bool
AutoResolveDomains bool
TLSConfig *tls.Config
}
// NewServer initializes and configures a new Server instance
@@ -110,6 +113,7 @@ func NewServer(cfg *Config) *BaseServer {
disableLegacyManagementPort: cfg.DisableLegacyManagementPort,
mgmtMetricsPort: cfg.MgmtMetricsPort,
autoResolveDomains: cfg.AutoResolveDomains,
tlsConfig: cfg.TLSConfig,
}
s.container[ContainerKeyBaseServer] = s
@@ -139,21 +143,9 @@ func (s *BaseServer) Start(ctx context.Context) error {
}
s.EphemeralManager().LoadInitialPeers(srvCtx)
var tlsConfig *tls.Config
tlsEnabled := false
if s.Config.HttpConfig.LetsEncryptDomain != "" {
s.certManager, err = encryption.CreateCertManager(s.Config.Datadir, s.Config.HttpConfig.LetsEncryptDomain)
if err != nil {
return fmt.Errorf("failed creating LetsEncrypt cert manager: %v", err)
}
tlsEnabled = true
} else if s.Config.HttpConfig.CertFile != "" && s.Config.HttpConfig.CertKey != "" {
tlsConfig, err = loadTLSConfig(s.Config.HttpConfig.CertFile, s.Config.HttpConfig.CertKey)
if err != nil {
log.WithContext(srvCtx).Errorf("cannot load TLS credentials: %v", err)
return err
}
tlsEnabled = true
tlsEnabled, err := s.setupTLS(srvCtx)
if err != nil {
return err
}
installationID, err := getInstallationID(srvCtx, s.Store())
@@ -215,8 +207,8 @@ func (s *BaseServer) Start(ctx context.Context) error {
log.WithContext(ctx).Infof("running HTTP server (LetsEncrypt challenge handler): %s", cml.Addr().String())
s.serveHTTP(ctx, cml, s.certManager.HTTPHandler(nil))
}
case tlsConfig != nil:
s.listener, err = tls.Listen("tcp", fmt.Sprintf(":%d", s.mgmtPort), tlsConfig)
case s.tlsConfig != nil:
s.listener, err = tls.Listen("tcp", fmt.Sprintf(":%d", s.mgmtPort), s.tlsConfig)
if err != nil {
return fmt.Errorf("failed creating TLS listener on port %d: %v", s.mgmtPort, err)
}
@@ -236,14 +228,59 @@ func (s *BaseServer) Start(ctx context.Context) error {
s.update.SetOnUpdateListener(func() {
log.WithContext(ctx).Infof("your management version, \"%s\", is outdated, a new management version is available. Learn more here: https://github.com/netbirdio/netbird/releases", version.NetbirdVersion())
})
s.startDomainCleanup(srvCtx)
return nil
}
func (s *BaseServer) startDomainCleanup(ctx context.Context) {
if s.domainCleanupStop != nil {
return
}
mgr := s.ReverseProxyDomainManager()
ctx, cancel := context.WithCancel(ctx)
done := make(chan struct{})
s.domainCleanupStop = func() {
cancel()
<-done
}
go func() {
defer close(done)
mgr.RunValidationCleanup(ctx)
}()
}
// setupTLS resolves the listener's TLS source: an injected config wins over the HttpConfig certificate settings
func (s *BaseServer) setupTLS(ctx context.Context) (bool, error) {
switch {
case s.tlsConfig != nil:
return true, nil
case s.Config.HttpConfig.LetsEncryptDomain != "":
certManager, err := encryption.CreateCertManager(s.Config.Datadir, s.Config.HttpConfig.LetsEncryptDomain)
if err != nil {
return false, fmt.Errorf("failed creating LetsEncrypt cert manager: %v", err)
}
s.certManager = certManager
return true, nil
case s.Config.HttpConfig.CertFile != "" && s.Config.HttpConfig.CertKey != "":
tlsConfig, err := loadTLSConfig(s.Config.HttpConfig.CertFile, s.Config.HttpConfig.CertKey)
if err != nil {
log.WithContext(ctx).Errorf("cannot load TLS credentials: %v", err)
return false, err
}
s.tlsConfig = tlsConfig
return true, nil
default:
return false, nil
}
}
// Stop attempts a graceful shutdown, waiting up to 5 seconds for active connections to finish
func (s *BaseServer) Stop() error {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if s.domainCleanupStop != nil {
s.domainCleanupStop()
}
s.IntegratedValidator().Stop(ctx)
if s.GeoLocationManager() != nil {
@@ -0,0 +1,135 @@
package grpc
import (
"context"
"time"
"github.com/netbirdio/netbird/encryption"
"github.com/netbirdio/netbird/management/internals/controllers/network_map"
"github.com/netbirdio/netbird/management/server/telemetry"
"github.com/netbirdio/netbird/shared/management/proto"
log "github.com/sirupsen/logrus"
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)
func PeerUpdateHandlerFactory(
peerKey wgtypes.Key,
updates chan *network_map.UpdateMessage,
secretsManager SecretsManager,
srv proto.ManagementService_SyncServer,
cleanupfunc func()) *PeerUpdateHandler {
return &PeerUpdateHandler{
peerKey: peerKey,
updates: updates,
secretsManager: secretsManager,
srv: srv,
encrypter: encryption.DefaultEncrypter{},
debouncer: NewUpdateDebouncer(1000 * time.Millisecond),
cleanupFunc: cleanupfunc,
}
}
// PeerUpdateHandler sends updates to the connected peer until the updates channel is closed.
// It implements a backpressure mechanism that sends the first update immediately,
// then debounces subsequent rapid updates, ensuring only the latest update is sent
// after a quiet period.
type PeerUpdateHandler struct {
peerKey wgtypes.Key
updates chan *network_map.UpdateMessage
appMetrics telemetry.AppMetrics
secretsManager SecretsManager
srv syncSender
encrypter encryption.Encrypter
debouncer Debouncer
cleanupFunc func()
}
func (pu *PeerUpdateHandler) WithMetrics(appMetrics telemetry.AppMetrics) *PeerUpdateHandler {
pu.appMetrics = appMetrics
return pu
}
//go:generate go tool mockgen -source=./peer_update_handler.go -destination=./sync_sender_mock.go -package=grpc
type syncSender interface {
Send(*proto.EncryptedMessage) error
Context() context.Context
}
func (pu *PeerUpdateHandler) HandleUpdates(ctx context.Context) error {
log.WithContext(ctx).Tracef("starting to handle updates for peer %s", pu.peerKey.String())
defer pu.debouncer.Stop()
for {
select {
// condition when there are some updates
// todo set the updates channel size to 1
case update, open := <-pu.updates:
if pu.appMetrics != nil {
pu.appMetrics.GRPCMetrics().UpdateChannelQueueLength(len(pu.updates) + 1)
}
if !open {
log.WithContext(ctx).Debugf("updates channel for peer %s was closed", pu.peerKey.String())
pu.cleanupFunc()
return nil
}
log.WithContext(ctx).Tracef("received an update for peer %s", pu.peerKey.String())
if pu.debouncer.ProcessUpdate(update) {
// Send immediately (first update or after quiet period)
if err := pu.SendUpdate(ctx, update); err != nil {
log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", pu.peerKey.String(), err)
return err
}
}
// Timer expired - quiet period reached, send pending updates if any
case <-pu.debouncer.TimerChannel():
pendingUpdates := pu.debouncer.GetPendingUpdates()
if len(pendingUpdates) == 0 {
continue
}
log.WithContext(ctx).Debugf("sending %d debounced update(s) for peer %s", len(pendingUpdates), pu.peerKey.String())
for _, pendingUpdate := range pendingUpdates {
if err := pu.SendUpdate(ctx, pendingUpdate); err != nil {
log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", pu.peerKey.String(), err)
return err
}
}
// condition when client <-> server connection has been terminated
case <-pu.srv.Context().Done():
// happens when connection drops, e.g. client disconnects
log.WithContext(ctx).Debugf("stream of peer %s has been closed", pu.peerKey.String())
pu.cleanupFunc()
return pu.srv.Context().Err()
}
}
}
func (pu *PeerUpdateHandler) SendUpdate(ctx context.Context, update *network_map.UpdateMessage) error {
key, err := pu.secretsManager.GetWGKey()
if err != nil {
pu.cleanupFunc()
return status.Errorf(codes.Internal, "failed processing update message")
}
encryptedResp, err := pu.encrypter.EncryptMessage(pu.peerKey, key, update.Update)
if err != nil {
pu.cleanupFunc()
return status.Errorf(codes.Internal, "failed processing update message")
}
err = pu.srv.Send(&proto.EncryptedMessage{
WgPubKey: key.PublicKey().String(),
Body: encryptedResp,
})
if err != nil {
pu.cleanupFunc()
return status.Errorf(codes.Internal, "failed sending update message")
}
log.WithContext(ctx).Tracef("sent an update to peer %s", pu.peerKey.String())
return nil
}
@@ -0,0 +1,155 @@
package grpc
import (
"context"
"fmt"
"sync"
"testing"
"time"
pb "github.com/golang/protobuf/proto" //nolint
"github.com/netbirdio/netbird/management/internals/controllers/network_map"
"github.com/netbirdio/netbird/shared/management/proto"
"github.com/stretchr/testify/assert"
"go.uber.org/mock/gomock"
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
)
func TestSendPeerUpdates_FirstUpdate(t *testing.T) {
ctrl := gomock.NewController(t)
secretsManager := NewMockSecretsManager(ctrl)
updateDebouncer := NewMockDebouncer(ctrl)
syncSender := NewMocksyncSender(ctrl)
pu := PeerUpdateHandler{
peerKey: mustGenerateKey(t),
updates: make(chan *network_map.UpdateMessage),
secretsManager: secretsManager,
encrypter: testEncrypter{},
debouncer: updateDebouncer,
srv: syncSender,
cleanupFunc: func() {},
}
msg := network_map.UpdateMessage{
Update: &proto.SyncResponse{Version: 1},
}
timeCh := make(chan time.Time)
srvCtx := context.TODO()
srvKey := mustGenerateKey(t)
// mock a first update, should send it right away
updateDebouncer.EXPECT().ProcessUpdate(gomock.Eq(&msg)).Return(true)
updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh)
syncSender.EXPECT().Context().AnyTimes().Return(srvCtx)
secretsManager.EXPECT().GetWGKey().Return(srvKey, nil)
syncSender.EXPECT().Send(pbMatcher{x: &proto.EncryptedMessage{WgPubKey: srvKey.PublicKey().String(), Body: mustMarshal(t, &msg)}})
updateDebouncer.EXPECT().Stop()
var wg sync.WaitGroup
wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck
pu.updates <- &msg
close(pu.updates)
wg.Wait()
}
func TestSendPeerUpdates_TimerUpdate(t *testing.T) {
ctrl := gomock.NewController(t)
secretsManager := NewMockSecretsManager(ctrl)
updateDebouncer := NewMockDebouncer(ctrl)
syncSender := NewMocksyncSender(ctrl)
pu := PeerUpdateHandler{
peerKey: mustGenerateKey(t),
updates: make(chan *network_map.UpdateMessage),
secretsManager: secretsManager,
encrypter: testEncrypter{},
debouncer: updateDebouncer,
srv: syncSender,
cleanupFunc: func() {},
}
msg := network_map.UpdateMessage{
Update: &proto.SyncResponse{Version: 1},
}
timeCh := make(chan time.Time)
srvCtx := context.TODO()
srvKey := mustGenerateKey(t)
updateDebouncer.EXPECT().GetPendingUpdates().Return([]*network_map.UpdateMessage{&msg})
updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh)
syncSender.EXPECT().Context().AnyTimes().Return(srvCtx)
secretsManager.EXPECT().GetWGKey().Return(srvKey, nil)
syncSender.EXPECT().Send(pbMatcher{x: &proto.EncryptedMessage{WgPubKey: srvKey.PublicKey().String(), Body: mustMarshal(t, &msg)}})
updateDebouncer.EXPECT().Stop()
var wg sync.WaitGroup
wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck
timeCh <- time.Now()
close(pu.updates)
wg.Wait()
}
func TestSendPeerUpdates_ServerContextDone(t *testing.T) {
ctrl := gomock.NewController(t)
secretsManager := NewMockSecretsManager(ctrl)
updateDebouncer := NewMockDebouncer(ctrl)
syncSender := NewMocksyncSender(ctrl)
pu := PeerUpdateHandler{
peerKey: mustGenerateKey(t),
updates: make(chan *network_map.UpdateMessage),
secretsManager: secretsManager,
encrypter: testEncrypter{},
debouncer: updateDebouncer,
srv: syncSender,
cleanupFunc: func() {},
}
timeCh := make(chan time.Time)
srvCtx, cancel := context.WithCancel(context.TODO())
updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh)
syncSender.EXPECT().Context().AnyTimes().Return(srvCtx)
updateDebouncer.EXPECT().Stop()
var wg sync.WaitGroup
wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck
cancel()
wg.Wait()
}
func mustGenerateKey(t *testing.T) wgtypes.Key {
t.Helper()
k, err := wgtypes.GenerateKey()
assert.NoError(t, err)
return k
}
func mustMarshal(t *testing.T, msg *network_map.UpdateMessage) []byte {
t.Helper()
r, err := pb.Marshal(msg.Update)
assert.NoError(t, err)
return r
}
type testEncrypter struct{}
func (testEncrypter) EncryptMessage(remotePubKey wgtypes.Key, ourPrivateKey wgtypes.Key, message pb.Message) ([]byte, error) {
return pb.Marshal(message)
}
type pbMatcher struct {
x pb.Message
}
func (pbm pbMatcher) Matches(x any) bool {
msg, ok := x.(pb.Message)
if !ok {
return false
}
return pb.Equal(pbm.x, msg)
}
func (pbm pbMatcher) String() string {
return fmt.Sprintf("is equal to %s (%T)", pbm.x, pbm.x)
}
+2 -86
View File
@@ -337,7 +337,8 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
s.syncSem.Add(-1)
return s.handleUpdates(ctx, accountID, peerKey, peer, updates, srv, syncStart)
return PeerUpdateHandlerFactory(peerKey, updates, s.secretsManager, srv, func() { s.cancelPeerRoutines(ctx, accountID, peer, syncStart) }).
WithMetrics(s.appMetrics).HandleUpdates(ctx)
}
func (s *Server) handleHandshake(ctx context.Context, srv proto.ManagementService_JobServer) (wgtypes.Key, error) {
@@ -404,91 +405,6 @@ func (s *Server) sendJobsLoop(ctx context.Context, accountID string, peerKey wgt
}
}
// handleUpdates sends updates to the connected peer until the updates channel is closed.
// It implements a backpressure mechanism that sends the first update immediately,
// then debounces subsequent rapid updates, ensuring only the latest update is sent
// after a quiet period.
func (s *Server) handleUpdates(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, updates chan *network_map.UpdateMessage, srv proto.ManagementService_SyncServer, streamStartTime time.Time) error {
log.WithContext(ctx).Tracef("starting to handle updates for peer %s", peerKey.String())
// Create a debouncer for this peer connection
debouncer := NewUpdateDebouncer(1000 * time.Millisecond)
defer debouncer.Stop()
for {
select {
// condition when there are some updates
// todo set the updates channel size to 1
case update, open := <-updates:
if s.appMetrics != nil {
s.appMetrics.GRPCMetrics().UpdateChannelQueueLength(len(updates) + 1)
}
if !open {
log.WithContext(ctx).Debugf("updates channel for peer %s was closed", peerKey.String())
s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
return nil
}
log.WithContext(ctx).Tracef("received an update for peer %s", peerKey.String())
if debouncer.ProcessUpdate(update) {
// Send immediately (first update or after quiet period)
if err := s.sendUpdate(ctx, accountID, peerKey, peer, update, srv, streamStartTime); err != nil {
log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", peerKey.String(), err)
return err
}
}
// Timer expired - quiet period reached, send pending updates if any
case <-debouncer.TimerChannel():
pendingUpdates := debouncer.GetPendingUpdates()
if len(pendingUpdates) == 0 {
continue
}
log.WithContext(ctx).Debugf("sending %d debounced update(s) for peer %s", len(pendingUpdates), peerKey.String())
for _, pendingUpdate := range pendingUpdates {
if err := s.sendUpdate(ctx, accountID, peerKey, peer, pendingUpdate, srv, streamStartTime); err != nil {
log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", peerKey.String(), err)
return err
}
}
// condition when client <-> server connection has been terminated
case <-srv.Context().Done():
// happens when connection drops, e.g. client disconnects
log.WithContext(ctx).Debugf("stream of peer %s has been closed", peerKey.String())
s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
return srv.Context().Err()
}
}
}
// sendUpdate encrypts the update message using the peer key and the server's wireguard key,
// then sends the encrypted message to the connected peer via the sync server.
func (s *Server) sendUpdate(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, update *network_map.UpdateMessage, srv proto.ManagementService_SyncServer, streamStartTime time.Time) error {
key, err := s.secretsManager.GetWGKey()
if err != nil {
s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
return status.Errorf(codes.Internal, "failed processing update message")
}
encryptedResp, err := encryption.EncryptMessage(peerKey, key, update.Update)
if err != nil {
s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
return status.Errorf(codes.Internal, "failed processing update message")
}
err = srv.Send(&proto.EncryptedMessage{
WgPubKey: key.PublicKey().String(),
Body: encryptedResp,
})
if err != nil {
s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
return status.Errorf(codes.Internal, "failed sending update message")
}
log.WithContext(ctx).Tracef("sent an update to peer %s", peerKey.String())
return nil
}
// sendJob encrypts the update message using the peer key and the server's wireguard key,
// then sends the encrypted message to the connected peer via the sync server.
func (s *Server) sendJob(ctx context.Context, peerKey wgtypes.Key, job *job.Event, srv proto.ManagementService_JobServer) error {
@@ -0,0 +1,70 @@
// Code generated by MockGen. DO NOT EDIT.
// Source: ./peer_update_handler.go
//
// Generated by this command:
//
// mockgen -source=./peer_update_handler.go -destination=./sync_sender_mock.go -package=grpc
//
// Package grpc is a generated GoMock package.
package grpc
import (
context "context"
reflect "reflect"
proto "github.com/netbirdio/netbird/shared/management/proto"
gomock "go.uber.org/mock/gomock"
)
// MocksyncSender is a mock of syncSender interface.
type MocksyncSender struct {
ctrl *gomock.Controller
recorder *MocksyncSenderMockRecorder
isgomock struct{}
}
// MocksyncSenderMockRecorder is the mock recorder for MocksyncSender.
type MocksyncSenderMockRecorder struct {
mock *MocksyncSender
}
// NewMocksyncSender creates a new mock instance.
func NewMocksyncSender(ctrl *gomock.Controller) *MocksyncSender {
mock := &MocksyncSender{ctrl: ctrl}
mock.recorder = &MocksyncSenderMockRecorder{mock}
return mock
}
// EXPECT returns an object that allows the caller to indicate expected use.
func (m *MocksyncSender) EXPECT() *MocksyncSenderMockRecorder {
return m.recorder
}
// Context mocks base method.
func (m *MocksyncSender) Context() context.Context {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "Context")
ret0, _ := ret[0].(context.Context)
return ret0
}
// Context indicates an expected call of Context.
func (mr *MocksyncSenderMockRecorder) Context() *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Context", reflect.TypeOf((*MocksyncSender)(nil).Context))
}
// Send mocks base method.
func (m *MocksyncSender) Send(arg0 *proto.EncryptedMessage) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "Send", arg0)
ret0, _ := ret[0].(error)
return ret0
}
// Send indicates an expected call of Send.
func (mr *MocksyncSenderMockRecorder) Send(arg0 any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Send", reflect.TypeOf((*MocksyncSender)(nil).Send), arg0)
}
@@ -25,6 +25,8 @@ import (
const defaultDuration = 12 * time.Hour
// SecretsManager used to manage TURN and relay secrets
//
//go:generate go tool mockgen -source=./token_mgr.go -destination=./token_mgr_mock.go -package=grpc
type SecretsManager interface {
GenerateTurnToken() (*Token, error)
GenerateRelayToken() (*Token, error)
@@ -0,0 +1,111 @@
// Code generated by MockGen. DO NOT EDIT.
// Source: ./token_mgr.go
//
// Generated by this command:
//
// mockgen -source=./token_mgr.go -destination=./token_mgr_mock.go -package=grpc
//
// Package grpc is a generated GoMock package.
package grpc
import (
context "context"
reflect "reflect"
gomock "go.uber.org/mock/gomock"
wgtypes "golang.zx2c4.com/wireguard/wgctrl/wgtypes"
)
// MockSecretsManager is a mock of SecretsManager interface.
type MockSecretsManager struct {
ctrl *gomock.Controller
recorder *MockSecretsManagerMockRecorder
isgomock struct{}
}
// MockSecretsManagerMockRecorder is the mock recorder for MockSecretsManager.
type MockSecretsManagerMockRecorder struct {
mock *MockSecretsManager
}
// NewMockSecretsManager creates a new mock instance.
func NewMockSecretsManager(ctrl *gomock.Controller) *MockSecretsManager {
mock := &MockSecretsManager{ctrl: ctrl}
mock.recorder = &MockSecretsManagerMockRecorder{mock}
return mock
}
// EXPECT returns an object that allows the caller to indicate expected use.
func (m *MockSecretsManager) EXPECT() *MockSecretsManagerMockRecorder {
return m.recorder
}
// CancelRefresh mocks base method.
func (m *MockSecretsManager) CancelRefresh(peerKey string) {
m.ctrl.T.Helper()
m.ctrl.Call(m, "CancelRefresh", peerKey)
}
// CancelRefresh indicates an expected call of CancelRefresh.
func (mr *MockSecretsManagerMockRecorder) CancelRefresh(peerKey any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CancelRefresh", reflect.TypeOf((*MockSecretsManager)(nil).CancelRefresh), peerKey)
}
// GenerateRelayToken mocks base method.
func (m *MockSecretsManager) GenerateRelayToken() (*Token, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GenerateRelayToken")
ret0, _ := ret[0].(*Token)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GenerateRelayToken indicates an expected call of GenerateRelayToken.
func (mr *MockSecretsManagerMockRecorder) GenerateRelayToken() *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GenerateRelayToken", reflect.TypeOf((*MockSecretsManager)(nil).GenerateRelayToken))
}
// GenerateTurnToken mocks base method.
func (m *MockSecretsManager) GenerateTurnToken() (*Token, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GenerateTurnToken")
ret0, _ := ret[0].(*Token)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GenerateTurnToken indicates an expected call of GenerateTurnToken.
func (mr *MockSecretsManagerMockRecorder) GenerateTurnToken() *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GenerateTurnToken", reflect.TypeOf((*MockSecretsManager)(nil).GenerateTurnToken))
}
// GetWGKey mocks base method.
func (m *MockSecretsManager) GetWGKey() (wgtypes.Key, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetWGKey")
ret0, _ := ret[0].(wgtypes.Key)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetWGKey indicates an expected call of GetWGKey.
func (mr *MockSecretsManagerMockRecorder) GetWGKey() *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWGKey", reflect.TypeOf((*MockSecretsManager)(nil).GetWGKey))
}
// SetupRefresh mocks base method.
func (m *MockSecretsManager) SetupRefresh(ctx context.Context, accountID, peerKey string) {
m.ctrl.T.Helper()
m.ctrl.Call(m, "SetupRefresh", ctx, accountID, peerKey)
}
// SetupRefresh indicates an expected call of SetupRefresh.
func (mr *MockSecretsManagerMockRecorder) SetupRefresh(ctx, accountID, peerKey any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetupRefresh", reflect.TypeOf((*MockSecretsManager)(nil).SetupRefresh), ctx, accountID, peerKey)
}
@@ -6,6 +6,14 @@ import (
"github.com/netbirdio/netbird/management/internals/controllers/network_map"
)
//go:generate go tool mockgen -source=./update_debouncer.go -destination=./update_debouncer_mock.go -package=grpc
type Debouncer interface {
Stop()
TimerChannel() <-chan time.Time
ProcessUpdate(update *network_map.UpdateMessage) bool
GetPendingUpdates() []*network_map.UpdateMessage
}
// UpdateDebouncer implements a backpressure mechanism that:
// - Sends the first update immediately
// - Coalesces rapid subsequent network map updates (only latest matters)
@@ -0,0 +1,96 @@
// Code generated by MockGen. DO NOT EDIT.
// Source: ./update_debouncer.go
//
// Generated by this command:
//
// mockgen -source=./update_debouncer.go -destination=./update_debouncer_mock.go -package=grpc
//
// Package grpc is a generated GoMock package.
package grpc
import (
reflect "reflect"
time "time"
network_map "github.com/netbirdio/netbird/management/internals/controllers/network_map"
gomock "go.uber.org/mock/gomock"
)
// MockDebouncer is a mock of Debouncer interface.
type MockDebouncer struct {
ctrl *gomock.Controller
recorder *MockDebouncerMockRecorder
isgomock struct{}
}
// MockDebouncerMockRecorder is the mock recorder for MockDebouncer.
type MockDebouncerMockRecorder struct {
mock *MockDebouncer
}
// NewMockDebouncer creates a new mock instance.
func NewMockDebouncer(ctrl *gomock.Controller) *MockDebouncer {
mock := &MockDebouncer{ctrl: ctrl}
mock.recorder = &MockDebouncerMockRecorder{mock}
return mock
}
// EXPECT returns an object that allows the caller to indicate expected use.
func (m *MockDebouncer) EXPECT() *MockDebouncerMockRecorder {
return m.recorder
}
// GetPendingUpdates mocks base method.
func (m *MockDebouncer) GetPendingUpdates() []*network_map.UpdateMessage {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetPendingUpdates")
ret0, _ := ret[0].([]*network_map.UpdateMessage)
return ret0
}
// GetPendingUpdates indicates an expected call of GetPendingUpdates.
func (mr *MockDebouncerMockRecorder) GetPendingUpdates() *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPendingUpdates", reflect.TypeOf((*MockDebouncer)(nil).GetPendingUpdates))
}
// ProcessUpdate mocks base method.
func (m *MockDebouncer) ProcessUpdate(update *network_map.UpdateMessage) bool {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "ProcessUpdate", update)
ret0, _ := ret[0].(bool)
return ret0
}
// ProcessUpdate indicates an expected call of ProcessUpdate.
func (mr *MockDebouncerMockRecorder) ProcessUpdate(update any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ProcessUpdate", reflect.TypeOf((*MockDebouncer)(nil).ProcessUpdate), update)
}
// Stop mocks base method.
func (m *MockDebouncer) Stop() {
m.ctrl.T.Helper()
m.ctrl.Call(m, "Stop")
}
// Stop indicates an expected call of Stop.
func (mr *MockDebouncerMockRecorder) Stop() *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Stop", reflect.TypeOf((*MockDebouncer)(nil).Stop))
}
// TimerChannel mocks base method.
func (m *MockDebouncer) TimerChannel() <-chan time.Time {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "TimerChannel")
ret0, _ := ret[0].(<-chan time.Time)
return ret0
}
// TimerChannel indicates an expected call of TimerChannel.
func (mr *MockDebouncerMockRecorder) TimerChannel() *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TimerChannel", reflect.TypeOf((*MockDebouncer)(nil).TimerChannel))
}