diff --git a/client/cmd/root.go b/client/cmd/root.go index f976ca895..3c49f8cca 100644 --- a/client/cmd/root.go +++ b/client/cmd/root.go @@ -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") } diff --git a/client/cmd/up_rosenpass_test.go b/client/cmd/up_rosenpass_test.go new file mode 100644 index 000000000..fd36df20f --- /dev/null +++ b/client/cmd/up_rosenpass_test.go @@ -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) + } +} diff --git a/client/ui/frontend/src/modules/settings/SettingsSecurity.tsx b/client/ui/frontend/src/modules/settings/SettingsSecurity.tsx index db1ad9c8f..eacf0ebfe 100644 --- a/client/ui/frontend/src/modules/settings/SettingsSecurity.tsx +++ b/client/ui/frontend/src/modules/settings/SettingsSecurity.tsx @@ -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")} + + {t("settings.security.rosenpass.deprecated")} + + + } /> )} {!hideRosenpassPermissive && ( diff --git a/client/ui/i18n/locales/en/common.json b/client/ui/i18n/locales/en/common.json index 94d741b3e..fb3af16ea 100644 --- a/client/ui/i18n/locales/en/common.json +++ b/client/ui/i18n/locales/en/common.json @@ -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)." diff --git a/go.mod b/go.mod index 699a553af..df3edbb38 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index de7b77649..ee70f2c09 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/management/internals/modules/reverseproxy/accesslogs/manager/repository.go b/management/internals/modules/reverseproxy/accesslogs/manager/repository.go index 3967442d8..1b2cd4976 100644 --- a/management/internals/modules/reverseproxy/accesslogs/manager/repository.go +++ b/management/internals/modules/reverseproxy/accesslogs/manager/repository.go @@ -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 { diff --git a/management/internals/modules/reverseproxy/accesslogs/manager/repository_test.go b/management/internals/modules/reverseproxy/accesslogs/manager/repository_test.go index 7fea12a85..6cbbd3c60 100644 --- a/management/internals/modules/reverseproxy/accesslogs/manager/repository_test.go +++ b/management/internals/modules/reverseproxy/accesslogs/manager/repository_test.go @@ -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 +} diff --git a/management/internals/server/boot.go b/management/internals/server/boot.go index 64d5ee4bc..d37784fae 100644 --- a/management/internals/server/boot.go +++ b/management/internals/server/boot.go @@ -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 }) } diff --git a/management/internals/server/container.go b/management/internals/server/container.go index ee9e65016..760f528a8 100644 --- a/management/internals/server/container.go +++ b/management/internals/server/container.go @@ -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 { diff --git a/management/internals/server/server.go b/management/internals/server/server.go index c95092d70..fa92be978 100644 --- a/management/internals/server/server.go +++ b/management/internals/server/server.go @@ -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) } } diff --git a/management/internals/shared/db/conn.go b/management/internals/shared/db/conn.go index b21b066e2..e12d8e8c8 100644 --- a/management/internals/shared/db/conn.go +++ b/management/internals/shared/db/conn.go @@ -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")) diff --git a/management/internals/shared/db/open.go b/management/internals/shared/db/open.go index f813b0ce8..a64a4003f 100644 --- a/management/internals/shared/db/open.go +++ b/management/internals/shared/db/open.go @@ -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 diff --git a/management/internals/shared/db/timeout.go b/management/internals/shared/db/timeout.go new file mode 100644 index 000000000..d6bc70296 --- /dev/null +++ b/management/internals/shared/db/timeout.go @@ -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) +} diff --git a/management/internals/shared/db/timeout_test.go b/management/internals/shared/db/timeout_test.go new file mode 100644 index 000000000..b9678082d --- /dev/null +++ b/management/internals/shared/db/timeout_test.go @@ -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 +} diff --git a/management/internals/shared/grpc/server.go b/management/internals/shared/grpc/server.go index ff466c5a6..09a4ddddc 100644 --- a/management/internals/shared/grpc/server.go +++ b/management/internals/shared/grpc/server.go @@ -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) diff --git a/management/internals/shared/grpc/server_disconnect_test.go b/management/internals/shared/grpc/server_disconnect_test.go index 7d186b715..c664502e7 100644 --- a/management/internals/shared/grpc/server_disconnect_test.go +++ b/management/internals/shared/grpc/server_disconnect_test.go @@ -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 +} diff --git a/management/server/account.go b/management/server/account.go index 21b9fa431..f62a83d40 100644 --- a/management/server/account.go +++ b/management/server/account.go @@ -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 { diff --git a/management/server/account/manager.go b/management/server/account/manager.go index 15fe11241..2f4df8a48 100644 --- a/management/server/account/manager.go +++ b/management/server/account/manager.go @@ -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) diff --git a/management/server/account/manager_mock.go b/management/server/account/manager_mock.go index b3868f2e2..e7b044bee 100644 --- a/management/server/account/manager_mock.go +++ b/management/server/account/manager_mock.go @@ -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. diff --git a/management/server/account_test.go b/management/server/account_test.go index c343575b8..d6971be85 100644 --- a/management/server/account_test.go +++ b/management/server/account_test.go @@ -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) } diff --git a/management/server/management_proto_test.go b/management/server/management_proto_test.go index 9ca94d2e6..8e495fa7e 100644 --- a/management/server/management_proto_test.go +++ b/management/server/management_proto_test.go @@ -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") +} diff --git a/management/server/mock_server/account_mock.go b/management/server/mock_server/account_mock.go index 50d18ca9b..69a0de251 100644 --- a/management/server/mock_server/account_mock.go +++ b/management/server/mock_server/account_mock.go @@ -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 diff --git a/management/server/peer.go b/management/server/peer.go index c3454179e..3a48c37fb 100644 --- a/management/server/peer.go +++ b/management/server/peer.go @@ -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() } diff --git a/management/server/peer_test.go b/management/server/peer_test.go index bf8093c1d..48b75edbd 100644 --- a/management/server/peer_test.go +++ b/management/server/peer_test.go @@ -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") +} diff --git a/management/server/store/sql_store_peer.go b/management/server/store/sql_store_peer.go index aab9d8fa1..ca0df0ff7 100644 --- a/management/server/store/sql_store_peer.go +++ b/management/server/store/sql_store_peer.go @@ -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{}). diff --git a/management/server/store/sql_store_peer_test.go b/management/server/store/sql_store_peer_test.go index 0b467e37b..b0f605661 100644 --- a/management/server/store/sql_store_peer_test.go +++ b/management/server/store/sql_store_peer_test.go @@ -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") +} diff --git a/management/server/store/store.go b/management/server/store/store.go index c35e84a6e..608120735 100644 --- a/management/server/store/store.go +++ b/management/server/store/store.go @@ -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 diff --git a/management/server/store/store_mock.go b/management/server/store/store_mock.go index bed6708d4..452793c5d 100644 --- a/management/server/store/store_mock.go +++ b/management/server/store/store_mock.go @@ -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() diff --git a/management/server/types/peer.go b/management/server/types/peer.go index 885d67bba..3793ccb83 100644 --- a/management/server/types/peer.go +++ b/management/server/types/peer.go @@ -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 } diff --git a/shared/relay/client/dialer/ws/dialopts_generic.go b/shared/relay/client/dialer/ws/dialopts_generic.go index 8008d89d3..5e1cf4c4f 100644 --- a/shared/relay/client/dialer/ws/dialopts_generic.go +++ b/shared/relay/client/dialer/ws/dialopts_generic.go @@ -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, } } diff --git a/shared/relay/client/dialer/ws/dialopts_js.go b/shared/relay/client/dialer/ws/dialopts_js.go index 5b11fe765..25c9b372c 100644 --- a/shared/relay/client/dialer/ws/dialopts_js.go +++ b/shared/relay/client/dialer/ws/dialopts_js.go @@ -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{} } diff --git a/shared/relay/client/dialer/ws/ws.go b/shared/relay/client/dialer/ws/ws.go index 6b310b73d..946d267fb 100644 --- a/shared/relay/client/dialer/ws/ws.go +++ b/shared/relay/client/dialer/ws/ws.go @@ -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. diff --git a/shared/relay/client/dialer/ws/ws_test.go b/shared/relay/client/dialer/ws/ws_test.go index 7357adbc0..5f04284f6 100644 --- a/shared/relay/client/dialer/ws/ws_test.go +++ b/shared/relay/client/dialer/ws/ws_test.go @@ -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") + } + }) + } +} diff --git a/util/duration.go b/util/duration.go index b657a582d..db4fb3eec 100644 --- a/util/duration.go +++ b/util/duration.go @@ -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 +} diff --git a/util/duration_test.go b/util/duration_test.go new file mode 100644 index 000000000..784bfcf2f --- /dev/null +++ b/util/duration_test.go @@ -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) + }) + } +}