core/server/auth/oidc.go

package auth

import (
	"context"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/big"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

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

	"switchboard-core/models"
	"switchboard-core/store"
)

// OIDCConfig holds OIDC-specific configuration from env vars.
type OIDCConfig struct {
	IssuerURL    string // e.g. "https://keycloak.corp/realms/switchboard"
	ExternalIssuerURL string // browser-facing URL (optional, defaults to IssuerURL)
	ClientID     string
	ClientSecret string
	RedirectURL  string // e.g. "https://switchboard.corp/api/v1/auth/oidc/callback"
	AutoActivate bool   // auto-activate new users (default true)
	DefaultTeam  string // team ID for auto-provisioned users (optional)
	DefaultRole  string // "user" (default) or "admin"
	RolesClaim   string // JWT claim for role mapping (default "realm_access.roles")
	GroupsClaim  string // JWT claim for group sync (default "groups")
	AdminRole    string // IdP role value that maps to admin (default "admin")
}

// OIDCProvider authenticates via OpenID Connect tokens from an external IdP.
//
// Supports two flows:
//  1. Authorization code flow: browser redirects to IdP, callback exchanges
//     code for tokens. Used by the login page.
//  2. Direct token validation: API clients present a bearer token from the
//     IdP, backend validates signature + claims. Used by programmatic access.
//
// JWKS keys are cached and refreshed periodically.
type OIDCProvider struct {
	cfg       OIDCConfig
	discovery *OIDCDiscovery
	keys      map[string]*rsa.PublicKey
	keysMu    sync.RWMutex
	keysAt    time.Time
	client    *http.Client
}

// OIDCDiscovery holds the OIDC well-known configuration.
type OIDCDiscovery struct {
	Issuer                string `json:"issuer"`
	AuthorizationEndpoint string `json:"authorization_endpoint"`
	TokenEndpoint         string `json:"token_endpoint"`
	JWKSUri               string `json:"jwks_uri"`
	UserinfoEndpoint      string `json:"userinfo_endpoint"`
}

type oidcJWKS struct {
	Keys []oidcJWK `json:"keys"`
}

type oidcJWK struct {
	Kid string `json:"kid"`
	Kty string `json:"kty"`
	Alg string `json:"alg"`
	Use string `json:"use"`
	N   string `json:"n"`
	E   string `json:"e"`
}

type oidcTokenResponse struct {
	AccessToken  string `json:"access_token"`
	IDToken      string `json:"id_token"`
	TokenType    string `json:"token_type"`
	ExpiresIn    int    `json:"expires_in"`
	RefreshToken string `json:"refresh_token"`
}

func NewOIDCProvider(cfg OIDCConfig) (*OIDCProvider, error) {
	if cfg.IssuerURL == "" {
		return nil, fmt.Errorf("OIDC_ISSUER_URL is required")
	}
	if cfg.ClientID == "" {
		return nil, fmt.Errorf("OIDC_CLIENT_ID is required")
	}
	if cfg.DefaultRole == "" {
		cfg.DefaultRole = models.UserRoleUser
	}
	if cfg.RolesClaim == "" {
		cfg.RolesClaim = "realm_access.roles"
	}
	if cfg.GroupsClaim == "" {
		cfg.GroupsClaim = "groups"
	}
	if cfg.AdminRole == "" {
		cfg.AdminRole = "admin"
	}

	p := &OIDCProvider{
		cfg:    cfg,
		keys:   make(map[string]*rsa.PublicKey),
		client: &http.Client{Timeout: 10 * time.Second},
	}

	// Fetch discovery document on startup
	if err := p.fetchDiscovery(); err != nil {
		return nil, fmt.Errorf("OIDC discovery failed: %w", err)
	}

	// Pre-fetch JWKS
	if err := p.refreshKeys(); err != nil {
		log.Printf("[auth/oidc] warn: initial JWKS fetch failed: %v (will retry on first request)", err)
	}

	return p, nil
}

func (p *OIDCProvider) Mode() Mode { return ModeOIDC }

func (p *OIDCProvider) SupportsRegistration() bool { return false }

func (p *OIDCProvider) Register(_ *gin.Context, _ store.Stores) (*Result, error) {
	return nil, ErrNotSupported
}

// Authenticate validates a bearer token from the OIDC IdP.
//
// For the authorization code callback flow, use ExchangeCode() first
// to get the token, then call Authenticate with the token in the
// Authorization header.
func (p *OIDCProvider) Authenticate(c *gin.Context, stores store.Stores) (*Result, error) {
	// Extract token from Authorization header
	header := c.GetHeader("Authorization")
	if header == "" {
		return nil, fmt.Errorf("%w: missing Authorization header", ErrInvalidCreds)
	}
	tokenString := strings.TrimPrefix(header, "Bearer ")
	if tokenString == header {
		return nil, fmt.Errorf("%w: invalid authorization format", ErrInvalidCreds)
	}

	// Validate and parse the token
	claims, err := p.validateToken(tokenString)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrInvalidCreds, err)
	}

	ctx := c.Request.Context()
	sub := claims.Subject

	// Look up existing user by external ID
	user, err := stores.Users.GetByExternalID(ctx, string(ModeOIDC), sub)
	if err == nil && user != nil {
		if !user.IsActive {
			return nil, ErrInactive
		}

		// Re-evaluate role on every login
		role := p.resolveRole(claims)
		if role != user.Role {
			stores.Users.Update(ctx, user.ID, map[string]interface{}{"role": role})
			user.Role = role
		}

		// Sync groups
		p.syncGroups(ctx, user.ID, claims, stores)

		log.Printf("[auth/oidc] existing user %s (sub=%s)", user.Username, sub)
		return &Result{User: user, IsNewUser: false, VaultHint: ""}, nil
	}

	// Auto-provision
	if !p.cfg.AutoActivate {
		return nil, fmt.Errorf("%w: auto-provision disabled", ErrInvalidCreds)
	}

	return p.autoProvision(ctx, claims, stores)
}

// ExchangeCode exchanges an authorization code for tokens.
// Called by the /auth/oidc/callback handler.
func (p *OIDCProvider) ExchangeCode(ctx context.Context, code, redirectURI string) (*oidcTokenResponse, error) {
	if p.discovery == nil {
		return nil, fmt.Errorf("OIDC discovery not loaded")
	}

	data := url.Values{
		"grant_type":   {"authorization_code"},
		"code":         {code},
		"redirect_uri": {redirectURI},
		"client_id":    {p.cfg.ClientID},
	}
	if p.cfg.ClientSecret != "" {
		data.Set("client_secret", p.cfg.ClientSecret)
	}

	resp, err := p.client.PostForm(p.discovery.TokenEndpoint, data)
	if err != nil {
		return nil, fmt.Errorf("token exchange failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("token endpoint returned %d: %s", resp.StatusCode, string(body))
	}

	var tokenResp oidcTokenResponse
	if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
		return nil, fmt.Errorf("decode token response: %w", err)
	}

	return &tokenResp, nil
}

// AuthorizationURL returns the IdP login URL for the authorization code flow.
func (p *OIDCProvider) AuthorizationURL(state, nonce, redirectURI string) string {
	if p.discovery == nil {
		return ""
	}
	params := url.Values{
		"response_type": {"code"},
		"client_id":     {p.cfg.ClientID},
		"redirect_uri":  {redirectURI},
		"scope":         {"openid profile email"},
		"state":         {state},
		"nonce":         {nonce},
	}
	authEndpoint := p.discovery.AuthorizationEndpoint
	if p.cfg.ExternalIssuerURL != "" {
    	// Rewrite internal hostname to external for browser redirect
    	authEndpoint = strings.Replace(authEndpoint,
        	strings.TrimRight(p.cfg.IssuerURL, "/"),
        	strings.TrimRight(p.cfg.ExternalIssuerURL, "/"), 1)
	}
	return authEndpoint + "?" + params.Encode()
}

// ── Token Validation ────────────────────────────────────────────────

type oidcClaims struct {
	jwt.RegisteredClaims
	Email             string      `json:"email"`
	PreferredUsername  string      `json:"preferred_username"`
	Name              string      `json:"name"`
	Groups            []string    `json:"groups"`
	RealmAccess       realmAccess `json:"realm_access"`
	ResourceAccess    interface{} `json:"resource_access"`
}

type realmAccess struct {
	Roles []string `json:"roles"`
}

func (p *OIDCProvider) validateToken(tokenString string) (*oidcClaims, error) {
	claims := &oidcClaims{}

	token, err := jwt.ParseWithClaims(tokenString, claims, func(t *jwt.Token) (interface{}, error) {
		// Ensure RSA signing
		if _, ok := t.Method.(*jwt.SigningMethodRSA); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"])
		}

		kid, _ := t.Header["kid"].(string)
		if kid == "" {
			return nil, fmt.Errorf("token has no kid header")
		}

		key := p.getKey(kid)
		if key == nil {
			// Key not found — try refreshing JWKS (key rotation)
			if err := p.refreshKeys(); err != nil {
				return nil, fmt.Errorf("JWKS refresh failed: %w", err)
			}
			key = p.getKey(kid)
		}
		if key == nil {
			return nil, fmt.Errorf("unknown signing key: %s", kid)
		}

		return key, nil
	})

	if err != nil {
		return nil, err
	}
	if !token.Valid {
		return nil, fmt.Errorf("token is invalid")
	}

	// Verify issuer — accept both internal and external (split-horizon).
	// In Docker/K8s, the backend reaches Keycloak via internal hostname
	// but the browser (and thus the token's iss claim) uses the external URL.
	issuer := strings.TrimRight(p.cfg.IssuerURL, "/")
	externalIssuer := strings.TrimRight(p.cfg.ExternalIssuerURL, "/")
	tokenIssuer := strings.TrimRight(claims.Issuer, "/")
	if tokenIssuer != issuer && (externalIssuer == "" || tokenIssuer != externalIssuer) {
		return nil, fmt.Errorf("issuer mismatch: got %s, want %s (or %s)", claims.Issuer, issuer, externalIssuer)
	}

	return claims, nil
}

func (p *OIDCProvider) resolveRole(claims *oidcClaims) string {
	// Check realm_access.roles for admin role
	for _, role := range claims.RealmAccess.Roles {
		if role == p.cfg.AdminRole {
			return models.UserRoleAdmin
		}
	}
	return p.cfg.DefaultRole
}

// syncGroups compares IdP groups claim with internal group memberships
// and adds/removes as needed. OIDC-synced groups have source='oidc'.
func (p *OIDCProvider) syncGroups(ctx context.Context, userID string, claims *oidcClaims, stores store.Stores) {
	idpGroups := claims.Groups
	if len(idpGroups) == 0 {
		// Also try extracting from a nested claim path
		idpGroups = p.extractGroups(claims)
	}
	if len(idpGroups) == 0 {
		return
	}

	// Get user's current group IDs
	currentGroupIDs, err := stores.Groups.GetUserGroupIDs(ctx, userID)
	if err != nil {
		log.Printf("[auth/oidc] warn: could not get user groups: %v", err)
		return
	}

	// Build set of current OIDC-sourced groups the user belongs to
	currentOIDCGroups := make(map[string]string) // group name → group ID
	for _, gid := range currentGroupIDs {
		g, err := stores.Groups.GetByID(ctx, gid)
		if err != nil {
			continue
		}
		// Only track OIDC-sourced groups — manual groups are untouched
		if g.Source == "oidc" {
			currentOIDCGroups[g.Name] = g.ID
		}
	}

	// Ensure each IdP group exists and user is a member
	idpGroupSet := make(map[string]bool)
	for _, gname := range idpGroups {
		gname = strings.TrimSpace(gname)
		if gname == "" {
			continue
		}
		idpGroupSet[gname] = true

		if _, already := currentOIDCGroups[gname]; already {
			continue // already a member
		}

		// Find or create the group
		groupID := p.findOrCreateOIDCGroup(ctx, gname, userID, stores)
		if groupID == "" {
			continue
		}

		// Add membership
		if err := stores.Groups.AddMember(ctx, groupID, userID, userID); err != nil {
			log.Printf("[auth/oidc] warn: could not add %s to group %s: %v", userID, gname, err)
		} else {
			log.Printf("[auth/oidc] added user %s to OIDC group %s", userID, gname)
		}
	}

	// Remove from OIDC groups no longer in the IdP claim
	for gname, gid := range currentOIDCGroups {
		if !idpGroupSet[gname] {
			if err := stores.Groups.RemoveMember(ctx, gid, userID); err != nil {
				log.Printf("[auth/oidc] warn: could not remove %s from group %s: %v", userID, gname, err)
			} else {
				log.Printf("[auth/oidc] removed user %s from OIDC group %s (no longer in IdP)", userID, gname)
			}
		}
	}
}

func (p *OIDCProvider) findOrCreateOIDCGroup(ctx context.Context, name, createdBy string, stores store.Stores) string {
	// Search existing OIDC-sourced groups by name
	groups, err := stores.Groups.ListAll(ctx)
	if err != nil {
		return ""
	}
	for _, g := range groups {
		if g.Name == name && g.Source == "oidc" {
			return g.ID
		}
	}

	// Create new OIDC group
	g := &models.Group{
		Name:        name,
		Description: "Synced from OIDC",
		Scope:       "global",
		Source:      "oidc",
		CreatedBy:   &createdBy,
	}
	if err := stores.Groups.Create(ctx, g); err != nil {
		log.Printf("[auth/oidc] warn: could not create group %s: %v", name, err)
		return ""
	}
	log.Printf("[auth/oidc] created OIDC group %s (%s)", name, g.ID)
	return g.ID
}

func (p *OIDCProvider) extractGroups(claims *oidcClaims) []string {
	// Keycloak puts groups in different places depending on mapper config.
	// Try the configured claim path as a fallback.
	// The claims.Groups field handles the simple case. This handles
	// nested paths like "resource_access.switchboard.roles".
	// For now, return empty — the direct Groups field covers the
	// standard Keycloak group mapper output.
	return nil
}

func (p *OIDCProvider) autoProvision(
	ctx context.Context,
	claims *oidcClaims,
	stores store.Stores,
) (*Result, error) {
	sub := claims.Subject

	username := claims.PreferredUsername
	if username == "" {
		username = claims.Email
	}
	if username == "" {
		username = sub
	}

	email := claims.Email
	if email == "" {
		email = models.HandleFromName(username) + "@oidc.local"
	}

	displayName := claims.Name
	if displayName == "" {
		displayName = username
	}

	role := p.resolveRole(claims)
	handle := UniqueHandle(ctx, stores.Users, models.HandleFromName(username))

	user := &models.User{
		Username:    strings.ToLower(username),
		Email:       strings.ToLower(email),
		DisplayName: displayName,
		Role:        role,
		IsActive:    true,
		AuthSource:  string(ModeOIDC),
		ExternalID:  &sub,
		Handle:      handle,
	}

	if err := stores.Users.Create(ctx, user); err != nil {
		return nil, fmt.Errorf("auto-provision failed: %w", err)
	}

	log.Printf("[auth/oidc] auto-provisioned user %s (sub=%s, email=%s)", user.ID, sub, email)

	// Auto-add to default team
	if p.cfg.DefaultTeam != "" {
		if err := stores.Teams.AddMember(ctx, p.cfg.DefaultTeam, user.ID, "member"); err != nil {
			log.Printf("[auth/oidc] warn: could not add %s to default team: %v", user.ID, err)
		}
	}

	// Sync groups
	p.syncGroups(ctx, user.ID, claims, stores)

	return &Result{User: user, IsNewUser: true, VaultHint: ""}, nil
}

// ── JWKS Management ─────────────────────────────────────────────────

func (p *OIDCProvider) fetchDiscovery() error {
	discoveryURL := strings.TrimRight(p.cfg.IssuerURL, "/") + "/.well-known/openid-configuration"

	resp, err := p.client.Get(discoveryURL)
	if err != nil {
		return fmt.Errorf("fetch %s: %w", discoveryURL, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("discovery returned %d: %s", resp.StatusCode, string(body))
	}

	var disc OIDCDiscovery
	if err := json.NewDecoder(resp.Body).Decode(&disc); err != nil {
		return fmt.Errorf("decode discovery: %w", err)
	}

	p.discovery = &disc
	log.Printf("[auth/oidc] discovery loaded: issuer=%s, jwks=%s", disc.Issuer, disc.JWKSUri)
	return nil
}

func (p *OIDCProvider) refreshKeys() error {
	if p.discovery == nil || p.discovery.JWKSUri == "" {
		return fmt.Errorf("no JWKS URI (discovery not loaded)")
	}

	resp, err := p.client.Get(p.discovery.JWKSUri)
	if err != nil {
		return fmt.Errorf("fetch JWKS: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("JWKS returned %d", resp.StatusCode)
	}

	var jwks oidcJWKS
	if err := json.NewDecoder(resp.Body).Decode(&jwks); err != nil {
		return fmt.Errorf("decode JWKS: %w", err)
	}

	newKeys := make(map[string]*rsa.PublicKey)
	for _, k := range jwks.Keys {
		if k.Kty != "RSA" || k.Use != "sig" {
			continue
		}
		pub, err := parseRSAPublicKey(k.N, k.E)
		if err != nil {
			log.Printf("[auth/oidc] warn: skip key %s: %v", k.Kid, err)
			continue
		}
		newKeys[k.Kid] = pub
	}

	p.keysMu.Lock()
	p.keys = newKeys
	p.keysAt = time.Now()
	p.keysMu.Unlock()

	log.Printf("[auth/oidc] JWKS refreshed: %d signing keys", len(newKeys))
	return nil
}

func (p *OIDCProvider) getKey(kid string) *rsa.PublicKey {
	p.keysMu.RLock()
	defer p.keysMu.RUnlock()
	return p.keys[kid]
}

// parseRSAPublicKey builds an RSA public key from base64url-encoded N and E.
func parseRSAPublicKey(nB64, eB64 string) (*rsa.PublicKey, error) {
	nBytes, err := base64.RawURLEncoding.DecodeString(nB64)
	if err != nil {
		return nil, fmt.Errorf("decode N: %w", err)
	}
	eBytes, err := base64.RawURLEncoding.DecodeString(eB64)
	if err != nil {
		return nil, fmt.Errorf("decode E: %w", err)
	}

	n := new(big.Int).SetBytes(nBytes)
	e := new(big.Int).SetBytes(eBytes)

	return &rsa.PublicKey{
		N: n,
		E: int(e.Int64()),
	}, nil
}

// Discovery exposes the cached discovery document for route handlers.
func (p *OIDCProvider) Discovery() *OIDCDiscovery {
	return p.discovery
}