Merge remote-tracking branch 'origin/main' into fix_debug_upload_url_from_mgmt

# Conflicts:
#	client/proto/daemon.pb.go
#	management/internals/controllers/network_map/controller/controller.go
#	management/internals/controllers/network_map/nmaptest/runner.go
#	management/internals/shared/grpc/server.go
#	management/server/account.go
#	shared/management/proto/management.pb.go
#	upload-server/server/server.go

management/internals/shared/db/conn.go

@@ -0,0 +1,121 @@
+package db
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"runtime"
+	"strconv"
+	"time"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+)
+
+const (
+	defaultTransactionTimeout = 5 * time.Minute
+	connMaxLifetime           = time.Hour
+	connMaxIdleTime           = 3 * time.Minute
+)
+
+// TxMetrics receives the duration of every committed top-level transaction.
+type TxMetrics interface {
+	CountTransactionDuration(duration time.Duration)
+}
+
+// Conn is the database connection shared by all repositories: one gorm handle,
+// the pgx pool of a Postgres deployment and the engine they talk to.
+type Conn struct {
+	db        *gorm.DB
+	pool      *pgxpool.Pool
+	engine    Engine
+	txTimeout time.Duration
+	metrics   TxMetrics
+}
+
+// NewConn takes ownership of an open gorm handle and pool once it returns
+// without error, applying the connection limits and transaction timeout
+// configured through the environment.
+func NewConn(ctx context.Context, gormDB *gorm.DB, engine Engine, pool *pgxpool.Pool) (*Conn, error) {
+	sqlDB, err := gormDB.DB()
+	if err != nil {
+		return nil, err
+	}
+
+	txTimeout := defaultTransactionTimeout
+	if v := os.Getenv("NB_STORE_TRANSACTION_TIMEOUT"); v != "" {
+		if parsed, err := time.ParseDuration(v); err == nil {
+			txTimeout = parsed
+		}
+	}
+	log.WithContext(ctx).Infof("Setting transaction timeout to %v", txTimeout)
+
+	conns := runtime.NumCPU()
+	configuredConns, err := strconv.Atoi(os.Getenv("NB_SQL_MAX_OPEN_CONNS"))
+	connsConfigured := err == nil
+	if connsConfigured {
+		conns = configuredConns
+	}
+	if engine == SqliteStoreEngine {
+		if connsConfigured {
+			log.WithContext(ctx).Warnf("setting NB_SQL_MAX_OPEN_CONNS is not supported for sqlite, using default value 1")
+		}
+		conns = 1
+	}
+
+	sqlDB.SetMaxOpenConns(conns)
+	sqlDB.SetMaxIdleConns(conns)
+	sqlDB.SetConnMaxLifetime(connMaxLifetime)
+	sqlDB.SetConnMaxIdleTime(connMaxIdleTime)
+
+	log.WithContext(ctx).Infof("Set max open db connections to %d, max idle to %d, max lifetime to %v, max idle time to %v",
+		conns, conns, connMaxLifetime, connMaxIdleTime)
+
+	return &Conn{db: gormDB, pool: pool, engine: engine, txTimeout: txTimeout}, nil
+}
+
+// DB returns the handle a query must run on: the transaction when tx is set,
+// otherwise the shared connection.
+func (c *Conn) DB(tx *Tx) *gorm.DB {
+	if tx != nil {
+		return tx.db
+	}
+	return c.db
+}
+
+// Pool returns the pgx pool for read paths that bypass gorm. It is nil on
+// engines other than Postgres and inside a transaction, where the pool would
+// not see the uncommitted writes.
+func (c *Conn) Pool(tx *Tx) *pgxpool.Pool {
+	if tx != nil {
+		return nil
+	}
+	return c.pool
+}
+
+func (c *Conn) Engine() Engine {
+	return c.engine
+}
+
+// SetTxMetrics registers the sink that receives transaction durations.
+func (c *Conn) SetTxMetrics(metrics TxMetrics) {
+	c.metrics = metrics
+}
+
+// AutoMigrate creates or updates the tables of the given models.
+func (c *Conn) AutoMigrate(models ...any) error {
+	return c.db.AutoMigrate(models...)
+}
+
+// Close releases the gorm connection and the pgx pool.
+func (c *Conn) Close() error {
+	if c.pool != nil {
+		c.pool.Close()
+	}
+	sqlDB, err := c.db.DB()
+	if err != nil {
+		return fmt.Errorf("get db: %w", err)
+	}
+	return sqlDB.Close()
+}

management/internals/shared/db/conn_test.go

@@ -0,0 +1,148 @@
+package db
+
+import (
+	"context"
+	"errors"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+)
+
+type testRow struct {
+	ID   uint `gorm:"primaryKey"`
+	Name string
+}
+
+func openTestConn(t *testing.T) *Conn {
+	t.Helper()
+	conn, err := OpenSqliteFile(context.Background(), t.TempDir(), SqliteFileName)
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, conn.Close()) })
+	require.NoError(t, conn.AutoMigrate(&testRow{}))
+	return conn
+}
+
+func countRows(t *testing.T, conn *Conn) int64 {
+	t.Helper()
+	var count int64
+	require.NoError(t, conn.DB(nil).Model(&testRow{}).Count(&count).Error)
+	return count
+}
+
+func TestNewConn_ReadsTransactionTimeoutFromEnv(t *testing.T) {
+	t.Setenv("NB_STORE_TRANSACTION_TIMEOUT", "1s")
+	conn := openTestConn(t)
+	assert.Equal(t, time.Second, conn.txTimeout)
+	assert.Equal(t, SqliteStoreEngine, conn.Engine())
+}
+
+func TestRunInTx_CommitsOnSuccess(t *testing.T) {
+	conn := openTestConn(t)
+
+	err := conn.RunInTx(context.Background(), func(tx *Tx) error {
+		return conn.DB(tx).Create(&testRow{Name: "a"}).Error
+	})
+	require.NoError(t, err)
+	assert.EqualValues(t, 1, countRows(t, conn))
+}
+
+func TestRunInTx_RollsBackOnError(t *testing.T) {
+	conn := openTestConn(t)
+	failure := errors.New("boom")
+
+	err := conn.RunInTx(context.Background(), func(tx *Tx) error {
+		require.NoError(t, conn.DB(tx).Create(&testRow{Name: "a"}).Error)
+		return failure
+	})
+	require.ErrorIs(t, err, failure)
+	assert.EqualValues(t, 0, countRows(t, conn))
+}
+
+func TestRunInTx_RollsBackOnPanic(t *testing.T) {
+	conn := openTestConn(t)
+
+	require.Panics(t, func() {
+		_ = conn.RunInTx(context.Background(), func(tx *Tx) error {
+			require.NoError(t, conn.DB(tx).Create(&testRow{Name: "a"}).Error)
+			panic("boom")
+		})
+	})
+	assert.EqualValues(t, 0, countRows(t, conn))
+}
+
+func TestRunInTx_FailsWhenTimeoutExceeded(t *testing.T) {
+	t.Setenv("NB_STORE_TRANSACTION_TIMEOUT", "50ms")
+	conn := openTestConn(t)
+
+	err := conn.RunInTx(context.Background(), func(tx *Tx) error {
+		time.Sleep(100 * time.Millisecond)
+		return conn.DB(tx).Create(&testRow{Name: "a"}).Error
+	})
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	assert.EqualValues(t, 0, countRows(t, conn))
+}
+
+func TestRunInTx_ReportsDurationToMetrics(t *testing.T) {
+	conn := openTestConn(t)
+	metrics := &recordingMetrics{}
+	conn.SetTxMetrics(metrics)
+
+	require.NoError(t, conn.RunInTx(context.Background(), func(*Tx) error { return nil }))
+	assert.Equal(t, 1, metrics.calls)
+}
+
+func TestConn_DBSelectsTransactionHandle(t *testing.T) {
+	conn := openTestConn(t)
+	assert.Same(t, conn.db, conn.DB(nil))
+
+	err := conn.RunInTx(context.Background(), func(tx *Tx) error {
+		assert.Same(t, tx.db, conn.DB(tx))
+		assert.NotSame(t, conn.db, conn.DB(tx))
+		return nil
+	})
+	require.NoError(t, err)
+}
+
+func TestConn_PoolIsUnavailableInsideTransaction(t *testing.T) {
+	conn := openTestConn(t)
+	conn.pool = &pgxpool.Pool{}
+	defer func() { conn.pool = nil }()
+
+	assert.Same(t, conn.pool, conn.Pool(nil))
+	err := conn.RunInTx(context.Background(), func(tx *Tx) error {
+		assert.Nil(t, conn.Pool(tx))
+		return nil
+	})
+	require.NoError(t, err)
+}
+
+type recordingMetrics struct {
+	calls int
+}
+
+func (m *recordingMetrics) CountTransactionDuration(time.Duration) {
+	m.calls++
+}
+
+func TestNewConn_MaxOpenConnsFromEnv(t *testing.T) {
+	t.Setenv("NB_SQL_MAX_OPEN_CONNS", "7")
+
+	gormDB, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "store.db")), GormConfig())
+	require.NoError(t, err)
+	conn, err := NewConn(context.Background(), gormDB, PostgresStoreEngine, nil)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+	sqlDB, err := conn.DB(nil).DB()
+	require.NoError(t, err)
+	assert.Equal(t, 7, sqlDB.Stats().MaxOpenConnections)
+
+	sqliteDB, err := openTestConn(t).DB(nil).DB()
+	require.NoError(t, err)
+	assert.Equal(t, 1, sqliteDB.Stats().MaxOpenConnections)
+}

management/internals/shared/db/dbtest/dbtest.go

@@ -0,0 +1,23 @@
+package dbtest
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/internals/shared/db"
+)
+
+// NewConn opens a fresh SQLite database in a temporary directory, migrates the
+// given models and closes the connection when the test ends. It ignores
+// NB_STORE_ENGINE_SQLITE_FILE, so a developer's configured database is never
+// touched, and is safe to call from parallel tests.
+func NewConn(t testing.TB, models ...any) *db.Conn {
+	t.Helper()
+	conn, err := db.OpenSqliteFile(context.Background(), t.TempDir(), db.SqliteFileName)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+	require.NoError(t, conn.AutoMigrate(models...))
+	return conn
+}

management/internals/shared/db/dbtest/dbtest_test.go

@@ -0,0 +1,31 @@
+package dbtest
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/internals/shared/db"
+)
+
+func TestNewConn_IgnoresSqliteFileOverride(t *testing.T) {
+	override := filepath.Join(t.TempDir(), "configured.db")
+	t.Setenv("NB_STORE_ENGINE_SQLITE_FILE", override)
+
+	conn := NewConn(t)
+
+	assert.Equal(t, db.SqliteStoreEngine, conn.Engine())
+	_, err := os.Stat(override)
+	require.ErrorIs(t, err, os.ErrNotExist)
+}
+
+func TestNewConn_Parallel(t *testing.T) {
+	t.Parallel()
+
+	conn := NewConn(t)
+
+	assert.Equal(t, db.SqliteStoreEngine, conn.Engine())
+}

management/internals/shared/db/engine.go

@@ -0,0 +1,10 @@
+package db
+
+// Engine identifies the SQL engine behind a Conn.
+type Engine string
+
+const (
+	SqliteStoreEngine   Engine = "sqlite"
+	PostgresStoreEngine Engine = "postgres"
+	MysqlStoreEngine    Engine = "mysql"
+)

management/internals/shared/db/lock.go

@@ -0,0 +1,12 @@
+package db
+
+// LockingStrength is the row lock a query holds until its transaction ends.
+type LockingStrength string
+
+const (
+	LockingStrengthUpdate      LockingStrength = "UPDATE"
+	LockingStrengthShare       LockingStrength = "SHARE"
+	LockingStrengthNoKeyUpdate LockingStrength = "NO KEY UPDATE"
+	LockingStrengthKeyShare    LockingStrength = "KEY SHARE"
+	LockingStrengthNone        LockingStrength = "NONE"
+)

management/internals/shared/db/open.go

@@ -0,0 +1,176 @@
+package db
+
+import (
+	"context"
+	"fmt"
+	"net/url"
+	"os"
+	"path/filepath"
+	"runtime"
+	"strings"
+	"time"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+	"gorm.io/driver/mysql"
+	"gorm.io/driver/postgres"
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
+)
+
+// SqliteFileName is the default SQLite database file inside the data directory.
+const SqliteFileName = "store.db"
+
+// PoolConfig sizes the pgx pool a Postgres deployment uses for the read paths
+// that bypass gorm.
+type PoolConfig struct {
+	MaxConns          int32
+	MinConns          int32
+	MaxConnLifetime   time.Duration
+	HealthCheckPeriod time.Duration
+}
+
+var DefaultPoolConfig = PoolConfig{
+	MaxConns:          30,
+	MinConns:          1,
+	MaxConnLifetime:   60 * time.Minute,
+	HealthCheckPeriod: time.Minute,
+}
+
+// GormConfig is the configuration every engine is opened with.
+func GormConfig() *gorm.Config {
+	return &gorm.Config{
+		Logger:          logger.Default.LogMode(logger.Silent),
+		CreateBatchSize: 400,
+	}
+}
+
+// OpenSqlite opens the SQLite database in dataDir, or the file named by
+// NB_STORE_ENGINE_SQLITE_FILE.
+func OpenSqlite(ctx context.Context, dataDir string) (*Conn, error) {
+	storeFile := SqliteFileName
+	if envFile, ok := os.LookupEnv("NB_STORE_ENGINE_SQLITE_FILE"); ok && envFile != "" {
+		storeFile = envFile
+	}
+	return OpenSqliteFile(ctx, dataDir, storeFile)
+}
+
+// OpenSqliteFile opens the SQLite database storeFile, resolved against dataDir
+// when relative. storeFile may carry SQLite URI query parameters.
+func OpenSqliteFile(ctx context.Context, dataDir, storeFile string) (*Conn, error) {
+	// Separate file path from any SQLite URI query parameters (e.g., "store.db?mode=rwc")
+	filePath, query, hasQuery := strings.Cut(storeFile, "?")
+
+	connStr := filePath
+	if !filepath.IsAbs(filePath) {
+		connStr = filepath.Join(dataDir, filePath)
+	}
+
+	// Compose query parameters. User-provided ?_busy_timeout (or its mattn alias
+	// ?_timeout) overrides our default; otherwise inject 30s so SQLite waits at
+	// most that long on a lock instead of blocking the only Go-side connection.
+	// mattn/go-sqlite3 applies PRAGMA from the DSN on every fresh connection, so
+	// the value survives ConnMaxIdleTime/ConnMaxLifetime recycling. cache=shared
+	// stays the default on non-Windows for the same reason as before.
+	parsed, _ := url.ParseQuery(query)
+	var defaults []string
+	if parsed.Get("_busy_timeout") == "" && parsed.Get("_timeout") == "" {
+		defaults = append(defaults, "_busy_timeout=30000")
+	}
+	if !hasQuery && runtime.GOOS != "windows" {
+		// To avoid `The process cannot access the file because it is being used by another process` on Windows
+		defaults = append(defaults, "cache=shared")
+	}
+	parts := defaults
+	if hasQuery {
+		parts = append(parts, query)
+	}
+	if len(parts) > 0 {
+		connStr += "?" + strings.Join(parts, "&")
+	}
+
+	gormDB, err := gorm.Open(sqlite.Open(connStr), GormConfig())
+	if err != nil {
+		return nil, err
+	}
+	conn, err := NewConn(ctx, gormDB, SqliteStoreEngine, nil)
+	if err != nil {
+		closeGorm(gormDB)
+		return nil, err
+	}
+	return conn, nil
+}
+
+// OpenPostgres opens a Postgres database through gorm and a pgx pool sized by pool.
+func OpenPostgres(ctx context.Context, dsn string, pool PoolConfig) (*Conn, error) {
+	gormDB, err := gorm.Open(postgres.Open(dsn), GormConfig())
+	if err != nil {
+		return nil, err
+	}
+	pgxPool, err := newPgxPool(ctx, dsn, pool)
+	if err != nil {
+		closeGorm(gormDB)
+		return nil, err
+	}
+	conn, err := NewConn(ctx, gormDB, PostgresStoreEngine, pgxPool)
+	if err != nil {
+		pgxPool.Close()
+		closeGorm(gormDB)
+		return nil, err
+	}
+	return conn, nil
+}
+
+// MysqlDSN adds the connection parameters every MySQL handle needs, keeping
+// the options already present in dsn.
+func MysqlDSN(dsn string) string {
+	separator := "?"
+	if strings.Contains(dsn, "?") {
+		separator = "&"
+	}
+	return dsn + separator + "charset=utf8&parseTime=True&loc=Local"
+}
+
+// OpenMysql opens a MySQL database through gorm.
+func OpenMysql(ctx context.Context, dsn string) (*Conn, error) {
+	gormDB, err := gorm.Open(mysql.Open(MysqlDSN(dsn)), GormConfig())
+	if err != nil {
+		return nil, err
+	}
+	conn, err := NewConn(ctx, gormDB, MysqlStoreEngine, nil)
+	if err != nil {
+		closeGorm(gormDB)
+		return nil, err
+	}
+	return conn, nil
+}
+
+func newPgxPool(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.MaxConns = cfg.MaxConns
+	config.MinConns = cfg.MinConns
+	config.MaxConnLifetime = cfg.MaxConnLifetime
+	config.HealthCheckPeriod = cfg.HealthCheckPeriod
+
+	pool, err := pgxpool.NewWithConfig(ctx, config)
+	if err != nil {
+		return nil, fmt.Errorf("unable to create connection pool: %w", err)
+	}
+
+	if err := pool.Ping(ctx); err != nil {
+		pool.Close()
+		return nil, fmt.Errorf("unable to ping database: %w", err)
+	}
+
+	return pool, nil
+}
+
+func closeGorm(gormDB *gorm.DB) {
+	if sqlDB, err := gormDB.DB(); err == nil {
+		_ = sqlDB.Close()
+	}
+}

management/internals/shared/db/open_test.go

@@ -0,0 +1,12 @@
+package db
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMysqlDSN(t *testing.T) {
+	assert.Equal(t, "user:pw@tcp(host:3306)/db?charset=utf8&parseTime=True&loc=Local", MysqlDSN("user:pw@tcp(host:3306)/db"))
+	assert.Equal(t, "user:pw@tcp(host:3306)/db?tls=true&charset=utf8&parseTime=True&loc=Local", MysqlDSN("user:pw@tcp(host:3306)/db?tls=true"))
+}

management/internals/shared/db/transaction.go

@@ -0,0 +1,105 @@
+package db
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"runtime/debug"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+)
+
+// Tx is an open transaction handed to repository calls; nil means autocommit.
+type Tx struct {
+	db *gorm.DB
+}
+
+// RunInTx runs fn in one transaction that commits when fn returns nil and rolls
+// back otherwise, bounded by the configured transaction timeout.
+func (c *Conn) RunInTx(ctx context.Context, fn func(tx *Tx) error) error {
+	timeoutCtx, cancel := context.WithTimeout(ctx, c.txTimeout)
+	defer cancel()
+
+	startTime := time.Now()
+	tx := c.db.WithContext(timeoutCtx).Begin()
+	if tx.Error != nil {
+		return tx.Error
+	}
+	defer func() {
+		if r := recover(); r != nil {
+			tx.Rollback()
+			panic(r)
+		}
+	}()
+
+	if err := c.applyStatementTimeouts(tx); err != nil {
+		tx.Rollback()
+		return err
+	}
+
+	err := c.withForeignKeyChecksDisabled(tx, func() error {
+		return fn(&Tx{db: tx})
+	})
+	if err != nil {
+		tx.Rollback()
+		c.logIfTimedOut(ctx, timeoutCtx, err, "transaction", startTime)
+		return err
+	}
+
+	if err := tx.Commit().Error; err != nil {
+		c.logIfTimedOut(ctx, timeoutCtx, err, "transaction commit", startTime)
+		return err
+	}
+
+	log.WithContext(ctx).Tracef("transaction took %v", time.Since(startTime))
+	if c.metrics != nil {
+		c.metrics.CountTransactionDuration(time.Since(startTime))
+	}
+	return nil
+}
+
+func (c *Conn) applyStatementTimeouts(tx *gorm.DB) error {
+	if c.engine != PostgresStoreEngine {
+		return nil
+	}
+	if err := tx.Exec("SET LOCAL statement_timeout = '1min'").Error; err != nil {
+		return fmt.Errorf("failed to set statement timeout: %w", err)
+	}
+	if err := tx.Exec("SET LOCAL lock_timeout = '1min'").Error; err != nil {
+		return fmt.Errorf("failed to set lock timeout: %w", err)
+	}
+	return nil
+}
+
+// withForeignKeyChecksDisabled runs fn with MySQL's FK checks off, which avoids
+// deadlocks on MySQL and Aurora without needing SUPER privilege. The setting is
+// session-scoped and survives a rollback, so it is turned back on whenever fn
+// returns or panics; otherwise the pooled connection would keep it disabled.
+func (c *Conn) withForeignKeyChecksDisabled(tx *gorm.DB, fn func() error) (err error) {
+	if c.engine != MysqlStoreEngine {
+		return fn()
+	}
+	if err := tx.Exec("SET FOREIGN_KEY_CHECKS = 0").Error; err != nil {
+		return fmt.Errorf("failed to disable FK checks: %w", err)
+	}
+	defer func() {
+		restoreErr := tx.Exec("SET FOREIGN_KEY_CHECKS = 1").Error
+		if restoreErr == nil {
+			return
+		}
+		if err == nil {
+			err = fmt.Errorf("failed to re-enable FK checks: %w", restoreErr)
+			return
+		}
+		log.WithContext(tx.Statement.Context).Warnf("failed to re-enable FK checks after failed transaction: %v", restoreErr)
+	}()
+	return fn()
+}
+
+func (c *Conn) logIfTimedOut(ctx, timeoutCtx context.Context, err error, phase string, startTime time.Time) {
+	if errors.Is(err, context.DeadlineExceeded) || errors.Is(timeoutCtx.Err(), context.DeadlineExceeded) {
+		log.WithContext(ctx).Warnf("%s exceeded %s timeout after %v, stack: %s", phase, c.txTimeout, time.Since(startTime), debug.Stack())
+	}
+}

management/internals/shared/grpc/components_encoder.go

@@ -28,10 +28,6 @@ type ComponentsEnvelopeInput struct {
 	// SshAuth.UserIDClaim when reconstructing the NetworkMap. Empty value
 	// is OK — client treats empty as "no SshAuth to build".
 	UserIDClaim string
-	// ProxyPatch carries pre-expanded NetworkMap fragments injected by
-	// external controllers (BYOP/port-forwarding). Nil when no proxy data
-	// is present; encoder skips the field in that case.
-	ProxyPatch *proto.ProxyPatch
 }
 
 // EncodeNetworkMapEnvelope converts NetworkMapComponents into the component
@@ -69,7 +65,6 @@ func EncodeNetworkMapEnvelope(in ComponentsEnvelopeInput) *proto.NetworkMapEnvel
 					DnsForwarderPort: in.DNSForwarderPort,
 					UserIdClaim:      in.UserIDClaim,
 					AccountSettings:  &proto.AccountSettingsCompact{},
-					ProxyPatch:       in.ProxyPatch,
 				},
 			},
 		}
@@ -101,7 +96,6 @@ func EncodeNetworkMapEnvelope(in ComponentsEnvelopeInput) *proto.NetworkMapEnvel
 		AccountSettings:     toAccountSettingsCompact(c.AccountSettings),
 		DnsForwarderPort:    in.DNSForwarderPort,
 		UserIdClaim:         in.UserIDClaim,
-		ProxyPatch:          in.ProxyPatch,
 		DnsSettings:         enc.encodeDNSSettings(c.DNSSettings),
 		DnsDomain:           in.DNSDomain,
 		CustomZoneDomain:    c.CustomZoneDomain,

management/internals/shared/grpc/components_encoder_test.go

@@ -713,66 +713,6 @@ func TestEncodeNetworkMapEnvelope_GroupIDToUserIDs(t *testing.T) {
 	assert.ElementsMatch(t, []string{"user-4"}, full.GroupIdToUserIds["group-users"].UserIds)
 }
 
-func TestToProxyPatch_EmptyInputReturnsNil(t *testing.T) {
-	assert.Nil(t, toProxyPatch(nil, "netbird.cloud", false, false, false))
-	assert.Nil(t, toProxyPatch(&types.NetworkMap{}, "netbird.cloud", false, false, false),
-		"empty NetworkMap (no peers, rules, routes etc) → nil patch so proto3 omits the field")
-}
-
-func TestToProxyPatch_PopulatesAllFields(t *testing.T) {
-	nm := &types.NetworkMap{
-		Peers: []*nmdata.Peer{{
-			ID: "ext-peer", Key: testWgKeyA, IP: netip.AddrFrom4([4]byte{100, 64, 0, 9}),
-			DNSLabel: "extpeer", Meta: nmdata.PeerSystemMeta{WtVersion: "0.40.0"},
-		}},
-		FirewallRules: []*types.FirewallRule{{
-			PeerIP: "100.64.0.9", Action: "accept", Direction: 0, Protocol: "tcp",
-		}},
-	}
-
-	patch := toProxyPatch(nm, "netbird.cloud", false, false, false)
-
-	require.NotNil(t, patch)
-	assert.Len(t, patch.Peers, 1)
-	assert.Len(t, patch.FirewallRules, 1)
-}
-
-// TestEncodeNetworkMapEnvelope_ProxyPatchPropagated covers the ProxyPatch
-// pass-through in both encoder branches (normal path + nil-Components
-// graceful-degrade). Guards against a regression that drops `ProxyPatch:`
-// from one of the envelope struct literals.
-func TestEncodeNetworkMapEnvelope_ProxyPatchPropagated(t *testing.T) {
-	patch := &proto.ProxyPatch{
-		ForwardingRules: []*proto.ForwardingRule{{
-			Protocol:          proto.RuleProtocol_TCP,
-			DestinationPort:   &proto.PortInfo{PortSelection: &proto.PortInfo_Port{Port: 80}},
-			TranslatedAddress: net.IPv4(10, 0, 0, 1).To4(),
-			TranslatedPort:    &proto.PortInfo{PortSelection: &proto.PortInfo_Port{Port: 8080}},
-		}},
-	}
-
-	t.Run("normal_path", func(t *testing.T) {
-		c := newTestComponents()
-		full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{
-			Components: c,
-			ProxyPatch: patch,
-		}).GetFull()
-
-		require.NotNil(t, full.ProxyPatch, "ProxyPatch must propagate through the normal encode path")
-		assert.Len(t, full.ProxyPatch.ForwardingRules, 1)
-	})
-
-	t.Run("empty_components_graceful_degrade", func(t *testing.T) {
-		full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{
-			Components: emptyNetworkMapComponents(),
-			ProxyPatch: patch,
-		}).GetFull()
-
-		require.NotNil(t, full.ProxyPatch, "ProxyPatch must propagate through the nil-Components branch too")
-		assert.Len(t, full.ProxyPatch.ForwardingRules, 1)
-	})
-}
-
 func TestEncodeNetworkMapEnvelope_NilComponentsGracefulDegrade(t *testing.T) {
 	// nil Components → minimal envelope, no crash. Matches the legacy
 	// behaviour for missing/unvalidated peers.

management/internals/shared/grpc/components_envelope_response.go

@@ -9,7 +9,6 @@ import (
 	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
 	"github.com/netbirdio/netbird/management/server/types"
 	sharedgrpc "github.com/netbirdio/netbird/shared/management/grpc"
-	"github.com/netbirdio/netbird/shared/management/networkmap"
 	nmdata "github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 	"github.com/netbirdio/netbird/shared/management/proto"
 )
@@ -34,7 +33,6 @@ func ToComponentSyncResponse(
 	turnCredentials *Token,
 	relayCredentials *Token,
 	components *types.NetworkMapComponents,
-	proxyPatch *types.NetworkMap,
 	dnsName string,
 	checks []*nmdata.PostureChecks,
 	settings *nmdata.AccountSettingsInfo,
@@ -52,9 +50,6 @@ func ToComponentSyncResponse(
 	enableSSH := computeSSHEnabledForPeer(components, peer)
 	peerConfig := toPeerConfig(peer, components.Network, dnsName, settings, httpConfig, deviceFlowConfig, enableSSH, components.ForceRoutingPeerDNSResolution, debugUploadDefault)
 
-	includeIPv6 := peer.SupportsIPv6() && peer.IPv6.IsValid()
-	useSourcePrefixes := peer.SupportsSourcePrefixes()
-
 	userIDClaim := auth.DefaultUserIDClaim
 	if httpConfig != nil && httpConfig.AuthUserIDClaim != "" {
 		userIDClaim = httpConfig.AuthUserIDClaim
@@ -66,7 +61,6 @@ func ToComponentSyncResponse(
 		DNSDomain:        dnsName,
 		DNSForwarderPort: dnsFwdPort,
 		UserIDClaim:      userIDClaim,
-		ProxyPatch:       toProxyPatch(proxyPatch, dnsName, includeIPv6, useSourcePrefixes, peer.ProxyMeta.Embedded),
 	})
 
 	resp := &proto.SyncResponse{
@@ -92,43 +86,6 @@ func ToComponentSyncResponse(
 	return resp
 }
 
-// toProxyPatch converts a proxy-injected *types.NetworkMap into the wire
-// patch the components envelope ships alongside. Returns nil when there are
-// no fragments to merge — proto3 omits a nil message field, so the receiver
-// sees no patch and skips the merge step entirely.
-//
-// We reuse the legacy proto-conversion helpers (toProtocolRoutes,
-// toProtocolFirewallRules, toProtocolRoutesFirewallRules,
-// appendRemotePeerConfig, ForwardingRule.ToProto) because the proxy
-// delivers fragments pre-expanded — there's no raw component shape to
-// derive them from. Components purity isn't violated: proxy data isn't
-// policy-graph-derived, it's externally injected post-Calculate, so the
-// client merges it on top of its locally-computed NetworkMap.
-func toProxyPatch(nm *types.NetworkMap, dnsName string, includeIPv6, useSourcePrefixes, localIsProxy bool) *proto.ProxyPatch {
-	if nm == nil {
-		return nil
-	}
-	if len(nm.Peers) == 0 && len(nm.OfflinePeers) == 0 && len(nm.FirewallRules) == 0 &&
-		len(nm.Routes) == 0 && len(nm.RoutesFirewallRules) == 0 && len(nm.ForwardingRules) == 0 {
-		return nil
-	}
-
-	patch := &proto.ProxyPatch{
-		Peers:              networkmap.AppendRemotePeerConfig(nil, nm.Peers, dnsName, includeIPv6, localIsProxy),
-		OfflinePeers:       networkmap.AppendRemotePeerConfig(nil, nm.OfflinePeers, dnsName, includeIPv6, localIsProxy),
-		FirewallRules:      networkmap.ToProtocolFirewallRules(nm.FirewallRules, includeIPv6, useSourcePrefixes),
-		Routes:             networkmap.ToProtocolRoutes(nm.Routes),
-		RouteFirewallRules: networkmap.ToProtocolRoutesFirewallRules(nm.RoutesFirewallRules),
-	}
-	if len(nm.ForwardingRules) > 0 {
-		patch.ForwardingRules = make([]*proto.ForwardingRule, 0, len(nm.ForwardingRules))
-		for _, r := range nm.ForwardingRules {
-			patch.ForwardingRules = append(patch.ForwardingRules, r.ToProto())
-		}
-	}
-	return patch
-}
-
 // computeSSHEnabledForPeer mirrors the SSH-server-activation bit that
 // Calculate() folds into NetworkMap.EnableSSH. Components-format peers
 // receive a freshly-computed PeerConfig.SshConfig.SshEnabled at sync time;

management/internals/shared/grpc/conversion.go

@@ -217,14 +217,6 @@ func ToSyncResponse(ctx context.Context, config *nbconfig.Config, httpConfig *nb
 	response.NetworkMap.RoutesFirewallRules = routesFirewallRules
 	response.NetworkMap.RoutesFirewallRulesIsEmpty = len(routesFirewallRules) == 0
 
-	if networkMap.ForwardingRules != nil {
-		forwardingRules := make([]*proto.ForwardingRule, 0, len(networkMap.ForwardingRules))
-		for _, rule := range networkMap.ForwardingRules {
-			forwardingRules = append(forwardingRules, rule.ToProto())
-		}
-		response.NetworkMap.ForwardingRules = forwardingRules
-	}
-
 	if networkMap.AuthorizedUsers != nil {
 		hashedUsers, machineUsers := networkmap.BuildAuthorizedUsersProto(ctx, networkMap.AuthorizedUsers)
 		userIDClaim := auth.DefaultUserIDClaim

management/internals/shared/grpc/loginfilter.go

@@ -14,6 +14,7 @@ const (
 	baseBlockDuration = 10 * time.Minute // Duration for which a peer is banned after exceeding the reconnection limit
 	reconnLimitForBan = 30               // Number of reconnections within the reconnTreshold that triggers a ban
 	metaChangeLimit   = 5                // Number of reconnections with different metadata that triggers a ban of one peer
+	maxBanLevel       = 6                // Highest ban level; the ban duration doubles per level up to this one
 )
 
 type lfConfig struct {
@@ -21,6 +22,7 @@ type lfConfig struct {
 	baseBlockDuration time.Duration
 	reconnLimitForBan int
 	metaChangeLimit   int
+	maxBanLevel       int
 }
 
 func initCfg() *lfConfig {
@@ -29,6 +31,7 @@ func initCfg() *lfConfig {
 		baseBlockDuration: baseBlockDuration,
 		reconnLimitForBan: reconnLimitForBan,
 		metaChangeLimit:   metaChangeLimit,
+		maxBanLevel:       maxBanLevel,
 	}
 }
 
@@ -102,11 +105,18 @@ func (l *loginFilter) addLogin(wgPubKey string, metaHash uint64) {
 		return
 	}
 
-	if state.isBanned && now.After(state.banExpiresAt) {
+	if state.isBanned {
+		if now.Before(state.banExpiresAt) {
+			return
+		}
 		state.isBanned = false
 	}
 
-	if state.banLevel > 0 && now.Sub(state.lastSeen) > (2*l.cfg.baseBlockDuration) {
+	quietSince := state.lastSeen
+	if state.banExpiresAt.After(quietSince) {
+		quietSince = state.banExpiresAt
+	}
+	if state.banLevel > 0 && now.Sub(quietSince) > (2*l.cfg.baseBlockDuration) {
 		state.banLevel = 0
 	}
 
@@ -124,10 +134,17 @@ func (l *loginFilter) addLogin(wgPubKey string, metaHash uint64) {
 		return
 	}
 
+	if now.Sub(state.sessionStart) >= l.cfg.reconnThreshold {
+		state.sessionStart = now
+		state.sessionCounter = 0
+	}
+
 	state.sessionCounter++
-	if state.sessionCounter > l.cfg.reconnLimitForBan && now.Sub(state.sessionStart) < l.cfg.reconnThreshold {
+	if state.sessionCounter > l.cfg.reconnLimitForBan {
 		state.isBanned = true
-		state.banLevel++
+		if state.banLevel < l.cfg.maxBanLevel {
+			state.banLevel++
+		}
 
 		backoffFactor := math.Pow(2, float64(state.banLevel-1))
 		duration := time.Duration(float64(l.cfg.baseBlockDuration) * backoffFactor)

management/internals/shared/grpc/loginfilter_test.go

@@ -20,6 +20,7 @@ func testAdvancedCfg() *lfConfig {
 		baseBlockDuration: 100 * time.Millisecond,
 		reconnLimitForBan: 3,
 		metaChangeLimit:   2,
+		maxBanLevel:       3,
 	}
 }
 
@@ -157,6 +158,187 @@ func (s *LoginFilterTestSuite) TestMetaChangeIsAllowedAfterWindowResets() {
 	s.Equal(1, s.filter.logged[pubKey].metaChangeCounter, "meta change counter should reset")
 }
 
+func (s *LoginFilterTestSuite) TestReconnectStormAfterQuietPeriodTriggersBan() {
+	pubKey := "PUB_KEY_A"
+	meta := uint64(1)
+	limit := s.filter.cfg.reconnLimitForBan
+
+	s.filter.addLogin(pubKey, meta)
+	s.Require().Contains(s.filter.logged, pubKey)
+	s.filter.logged[pubKey].sessionStart = time.Now().Add(-(s.filter.cfg.reconnThreshold + time.Second))
+
+	s.filter.addLogin(pubKey, meta)
+	s.Equal(1, s.filter.logged[pubKey].sessionCounter, "expired window should restart the count")
+
+	for i := 1; i < limit; i++ {
+		s.filter.addLogin(pubKey, meta)
+	}
+	s.True(s.filter.allowLogin(pubKey, meta))
+	s.False(s.filter.logged[pubKey].isBanned)
+
+	s.filter.addLogin(pubKey, meta)
+
+	s.False(s.filter.allowLogin(pubKey, meta))
+	s.True(s.filter.logged[pubKey].isBanned)
+}
+
+func (s *LoginFilterTestSuite) TestReconnectStormAfterBanExpiresTriggersBanAgain() {
+	pubKey := "PUB_KEY_A"
+	meta := uint64(1)
+	limit := s.filter.cfg.reconnLimitForBan
+
+	for i := 0; i <= limit; i++ {
+		s.filter.addLogin(pubKey, meta)
+	}
+	s.Require().Contains(s.filter.logged, pubKey)
+	s.Require().True(s.filter.logged[pubKey].isBanned)
+
+	expired := time.Now().Add(-(s.filter.cfg.baseBlockDuration + time.Second))
+	s.filter.logged[pubKey].banExpiresAt = expired
+	s.filter.logged[pubKey].sessionStart = expired
+
+	for i := 0; i <= limit; i++ {
+		s.filter.addLogin(pubKey, meta)
+	}
+
+	s.True(s.filter.logged[pubKey].isBanned)
+	s.Equal(2, s.filter.logged[pubKey].banLevel)
+}
+
+func (s *LoginFilterTestSuite) TestSlowReconnectsAcrossWindowsDoNotBan() {
+	pubKey := "PUB_KEY_A"
+	meta := uint64(1)
+	limit := s.filter.cfg.reconnLimitForBan
+
+	for i := 0; i < limit; i++ {
+		s.filter.addLogin(pubKey, meta)
+	}
+	s.Require().Contains(s.filter.logged, pubKey)
+	s.filter.logged[pubKey].sessionStart = time.Now().Add(-(s.filter.cfg.reconnThreshold + time.Second))
+
+	for i := 0; i < limit; i++ {
+		s.filter.addLogin(pubKey, meta)
+	}
+
+	s.True(s.filter.allowLogin(pubKey, meta))
+	s.False(s.filter.logged[pubKey].isBanned)
+}
+
+func (s *LoginFilterTestSuite) TestBanLevelEscalatesWhenStormResumesRightAfterBan() {
+	pubKey := "PUB_KEY_A"
+	meta := uint64(1)
+	limit := s.filter.cfg.reconnLimitForBan
+	banTime := time.Now().Add(-3 * s.filter.cfg.baseBlockDuration)
+
+	s.filter.logged[pubKey] = &peerState{
+		currentHash:  meta,
+		isBanned:     true,
+		banLevel:     1,
+		banExpiresAt: time.Now().Add(-time.Millisecond),
+		sessionStart: banTime,
+		lastSeen:     banTime,
+	}
+
+	for i := 0; i <= limit; i++ {
+		s.filter.addLogin(pubKey, meta)
+	}
+
+	s.True(s.filter.logged[pubKey].isBanned)
+	s.Equal(2, s.filter.logged[pubKey].banLevel)
+}
+
+func (s *LoginFilterTestSuite) TestBanLevelResetsAfterQuietPeriodFollowingBan() {
+	pubKey := "PUB_KEY_A"
+	meta := uint64(1)
+	quiet := 2*s.filter.cfg.baseBlockDuration + time.Second
+
+	s.filter.logged[pubKey] = &peerState{
+		currentHash:  meta,
+		banLevel:     2,
+		banExpiresAt: time.Now().Add(-s.filter.cfg.baseBlockDuration),
+		lastSeen:     time.Now().Add(-2 * quiet),
+	}
+
+	s.filter.addLogin(pubKey, meta)
+	s.Equal(2, s.filter.logged[pubKey].banLevel, "ban ended more recently than the quiet period")
+
+	s.filter.logged[pubKey].banExpiresAt = time.Now().Add(-quiet)
+	s.filter.logged[pubKey].lastSeen = time.Now().Add(-2 * quiet)
+
+	s.filter.addLogin(pubKey, meta)
+	s.Equal(0, s.filter.logged[pubKey].banLevel)
+}
+
+func (s *LoginFilterTestSuite) TestBanDurationIsCappedAtMaxLevel() {
+	pubKey := "PUB_KEY_A"
+	meta := uint64(1)
+	limit := s.filter.cfg.reconnLimitForBan
+	maxLevel := s.filter.cfg.maxBanLevel
+
+	s.filter.logged[pubKey] = &peerState{
+		currentHash:  meta,
+		banLevel:     maxLevel,
+		sessionStart: time.Now(),
+		lastSeen:     time.Now(),
+	}
+
+	for i := 0; i <= limit; i++ {
+		s.filter.addLogin(pubKey, meta)
+	}
+
+	s.True(s.filter.logged[pubKey].isBanned)
+	s.Equal(maxLevel, s.filter.logged[pubKey].banLevel)
+	expected := s.filter.cfg.baseBlockDuration << (maxLevel - 1)
+	s.InDelta(expected, s.filter.logged[pubKey].banExpiresAt.Sub(s.filter.logged[pubKey].lastSeen), float64(time.Millisecond))
+}
+
+func (s *LoginFilterTestSuite) TestEstablishedPeerReconnectingOnceIsAllowed() {
+	pubKey := "PUB_KEY_A"
+	meta := uint64(1)
+	longAgo := time.Now().Add(-time.Hour)
+
+	s.filter.logged[pubKey] = &peerState{
+		currentHash:           meta,
+		sessionCounter:        1,
+		sessionStart:          longAgo,
+		lastSeen:              longAgo,
+		metaChangeWindowStart: longAgo,
+		metaChangeCounter:     1,
+	}
+
+	s.True(s.filter.allowLogin(pubKey, meta))
+	s.filter.addLogin(pubKey, meta)
+
+	s.True(s.filter.allowLogin(pubKey, meta))
+	s.False(s.filter.logged[pubKey].isBanned)
+	s.Equal(1, s.filter.logged[pubKey].sessionCounter)
+}
+
+func (s *LoginFilterTestSuite) TestLoginsDuringActiveBanDoNotExtendIt() {
+	pubKey := "PUB_KEY_A"
+	meta := uint64(1)
+	limit := s.filter.cfg.reconnLimitForBan
+
+	for i := 0; i <= limit; i++ {
+		s.filter.addLogin(pubKey, meta)
+	}
+	s.Require().Contains(s.filter.logged, pubKey)
+	s.Require().True(s.filter.logged[pubKey].isBanned)
+	expiresAt := time.Now().Add(time.Hour)
+	s.filter.logged[pubKey].banExpiresAt = expiresAt
+	lastSeen := s.filter.logged[pubKey].lastSeen
+
+	for i := 0; i <= limit; i++ {
+		s.filter.addLogin(pubKey, meta)
+	}
+
+	s.True(s.filter.logged[pubKey].isBanned)
+	s.Equal(1, s.filter.logged[pubKey].banLevel)
+	s.Equal(expiresAt, s.filter.logged[pubKey].banExpiresAt)
+	s.Equal(lastSeen, s.filter.logged[pubKey].lastSeen)
+	s.Equal(0, s.filter.logged[pubKey].sessionCounter)
+}
+
 func BenchmarkHashingMethods(b *testing.B) {
 	meta := nbpeer.PeerSystemMeta{
 		WtVersion:          "1.25.1",

management/internals/shared/grpc/pkce_flow_test.go

@@ -0,0 +1,87 @@
+package grpc
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/netbirdio/netbird/shared/management/client/common"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+func TestApplySessionExtendFlowPolicy(t *testing.T) {
+	tests := []struct {
+		name               string
+		flow               *proto.PKCEAuthorizationFlow
+		sessionExtend      bool
+		disablePromptLogin bool
+		loginFlag          uint32
+	}{
+		{
+			name: "extend replaces max_age=0 so login_hint is honoured",
+			flow: &proto.PKCEAuthorizationFlow{
+				ProviderConfig: &proto.ProviderConfig{
+					DisablePromptLogin: false,
+					LoginFlag:          uint32(common.LoginFlagMaxAge0),
+				},
+			},
+			sessionExtend:      true,
+			disablePromptLogin: false,
+			loginFlag:          uint32(common.LoginFlagPromptLogin),
+		},
+		{
+			name: "extend replaces the none flag so the extend is not silent",
+			flow: &proto.PKCEAuthorizationFlow{
+				ProviderConfig: &proto.ProviderConfig{
+					DisablePromptLogin: false,
+					LoginFlag:          uint32(common.LoginFlagNone),
+				},
+			},
+			sessionExtend:      true,
+			disablePromptLogin: false,
+			loginFlag:          uint32(common.LoginFlagPromptLogin),
+		},
+		{
+			name: "extend respects DisablePromptLogin",
+			flow: &proto.PKCEAuthorizationFlow{
+				ProviderConfig: &proto.ProviderConfig{
+					DisablePromptLogin: true,
+					LoginFlag:          uint32(common.LoginFlagMaxAge0),
+				},
+			},
+			sessionExtend:      true,
+			disablePromptLogin: true,
+			loginFlag:          uint32(common.LoginFlagMaxAge0),
+		},
+		{
+			name: "login keeps the configured flow untouched",
+			flow: &proto.PKCEAuthorizationFlow{
+				ProviderConfig: &proto.ProviderConfig{
+					DisablePromptLogin: false,
+					LoginFlag:          uint32(common.LoginFlagMaxAge0),
+				},
+			},
+			sessionExtend:      false,
+			disablePromptLogin: false,
+			loginFlag:          uint32(common.LoginFlagMaxAge0),
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			applySessionExtendFlowPolicy(tc.flow, tc.sessionExtend)
+			cfg := tc.flow.GetProviderConfig()
+			assert.Equal(t, tc.disablePromptLogin, cfg.GetDisablePromptLogin())
+			assert.Equal(t, tc.loginFlag, cfg.GetLoginFlag())
+		})
+	}
+}
+
+// A provider config is not guaranteed to be present on the response; clearing
+// the flag must not panic when the validator returned an empty flow.
+func TestApplySessionExtendFlowPolicyWithoutProviderConfig(t *testing.T) {
+	assert.NotPanics(t, func() {
+		applySessionExtendFlowPolicy(&proto.PKCEAuthorizationFlow{}, true)
+		applySessionExtendFlowPolicy(nil, true)
+	})
+}

management/internals/shared/grpc/pkce_verifier.go

@@ -1,55 +0,0 @@
-package grpc
-
-import (
-	"context"
-	"fmt"
-	"time"
-
-	"github.com/eko/gocache/lib/v4/store"
-	log "github.com/sirupsen/logrus"
-
-	nbcache "github.com/netbirdio/netbird/management/server/cache"
-)
-
-// PKCEVerifierStore manages PKCE verifiers for OAuth flows.
-// Supports both in-memory and Redis storage via NB_IDP_CACHE_REDIS_ADDRESS env var.
-type PKCEVerifierStore struct {
-	cache nbcache.Store
-	ctx   context.Context
-}
-
-// NewPKCEVerifierStore creates a PKCE verifier store using the provided shared cache store.
-func NewPKCEVerifierStore(ctx context.Context, cacheStore nbcache.Store) *PKCEVerifierStore {
-	return &PKCEVerifierStore{
-		cache: cacheStore,
-		ctx:   ctx,
-	}
-}
-
-// Store saves a PKCE verifier associated with an OAuth state parameter.
-// The verifier is stored with the specified TTL and will be automatically deleted after expiration.
-func (s *PKCEVerifierStore) Store(state, verifier string, ttl time.Duration) error {
-	if err := s.cache.Set(s.ctx, state, verifier, store.WithExpiration(ttl)); err != nil {
-		return fmt.Errorf("failed to store PKCE verifier: %w", err)
-	}
-
-	log.Debugf("Stored PKCE verifier for state (expires in %s)", ttl)
-	return nil
-}
-
-// LoadAndDelete retrieves and removes a PKCE verifier for the given state.
-// Returns the verifier and true if found, or empty string and false if not found.
-// This enforces single-use semantics for PKCE verifiers.
-func (s *PKCEVerifierStore) LoadAndDelete(state string) (string, bool) {
-	verifier, found, err := s.cache.GetDel(s.ctx, state)
-	if err != nil {
-		log.Warnf("Failed to consume PKCE verifier: %v", err)
-		return "", false
-	}
-	if !found {
-		log.Debug("PKCE verifier not found for state")
-		return "", false
-	}
-
-	return verifier, true
-}

management/internals/shared/grpc/proxy.go

@@ -27,8 +27,6 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
-	"github.com/netbirdio/netbird/shared/management/domain"
-
 	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
 	"github.com/netbirdio/netbird/management/internals/modules/peers"
 	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs"
@@ -42,6 +40,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/users"
 	proxyauth "github.com/netbirdio/netbird/proxy/auth"
 	"github.com/netbirdio/netbird/shared/hash/argon2id"
+	"github.com/netbirdio/netbird/shared/management/domain"
 	"github.com/netbirdio/netbird/shared/management/proto"
 	nbstatus "github.com/netbirdio/netbird/shared/management/status"
 )
@@ -142,8 +141,8 @@ type ProxyServiceServer struct {
 	// OIDC configuration for proxy authentication
 	oidcConfig ProxyOIDCConfig
 
-	// Store for PKCE verifiers
-	pkceVerifierStore *PKCEVerifierStore
+	// singleUseStore backs both PKCE verifiers and OIDC session exchange codes.
+	singleUseStore *SingleUseStore
 
 	// tokenTTL is the lifetime of one-time tokens generated for proxy
 	// authentication. Defaults to defaultProxyTokenTTL when zero.
@@ -158,6 +157,13 @@ type ProxyServiceServer struct {
 
 const pkceVerifierTTL = 10 * time.Minute
 
+const sessionCodeTTL = 60 * time.Second
+
+const sessionCodeCacheNamespace = "proxy:session"
+
+// The signed nonce binds the handoff mode without changing the state format.
+const sessionCodeNoncePrefix = "code."
+
 const defaultProxyTokenTTL = 5 * time.Minute
 
 const defaultSnapshotBatchSize = 500
@@ -208,13 +214,13 @@ func enforceAccountScope(ctx context.Context, requestAccountID string) error {
 }
 
 // NewProxyServiceServer creates a new proxy service server.
-func NewProxyServiceServer(accessLogMgr accesslogs.Manager, tokenStore *OneTimeTokenStore, pkceStore *PKCEVerifierStore, oidcConfig ProxyOIDCConfig, peersManager peers.Manager, usersManager users.Manager, idpManager idp.Manager, proxyMgr proxy.Manager, tokenChecker ProxyTokenChecker) *ProxyServiceServer {
+func NewProxyServiceServer(accessLogMgr accesslogs.Manager, tokenStore *OneTimeTokenStore, singleUseStore *SingleUseStore, oidcConfig ProxyOIDCConfig, peersManager peers.Manager, usersManager users.Manager, idpManager idp.Manager, proxyMgr proxy.Manager, tokenChecker ProxyTokenChecker) *ProxyServiceServer {
 	ctx, cancel := context.WithCancel(context.Background())
 	s := &ProxyServiceServer{
 		accessLogManager:  accessLogMgr,
 		oidcConfig:        oidcConfig,
 		tokenStore:        tokenStore,
-		pkceVerifierStore: pkceStore,
+		singleUseStore:    singleUseStore,
 		peersManager:      peersManager,
 		usersManager:      usersManager,
 		idpManager:        idpManager,
@@ -306,6 +312,16 @@ func (s *ProxyServiceServer) proxyConnectAuthorizer() ProxyConnectAuthorizer {
 	return s.connectAuthorizer
 }
 
+// GenerateSessionCode creates a single-use code for the given session token.
+func (s *ProxyServiceServer) GenerateSessionCode(sessionToken string) (code string, ok bool) {
+	code, err := s.singleUseStore.Generate(sessionCodeCacheNamespace, sessionToken, sessionCodeTTL)
+	if err != nil {
+		log.WithError(err).Error("failed to generate proxy session code")
+		return "", false
+	}
+	return code, true
+}
+
 // CheckLLMPolicyLimits is the pre-flight policy gate the proxy calls before
 // forwarding an LLM request upstream. Delegates to the agent-network selector,
 // which scores applicable policies by remaining headroom and returns the
@@ -1536,18 +1552,20 @@ func (s *ProxyServiceServer) GetOIDCURL(ctx context.Context, req *proto.GetOIDCU
 		log.WithContext(ctx).Errorf("failed to get account services: %v", err)
 		return nil, status.Errorf(codes.FailedPrecondition, "get account services: %v", err)
 	}
-	var found bool
+	var matchedService *rpservice.Service
 	for _, service := range services {
 		if service.Domain == redirectURL.Hostname() {
-			found = true
+			matchedService = service
 			break
 		}
 	}
-	if !found {
+	if matchedService == nil {
 		log.WithContext(ctx).Debugf("OIDC redirect URL %q does not match any service domain", redirectURL.Hostname())
 		return nil, status.Errorf(codes.FailedPrecondition, "service not found in store")
 	}
 
+	useSessionCode := s.proxyManager.ClusterSupportsSessionCode(ctx, matchedService.ProxyCluster)
+
 	provider, err := oidc.NewProvider(ctx, s.oidcConfig.Issuer)
 	if err != nil {
 		log.WithContext(ctx).Errorf("failed to create OIDC provider: %v", err)
@@ -1567,15 +1585,18 @@ func (s *ProxyServiceServer) GetOIDCURL(ctx context.Context, req *proto.GetOIDCU
 		return nil, status.Errorf(codes.Internal, "generate nonce: %v", err)
 	}
 	nonceB64 := base64.URLEncoding.EncodeToString(nonce)
+	if useSessionCode {
+		nonceB64 = sessionCodeNoncePrefix + nonceB64
+	}
 
 	// Using an HMAC here to avoid redirection state being modified.
-	// State format: base64(redirectURL)|nonce|hmac(redirectURL|nonce)
+	// State format: base64(redirectURL)|[code.]nonce|hmac(redirectURL|nonce)
 	payload := redirectURL.String() + "|" + nonceB64
 	hmacSum := s.generateHMAC(payload)
 	state := fmt.Sprintf("%s|%s|%s", base64.URLEncoding.EncodeToString([]byte(redirectURL.String())), nonceB64, hmacSum)
 
 	codeVerifier := oauth2.GenerateVerifier()
-	if err := s.pkceVerifierStore.Store(state, codeVerifier, pkceVerifierTTL); err != nil {
+	if err := s.singleUseStore.Store(state, codeVerifier, pkceVerifierTTL); err != nil {
 		log.WithContext(ctx).Errorf("failed to store PKCE verifier: %v", err)
 		return nil, status.Errorf(codes.Internal, "store PKCE verifier: %v", err)
 	}
@@ -1612,15 +1633,12 @@ func (s *ProxyServiceServer) generateHMAC(input string) string {
 	return hex.EncodeToString(mac.Sum(nil))
 }
 
-// ValidateState validates the state parameter from an OAuth callback.
-// Returns the original redirect URL if valid, or an error if invalid.
-// The HMAC is verified before consuming the PKCE verifier to prevent
-// an attacker from invalidating a legitimate user's auth flow.
-func (s *ProxyServiceServer) ValidateState(state string) (verifier, redirectURL string, err error) {
-	// State format: base64(redirectURL)|nonce|hmac(redirectURL|nonce)
+// ValidateState validates and consumes an OIDC state.
+func (s *ProxyServiceServer) ValidateState(state string) (verifier, redirectURL string, useSessionCode bool, err error) {
+	// State format: base64(redirectURL)|[code.]nonce|hmac(redirectURL|nonce)
 	parts := strings.Split(state, "|")
 	if len(parts) != 3 {
-		return "", "", errors.New("invalid state format")
+		return "", "", false, errors.New("invalid state format")
 	}
 
 	encodedURL := parts[0]
@@ -1629,7 +1647,7 @@ func (s *ProxyServiceServer) ValidateState(state string) (verifier, redirectURL
 
 	redirectURLBytes, err := base64.URLEncoding.DecodeString(encodedURL)
 	if err != nil {
-		return "", "", fmt.Errorf("invalid state encoding: %w", err)
+		return "", "", false, fmt.Errorf("invalid state encoding: %w", err)
 	}
 	redirectURL = string(redirectURLBytes)
 
@@ -1637,16 +1655,17 @@ func (s *ProxyServiceServer) ValidateState(state string) (verifier, redirectURL
 	expectedHMAC := s.generateHMAC(payload)
 
 	if !hmac.Equal([]byte(providedHMAC), []byte(expectedHMAC)) {
-		return "", "", errors.New("invalid state signature")
+		return "", "", false, errors.New("invalid state signature")
 	}
+	useSessionCode = strings.HasPrefix(nonce, sessionCodeNoncePrefix)
 
 	// Consume the PKCE verifier only after HMAC validation passes.
-	verifier, ok := s.pkceVerifierStore.LoadAndDelete(state)
+	verifier, ok := s.singleUseStore.LoadAndDelete(state)
 	if !ok {
-		return "", "", errors.New("no verifier for state")
+		return "", "", false, errors.New("no verifier for state")
 	}
 
-	return verifier, redirectURL, nil
+	return verifier, redirectURL, useSessionCode, nil
 }
 
 // Denied reasons reported to the proxy when access is refused because of the
@@ -1704,6 +1723,40 @@ func sameAccount(userAccountID, serviceAccountID string) bool {
 	return userAccountID != "" && serviceAccountID != "" && userAccountID == serviceAccountID
 }
 
+// resolveUserEmail returns the user's stored email, else the email claim of
+// an already validated session token, else the IdP's email. The claim was
+// resolved the same way when the token was signed, so reusing it saves the
+// IdP round-trip on every validation. An IdP failure yields an empty string.
+func (s *ProxyServiceServer) resolveUserEmail(ctx context.Context, accountID string, user *types.User, tokenEmail string) string {
+	if user.Email != "" {
+		return user.Email
+	}
+	if tokenEmail != "" {
+		return tokenEmail
+	}
+	return s.idpEmail(ctx, accountID, user.Id)
+}
+
+// idpEmail looks up a user's email at the IdP, scoped to the given account.
+// It returns an empty string when no IdP is configured, the lookup fails or
+// the IdP knows no email for the user.
+func (s *ProxyServiceServer) idpEmail(ctx context.Context, accountID, userID string) string {
+	if s.idpManager == nil {
+		return ""
+	}
+
+	userData, err := s.idpManager.GetUserDataByID(ctx, userID, idp.AppMetadata{WTAccountID: accountID})
+	if err != nil {
+		log.WithContext(ctx).WithFields(log.Fields{"user_id": userID, "error": err.Error()}).Debug("IdP email lookup failed; using stored identity")
+		return ""
+	}
+	if userData == nil {
+		return ""
+	}
+
+	return userData.Email
+}
+
 // GenerateSessionToken creates a signed session JWT for the given domain and
 // user. The user's group memberships are embedded in the token so policy-aware
 // middlewares on the proxy can authorise without an extra management round-trip.
@@ -1759,7 +1812,7 @@ func (s *ProxyServiceServer) GenerateSessionToken(ctx context.Context, domain, u
 	token, err := sessionkey.SignToken(
 		service.SessionPrivateKey,
 		userID,
-		user.Email,
+		s.resolveUserEmail(ctx, service.AccountID, user, ""),
 		domain,
 		method,
 		groupIDs,
@@ -1851,7 +1904,20 @@ func (s *ProxyServiceServer) getAccountServiceByDomain(ctx context.Context, acco
 // ValidateSession validates a session token and checks if the user has access to the domain.
 func (s *ProxyServiceServer) ValidateSession(ctx context.Context, req *proto.ValidateSessionRequest) (*proto.ValidateSessionResponse, error) {
 	domain := req.GetDomain()
-	sessionToken := req.GetSessionToken()
+	sessionToken := req.GetSessionToken() //nolint:staticcheck
+
+	// A one-time code from the OIDC callback is redeemed here for the durable
+	// token, so the token never travels in a redirect URL. The redeemed token
+	// is returned to the proxy (mintedToken) to install as the session cookie.
+	mintedToken := ""
+	if code := req.GetSessionCode(); code != "" {
+		redeemed, found := s.singleUseStore.LoadAndDelete(singleUseCacheKey(sessionCodeCacheNamespace, code))
+		if !found {
+			return deniedSessionResponse("invalid or expired session code"), nil
+		}
+		sessionToken = redeemed
+		mintedToken = redeemed
+	}
 
 	if domain == "" || sessionToken == "" {
 		return deniedSessionResponse("missing domain or session_token"), nil
@@ -1871,7 +1937,7 @@ func (s *ProxyServiceServer) ValidateSession(ctx context.Context, req *proto.Val
 		return nil, err
 	}
 
-	userID, reason := sessionTokenSubject(domain, service, sessionToken)
+	userID, tokenEmail, reason := sessionTokenSubject(domain, service, sessionToken)
 	if reason != "" {
 		return deniedSessionResponse(reason), nil
 	}
@@ -1900,12 +1966,13 @@ func (s *ProxyServiceServer) ValidateSession(ctx context.Context, req *proto.Val
 	}
 
 	groupIDs, groupNames := pairGroupIDsAndNames(userGroups)
+	userEmail := s.resolveUserEmail(ctx, service.AccountID, user, tokenEmail)
 
 	if reason := s.accountUserDeniedReason(domain, service, user); reason != "" {
 		return &proto.ValidateSessionResponse{
 			Valid:          false,
 			UserId:         user.Id,
-			UserEmail:      user.Email,
+			UserEmail:      userEmail,
 			DeniedReason:   reason,
 			PeerGroupIds:   groupIDs,
 			PeerGroupNames: groupNames,
@@ -1915,15 +1982,16 @@ func (s *ProxyServiceServer) ValidateSession(ctx context.Context, req *proto.Val
 	log.WithFields(log.Fields{
 		"domain":  domain,
 		"user_id": userID,
-		"email":   user.Email,
+		"email":   userEmail,
 	}).Debug("ValidateSession: access granted")
 
 	return &proto.ValidateSessionResponse{
 		Valid:          true,
 		UserId:         user.Id,
-		UserEmail:      user.Email,
+		UserEmail:      userEmail,
 		PeerGroupIds:   groupIDs,
 		PeerGroupNames: groupNames,
+		SessionToken:   mintedToken,
 	}, nil
 }
 
@@ -1937,28 +2005,28 @@ func deniedSessionResponse(reason string) *proto.ValidateSessionResponse {
 }
 
 // sessionTokenSubject verifies the session token against the service's session
-// key and returns the user it was minted for, or the reason it cannot be
-// trusted.
-func sessionTokenSubject(domain string, service *rpservice.Service, sessionToken string) (userID, deniedReason string) {
+// key and returns the user and email it was minted for, or the reason it
+// cannot be trusted.
+func sessionTokenSubject(domain string, service *rpservice.Service, sessionToken string) (userID, email, deniedReason string) {
 	pubKeyBytes, err := base64.StdEncoding.DecodeString(service.SessionPublicKey)
 	if err != nil {
 		log.WithFields(log.Fields{
 			"domain": domain,
 			"error":  err.Error(),
 		}).Error("ValidateSession: decode public key")
-		return "", "invalid_service_config"
+		return "", "", "invalid_service_config"
 	}
 
-	userID, _, _, _, _, err = proxyauth.ValidateSessionJWT(sessionToken, domain, pubKeyBytes)
+	userID, email, _, _, _, err = proxyauth.ValidateSessionJWT(sessionToken, domain, pubKeyBytes)
 	if err != nil {
 		log.WithFields(log.Fields{
 			"domain": domain,
 			"error":  err.Error(),
 		}).Debug("ValidateSession: invalid session token")
-		return "", "invalid_token"
+		return "", "", "invalid_token"
 	}
 
-	return userID, ""
+	return userID, email, ""
 }
 
 // accountUserDeniedReason gates a user of the service's own account, returning
@@ -2239,12 +2307,8 @@ func (s *ProxyServiceServer) getTunnelPeerInfo(ctx context.Context, domain strin
 	// IdP enrichment wins when available — the stored email column is a
 	// best-effort cache and is frequently empty for OIDC users. Enrichment
 	// failures must never fail the RPC; we simply keep the stored/peer identity.
-	if s.idpManager != nil {
-		if ud, uerr := s.idpManager.GetUserDataByID(ctx, peer.UserID, idp.AppMetadata{WTAccountID: service.AccountID}); uerr == nil && ud != nil && ud.Email != "" {
-			displayIdentity = ud.Email
-		} else if uerr != nil {
-			log.WithFields(log.Fields{"domain": domain, "user_id": peer.UserID, "error": uerr.Error()}).Debug("ValidateTunnelPeer: IdP user enrichment failed; using stored/peer identity")
-		}
+	if email := s.idpEmail(ctx, service.AccountID, peer.UserID); email != "" {
+		displayIdentity = email
 	}
 
 	return principalID, displayIdentity

management/internals/shared/grpc/proxy_test.go

@@ -129,11 +129,11 @@ func drainEmpty(ch chan *proto.GetMappingUpdateResponse) bool {
 func TestSendServiceUpdateToCluster_UniqueTokensPerProxy(t *testing.T) {
 	ctx := context.Background()
 	tokenStore := NewOneTimeTokenStore(ctx, testCacheStore(t))
-	pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t))
+	singleUseStore := NewSingleUseStore(ctx, testCacheStore(t))
 
 	s := &ProxyServiceServer{
-		tokenStore:        tokenStore,
-		pkceVerifierStore: pkceStore,
+		tokenStore:     tokenStore,
+		singleUseStore: singleUseStore,
 	}
 	s.SetProxyController(newTestProxyController())
 
@@ -186,11 +186,11 @@ func TestSendServiceUpdateToCluster_UniqueTokensPerProxy(t *testing.T) {
 func TestSendServiceUpdateToCluster_DeleteNoToken(t *testing.T) {
 	ctx := context.Background()
 	tokenStore := NewOneTimeTokenStore(ctx, testCacheStore(t))
-	pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t))
+	singleUseStore := NewSingleUseStore(ctx, testCacheStore(t))
 
 	s := &ProxyServiceServer{
-		tokenStore:        tokenStore,
-		pkceVerifierStore: pkceStore,
+		tokenStore:     tokenStore,
+		singleUseStore: singleUseStore,
 	}
 	s.SetProxyController(newTestProxyController())
 
@@ -220,11 +220,11 @@ func TestSendServiceUpdateToCluster_DeleteNoToken(t *testing.T) {
 func TestSendServiceUpdate_UniqueTokensPerProxy(t *testing.T) {
 	ctx := context.Background()
 	tokenStore := NewOneTimeTokenStore(ctx, testCacheStore(t))
-	pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t))
+	singleUseStore := NewSingleUseStore(ctx, testCacheStore(t))
 
 	s := &ProxyServiceServer{
-		tokenStore:        tokenStore,
-		pkceVerifierStore: pkceStore,
+		tokenStore:     tokenStore,
+		singleUseStore: singleUseStore,
 	}
 	s.SetProxyController(newTestProxyController())
 
@@ -272,13 +272,13 @@ func generateState(s *ProxyServiceServer, redirectURL string) string {
 
 func TestOAuthState_NeverTheSame(t *testing.T) {
 	ctx := context.Background()
-	pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t))
+	singleUseStore := NewSingleUseStore(ctx, testCacheStore(t))
 
 	s := &ProxyServiceServer{
 		oidcConfig: ProxyOIDCConfig{
 			HMACKey: []byte("test-hmac-key"),
 		},
-		pkceVerifierStore: pkceStore,
+		singleUseStore: singleUseStore,
 	}
 
 	redirectURL := "https://app.example.com/callback"
@@ -300,20 +300,20 @@ func TestOAuthState_NeverTheSame(t *testing.T) {
 
 func TestValidateState_RejectsOldTwoPartFormat(t *testing.T) {
 	ctx := context.Background()
-	pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t))
+	singleUseStore := NewSingleUseStore(ctx, testCacheStore(t))
 
 	s := &ProxyServiceServer{
 		oidcConfig: ProxyOIDCConfig{
 			HMACKey: []byte("test-hmac-key"),
 		},
-		pkceVerifierStore: pkceStore,
+		singleUseStore: singleUseStore,
 	}
 
 	// Old format had only 2 parts: base64(url)|hmac
-	err := s.pkceVerifierStore.Store("base64url|hmac", "test", 10*time.Minute)
+	err := s.singleUseStore.Store("base64url|hmac", "test", 10*time.Minute)
 	require.NoError(t, err)
 
-	_, _, err = s.ValidateState("base64url|hmac")
+	_, _, _, err = s.ValidateState("base64url|hmac")
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "invalid state format")
 }
@@ -372,24 +372,48 @@ func TestEnforceAccountScope_AllowsNoTokenInContext(t *testing.T) {
 
 func TestValidateState_RejectsInvalidHMAC(t *testing.T) {
 	ctx := context.Background()
-	pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t))
+	singleUseStore := NewSingleUseStore(ctx, testCacheStore(t))
 
 	s := &ProxyServiceServer{
 		oidcConfig: ProxyOIDCConfig{
 			HMACKey: []byte("test-hmac-key"),
 		},
-		pkceVerifierStore: pkceStore,
+		singleUseStore: singleUseStore,
 	}
 
 	// Store with tampered HMAC
-	err := s.pkceVerifierStore.Store("dGVzdA==|nonce|wrong-hmac", "test", 10*time.Minute)
+	err := s.singleUseStore.Store("dGVzdA==|nonce|wrong-hmac", "test", 10*time.Minute)
 	require.NoError(t, err)
 
-	_, _, err = s.ValidateState("dGVzdA==|nonce|wrong-hmac")
+	_, _, _, err = s.ValidateState("dGVzdA==|nonce|wrong-hmac")
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "invalid state signature")
 }
 
+func TestSessionCodeCannotConsumeOIDCState(t *testing.T) {
+	const verifier = "pkce-verifier"
+
+	store := NewSingleUseStore(context.Background(), testCacheStore(t))
+	server := &ProxyServiceServer{
+		oidcConfig: ProxyOIDCConfig{
+			HMACKey: []byte("test-hmac-key"),
+		},
+		singleUseStore: store,
+	}
+	state := generateState(server, "https://service.example.com/callback")
+	require.NoError(t, store.Store(state, verifier, time.Minute))
+
+	response, err := server.ValidateSession(context.Background(), &proto.ValidateSessionRequest{
+		SessionCode: state,
+	})
+	require.NoError(t, err)
+	assert.False(t, response.GetValid())
+
+	gotVerifier, _, _, err := server.ValidateState(state)
+	require.NoError(t, err)
+	assert.Equal(t, verifier, gotVerifier)
+}
+
 func TestSendServiceUpdateToCluster_FiltersOnCapability(t *testing.T) {
 	tokenStore := NewOneTimeTokenStore(context.Background(), testCacheStore(t))
 

management/internals/shared/grpc/server.go

@@ -313,7 +313,7 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
 	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)
+		s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil)
 		return err
 	}
 
@@ -321,7 +321,7 @@ 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)
+		s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil)
 		return err
 	}
 
@@ -337,7 +337,7 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
 
 	s.syncSem.Add(-1)
 
-	return PeerUpdateHandlerFactory(peerKey, updates, s.secretsManager, srv, func() { s.cancelPeerRoutines(ctx, accountID, peer, syncStart) }).
+	return PeerUpdateHandlerFactory(peerKey, updates, s.secretsManager, srv, func() { s.cancelPeerRoutines(ctx, accountID, peer, syncStart, updates) }).
 		WithMetrics(s.appMetrics).HandleUpdates(ctx)
 }
 
@@ -383,7 +383,7 @@ func (s *Server) startResponseReceiver(ctx context.Context, srv proto.Management
 
 func (s *Server) sendJobsLoop(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, updates *job.Channel, srv proto.ManagementService_JobServer) error {
 	// todo figure out better error handling strategy
-	defer s.jobManager.CloseChannel(ctx, accountID, peer.ID)
+	defer s.jobManager.CloseChannel(ctx, accountID, peer.ID, updates)
 
 	for {
 		event, err := updates.Event(ctx)
@@ -430,20 +430,26 @@ func (s *Server) sendJob(ctx context.Context, peerKey wgtypes.Key, job *job.Even
 	return nil
 }
 
-func (s *Server) cancelPeerRoutines(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time) {
+func (s *Server) cancelPeerRoutines(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time, session chan *network_map.UpdateMessage) {
 	uncanceledCTX := context.WithoutCancel(ctx)
 	unlock := s.acquirePeerLockByUID(uncanceledCTX, peer.Key)
 	defer unlock()
 
-	s.cancelPeerRoutinesWithoutLock(uncanceledCTX, accountID, peer, streamStartTime)
+	s.cancelPeerRoutinesWithoutLock(uncanceledCTX, accountID, peer, streamStartTime, session)
 }
 
-func (s *Server) cancelPeerRoutinesWithoutLock(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time) {
+// 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)
 	}
-	s.networkMapController.OnPeerDisconnected(ctx, accountID, peer.ID)
+	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)
+		return
+	}
 	s.secretsManager.CancelRefresh(peer.ID)
 
 	log.WithContext(ctx).Debugf("peer %s has been disconnected", peer.Key)
@@ -955,12 +961,12 @@ func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer
 		// stops doing duplicate work. Deferred until the client-side
 		// decoder lands and there's a real deployment of capability=3 peers
 		// worth optimizing for.
-		freshPeer, components, proxyPatch, freshPostureChecks, freshDnsFwdPort, err := s.networkMapController.GetValidatedPeerWithComponents(ctx, false, peer.AccountID, peer)
+		freshPeer, components, freshPostureChecks, freshDnsFwdPort, err := s.networkMapController.GetValidatedPeerWithComponents(ctx, false, peer.AccountID, peer)
 		if err != nil {
 			log.WithContext(ctx).Errorf("failed to build components for peer %s on initial sync: %v", peer.ID, err)
 			return status.Errorf(codes.Internal, "failed to build initial sync envelope")
 		}
-		plainResp = ToComponentSyncResponse(ctx, s.config, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, types.TwinPeer(freshPeer), turnToken, relayToken, components, proxyPatch, dnsName, freshPostureChecks, types.TwinAccountSettings(settings), settings.Extra, peerGroups, freshDnsFwdPort, s.config.DebugUpload.URL)
+		plainResp = ToComponentSyncResponse(ctx, s.config, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, types.TwinPeer(freshPeer), turnToken, relayToken, components, dnsName, freshPostureChecks, types.TwinAccountSettings(settings), settings.Extra, peerGroups, freshDnsFwdPort, s.config.DebugUpload.URL)
 	} else {
 		plainResp = ToSyncResponse(ctx, s.config, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, types.TwinPeer(peer), turnToken, relayToken, networkMap, dnsName, postureChecks, nil, types.TwinAccountSettings(settings), settings.Extra, peerGroups, dnsFwdPort, s.config.DebugUpload.URL)
 	}
@@ -1088,7 +1094,8 @@ func (s *Server) GetPKCEAuthorizationFlow(ctx context.Context, req *proto.Encryp
 		return nil, status.Errorf(codes.Internal, "failed to get server key")
 	}
 
-	err = encryption.DecryptMessage(peerKey, key, req.Body, &proto.PKCEAuthorizationFlowRequest{})
+	flowReq := &proto.PKCEAuthorizationFlowRequest{}
+	err = encryption.DecryptMessage(peerKey, key, req.Body, flowReq)
 	if err != nil {
 		errMSG := fmt.Sprintf("error while decrypting peer's message with Wireguard public key %s.", req.WgPubKey)
 		log.WithContext(ctx).Warn(errMSG)
@@ -1132,6 +1139,7 @@ func (s *Server) GetPKCEAuthorizationFlow(ctx context.Context, req *proto.Encryp
 	}
 
 	flowInfoResp := s.integratedPeerValidator.ValidateFlowResponse(ctx, peerKey.String(), initInfoFlow)
+	applySessionExtendFlowPolicy(flowInfoResp, flowReq.GetSessionExtend())
 
 	encryptedResp, err := encryption.EncryptMessage(peerKey, key, flowInfoResp)
 	if err != nil {
@@ -1144,6 +1152,40 @@ func (s *Server) GetPKCEAuthorizationFlow(ctx context.Context, req *proto.Encryp
 	}, nil
 }
 
+// applySessionExtendFlowPolicy forces a prompt=login flow for a session extend.
+//
+// An extend renews the session of one specific peer, so its token has to come
+// from the account that peer is registered under. A flow that does not prompt
+// leaves the choice to the IdP, which answers a silent authorization from any
+// session it already holds — not necessarily this peer's account when several
+// are signed in, and login_hint is a suggestion the IdP may ignore. The token
+// then fails the jwt.UserID == peer.UserID check in ExtendAuthSession, and the
+// user is given no opportunity to pick a different account.
+//
+// LoginFlagPromptLogin rather than max_age=0: both re-authenticate, but with
+// prompt=login the IdP honours login_hint and offers the peer's own account,
+// whereas max_age=0 leaves the user to find it among every account signed in.
+//
+// DisablePromptLogin is left alone. It is set for IdPs that break on
+// prompt=login — Authentik triggers a double authentication, and social logins
+// fail outright — so overriding it would trade a recoverable session extend for
+// a login that cannot complete at all. Those deployments keep the silent flow
+// and, with several accounts signed in, an extend answered from the wrong one
+// still fails the user match.
+//
+// Called after ValidateFlowResponse so that a per-peer override cannot reinstate
+// the silent flow for an extend.
+func applySessionExtendFlowPolicy(flow *proto.PKCEAuthorizationFlow, sessionExtend bool) {
+	if !sessionExtend {
+		return
+	}
+	cfg := flow.GetProviderConfig()
+	if cfg == nil || cfg.GetDisablePromptLogin() {
+		return
+	}
+	cfg.LoginFlag = uint32(common.LoginFlagPromptLogin)
+}
+
 // SyncMeta endpoint is used to synchronize peer's system metadata and notifies the connected,
 // peer's under the same account of any updates.
 func (s *Server) SyncMeta(ctx context.Context, req *proto.EncryptedMessage) (*proto.Empty, error) {

management/internals/shared/grpc/server_disconnect_test.go

@@ -0,0 +1,48 @@
+package grpc
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"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"
+)
+
+func TestCancelPeerRoutines_SessionOwnership(t *testing.T) {
+	peer := &nbpeer.Peer{ID: "peer-1", Key: "peer-key"}
+	streamStart := time.Unix(1700000000, 0)
+	session := make(chan *network_map.UpdateMessage)
+
+	tests := []struct {
+		name          string
+		session       chan *network_map.UpdateMessage
+		ownsPeer      bool
+		cancelRefresh bool
+	}{
+		{name: "owning session tears everything down", session: session, ownsPeer: true, cancelRefresh: true},
+		{name: "stale session keeps the newer session's refresh", session: session, ownsPeer: false, cancelRefresh: false},
+		{name: "failed sync without a channel closes the older session", session: nil, ownsPeer: true, cancelRefresh: 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}
+
+			accountManager.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.Key, streamStart).Return(nil)
+			controller.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.ID, tt.session).Return(tt.ownsPeer)
+			if tt.cancelRefresh {
+				secretsManager.EXPECT().CancelRefresh(peer.ID)
+			}
+
+			s.cancelPeerRoutines(context.Background(), "account-1", peer, streamStart, tt.session)
+		})
+	}
+}

management/internals/shared/grpc/single_use.go

@@ -0,0 +1,67 @@
+package grpc
+
+import (
+	"context"
+	"crypto/rand"
+	"encoding/base64"
+	"fmt"
+	"time"
+
+	"github.com/eko/gocache/lib/v4/store"
+	log "github.com/sirupsen/logrus"
+
+	nbcache "github.com/netbirdio/netbird/management/server/cache"
+)
+
+// SingleUseStore stores short-lived values that can be retrieved only once.
+type SingleUseStore struct {
+	cache nbcache.Store
+	ctx   context.Context
+}
+
+// NewSingleUseStore creates a single-use value store over the shared cache.
+func NewSingleUseStore(ctx context.Context, cacheStore nbcache.Store) *SingleUseStore {
+	return &SingleUseStore{
+		cache: cacheStore,
+		ctx:   ctx,
+	}
+}
+
+// Store saves value under key with the given TTL, after which it is evicted.
+func (s *SingleUseStore) Store(key, value string, ttl time.Duration) error {
+	if err := s.cache.Set(s.ctx, key, value, store.WithExpiration(ttl)); err != nil {
+		return fmt.Errorf("store single-use value: %w", err)
+	}
+	return nil
+}
+
+// Generate stores a value under a namespaced random key and returns the random key.
+func (s *SingleUseStore) Generate(namespace, value string, ttl time.Duration) (string, error) {
+	buf := make([]byte, 32)
+	if _, err := rand.Read(buf); err != nil {
+		return "", fmt.Errorf("generate single-use key: %w", err)
+	}
+
+	key := base64.RawURLEncoding.EncodeToString(buf)
+	if err := s.Store(singleUseCacheKey(namespace, key), value, ttl); err != nil {
+		return "", err
+	}
+	return key, nil
+}
+
+func singleUseCacheKey(namespace, key string) string {
+	return namespace + ":" + key
+}
+
+// LoadAndDelete retrieves and removes the value for a key.
+func (s *SingleUseStore) LoadAndDelete(key string) (string, bool) {
+	value, found, err := s.cache.GetDel(s.ctx, key)
+	if err != nil {
+		log.Warnf("failed to consume single-use value: %v", err)
+		return "", false
+	}
+	if !found {
+		return "", false
+	}
+	return value, true
+}

management/internals/shared/grpc/single_use_test.go

@@ -6,7 +6,7 @@ import (
 	"time"
 )
 
-func TestPKCEVerifierStoreLoadAndDelete(t *testing.T) {
+func TestSingleUseStoreLoadAndDelete(t *testing.T) {
 	const (
 		state    = "state"
 		verifier = "verifier"
@@ -14,7 +14,7 @@ func TestPKCEVerifierStoreLoadAndDelete(t *testing.T) {
 	)
 
 	t.Run("exactly one concurrent caller consumes the verifier", func(t *testing.T) {
-		store := NewPKCEVerifierStore(context.Background(), testCacheStore(t))
+		store := NewSingleUseStore(context.Background(), testCacheStore(t))
 		if err := store.Store(state, verifier, time.Minute); err != nil {
 			t.Fatalf("couldn't store PKCE verifier: %s", err)
 		}
@@ -50,7 +50,7 @@ func TestPKCEVerifierStoreLoadAndDelete(t *testing.T) {
 	})
 
 	t.Run("replayed state is rejected", func(t *testing.T) {
-		store := NewPKCEVerifierStore(context.Background(), testCacheStore(t))
+		store := NewSingleUseStore(context.Background(), testCacheStore(t))
 		if err := store.Store(state, verifier, time.Minute); err != nil {
 			t.Fatalf("couldn't store PKCE verifier: %s", err)
 		}
@@ -64,7 +64,7 @@ func TestPKCEVerifierStoreLoadAndDelete(t *testing.T) {
 	})
 
 	t.Run("unknown state is rejected", func(t *testing.T) {
-		store := NewPKCEVerifierStore(context.Background(), testCacheStore(t))
+		store := NewSingleUseStore(context.Background(), testCacheStore(t))
 
 		if got, found := store.LoadAndDelete("never-stored"); found {
 			t.Fatalf("unknown state should not resolve, got %q", got)
@@ -72,7 +72,7 @@ func TestPKCEVerifierStoreLoadAndDelete(t *testing.T) {
 	})
 
 	t.Run("expired verifier is rejected", func(t *testing.T) {
-		store := NewPKCEVerifierStore(context.Background(), testCacheStore(t))
+		store := NewSingleUseStore(context.Background(), testCacheStore(t))
 		if err := store.Store(state, verifier, 50*time.Millisecond); err != nil {
 			t.Fatalf("couldn't store PKCE verifier: %s", err)
 		}
@@ -83,3 +83,40 @@ func TestPKCEVerifierStoreLoadAndDelete(t *testing.T) {
 		}
 	})
 }
+
+func TestSingleUseStore_GenerateAndConsumeOnce(t *testing.T) {
+	const namespace = "test"
+	s := NewSingleUseStore(context.Background(), testCacheStore(t))
+
+	key, err := s.Generate(namespace, "the-value", time.Minute)
+	if err != nil {
+		t.Fatalf("generate: %v", err)
+	}
+	if key == "" || key == "the-value" {
+		t.Fatalf("unexpected key %q", key)
+	}
+
+	value, found := s.LoadAndDelete(singleUseCacheKey(namespace, key))
+	if !found || value != "the-value" {
+		t.Fatalf("expected to load the stored value, got %q found=%v", value, found)
+	}
+
+	if _, found := s.LoadAndDelete(singleUseCacheKey(namespace, key)); found {
+		t.Fatal("value must be consumed on first LoadAndDelete")
+	}
+}
+
+func TestSingleUseStore_GenerateUniqueKeys(t *testing.T) {
+	s := NewSingleUseStore(context.Background(), testCacheStore(t))
+	a, err := s.Generate("test", "v", time.Minute)
+	if err != nil {
+		t.Fatalf("generate a: %v", err)
+	}
+	b, err := s.Generate("test", "v", time.Minute)
+	if err != nil {
+		t.Fatalf("generate b: %v", err)
+	}
+	if a == b {
+		t.Fatal("generated keys must be distinct")
+	}
+}

management/internals/shared/grpc/validate_session_test.go

@@ -1,5 +1,3 @@
-//go:build integration
-
 package grpc
 
 import (
@@ -7,6 +5,7 @@ import (
 	"crypto/ed25519"
 	"crypto/rand"
 	"encoding/base64"
+	"errors"
 	"testing"
 	"time"
 
@@ -40,9 +39,9 @@ func setupValidateSessionTest(t *testing.T) *validateSessionTestSetup {
 	proxyManager := &testValidateSessionProxyManager{}
 
 	tokenStore := NewOneTimeTokenStore(ctx, testCacheStore(t))
-	pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t))
+	singleUseStore := NewSingleUseStore(ctx, testCacheStore(t))
 
-	proxyService := NewProxyServiceServer(nil, tokenStore, pkceStore, ProxyOIDCConfig{}, nil, usersManager, nil, proxyManager, nil)
+	proxyService := NewProxyServiceServer(nil, tokenStore, singleUseStore, ProxyOIDCConfig{}, nil, usersManager, nil, proxyManager, nil)
 	proxyService.SetServiceManager(serviceManager)
 
 	createTestProxies(t, ctx, testStore)
@@ -179,7 +178,12 @@ func generateSessionKeyPair(t *testing.T) (string, string) {
 
 func createSessionToken(t *testing.T, privKeyB64, userID, domain string) string {
 	t.Helper()
-	token, err := sessionkey.SignToken(privKeyB64, userID, "", domain, auth.MethodOIDC, nil, nil, time.Hour)
+	return createSessionTokenWithEmail(t, privKeyB64, userID, "", domain)
+}
+
+func createSessionTokenWithEmail(t *testing.T, privKeyB64, userID, email, domain string) string {
+	t.Helper()
+	token, err := sessionkey.SignToken(privKeyB64, userID, email, domain, auth.MethodOIDC, nil, nil, time.Hour)
 	require.NoError(t, err)
 	return token
 }
@@ -195,7 +199,7 @@ func TestValidateSession_UserAllowed(t *testing.T) {
 
 	resp, err := setup.proxyService.ValidateSession(context.Background(), &proto.ValidateSessionRequest{
 		Domain:       "test-proxy.example.com",
-		SessionToken: token,
+		SessionToken: token, //nolint:staticcheck
 	})
 
 	require.NoError(t, err)
@@ -216,7 +220,7 @@ func TestValidateSession_UserNotInAllowedGroup(t *testing.T) {
 
 	resp, err := setup.proxyService.ValidateSession(context.Background(), &proto.ValidateSessionRequest{
 		Domain:       "restricted-proxy.example.com",
-		SessionToken: token,
+		SessionToken: token, //nolint:staticcheck
 	})
 
 	require.NoError(t, err)
@@ -240,7 +244,7 @@ func TestValidateSession_PendingApprovalUserDenied(t *testing.T) {
 
 	resp, err := setup.proxyService.ValidateSession(context.Background(), &proto.ValidateSessionRequest{
 		Domain:       "restricted-proxy.example.com",
-		SessionToken: token,
+		SessionToken: token, //nolint:staticcheck
 	})
 
 	require.NoError(t, err)
@@ -265,7 +269,7 @@ func TestValidateSession_PendingApprovalUserInAllUsersGroupDenied(t *testing.T)
 
 	resp, err := setup.proxyService.ValidateSession(context.Background(), &proto.ValidateSessionRequest{
 		Domain:       "all-users-proxy.example.com",
-		SessionToken: token,
+		SessionToken: token, //nolint:staticcheck
 	})
 
 	require.NoError(t, err)
@@ -288,7 +292,7 @@ func TestValidateSession_BlockedUserDenied(t *testing.T) {
 
 	resp, err := setup.proxyService.ValidateSession(context.Background(), &proto.ValidateSessionRequest{
 		Domain:       "restricted-proxy.example.com",
-		SessionToken: token,
+		SessionToken: token, //nolint:staticcheck
 	})
 
 	require.NoError(t, err)
@@ -312,7 +316,7 @@ func TestValidateSession_UserAllowedAfterApproval(t *testing.T) {
 	token := createSessionToken(t, proxy.SessionPrivateKey, pendingUserID, "restricted-proxy.example.com")
 	req := &proto.ValidateSessionRequest{
 		Domain:       "restricted-proxy.example.com",
-		SessionToken: token,
+		SessionToken: token, //nolint:staticcheck,
 	}
 
 	resp, err := setup.proxyService.ValidateSession(ctx, req)
@@ -345,7 +349,7 @@ func TestValidateSession_UserInDifferentAccount(t *testing.T) {
 
 	resp, err := setup.proxyService.ValidateSession(context.Background(), &proto.ValidateSessionRequest{
 		Domain:       "test-proxy.example.com",
-		SessionToken: token,
+		SessionToken: token, //nolint:staticcheck,
 	})
 
 	require.NoError(t, err)
@@ -364,7 +368,7 @@ func TestValidateSession_UserNotFound(t *testing.T) {
 
 	resp, err := setup.proxyService.ValidateSession(context.Background(), &proto.ValidateSessionRequest{
 		Domain:       "test-proxy.example.com",
-		SessionToken: token,
+		SessionToken: token, //nolint:staticcheck,
 	})
 
 	require.NoError(t, err)
@@ -383,7 +387,7 @@ func TestValidateSession_ProxyNotFound(t *testing.T) {
 
 	resp, err := setup.proxyService.ValidateSession(context.Background(), &proto.ValidateSessionRequest{
 		Domain:       "unknown-proxy.example.com",
-		SessionToken: token,
+		SessionToken: token, //nolint:staticcheck,
 	})
 
 	require.NoError(t, err)
@@ -397,7 +401,7 @@ func TestValidateSession_InvalidToken(t *testing.T) {
 
 	resp, err := setup.proxyService.ValidateSession(context.Background(), &proto.ValidateSessionRequest{
 		Domain:       "test-proxy.example.com",
-		SessionToken: "invalid-token",
+		SessionToken: "invalid-token", //nolint:staticcheck,
 	})
 
 	require.NoError(t, err)
@@ -410,7 +414,7 @@ func TestValidateSession_MissingDomain(t *testing.T) {
 	defer setup.cleanup()
 
 	resp, err := setup.proxyService.ValidateSession(context.Background(), &proto.ValidateSessionRequest{
-		SessionToken: "some-token",
+		SessionToken: "some-token", //nolint:staticcheck,
 	})
 
 	require.NoError(t, err)
@@ -491,15 +495,15 @@ func (m *testValidateSessionServiceManager) GetAllServices(_ context.Context, _,
 }
 
 func (m *testValidateSessionServiceManager) GetService(_ context.Context, _, _, _ string) (*service.Service, error) {
-	return nil, nil
+	return nil, nil //nolint:nilnil
 }
 
 func (m *testValidateSessionServiceManager) CreateService(_ context.Context, _, _ string, _ *service.Service) (*service.Service, error) {
-	return nil, nil
+	return nil, nil //nolint:nilnil
 }
 
 func (m *testValidateSessionServiceManager) UpdateService(_ context.Context, _, _ string, _ *service.Service) (*service.Service, error) {
-	return nil, nil
+	return nil, nil //nolint:nilnil
 }
 
 func (m *testValidateSessionServiceManager) DeleteService(_ context.Context, _, _, _ string) error {
@@ -543,7 +547,7 @@ func (m *testValidateSessionServiceManager) GetServiceIDByTargetID(_ context.Con
 }
 
 func (m *testValidateSessionServiceManager) CreateServiceFromPeer(_ context.Context, _, _ string, _ *service.ExposeServiceRequest) (*service.ExposeServiceResponse, error) {
-	return nil, nil
+	return nil, nil //nolint:nilnil
 }
 
 func (m *testValidateSessionServiceManager) RenewServiceFromPeer(_ context.Context, _, _, _ string) error {
@@ -571,7 +575,7 @@ func (m *testValidateSessionServiceManager) DeleteAccountCluster(_ context.Conte
 type testValidateSessionProxyManager struct{}
 
 func (m *testValidateSessionProxyManager) Connect(_ context.Context, _, _, _, _, _ string, _ *string, _ *proxy.Capabilities) (*proxy.Proxy, error) {
-	return nil, nil
+	return nil, nil //nolint:nilnil
 }
 
 func (m *testValidateSessionProxyManager) Disconnect(_ context.Context, _, _ string) error {
@@ -603,7 +607,7 @@ func (m *testValidateSessionProxyManager) CleanupStale(_ context.Context, _ time
 }
 
 func (m *testValidateSessionProxyManager) GetAccountProxy(_ context.Context, _ string) (*proxy.Proxy, error) {
-	return nil, nil
+	return nil, nil //nolint:nilnil
 }
 
 func (m *testValidateSessionProxyManager) CountAccountProxies(_ context.Context, _ string) (int64, error) {
@@ -634,6 +638,14 @@ func (m *testValidateSessionProxyManager) ClusterSupportsPrivate(_ context.Conte
 	return nil
 }
 
+func (m *testValidateSessionProxyManager) ClusterAllProxiesPrivate(_ context.Context, _ string) *bool {
+	return nil
+}
+
+func (m *testValidateSessionProxyManager) ClusterSupportsSessionCode(_ context.Context, _ string) bool {
+	return false
+}
+
 type testValidateSessionUsersManager struct {
 	store store.Store
 }
@@ -662,3 +674,169 @@ func (m *testValidateSessionUsersManager) GetUserWithGroups(ctx context.Context,
 	}
 	return user, groups, nil
 }
+
+func TestValidateSession_RedeemsSessionCode(t *testing.T) {
+	setup := setupValidateSessionTest(t)
+	defer setup.cleanup()
+
+	proxy, err := setup.store.GetServiceByID(context.Background(), store.LockingStrengthNone, "testAccountId", "testProxyId")
+	require.NoError(t, err)
+
+	token := createSessionToken(t, proxy.SessionPrivateKey, "allowedUserId", "test-proxy.example.com")
+	code, ok := setup.proxyService.GenerateSessionCode(token)
+	require.True(t, ok)
+	require.NotEqual(t, token, code, "code must not be the token itself")
+
+	resp, err := setup.proxyService.ValidateSession(context.Background(), &proto.ValidateSessionRequest{
+		Domain:      "test-proxy.example.com",
+		SessionCode: code,
+	})
+	require.NoError(t, err)
+	assert.True(t, resp.Valid, "redeemed code should authorize the user")
+	assert.Equal(t, "allowedUserId", resp.UserId)
+	assert.Equal(t, token, resp.GetSessionToken(), "response must carry the durable token for the cookie")
+
+	// Single-use: the same code must not redeem again.
+	resp2, err := setup.proxyService.ValidateSession(context.Background(), &proto.ValidateSessionRequest{
+		Domain:      "test-proxy.example.com",
+		SessionCode: code,
+	})
+	require.NoError(t, err)
+	assert.False(t, resp2.Valid, "a consumed code must be rejected")
+	assert.Empty(t, resp2.GetSessionToken())
+}
+
+func TestValidateSession_InvalidSessionCode(t *testing.T) {
+	setup := setupValidateSessionTest(t)
+	defer setup.cleanup()
+
+	resp, err := setup.proxyService.ValidateSession(context.Background(), &proto.ValidateSessionRequest{
+		Domain:      "test-proxy.example.com",
+		SessionCode: "does-not-exist",
+	})
+	require.NoError(t, err)
+	assert.False(t, resp.Valid)
+	assert.Empty(t, resp.GetSessionToken())
+}
+
+// TestGenerateSessionToken_EmailFallsBackToIdP covers users whose stored email
+// column is empty: the token's email claim comes from the IdP so the proxy can
+// still stamp X-NetBird-User.
+func TestGenerateSessionToken_EmailFallsBackToIdP(t *testing.T) {
+	setup := setupValidateSessionTest(t)
+	defer setup.cleanup()
+
+	ctx := context.Background()
+	svc, err := setup.store.GetServiceByID(ctx, store.LockingStrengthNone, "testAccountId", "restrictedProxyId")
+	require.NoError(t, err)
+
+	idpMock := &mockTunnelIdpManager{email: "idp@example.com", hasData: true}
+	setup.proxyService.idpManager = idpMock
+
+	token, err := setup.proxyService.GenerateSessionToken(ctx, "restricted-proxy.example.com", "allowedUserId", auth.MethodOIDC)
+	require.NoError(t, err)
+
+	pubKey, err := base64.StdEncoding.DecodeString(svc.SessionPublicKey)
+	require.NoError(t, err)
+
+	_, email, _, _, _, err := auth.ValidateSessionJWT(token, "restricted-proxy.example.com", pubKey)
+	require.NoError(t, err)
+	assert.Equal(t, "idp@example.com", email)
+	require.Len(t, idpMock.gotMeta, 1)
+	assert.Equal(t, "testAccountId", idpMock.gotMeta[0].WTAccountID, "IdP lookup must be scoped to the service's account")
+}
+
+// TestGenerateSessionToken_IdPFailureStillIssuesToken pins the fail-soft
+// behaviour: an unreachable IdP costs the email claim, not the login.
+func TestGenerateSessionToken_IdPFailureStillIssuesToken(t *testing.T) {
+	setup := setupValidateSessionTest(t)
+	defer setup.cleanup()
+
+	ctx := context.Background()
+	svc, err := setup.store.GetServiceByID(ctx, store.LockingStrengthNone, "testAccountId", "restrictedProxyId")
+	require.NoError(t, err)
+
+	setup.proxyService.idpManager = &mockTunnelIdpManager{err: errors.New("idp down")}
+
+	token, err := setup.proxyService.GenerateSessionToken(ctx, "restricted-proxy.example.com", "allowedUserId", auth.MethodOIDC)
+	require.NoError(t, err)
+
+	pubKey, err := base64.StdEncoding.DecodeString(svc.SessionPublicKey)
+	require.NoError(t, err)
+
+	userID, email, _, _, _, err := auth.ValidateSessionJWT(token, "restricted-proxy.example.com", pubKey)
+	require.NoError(t, err)
+	assert.Equal(t, "allowedUserId", userID)
+	assert.Empty(t, email)
+}
+
+func TestValidateSession_EmailResolution(t *testing.T) {
+	tests := []struct {
+		name        string
+		storedEmail string
+		tokenEmail  string
+		idp         *mockTunnelIdpManager
+		expectEmail string
+		expectCalls int
+	}{
+		{
+			name:        "stored email wins without IdP call",
+			storedEmail: "stored@example.com",
+			tokenEmail:  "claim@example.com",
+			idp:         &mockTunnelIdpManager{email: "idp@example.com", hasData: true},
+			expectEmail: "stored@example.com",
+		},
+		{
+			name:        "token email claim is reused without IdP call",
+			tokenEmail:  "claim@example.com",
+			idp:         &mockTunnelIdpManager{email: "idp@example.com", hasData: true},
+			expectEmail: "claim@example.com",
+		},
+		{
+			name:        "empty stored email and claim fall back to IdP",
+			idp:         &mockTunnelIdpManager{email: "idp@example.com", hasData: true},
+			expectEmail: "idp@example.com",
+			expectCalls: 1,
+		},
+		{
+			name:        "IdP error leaves email empty but session valid",
+			idp:         &mockTunnelIdpManager{err: errors.New("idp down")},
+			expectCalls: 1,
+		},
+		{
+			name:        "IdP without data leaves email empty",
+			idp:         &mockTunnelIdpManager{},
+			expectCalls: 1,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			setup := setupValidateSessionTest(t)
+			defer setup.cleanup()
+
+			ctx := context.Background()
+			user, err := setup.store.GetUserByUserID(ctx, store.LockingStrengthNone, "allowedUserId")
+			require.NoError(t, err)
+			user.Email = tt.storedEmail
+			require.NoError(t, setup.store.SaveUser(ctx, user))
+
+			setup.proxyService.idpManager = tt.idp
+
+			svc, err := setup.store.GetServiceByID(ctx, store.LockingStrengthNone, "testAccountId", "testProxyId")
+			require.NoError(t, err)
+			token := createSessionTokenWithEmail(t, svc.SessionPrivateKey, "allowedUserId", tt.tokenEmail, "test-proxy.example.com")
+			code, ok := setup.proxyService.GenerateSessionCode(token)
+			require.True(t, ok)
+
+			resp, err := setup.proxyService.ValidateSession(ctx, &proto.ValidateSessionRequest{
+				Domain:      "test-proxy.example.com",
+				SessionCode: code,
+			})
+			require.NoError(t, err)
+			assert.True(t, resp.GetValid())
+			assert.Equal(t, tt.expectEmail, resp.GetUserEmail())
+			assert.Equal(t, tt.expectCalls, tt.idp.gotCalls)
+		})
+	}
+}