The file /Users/misha/Documents/GolandProjects/netbird/netbird/idp/dex/provider.go has been updated successfully.

// Package dex provides an embedded Dex OIDC identity provider.
package dex

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	dexapi "github.com/dexidp/dex/api/v2"
	"github.com/dexidp/dex/server"
	"github.com/dexidp/dex/storage"
	"github.com/dexidp/dex/storage/sql"
	"github.com/google/uuid"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/crypto/bcrypt"
	"google.golang.org/grpc"
)

// Config matches what management/internals/server/server.go expects
type Config struct {
	Issuer  string
	Port    int
	DataDir string
	DevMode bool

	// GRPCAddr is the address for the gRPC API (e.g., ":5557"). Empty disables gRPC.
	GRPCAddr string
}

// Provider wraps a Dex server
type Provider struct {
	config       *Config
	yamlConfig   *YAMLConfig
	dexServer    *server.Server
	httpServer   *http.Server
	listener     net.Listener
	grpcServer   *grpc.Server
	grpcListener net.Listener
	storage      storage.Storage
	logger       *slog.Logger
	mu           sync.Mutex
	running      bool
}

// NewProvider creates and initializes the Dex server
func NewProvider(ctx context.Context, config *Config) (*Provider, error) {
	if config.Issuer == "" {
		return nil, fmt.Errorf("issuer is required")
	}
	if config.Port <= 0 {
		return nil, fmt.Errorf("invalid port")
	}
	if config.DataDir == "" {
		return nil, fmt.Errorf("data directory is required")
	}

	logger := slog.New(slog.NewTextHandler(os.Stderr, nil))

	// Ensure data directory exists
	if err := os.MkdirAll(config.DataDir, 0700); err != nil {
		return nil, fmt.Errorf("failed to create data directory: %w", err)
	}

	// Initialize SQLite storage
	dbPath := filepath.Join(config.DataDir, "oidc.db")
	sqliteConfig := &sql.SQLite3{File: dbPath}
	stor, err := sqliteConfig.Open(logger)
	if err != nil {
		return nil, fmt.Errorf("failed to open storage: %w", err)
	}

	// Ensure a local connector exists (for password authentication)
	if err := ensureLocalConnector(ctx, stor); err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to ensure local connector: %w", err)
	}

	// Ensure issuer ends with /oauth2 for proper path mounting
	issuer := strings.TrimSuffix(config.Issuer, "/")
	if !strings.HasSuffix(issuer, "/oauth2") {
		issuer += "/oauth2"
	}

	// Build refresh token policy (required to avoid nil pointer panics)
	refreshPolicy, err := server.NewRefreshTokenPolicy(logger, false, "", "", "")
	if err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to create refresh token policy: %w", err)
	}

	// Build Dex server config - use Dex's types directly
	dexConfig := server.Config{
		Issuer:                     issuer,
		Storage:                    stor,
		SkipApprovalScreen:         true,
		SupportedResponseTypes:     []string{"code"},
		ContinueOnConnectorFailure: true,
		Logger:                     logger,
		PrometheusRegistry:         prometheus.NewRegistry(),
		RotateKeysAfter:            6 * time.Hour,
		IDTokensValidFor:           24 * time.Hour,
		RefreshTokenPolicy:         refreshPolicy,
		Web: server.WebConfig{
			Issuer: "NetBird",
		},
	}

	dexSrv, err := server.NewServer(ctx, dexConfig)
	if err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to create dex server: %w", err)
	}

	return &Provider{
		config:    config,
		dexServer: dexSrv,
		storage:   stor,
		logger:    logger,
	}, nil
}

// NewProviderFromYAML creates and initializes the Dex server from a YAMLConfig
func NewProviderFromYAML(ctx context.Context, yamlConfig *YAMLConfig) (*Provider, error) {
	// Configure log level from config, default to WARN to avoid logging sensitive data (emails)
	logLevel := slog.LevelWarn
	if yamlConfig.Logger.Level != "" {
		switch strings.ToLower(yamlConfig.Logger.Level) {
		case "debug":
			logLevel = slog.LevelDebug
		case "info":
			logLevel = slog.LevelInfo
		case "warn", "warning":
			logLevel = slog.LevelWarn
		case "error":
			logLevel = slog.LevelError
		}
	}
	logger := slog.New(NewLogrusHandler(logLevel))

	stor, err := yamlConfig.Storage.OpenStorage(logger)
	if err != nil {
		return nil, fmt.Errorf("failed to open storage: %w", err)
	}

	if err := initializeStorage(ctx, stor, yamlConfig); err != nil {
		stor.Close()
		return nil, err
	}

	dexConfig := buildDexConfig(yamlConfig, stor, logger)
	dexConfig.RefreshTokenPolicy, err = yamlConfig.GetRefreshTokenPolicy(logger)
	if err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to create refresh token policy: %w", err)
	}

	dexSrv, err := server.NewServer(ctx, dexConfig)
	if err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to create dex server: %w", err)
	}

	return &Provider{
		config:     &Config{Issuer: yamlConfig.Issuer, GRPCAddr: yamlConfig.GRPC.Addr},
		yamlConfig: yamlConfig,
		dexServer:  dexSrv,
		storage:    stor,
		logger:     logger,
	}, nil
}

// initializeStorage sets up connectors, passwords, and clients in storage
func initializeStorage(ctx context.Context, stor storage.Storage, cfg *YAMLConfig) error {
	if cfg.EnablePasswordDB {
		if err := ensureLocalConnector(ctx, stor); err != nil {
			return fmt.Errorf("failed to ensure local connector: %w", err)
		}
	}
	if err := ensureStaticPasswords(ctx, stor, cfg.StaticPasswords); err != nil {
		return err
	}
	if err := ensureStaticClients(ctx, stor, cfg.StaticClients); err != nil {
		return err
	}
	return ensureStaticConnectors(ctx, stor, cfg.StaticConnectors)
}

// ensureStaticPasswords creates or updates static passwords in storage
func ensureStaticPasswords(ctx context.Context, stor storage.Storage, passwords []Password) error {
	for _, pw := range passwords {
		existing, err := stor.GetPassword(ctx, pw.Email)
		if errors.Is(err, storage.ErrNotFound) {
			if err := stor.CreatePassword(ctx, storage.Password(pw)); err != nil {
				return fmt.Errorf("failed to create password for %s: %w", pw.Email, err)
			}
			continue
		}
		if err != nil {
			return fmt.Errorf("failed to get password for %s: %w", pw.Email, err)
		}
		if string(existing.Hash) != string(pw.Hash) {
			if err := stor.UpdatePassword(ctx, pw.Email, func(old storage.Password) (storage.Password, error) {
				old.Hash = pw.Hash
				old.Username = pw.Username
				return old, nil
			}); err != nil {
				return fmt.Errorf("failed to update password for %s: %w", pw.Email, err)
			}
		}
	}
	return nil
}

// ensureStaticClients creates or updates static clients in storage
func ensureStaticClients(ctx context.Context, stor storage.Storage, clients []storage.Client) error {
	for _, client := range clients {
		_, err := stor.GetClient(ctx, client.ID)
		if errors.Is(err, storage.ErrNotFound) {
			if err := stor.CreateClient(ctx, client); err != nil {
				return fmt.Errorf("failed to create client %s: %w", client.ID, err)
			}
			continue
		}
		if err != nil {
			return fmt.Errorf("failed to get client %s: %w", client.ID, err)
		}
		if err := stor.UpdateClient(ctx, client.ID, func(old storage.Client) (storage.Client, error) {
			old.RedirectURIs = client.RedirectURIs
			old.Name = client.Name
			old.Public = client.Public
			return old, nil
		}); err != nil {
			return fmt.Errorf("failed to update client %s: %w", client.ID, err)
		}
	}
	return nil
}

// buildDexConfig creates a server.Config with defaults applied
func buildDexConfig(yamlConfig *YAMLConfig, stor storage.Storage, logger *slog.Logger) server.Config {
	cfg := yamlConfig.ToServerConfig(stor, logger)
	cfg.PrometheusRegistry = prometheus.NewRegistry()
	if cfg.RotateKeysAfter == 0 {
		cfg.RotateKeysAfter = 24 * 30 * time.Hour
	}
	if cfg.IDTokensValidFor == 0 {
		cfg.IDTokensValidFor = 24 * time.Hour
	}
	if cfg.Web.Issuer == "" {
		cfg.Web.Issuer = "NetBird"
	}
	if len(cfg.SupportedResponseTypes) == 0 {
		cfg.SupportedResponseTypes = []string{"code"}
	}
	cfg.ContinueOnConnectorFailure = true
	return cfg
}

// Start starts the HTTP server and optionally the gRPC API server
func (p *Provider) Start(_ context.Context) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.running {
		return fmt.Errorf("already running")
	}

	// Determine listen address from config
	var addr string
	if p.yamlConfig != nil {
		addr = p.yamlConfig.Web.HTTP
		if addr == "" {
			addr = p.yamlConfig.Web.HTTPS
		}
	} else if p.config != nil && p.config.Port > 0 {
		addr = fmt.Sprintf(":%d", p.config.Port)
	}
	if addr == "" {
		return fmt.Errorf("no listen address configured")
	}

	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return fmt.Errorf("failed to listen on %s: %w", addr, err)
	}
	p.listener = listener

	// Mount Dex at /oauth2/ path for reverse proxy compatibility
	// Don't strip the prefix - Dex's issuer includes /oauth2 so it expects the full path
	mux := http.NewServeMux()
	mux.Handle("/oauth2/", p.dexServer)

	p.httpServer = &http.Server{Handler: mux}
	p.running = true

	go func() {
		if err := p.httpServer.Serve(listener); err != nil && err != http.ErrServerClosed {
			p.logger.Error("http server error", "error", err)
		}
	}()

	// Start gRPC API server if configured
	if p.config.GRPCAddr != "" {
		if err := p.startGRPCServer(); err != nil {
			// Clean up HTTP server on failure
			_ = p.httpServer.Close()
			_ = p.listener.Close()
			return fmt.Errorf("failed to start gRPC server: %w", err)
		}
	}

	p.logger.Info("HTTP server started", "addr", addr)
	return nil
}

// startGRPCServer starts the gRPC API server using Dex's built-in API
func (p *Provider) startGRPCServer() error {
	grpcListener, err := net.Listen("tcp", p.config.GRPCAddr)
	if err != nil {
		return fmt.Errorf("failed to listen on %s: %w", p.config.GRPCAddr, err)
	}
	p.grpcListener = grpcListener

	p.grpcServer = grpc.NewServer()
	// Use Dex's built-in API server implementation
	// server.NewAPI(storage, logger, version, dexServer)
	dexapi.RegisterDexServer(p.grpcServer, server.NewAPI(p.storage, p.logger, "netbird-dex", p.dexServer))

	go func() {
		if err := p.grpcServer.Serve(grpcListener); err != nil {
			p.logger.Error("grpc server error", "error", err)
		}
	}()

	p.logger.Info("gRPC API server started", "addr", p.config.GRPCAddr)
	return nil
}

// Stop gracefully shuts down
func (p *Provider) Stop(ctx context.Context) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if !p.running {
		return nil
	}

	var errs []error

	// Stop gRPC server first
	if p.grpcServer != nil {
		p.grpcServer.GracefulStop()
		p.grpcServer = nil
	}
	if p.grpcListener != nil {
		p.grpcListener.Close()
		p.grpcListener = nil
	}

	if p.httpServer != nil {
		if err := p.httpServer.Shutdown(ctx); err != nil {
			errs = append(errs, err)
		}
	}

	// Explicitly close listener as fallback (Shutdown should do this, but be safe)
	if p.listener != nil {
		if err := p.listener.Close(); err != nil {
			// Ignore "use of closed network connection" - expected after Shutdown
			if !strings.Contains(err.Error(), "use of closed") {
				errs = append(errs, err)
			}
		}
		p.listener = nil
	}

	if p.storage != nil {
		if err := p.storage.Close(); err != nil {
			errs = append(errs, err)
		}
	}

	p.httpServer = nil
	p.running = false

	if len(errs) > 0 {
		return fmt.Errorf("shutdown errors: %v", errs)
	}
	return nil
}

// EnsureDefaultClients creates dashboard and CLI OAuth clients
// Uses Dex's storage.Client directly - no custom wrappers
func (p *Provider) EnsureDefaultClients(ctx context.Context, dashboardURIs, cliURIs []string) error {
	clients := []storage.Client{
		{
			ID:           "netbird-dashboard",
			Name:         "NetBird Dashboard",
			RedirectURIs: dashboardURIs,
			Public:       true,
		},
		{
			ID:           "netbird-cli",
			Name:         "NetBird CLI",
			RedirectURIs: cliURIs,
			Public:       true,
		},
	}

	for _, client := range clients {
		_, err := p.storage.GetClient(ctx, client.ID)
		if err == storage.ErrNotFound {
			if err := p.storage.CreateClient(ctx, client); err != nil {
				return fmt.Errorf("failed to create client %s: %w", client.ID, err)
			}
			continue
		}
		if err != nil {
			return fmt.Errorf("failed to get client %s: %w", client.ID, err)
		}
		// Update if exists
		if err := p.storage.UpdateClient(ctx, client.ID, func(old storage.Client) (storage.Client, error) {
			old.RedirectURIs = client.RedirectURIs
			return old, nil
		}); err != nil {
			return fmt.Errorf("failed to update client %s: %w", client.ID, err)
		}
	}

	p.logger.Info("default OIDC clients ensured")
	return nil
}

// Storage returns the underlying Dex storage for direct access
// Users can use storage.Client, storage.Password, storage.Connector directly
func (p *Provider) Storage() storage.Storage {
	return p.storage
}

// Handler returns the Dex server as an http.Handler for embedding in another server.
// The handler expects requests with path prefix "/oauth2/".
func (p *Provider) Handler() http.Handler {
	return p.dexServer
}

// CreateUser creates a new user with the given email, username, and password.
// Returns the encoded user ID in Dex's format (base64-encoded protobuf with connector ID).
func (p *Provider) CreateUser(ctx context.Context, email, username, password string) (string, error) {
	hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		return "", fmt.Errorf("failed to hash password: %w", err)
	}

	userID := uuid.New().String()
	err = p.storage.CreatePassword(ctx, storage.Password{
		Email:    email,
		Username: username,
		UserID:   userID,
		Hash:     hash,
	})
	if err != nil {
		return "", err
	}

	// Encode the user ID in Dex's format: base64(protobuf{user_id, connector_id})
	// This matches the format Dex uses in JWT tokens
	encodedID := EncodeDexUserID(userID, "local")
	return encodedID, nil
}

// EncodeDexUserID encodes user ID and connector ID into Dex's base64-encoded protobuf format.
// Dex uses this format for the 'sub' claim in JWT tokens.
// Format: base64(protobuf message with field 1 = user_id, field 2 = connector_id)
func EncodeDexUserID(userID, connectorID string) string {
	// Manually encode protobuf: field 1 (user_id) and field 2 (connector_id)
	// Wire type 2 (length-delimited) for strings
	var buf []byte

	// Field 1: user_id (tag = 0x0a = field 1, wire type 2)
	buf = append(buf, 0x0a)
	buf = append(buf, byte(len(userID)))
	buf = append(buf, []byte(userID)...)

	// Field 2: connector_id (tag = 0x12 = field 2, wire type 2)
	buf = append(buf, 0x12)
	buf = append(buf, byte(len(connectorID)))
	buf = append(buf, []byte(connectorID)...)

	return base64.RawStdEncoding.EncodeToString(buf)
}

// DecodeDexUserID decodes Dex's base64-encoded user ID back to the raw user ID and connector ID.
func DecodeDexUserID(encodedID string) (userID, connectorID string, err error) {
	// Try RawStdEncoding first, then StdEncoding (with padding)
	buf, err := base64.RawStdEncoding.DecodeString(encodedID)
	if err != nil {
		buf, err = base64.StdEncoding.DecodeString(encodedID)
		if err != nil {
			return "", "", fmt.Errorf("failed to decode base64: %w", err)
		}
	}

	// Parse protobuf manually
	i := 0
	for i < len(buf) {
		if i >= len(buf) {
			break
		}
		tag := buf[i]
		i++

		fieldNum := tag >> 3
		wireType := tag & 0x07

		if wireType != 2 { // We only expect length-delimited strings
			return "", "", fmt.Errorf("unexpected wire type %d", wireType)
		}

		if i >= len(buf) {
			return "", "", fmt.Errorf("truncated message")
		}
		length := int(buf[i])
		i++

		if i+length > len(buf) {
			return "", "", fmt.Errorf("truncated string field")
		}
		value := string(buf[i : i+length])
		i += length

		switch fieldNum {
		case 1:
			userID = value
		case 2:
			connectorID = value
		}
	}

	return userID, connectorID, nil
}

// GetUser returns a user by email
func (p *Provider) GetUser(ctx context.Context, email string) (storage.Password, error) {
	return p.storage.GetPassword(ctx, email)
}

// GetUserByID returns a user by user ID.
// The userID can be either an encoded Dex ID (base64 protobuf) or a raw UUID.
// Note: This requires iterating through all users since dex storage doesn't index by userID.
func (p *Provider) GetUserByID(ctx context.Context, userID string) (storage.Password, error) {
	// Try to decode the user ID in case it's encoded
	rawUserID, _, err := DecodeDexUserID(userID)
	if err != nil {
		// If decoding fails, assume it's already a raw UUID
		rawUserID = userID
	}

	users, err := p.storage.ListPasswords(ctx)
	if err != nil {
		return storage.Password{}, fmt.Errorf("failed to list users: %w", err)
	}
	for _, user := range users {
		if user.UserID == rawUserID {
			return user, nil
		}
	}
	return storage.Password{}, storage.ErrNotFound
}

// DeleteUser removes a user by email
func (p *Provider) DeleteUser(ctx context.Context, email string) error {
	return p.storage.DeletePassword(ctx, email)
}

// ListUsers returns all users
func (p *Provider) ListUsers(ctx context.Context) ([]storage.Password, error) {
	return p.storage.ListPasswords(ctx)
}

// UpdateUserPassword updates the password for a user identified by userID.
// The userID can be either an encoded Dex ID (base64 protobuf) or a raw UUID.
// It verifies the current password before updating.
func (p *Provider) UpdateUserPassword(ctx context.Context, userID string, oldPassword, newPassword string) error {
	// Get the user by ID to find their email
	user, err := p.GetUserByID(ctx, userID)
	if err != nil {
		return fmt.Errorf("failed to get user: %w", err)
	}

	// Verify old password
	if err := bcrypt.CompareHashAndPassword(user.Hash, []byte(oldPassword)); err != nil {
		return fmt.Errorf("current password is incorrect")
	}

	// Hash the new password
	newHash, err := bcrypt.GenerateFromPassword([]byte(newPassword), bcrypt.DefaultCost)
	if err != nil {
		return fmt.Errorf("failed to hash new password: %w", err)
	}

	// Update the password in storage
	err = p.storage.UpdatePassword(ctx, user.Email, func(old storage.Password) (storage.Password, error) {
		old.Hash = newHash
		return old, nil
	})
	if err != nil {
		return fmt.Errorf("failed to update password: %w", err)
	}

	return nil
}

// GetIssuer returns the OIDC issuer URL.
func (p *Provider) GetIssuer() string {
	if p.config == nil {
		return ""
	}
	issuer := strings.TrimSuffix(p.config.Issuer, "/")
	if !strings.HasSuffix(issuer, "/oauth2") {
		issuer += "/oauth2"
	}
	return issuer
}

// GetKeysLocation returns the JWKS endpoint URL for token validation.
func (p *Provider) GetKeysLocation() string {
	issuer := p.GetIssuer()
	if issuer == "" {
		return ""
	}
	return issuer + "/keys"
}

// GetTokenEndpoint returns the OAuth2 token endpoint URL.
func (p *Provider) GetTokenEndpoint() string {
	issuer := p.GetIssuer()
	if issuer == "" {
		return ""
	}
	return issuer + "/token"
}

// GetDeviceAuthEndpoint returns the OAuth2 device authorization endpoint URL.
func (p *Provider) GetDeviceAuthEndpoint() string {
	issuer := p.GetIssuer()
	if issuer == "" {
		return ""
	}
	return issuer + "/device/code"
}

// GetAuthorizationEndpoint returns the OAuth2 authorization endpoint URL.
func (p *Provider) GetAuthorizationEndpoint() string {
	issuer := p.GetIssuer()
	if issuer == "" {
		return ""
	}
	return issuer + "/auth"
}

// Package dex provides an embedded Dex OIDC identity provider.
package dex

import (
	"context"
	"crypto/ecdsa"
	"crypto/rsa"
	"encoding/base64"
	"errors"
	"fmt"
	"log/slog"
	"math/big"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	dexapi "github.com/dexidp/dex/api/v2"
	"github.com/dexidp/dex/server"
	"github.com/dexidp/dex/storage"
	"github.com/dexidp/dex/storage/sql"
	jose "github.com/go-jose/go-jose/v4"
	"github.com/google/uuid"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/crypto/bcrypt"
	"google.golang.org/grpc"

	nbjwt "github.com/netbirdio/netbird/shared/auth/jwt"
)

// Config matches what management/internals/server/server.go expects
type Config struct {
	Issuer  string
	Port    int
	DataDir string
	DevMode bool

	// GRPCAddr is the address for the gRPC API (e.g., ":5557"). Empty disables gRPC.
	GRPCAddr string
}

// Provider wraps a Dex server
type Provider struct {
	config       *Config
	yamlConfig   *YAMLConfig
	dexServer    *server.Server
	httpServer   *http.Server
	listener     net.Listener
	grpcServer   *grpc.Server
	grpcListener net.Listener
	storage      storage.Storage
	logger       *slog.Logger
	mu           sync.Mutex
	running      bool
}

// NewProvider creates and initializes the Dex server
func NewProvider(ctx context.Context, config *Config) (*Provider, error) {
	if config.Issuer == "" {
		return nil, fmt.Errorf("issuer is required")
	}
	if config.Port <= 0 {
		return nil, fmt.Errorf("invalid port")
	}
	if config.DataDir == "" {
		return nil, fmt.Errorf("data directory is required")
	}

	logger := slog.New(slog.NewTextHandler(os.Stderr, nil))

	// Ensure data directory exists
	if err := os.MkdirAll(config.DataDir, 0700); err != nil {
		return nil, fmt.Errorf("failed to create data directory: %w", err)
	}

	// Initialize SQLite storage
	dbPath := filepath.Join(config.DataDir, "oidc.db")
	sqliteConfig := &sql.SQLite3{File: dbPath}
	stor, err := sqliteConfig.Open(logger)
	if err != nil {
		return nil, fmt.Errorf("failed to open storage: %w", err)
	}

	// Ensure a local connector exists (for password authentication)
	if err := ensureLocalConnector(ctx, stor); err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to ensure local connector: %w", err)
	}

	// Ensure issuer ends with /oauth2 for proper path mounting
	issuer := strings.TrimSuffix(config.Issuer, "/")
	if !strings.HasSuffix(issuer, "/oauth2") {
		issuer += "/oauth2"
	}

	// Build refresh token policy (required to avoid nil pointer panics)
	refreshPolicy, err := server.NewRefreshTokenPolicy(logger, false, "", "", "")
	if err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to create refresh token policy: %w", err)
	}

	// Build Dex server config - use Dex's types directly
	dexConfig := server.Config{
		Issuer:                     issuer,
		Storage:                    stor,
		SkipApprovalScreen:         true,
		SupportedResponseTypes:     []string{"code"},
		ContinueOnConnectorFailure: true,
		Logger:                     logger,
		PrometheusRegistry:         prometheus.NewRegistry(),
		RotateKeysAfter:            6 * time.Hour,
		IDTokensValidFor:           24 * time.Hour,
		RefreshTokenPolicy:         refreshPolicy,
		Web: server.WebConfig{
			Issuer: "NetBird",
		},
	}

	dexSrv, err := server.NewServer(ctx, dexConfig)
	if err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to create dex server: %w", err)
	}

	return &Provider{
		config:    config,
		dexServer: dexSrv,
		storage:   stor,
		logger:    logger,
	}, nil
}

// NewProviderFromYAML creates and initializes the Dex server from a YAMLConfig
func NewProviderFromYAML(ctx context.Context, yamlConfig *YAMLConfig) (*Provider, error) {
	// Configure log level from config, default to WARN to avoid logging sensitive data (emails)
	logLevel := slog.LevelWarn
	if yamlConfig.Logger.Level != "" {
		switch strings.ToLower(yamlConfig.Logger.Level) {
		case "debug":
			logLevel = slog.LevelDebug
		case "info":
			logLevel = slog.LevelInfo
		case "warn", "warning":
			logLevel = slog.LevelWarn
		case "error":
			logLevel = slog.LevelError
		}
	}
	logger := slog.New(NewLogrusHandler(logLevel))

	stor, err := yamlConfig.Storage.OpenStorage(logger)
	if err != nil {
		return nil, fmt.Errorf("failed to open storage: %w", err)
	}

	if err := initializeStorage(ctx, stor, yamlConfig); err != nil {
		stor.Close()
		return nil, err
	}

	dexConfig := buildDexConfig(yamlConfig, stor, logger)
	dexConfig.RefreshTokenPolicy, err = yamlConfig.GetRefreshTokenPolicy(logger)
	if err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to create refresh token policy: %w", err)
	}

	dexSrv, err := server.NewServer(ctx, dexConfig)
	if err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to create dex server: %w", err)
	}

	return &Provider{
		config:     &Config{Issuer: yamlConfig.Issuer, GRPCAddr: yamlConfig.GRPC.Addr},
		yamlConfig: yamlConfig,
		dexServer:  dexSrv,
		storage:    stor,
		logger:     logger,
	}, nil
}

// initializeStorage sets up connectors, passwords, and clients in storage
func initializeStorage(ctx context.Context, stor storage.Storage, cfg *YAMLConfig) error {
	if cfg.EnablePasswordDB {
		if err := ensureLocalConnector(ctx, stor); err != nil {
			return fmt.Errorf("failed to ensure local connector: %w", err)
		}
	}
	if err := ensureStaticPasswords(ctx, stor, cfg.StaticPasswords); err != nil {
		return err
	}
	if err := ensureStaticClients(ctx, stor, cfg.StaticClients); err != nil {
		return err
	}
	return ensureStaticConnectors(ctx, stor, cfg.StaticConnectors)
}

// ensureStaticPasswords creates or updates static passwords in storage
func ensureStaticPasswords(ctx context.Context, stor storage.Storage, passwords []Password) error {
	for _, pw := range passwords {
		existing, err := stor.GetPassword(ctx, pw.Email)
		if errors.Is(err, storage.ErrNotFound) {
			if err := stor.CreatePassword(ctx, storage.Password(pw)); err != nil {
				return fmt.Errorf("failed to create password for %s: %w", pw.Email, err)
			}
			continue
		}
		if err != nil {
			return fmt.Errorf("failed to get password for %s: %w", pw.Email, err)
		}
		if string(existing.Hash) != string(pw.Hash) {
			if err := stor.UpdatePassword(ctx, pw.Email, func(old storage.Password) (storage.Password, error) {
				old.Hash = pw.Hash
				old.Username = pw.Username
				return old, nil
			}); err != nil {
				return fmt.Errorf("failed to update password for %s: %w", pw.Email, err)
			}
		}
	}
	return nil
}

// ensureStaticClients creates or updates static clients in storage
func ensureStaticClients(ctx context.Context, stor storage.Storage, clients []storage.Client) error {
	for _, client := range clients {
		_, err := stor.GetClient(ctx, client.ID)
		if errors.Is(err, storage.ErrNotFound) {
			if err := stor.CreateClient(ctx, client); err != nil {
				return fmt.Errorf("failed to create client %s: %w", client.ID, err)
			}
			continue
		}
		if err != nil {
			return fmt.Errorf("failed to get client %s: %w", client.ID, err)
		}
		if err := stor.UpdateClient(ctx, client.ID, func(old storage.Client) (storage.Client, error) {
			old.RedirectURIs = client.RedirectURIs
			old.Name = client.Name
			old.Public = client.Public
			return old, nil
		}); err != nil {
			return fmt.Errorf("failed to update client %s: %w", client.ID, err)
		}
	}
	return nil
}

// buildDexConfig creates a server.Config with defaults applied
func buildDexConfig(yamlConfig *YAMLConfig, stor storage.Storage, logger *slog.Logger) server.Config {
	cfg := yamlConfig.ToServerConfig(stor, logger)
	cfg.PrometheusRegistry = prometheus.NewRegistry()
	if cfg.RotateKeysAfter == 0 {
		cfg.RotateKeysAfter = 24 * 30 * time.Hour
	}
	if cfg.IDTokensValidFor == 0 {
		cfg.IDTokensValidFor = 24 * time.Hour
	}
	if cfg.Web.Issuer == "" {
		cfg.Web.Issuer = "NetBird"
	}
	if len(cfg.SupportedResponseTypes) == 0 {
		cfg.SupportedResponseTypes = []string{"code"}
	}
	cfg.ContinueOnConnectorFailure = true
	return cfg
}

// Start starts the HTTP server and optionally the gRPC API server
func (p *Provider) Start(_ context.Context) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.running {
		return fmt.Errorf("already running")
	}

	// Determine listen address from config
	var addr string
	if p.yamlConfig != nil {
		addr = p.yamlConfig.Web.HTTP
		if addr == "" {
			addr = p.yamlConfig.Web.HTTPS
		}
	} else if p.config != nil && p.config.Port > 0 {
		addr = fmt.Sprintf(":%d", p.config.Port)
	}
	if addr == "" {
		return fmt.Errorf("no listen address configured")
	}

	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return fmt.Errorf("failed to listen on %s: %w", addr, err)
	}
	p.listener = listener

	// Mount Dex at /oauth2/ path for reverse proxy compatibility
	// Don't strip the prefix - Dex's issuer includes /oauth2 so it expects the full path
	mux := http.NewServeMux()
	mux.Handle("/oauth2/", p.dexServer)

	p.httpServer = &http.Server{Handler: mux}
	p.running = true

	go func() {
		if err := p.httpServer.Serve(listener); err != nil && err != http.ErrServerClosed {
			p.logger.Error("http server error", "error", err)
		}
	}()

	// Start gRPC API server if configured
	if p.config.GRPCAddr != "" {
		if err := p.startGRPCServer(); err != nil {
			// Clean up HTTP server on failure
			_ = p.httpServer.Close()
			_ = p.listener.Close()
			return fmt.Errorf("failed to start gRPC server: %w", err)
		}
	}

	p.logger.Info("HTTP server started", "addr", addr)
	return nil
}

// startGRPCServer starts the gRPC API server using Dex's built-in API
func (p *Provider) startGRPCServer() error {
	grpcListener, err := net.Listen("tcp", p.config.GRPCAddr)
	if err != nil {
		return fmt.Errorf("failed to listen on %s: %w", p.config.GRPCAddr, err)
	}
	p.grpcListener = grpcListener

	p.grpcServer = grpc.NewServer()
	// Use Dex's built-in API server implementation
	// server.NewAPI(storage, logger, version, dexServer)
	dexapi.RegisterDexServer(p.grpcServer, server.NewAPI(p.storage, p.logger, "netbird-dex", p.dexServer))

	go func() {
		if err := p.grpcServer.Serve(grpcListener); err != nil {
			p.logger.Error("grpc server error", "error", err)
		}
	}()

	p.logger.Info("gRPC API server started", "addr", p.config.GRPCAddr)
	return nil
}

// Stop gracefully shuts down
func (p *Provider) Stop(ctx context.Context) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if !p.running {
		return nil
	}

	var errs []error

	// Stop gRPC server first
	if p.grpcServer != nil {
		p.grpcServer.GracefulStop()
		p.grpcServer = nil
	}
	if p.grpcListener != nil {
		p.grpcListener.Close()
		p.grpcListener = nil
	}

	if p.httpServer != nil {
		if err := p.httpServer.Shutdown(ctx); err != nil {
			errs = append(errs, err)
		}
	}

	// Explicitly close listener as fallback (Shutdown should do this, but be safe)
	if p.listener != nil {
		if err := p.listener.Close(); err != nil {
			// Ignore "use of closed network connection" - expected after Shutdown
			if !strings.Contains(err.Error(), "use of closed") {
				errs = append(errs, err)
			}
		}
		p.listener = nil
	}

	if p.storage != nil {
		if err := p.storage.Close(); err != nil {
			errs = append(errs, err)
		}
	}

	p.httpServer = nil
	p.running = false

	if len(errs) > 0 {
		return fmt.Errorf("shutdown errors: %v", errs)
	}
	return nil
}

// EnsureDefaultClients creates dashboard and CLI OAuth clients
// Uses Dex's storage.Client directly - no custom wrappers
func (p *Provider) EnsureDefaultClients(ctx context.Context, dashboardURIs, cliURIs []string) error {
	clients := []storage.Client{
		{
			ID:           "netbird-dashboard",
			Name:         "NetBird Dashboard",
			RedirectURIs: dashboardURIs,
			Public:       true,
		},
		{
			ID:           "netbird-cli",
			Name:         "NetBird CLI",
			RedirectURIs: cliURIs,
			Public:       true,
		},
	}

	for _, client := range clients {
		_, err := p.storage.GetClient(ctx, client.ID)
		if err == storage.ErrNotFound {
			if err := p.storage.CreateClient(ctx, client); err != nil {
				return fmt.Errorf("failed to create client %s: %w", client.ID, err)
			}
			continue
		}
		if err != nil {
			return fmt.Errorf("failed to get client %s: %w", client.ID, err)
		}
		// Update if exists
		if err := p.storage.UpdateClient(ctx, client.ID, func(old storage.Client) (storage.Client, error) {
			old.RedirectURIs = client.RedirectURIs
			return old, nil
		}); err != nil {
			return fmt.Errorf("failed to update client %s: %w", client.ID, err)
		}
	}

	p.logger.Info("default OIDC clients ensured")
	return nil
}

// Storage returns the underlying Dex storage for direct access
// Users can use storage.Client, storage.Password, storage.Connector directly
func (p *Provider) Storage() storage.Storage {
	return p.storage
}

// Handler returns the Dex server as an http.Handler for embedding in another server.
// The handler expects requests with path prefix "/oauth2/".
func (p *Provider) Handler() http.Handler {
	return p.dexServer
}

// CreateUser creates a new user with the given email, username, and password.
// Returns the encoded user ID in Dex's format (base64-encoded protobuf with connector ID).
func (p *Provider) CreateUser(ctx context.Context, email, username, password string) (string, error) {
	hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		return "", fmt.Errorf("failed to hash password: %w", err)
	}

	userID := uuid.New().String()
	err = p.storage.CreatePassword(ctx, storage.Password{
		Email:    email,
		Username: username,
		UserID:   userID,
		Hash:     hash,
	})
	if err != nil {
		return "", err
	}

	// Encode the user ID in Dex's format: base64(protobuf{user_id, connector_id})
	// This matches the format Dex uses in JWT tokens
	encodedID := EncodeDexUserID(userID, "local")
	return encodedID, nil
}

// EncodeDexUserID encodes user ID and connector ID into Dex's base64-encoded protobuf format.
// Dex uses this format for the 'sub' claim in JWT tokens.
// Format: base64(protobuf message with field 1 = user_id, field 2 = connector_id)
func EncodeDexUserID(userID, connectorID string) string {
	// Manually encode protobuf: field 1 (user_id) and field 2 (connector_id)
	// Wire type 2 (length-delimited) for strings
	var buf []byte

	// Field 1: user_id (tag = 0x0a = field 1, wire type 2)
	buf = append(buf, 0x0a)
	buf = append(buf, byte(len(userID)))
	buf = append(buf, []byte(userID)...)

	// Field 2: connector_id (tag = 0x12 = field 2, wire type 2)
	buf = append(buf, 0x12)
	buf = append(buf, byte(len(connectorID)))
	buf = append(buf, []byte(connectorID)...)

	return base64.RawStdEncoding.EncodeToString(buf)
}

// DecodeDexUserID decodes Dex's base64-encoded user ID back to the raw user ID and connector ID.
func DecodeDexUserID(encodedID string) (userID, connectorID string, err error) {
	// Try RawStdEncoding first, then StdEncoding (with padding)
	buf, err := base64.RawStdEncoding.DecodeString(encodedID)
	if err != nil {
		buf, err = base64.StdEncoding.DecodeString(encodedID)
		if err != nil {
			return "", "", fmt.Errorf("failed to decode base64: %w", err)
		}
	}

	// Parse protobuf manually
	i := 0
	for i < len(buf) {
		if i >= len(buf) {
			break
		}
		tag := buf[i]
		i++

		fieldNum := tag >> 3
		wireType := tag & 0x07

		if wireType != 2 { // We only expect length-delimited strings
			return "", "", fmt.Errorf("unexpected wire type %d", wireType)
		}

		if i >= len(buf) {
			return "", "", fmt.Errorf("truncated message")
		}
		length := int(buf[i])
		i++

		if i+length > len(buf) {
			return "", "", fmt.Errorf("truncated string field")
		}
		value := string(buf[i : i+length])
		i += length

		switch fieldNum {
		case 1:
			userID = value
		case 2:
			connectorID = value
		}
	}

	return userID, connectorID, nil
}

// GetUser returns a user by email
func (p *Provider) GetUser(ctx context.Context, email string) (storage.Password, error) {
	return p.storage.GetPassword(ctx, email)
}

// GetUserByID returns a user by user ID.
// The userID can be either an encoded Dex ID (base64 protobuf) or a raw UUID.
// Note: This requires iterating through all users since dex storage doesn't index by userID.
func (p *Provider) GetUserByID(ctx context.Context, userID string) (storage.Password, error) {
	// Try to decode the user ID in case it's encoded
	rawUserID, _, err := DecodeDexUserID(userID)
	if err != nil {
		// If decoding fails, assume it's already a raw UUID
		rawUserID = userID
	}

	users, err := p.storage.ListPasswords(ctx)
	if err != nil {
		return storage.Password{}, fmt.Errorf("failed to list users: %w", err)
	}
	for _, user := range users {
		if user.UserID == rawUserID {
			return user, nil
		}
	}
	return storage.Password{}, storage.ErrNotFound
}

// DeleteUser removes a user by email
func (p *Provider) DeleteUser(ctx context.Context, email string) error {
	return p.storage.DeletePassword(ctx, email)
}

// ListUsers returns all users
func (p *Provider) ListUsers(ctx context.Context) ([]storage.Password, error) {
	return p.storage.ListPasswords(ctx)
}

// UpdateUserPassword updates the password for a user identified by userID.
// The userID can be either an encoded Dex ID (base64 protobuf) or a raw UUID.
// It verifies the current password before updating.
func (p *Provider) UpdateUserPassword(ctx context.Context, userID string, oldPassword, newPassword string) error {
	// Get the user by ID to find their email
	user, err := p.GetUserByID(ctx, userID)
	if err != nil {
		return fmt.Errorf("failed to get user: %w", err)
	}

	// Verify old password
	if err := bcrypt.CompareHashAndPassword(user.Hash, []byte(oldPassword)); err != nil {
		return fmt.Errorf("current password is incorrect")
	}

	// Hash the new password
	newHash, err := bcrypt.GenerateFromPassword([]byte(newPassword), bcrypt.DefaultCost)
	if err != nil {
		return fmt.Errorf("failed to hash new password: %w", err)
	}

	// Update the password in storage
	err = p.storage.UpdatePassword(ctx, user.Email, func(old storage.Password) (storage.Password, error) {
		old.Hash = newHash
		return old, nil
	})
	if err != nil {
		return fmt.Errorf("failed to update password: %w", err)
	}

	return nil
}

// GetIssuer returns the OIDC issuer URL.
func (p *Provider) GetIssuer() string {
	if p.config == nil {
		return ""
	}
	issuer := strings.TrimSuffix(p.config.Issuer, "/")
	if !strings.HasSuffix(issuer, "/oauth2") {
		issuer += "/oauth2"
	}
	return issuer
}

// GetKeysLocation returns the JWKS endpoint URL for token validation.
func (p *Provider) GetKeysLocation() string {
	issuer := p.GetIssuer()
	if issuer == "" {
		return ""
	}
	return issuer + "/keys"
}

// GetTokenEndpoint returns the OAuth2 token endpoint URL.
func (p *Provider) GetTokenEndpoint() string {
	issuer := p.GetIssuer()
	if issuer == "" {
		return ""
	}
	return issuer + "/token"
}

// GetDeviceAuthEndpoint returns the OAuth2 device authorization endpoint URL.
func (p *Provider) GetDeviceAuthEndpoint() string {
	issuer := p.GetIssuer()
	if issuer == "" {
		return ""
	}
	return issuer + "/device/code"
}

// GetAuthorizationEndpoint returns the OAuth2 authorization endpoint URL.
func (p *Provider) GetAuthorizationEndpoint() string {
	issuer := p.GetIssuer()
	if issuer == "" {
		return ""
	}
	return issuer + "/auth"
}

import (
"context"
+  "crypto/ecdsa"
+  "crypto/rsa"
"encoding/base64"
"errors"
"fmt"
"log/slog"
+  "math/big"
"net"
"net/http"
"os"

+  nbjwt "github.com/netbirdio/netbird/shared/auth/jwt"


Now add the `GetJWKS` method to the Provider. I'll add it after the `GetAuthorizationEndpoint` method.


Edit

// GetAuthorizationEndpoint returns the OAuth2 authorization endpoint URL.
func (p *Provider) GetAuthorizationEndpoint() string {
	issuer := p.GetIssuer()
	if issuer == "" {
		return ""
	}
	return issuer + "/auth"
}

// GetAuthorizationEndpoint returns the OAuth2 authorization endpoint URL.
func (p *Provider) GetAuthorizationEndpoint() string {
	issuer := p.GetIssuer()
	if issuer == "" {
		return ""
	}
	return issuer + "/auth"
}

// GetJWKS reads signing keys directly from Dex storage and returns them as Jwks.
// This avoids HTTP round-trips when the embedded IDP is co-located with the management server.
func (p *Provider) GetJWKS(ctx context.Context) (*nbjwt.Jwks, error) {
	keys, err := p.storage.GetKeys(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to get keys from dex storage: %w", err)
	}

	jwks := &nbjwt.Jwks{}

	if keys.SigningKeyPub != nil {
		jwk, err := joseKeyToJSONWebKey(keys.SigningKeyPub)
		if err == nil {
			jwks.Keys = append(jwks.Keys, jwk)
		}
	}

	for _, vk := range keys.VerificationKeys {
		if vk.PublicKey != nil {
			jwk, err := joseKeyToJSONWebKey(vk.PublicKey)
			if err == nil {
				jwks.Keys = append(jwks.Keys, jwk)
			}
		}
	}

	if !keys.NextRotation.IsZero() {
		jwks.ExpiresInTime = keys.NextRotation
	} else {
		jwks.ExpiresInTime = time.Now().Add(1 * time.Hour)
	}

	return jwks, nil
}

// joseKeyToJSONWebKey converts a go-jose JSONWebKey to our JSONWebKey representation.
func joseKeyToJSONWebKey(key *jose.JSONWebKey) (nbjwt.JSONWebKey, error) {
	jwk := nbjwt.JSONWebKey{
		Kid: key.KeyID,
		Use: key.Use,
	}

	switch pub := key.Key.(type) {
	case *rsa.PublicKey:
		jwk.Kty = "RSA"
		jwk.N = base64.RawURLEncoding.EncodeToString(pub.N.Bytes())
		jwk.E = base64.RawURLEncoding.EncodeToString(big.NewInt(int64(pub.E)).Bytes())
	case *ecdsa.PublicKey:
		jwk.Kty = "EC"
		jwk.X = base64.RawURLEncoding.EncodeToString(pub.X.Bytes())
		jwk.Y = base64.RawURLEncoding.EncodeToString(pub.Y.Bytes())
		jwk.Crv = pub.Curve.Params().Name
	default:
		return jwk, fmt.Errorf("unsupported key type: %T", key.Key)
	}

	return jwk, nil
}


// Package dex provides an embedded Dex OIDC identity provider.
package dex

import (
	"context"
	"crypto/ecdsa"
	"crypto/rsa"
	"encoding/base64"
	"errors"
	"fmt"
	"log/slog"
	"math/big"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	dexapi "github.com/dexidp/dex/api/v2"
	"github.com/dexidp/dex/server"
	"github.com/dexidp/dex/storage"
	"github.com/dexidp/dex/storage/sql"
	jose "github.com/go-jose/go-jose/v4"
	"github.com/google/uuid"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/crypto/bcrypt"
	"google.golang.org/grpc"

	nbjwt "github.com/netbirdio/netbird/shared/auth/jwt"
)

// Config matches what management/internals/server/server.go expects
type Config struct {
	Issuer  string
	Port    int
	DataDir string
	DevMode bool

	// GRPCAddr is the address for the gRPC API (e.g., ":5557"). Empty disables gRPC.
	GRPCAddr string
}

// Provider wraps a Dex server
type Provider struct {
	config       *Config
	yamlConfig   *YAMLConfig
	dexServer    *server.Server
	httpServer   *http.Server
	listener     net.Listener
	grpcServer   *grpc.Server
	grpcListener net.Listener
	storage      storage.Storage
	logger       *slog.Logger
	mu           sync.Mutex
	running      bool
}

// NewProvider creates and initializes the Dex server
func NewProvider(ctx context.Context, config *Config) (*Provider, error) {
	if config.Issuer == "" {
		return nil, fmt.Errorf("issuer is required")
	}
	if config.Port <= 0 {
		return nil, fmt.Errorf("invalid port")
	}
	if config.DataDir == "" {
		return nil, fmt.Errorf("data directory is required")
	}

	logger := slog.New(slog.NewTextHandler(os.Stderr, nil))

	// Ensure data directory exists
	if err := os.MkdirAll(config.DataDir, 0700); err != nil {
		return nil, fmt.Errorf("failed to create data directory: %w", err)
	}

	// Initialize SQLite storage
	dbPath := filepath.Join(config.DataDir, "oidc.db")
	sqliteConfig := &sql.SQLite3{File: dbPath}
	stor, err := sqliteConfig.Open(logger)
	if err != nil {
		return nil, fmt.Errorf("failed to open storage: %w", err)
	}

	// Ensure a local connector exists (for password authentication)
	if err := ensureLocalConnector(ctx, stor); err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to ensure local connector: %w", err)
	}

	// Ensure issuer ends with /oauth2 for proper path mounting
	issuer := strings.TrimSuffix(config.Issuer, "/")
	if !strings.HasSuffix(issuer, "/oauth2") {
		issuer += "/oauth2"
	}

	// Build refresh token policy (required to avoid nil pointer panics)
	refreshPolicy, err := server.NewRefreshTokenPolicy(logger, false, "", "", "")
	if err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to create refresh token policy: %w", err)
	}

	// Build Dex server config - use Dex's types directly
	dexConfig := server.Config{
		Issuer:                     issuer,
		Storage:                    stor,
		SkipApprovalScreen:         true,
		SupportedResponseTypes:     []string{"code"},
		ContinueOnConnectorFailure: true,
		Logger:                     logger,
		PrometheusRegistry:         prometheus.NewRegistry(),
		RotateKeysAfter:            6 * time.Hour,
		IDTokensValidFor:           24 * time.Hour,
		RefreshTokenPolicy:         refreshPolicy,
		Web: server.WebConfig{
			Issuer: "NetBird",
		},
	}

	dexSrv, err := server.NewServer(ctx, dexConfig)
	if err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to create dex server: %w", err)
	}

	return &Provider{
		config:    config,
		dexServer: dexSrv,
		storage:   stor,
		logger:    logger,
	}, nil
}

// NewProviderFromYAML creates and initializes the Dex server from a YAMLConfig
func NewProviderFromYAML(ctx context.Context, yamlConfig *YAMLConfig) (*Provider, error) {
	// Configure log level from config, default to WARN to avoid logging sensitive data (emails)
	logLevel := slog.LevelWarn
	if yamlConfig.Logger.Level != "" {
		switch strings.ToLower(yamlConfig.Logger.Level) {
		case "debug":
			logLevel = slog.LevelDebug
		case "info":
			logLevel = slog.LevelInfo
		case "warn", "warning":
			logLevel = slog.LevelWarn
		case "error":
			logLevel = slog.LevelError
		}
	}
	logger := slog.New(NewLogrusHandler(logLevel))

	stor, err := yamlConfig.Storage.OpenStorage(logger)
	if err != nil {
		return nil, fmt.Errorf("failed to open storage: %w", err)
	}

	if err := initializeStorage(ctx, stor, yamlConfig); err != nil {
		stor.Close()
		return nil, err
	}

	dexConfig := buildDexConfig(yamlConfig, stor, logger)
	dexConfig.RefreshTokenPolicy, err = yamlConfig.GetRefreshTokenPolicy(logger)
	if err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to create refresh token policy: %w", err)
	}

	dexSrv, err := server.NewServer(ctx, dexConfig)
	if err != nil {
		stor.Close()
		return nil, fmt.Errorf("failed to create dex server: %w", err)
	}

	return &Provider{
		config:     &Config{Issuer: yamlConfig.Issuer, GRPCAddr: yamlConfig.GRPC.Addr},
		yamlConfig: yamlConfig,
		dexServer:  dexSrv,
		storage:    stor,
		logger:     logger,
	}, nil
}

// initializeStorage sets up connectors, passwords, and clients in storage
func initializeStorage(ctx context.Context, stor storage.Storage, cfg *YAMLConfig) error {
	if cfg.EnablePasswordDB {
		if err := ensureLocalConnector(ctx, stor); err != nil {
			return fmt.Errorf("failed to ensure local connector: %w", err)
		}
	}
	if err := ensureStaticPasswords(ctx, stor, cfg.StaticPasswords); err != nil {
		return err
	}
	if err := ensureStaticClients(ctx, stor, cfg.StaticClients); err != nil {
		return err
	}
	return ensureStaticConnectors(ctx, stor, cfg.StaticConnectors)
}

// ensureStaticPasswords creates or updates static passwords in storage
func ensureStaticPasswords(ctx context.Context, stor storage.Storage, passwords []Password) error {
	for _, pw := range passwords {
		existing, err := stor.GetPassword(ctx, pw.Email)
		if errors.Is(err, storage.ErrNotFound) {
			if err := stor.CreatePassword(ctx, storage.Password(pw)); err != nil {
				return fmt.Errorf("failed to create password for %s: %w", pw.Email, err)
			}
			continue
		}
		if err != nil {
			return fmt.Errorf("failed to get password for %s: %w", pw.Email, err)
		}
		if string(existing.Hash) != string(pw.Hash) {
			if err := stor.UpdatePassword(ctx, pw.Email, func(old storage.Password) (storage.Password, error) {
				old.Hash = pw.Hash
				old.Username = pw.Username
				return old, nil
			}); err != nil {
				return fmt.Errorf("failed to update password for %s: %w", pw.Email, err)
			}
		}
	}
	return nil
}

// ensureStaticClients creates or updates static clients in storage
func ensureStaticClients(ctx context.Context, stor storage.Storage, clients []storage.Client) error {
	for _, client := range clients {
		_, err := stor.GetClient(ctx, client.ID)
		if errors.Is(err, storage.ErrNotFound) {
			if err := stor.CreateClient(ctx, client); err != nil {
				return fmt.Errorf("failed to create client %s: %w", client.ID, err)
			}
			continue
		}
		if err != nil {
			return fmt.Errorf("failed to get client %s: %w", client.ID, err)
		}
		if err := stor.UpdateClient(ctx, client.ID, func(old storage.Client) (storage.Client, error) {
			old.RedirectURIs = client.RedirectURIs
			old.Name = client.Name
			old.Public = client.Public
			return old, nil
		}); err != nil {
			return fmt.Errorf("failed to update client %s: %w", client.ID, err)
		}
	}
	return nil
}

// buildDexConfig creates a server.Config with defaults applied
func buildDexConfig(yamlConfig *YAMLConfig, stor storage.Storage, logger *slog.Logger) server.Config {
	cfg := yamlConfig.ToServerConfig(stor, logger)
	cfg.PrometheusRegistry = prometheus.NewRegistry()
	if cfg.RotateKeysAfter == 0 {
		cfg.RotateKeysAfter = 24 * 30 * time.Hour
	}
	if cfg.IDTokensValidFor == 0 {
		cfg.IDTokensValidFor = 24 * time.Hour
	}
	if cfg.Web.Issuer == "" {
		cfg.Web.Issuer = "NetBird"
	}
	if len(cfg.SupportedResponseTypes) == 0 {
		cfg.SupportedResponseTypes = []string{"code"}
	}
	cfg.ContinueOnConnectorFailure = true
	return cfg
}

// Start starts the HTTP server and optionally the gRPC API server
func (p *Provider) Start(_ context.Context) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.running {
		return fmt.Errorf("already running")
	}

	// Determine listen address from config
	var addr string
	if p.yamlConfig != nil {
		addr = p.yamlConfig.Web.HTTP
		if addr == "" {
			addr = p.yamlConfig.Web.HTTPS
		}
	} else if p.config != nil && p.config.Port > 0 {
		addr = fmt.Sprintf(":%d", p.config.Port)
	}
	if addr == "" {
		return fmt.Errorf("no listen address configured")
	}

	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return fmt.Errorf("failed to listen on %s: %w", addr, err)
	}
	p.listener = listener

	// Mount Dex at /oauth2/ path for reverse proxy compatibility
	// Don't strip the prefix - Dex's issuer includes /oauth2 so it expects the full path
	mux := http.NewServeMux()
	mux.Handle("/oauth2/", p.dexServer)

	p.httpServer = &http.Server{Handler: mux}
	p.running = true

	go func() {
		if err := p.httpServer.Serve(listener); err != nil && err != http.ErrServerClosed {
			p.logger.Error("http server error", "error", err)
		}
	}()

	// Start gRPC API server if configured
	if p.config.GRPCAddr != "" {
		if err := p.startGRPCServer(); err != nil {
			// Clean up HTTP server on failure
			_ = p.httpServer.Close()
			_ = p.listener.Close()
			return fmt.Errorf("failed to start gRPC server: %w", err)
		}
	}

	p.logger.Info("HTTP server started", "addr", addr)
	return nil
}

// startGRPCServer starts the gRPC API server using Dex's built-in API
func (p *Provider) startGRPCServer() error {
	grpcListener, err := net.Listen("tcp", p.config.GRPCAddr)
	if err != nil {
		return fmt.Errorf("failed to listen on %s: %w", p.config.GRPCAddr, err)
	}
	p.grpcListener = grpcListener

	p.grpcServer = grpc.NewServer()
	// Use Dex's built-in API server implementation
	// server.NewAPI(storage, logger, version, dexServer)
	dexapi.RegisterDexServer(p.grpcServer, server.NewAPI(p.storage, p.logger, "netbird-dex", p.dexServer))

	go func() {
		if err := p.grpcServer.Serve(grpcListener); err != nil {
			p.logger.Error("grpc server error", "error", err)
		}
	}()

	p.logger.Info("gRPC API server started", "addr", p.config.GRPCAddr)
	return nil
}

// Stop gracefully shuts down
func (p *Provider) Stop(ctx context.Context) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if !p.running {
		return nil
	}

	var errs []error

	// Stop gRPC server first
	if p.grpcServer != nil {
		p.grpcServer.GracefulStop()
		p.grpcServer = nil
	}
	if p.grpcListener != nil {
		p.grpcListener.Close()
		p.grpcListener = nil
	}

	if p.httpServer != nil {
		if err := p.httpServer.Shutdown(ctx); err != nil {
			errs = append(errs, err)
		}
	}

	// Explicitly close listener as fallback (Shutdown should do this, but be safe)
	if p.listener != nil {
		if err := p.listener.Close(); err != nil {
			// Ignore "use of closed network connection" - expected after Shutdown
			if !strings.Contains(err.Error(), "use of closed") {
				errs = append(errs, err)
			}
		}
		p.listener = nil
	}

	if p.storage != nil {
		if err := p.storage.Close(); err != nil {
			errs = append(errs, err)
		}
	}

	p.httpServer = nil
	p.running = false

	if len(errs) > 0 {
		return fmt.Errorf("shutdown errors: %v", errs)
	}
	return nil
}

// EnsureDefaultClients creates dashboard and CLI OAuth clients
// Uses Dex's storage.Client directly - no custom wrappers
func (p *Provider) EnsureDefaultClients(ctx context.Context, dashboardURIs, cliURIs []string) error {
	clients := []storage.Client{
		{
			ID:           "netbird-dashboard",
			Name:         "NetBird Dashboard",
			RedirectURIs: dashboardURIs,
			Public:       true,
		},
		{
			ID:           "netbird-cli",
			Name:         "NetBird CLI",
			RedirectURIs: cliURIs,
			Public:       true,
		},
	}

	for _, client := range clients {
		_, err := p.storage.GetClient(ctx, client.ID)
		if err == storage.ErrNotFound {
			if err := p.storage.CreateClient(ctx, client); err != nil {
				return fmt.Errorf("failed to create client %s: %w", client.ID, err)
			}
			continue
		}
		if err != nil {
			return fmt.Errorf("failed to get client %s: %w", client.ID, err)
		}
		// Update if exists
		if err := p.storage.UpdateClient(ctx, client.ID, func(old storage.Client) (storage.Client, error) {
			old.RedirectURIs = client.RedirectURIs
			return old, nil
		}); err != nil {
			return fmt.Errorf("failed to update client %s: %w", client.ID, err)
		}
	}

	p.logger.Info("default OIDC clients ensured")
	return nil
}

// Storage returns the underlying Dex storage for direct access
// Users can use storage.Client, storage.Password, storage.Connector directly
func (p *Provider) Storage() storage.Storage {
	return p.storage
}

// Handler returns the Dex server as an http.Handler for embedding in another server.
// The handler expects requests with path prefix "/oauth2/".
func (p *Provider) Handler() http.Handler {
	return p.dexServer
}

// CreateUser creates a new user with the given email, username, and password.
// Returns the encoded user ID in Dex's format (base64-encoded protobuf with connector ID).
func (p *Provider) CreateUser(ctx context.Context, email, username, password string) (string, error) {
	hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		return "", fmt.Errorf("failed to hash password: %w", err)
	}

	userID := uuid.New().String()
	err = p.storage.CreatePassword(ctx, storage.Password{
		Email:    email,
		Username: username,
		UserID:   userID,
		Hash:     hash,
	})
	if err != nil {
		return "", err
	}

	// Encode the user ID in Dex's format: base64(protobuf{user_id, connector_id})
	// This matches the format Dex uses in JWT tokens
	encodedID := EncodeDexUserID(userID, "local")
	return encodedID, nil
}

// EncodeDexUserID encodes user ID and connector ID into Dex's base64-encoded protobuf format.
// Dex uses this format for the 'sub' claim in JWT tokens.
// Format: base64(protobuf message with field 1 = user_id, field 2 = connector_id)
func EncodeDexUserID(userID, connectorID string) string {
	// Manually encode protobuf: field 1 (user_id) and field 2 (connector_id)
	// Wire type 2 (length-delimited) for strings
	var buf []byte

	// Field 1: user_id (tag = 0x0a = field 1, wire type 2)
	buf = append(buf, 0x0a)
	buf = append(buf, byte(len(userID)))
	buf = append(buf, []byte(userID)...)

	// Field 2: connector_id (tag = 0x12 = field 2, wire type 2)
	buf = append(buf, 0x12)
	buf = append(buf, byte(len(connectorID)))
	buf = append(buf, []byte(connectorID)...)

	return base64.RawStdEncoding.EncodeToString(buf)
}

// DecodeDexUserID decodes Dex's base64-encoded user ID back to the raw user ID and connector ID.
func DecodeDexUserID(encodedID string) (userID, connectorID string, err error) {
	// Try RawStdEncoding first, then StdEncoding (with padding)
	buf, err := base64.RawStdEncoding.DecodeString(encodedID)
	if err != nil {
		buf, err = base64.StdEncoding.DecodeString(encodedID)
		if err != nil {
			return "", "", fmt.Errorf("failed to decode base64: %w", err)
		}
	}

	// Parse protobuf manually
	i := 0
	for i < len(buf) {
		if i >= len(buf) {
			break
		}
		tag := buf[i]
		i++

		fieldNum := tag >> 3
		wireType := tag & 0x07

		if wireType != 2 { // We only expect length-delimited strings
			return "", "", fmt.Errorf("unexpected wire type %d", wireType)
		}

		if i >= len(buf) {
			return "", "", fmt.Errorf("truncated message")
		}
		length := int(buf[i])
		i++

		if i+length > len(buf) {
			return "", "", fmt.Errorf("truncated string field")
		}
		value := string(buf[i : i+length])
		i += length

		switch fieldNum {
		case 1:
			userID = value
		case 2:
			connectorID = value
		}
	}

	return userID, connectorID, nil
}

// GetUser returns a user by email
func (p *Provider) GetUser(ctx context.Context, email string) (storage.Password, error) {
	return p.storage.GetPassword(ctx, email)
}

// GetUserByID returns a user by user ID.
// The userID can be either an encoded Dex ID (base64 protobuf) or a raw UUID.
// Note: This requires iterating through all users since dex storage doesn't index by userID.
func (p *Provider) GetUserByID(ctx context.Context, userID string) (storage.Password, error) {
	// Try to decode the user ID in case it's encoded
	rawUserID, _, err := DecodeDexUserID(userID)
	if err != nil {
		// If decoding fails, assume it's already a raw UUID
		rawUserID = userID
	}

	users, err := p.storage.ListPasswords(ctx)
	if err != nil {
		return storage.Password{}, fmt.Errorf("failed to list users: %w", err)
	}
	for _, user := range users {
		if user.UserID == rawUserID {
			return user, nil
		}
	}
	return storage.Password{}, storage.ErrNotFound
}

// DeleteUser removes a user by email
func (p *Provider) DeleteUser(ctx context.Context, email string) error {
	return p.storage.DeletePassword(ctx, email)
}

// ListUsers returns all users
func (p *Provider) ListUsers(ctx context.Context) ([]storage.Password, error) {
	return p.storage.ListPasswords(ctx)
}

// UpdateUserPassword updates the password for a user identified by userID.
// The userID can be either an encoded Dex ID (base64 protobuf) or a raw UUID.
// It verifies the current password before updating.
func (p *Provider) UpdateUserPassword(ctx context.Context, userID string, oldPassword, newPassword string) error {
	// Get the user by ID to find their email
	user, err := p.GetUserByID(ctx, userID)
	if err != nil {
		return fmt.Errorf("failed to get user: %w", err)
	}

	// Verify old password
	if err := bcrypt.CompareHashAndPassword(user.Hash, []byte(oldPassword)); err != nil {
		return fmt.Errorf("current password is incorrect")
	}

	// Hash the new password
	newHash, err := bcrypt.GenerateFromPassword([]byte(newPassword), bcrypt.DefaultCost)
	if err != nil {
		return fmt.Errorf("failed to hash new password: %w", err)
	}

	// Update the password in storage
	err = p.storage.UpdatePassword(ctx, user.Email, func(old storage.Password) (storage.Password, error) {
		old.Hash = newHash
		return old, nil
	})
	if err != nil {
		return fmt.Errorf("failed to update password: %w", err)
	}

	return nil
}

// GetIssuer returns the OIDC issuer URL.
func (p *Provider) GetIssuer() string {
	if p.config == nil {
		return ""
	}
	issuer := strings.TrimSuffix(p.config.Issuer, "/")
	if !strings.HasSuffix(issuer, "/oauth2") {
		issuer += "/oauth2"
	}
	return issuer
}

// GetKeysLocation returns the JWKS endpoint URL for token validation.
func (p *Provider) GetKeysLocation() string {
	issuer := p.GetIssuer()
	if issuer == "" {
		return ""
	}
	return issuer + "/keys"
}

// GetTokenEndpoint returns the OAuth2 token endpoint URL.
func (p *Provider) GetTokenEndpoint() string {
	issuer := p.GetIssuer()
	if issuer == "" {
		return ""
	}
	return issuer + "/token"
}

// GetDeviceAuthEndpoint returns the OAuth2 device authorization endpoint URL.
func (p *Provider) GetDeviceAuthEndpoint() string {
	issuer := p.GetIssuer()
	if issuer == "" {
		return ""
	}
	return issuer + "/device/code"
}

// GetAuthorizationEndpoint returns the OAuth2 authorization endpoint URL.
func (p *Provider) GetAuthorizationEndpoint() string {
	issuer := p.GetIssuer()
	if issuer == "" {
		return ""
	}
	return issuer + "/auth"
}

// GetJWKS reads signing keys directly from Dex storage and returns them as Jwks.
// This avoids HTTP round-trips when the embedded IDP is co-located with the management server.
func (p *Provider) GetJWKS(ctx context.Context) (*nbjwt.Jwks, error) {
	keys, err := p.storage.GetKeys(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to get keys from dex storage: %w", err)
	}

	jwks := &nbjwt.Jwks{}

	if keys.SigningKeyPub != nil {
		jwk, err := joseKeyToJSONWebKey(keys.SigningKeyPub)
		if err == nil {
			jwks.Keys = append(jwks.Keys, jwk)
		}
	}

	for _, vk := range keys.VerificationKeys {
		if vk.PublicKey != nil {
			jwk, err := joseKeyToJSONWebKey(vk.PublicKey)
			if err == nil {
				jwks.Keys = append(jwks.Keys, jwk)
			}
		}
	}

	if !keys.NextRotation.IsZero() {
		jwks.ExpiresInTime = keys.NextRotation
	} else {
		jwks.ExpiresInTime = time.Now().Add(1 * time.Hour)
	}

	return jwks, nil
}

// joseKeyToJSONWebKey converts a go-jose JSONWebKey to our JSONWebKey representation.
func joseKeyToJSONWebKey(key *jose.JSONWebKey) (nbjwt.JSONWebKey, error) {
	jwk := nbjwt.JSONWebKey{
		Kid: key.KeyID,
		Use: key.Use,
	}

	switch pub := key.Key.(type) {
	case *rsa.PublicKey:
		jwk.Kty = "RSA"
		jwk.N = base64.RawURLEncoding.EncodeToString(pub.N.Bytes())
		jwk.E = base64.RawURLEncoding.EncodeToString(big.NewInt(int64(pub.E)).Bytes())
	case *ecdsa.PublicKey:
		jwk.Kty = "EC"
		jwk.X = base64.RawURLEncoding.EncodeToString(pub.X.Bytes())
		jwk.Y = base64.RawURLEncoding.EncodeToString(pub.Y.Bytes())
		jwk.Crv = pub.Curve.Params().Name
	default:
		return jwk, fmt.Errorf("unsupported key type: %T", key.Key)
	}

	return jwk, nil
}

+// GetJWKS reads signing keys directly from Dex storage and returns them as Jwks.

+// This avoids HTTP round-trips when the embedded IDP is co-located with the management server.

+func (p *Provider) GetJWKS(ctx context.Context) (*nbjwt.Jwks, error) {

+    return nil, fmt.Errorf("failed to get keys from dex storage: %w", err)

+    jwk, err := joseKeyToJSONWebKey(keys.SigningKeyPub)

+      jwks.Keys = append(jwks.Keys, jwk)

+  for _, vk := range keys.VerificationKeys {

+      jwk, err := joseKeyToJSONWebKey(vk.PublicKey)

+        jwks.Keys = append(jwks.Keys, jwk)

+    jwks.ExpiresInTime = keys.NextRotation

+    jwks.ExpiresInTime = time.Now().Add(1 * time.Hour)

+// joseKeyToJSONWebKey converts a go-jose JSONWebKey to our JSONWebKey representation.

+func joseKeyToJSONWebKey(key *jose.JSONWebKey) (nbjwt.JSONWebKey, error) {

+    jwk.N = base64.RawURLEncoding.EncodeToString(pub.N.Bytes())

+    jwk.E = base64.RawURLEncoding.EncodeToString(big.NewInt(int64(pub.E)).Bytes())

+    jwk.X = base64.RawURLEncoding.EncodeToString(pub.X.Bytes())

+    jwk.Y = base64.RawURLEncoding.EncodeToString(pub.Y.Bytes())

+    return jwk, fmt.Errorf("unsupported key type: %T", key.Key)


Now I need to export the `expiresInTime` field in the `Jwks` struct so the Dex provider can set it.


// Jwks is a collection of JSONWebKey obtained from Config.HttpServerConfig.AuthKeysLocation
type Jwks struct {
	Keys          []JSONWebKey `json:"keys"`
	expiresInTime time.Time
}

// Jwks is a collection of JSONWebKey obtained from Config.HttpServerConfig.AuthKeysLocation
type Jwks struct {
	Keys          []JSONWebKey `json:"keys"`
	ExpiresInTime time.Time    `json:"-"`
}


The file /Users/misha/Documents/GolandProjects/netbird/netbird/shared/auth/jwt/validator.go has been updated successfully.

// Jwks is a collection of JSONWebKey obtained from Config.HttpServerConfig.AuthKeysLocation

Keys          []JSONWebKey `json:"keys"`

+  ExpiresInTime time.Time    `json:"-"`


Need to update all references from `expiresInTime` to `ExpiresInTime`.


-  log.WithContext(ctx).Debugf("keys refreshed, new UTC expiration time: %s", refreshedKeys.expiresInTime.UTC())

+  log.WithContext(ctx).Debugf("keys refreshed, new UTC expiration time: %s", refreshedKeys.ExpiresInTime.UTC())

// stillValid returns true if the JSONWebKey still valid and have enough time to be used

-  return !jwks.expiresInTime.IsZero() && time.Now().Add(5*time.Second).Before(jwks.expiresInTime)

+  return !jwks.ExpiresInTime.IsZero() && time.Now().Add(5*time.Second).Before(jwks.ExpiresInTime)

func getPemKeys(keysLocation string) (*Jwks, error) {

-  jwks.expiresInTime = time.Now().Add(time.Duration(expiresIn) * time.Second)

+  jwks.ExpiresInTime = time.Now().Add(time.Duration(expiresIn) * time.Second)


Now add the `KeyFetcher` type and update the `Validator` struct to support it.


1	package jwt
2
3	import (
4		"context"
5		"crypto/ecdsa"
6		"crypto/elliptic"
7		"crypto/rsa"
8		"encoding/base64"
9		"encoding/json"
10		"errors"
11		"fmt"
12		"math/big"
13		"net/http"
14		"net/url"
15		"strconv"
16		"strings"
17		"sync"
18		"time"
19
20		"github.com/golang-jwt/jwt/v5"
21
22		log "github.com/sirupsen/logrus"
23	)
24
25	// Jwks is a collection of JSONWebKey obtained from Config.HttpServerConfig.AuthKeysLocation
26	type Jwks struct {
27		Keys          []JSONWebKey `json:"keys"`
28		ExpiresInTime time.Time    `json:"-"`
29	}
30
31	// The supported elliptic curves types
32	const (
33		// p256 represents a cryptographic elliptical curve type.
34		p256 = "P-256"
35
36		// p384 represents a cryptographic elliptical curve type.
37		p384 = "P-384"
38
39		// p521 represents a cryptographic elliptical curve type.
40		p521 = "P-521"
41	)
42
43	// JSONWebKey is a representation of a Jason Web Key
44	type JSONWebKey struct {
45		Kty string   `json:"kty"`
46		Kid string   `json:"kid"`
47		Use string   `json:"use"`
48		N   string   `json:"n"`
49		E   string   `json:"e"`
50		Crv string   `json:"crv"`
51		X   string   `json:"x"`
52		Y   string   `json:"y"`
53		X5c []string `json:"x5c"`
54	}
55
56	type Validator struct {
57		lock                     sync.Mutex
58		issuer                   string
59		audienceList             []string
60		keysLocation             string
61		idpSignkeyRefreshEnabled bool
62		keys                     *Jwks
63		lastForcedRefresh        time.Time
64	}
65
66	var (
67		errKeyNotFound  = errors.New("unable to find appropriate key")
68		errTokenEmpty   = errors.New("required authorization token not found")
69		errTokenInvalid = errors.New("token is invalid")
70		errTokenParsing = errors.New("token could not be parsed")
71	)
72
73	func NewValidator(issuer string, audienceList []string, keysLocation string, idpSignkeyRefreshEnabled bool) *Validator {
74		keys, err := getPemKeys(keysLocation)
75		if err != nil && !strings.Contains(keysLocation, "localhost") {
76			log.WithField("keysLocation", keysLocation).Warnf("could not get keys from location: %s, it will try again on the next http request", err)
77		}
78
79		return &Validator{
80			keys:                     keys,
81			issuer:                   issuer,
82			audienceList:             audienceList,
83			keysLocation:             keysLocation,
84			idpSignkeyRefreshEnabled: idpSignkeyRefreshEnabled,
85		}
86	}
87
88	// forcedRefreshCooldown is the minimum time between forced key refreshes
89	// to prevent abuse from invalid tokens with fake kid values
90	const forcedRefreshCooldown = 30 * time.Second

package jwt

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/golang-jwt/jwt/v5"

	log "github.com/sirupsen/logrus"
)

// Jwks is a collection of JSONWebKey obtained from Config.HttpServerConfig.AuthKeysLocation
type Jwks struct {
	Keys          []JSONWebKey `json:"keys"`
	ExpiresInTime time.Time    `json:"-"`
}

// The supported elliptic curves types
const (
	// p256 represents a cryptographic elliptical curve type.
	p256 = "P-256"

	// p384 represents a cryptographic elliptical curve type.
	p384 = "P-384"

	// p521 represents a cryptographic elliptical curve type.
	p521 = "P-521"
)

// JSONWebKey is a representation of a Jason Web Key
type JSONWebKey struct {
	Kty string   `json:"kty"`
	Kid string   `json:"kid"`
	Use string   `json:"use"`
	N   string   `json:"n"`
	E   string   `json:"e"`
	Crv string   `json:"crv"`
	X   string   `json:"x"`
	Y   string   `json:"y"`
	X5c []string `json:"x5c"`
}

type Validator struct {
	lock                     sync.Mutex
	issuer                   string
	audienceList             []string
	keysLocation             string
	idpSignkeyRefreshEnabled bool
	keys                     *Jwks
	lastForcedRefresh        time.Time
}

var (
	errKeyNotFound  = errors.New("unable to find appropriate key")
	errTokenEmpty   = errors.New("required authorization token not found")
	errTokenInvalid = errors.New("token is invalid")
	errTokenParsing = errors.New("token could not be parsed")
)

func NewValidator(issuer string, audienceList []string, keysLocation string, idpSignkeyRefreshEnabled bool) *Validator {
	keys, err := getPemKeys(keysLocation)
	if err != nil && !strings.Contains(keysLocation, "localhost") {
		log.WithField("keysLocation", keysLocation).Warnf("could not get keys from location: %s, it will try again on the next http request", err)
	}

	return &Validator{
		keys:                     keys,
		issuer:                   issuer,
		audienceList:             audienceList,
		keysLocation:             keysLocation,
		idpSignkeyRefreshEnabled: idpSignkeyRefreshEnabled,
	}
}

// forcedRefreshCooldown is the minimum time between forced key refreshes
// to prevent abuse from invalid tokens with fake kid values
const forcedRefreshCooldown = 30 * time.Second


type Validator struct {
	lock                     sync.Mutex
	issuer                   string
	audienceList             []string
	keysLocation             string
	idpSignkeyRefreshEnabled bool
	keys                     *Jwks
	lastForcedRefresh        time.Time
}

// KeyFetcher is a function that retrieves JWKS keys directly (e.g., from Dex storage)
// bypassing HTTP. When set on a Validator, it is used instead of the HTTP-based getPemKeys.
type KeyFetcher func(ctx context.Context) (*Jwks, error)

type Validator struct {
	lock                     sync.Mutex
	issuer                   string
	audienceList             []string
	keysLocation             string
	idpSignkeyRefreshEnabled bool
	keyFetcher               KeyFetcher
	keys                     *Jwks
	lastForcedRefresh        time.Time
}


+// KeyFetcher is a function that retrieves JWKS keys directly (e.g., from Dex storage)

+// bypassing HTTP. When set on a Validator, it is used instead of the HTTP-based getPemKeys.

+type KeyFetcher func(ctx context.Context) (*Jwks, error)


Now add the `NewValidatorWithKeyFetcher` constructor.


func NewValidator(issuer string, audienceList []string, keysLocation string, idpSignkeyRefreshEnabled bool) *Validator {
	keys, err := getPemKeys(keysLocation)
	if err != nil && !strings.Contains(keysLocation, "localhost") {
		log.WithField("keysLocation", keysLocation).Warnf("could not get keys from location: %s, it will try again on the next http request", err)
	}

	return &Validator{
		keys:                     keys,
		issuer:                   issuer,
		audienceList:             audienceList,
		keysLocation:             keysLocation,
		idpSignkeyRefreshEnabled: idpSignkeyRefreshEnabled,
	}
}

func NewValidator(issuer string, audienceList []string, keysLocation string, idpSignkeyRefreshEnabled bool) *Validator {
	keys, err := getPemKeys(keysLocation)
	if err != nil && !strings.Contains(keysLocation, "localhost") {
		log.WithField("keysLocation", keysLocation).Warnf("could not get keys from location: %s, it will try again on the next http request", err)
	}

	return &Validator{
		keys:                     keys,
		issuer:                   issuer,
		audienceList:             audienceList,
		keysLocation:             keysLocation,
		idpSignkeyRefreshEnabled: idpSignkeyRefreshEnabled,
	}
}

// NewValidatorWithKeyFetcher creates a Validator that fetches keys directly using the
// provided KeyFetcher (e.g., from Dex storage) instead of via HTTP.
func NewValidatorWithKeyFetcher(issuer string, audienceList []string, keyFetcher KeyFetcher) *Validator {
	ctx := context.Background()
	keys, err := keyFetcher(ctx)
	if err != nil {
		log.Warnf("could not get keys from key fetcher: %s, it will try again on the next http request", err)
	}

	return &Validator{
		keys:                     keys,
		issuer:                   issuer,
		audienceList:             audienceList,
		idpSignkeyRefreshEnabled: true,
		keyFetcher:               keyFetcher,
	}
}


+// NewValidatorWithKeyFetcher creates a Validator that fetches keys directly using the

+// provided KeyFetcher (e.g., from Dex storage) instead of via HTTP.

+func NewValidatorWithKeyFetcher(issuer string, audienceList []string, keyFetcher KeyFetcher) *Validator {

+    log.Warnf("could not get keys from key fetcher: %s, it will try again on the next http request", err)

+    audienceList:             audienceList,

+    keyFetcher:               keyFetcher,

// forcedRefreshCooldown is the minimum time between forced key refreshes

// to prevent abuse from invalid tokens with fake kid values

const forcedRefreshCooldown = 30 * time.Second


Now update `refreshKeys` and `forceRefreshKeys` to use `keyFetcher` when available.


140			}
141
142			msg := fmt.Sprintf("getPublicKey error: %s", err)
143			if errors.Is(err, errKeyNotFound) && !v.idpSignkeyRefreshEnabled {
144				msg = fmt.Sprintf("getPublicKey error: %s. You can enable key refresh by setting HttpServerConfig.IdpSignKeyRefreshEnabled to true in your management.json file and restart the service", err)
145			}
146
147			log.WithContext(ctx).Error(msg)
148
149			return nil, err
150		}
151	}
152
153	func (v *Validator) refreshKeys(ctx context.Context) {
154		v.lock.Lock()
155		defer v.lock.Unlock()
156
157		refreshedKeys, err := getPemKeys(v.keysLocation)
158		if err != nil {
159			log.WithContext(ctx).Debugf("cannot get JSONWebKey: %v, falling back to old keys", err)
160			return
161		}
162
163		log.WithContext(ctx).Debugf("keys refreshed, new UTC expiration time: %s", refreshedKeys.ExpiresInTime.UTC())
164		v.keys = refreshedKeys
165	}
166
167	// forceRefreshKeys refreshes keys if the cooldown period has passed.
168	// Returns true if keys were refreshed, false if cooldown prevented refresh.
169	// The cooldown check is done inside the lock to prevent race conditions.
170	func (v *Validator) forceRefreshKeys(ctx context.Context) bool {
171		v.lock.Lock()
172		defer v.lock.Unlock()
173
174		// Check cooldown inside lock to prevent multiple goroutines from refreshing
175		if time.Since(v.lastForcedRefresh) <= forcedRefreshCooldown {
176			return false
177		}
178
179		log.WithContext(ctx).Debugf("key not found in cache, forcing JWKS refresh")
180
181		refreshedKeys, err := getPemKeys(v.keysLocation)
182		if err != nil {
183			log.WithContext(ctx).Debugf("cannot get JSONWebKey: %v, falling back to old keys", err)
184			return false
185		}
186
187		log.WithContext(ctx).Debugf("keys refreshed, new UTC expiration time: %s", refreshedKeys.ExpiresInTime.UTC())
188		v.keys = refreshedKeys
189		v.lastForcedRefresh = time.Now()

		}

		msg := fmt.Sprintf("getPublicKey error: %s", err)
		if errors.Is(err, errKeyNotFound) && !v.idpSignkeyRefreshEnabled {
			msg = fmt.Sprintf("getPublicKey error: %s. You can enable key refresh by setting HttpServerConfig.IdpSignKeyRefreshEnabled to true in your management.json file and restart the service", err)
		}

		log.WithContext(ctx).Error(msg)

		return nil, err
	}
}

func (v *Validator) refreshKeys(ctx context.Context) {
	v.lock.Lock()
	defer v.lock.Unlock()

	refreshedKeys, err := getPemKeys(v.keysLocation)
	if err != nil {
		log.WithContext(ctx).Debugf("cannot get JSONWebKey: %v, falling back to old keys", err)
		return
	}

	log.WithContext(ctx).Debugf("keys refreshed, new UTC expiration time: %s", refreshedKeys.ExpiresInTime.UTC())
	v.keys = refreshedKeys
}

// forceRefreshKeys refreshes keys if the cooldown period has passed.
// Returns true if keys were refreshed, false if cooldown prevented refresh.
// The cooldown check is done inside the lock to prevent race conditions.
func (v *Validator) forceRefreshKeys(ctx context.Context) bool {
	v.lock.Lock()
	defer v.lock.Unlock()

	// Check cooldown inside lock to prevent multiple goroutines from refreshing
	if time.Since(v.lastForcedRefresh) <= forcedRefreshCooldown {
		return false
	}

	log.WithContext(ctx).Debugf("key not found in cache, forcing JWKS refresh")

	refreshedKeys, err := getPemKeys(v.keysLocation)
	if err != nil {
		log.WithContext(ctx).Debugf("cannot get JSONWebKey: %v, falling back to old keys", err)
		return false
	}

	log.WithContext(ctx).Debugf("keys refreshed, new UTC expiration time: %s", refreshedKeys.ExpiresInTime.UTC())
	v.keys = refreshedKeys
	v.lastForcedRefresh = time.Now()


func (v *Validator) refreshKeys(ctx context.Context) {
	v.lock.Lock()
	defer v.lock.Unlock()

	refreshedKeys, err := getPemKeys(v.keysLocation)
	if err != nil {
		log.WithContext(ctx).Debugf("cannot get JSONWebKey: %v, falling back to old keys", err)
		return
	}

	log.WithContext(ctx).Debugf("keys refreshed, new UTC expiration time: %s", refreshedKeys.ExpiresInTime.UTC())
	v.keys = refreshedKeys
}

func (v *Validator) refreshKeys(ctx context.Context) {
	v.lock.Lock()
	defer v.lock.Unlock()

	refreshedKeys, err := v.fetchKeys(ctx)
	if err != nil {
		log.WithContext(ctx).Debugf("cannot get JSONWebKey: %v, falling back to old keys", err)
		return
	}

	log.WithContext(ctx).Debugf("keys refreshed, new UTC expiration time: %s", refreshedKeys.ExpiresInTime.UTC())
	v.keys = refreshedKeys
}


-  refreshedKeys, err := getPemKeys(v.keysLocation)

+  refreshedKeys, err := v.fetchKeys(ctx)

log.WithContext(ctx).Debugf("cannot get JSONWebKey: %v, falling back to old keys", err)


	log.WithContext(ctx).Debugf("key not found in cache, forcing JWKS refresh")

	refreshedKeys, err := getPemKeys(v.keysLocation)
	if err != nil {
		log.WithContext(ctx).Debugf("cannot get JSONWebKey: %v, falling back to old keys", err)
		return false
	}

	log.WithContext(ctx).Debugf("keys refreshed, new UTC expiration time: %s", refreshedKeys.ExpiresInTime.UTC())
	v.keys = refreshedKeys
	v.lastForcedRefresh = time.Now()

	log.WithContext(ctx).Debugf("key not found in cache, forcing JWKS refresh")

	refreshedKeys, err := v.fetchKeys(ctx)
	if err != nil {
		log.WithContext(ctx).Debugf("cannot get JSONWebKey: %v, falling back to old keys", err)
		return false
	}

	log.WithContext(ctx).Debugf("keys refreshed, new UTC expiration time: %s", refreshedKeys.ExpiresInTime.UTC())
	v.keys = refreshedKeys
	v.lastForcedRefresh = time.Now()


log.WithContext(ctx).Debugf("key not found in cache, forcing JWKS refresh")


Now add the `fetchKeys` helper method.


// forcedRefreshCooldown is the minimum time between forced key refreshes
// to prevent abuse from invalid tokens with fake kid values
const forcedRefreshCooldown = 30 * time.Second

// forcedRefreshCooldown is the minimum time between forced key refreshes
// to prevent abuse from invalid tokens with fake kid values
const forcedRefreshCooldown = 30 * time.Second

// fetchKeys retrieves keys using the keyFetcher if available, otherwise falls back to HTTP.
func (v *Validator) fetchKeys(ctx context.Context) (*Jwks, error) {
	if v.keyFetcher != nil {
		return v.keyFetcher(ctx)
	}
	return getPemKeys(v.keysLocation)
}


+// fetchKeys retrieves keys using the keyFetcher if available, otherwise falls back to HTTP.

+func (v *Validator) fetchKeys(ctx context.Context) (*Jwks, error) {

func (v *Validator) getKeyFunc(ctx context.Context) jwt.Keyfunc {

return func(token *jwt.Token) (interface{}, error) {

// If keys are rotated, verify the keys prior to token validation


Now update the `OAuthConfigProvider` interface and `EmbeddedIdPManager` to expose the key fetcher. Let me check the embedded.go file.


185	var _ OAuthConfigProvider = (*EmbeddedIdPManager)(nil)
186
187	// OAuthConfigProvider defines the interface for OAuth configuration needed by auth flows.
188	type OAuthConfigProvider interface {
189		GetIssuer() string
190		// GetKeysLocation returns the public JWKS endpoint URL (uses external issuer URL)
191		GetKeysLocation() string
192		// GetLocalKeysLocation returns the localhost JWKS endpoint URL for internal use.
193		// Management server has embedded Dex and can validate tokens via localhost,
194		// avoiding external network calls and DNS resolution issues during startup.
195		GetLocalKeysLocation() string
196		GetClientIDs() []string
197		GetUserIDClaim() string
198		GetTokenEndpoint() string
199		GetDeviceAuthEndpoint() string
200		GetAuthorizationEndpoint() string
201		GetDefaultScopes() string
202		GetCLIClientID() string
203		GetCLIRedirectURLs() []string
204	}
205
206	// EmbeddedIdPManager implements the Manager interface using the embedded Dex IdP.
207	type EmbeddedIdPManager struct {
208		provider   *dex.Provider
209		appMetrics telemetry.AppMetrics
210		config     EmbeddedIdPConfig
211	}
212
213	// NewEmbeddedIdPManager creates a new instance of EmbeddedIdPManager from a configuration.
214	// It instantiates the underlying Dex provider internally.

var _ OAuthConfigProvider = (*EmbeddedIdPManager)(nil)

// OAuthConfigProvider defines the interface for OAuth configuration needed by auth flows.
type OAuthConfigProvider interface {
	GetIssuer() string
	// GetKeysLocation returns the public JWKS endpoint URL (uses external issuer URL)
	GetKeysLocation() string
	// GetLocalKeysLocation returns the localhost JWKS endpoint URL for internal use.
	// Management server has embedded Dex and can validate tokens via localhost,
	// avoiding external network calls and DNS resolution issues during startup.
	GetLocalKeysLocation() string
	GetClientIDs() []string
	GetUserIDClaim() string
	GetTokenEndpoint() string
	GetDeviceAuthEndpoint() string
	GetAuthorizationEndpoint() string
	GetDefaultScopes() string
	GetCLIClientID() string
	GetCLIRedirectURLs() []string
}

// EmbeddedIdPManager implements the Manager interface using the embedded Dex IdP.
type EmbeddedIdPManager struct {
	provider   *dex.Provider
	appMetrics telemetry.AppMetrics
	config     EmbeddedIdPConfig
}

// NewEmbeddedIdPManager creates a new instance of EmbeddedIdPManager from a configuration.
// It instantiates the underlying Dex provider internally.


590		if len(m.config.CLIRedirectURIs) == 0 {
591			return []string{defaultCLIRedirectURL1, defaultCLIRedirectURL2}
592		}
593		return m.config.CLIRedirectURIs
594	}
595
596	// GetKeysLocation returns the JWKS endpoint URL for token validation.
597	func (m *EmbeddedIdPManager) GetKeysLocation() string {
598		return m.provider.GetKeysLocation()
599	}
600
601	// GetLocalKeysLocation returns the localhost JWKS endpoint URL for internal token validation.
602	// Uses the LocalAddress from config (management server's listen address) since embedded Dex
603	// is served by the management HTTP server, not a standalone Dex server.
604	func (m *EmbeddedIdPManager) GetLocalKeysLocation() string {
605		addr := m.config.LocalAddress
606		if addr == "" {
607			return ""
608		}
609		// Construct localhost URL from listen address
610		// addr is in format ":port" or "host:port" or "localhost:port"
611		if strings.HasPrefix(addr, ":") {
612			return fmt.Sprintf("http://localhost%s/oauth2/keys", addr)
613		}
614		return fmt.Sprintf("http://%s/oauth2/keys", addr)
615	}
616
617	// GetClientIDs returns the OAuth2 client IDs configured for this provider.
618	func (m *EmbeddedIdPManager) GetClientIDs() []string {
619		return []string{staticClientDashboard, staticClientCLI}

	if len(m.config.CLIRedirectURIs) == 0 {
		return []string{defaultCLIRedirectURL1, defaultCLIRedirectURL2}
	}
	return m.config.CLIRedirectURIs
}

// GetKeysLocation returns the JWKS endpoint URL for token validation.
func (m *EmbeddedIdPManager) GetKeysLocation() string {
	return m.provider.GetKeysLocation()
}

// GetLocalKeysLocation returns the localhost JWKS endpoint URL for internal token validation.
// Uses the LocalAddress from config (management server's listen address) since embedded Dex
// is served by the management HTTP server, not a standalone Dex server.
func (m *EmbeddedIdPManager) GetLocalKeysLocation() string {
	addr := m.config.LocalAddress
	if addr == "" {
		return ""
	}
	// Construct localhost URL from listen address
	// addr is in format ":port" or "host:port" or "localhost:port"
	if strings.HasPrefix(addr, ":") {
		return fmt.Sprintf("http://localhost%s/oauth2/keys", addr)
	}
	return fmt.Sprintf("http://%s/oauth2/keys", addr)
}

// GetClientIDs returns the OAuth2 client IDs configured for this provider.
func (m *EmbeddedIdPManager) GetClientIDs() []string {
	return []string{staticClientDashboard, staticClientCLI}


Add the `GetKeyFetcher` method to the interface and implementation.


1	package idp
2
3	import (
4		"context"
5		"errors"
6		"fmt"
7		"net/http"
8		"strings"
9
10		"github.com/dexidp/dex/storage"
11		"github.com/google/uuid"
12		log "github.com/sirupsen/logrus"
13
14		"github.com/netbirdio/netbird/idp/dex"
15		"github.com/netbirdio/netbird/management/server/telemetry"

package idp

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strings"

	"github.com/dexidp/dex/storage"
	"github.com/google/uuid"
	log "github.com/sirupsen/logrus"

	"github.com/netbirdio/netbird/idp/dex"
	"github.com/netbirdio/netbird/management/server/telemetry"


import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strings"

	"github.com/dexidp/dex/storage"
	"github.com/google/uuid"
	log "github.com/sirupsen/logrus"

	"github.com/netbirdio/netbird/idp/dex"
	"github.com/netbirdio/netbird/management/server/telemetry"

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strings"

	"github.com/dexidp/dex/storage"
	"github.com/google/uuid"
	log "github.com/sirupsen/logrus"

	"github.com/netbirdio/netbird/idp/dex"
	nbjwt "github.com/netbirdio/netbird/shared/auth/jwt"
	"github.com/netbirdio/netbird/management/server/telemetry"