674	}
675
676	// GetJWKS reads signing keys directly from Dex storage and returns them as Jwks.
677	// This avoids HTTP round-trips when the embedded IDP is co-located with the management server.
678	func (p *Provider) GetJWKS(ctx context.Context) (*nbjwt.Jwks, error) {
679		keys, err := p.storage.GetKeys(ctx)
680		if err != nil {
681			return nil, fmt.Errorf("failed to get keys from dex storage: %w", err)
682		}
683
684		jwks := &nbjwt.Jwks{}
685
686		if keys.SigningKeyPub != nil {
687			jwk, err := joseKeyToJSONWebKey(keys.SigningKeyPub)
688			if err == nil {
689				jwks.Keys = append(jwks.Keys, jwk)
690			}
691		}
692
693		for _, vk := range keys.VerificationKeys {
694			if vk.PublicKey != nil {
695				jwk, err := joseKeyToJSONWebKey(vk.PublicKey)
696				if err == nil {
697					jwks.Keys = append(jwks.Keys, jwk)
698				}
699			}
700		}
701
702		if !keys.NextRotation.IsZero() {
703			jwks.ExpiresInTime = keys.NextRotation
704		} else {
705			jwks.ExpiresInTime = time.Now().Add(1 * time.Hour)
706		}
707
708		return jwks, nil
709	}
710
711	// joseKeyToJSONWebKey converts a go-jose JSONWebKey to our JSONWebKey representation.
712	func joseKeyToJSONWebKey(key *jose.JSONWebKey) (nbjwt.JSONWebKey, error) {
713		jwk := nbjwt.JSONWebKey{
714			Kid: key.KeyID,
715			Use: key.Use,
716		}
717
718		switch pub := key.Key.(type) {
719		case *rsa.PublicKey:
720			jwk.Kty = "RSA"
721			jwk.N = base64.RawURLEncoding.EncodeToString(pub.N.Bytes())
722			jwk.E = base64.RawURLEncoding.EncodeToString(big.NewInt(int64(pub.E)).Bytes())
723		case *ecdsa.PublicKey:
724			jwk.Kty = "EC"
725			jwk.X = base64.RawURLEncoding.EncodeToString(pub.X.Bytes())
726			jwk.Y = base64.RawURLEncoding.EncodeToString(pub.Y.Bytes())
727			jwk.Crv = pub.Curve.Params().Name
728		default:
729			return jwk, fmt.Errorf("unsupported key type: %T", key.Key)
730		}
731
732		return jwk, nil
733	}
734

}

// 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
}


The issue: `GetJWKS` includes **all** verification keys regardless of whether they've expired. In the HTTP path, Dex's `/keys` endpoint already filters expired verification keys server-side. When reading directly from storage, we need to do that filtering ourselves.

Let me check the Dex `VerificationKey` struct again - it has an `Expiry` field.


Edit

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

	now := time.Now()
	for _, vk := range keys.VerificationKeys {
		if vk.PublicKey == nil {
			continue
		}
		// Skip expired verification keys, matching the filtering that
		// Dex's HTTP /keys endpoint performs server-side.
		if !vk.Expiry.IsZero() && now.After(vk.Expiry) {
			continue
		}
		jwk, err := joseKeyToJSONWebKey(vk.PublicKey)
		if err == nil {
			jwks.Keys = append(jwks.Keys, jwk)
		}
	}


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"
	"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
}

// 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)
		}
	}

	now := time.Now()
	for _, vk := range keys.VerificationKeys {
		if vk.PublicKey == nil {
			continue
		}
		// Skip expired verification keys, matching the filtering that
		// Dex's HTTP /keys endpoint performs server-side.
		if !vk.Expiry.IsZero() && now.After(vk.Expiry) {
			continue
		}
		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
}

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

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

+    // Skip expired verification keys, matching the filtering that

+    // Dex's HTTP /keys endpoint performs server-side.

+    if !vk.Expiry.IsZero() && now.After(vk.Expiry) {

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

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


Now let me also verify the `ExpiresInTime` logic is sound. In the HTTP path, `getPemKeys` sets `ExpiresInTime` from the `Cache-Control: max-age` header. The `stillValid()` check triggers a refresh when `ExpiresInTime` is approaching. Using `NextRotation` is semantically correct — it's when Dex will rotate the signing key, so we should re-fetch around that time. But let's also handle the case where keys exist but `NextRotation` is zero (e.g., keys were just created and rotation hasn't been scheduled yet) more conservatively.


676	// GetJWKS reads signing keys directly from Dex storage and returns them as Jwks.
677	// This avoids HTTP round-trips when the embedded IDP is co-located with the management server.
678	func (p *Provider) GetJWKS(ctx context.Context) (*nbjwt.Jwks, error) {
679		keys, err := p.storage.GetKeys(ctx)
680		if err != nil {
681			return nil, fmt.Errorf("failed to get keys from dex storage: %w", err)
682		}
683
684		jwks := &nbjwt.Jwks{}
685
686		if keys.SigningKeyPub != nil {
687			jwk, err := joseKeyToJSONWebKey(keys.SigningKeyPub)
688			if err == nil {
689				jwks.Keys = append(jwks.Keys, jwk)
690			}
691		}
692
693		now := time.Now()
694		for _, vk := range keys.VerificationKeys {
695			if vk.PublicKey == nil {
696				continue
697			}
698			// Skip expired verification keys, matching the filtering that
699			// Dex's HTTP /keys endpoint performs server-side.
700			if !vk.Expiry.IsZero() && now.After(vk.Expiry) {
701				continue
702			}
703			jwk, err := joseKeyToJSONWebKey(vk.PublicKey)
704			if err == nil {
705				jwks.Keys = append(jwks.Keys, jwk)
706			}
707		}
708
709		if !keys.NextRotation.IsZero() {
710			jwks.ExpiresInTime = keys.NextRotation
711		} else {
712			jwks.ExpiresInTime = time.Now().Add(1 * time.Hour)
713		}
714
715		return jwks, nil
716	}
717
718	// joseKeyToJSONWebKey converts a go-jose JSONWebKey to our JSONWebKey representation.
719	func joseKeyToJSONWebKey(key *jose.JSONWebKey) (nbjwt.JSONWebKey, error) {
720		jwk := nbjwt.JSONWebKey{

// 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)
		}
	}

	now := time.Now()
	for _, vk := range keys.VerificationKeys {
		if vk.PublicKey == nil {
			continue
		}
		// Skip expired verification keys, matching the filtering that
		// Dex's HTTP /keys endpoint performs server-side.
		if !vk.Expiry.IsZero() && now.After(vk.Expiry) {
			continue
		}
		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{


The logic looks correct now. Let me also verify there are no returned empty JWKS (no keys at all) which would cause issues — `stillValid()` is called on the result, and `getPublicKey` iterates `Keys`. An empty `Keys` slice would just trigger `errKeyNotFound`, which then triggers `forceRefreshKeys` — that's the correct fallback behavior.

Let me also verify the `stillValid()` interaction: when `ExpiresInTime = NextRotation` and we're past that time, `stillValid()` returns false, triggering `refreshKeys` which calls `GetJWKS` again — getting the newly rotated key. That's correct.


Bash

go build ./idp/dex/... 2>&1 && go build ./management/... 2>&1

Build affected packages


(Bash completed with no output)


go test ./shared/auth/jwt/... ./management/server/auth/... -count=1 -timeout 60s 2>&1 | tail -10

Run tests


ok  	github.com/netbirdio/netbird/shared/auth/jwt	0.614s
ok  	github.com/netbirdio/netbird/management/server/auth	0.803s


All good. The fix filters out expired verification keys in `GetJWKS`, matching what Dex's HTTP `/keys` endpoint does server-side. The expiration flow:

- `ExpiresInTime` is set to `NextRotation` (when Dex will rotate the signing key)
- `stillValid()` returns `false` when we're within 5 seconds of that time, triggering `refreshKeys`
- `refreshKeys` calls `GetJWKS` again, getting the newly rotated key and fresh `NextRotation`
- Expired `VerificationKeys` (old rotated keys past their expiry) are now filtered out, just as the HTTP endpoint would


use the same code as in dex itself


Let me find how Dex's `/keys` endpoint builds the JWKS response.


Agent

Find Dex keys endpoint code

Prompt:
In the Go module cache or vendor directory, find the Dex server code that handles the /keys endpoint and builds the JWKS JSON response. Look for how it converts storage.Keys into the JSON response with jose.JSONWebKey. Check:
1. The dexidp/dex/server package - look for "keys" handler, handlePublicKeys, or similar
2. How it serializes the keys to JSON (it likely uses go-jose's JSONWebKey.MarshalJSON)

The module cache is likely at ~/go/pkg/mod/github.com/dexidp/

Be thorough - I need the exact code that constructs the JWKS response.


The user doesn't want to proceed with this tool use. The tool use was rejected (eg. if it was a file edit, the new_string was NOT written to the file). STOP what you are doing and wait for the user to tell you how to proceed.

Note: The user's next message may contain a correction or preference. Pay close attention — if they explain what went wrong or how they'd prefer you to work, consider saving that to memory for future sessions.


[Request interrupted by user for tool use]


just use ../dexidp folder


Bash

ls /Users/misha/Documents/GolandProjects/netbird/netbird/../dexidp/ 2>/dev/null || ls /Users/misha/Documents/GolandProjects/netbird/dexidp/ 2>/dev/null || ls /Users/misha/Documents/GolandProjects/dexidp/ 2>/dev/null

Find dexidp folder


ADOPTERS.md
CONTRIBUTING.md
Dockerfile
LICENSE
MAINTAINERS.md
Makefile
README.md
api
cmd
config.dev.yaml
config.docker.yaml
config.yaml.dist
connector
docker-compose.override.yaml.dist
docker-compose.test.yaml
docker-compose.yaml
docs
examples
flake.lock
flake.nix
go.mod
go.sum
pkg
scripts
server
storage
web


<tool_use_error>Path does not exist: /Users/misha/Documents/GolandProjects/dexidp/server. Note: your current working directory is /Users/misha/Documents/GolandProjects/netbird/netbird.</tool_use_error>

Error: Path does not exist: /Users/misha/Documents/GolandProjects/dexidp/server. Note: your current working directory is /Users/misha/Documents/GolandProjects/netbird/netbird.


No matches found


33-	codeChallengeMethodS256  = "S256"
34-)
35-
36:func (s *Server) handlePublicKeys(w http.ResponseWriter, r *http.Request) {
37-	ctx := r.Context()
38-	/ TODO(ericchiang): Cache this.
39:	keys, err := s.signer.ValidationKeys(ctx)
40-	if err != nil {
41:		s.logger.ErrorContext(r.Context(), "failed to get keys", "err", err)
42-		s.renderError(r, w, http.StatusInternalServerError, "Internal server error.")
43-		return
44-	}
45-
46:	if len(keys) == 0 {
47:		s.logger.ErrorContext(r.Context(), "no public keys found.")
48-		s.renderError(r, w, http.StatusInternalServerError, "Internal server error.")
49-		return
50-	}
51-
52:	jwks := jose.JSONWebKeySet{
53:		Keys: make([]jose.JSONWebKey, len(keys)),
54-	}
55:	for i, key := range keys {
56:		jwks.Keys[i] = *key
57-	}
58-
59:	data, err := json.MarshalIndent(jwks, "", "  ")
60-	if err != nil {
61-		s.logger.ErrorContext(r.Context(), "failed to marshal discovery data", "err", err)
62-		s.renderError(r, w, http.StatusInternalServerError, "Internal server error.")
--
77-	Issuer            string   `json:"issuer"`
78-	Auth              string   `json:"authorization_endpoint"`
79-	Token             string   `json:"token_endpoint"`
80:	Keys              string   `json:"jwks_uri"`
81-	UserInfo          string   `json:"userinfo_endpoint"`
82-	DeviceEndpoint    string   `json:"device_authorization_endpoint"`
83-	Introspect        string   `json:"introspection_endpoint"`
--
111-		Issuer:            s.issuerURL.String(),
112-		Auth:              s.absURL("/auth"),
113-		Token:             s.absURL("/token"),
114:		Keys:              s.absURL("/keys"),
115-		UserInfo:          s.absURL("/userinfo"),
116-		DeviceEndpoint:    s.absURL("/device/code"),
117-		Introspect:        s.absURL("/token/introspect"),


It uses `s.signer.ValidationKeys(ctx)`. Let me check that.


1	package signer
2
3	import (
4		"context"
5		"errors"
6		"fmt"
7		"log/slog"
8		"time"
9
10		"github.com/go-jose/go-jose/v4"
11
12		"github.com/dexidp/dex/storage"
13	)
14
15	// LocalConfig holds configuration for the local signer.
16	type LocalConfig struct {
17		// KeysRotationPeriod defines the duration of time after which the signing keys will be rotated.
18		KeysRotationPeriod string `json:"keysRotationPeriod"`
19		// Algorithm defines the signing algorithm used for newly generated local keys.
20		// Changing it does not replace the current signing key immediately. The new
21		// algorithm is applied when Dex generates the next signing key during rotation.
22		// Supported values are RS256 and ES256.
23		Algorithm jose.SignatureAlgorithm `json:"algorithm"`
24	}
25
26	func (c *LocalConfig) signingAlgorithm() (jose.SignatureAlgorithm, error) {
27		if c.Algorithm == "" {
28			return jose.RS256, nil
29		}
30		if c.Algorithm == jose.RS256 || c.Algorithm == jose.ES256 {
31			return c.Algorithm, nil
32		}
33		return "", fmt.Errorf("unsupported local signer algorithm %q", c.Algorithm)
34	}
35
36	// Open creates a new local signer.
37	func (c *LocalConfig) Open(_ context.Context, s storage.Storage, idTokenValidFor time.Duration, now func() time.Time, logger *slog.Logger) (Signer, error) {
38		rotateKeysAfter, err := time.ParseDuration(c.KeysRotationPeriod)
39		if err != nil {
40			return nil, fmt.Errorf("invalid config value %q for local signer rotation period: %v", c.KeysRotationPeriod, err)
41		}
42
43		alg, err := c.signingAlgorithm()
44		if err != nil {
45			return nil, fmt.Errorf("invalid config value %q for local signer algorithm: %v", c.Algorithm, err)
46		}
47
48		strategy, err := rotationStrategyForAlgorithm(rotateKeysAfter, idTokenValidFor, alg)
49		if err != nil {
50			return nil, err
51		}
52		r := &keyRotator{s, strategy, now, logger}
53		return &localSigner{
54			storage: s,
55			rotator: r,
56			logger:  logger,
57		}, nil
58	}
59
60	// localSigner signs payloads using keys stored in the Dex storage.
61	// It manages key rotation and storage using the existing keyRotator logic.
62	type localSigner struct {
63		storage storage.Storage
64		rotator *keyRotator
65		logger  *slog.Logger
66	}
67
68	// Start begins key rotation in a new goroutine, closing once the context is canceled.
69	//
70	// The method blocks until after the first attempt to rotate keys has completed. That way
71	// healthy storages will return from this call with valid keys.
72	func (l *localSigner) Start(ctx context.Context) {
73		// Try to rotate immediately so properly configured storages will have keys.
74		if err := l.rotator.rotate(); err != nil {
75			l.logRotateError(err)
76		}
77
78		go func() {
79			for {
80				select {
81				case <-ctx.Done():
82					return
83				case <-time.After(time.Second * 30):
84					if err := l.rotator.rotate(); err != nil {
85						l.logRotateError(err)
86					}
87				}
88			}
89		}()
90	}
91
92	func (l *localSigner) logRotateError(err error) {
93		if errors.Is(err, errAlreadyRotated) {
94			l.logger.Info("key rotation not needed", "err", err)
95			return
96		}
97		l.logger.Error("failed to rotate keys", "err", err)
98	}
99
100	func (l *localSigner) Sign(ctx context.Context, payload []byte) (string, error) {
101		keys, err := l.storage.GetKeys(ctx)
102		if err != nil {
103			return "", fmt.Errorf("failed to get keys: %v", err)
104		}
105
106		signingKey := keys.SigningKey
107		if signingKey == nil {
108			return "", fmt.Errorf("no key to sign payload with")
109		}
110		signingAlg, err := signatureAlgorithm(signingKey)
111		if err != nil {
112			return "", err
113		}
114
115		return signPayload(signingKey, signingAlg, payload)
116	}
117
118	func (l *localSigner) ValidationKeys(ctx context.Context) ([]*jose.JSONWebKey, error) {
119		keys, err := l.storage.GetKeys(ctx)
120		if err != nil {
121			return nil, fmt.Errorf("failed to get keys: %v", err)
122		}
123
124		if keys.SigningKeyPub == nil {
125			return nil, fmt.Errorf("no public keys found")
126		}
127
128		jwks := make([]*jose.JSONWebKey, len(keys.VerificationKeys)+1)
129		jwks[0] = keys.SigningKeyPub
130		for i, verificationKey := range keys.VerificationKeys {
131			jwks[i+1] = verificationKey.PublicKey
132		}
133		return jwks, nil
134	}
135
136	func (l *localSigner) Algorithm(ctx context.Context) (jose.SignatureAlgorithm, error) {
137		keys, err := l.storage.GetKeys(ctx)
138		if err != nil && !errors.Is(err, storage.ErrNotFound) {
139			return "", fmt.Errorf("failed to get keys: %v", err)
140		}
141		if keys.SigningKey == nil {
142			return l.rotator.strategy.algorithm, nil
143		}
144		return signatureAlgorithm(keys.SigningKey)
145	}
146

package signer

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"time"

	"github.com/go-jose/go-jose/v4"

	"github.com/dexidp/dex/storage"
)

// LocalConfig holds configuration for the local signer.
type LocalConfig struct {
	// KeysRotationPeriod defines the duration of time after which the signing keys will be rotated.
	KeysRotationPeriod string `json:"keysRotationPeriod"`
	// Algorithm defines the signing algorithm used for newly generated local keys.
	// Changing it does not replace the current signing key immediately. The new
	// algorithm is applied when Dex generates the next signing key during rotation.
	// Supported values are RS256 and ES256.
	Algorithm jose.SignatureAlgorithm `json:"algorithm"`
}

func (c *LocalConfig) signingAlgorithm() (jose.SignatureAlgorithm, error) {
	if c.Algorithm == "" {
		return jose.RS256, nil
	}
	if c.Algorithm == jose.RS256 || c.Algorithm == jose.ES256 {
		return c.Algorithm, nil
	}
	return "", fmt.Errorf("unsupported local signer algorithm %q", c.Algorithm)
}

// Open creates a new local signer.
func (c *LocalConfig) Open(_ context.Context, s storage.Storage, idTokenValidFor time.Duration, now func() time.Time, logger *slog.Logger) (Signer, error) {
	rotateKeysAfter, err := time.ParseDuration(c.KeysRotationPeriod)
	if err != nil {
		return nil, fmt.Errorf("invalid config value %q for local signer rotation period: %v", c.KeysRotationPeriod, err)
	}

	alg, err := c.signingAlgorithm()
	if err != nil {
		return nil, fmt.Errorf("invalid config value %q for local signer algorithm: %v", c.Algorithm, err)
	}

	strategy, err := rotationStrategyForAlgorithm(rotateKeysAfter, idTokenValidFor, alg)
	if err != nil {
		return nil, err
	}
	r := &keyRotator{s, strategy, now, logger}
	return &localSigner{
		storage: s,
		rotator: r,
		logger:  logger,
	}, nil
}

// localSigner signs payloads using keys stored in the Dex storage.
// It manages key rotation and storage using the existing keyRotator logic.
type localSigner struct {
	storage storage.Storage
	rotator *keyRotator
	logger  *slog.Logger
}

// Start begins key rotation in a new goroutine, closing once the context is canceled.
//
// The method blocks until after the first attempt to rotate keys has completed. That way
// healthy storages will return from this call with valid keys.
func (l *localSigner) Start(ctx context.Context) {
	// Try to rotate immediately so properly configured storages will have keys.
	if err := l.rotator.rotate(); err != nil {
		l.logRotateError(err)
	}

	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case <-time.After(time.Second * 30):
				if err := l.rotator.rotate(); err != nil {
					l.logRotateError(err)
				}
			}
		}
	}()
}

func (l *localSigner) logRotateError(err error) {
	if errors.Is(err, errAlreadyRotated) {
		l.logger.Info("key rotation not needed", "err", err)
		return
	}
	l.logger.Error("failed to rotate keys", "err", err)
}

func (l *localSigner) Sign(ctx context.Context, payload []byte) (string, error) {
	keys, err := l.storage.GetKeys(ctx)
	if err != nil {
		return "", fmt.Errorf("failed to get keys: %v", err)
	}

	signingKey := keys.SigningKey
	if signingKey == nil {
		return "", fmt.Errorf("no key to sign payload with")
	}
	signingAlg, err := signatureAlgorithm(signingKey)
	if err != nil {
		return "", err
	}

	return signPayload(signingKey, signingAlg, payload)
}

func (l *localSigner) ValidationKeys(ctx context.Context) ([]*jose.JSONWebKey, error) {
	keys, err := l.storage.GetKeys(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to get keys: %v", err)
	}

	if keys.SigningKeyPub == nil {
		return nil, fmt.Errorf("no public keys found")
	}

	jwks := make([]*jose.JSONWebKey, len(keys.VerificationKeys)+1)
	jwks[0] = keys.SigningKeyPub
	for i, verificationKey := range keys.VerificationKeys {
		jwks[i+1] = verificationKey.PublicKey
	}
	return jwks, nil
}

func (l *localSigner) Algorithm(ctx context.Context) (jose.SignatureAlgorithm, error) {
	keys, err := l.storage.GetKeys(ctx)
	if err != nil && !errors.Is(err, storage.ErrNotFound) {
		return "", fmt.Errorf("failed to get keys: %v", err)
	}
	if keys.SigningKey == nil {
		return l.rotator.strategy.algorithm, nil
	}
	return signatureAlgorithm(keys.SigningKey)
}


Now I see it. Dex's `ValidationKeys` (line 118-134) does **no** expiry filtering — it just returns `SigningKeyPub` + all `VerificationKeys`. Then `handlePublicKeys` serializes them using `jose.JSONWebKeySet` which calls `jose.JSONWebKey.MarshalJSON()` — go-jose handles the serialization natively.

So the right approach is: use the same code as Dex — get keys from storage, use `jose.JSONWebKeySet` to marshal to JSON, then unmarshal into our `Jwks`. Or even better, just do what Dex does and build the key set the same way.


674	}
675
676	// GetJWKS reads signing keys directly from Dex storage and returns them as Jwks.
677	// This avoids HTTP round-trips when the embedded IDP is co-located with the management server.
678	func (p *Provider) GetJWKS(ctx context.Context) (*nbjwt.Jwks, error) {
679		keys, err := p.storage.GetKeys(ctx)
680		if err != nil {
681			return nil, fmt.Errorf("failed to get keys from dex storage: %w", err)
682		}
683
684		jwks := &nbjwt.Jwks{}
685
686		if keys.SigningKeyPub != nil {
687			jwk, err := joseKeyToJSONWebKey(keys.SigningKeyPub)
688			if err == nil {
689				jwks.Keys = append(jwks.Keys, jwk)
690			}
691		}
692
693		now := time.Now()
694		for _, vk := range keys.VerificationKeys {
695			if vk.PublicKey == nil {
696				continue
697			}
698			// Skip expired verification keys, matching the filtering that
699			// Dex's HTTP /keys endpoint performs server-side.
700			if !vk.Expiry.IsZero() && now.After(vk.Expiry) {
701				continue
702			}
703			jwk, err := joseKeyToJSONWebKey(vk.PublicKey)
704			if err == nil {
705				jwks.Keys = append(jwks.Keys, jwk)
706			}
707		}
708
709		if !keys.NextRotation.IsZero() {
710			jwks.ExpiresInTime = keys.NextRotation
711		} else {
712			jwks.ExpiresInTime = time.Now().Add(1 * time.Hour)
713		}
714
715		return jwks, nil
716	}
717
718	// joseKeyToJSONWebKey converts a go-jose JSONWebKey to our JSONWebKey representation.
719	func joseKeyToJSONWebKey(key *jose.JSONWebKey) (nbjwt.JSONWebKey, error) {
720		jwk := nbjwt.JSONWebKey{
721			Kid: key.KeyID,
722			Use: key.Use,
723		}
724
725		switch pub := key.Key.(type) {
726		case *rsa.PublicKey:
727			jwk.Kty = "RSA"
728			jwk.N = base64.RawURLEncoding.EncodeToString(pub.N.Bytes())
729			jwk.E = base64.RawURLEncoding.EncodeToString(big.NewInt(int64(pub.E)).Bytes())
730		case *ecdsa.PublicKey:
731			jwk.Kty = "EC"
732			jwk.X = base64.RawURLEncoding.EncodeToString(pub.X.Bytes())
733			jwk.Y = base64.RawURLEncoding.EncodeToString(pub.Y.Bytes())

}

// 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)
		}
	}

	now := time.Now()
	for _, vk := range keys.VerificationKeys {
		if vk.PublicKey == nil {
			continue
		}
		// Skip expired verification keys, matching the filtering that
		// Dex's HTTP /keys endpoint performs server-side.
		if !vk.Expiry.IsZero() && now.After(vk.Expiry) {
			continue
		}
		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())


Now I'll replace `GetJWKS` and `joseKeyToJSONWebKey` with the same approach Dex uses: get keys from storage identically to `ValidationKeys`, marshal via `jose.JSONWebKeySet`, unmarshal into our `Jwks`.


Edit

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"
)

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"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"
	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"
)