sendnrw/netbird
2
0
Fork 0
mirror of https://github.com/netbirdio/netbird.git synced 2026-10-28 15:29:09 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/main' into dmitri-catch-disconnected-peer

Browse Source
This commit is contained in:
pascal
2026-10-28 15:19:14 +01:00
parent d8f8bd7271 1acdebf395
commit 264a0865d1
36 changed files with 1637 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files
client/cmd/root.go
+4 -2
View File
@@ -220,13 +220,15 @@ func init() {
`An empty string "" clears the previous configuration. `+
`E.g. --dns-resolver-address 127.0.0.1:5053 or --dns-resolver-address ""`,
)
upCmd.PersistentFlags().BoolVar(&rosenpassEnabled, enableRosenpassFlag, false, "[Experimental] Enable Rosenpass feature. If enabled, the connection will be post-quantum secured via Rosenpass.")
upCmd.PersistentFlags().BoolVar(&rosenpassPermissive, rosenpassPermissiveFlag, false, "[Experimental] Enable Rosenpass in permissive mode to allow this peer to accept WireGuard connections without requiring Rosenpass functionality from peers that do not have Rosenpass enabled.")
upCmd.PersistentFlags().BoolVar(&rosenpassEnabled, enableRosenpassFlag, false, "[Deprecated] Enable Rosenpass feature. If enabled, the connection will be post-quantum secured via Rosenpass.")
upCmd.PersistentFlags().BoolVar(&rosenpassPermissive, rosenpassPermissiveFlag, false, "[Deprecated] Enable Rosenpass in permissive mode to allow this peer to accept WireGuard connections without requiring Rosenpass functionality from peers that do not have Rosenpass enabled.")
upCmd.PersistentFlags().BoolVar(&autoConnectDisabled, disableAutoConnectFlag, false, "Disables auto-connect feature. If enabled, then the client won't connect automatically when the service starts.")
upCmd.PersistentFlags().BoolVar(&localMetricsEnabled, enableLocalMetricsFlag, false, "Enables a local Prometheus /metrics endpoint exposing connection state (peers, latency, P2P vs relay).")
upCmd.PersistentFlags().StringVar(&localMetricsAddr, localMetricsAddressFlag, localmetrics.DefaultListenAddress, "Listen address of the local Prometheus /metrics endpoint.")
upCmd.PersistentFlags().BoolVar(&lazyConnEnabled, enableLazyConnectionFlag, false, "Deprecated: no longer used. Lazy connections are controlled by the server and the NB_LAZY_CONN environment variable.")
_ = upCmd.PersistentFlags().MarkDeprecated(enableLazyConnectionFlag, "no longer used; lazy connections are controlled by the server and the NB_LAZY_CONN environment variable")
_ = upCmd.PersistentFlags().MarkDeprecated(enableRosenpassFlag, "Rosenpass support may be removed in a future release")
_ = upCmd.PersistentFlags().MarkDeprecated(rosenpassPermissiveFlag, "Rosenpass support may be removed in a future release")
}
client/cmd/up_rosenpass_test.go
+34
View File
@@ -0,0 +1,34 @@
package cmd
import (
"bytes"
"testing"
"github.com/spf13/pflag"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestRosenpassFlagsWarnDeprecated(t *testing.T) {
origEnabled, origPermissive := rosenpassEnabled, rosenpassPermissive
t.Cleanup(func() {
rosenpassEnabled, rosenpassPermissive = origEnabled, origPermissive
})
for _, name := range []string{enableRosenpassFlag, rosenpassPermissiveFlag} {
// Copy the production flag into a private set so Set marks only the copy as
// Changed, leaving upCmd's shared flag set untouched for other tests.
prod := upCmd.PersistentFlags().Lookup(name)
require.NotNil(t, prod, "flag --%s must be registered on up", name)
flagCopy := *prod
var out bytes.Buffer
flags := pflag.NewFlagSet("rosenpass-test", pflag.ContinueOnError)
flags.SetOutput(&out)
flags.AddFlag(&flagCopy)
require.NoError(t, flags.Set(name, "true"))
assert.Contains(t, out.String(), "Flag --"+name+" has been deprecated", "setting --%s should print the deprecation notice", name)
assert.False(t, prod.Changed, "the shared --%s flag must stay unchanged", name)
}
}
client/ui/frontend/src/modules/settings/SettingsSecurity.tsx
+8 -1
View File
@@ -42,7 +42,14 @@ export function SettingsSecurity() {
if (!v) setField("rosenpassPermissive", false);
}}
label={t("settings.security.rosenpass.label")}
helpText={t("settings.security.rosenpass.help")}
helpText={
<>
{t("settings.security.rosenpass.help")}
<span className={"mt-1 block"}>
{t("settings.security.rosenpass.deprecated")}
</span>
</>
}
/>
)}
{!hideRosenpassPermissive && (
client/ui/i18n/locales/en/common.json
+4
View File
@@ -899,6 +899,10 @@
"message": "Add a post-quantum key exchange via Rosenpass on top of WireGuard®.",
"description": "Helper text: adds a post-quantum key exchange via Rosenpass on top of WireGuard®. 'Rosenpass' and 'WireGuard®' are product names — do not translate; keep the ® symbol."
},
"settings.security.rosenpass.deprecated": {
"message": "Deprecated: Rosenpass support may be removed in a future release.",
"description": "Notice shown under the quantum-resistance toggle: the Rosenpass feature is deprecated and may be removed later. 'Rosenpass' is a product name — do not translate."
},
"settings.security.rosenpassPermissive.label": {
"message": "Enable Permissive Mode",
"description": "Toggle label: enable permissive mode (for quantum-resistance)."
go.mod
+1 -1
View File
@@ -369,7 +369,7 @@ replace github.com/cloudflare/circl => codeberg.org/cunicu/circl v0.0.0-20230801
replace github.com/mailru/easyjson => github.com/netbirdio/easyjson v0.9.0
replace github.com/wailsapp/wails/v3 => github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20261009120932-4e087a8bfd6d
replace github.com/wailsapp/wails/v3 => github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20261028091752-cfa103160796
tool (
github.com/goreleaser/chglog/cmd/chglog
go.sum
+2 -2
View File
@@ -538,8 +538,8 @@ github.com/netbirdio/management-integrations/integrations v0.0.0-20261021190039-
github.com/netbirdio/management-integrations/integrations v0.0.0-20261021190039-6c1585eb11de/go.mod h1:LLvpw/iWbXX9+7HF7ueqkyE7SoFaeqoCmbskOQL9gbs=
github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20250805121659-6b4ac470ca45 h1:ujgviVYmx243Ksy7NdSwrdGPSRNE3pb8kEDSpH0QuAQ=
github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20250805121659-6b4ac470ca45/go.mod h1:5/sjFmLb8O96B5737VCqhHyGRzNFIaN/Bu7ZodXc3qQ=
github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20261009120932-4e087a8bfd6d h1:DV0FHFhNAaU69JHXKclYQWr/ZJoc2BA/ZpqlnX8l22U=
github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20261009120932-4e087a8bfd6d/go.mod h1:62UsqQRqanuCbFi6xElIw3/2X+yGF+TKlyiWImA5taU=
github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20261028091752-cfa103160796 h1:+21hEHg5uXZiYP7vpEH3Eo3jk5anmynYRUq5dPGErxw=
github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20261028091752-cfa103160796/go.mod h1:1U8Kg7QyeO09HanpT48heL7yOCPaYfh9ErFVzmxfiLQ=
github.com/netbirdio/wireguard-go v0.0.0-20261021152105-0de46b309cdb h1:3X7pVeZeneq4MK7hvvqeSGi1ir6OElgIqHd6cI4KsXo=
github.com/netbirdio/wireguard-go v0.0.0-20261021152105-0de46b309cdb/go.mod h1:Vk/c438m8FV2hHcT/6Efp8xe/JhHTzf3eUcuaohKKkw=
github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
management/internals/modules/reverseproxy/accesslogs/manager/repository.go
+28 -9
View File
@@ -14,18 +14,21 @@ import (
"github.com/netbirdio/netbird/shared/management/status"
)
const deleteBatchSize = 10000
type sqlRepository struct {
conn *db.Conn
db *gorm.DB
conn *db.Conn
db *gorm.DB
deleteBatchSize int
}
// NewRepository returns the access log repository backed by conn.
func NewRepository(conn *db.Conn) accesslogs.Repository {
return &sqlRepository{conn: conn, db: conn.DB(nil)}
return &sqlRepository{conn: conn, db: conn.DB(nil), deleteBatchSize: deleteBatchSize}
}
func (r *sqlRepository) WithTx(tx *db.Tx) accesslogs.Repository {
return &sqlRepository{conn: r.conn, db: r.conn.DB(tx)}
return &sqlRepository{conn: r.conn, db: r.conn.DB(tx), deleteBatchSize: r.deleteBatchSize}
}
func (r *sqlRepository) Create(ctx context.Context, entry *accesslogs.AccessLogEntry) error {
@@ -73,12 +76,28 @@ func (r *sqlRepository) ListByAccount(ctx context.Context, lockStrength db.Locki
}
func (r *sqlRepository) DeleteOlderThan(ctx context.Context, olderThan time.Time) (int64, error) {
result := r.db.Where("timestamp < ?", olderThan).Delete(&accesslogs.AccessLogEntry{})
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to delete old access logs: %v", result.Error)
return 0, status.Errorf(status.Internal, "failed to delete old access logs")
var deleted int64
for {
if err := ctx.Err(); err != nil {
return deleted, err
}
batch := r.db.Model(&accesslogs.AccessLogEntry{}).
Select("id").
Where("timestamp < ?", olderThan).
Order("timestamp").
Limit(r.deleteBatchSize)
result := r.db.Where("id IN (?)", batch).Delete(&accesslogs.AccessLogEntry{})
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to delete old access logs: %v", result.Error)
return deleted, status.Errorf(status.Internal, "failed to delete old access logs")
}
deleted += result.RowsAffected
if result.RowsAffected < int64(r.deleteBatchSize) {
return deleted, nil
}
}
return result.RowsAffected, nil
}
func applyFilters(query *gorm.DB, filter accesslogs.AccessLogFilter) *gorm.DB {
management/internals/modules/reverseproxy/accesslogs/manager/repository_test.go
+70
View File
@@ -3,6 +3,8 @@ package manager
import (
"context"
"errors"
"fmt"
"os"
"testing"
"time"
@@ -12,6 +14,7 @@ import (
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs"
"github.com/netbirdio/netbird/management/internals/shared/db"
"github.com/netbirdio/netbird/management/internals/shared/db/dbtest"
"github.com/netbirdio/netbird/management/server/testutil"
)
func newTestRepository(t *testing.T) (accesslogs.Repository, *db.Conn) {
@@ -157,3 +160,70 @@ func TestSqlRepository_ListByAccount_Search(t *testing.T) {
assert.Equal(t, []string{"req-api", "req-web"}, search("example.com"), "search should list every match newest first")
assert.Empty(t, search("nothing-matches"))
}
func TestSqlRepository_DeleteOlderThan_DeletesInBatches(t *testing.T) {
cases := map[string]struct{ old, recent int }{
"partial last batch": {old: 5, recent: 2},
"exact multiple of batch": {old: 4, recent: 1},
"nothing to delete": {old: 0, recent: 3},
}
for name, tc := range cases {
t.Run(name, func(t *testing.T) {
repo, _ := newTestRepository(t)
assertBatchedDelete(t, repo, tc.old, tc.recent)
})
}
}
func TestSqlRepository_DeleteOlderThan_DeletesInBatchesOnPostgres(t *testing.T) {
if os.Getenv("NETBIRD_STORE_ENGINE") != string(db.PostgresStoreEngine) {
t.Skip("requires NETBIRD_STORE_ENGINE=postgres")
}
cleanup, dsn, err := testutil.CreatePostgresTestContainer()
require.NoError(t, err)
t.Cleanup(cleanup)
conn, err := db.OpenPostgres(context.Background(), dsn, db.DefaultPoolConfig)
require.NoError(t, err)
t.Cleanup(func() { _ = conn.Close() })
require.NoError(t, conn.AutoMigrate(&accesslogs.AccessLogEntry{}))
require.NoError(t, conn.DB(nil).Exec("DELETE FROM access_log_entries").Error)
assertBatchedDelete(t, NewRepository(conn), 5, 2)
}
func TestSqlRepository_DeleteOlderThan_StopsWhenContextDone(t *testing.T) {
repo, _ := newTestRepository(t)
require.NoError(t, repo.Create(context.Background(), newEntry("old", "acc", "GET", 48*time.Hour)))
ctx, cancel := context.WithCancel(context.Background())
cancel()
deleted, err := repo.DeleteOlderThan(ctx, time.Now().Add(-24*time.Hour))
require.ErrorIs(t, err, context.Canceled)
assert.Zero(t, deleted)
assert.EqualValues(t, 1, countEntries(t, repo), "nothing should be deleted once the context is done")
}
func assertBatchedDelete(t *testing.T, repo accesslogs.Repository, old, recent int) {
t.Helper()
repo.(*sqlRepository).deleteBatchSize = 2
ctx := context.Background()
for i := 0; i < old; i++ {
require.NoError(t, repo.Create(ctx, newEntry(fmt.Sprintf("old-%d", i), "acc", "GET", 48*time.Hour)))
}
for i := 0; i < recent; i++ {
require.NoError(t, repo.Create(ctx, newEntry(fmt.Sprintf("new-%d", i), "acc", "GET", time.Hour)))
}
deleted, err := repo.DeleteOlderThan(ctx, time.Now().Add(-24*time.Hour))
require.NoError(t, err)
assert.EqualValues(t, old, deleted, "every old entry should be deleted across batches")
assert.EqualValues(t, recent, countEntries(t, repo), "recent entries should be kept")
}
func countEntries(t *testing.T, repo accesslogs.Repository) int64 {
t.Helper()
_, total, err := repo.ListByAccount(context.Background(), db.LockingStrengthNone, "acc", accesslogs.AccessLogFilter{Page: 1, PageSize: 10})
require.NoError(t, err)
return total
}
management/internals/server/boot.go
+19 -9
View File
@@ -269,6 +269,9 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
grpc.KeepaliveParams(kasp),
grpc.ChainUnaryInterceptor(realip.UnaryServerInterceptorOpts(realipOpts...), unaryInterceptor, proxyUnary),
grpc.ChainStreamInterceptor(realip.StreamServerInterceptorOpts(realipOpts...), streamInterceptor, proxyStream),
// Stop waits for the sync handlers, so their teardown has queued every
// disconnect before FlushDisconnects writes them.
grpc.WaitForHandlers(true),
}
// Append interceptors contributed by registered gRPC extensions. These
@@ -292,6 +295,21 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
}
gRPCAPIHandler := grpc.NewServer(gRPCOpts...)
mgmtProto.RegisterManagementServiceServer(gRPCAPIHandler, s.ManagementGRPCServer())
mgmtProto.RegisterProxyServiceServer(gRPCAPIHandler, s.ReverseProxyGRPCServer())
log.Info("ProxyService registered on gRPC server")
// Register services contributed by external modules via the extension seam.
registerExtensions(gRPCAPIHandler, s.grpcExtensions)
return gRPCAPIHandler
})
}
// ManagementGRPCServer returns the management gRPC service, created once and served by GRPCServer.
func (s *BaseServer) ManagementGRPCServer() *nbgrpc.Server {
return Create(s, func() *nbgrpc.Server {
srv, err := nbgrpc.NewServer(s.Config, s.AccountManager(), s.SettingsManager(), s.JobManager(), s.SecretsManager(), s.Metrics(), s.AuthManager(), s.IntegratedValidator(), s.NetworkMapController(), s.OAuthConfigProvider(), s.SessionStore())
if err != nil {
log.Fatalf("failed to create management server: %v", err)
@@ -301,15 +319,7 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
if serviceMgr != nil {
serviceMgr.StartExposeReaper(context.Background())
}
mgmtProto.RegisterManagementServiceServer(gRPCAPIHandler, srv)
mgmtProto.RegisterProxyServiceServer(gRPCAPIHandler, s.ReverseProxyGRPCServer())
log.Info("ProxyService registered on gRPC server")
// Register services contributed by external modules via the extension seam.
registerExtensions(gRPCAPIHandler, s.grpcExtensions)
return gRPCAPIHandler
return srv
})
}
management/internals/server/container.go
+33 -4
View File
@@ -1,6 +1,9 @@
package server
import "fmt"
import (
"fmt"
"reflect"
)
// Create a dependency and add it to the BaseServer's container. A string key identifier will be based on its type definition.
func Create[T any](s Server, createFunc func() T) T {
@@ -33,15 +36,41 @@ func InjectNamed[T any](c Server, name string, thing T) {
}
func maybeCreate[T any](s Server, createFunc func() T) (result T, isNew bool) {
key := fmt.Sprintf("%T", (*T)(nil))[1:]
return maybeCreateKeyed(s, key, createFunc)
return maybeCreateKeyed(s, typeKey(reflect.TypeFor[T]()), createFunc)
}
func maybeCreateNamed[T any](s Server, name string, createFunc func() T) (result T, isNew bool) {
key := fmt.Sprintf("%T:%s", (*T)(nil), name)[1:]
key := fmt.Sprintf("%s:%s", typeKey(reflect.TypeFor[T]()), name)
return maybeCreateKeyed(s, key, createFunc)
}
// typeKey names t by the full package path of every named type it is built from, so
// same-named types of different packages, such as google.golang.org/grpc.Server and a
// local grpc.Server, never share a container key.
func typeKey(t reflect.Type) string {
if t.Name() != "" {
if t.PkgPath() == "" {
return t.Name()
}
return t.PkgPath() + "." + t.Name()
}
switch t.Kind() {
case reflect.Pointer:
return "*" + typeKey(t.Elem())
case reflect.Slice:
return "[]" + typeKey(t.Elem())
case reflect.Array:
return fmt.Sprintf("[%d]%s", t.Len(), typeKey(t.Elem()))
case reflect.Map:
return "map[" + typeKey(t.Key()) + "]" + typeKey(t.Elem())
case reflect.Chan:
return t.ChanDir().String() + " " + typeKey(t.Elem())
default:
return t.String()
}
}
func maybeCreateKeyed[T any](s Server, key string, createFunc func() T) (result T, isNew bool) {
if t, ok := s.GetContainer(key); ok {
if t == nil {
management/internals/server/server.go
+46 -7
View File
@@ -3,6 +3,7 @@ package server
import (
"context"
"crypto/tls"
"errors"
"fmt"
"net"
"net/http"
@@ -40,6 +41,15 @@ const (
ContainerKeyBaseServer = "baseServer"
applicationName = "management"
// shutdownTimeout stays below the 10s that systemd (TimeoutStopSec) and Docker grant
// before killing the process.
shutdownTimeout = 9 * time.Second
grpcStopTimeout = 3 * time.Second
// shutdownCloseReserve is kept from the flush for closing the stores.
shutdownCloseReserve = 2 * time.Second
// shutdownIdPReserve is kept from the extension shutdown hooks for stopping the embedded IdP.
shutdownIdPReserve = 1 * time.Second
)
type Server interface {
@@ -297,9 +307,13 @@ func (s *BaseServer) setupTLS(ctx context.Context) (bool, error) {
}
}
// Stop attempts a graceful shutdown, waiting up to 5 seconds for active connections to finish
// Stop attempts a graceful shutdown within shutdownTimeout. Writing the disconnects of the
// drained streams may use the budget left after stopping the services and gRPC streams,
// except for shutdownCloseReserve, which stays for closing the stores; the extension hooks
// leave shutdownIdPReserve for stopping the embedded IdP. It returns an error when the
// disconnects could not all be written or the shutdown ran out of time.
func (s *BaseServer) Stop() error {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
defer cancel()
defer s.RunStopHandlers()
if s.domainCleanupStop != nil {
@@ -318,12 +332,20 @@ func (s *BaseServer) Stop() error {
if s.certListener != nil {
_ = s.certListener.Close()
}
s.GRPCServer().Stop()
s.ManagementGRPCServer().StartDraining()
s.stopGRPCServer()
deadline, _ := ctx.Deadline()
flushCtx, cancelFlush := context.WithDeadline(ctx, deadline.Add(-shutdownCloseReserve))
flushErr := s.ManagementGRPCServer().FlushDisconnects(flushCtx)
cancelFlush()
if s.proxyAuthClose != nil {
s.proxyAuthClose()
s.proxyAuthClose = nil
}
runExtensionShutdownHooks(ctx, s.grpcExtensions)
hooksCtx, cancelHooks := context.WithDeadline(ctx, deadline.Add(-shutdownIdPReserve))
runExtensionShutdownHooks(hooksCtx, s.grpcExtensions)
cancelHooks()
_ = s.Store().Close(ctx)
_ = s.EventStore().Close(ctx)
if s.update != nil {
@@ -338,12 +360,29 @@ func (s *BaseServer) Stop() error {
case err := <-s.Errors():
if err != nil {
log.WithContext(ctx).Infof("stopped Management Service: %v", err)
return nil
return flushErr
}
log.WithContext(ctx).Infof("stopped Management Service")
return nil
return flushErr
case <-ctx.Done():
return ctx.Err()
return errors.Join(flushErr, ctx.Err())
}
}
// stopGRPCServer stops the gRPC server, which waits for the stream handlers to finish
// their teardown, but gives up waiting after grpcStopTimeout so a stuck handler cannot
// block the shutdown.
func (s *BaseServer) stopGRPCServer() {
stopped := make(chan struct{})
go func() {
s.GRPCServer().Stop()
close(stopped)
}()
select {
case <-stopped:
case <-time.After(grpcStopTimeout):
log.Warnf("gRPC server did not stop within %s, continuing shutdown", grpcStopTimeout)
}
}
management/internals/shared/db/conn.go
+6 -6
View File
@@ -11,6 +11,8 @@ import (
"github.com/jackc/pgx/v5/pgxpool"
log "github.com/sirupsen/logrus"
"gorm.io/gorm"
"github.com/netbirdio/netbird/util"
)
const (
@@ -64,13 +66,11 @@ func NewConn(ctx context.Context, gormDB *gorm.DB, engine Engine, pool *pgxpool.
opt(&options)
}
txTimeout := defaultTransactionTimeout
if v := os.Getenv("NB_STORE_TRANSACTION_TIMEOUT"); v != "" {
if parsed, err := time.ParseDuration(v); err == nil {
txTimeout = parsed
}
}
txTimeout := util.DurationFromEnv(transactionTimeoutEnv, defaultTransactionTimeout)
log.WithContext(ctx).Infof("Setting transaction timeout to %v", txTimeout)
if d := util.DurationFromEnv(statementTimeoutEnv, 0); engine == PostgresStoreEngine && d > 0 {
log.WithContext(ctx).Infof("Setting statement timeout to %v", d)
}
conns := runtime.NumCPU()
configuredConns, err := strconv.Atoi(os.Getenv("NB_SQL_MAX_OPEN_CONNS"))
management/internals/shared/db/open.go
+23 -4
View File
@@ -10,7 +10,9 @@ import (
"strings"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/jackc/pgx/v5/stdlib"
"gorm.io/driver/postgres"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
@@ -100,9 +102,10 @@ func OpenSqliteFile(ctx context.Context, dataDir, storeFile string) (*Conn, erro
return conn, nil
}
// OpenPostgres opens a Postgres database through gorm and a pgx pool sized by pool.
// OpenPostgres opens a Postgres database through gorm and a pgx pool sized by
// pool, applying the configured timeouts to every session.
func OpenPostgres(ctx context.Context, dsn string, pool PoolConfig, opts ...Option) (*Conn, error) {
gormDB, err := gorm.Open(postgres.Open(dsn), GormConfig())
gormDB, err := openGormPostgres(ctx, dsn)
if err != nil {
return nil, err
}
@@ -123,7 +126,7 @@ func OpenPostgres(ctx context.Context, dsn string, pool PoolConfig, opts ...Opti
// OpenPostgresWithoutPool opens a Postgres database through gorm alone, for
// repositories that never take the pgx read paths.
func OpenPostgresWithoutPool(ctx context.Context, dsn string, opts ...Option) (*Conn, error) {
gormDB, err := gorm.Open(postgres.Open(dsn), GormConfig())
gormDB, err := openGormPostgres(ctx, dsn)
if err != nil {
return nil, err
}
@@ -135,13 +138,29 @@ func OpenPostgresWithoutPool(ctx context.Context, dsn string, opts ...Option) (*
return conn, nil
}
// OpenPgxPool opens a pgx pool to dsn sized by cfg and verifies it with a ping.
func openGormPostgres(ctx context.Context, dsn string) (*gorm.DB, error) {
config, err := pgx.ParseConfig(dsn)
if err != nil {
return nil, fmt.Errorf("parse database config: %w", err)
}
sqlDB := stdlib.OpenDB(*config, stdlib.OptionAfterConnect(sessionTimeouts()))
gormDB, err := gorm.Open(postgres.New(postgres.Config{Conn: sqlDB}), GormConfig())
if err != nil {
_ = sqlDB.Close()
return nil, err
}
return gormDB, nil
}
// OpenPgxPool opens a pgx pool to dsn sized by cfg, applying the configured
// timeouts to every session, and verifies it with a ping.
func OpenPgxPool(ctx context.Context, dsn string, cfg PoolConfig) (*pgxpool.Pool, error) {
config, err := pgxpool.ParseConfig(dsn)
if err != nil {
return nil, fmt.Errorf("unable to parse database config: %w", err)
}
config.AfterConnect = sessionTimeouts()
config.MaxConns = cfg.MaxConns
config.MinConns = cfg.MinConns
config.MaxConnLifetime = cfg.MaxConnLifetime
management/internals/shared/db/timeout.go
+52
View File
@@ -0,0 +1,52 @@
package db
import (
"context"
"fmt"
"math"
"strconv"
"strings"
"time"
"github.com/jackc/pgx/v5"
"github.com/netbirdio/netbird/util"
)
const (
transactionTimeoutEnv = "NB_STORE_TRANSACTION_TIMEOUT"
statementTimeoutEnv = "NB_STORE_STATEMENT_TIMEOUT"
)
// sessionTimeouts sets the configured timeouts on each new Postgres session with
// SET, since poolers such as PgBouncer reject them as startup parameters.
func sessionTimeouts() func(context.Context, *pgx.Conn) error {
var settings []string
if d := util.DurationFromEnv(statementTimeoutEnv, 0); d > 0 {
settings = append(settings, "SET statement_timeout = "+postgresMillis(d))
}
if d := util.DurationFromEnv(transactionTimeoutEnv, defaultTransactionTimeout); d > 0 {
settings = append(settings, "SET idle_in_transaction_session_timeout = "+postgresMillis(d))
}
query := strings.Join(settings, "; ")
return func(ctx context.Context, conn *pgx.Conn) error {
if query == "" {
return nil
}
if _, err := conn.Exec(ctx, query); err != nil {
return fmt.Errorf("set session timeouts: %w", err)
}
return nil
}
}
// postgresMillis converts d to milliseconds for a Postgres timeout, rounding up
// so a positive d never becomes 0 (disabled) and capping at the server maximum.
func postgresMillis(d time.Duration) string {
ms := d.Milliseconds()
if d%time.Millisecond != 0 {
ms++
}
return strconv.FormatInt(min(ms, math.MaxInt32), 10)
}
management/internals/shared/db/timeout_test.go
+146
View File
@@ -0,0 +1,146 @@
package db
import (
"context"
"errors"
"os"
"testing"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
"github.com/netbirdio/netbird/management/server/testutil"
)
// queryCanceledCode is the SQLSTATE Postgres reports when statement_timeout fires.
const queryCanceledCode = "57014"
func TestPostgresMillis(t *testing.T) {
cases := map[time.Duration]string{
time.Minute: "60000",
1500 * time.Microsecond: "2",
500 * time.Microsecond: "1",
720 * time.Hour: "2147483647",
}
for d, want := range cases {
assert.Equal(t, want, postgresMillis(d), "milliseconds for %v", d)
}
}
func TestSessionTimeouts_NothingToApply(t *testing.T) {
t.Setenv(statementTimeoutEnv, "0")
t.Setenv(transactionTimeoutEnv, "0")
ctx := context.Background()
assert.NoError(t, sessionTimeouts()(ctx, nil), "the hook should not touch the session when both timeouts are disabled")
}
func TestOpenPostgres_BoundsEverySession(t *testing.T) {
t.Setenv(statementTimeoutEnv, "500ms")
t.Setenv(transactionTimeoutEnv, "3m")
conn := openPostgresTestConn(t)
ctx := context.Background()
requireStatementTimeout(t, conn.DB(nil).Exec("SELECT pg_sleep(1)").Error)
_, err := conn.Pool(nil).Exec(ctx, "SELECT pg_sleep(1)")
requireStatementTimeout(t, err)
var inTx string
require.NoError(t, conn.RunInTx(ctx, func(tx *Tx) error {
return conn.DB(tx).Raw("SHOW statement_timeout").Scan(&inTx).Error
}))
assert.Equal(t, "1min", inTx, "transactions should keep their own 1 minute bound")
assert.Equal(t, "3min", showSetting(t, conn.DB(nil), "idle_in_transaction_session_timeout"),
"the idle transaction timeout should follow the transaction timeout")
}
func TestOpenPostgres_NoStatementTimeoutKeepsServerSetting(t *testing.T) {
for name, value := range map[string]string{"unset": "", "zero": "0"} {
t.Run(name, func(t *testing.T) {
t.Setenv(statementTimeoutEnv, value)
if value == "" {
require.NoError(t, os.Unsetenv(statementTimeoutEnv))
}
ctx := context.Background()
dsn := postgresTestDSN(t)
admin, err := pgx.Connect(ctx, dsn)
require.NoError(t, err)
t.Cleanup(func() { _ = admin.Close(ctx) })
_, err = admin.Exec(ctx, "ALTER ROLE CURRENT_USER SET statement_timeout = '7s'")
require.NoError(t, err)
conn, err := OpenPostgres(ctx, dsn, DefaultPoolConfig)
require.NoError(t, err)
t.Cleanup(func() { _ = conn.Close() })
assert.Equal(t, "7s", showSetting(t, conn.DB(nil), "statement_timeout"), "sessions should keep the server setting")
})
}
}
func TestOpenPostgres_CapsTimeoutsAtServerMaximum(t *testing.T) {
t.Setenv(transactionTimeoutEnv, "720h")
conn := openPostgresTestConn(t)
assert.Equal(t, "2147483647ms", showSetting(t, conn.DB(nil), "idle_in_transaction_session_timeout"),
"a timeout above the server maximum should be capped instead of failing the connection")
}
func TestOpenPostgresWithoutPool_BoundsSessions(t *testing.T) {
t.Setenv(statementTimeoutEnv, "500ms")
conn, err := OpenPostgresWithoutPool(context.Background(), postgresTestDSN(t))
require.NoError(t, err)
t.Cleanup(func() { _ = conn.Close() })
requireStatementTimeout(t, conn.DB(nil).Exec("SELECT pg_sleep(1)").Error)
}
func TestOpenPgxPool_BoundsSessions(t *testing.T) {
t.Setenv(statementTimeoutEnv, "500ms")
ctx := context.Background()
pool, err := OpenPgxPool(ctx, postgresTestDSN(t), DefaultPoolConfig)
require.NoError(t, err)
t.Cleanup(pool.Close)
_, err = pool.Exec(ctx, "SELECT pg_sleep(1)")
requireStatementTimeout(t, err)
}
func postgresTestDSN(t *testing.T) string {
t.Helper()
if os.Getenv("NETBIRD_STORE_ENGINE") != string(PostgresStoreEngine) {
t.Skip("requires NETBIRD_STORE_ENGINE=postgres")
}
cleanup, dsn, err := testutil.CreatePostgresTestContainer()
require.NoError(t, err)
t.Cleanup(cleanup)
return dsn
}
func openPostgresTestConn(t *testing.T) *Conn {
t.Helper()
conn, err := OpenPostgres(context.Background(), postgresTestDSN(t), DefaultPoolConfig)
require.NoError(t, err)
t.Cleanup(func() { _ = conn.Close() })
return conn
}
func requireStatementTimeout(t *testing.T, err error) {
t.Helper()
var pgErr *pgconn.PgError
require.True(t, errors.As(err, &pgErr), "expected a Postgres error, got %v", err)
assert.Equal(t, queryCanceledCode, pgErr.Code, "the statement should be cancelled by statement_timeout")
}
func showSetting(t *testing.T, db *gorm.DB, name string) string {
t.Helper()
var value string
require.NoError(t, db.Raw("SHOW "+name).Scan(&value).Error)
return value
}
management/internals/shared/grpc/server.go
+111 -22
View File
@@ -91,6 +91,12 @@ type Server struct {
reverseProxyManager rpservice.Manager
reverseProxyMu sync.RWMutex
// drainMu guards draining and pendingDisconnects. While draining, ending sync streams
// queue their disconnect for FlushDisconnects instead of writing it.
drainMu sync.Mutex
draining bool
pendingDisconnects []types.PeerSession
}
// NewServer creates a new Management server
@@ -313,7 +319,14 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
metahash := metaHash(peerMeta)
s.loginFilter.addLogin(peerKey.String(), metahash)
peer, netMap, postureChecks, dnsFwdPort, err := s.accountManager.SyncAndMarkPeer(ctx, accountID, peerKey.String(), peerMeta, realIP, syncStart)
settings, err := s.settingsManager.GetSettings(ctx, accountID, activity.SystemInitiator)
if err != nil {
log.WithContext(ctx).Debugf("error while getting settings for peer %s: %v", peerKey.String(), err)
s.syncSem.Add(-1)
return status.Errorf(codes.Internal, "error handling request")
}
peer, netMap, postureChecks, dnsFwdPort, err := s.accountManager.SyncAndMarkPeer(ctx, accountID, peerKey.String(), peerMeta, realIP, syncStart, settings)
if err != nil {
log.WithContext(ctx).Debugf("error while syncing peer %s: %v", peerKey.String(), err)
s.syncSem.Add(-1)
@@ -321,11 +334,11 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
}
trackChallenge := func() { s.accountManager.TrackCertificateChallenges(ctx, accountID, peer.ID, syncStart) }
err = s.sendInitialSync(ctx, peerKey, peer, netMap, postureChecks, srv, dnsFwdPort, trackChallenge)
err = s.sendInitialSync(ctx, peerKey, peer, netMap, postureChecks, srv, dnsFwdPort, settings, trackChallenge)
if err != nil {
log.WithContext(ctx).Debugf("error while sending initial sync for %s: %v", peerKey.String(), err)
s.syncSem.Add(-1)
s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil)
s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil, false)
return err
}
@@ -333,10 +346,17 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
if err != nil {
log.WithContext(ctx).Debugf("error while notify peer connected for %s: %v", peerKey.String(), err)
s.syncSem.Add(-1)
s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil)
s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil, false)
return err
}
if err = s.accountManager.CheckPeerMaySync(ctx, accountID, peerKey.String(), settings); err != nil {
log.WithContext(ctx).Debugf("peer %s may no longer sync after connecting: %v", peerKey.String(), err)
s.syncSem.Add(-1)
s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, updates, false)
return mapError(ctx, err)
}
s.secretsManager.SetupRefresh(ctx, accountID, peer.ID)
unlock()
@@ -455,21 +475,28 @@ func (s *Server) heartbeatPeerSession(ctx context.Context, accountID, peerID str
}
}
// cancelPeerRoutines tears down the stream under the peer lock. While draining, the
// disconnect is queued before waiting for the lock, so the shutdown flush covers the
// stream even when the lock is held past the gRPC stop timeout.
func (s *Server) cancelPeerRoutines(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time, session chan *network_map.UpdateMessage) {
uncanceledCTX := context.WithoutCancel(ctx)
disconnectQueued := s.queueDisconnect(accountID, peer, streamStartTime)
unlock := s.acquirePeerLockByUID(uncanceledCTX, peer.Key)
defer unlock()
s.cancelPeerRoutinesWithoutLock(uncanceledCTX, accountID, peer, streamStartTime, session)
s.cancelPeerRoutinesWithoutLock(uncanceledCTX, accountID, peer, streamStartTime, session, disconnectQueued)
}
// cancelPeerRoutinesWithoutLock tears down the stream of the session identified by streamStartTime
// and its updates channel. A nil session means the stream failed before it registered a channel;
// the controller then closes any channel still registered.
func (s *Server) cancelPeerRoutinesWithoutLock(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time, session chan *network_map.UpdateMessage) {
err := s.accountManager.OnPeerDisconnected(ctx, accountID, peer.Key, streamStartTime)
if err != nil {
log.WithContext(ctx).Errorf("failed to disconnect peer %s properly: %v", peer.Key, err)
// the controller then closes any channel still registered. disconnectQueued tells that the
// caller already queued the disconnect for the shutdown flush.
func (s *Server) cancelPeerRoutinesWithoutLock(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time, session chan *network_map.UpdateMessage, disconnectQueued bool) {
if !disconnectQueued && !s.queueDisconnect(accountID, peer, streamStartTime) {
err := s.accountManager.OnPeerDisconnected(ctx, accountID, peer.Key, streamStartTime)
if err != nil {
log.WithContext(ctx).Errorf("failed to disconnect peer %s properly: %v", peer.Key, err)
}
}
if !s.networkMapController.OnPeerDisconnected(ctx, accountID, peer.ID, session) {
log.WithContext(ctx).Debugf("skipped peer routines teardown for %s: a newer session owns the peer", peer.Key)
@@ -481,6 +508,60 @@ func (s *Server) cancelPeerRoutinesWithoutLock(ctx context.Context, accountID st
log.WithContext(ctx).Debugf("peer %s has been disconnected", peer.Key)
}
// queueDisconnect queues the disconnect of the stream for FlushDisconnects while the
// server drains and reports whether it did. Embedded proxy peers are never queued
// because their disconnect refreshes the network maps of the peers using them.
func (s *Server) queueDisconnect(accountID string, peer *nbpeer.Peer, streamStartTime time.Time) bool {
if peer.ProxyMeta.Embedded {
return false
}
s.drainMu.Lock()
defer s.drainMu.Unlock()
if !s.draining {
return false
}
s.pendingDisconnects = append(s.pendingDisconnects, types.PeerSession{
AccountID: accountID,
PeerID: peer.ID,
SessionStartedAt: streamStartTime.UnixNano(),
})
return true
}
// StartDraining makes ending sync streams queue their disconnect for FlushDisconnects
// instead of writing one each.
func (s *Server) StartDraining() {
s.drainMu.Lock()
defer s.drainMu.Unlock()
s.draining = true
}
// FlushDisconnects stops draining and marks the queued sessions disconnected in batched
// writes. Streams ending afterwards write their own disconnect again. It returns an error
// when not every queued session could be written.
func (s *Server) FlushDisconnects(ctx context.Context) error {
s.drainMu.Lock()
sessions := s.pendingDisconnects
s.pendingDisconnects = nil
s.draining = false
s.drainMu.Unlock()
if len(sessions) == 0 {
return nil
}
err := s.accountManager.MarkPeersDisconnected(ctx, sessions)
if err != nil && ctx.Err() == nil {
// Retrying the whole batch is safe: rows already marked no longer match their session token.
log.WithContext(ctx).Warnf("retrying to mark %d peers disconnected on shutdown: %v", len(sessions), err)
err = s.accountManager.MarkPeersDisconnected(ctx, sessions)
}
if err != nil {
return fmt.Errorf("mark %d peers disconnected on shutdown: %w", len(sessions), err)
}
return nil
}
func (s *Server) validateToken(ctx context.Context, peerKey, jwtToken string) (string, error) {
if s.authManager == nil {
return "", status.Errorf(codes.Internal, "missing auth manager")
@@ -701,6 +782,7 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
//nolint
ctx = context.WithValue(ctx, nbContext.PeerIDKey, peerKey.String())
accountID, err := s.accountManager.GetAccountIDForPeerKey(ctx, peerKey.String())
accountKnown := err == nil
if err != nil {
// this case should not happen and already indicates an issue but we don't want the system to fail due to being unable to log in detail
accountID = "UNKNOWN"
@@ -727,6 +809,15 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
sshKey = loginReq.GetPeerKeys().GetSshPubKey()
}
var settings *types.Settings
if accountKnown {
settings, err = s.settingsManager.GetSettings(ctx, accountID, activity.SystemInitiator)
if err != nil {
log.WithContext(ctx).Warnf("failed getting settings for peer %s: %s", peerKey, err)
return nil, status.Errorf(codes.Internal, "failed logging in peer")
}
}
peer, network, postureChecks, enableSSH, err := s.accountManager.LoginPeer(ctx, types.PeerLogin{
WireGuardPubKey: peerKey.String(),
SSHKey: cleanClientString(string(sshKey)),
@@ -735,6 +826,7 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
SetupKey: loginReq.GetSetupKey(),
ConnectionIP: realIP,
ExtraDNSLabels: loginReq.GetDnsLabels(),
Settings: settings,
})
if err != nil {
if errors.Is(err, internalStatus.ErrNoAuthMethodProvided) {
@@ -745,7 +837,7 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
return nil, mapError(ctx, err)
}
loginResp, err := s.prepareLoginResponse(ctx, peer, network, postureChecks, enableSSH)
loginResp, err := s.prepareLoginResponse(ctx, peer, network, postureChecks, enableSSH, settings)
if err != nil {
log.WithContext(ctx).Warnf("failed preparing login response for peer %s: %s", peerKey, err)
return nil, status.Errorf(codes.Internal, "failed logging in peer")
@@ -850,7 +942,7 @@ func (s *Server) ExtendAuthSession(ctx context.Context, req *proto.EncryptedMess
}, nil
}
func (s *Server) prepareLoginResponse(ctx context.Context, peer *nbpeer.Peer, network *types.Network, postureChecks []*nmdata.PostureChecks, enableSSH bool) (*proto.LoginResponse, error) {
func (s *Server) prepareLoginResponse(ctx context.Context, peer *nbpeer.Peer, network *types.Network, postureChecks []*nmdata.PostureChecks, enableSSH bool, settings *types.Settings) (*proto.LoginResponse, error) {
var relayToken *Token
var err error
if s.config.Relay != nil && len(s.config.Relay.Addresses) > 0 {
@@ -860,10 +952,12 @@ func (s *Server) prepareLoginResponse(ctx context.Context, peer *nbpeer.Peer, ne
}
}
settings, err := s.settingsManager.GetSettings(ctx, peer.AccountID, activity.SystemInitiator)
if err != nil {
log.WithContext(ctx).Warnf("failed getting settings for peer %s: %s", peer.Key, err)
return nil, status.Errorf(codes.Internal, "failed getting settings")
if settings == nil {
settings, err = s.settingsManager.GetSettings(ctx, peer.AccountID, activity.SystemInitiator)
if err != nil {
log.WithContext(ctx).Warnf("failed getting settings for peer %s: %s", peer.Key, err)
return nil, status.Errorf(codes.Internal, "failed getting settings")
}
}
// if peer has reached this point then it has logged in
@@ -938,7 +1032,7 @@ func (s *Server) IsHealthy(ctx context.Context, req *proto.Empty) (*proto.Empty,
}
// sendInitialSync sends initial proto.SyncResponse to the peer requesting synchronization
func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer *nbpeer.Peer, networkMap *types.NetworkMap, postureChecks []*nmdata.PostureChecks, srv proto.ManagementService_SyncServer, dnsFwdPort int64, onChallengeStamped func()) error {
func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer *nbpeer.Peer, networkMap *types.NetworkMap, postureChecks []*nmdata.PostureChecks, srv proto.ManagementService_SyncServer, dnsFwdPort int64, settings *types.Settings, onChallengeStamped func()) error {
var err error
var turnToken *Token
@@ -957,11 +1051,6 @@ func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer
}
}
settings, err := s.settingsManager.GetSettings(ctx, peer.AccountID, activity.SystemInitiator)
if err != nil {
return status.Errorf(codes.Internal, "error handling request")
}
peerGroups, err := s.accountManager.GetStore().GetPeerGroupIDs(ctx, store.LockingStrengthNone, peer.AccountID, peer.ID)
if err != nil {
return status.Errorf(codes.Internal, "failed to get peer groups %s", err)
management/internals/shared/grpc/server_disconnect_test.go
+145
View File
@@ -2,14 +2,18 @@ package grpc
import (
"context"
"errors"
"sync"
"testing"
"time"
"github.com/stretchr/testify/require"
"go.uber.org/mock/gomock"
"github.com/netbirdio/netbird/management/internals/controllers/network_map"
"github.com/netbirdio/netbird/management/server/account"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
)
func TestCancelPeerRoutines_SessionOwnership(t *testing.T) {
@@ -49,3 +53,144 @@ func TestCancelPeerRoutines_SessionOwnership(t *testing.T) {
})
}
}
func TestCancelPeerRoutines_DrainingQueuesDisconnect(t *testing.T) {
streamStart := time.Unix(1700000000, 0)
tests := []struct {
name string
embedded bool
draining bool
flushBefore bool
wantQueued bool
wantOwnWrite bool
}{
{name: "regular peer while draining is queued", draining: true, wantQueued: true},
{name: "embedded proxy peer while draining writes its own disconnect", embedded: true, draining: true, wantOwnWrite: true},
{name: "regular peer without draining writes its own disconnect", wantOwnWrite: true},
{name: "regular peer after the flush writes its own disconnect", draining: true, flushBefore: true, wantOwnWrite: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctrl := gomock.NewController(t)
accountManager := account.NewMockManager(ctrl)
controller := network_map.NewMockController(ctrl)
secretsManager := NewMockSecretsManager(ctrl)
s := &Server{accountManager: accountManager, networkMapController: controller, secretsManager: secretsManager}
peer := &nbpeer.Peer{ID: "peer-1", Key: "peer-key"}
peer.ProxyMeta.Embedded = tt.embedded
if tt.draining {
s.StartDraining()
}
if tt.flushBefore {
require.NoError(t, s.FlushDisconnects(context.Background()))
}
if tt.wantOwnWrite {
accountManager.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.Key, streamStart).Return(nil)
}
controller.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.ID, gomock.Any()).Return(true)
secretsManager.EXPECT().CancelRefresh(peer.ID)
accountManager.EXPECT().UntrackCertificateChallenges("account-1", peer.ID, streamStart)
s.cancelPeerRoutines(context.Background(), "account-1", peer, streamStart, nil)
if tt.wantQueued {
accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), []types.PeerSession{
{AccountID: "account-1", PeerID: peer.ID, SessionStartedAt: streamStart.UnixNano()},
}).Return(nil)
}
require.NoError(t, s.FlushDisconnects(context.Background()))
})
}
}
func TestFlushDisconnects_FailedBatchStopsDraining(t *testing.T) {
ctrl := gomock.NewController(t)
accountManager := account.NewMockManager(ctrl)
controller := network_map.NewMockController(ctrl)
secretsManager := NewMockSecretsManager(ctrl)
s := &Server{accountManager: accountManager, networkMapController: controller, secretsManager: secretsManager}
peer := &nbpeer.Peer{ID: "peer-1", Key: "peer-key"}
first := time.Unix(1700000000, 0)
second := first.Add(time.Second)
controller.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.ID, gomock.Any()).Return(true).Times(2)
secretsManager.EXPECT().CancelRefresh(peer.ID).Times(2)
accountManager.EXPECT().UntrackCertificateChallenges("account-1", peer.ID, gomock.Any()).Times(2)
s.StartDraining()
s.cancelPeerRoutines(context.Background(), "account-1", peer, first, nil)
accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), gomock.Any()).Return(errors.New("db down")).Times(2)
require.Error(t, s.FlushDisconnects(context.Background()))
accountManager.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.Key, second).Return(nil)
s.cancelPeerRoutines(context.Background(), "account-1", peer, second, nil)
}
func TestFlushDisconnects_RetriesFailedBatch(t *testing.T) {
ctrl := gomock.NewController(t)
accountManager := account.NewMockManager(ctrl)
controller := network_map.NewMockController(ctrl)
secretsManager := NewMockSecretsManager(ctrl)
s := &Server{accountManager: accountManager, networkMapController: controller, secretsManager: secretsManager}
peer := &nbpeer.Peer{ID: "peer-1", Key: "peer-key"}
streamStart := time.Unix(1700000000, 0)
controller.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.ID, gomock.Any()).Return(true)
secretsManager.EXPECT().CancelRefresh(peer.ID)
accountManager.EXPECT().UntrackCertificateChallenges("account-1", peer.ID, streamStart)
s.StartDraining()
s.cancelPeerRoutines(context.Background(), "account-1", peer, streamStart, nil)
want := []types.PeerSession{{AccountID: "account-1", PeerID: peer.ID, SessionStartedAt: streamStart.UnixNano()}}
gomock.InOrder(
accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), want).Return(errors.New("deadlock detected")),
accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), want).Return(nil),
)
require.NoError(t, s.FlushDisconnects(context.Background()))
}
func TestCancelPeerRoutines_QueuesDisconnectBeforeWaitingForPeerLock(t *testing.T) {
ctrl := gomock.NewController(t)
accountManager := account.NewMockManager(ctrl)
controller := network_map.NewMockController(ctrl)
secretsManager := NewMockSecretsManager(ctrl)
s := &Server{accountManager: accountManager, networkMapController: controller, secretsManager: secretsManager}
peer := &nbpeer.Peer{ID: "peer-1", Key: "peer-key"}
streamStart := time.Unix(1700000000, 0)
controller.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.ID, gomock.Any()).Return(true)
secretsManager.EXPECT().CancelRefresh(peer.ID)
accountManager.EXPECT().UntrackCertificateChallenges("account-1", peer.ID, streamStart)
unlock := s.acquirePeerLockByUID(context.Background(), peer.Key)
var unlockOnce sync.Once
t.Cleanup(func() { unlockOnce.Do(unlock) })
s.StartDraining()
done := make(chan struct{})
go func() {
defer close(done)
s.cancelPeerRoutines(context.Background(), "account-1", peer, streamStart, nil)
}()
require.Eventually(t, func() bool {
s.drainMu.Lock()
defer s.drainMu.Unlock()
return len(s.pendingDisconnects) == 1
}, time.Second, 5*time.Millisecond, "the disconnect must be queued while the teardown waits for the peer lock")
accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), []types.PeerSession{
{AccountID: "account-1", PeerID: peer.ID, SessionStartedAt: streamStart.UnixNano()},
}).Return(nil)
require.NoError(t, s.FlushDisconnects(context.Background()))
unlockOnce.Do(unlock)
<-done
}
management/server/account.go
+14 -3
View File
@@ -2033,13 +2033,13 @@ func domainIsUpToDate(domain string, domainCategory string, userAuth auth.UserAu
// derived from syncTime (the moment the gRPC stream opened). Any
// concurrent stream that started earlier loses the optimistic-lock race
// in MarkPeerConnected and bails without writing.
func (am *DefaultAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
peer, netMap, postureChecks, dnsfwdPort, err := am.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peerPubKey, Meta: meta, RealIP: realIP}, accountID)
func (am *DefaultAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
peer, netMap, postureChecks, dnsfwdPort, err := am.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peerPubKey, Meta: meta, RealIP: realIP, Settings: settings}, accountID)
if err != nil {
return nil, nil, nil, 0, fmt.Errorf("error syncing peer: %w", err)
}
if err := am.MarkPeerConnected(ctx, peerPubKey, accountID, syncTime.UnixNano(), netMap); err != nil {
if err := am.markPeerConnected(ctx, peerPubKey, accountID, syncTime.UnixNano(), netMap, settings); err != nil {
log.WithContext(ctx).Warnf("failed marking peer as connected %s %v", peerPubKey, err)
}
@@ -2059,6 +2059,17 @@ func (am *DefaultAccountManager) OnPeerDisconnected(ctx context.Context, account
return nil
}
// MarkPeersDisconnected marks the peers of the given sessions disconnected in batched writes,
// skipping every peer whose stored session differs from the given one.
func (am *DefaultAccountManager) MarkPeersDisconnected(ctx context.Context, sessions []types.PeerSession) error {
marked, err := am.Store.MarkPeersDisconnectedIfSameSession(ctx, sessions)
if err != nil {
return fmt.Errorf("mark peers disconnected: %w", err)
}
log.WithContext(ctx).Infof("marked %d of %d peers as disconnected", marked, len(sessions))
return nil
}
func (am *DefaultAccountManager) SyncPeerMeta(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) error {
accountID, err := am.Store.GetAccountIDByPeerPubKey(ctx, peerPubKey)
if err != nil {
management/server/account/manager.go
+3 -1
View File
@@ -124,8 +124,10 @@ type Manager interface {
UpdateIntegratedValidator(ctx context.Context, accountID, userID, validator string, groups []string) error
GroupValidation(ctx context.Context, accountId string, groups []string) (bool, error)
GetValidatedPeers(ctx context.Context, accountID string) (map[string]struct{}, map[string]string, error)
SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error)
SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error)
OnPeerDisconnected(ctx context.Context, accountID string, peerPubKey string, streamStartTime time.Time) error
CheckPeerMaySync(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error
MarkPeersDisconnected(ctx context.Context, sessions []types.PeerSession) error
SyncPeerMeta(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) error
FindExistingPostureCheck(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error)
GetAccountIDForPeerKey(ctx context.Context, peerKey string) (string, error)
management/server/account/manager_mock.go
+32 -4
View File
@@ -146,6 +146,20 @@ func (mr *MockManagerMockRecorder) BuildUserInfosForAccount(ctx, accountID, init
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BuildUserInfosForAccount", reflect.TypeOf((*MockManager)(nil).BuildUserInfosForAccount), ctx, accountID, initiatorUserID, accountUsers)
}
// CheckPeerMaySync mocks base method.
func (m *MockManager) CheckPeerMaySync(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "CheckPeerMaySync", ctx, accountID, peerPubKey, settings)
ret0, _ := ret[0].(error)
return ret0
}
// CheckPeerMaySync indicates an expected call of CheckPeerMaySync.
func (mr *MockManagerMockRecorder) CheckPeerMaySync(ctx, accountID, peerPubKey, settings any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckPeerMaySync", reflect.TypeOf((*MockManager)(nil).CheckPeerMaySync), ctx, accountID, peerPubKey, settings)
}
// CreateGroup mocks base method.
func (m *MockManager) CreateGroup(ctx context.Context, accountID, userID string, group *types.Group) error {
m.ctrl.T.Helper()
@@ -1369,6 +1383,20 @@ func (mr *MockManagerMockRecorder) MarkPeerDisconnected(ctx, peerKey, accountID,
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeerDisconnected", reflect.TypeOf((*MockManager)(nil).MarkPeerDisconnected), ctx, peerKey, accountID, sessionStartedAt)
}
// MarkPeersDisconnected mocks base method.
func (m *MockManager) MarkPeersDisconnected(ctx context.Context, sessions []types.PeerSession) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "MarkPeersDisconnected", ctx, sessions)
ret0, _ := ret[0].(error)
return ret0
}
// MarkPeersDisconnected indicates an expected call of MarkPeersDisconnected.
func (mr *MockManagerMockRecorder) MarkPeersDisconnected(ctx, sessions any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeersDisconnected", reflect.TypeOf((*MockManager)(nil).MarkPeersDisconnected), ctx, sessions)
}
// OnPeerDisconnected mocks base method.
func (m *MockManager) OnPeerDisconnected(ctx context.Context, accountID, peerPubKey string, streamStartTime time.Time) error {
m.ctrl.T.Helper()
@@ -1569,9 +1597,9 @@ func (mr *MockManagerMockRecorder) StoreEvent(ctx, initiatorID, targetID, accoun
}
// SyncAndMarkPeer mocks base method.
func (m *MockManager) SyncAndMarkPeer(ctx context.Context, accountID, peerPubKey string, meta peer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*peer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
func (m *MockManager) SyncAndMarkPeer(ctx context.Context, accountID, peerPubKey string, meta peer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*peer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "SyncAndMarkPeer", ctx, accountID, peerPubKey, meta, realIP, syncTime)
ret := m.ctrl.Call(m, "SyncAndMarkPeer", ctx, accountID, peerPubKey, meta, realIP, syncTime, settings)
ret0, _ := ret[0].(*peer.Peer)
ret1, _ := ret[1].(*types.NetworkMap)
ret2, _ := ret[2].([]*nmdata.PostureChecks)
@@ -1581,9 +1609,9 @@ func (m *MockManager) SyncAndMarkPeer(ctx context.Context, accountID, peerPubKey
}
// SyncAndMarkPeer indicates an expected call of SyncAndMarkPeer.
func (mr *MockManagerMockRecorder) SyncAndMarkPeer(ctx, accountID, peerPubKey, meta, realIP, syncTime any) *gomock.Call {
func (mr *MockManagerMockRecorder) SyncAndMarkPeer(ctx, accountID, peerPubKey, meta, realIP, syncTime, settings any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SyncAndMarkPeer", reflect.TypeOf((*MockManager)(nil).SyncAndMarkPeer), ctx, accountID, peerPubKey, meta, realIP, syncTime)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SyncAndMarkPeer", reflect.TypeOf((*MockManager)(nil).SyncAndMarkPeer), ctx, accountID, peerPubKey, meta, realIP, syncTime, settings)
}
// SyncPeer mocks base method.
management/server/account_test.go
+1 -1
View File
@@ -3977,7 +3977,7 @@ func BenchmarkSyncAndMarkPeer(b *testing.B) {
b.ResetTimer()
start := time.Now()
for i := 0; i < b.N; i++ {
_, _, _, _, err := manager.SyncAndMarkPeer(context.Background(), account.Id, account.Peers["peer-1"].Key, nbpeer.PeerSystemMeta{Hostname: strconv.Itoa(i)}, net.IP{1, 1, 1, 1}, time.Now().UTC())
_, _, _, _, err := manager.SyncAndMarkPeer(context.Background(), account.Id, account.Peers["peer-1"].Key, nbpeer.PeerSystemMeta{Hostname: strconv.Itoa(i)}, net.IP{1, 1, 1, 1}, time.Now().UTC(), nil)
assert.NoError(b, err)
}
management/server/management_proto_test.go
+129
View File
@@ -17,8 +17,10 @@ import (
"go.uber.org/mock/gomock"
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/keepalive"
"google.golang.org/grpc/status"
"github.com/netbirdio/netbird/encryption"
"github.com/netbirdio/netbird/formatter/hook"
@@ -824,3 +826,130 @@ func testLoginPerformance(t *testing.T, loginCalls []func() error) {
tAvg := tSum / time.Duration(len(durations))
t.Logf("Min: %v, Max: %v, Avg: %v", tMin, tMax, tAvg)
}
type expireOnConnectStore struct {
store.Store
once sync.Once
err error
}
func (s *expireOnConnectStore) MarkPeerConnectedIfNewerSession(ctx context.Context, accountID, peerID string, newSessionStartedAt int64) (bool, error) {
updated, err := s.Store.MarkPeerConnectedIfNewerSession(ctx, accountID, peerID, newSessionStartedAt)
if err != nil {
return updated, err
}
s.once.Do(func() {
peer, getErr := s.Store.GetPeerByID(ctx, store.LockingStrengthNone, accountID, peerID)
if getErr != nil {
s.err = getErr
return
}
peer.Status.LoginExpired = true
s.err = s.Store.SavePeerStatus(ctx, accountID, peerID, *peer.Status)
})
return updated, err
}
func Test_SyncRejectsPeerExpiredWhileConnecting(t *testing.T) {
mgmtServer, am, mgmtAddr, cleanup, err := startManagementForTest(t, "testdata/store_with_expired_peers.sql", &config.Config{
Datadir: t.TempDir(),
})
defer cleanup()
require.NoError(t, err)
defer mgmtServer.GracefulStop()
client, clientConn, err := createRawClient(mgmtAddr)
require.NoError(t, err)
defer clientConn.Close()
key, err := wgtypes.GeneratePrivateKey()
require.NoError(t, err)
_, err = loginPeerWithValidSetupKey(key, client)
require.NoError(t, err)
serverKey, err := getServerKey(client)
require.NoError(t, err)
decorated := &expireOnConnectStore{Store: am.Store}
am.Store = decorated
message, err := encryption.EncryptMessage(*serverKey, key, &mgmtProto.SyncRequest{Meta: &mgmtProto.PeerSystemMeta{}})
require.NoError(t, err)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
sync, err := client.Sync(ctx, &mgmtProto.EncryptedMessage{
WgPubKey: key.PublicKey().String(),
Body: message,
})
require.NoError(t, err)
resp := &mgmtProto.EncryptedMessage{}
for err == nil {
err = sync.RecvMsg(resp)
}
require.Equal(t, codes.PermissionDenied, status.Code(err), "sync stream ended with: %v", err)
require.NoError(t, decorated.err)
peer, err := am.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, key.PublicKey().String())
require.NoError(t, err)
require.True(t, peer.Status.LoginExpired)
require.Eventually(t, func() bool {
peer, err = am.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, key.PublicKey().String())
return err == nil && !peer.Status.Connected
}, 5*time.Second, 50*time.Millisecond)
}
type settingsCountingStore struct {
store.Store
reads atomic.Int32
}
func (s *settingsCountingStore) GetAccountSettings(ctx context.Context, lockStrength store.LockingStrength, accountID string) (*types.Settings, error) {
s.reads.Add(1)
return s.Store.GetAccountSettings(ctx, lockStrength, accountID)
}
func Test_LoginAndSyncReuseRequestSettings(t *testing.T) {
mgmtServer, am, mgmtAddr, cleanup, err := startManagementForTest(t, "testdata/store_with_expired_peers.sql", &config.Config{
Datadir: t.TempDir(),
})
defer cleanup()
require.NoError(t, err)
defer mgmtServer.GracefulStop()
client, clientConn, err := createRawClient(mgmtAddr)
require.NoError(t, err)
defer clientConn.Close()
key, err := wgtypes.GeneratePrivateKey()
require.NoError(t, err)
_, err = loginPeerWithValidSetupKey(key, client)
require.NoError(t, err)
counting := &settingsCountingStore{Store: am.Store}
am.Store = counting
_, err = loginPeerWithValidSetupKey(key, client)
require.NoError(t, err)
require.Zero(t, counting.reads.Load(), "login must not read settings through the account manager")
serverKey, err := getServerKey(client)
require.NoError(t, err)
message, err := encryption.EncryptMessage(*serverKey, key, &mgmtProto.SyncRequest{Meta: &mgmtProto.PeerSystemMeta{}})
require.NoError(t, err)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
sync, err := client.Sync(ctx, &mgmtProto.EncryptedMessage{WgPubKey: key.PublicKey().String(), Body: message})
require.NoError(t, err)
require.NoError(t, sync.RecvMsg(&mgmtProto.EncryptedMessage{}))
require.Eventually(t, func() bool {
peer, err := counting.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, key.PublicKey().String())
return err == nil && peer.Status.Connected
}, 5*time.Second, 50*time.Millisecond)
require.Zero(t, counting.reads.Load(), "sync must not read settings through the account manager")
}
management/server/mock_server/account_mock.go
+19 -3
View File
@@ -42,7 +42,9 @@ type MockAccountManager struct {
GetPeersFunc func(ctx context.Context, accountID, userID, nameFilter, ipFilter, macFilter string) ([]*nbpeer.Peer, int, error)
MarkPeerConnectedFunc func(ctx context.Context, peerKey string, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error
MarkPeerDisconnectedFunc func(ctx context.Context, peerKey string, accountID string, sessionStartedAt int64) error
SyncAndMarkPeerFunc func(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error)
MarkPeersDisconnectedFunc func(ctx context.Context, sessions []types.PeerSession) error
CheckPeerMaySyncFunc func(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error
SyncAndMarkPeerFunc func(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error)
DeletePeerFunc func(ctx context.Context, accountID, peerKey, userID string) error
GetNetworkMapFunc func(ctx context.Context, peerKey string) (*types.NetworkMap, error)
GetPeerNetworkFunc func(ctx context.Context, peerKey string) (*types.Network, error)
@@ -242,13 +244,27 @@ func (am *MockAccountManager) DeleteSetupKey(ctx context.Context, accountID, use
return status.Errorf(codes.Unimplemented, "method DeleteSetupKey is not implemented")
}
func (am *MockAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
func (am *MockAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time, settings *types.Settings) (*nbpeer.Peer, *types.NetworkMap, []*nmdata.PostureChecks, int64, error) {
if am.SyncAndMarkPeerFunc != nil {
return am.SyncAndMarkPeerFunc(ctx, accountID, peerPubKey, meta, realIP, syncTime)
return am.SyncAndMarkPeerFunc(ctx, accountID, peerPubKey, meta, realIP, syncTime, settings)
}
return nil, nil, nil, 0, status.Errorf(codes.Unimplemented, "method MarkPeerConnected is not implemented")
}
func (am *MockAccountManager) MarkPeersDisconnected(ctx context.Context, sessions []types.PeerSession) error {
if am.MarkPeersDisconnectedFunc != nil {
return am.MarkPeersDisconnectedFunc(ctx, sessions)
}
return nil
}
func (am *MockAccountManager) CheckPeerMaySync(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error {
if am.CheckPeerMaySyncFunc != nil {
return am.CheckPeerMaySyncFunc(ctx, accountID, peerPubKey, settings)
}
return status.Errorf(codes.Unimplemented, "method CheckPeerMaySync is not implemented")
}
func (am *MockAccountManager) OnPeerDisconnected(ctx context.Context, accountID string, peerPubKey string, streamStartTime time.Time) error {
// Mirror DefaultAccountManager.OnPeerDisconnected: drive the fencing
// hook so tests that inject MarkPeerDisconnectedFunc actually observe
management/server/peer.go
+183 -89
View File
@@ -88,6 +88,12 @@ func (am *DefaultAccountManager) GetPeers(ctx context.Context, accountID, userID
// Disconnects use MarkPeerDisconnected and require the session to match
// exactly; see PeerStatus.SessionStartedAt for the protocol.
func (am *DefaultAccountManager) MarkPeerConnected(ctx context.Context, peerPubKey string, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error {
return am.markPeerConnected(ctx, peerPubKey, accountID, sessionStartedAt, nmap, nil)
}
// markPeerConnected is MarkPeerConnected with the account settings the caller already
// read; nil settings are loaded when needed.
func (am *DefaultAccountManager) markPeerConnected(ctx context.Context, peerPubKey string, accountID string, sessionStartedAt int64, nmap *types.NetworkMap, settings *types.Settings) error {
start := time.Now()
defer func() {
am.metrics.AccountManagerMetrics().RecordPeerStatusUpdateDuration(telemetry.PeerStatusConnect, time.Since(start))
@@ -119,19 +125,13 @@ func (am *DefaultAccountManager) MarkPeerConnected(ctx context.Context, peerPubK
log.WithContext(ctx).Debugf("mark peer %s connected", peer.ID)
if err = am.schedulePeerExpirations(ctx, accountID, peer); err != nil {
if err = am.schedulePeerExpirations(ctx, accountID, peer, settings); err != nil {
return err
}
// A login-expired peer reconnecting, or an embedded proxy peer flipping to
// connected (which triggers SynthesizePrivateServiceZones), must refresh the
// peers reachable from it. The embedded-proxy fan-out tolerates a dispatch error.
if peer.Status != nil && peer.Status.LoginExpired {
affectedPeerIDs := am.markConnectedAffectedPeers(ctx, accountID, peer.ID, nmap)
if err = am.networkMapController.OnPeersUpdated(ctx, accountID, []string{peer.ID}, affectedPeerIDs); err != nil {
return fmt.Errorf("notify network map controller of peer update: %w", err)
}
}
// An embedded proxy peer flipping to connected (which triggers
// SynthesizePrivateServiceZones) must refresh the peers reachable from it.
// The fan-out tolerates a dispatch error.
if peer.ProxyMeta.Embedded {
affectedPeerIDs := am.markConnectedAffectedPeers(ctx, accountID, peer.ID, nmap)
if err := am.networkMapController.OnPeersUpdated(ctx, accountID, []string{peer.ID}, affectedPeerIDs); err != nil {
@@ -166,11 +166,11 @@ func (am *DefaultAccountManager) writePeerStatus(ctx context.Context, accountID
// schedulePeerExpirations arms the account's login/inactivity expiration
// timers for an SSO peer that just connected.
func (am *DefaultAccountManager) schedulePeerExpirations(ctx context.Context, accountID string, peer *nbpeer.Peer) error {
func (am *DefaultAccountManager) schedulePeerExpirations(ctx context.Context, accountID string, peer *nbpeer.Peer, settings *types.Settings) error {
if !peer.AddedWithSSOLogin() {
return nil
}
settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
settings, err := am.accountSettings(ctx, accountID, settings)
if err != nil {
return err
}
@@ -1087,37 +1087,22 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
var metaDiff nbpeer.MetaDiff
var err error
settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
settings, err := am.accountSettings(ctx, accountID, sync.Settings)
if err != nil {
return nil, nil, nil, 0, err
}
err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
peer, err = transaction.GetPeerByPeerPubKey(ctx, store.LockingStrengthUpdate, sync.WireGuardPubKey)
if err != nil {
return status.NewPeerNotRegisteredError()
}
peer, metaDiff, ipv6CapabilityChanged, err = am.prepareSyncPeer(ctx, am.Store, store.LockingStrengthNone, sync, settings)
if err != nil {
return nil, nil, nil, 0, err
}
if peer.UserID != "" {
user, err := transaction.GetUserByUserID(ctx, store.LockingStrengthNone, peer.UserID)
if err != nil {
if metaDiff.Updated() {
err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
peer, metaDiff, ipv6CapabilityChanged, err = am.prepareSyncPeer(ctx, transaction, store.LockingStrengthUpdate, sync, settings)
if err != nil || !metaDiff.Updated() {
return err
}
if err = checkIfPeerOwnerIsBlocked(peer, user); err != nil {
return err
}
}
if peerLoginExpired(ctx, peer, settings) {
return status.NewPeerLoginExpiredError()
}
oldHasIPv6Cap := peer.HasCapability(nbpeer.PeerCapabilityIPv6Overlay)
newLocation := am.resolvePeerLocation(ctx, peer, sync.RealIP)
metaDiff = peer.UpdateMetaIfNew(ctx, sync.Meta, newLocation)
ipv6CapabilityChanged = oldHasIPv6Cap != peer.HasCapability(nbpeer.PeerCapabilityIPv6Overlay)
if metaDiff.Updated() {
am.metrics.AccountManagerMetrics().CountPeerMetUpdate()
log.WithContext(ctx).Tracef("peer %s metadata updated", peer.ID)
if err = transaction.SavePeer(ctx, accountID, peer); err != nil {
@@ -1126,11 +1111,11 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
if err = transaction.IncrementNetworkSerial(ctx, accountID); err != nil {
return fmt.Errorf("increment network serial: %w", err)
}
return nil
})
if err != nil {
return nil, nil, nil, 0, err
}
return nil
})
if err != nil {
return nil, nil, nil, 0, err
}
peerGroupIDs, err := getPeerGroupIDs(ctx, am.Store, accountID, peer.ID)
@@ -1160,6 +1145,82 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
return peer, nmap, resPostureChecks, dnsFwdPort, nil
}
// prepareSyncPeer loads the syncing peer with lockStrength, checks that it may sync and
// applies the reported meta to it in memory. The returned diff tells whether the peer
// has to be saved.
func (am *DefaultAccountManager) prepareSyncPeer(ctx context.Context, st store.Store, lockStrength store.LockingStrength, sync types.PeerSync, settings *types.Settings) (*nbpeer.Peer, nbpeer.MetaDiff, bool, error) {
peer, err := st.GetPeerByPeerPubKey(ctx, lockStrength, sync.WireGuardPubKey)
if err != nil {
return nil, nbpeer.MetaDiff{}, false, peerLookupError(err)
}
if err = checkPeerMaySync(ctx, st, peer, settings); err != nil {
return nil, nbpeer.MetaDiff{}, false, err
}
oldHasIPv6Cap := peer.HasCapability(nbpeer.PeerCapabilityIPv6Overlay)
newLocation := am.resolvePeerLocation(ctx, peer, sync.RealIP)
metaDiff := peer.UpdateMetaIfNew(ctx, sync.Meta, newLocation)
ipv6CapabilityChanged := oldHasIPv6Cap != peer.HasCapability(nbpeer.PeerCapabilityIPv6Overlay)
return peer, metaDiff, ipv6CapabilityChanged, nil
}
// CheckPeerMaySync reads the peer again and rejects it like SyncPeer does when its owner
// is blocked or its login expired. Nil settings are loaded from the store. The gRPC
// server calls it once the sync stream can receive updates, so an expiry committed
// while the stream was being set up is not missed.
func (am *DefaultAccountManager) CheckPeerMaySync(ctx context.Context, accountID, peerPubKey string, settings *types.Settings) error {
settings, err := am.accountSettings(ctx, accountID, settings)
if err != nil {
return err
}
peer, err := am.Store.GetPeerByPeerPubKey(ctx, store.LockingStrengthNone, peerPubKey)
if err != nil {
return peerLookupError(err)
}
return checkPeerMaySync(ctx, am.Store, peer, settings)
}
// peerLookupError maps a missing peer to the not-registered error that makes the client
// register again, and keeps any other store error so a transient failure is retried.
func peerLookupError(err error) error {
if s, ok := status.FromError(err); ok && s.Type() == status.NotFound {
return status.NewPeerNotRegisteredError()
}
return err
}
// accountSettings returns settings when the caller already read them for this request
// and loads them from the store otherwise.
func (am *DefaultAccountManager) accountSettings(ctx context.Context, accountID string, settings *types.Settings) (*types.Settings, error) {
if settings != nil {
return settings, nil
}
return am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
}
// checkPeerMaySync rejects a peer whose owner is blocked or whose login expired.
func checkPeerMaySync(ctx context.Context, st store.Store, peer *nbpeer.Peer, settings *types.Settings) error {
if peer.UserID != "" {
user, err := st.GetUserByUserID(ctx, store.LockingStrengthNone, peer.UserID)
if err != nil {
return err
}
if err = checkIfPeerOwnerIsBlocked(peer, user); err != nil {
return err
}
}
if peerLoginExpired(ctx, peer, settings) {
return status.NewPeerLoginExpiredError()
}
return nil
}
// metaDiffAffectsPosture reports whether the meta change flips the verdict of any of
// the peer's posture checks, replaying them against the old and new state.
func metaDiffAffectsPosture(diff *nbpeer.MetaDiff, checks []*nmdata.PostureChecks) bool {
@@ -1206,8 +1267,8 @@ func (am *DefaultAccountManager) syncPeerAffectedPeers(ctx context.Context, acco
return affectedPeerIDsFromNetworkMap(nmap, peerID)
}
// markConnectedAffectedPeers resolves the peers affected when a peer connects
// (login-expiry reconnect or embedded-proxy connect). The connecting peer's
// markConnectedAffectedPeers resolves the peers affected when an embedded proxy
// peer connects. The connecting peer's
// network map already lists them bidirectionally — the synthesized
// private-service policy puts proxy access-group members in the proxy peer's own
// map, and these edges carry no source-posture gate. An invalid peer has an
@@ -1249,8 +1310,13 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login types.Peer
// when the client sends a login request with a JWT which is used to get the user ID,
// it means that the client has already checked if it needs login and had been through the SSO flow
// so, we can skip this check and directly proceed with the login
settings, err := am.accountSettings(ctx, accountID, login.Settings)
if err != nil {
return nil, nil, nil, false, err
}
if login.UserID == "" {
err = am.checkIFPeerNeedsLoginWithoutLock(ctx, accountID, login)
err = am.checkIFPeerNeedsLoginWithoutLock(ctx, login, settings)
if err != nil {
return nil, nil, nil, false, err
}
@@ -1260,56 +1326,55 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login types.Peer
var shouldStorePeer, shouldUpdatePeers bool
var peerGroupIDs []string
settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
peer, err = am.Store.GetPeerByPeerPubKey(ctx, store.LockingStrengthNone, login.WireGuardPubKey)
if err != nil {
return nil, nil, nil, false, err
}
if err = am.checkPeerLogin(ctx, am.Store, peer, login); err != nil {
return nil, nil, nil, false, err
}
err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
peer, err = transaction.GetPeerByPeerPubKey(ctx, store.LockingStrengthUpdate, login.WireGuardPubKey)
if err != nil {
return err
}
if login.UserID != "" {
if peer.UserID != login.UserID {
log.Warnf("user mismatch when logging in peer %s: peer user %s, login user %s ", peer.ID, peer.UserID, login.UserID)
return status.NewPeerLoginMismatchError()
}
changed, err := am.handleUserPeer(ctx, transaction, peer, settings)
if peerLoginChangesPeer(ctx, peer, login, settings) {
err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
peer, err = transaction.GetPeerByPeerPubKey(ctx, store.LockingStrengthUpdate, login.WireGuardPubKey)
if err != nil {
return err
}
if changed {
shouldStorePeer = true
shouldUpdatePeers = true
}
}
if peer.SSHKey != login.SSHKey {
peer.SSHKey = login.SSHKey
shouldStorePeer = true
}
if !peer.AllowExtraDNSLabels && len(login.ExtraDNSLabels) > 0 {
return status.Errorf(status.PreconditionFailed, "couldn't login peer: setup key doesn't allow extra DNS labels")
}
if shouldStorePeer {
if err = transaction.SavePeer(ctx, accountID, peer); err != nil {
if err = am.checkPeerLogin(ctx, transaction, peer, login); err != nil {
return err
}
if err = transaction.IncrementNetworkSerial(ctx, accountID); err != nil {
return fmt.Errorf("increment network serial: %w", err)
}
}
return nil
})
if err != nil {
return nil, nil, nil, false, err
if login.UserID != "" {
changed, err := am.handleUserPeer(ctx, transaction, peer, settings)
if err != nil {
return err
}
if changed {
shouldStorePeer = true
shouldUpdatePeers = true
}
}
if peer.SSHKey != login.SSHKey {
peer.SSHKey = login.SSHKey
shouldStorePeer = true
}
if shouldStorePeer {
if err = transaction.SavePeer(ctx, accountID, peer); err != nil {
return err
}
if err = transaction.IncrementNetworkSerial(ctx, accountID); err != nil {
return fmt.Errorf("increment network serial: %w", err)
}
}
return nil
})
if err != nil {
return nil, nil, nil, false, err
}
}
// This is needed to keep in memory for the peer config. Otherwise browser client will end in a retry loop
@@ -1341,6 +1406,40 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login types.Peer
return peer, network, postureChecks, enableSSH, nil
}
// checkPeerLogin rejects a login whose user does not own the peer, whose owner is
// blocked, or that requests extra DNS labels the peer is not allowed to use.
func (am *DefaultAccountManager) checkPeerLogin(ctx context.Context, st store.Store, peer *nbpeer.Peer, login types.PeerLogin) error {
if login.UserID != "" {
if peer.UserID != login.UserID {
log.Warnf("user mismatch when logging in peer %s: peer user %s, login user %s ", peer.ID, peer.UserID, login.UserID)
return status.NewPeerLoginMismatchError()
}
user, err := st.GetUserByUserID(ctx, store.LockingStrengthNone, peer.UserID)
if err != nil {
return err
}
if err = checkIfPeerOwnerIsBlocked(peer, user); err != nil {
return err
}
}
if !peer.AllowExtraDNSLabels && len(login.ExtraDNSLabels) > 0 {
return status.Errorf(status.PreconditionFailed, "couldn't login peer: setup key doesn't allow extra DNS labels")
}
return nil
}
// peerLoginChangesPeer reports whether logging in writes the peer: an SSO
// re-login of an expired peer or a new SSH key.
func peerLoginChangesPeer(ctx context.Context, peer *nbpeer.Peer, login types.PeerLogin, settings *types.Settings) bool {
if peer.SSHKey != login.SSHKey {
return true
}
return login.UserID != "" && peerLoginExpired(ctx, peer, settings)
}
// ExtendPeerSession refreshes the peer's SSO session deadline by updating
// LastLogin after a successful JWT validation. The tunnel is untouched: no
// network map sync, no peer reconnect.
@@ -1507,7 +1606,7 @@ func processPeerPostureChecks(policy *types.Policy, peerID string, peerGroupIDs
// The NetBird client doesn't have a way to check if the peer needs login besides sending a login request
// with no JWT token and usually no setup-key. As the client can send up to two login request to check if it is expired
// and before starting the engine, we do the checks without an account lock to avoid piling up requests.
func (am *DefaultAccountManager) checkIFPeerNeedsLoginWithoutLock(ctx context.Context, accountID string, login types.PeerLogin) error {
func (am *DefaultAccountManager) checkIFPeerNeedsLoginWithoutLock(ctx context.Context, login types.PeerLogin, settings *types.Settings) error {
peer, err := am.Store.GetPeerByPeerPubKey(ctx, store.LockingStrengthNone, login.WireGuardPubKey)
if err != nil {
return err
@@ -1519,11 +1618,6 @@ func (am *DefaultAccountManager) checkIFPeerNeedsLoginWithoutLock(ctx context.Co
return nil
}
settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
if err != nil {
return err
}
if peerLoginExpired(ctx, peer, settings) {
return status.NewPeerLoginExpiredError()
}
management/server/peer_test.go
+134
View File
@@ -17,6 +17,7 @@ import (
"strconv"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
@@ -3323,3 +3324,136 @@ func TestPeerWrites_BumpNetworkSerial(t *testing.T) {
assert.Greater(t, networkSerial(t, manager, account.Id), afterDisconnect, "embedded proxy connect must bump")
})
}
type txCountingStore struct {
store.Store
txs atomic.Int32
}
func (s *txCountingStore) ExecuteInTransaction(ctx context.Context, f func(store.Store) error) error {
s.txs.Add(1)
return s.Store.ExecuteInTransaction(ctx, f)
}
func TestPeerSyncAndLogin_LockOnlyOnChange(t *testing.T) {
manager, _, account, peer1, _, _ := setupNetworkMapTest(t)
ctx := context.Background()
counting := &txCountingStore{Store: manager.Store}
manager.Store = counting
current, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
require.NoError(t, err)
t.Run("SyncPeer without meta change opens no transaction", func(t *testing.T) {
counting.txs.Store(0)
_, _, _, _, err := manager.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peer1.Key, Meta: current.Meta}, account.Id)
require.NoError(t, err)
assert.Zero(t, counting.txs.Load())
})
t.Run("SyncPeer with meta change saves in one transaction", func(t *testing.T) {
counting.txs.Store(0)
meta := current.Meta
meta.Hostname = "lock-on-change-host"
_, _, _, _, err := manager.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peer1.Key, Meta: meta}, account.Id)
require.NoError(t, err)
assert.Equal(t, int32(1), counting.txs.Load())
saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
require.NoError(t, err)
assert.Equal(t, "lock-on-change-host", saved.Meta.Hostname)
})
t.Run("LoginPeer without change opens no transaction", func(t *testing.T) {
saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
require.NoError(t, err)
counting.txs.Store(0)
_, _, _, _, err = manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, SSHKey: saved.SSHKey, Meta: saved.Meta})
require.NoError(t, err)
assert.Zero(t, counting.txs.Load())
})
t.Run("LoginPeer with new ssh key saves in one transaction", func(t *testing.T) {
counting.txs.Store(0)
_, _, _, _, err := manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, SSHKey: "ssh-ed25519 AAAAlockonchange", Meta: current.Meta})
require.NoError(t, err)
assert.Equal(t, int32(1), counting.txs.Load())
saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
require.NoError(t, err)
assert.Equal(t, "ssh-ed25519 AAAAlockonchange", saved.SSHKey)
})
t.Run("LoginPeer rejects a user mismatch without a transaction", func(t *testing.T) {
counting.txs.Store(0)
_, _, _, _, err := manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, UserID: "someone-else", SSHKey: "ssh-ed25519 AAAAother", Meta: current.Meta})
assert.Equal(t, status.NewPeerLoginMismatchError(), err)
assert.Zero(t, counting.txs.Load())
saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
require.NoError(t, err)
assert.Equal(t, "ssh-ed25519 AAAAlockonchange", saved.SSHKey, "a rejected login must not store its ssh key")
})
t.Run("LoginPeer rejects disallowed extra DNS labels without a transaction", func(t *testing.T) {
counting.txs.Store(0)
_, _, _, _, err := manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, SSHKey: "ssh-ed25519 AAAAlabels", ExtraDNSLabels: []string{"extra"}, Meta: current.Meta})
errStatus, ok := status.FromError(err)
require.True(t, ok, "expected a status error, got %v", err)
assert.Equal(t, status.PreconditionFailed, errStatus.Type())
assert.Zero(t, counting.txs.Load())
})
}
func TestCheckPeerMaySync(t *testing.T) {
manager, _, account, peer1, peer2, _ := setupNetworkMapTest(t)
ctx := context.Background()
require.NoError(t, manager.CheckPeerMaySync(ctx, account.Id, peer1.Key, nil))
expired, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
require.NoError(t, err)
expired.Status.LoginExpired = true
require.NoError(t, manager.Store.SavePeer(ctx, account.Id, expired))
assert.Equal(t, status.NewPeerLoginExpiredError(), manager.CheckPeerMaySync(ctx, account.Id, peer1.Key, nil))
owner := types.NewRegularUser("blocked-owner", "", "")
owner.AccountID = account.Id
owner.Blocked = true
require.NoError(t, manager.Store.SaveUser(ctx, owner))
blocked, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer2.ID)
require.NoError(t, err)
blocked.UserID = owner.Id
require.NoError(t, manager.Store.SavePeer(ctx, account.Id, blocked))
err = manager.CheckPeerMaySync(ctx, account.Id, peer2.Key, nil)
errStatus, ok := status.FromError(err)
require.True(t, ok, "expected a status error, got %v", err)
assert.Equal(t, status.PermissionDenied, errStatus.Type())
assert.Equal(t, status.NewPeerNotRegisteredError(), manager.CheckPeerMaySync(ctx, account.Id, "unknown-key", nil))
}
type failingPeerLookupStore struct {
store.Store
}
func (s *failingPeerLookupStore) GetPeerByPeerPubKey(context.Context, store.LockingStrength, string) (*nbpeer.Peer, error) {
return nil, status.Errorf(status.Internal, "issue getting peer from store: connection reset")
}
func TestPeerLookupErrorKeepsTransientStoreErrors(t *testing.T) {
manager, _, account, peer1, _, _ := setupNetworkMapTest(t)
ctx := context.Background()
manager.Store = &failingPeerLookupStore{Store: manager.Store}
err := manager.CheckPeerMaySync(ctx, account.Id, peer1.Key, nil)
errStatus, ok := status.FromError(err)
require.True(t, ok, "expected a status error, got %v", err)
assert.Equal(t, status.Internal, errStatus.Type(), "a store failure must not look like an unregistered peer")
_, _, _, _, err = manager.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peer1.Key}, account.Id)
errStatus, ok = status.FromError(err)
require.True(t, ok, "expected a status error, got %v", err)
assert.Equal(t, status.Internal, errStatus.Type(), "a store failure must not look like an unregistered peer")
}
management/server/store/sql_store_peer.go
+45 -1
View File
@@ -106,7 +106,6 @@ func (s *SqlStore) MarkPeerConnectedIfNewerSession(ctx context.Context, accountI
"peer_status_connected": true,
"peer_status_last_seen": gorm.Expr("CURRENT_TIMESTAMP"),
"peer_status_session_started_at": newSessionStartedAt,
"peer_status_login_expired": false,
})
if result.Error != nil {
return false, status.Errorf(status.Internal, "mark peer connected: %v", result.Error)
@@ -142,6 +141,51 @@ func (s *SqlStore) MarkPeerDisconnectedIfSameSession(ctx context.Context, accoun
return result.RowsAffected > 0, nil
}
// markPeersDisconnectedChunkSize bounds the sessions per statement so the
// bound parameter count stays below the database limits.
const markPeersDisconnectedChunkSize = 1000
// MarkPeersDisconnectedIfSameSession is the batched form of
// MarkPeerDisconnectedIfSameSession: one UPDATE per chunk of sessions, each
// peer guarded by its own (account, peer, session) tuple. A failed chunk does
// not stop the others, only a done ctx does; the count then covers the chunks
// that succeeded and the error joins the failures.
func (s *SqlStore) MarkPeersDisconnectedIfSameSession(ctx context.Context, sessions []types.PeerSession) (int64, error) {
var total int64
var errs []error
for start := 0; start < len(sessions); start += markPeersDisconnectedChunkSize {
if err := ctx.Err(); err != nil {
errs = append(errs, status.Errorf(status.Internal, "mark peers disconnected: %v", err))
break
}
end := min(start+markPeersDisconnectedChunkSize, len(sessions))
tuples := make([][]any, 0, end-start)
for _, session := range sessions[start:end] {
if session.SessionStartedAt == 0 {
continue
}
tuples = append(tuples, []any{session.AccountID, session.PeerID, session.SessionStartedAt})
}
if len(tuples) == 0 {
continue
}
result := s.db.WithContext(ctx).
Model(&nbpeer.Peer{}).
Where("(account_id, id, peer_status_session_started_at) IN ?", tuples).
Updates(map[string]any{
"peer_status_connected": false,
"peer_status_last_seen": gorm.Expr("CURRENT_TIMESTAMP"),
"peer_status_session_started_at": int64(0),
})
if result.Error != nil {
errs = append(errs, status.Errorf(status.Internal, "mark peers disconnected: %v", result.Error))
continue
}
total += result.RowsAffected
}
return total, errors.Join(errs...)
}
// ApproveAccountPeers marks all peers that currently require approval in the given account as approved.
func (s *SqlStore) ApproveAccountPeers(ctx context.Context, accountID string) (int, error) {
result := s.db.Model(&nbpeer.Peer{}).
management/server/store/sql_store_peer_test.go
+149
View File
@@ -3,16 +3,20 @@ package store
import (
"context"
"encoding/binary"
"errors"
"fmt"
"net"
"net/netip"
"reflect"
"sort"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
@@ -941,3 +945,148 @@ func TestSqlStore_ApproveAccountPeers(t *testing.T) {
})
})
}
func TestSqlStore_MarkPeersDisconnectedIfSameSession(t *testing.T) {
ctx := context.Background()
store, cleanUp, err := NewTestStoreFromSQL(ctx, "../testdata/extended-store.sql", t.TempDir())
t.Cleanup(cleanUp)
require.NoError(t, err)
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
sqlStore := store.(*SqlStore)
sessionOf := map[string]int64{"peer-a": 101, "peer-b": 102, "peer-c": 103, "peer-d": 104}
for id, session := range sessionOf {
require.NoError(t, store.AddPeerToAccount(ctx, &nbpeer.Peer{
ID: id,
AccountID: accountID,
Key: "key-" + id,
DNSLabel: id,
IP: netip.AddrFrom4([4]byte{10, 9, 0, byte(session)}),
IPv6: netip.MustParseAddr(fmt.Sprintf("fd00::9:%d", session)),
}))
require.NoError(t, sqlStore.db.Model(&nbpeer.Peer{}).Where("id = ?", id).Updates(map[string]any{
"peer_status_connected": true,
"peer_status_session_started_at": session,
"peer_status_last_seen": time.Now().UTC(),
}).Error)
}
sessions := []types.PeerSession{{AccountID: accountID, PeerID: "peer-a", SessionStartedAt: 101}}
for i := 0; i < markPeersDisconnectedChunkSize+5; i++ {
sessions = append(sessions, types.PeerSession{AccountID: accountID, PeerID: fmt.Sprintf("missing-%d", i), SessionStartedAt: 1})
}
sessions = append(sessions,
types.PeerSession{AccountID: accountID, PeerID: "peer-b", SessionStartedAt: 102},
types.PeerSession{AccountID: accountID, PeerID: "peer-c", SessionStartedAt: 103},
types.PeerSession{AccountID: accountID, PeerID: "peer-d", SessionStartedAt: 999},
types.PeerSession{AccountID: accountID, PeerID: "peer-d", SessionStartedAt: 0},
)
marked, err := store.MarkPeersDisconnectedIfSameSession(ctx, sessions)
require.NoError(t, err)
assert.Equal(t, int64(3), marked)
for id, wantConnected := range map[string]bool{"peer-a": false, "peer-b": false, "peer-c": false, "peer-d": true} {
peer, err := store.GetPeerByID(ctx, LockingStrengthNone, accountID, id)
require.NoError(t, err)
assert.Equal(t, wantConnected, peer.Status.Connected, id)
wantSession := int64(0)
if wantConnected {
wantSession = sessionOf[id]
}
assert.Equal(t, wantSession, peer.Status.SessionStartedAt, id)
}
}
func TestSqlStore_MarkPeersDisconnectedIfSameSessionContinuesAfterFailedChunk(t *testing.T) {
ctx := context.Background()
store, cleanUp, err := NewTestStoreFromSQL(ctx, "../testdata/extended-store.sql", t.TempDir())
t.Cleanup(cleanUp)
require.NoError(t, err)
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
sqlStore := store.(*SqlStore)
sessionOf := map[string]int64{"peer-a": 101, "peer-b": 102, "peer-c": 103}
for id, session := range sessionOf {
require.NoError(t, store.AddPeerToAccount(ctx, &nbpeer.Peer{
ID: id,
AccountID: accountID,
Key: "key-" + id,
DNSLabel: id,
IP: netip.AddrFrom4([4]byte{10, 9, 0, byte(session)}),
IPv6: netip.MustParseAddr(fmt.Sprintf("fd00::9:%d", session)),
}))
require.NoError(t, sqlStore.db.Model(&nbpeer.Peer{}).Where("id = ?", id).Updates(map[string]any{
"peer_status_connected": true,
"peer_status_session_started_at": session,
"peer_status_last_seen": time.Now().UTC(),
}).Error)
}
sessions := []types.PeerSession{{AccountID: accountID, PeerID: "peer-a", SessionStartedAt: 101}}
for i := 0; i < markPeersDisconnectedChunkSize-1; i++ {
sessions = append(sessions, types.PeerSession{AccountID: accountID, PeerID: fmt.Sprintf("missing-%d", i), SessionStartedAt: 1})
}
sessions = append(sessions,
types.PeerSession{AccountID: accountID, PeerID: "peer-b", SessionStartedAt: 102},
types.PeerSession{AccountID: accountID, PeerID: "peer-c", SessionStartedAt: 103},
)
chunkErr := errors.New("deadlock detected")
var updates atomic.Int32
require.NoError(t, sqlStore.db.Callback().Update().Before("gorm:update").Register("test:fail_first_chunk", func(db *gorm.DB) {
if updates.Add(1) == 1 {
_ = db.AddError(chunkErr)
}
}))
marked, err := store.MarkPeersDisconnectedIfSameSession(ctx, sessions)
require.Error(t, err)
assert.ErrorContains(t, err, chunkErr.Error())
assert.Equal(t, int64(2), marked, "only the second chunk is counted")
assert.Equal(t, int32(2), updates.Load(), "the second chunk must still run")
for id, wantConnected := range map[string]bool{"peer-a": true, "peer-b": false, "peer-c": false} {
peer, err := store.GetPeerByID(ctx, LockingStrengthNone, accountID, id)
require.NoError(t, err)
assert.Equal(t, wantConnected, peer.Status.Connected, id)
}
cancelled, cancel := context.WithCancel(ctx)
cancel()
updates.Store(1)
marked, err = store.MarkPeersDisconnectedIfSameSession(cancelled, sessions)
assert.Zero(t, marked)
assert.ErrorContains(t, err, context.Canceled.Error())
assert.Equal(t, int32(1), updates.Load(), "no chunk runs once the context is done")
assert.Equal(t, 1, strings.Count(err.Error(), context.Canceled.Error()), "a done context is reported once")
}
func TestSqlStore_MarkPeerConnectedKeepsLoginExpired(t *testing.T) {
ctx := context.Background()
store, cleanUp, err := NewTestStoreFromSQL(ctx, "../testdata/extended-store.sql", t.TempDir())
t.Cleanup(cleanUp)
require.NoError(t, err)
accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
require.NoError(t, store.AddPeerToAccount(ctx, &nbpeer.Peer{
ID: "expired-peer",
AccountID: accountID,
Key: "key-expired-peer",
DNSLabel: "expired-peer",
IP: netip.AddrFrom4([4]byte{10, 9, 1, 1}),
IPv6: netip.MustParseAddr("fd00::9:101"),
Status: &nbpeer.PeerStatus{LoginExpired: true},
}))
updated, err := store.MarkPeerConnectedIfNewerSession(ctx, accountID, "expired-peer", 1)
require.NoError(t, err)
require.True(t, updated)
peer, err := store.GetPeerByID(ctx, LockingStrengthNone, accountID, "expired-peer")
require.NoError(t, err)
assert.True(t, peer.Status.Connected)
assert.True(t, peer.Status.LoginExpired, "connecting must not clear a login expiry")
}
management/server/store/store.go
+5
View File
@@ -201,6 +201,11 @@ type Store interface {
// sessionStartedAt and the peer's login has not expired. Returns true when
// the update happened.
HeartbeatPeerSession(ctx context.Context, accountID, peerID string, sessionStartedAt int64) (bool, error)
// MarkPeersDisconnectedIfSameSession applies the semantics of
// MarkPeerDisconnectedIfSameSession to many peers in batched statements.
// Sessions with a zero SessionStartedAt are skipped. Returns the number
// of peers that were marked disconnected.
MarkPeersDisconnectedIfSameSession(ctx context.Context, sessions []types.PeerSession) (int64, error)
ApproveAccountPeers(ctx context.Context, accountID string) (int, error)
DeletePeer(ctx context.Context, accountID string, peerID string) error
management/server/store/store_mock.go
+15
View File
@@ -3434,6 +3434,21 @@ func (mr *MockStoreMockRecorder) MarkPeerDisconnectedIfSameSession(ctx, accountI
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeerDisconnectedIfSameSession", reflect.TypeOf((*MockStore)(nil).MarkPeerDisconnectedIfSameSession), ctx, accountID, peerID, sessionStartedAt)
}
// MarkPeersDisconnectedIfSameSession mocks base method.
func (m *MockStore) MarkPeersDisconnectedIfSameSession(ctx context.Context, sessions []types3.PeerSession) (int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "MarkPeersDisconnectedIfSameSession", ctx, sessions)
ret0, _ := ret[0].(int64)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// MarkPeersDisconnectedIfSameSession indicates an expected call of MarkPeersDisconnectedIfSameSession.
func (mr *MockStoreMockRecorder) MarkPeersDisconnectedIfSameSession(ctx, sessions any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeersDisconnectedIfSameSession", reflect.TypeOf((*MockStore)(nil).MarkPeersDisconnectedIfSameSession), ctx, sessions)
}
// MarkPendingJobsAsFailed mocks base method.
func (m *MockStore) MarkPendingJobsAsFailed(ctx context.Context, accountID, peerID, jobID, reason string) error {
m.ctrl.T.Helper()
management/server/types/peer.go
+13
View File
@@ -18,6 +18,9 @@ type PeerSync struct {
// UpdateAccountPeers indicate updating account peers,
// which occurs when the peer's metadata is updated
UpdateAccountPeers bool
// Settings are the account settings the caller already read for this request.
// When nil they are loaded from the store.
Settings *Settings
}
// PeerLogin used as a data object between the gRPC API and Manager on Login request.
@@ -37,4 +40,14 @@ type PeerLogin struct {
// ExtraDNSLabels is a list of extra DNS labels that the peer wants to use
ExtraDNSLabels []string
// Settings are the account settings the caller already read for this request.
// When nil they are loaded from the store.
Settings *Settings
}
// PeerSession identifies the sync session a peer is currently owned by.
type PeerSession struct {
AccountID string
PeerID string
SessionStartedAt int64
}
shared/relay/client/dialer/ws/dialopts_generic.go
+2 -1
View File
@@ -8,8 +8,9 @@ import (
"github.com/coder/websocket"
)
func createDialOptions(serverName string, underlyingOut *net.Conn) *websocket.DialOptions {
func createDialOptions(serverName, host string, underlyingOut *net.Conn) *websocket.DialOptions {
return &websocket.DialOptions{
HTTPClient: httpClientNbDialer(serverName, underlyingOut),
Host: host,
}
}
shared/relay/client/dialer/ws/dialopts_js.go
+1 -1
View File
@@ -8,7 +8,7 @@ import (
"github.com/coder/websocket"
)
func createDialOptions(_ string, _ *net.Conn) *websocket.DialOptions {
func createDialOptions(_, _ string, _ *net.Conn) *websocket.DialOptions {
// WASM version doesn't support HTTPClient or custom TLS config.
return &websocket.DialOptions{}
}
shared/relay/client/dialer/ws/ws.go
+22 -1
View File
@@ -31,8 +31,13 @@ func (d Dialer) Dial(ctx context.Context, address, serverName string) (net.Conn,
return nil, err
}
host, err := requestHost(wsURL, serverName)
if err != nil {
return nil, err
}
var underlying net.Conn
opts := createDialOptions(serverName, &underlying)
opts := createDialOptions(serverName, host, &underlying)
wsConn, resp, err := websocket.Dial(ctx, wsURL, opts)
if err != nil {
@@ -78,6 +83,22 @@ func prepareURL(address string) (string, error) {
return parsed.String(), nil
}
// requestHost keeps the server name in the Host header when dialing by IP, so
// proxies that route by host still match.
func requestHost(wsURL, serverName string) (string, error) {
if serverName == "" {
return "", nil
}
parsed, err := url.Parse(wsURL)
if err != nil {
return "", fmt.Errorf("parse relay URL %q: %w", wsURL, err)
}
if port := parsed.Port(); port != "" {
return net.JoinHostPort(serverName, port), nil
}
return serverName, nil
}
// httpClientNbDialer builds the http client used by the websocket library.
// underlyingOut, when non-nil, is populated with the raw conn from the
// transport's DialContext so the caller can read its RemoteAddr.
shared/relay/client/dialer/ws/ws_test.go
+97
View File
@@ -1,7 +1,12 @@
package ws
import (
"context"
"net"
"net/http"
"net/http/httptest"
"testing"
"time"
)
func TestPrepareURL(t *testing.T) {
@@ -74,3 +79,95 @@ func TestPrepareURL(t *testing.T) {
})
}
}
func TestRequestHost(t *testing.T) {
tests := []struct {
name string
wsURL string
serverName string
want string
}{
{
name: "no server name keeps the URL host",
wsURL: "wss://1.2.3.4:443/relay",
want: "",
},
{
name: "IP with port",
wsURL: "wss://1.2.3.4:443/relay",
serverName: "relay.example.com",
want: "relay.example.com:443",
},
{
name: "IP with non-standard port",
wsURL: "ws://1.2.3.4:45678/relay",
serverName: "relay.example.com",
want: "relay.example.com:45678",
},
{
name: "IP without port",
wsURL: "wss://1.2.3.4/relay",
serverName: "relay.example.com",
want: "relay.example.com",
},
{
name: "IPv6 with port",
wsURL: "wss://[2001:db8::1]:443/relay",
serverName: "relay.example.com",
want: "relay.example.com:443",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := requestHost(tt.wsURL, tt.serverName)
if err != nil {
t.Fatalf("requestHost(%q, %q) err = %v", tt.wsURL, tt.serverName, err)
}
if got != tt.want {
t.Errorf("requestHost(%q, %q) = %q, want %q", tt.wsURL, tt.serverName, got, tt.want)
}
})
}
}
func TestDialByIPSendsServerNameAsHost(t *testing.T) {
hosts := make(chan string, 1)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hosts <- r.Host
http.NotFound(w, r)
}))
defer srv.Close()
_, port, err := net.SplitHostPort(srv.Listener.Addr().String())
if err != nil {
t.Fatal(err)
}
tests := []struct {
name string
serverName string
want string
}{
{name: "dial by IP with server name", serverName: "relay.example.com", want: "relay.example.com:" + port},
{name: "dial without server name", want: "127.0.0.1:" + port},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if _, err := (Dialer{}).Dial(ctx, "rel://127.0.0.1:"+port, tt.serverName); err == nil {
t.Fatal("expected the handshake to fail against a plain HTTP server")
}
select {
case got := <-hosts:
if got != tt.want {
t.Errorf("server saw Host %q, want %q", got, tt.want)
}
case <-ctx.Done():
t.Fatal("server never received the handshake request")
}
})
}
}
util/duration.go
+16
View File
@@ -3,7 +3,10 @@ package util
import (
"encoding/json"
"errors"
"os"
"time"
log "github.com/sirupsen/logrus"
)
// Duration is used strictly for JSON requests/responses due to duration marshalling issues
@@ -35,3 +38,16 @@ func (d *Duration) UnmarshalJSON(b []byte) error {
return errors.New("invalid duration")
}
}
func DurationFromEnv(key string, fallback time.Duration) time.Duration {
val := os.Getenv(key)
if val == "" {
return fallback
}
parsed, err := time.ParseDuration(val)
if err != nil || parsed < 0 {
log.Warnf("ignoring invalid %s value %q, using %v", key, val, fallback)
return fallback
}
return parsed
}
util/duration_test.go
+25
View File
@@ -0,0 +1,25 @@
package util
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
)
func TestDurationFromEnv(t *testing.T) {
const key = "NB_TEST_DURATION"
cases := map[string]time.Duration{
"": time.Minute,
"30s": 30 * time.Second,
"0": 0,
"soon": time.Minute,
"-1s": time.Minute,
}
for value, want := range cases {
t.Run(value, func(t *testing.T) {
t.Setenv(key, value)
assert.Equal(t, want, DurationFromEnv(key, time.Minute), "duration for %q", value)
})
}
}