Archivum/backend/internal/auth/oidc.go

package auth

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"strings"
	"sync"
	"time"

	"github.com/brasse-b/archivum/internal/config"
	"github.com/coreos/go-oidc/v3/oidc"
	"golang.org/x/oauth2"
)

// OIDCUser is the identity extracted from a verified id_token / userinfo.
type OIDCUser struct {
	Username string
	Email    string
	Groups   []string
}

var ErrOIDCNotConfigured = errors.New("OIDC not configured")

type oidcState struct {
	expiresAt time.Time
}

// OIDCProvider wraps an OpenID Connect provider and OAuth2 client and manages
// the short-lived CSRF state values used during the auth-code flow. It can be
// reconfigured at runtime (after an admin changes the OIDC settings).
type OIDCProvider struct {
	mu       sync.Mutex
	cfg      config.OIDCConfig
	redirect string
	oauth    *oauth2.Config
	provider *oidc.Provider
	verifier *oidc.IDTokenVerifier
	ready    bool
	states   map[string]oidcState
}

func NewOIDCProvider() *OIDCProvider {
	return &OIDCProvider{states: make(map[string]oidcState)}
}

// Configure (re)initialises the provider from config. It performs OIDC
// discovery against the issuer, which requires network access to the IdP.
// Returns an error if discovery fails; the provider is then left disabled.
func (p *OIDCProvider) Configure(ctx context.Context, cfg config.OIDCConfig, redirectURL string) error {
	cfg.Normalize()

	p.mu.Lock()
	p.cfg = cfg
	p.redirect = redirectURL
	p.ready = false
	p.provider = nil
	p.verifier = nil
	p.oauth = nil
	p.mu.Unlock()

	if !cfg.Enabled || cfg.Issuer == "" || cfg.ClientID == "" || redirectURL == "" {
		return nil
	}

	// go-oidc strictly checks that the issuer we pass matches the issuer field
	// in the discovery document. Providers differ on the trailing slash
	// (Authentik returns ".../application/o/<slug>/" WITH a slash), so try the
	// value verbatim first and then the alternate slash form.
	issuer := strings.TrimSpace(cfg.Issuer)
	provider, err := oidc.NewProvider(ctx, issuer)
	if err != nil {
		var alt string
		if strings.HasSuffix(issuer, "/") {
			alt = strings.TrimRight(issuer, "/")
		} else {
			alt = issuer + "/"
		}
		if provider2, err2 := oidc.NewProvider(ctx, alt); err2 == nil {
			provider, err = provider2, nil
		}
	}
	if err != nil {
		return err
	}

	oauthCfg := &oauth2.Config{
		ClientID:     cfg.ClientID,
		ClientSecret: cfg.ClientSecret,
		Endpoint:     provider.Endpoint(),
		RedirectURL:  redirectURL,
		Scopes:       []string{oidc.ScopeOpenID, "profile", "email", cfg.GroupsClaim},
	}
	verifier := provider.Verifier(&oidc.Config{ClientID: cfg.ClientID})

	p.mu.Lock()
	p.provider = provider
	p.oauth = oauthCfg
	p.verifier = verifier
	p.ready = true
	p.mu.Unlock()
	return nil
}

// Enabled reports whether the provider is configured and ready.
func (p *OIDCProvider) Enabled() bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.ready
}

// AuthURL creates a fresh state value and returns the authorization URL to
// redirect the browser to.
func (p *OIDCProvider) AuthURL() (string, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if !p.ready {
		return "", ErrOIDCNotConfigured
	}
	state := randToken()
	p.states[state] = oidcState{expiresAt: time.Now().Add(10 * time.Minute)}
	p.gcStatesLocked()
	return p.oauth.AuthCodeURL(state), nil
}

// Exchange validates the state, swaps the code for tokens, verifies the
// id_token and returns the resulting identity (with groups). If the id_token
// carries no groups claim it falls back to the userinfo endpoint.
func (p *OIDCProvider) Exchange(ctx context.Context, state, code string) (*OIDCUser, error) {
	p.mu.Lock()
	if !p.ready {
		p.mu.Unlock()
		return nil, ErrOIDCNotConfigured
	}
	st, ok := p.states[state]
	if ok {
		delete(p.states, state)
	}
	oauthCfg := p.oauth
	verifier := p.verifier
	provider := p.provider
	cfg := p.cfg
	p.mu.Unlock()

	if !ok || time.Now().After(st.expiresAt) {
		return nil, errors.New("invalid or expired state")
	}

	tok, err := oauthCfg.Exchange(ctx, code)
	if err != nil {
		return nil, err
	}
	rawID, ok := tok.Extra("id_token").(string)
	if !ok {
		return nil, errors.New("no id_token in token response")
	}
	idTok, err := verifier.Verify(ctx, rawID)
	if err != nil {
		return nil, err
	}

	var claims map[string]interface{}
	if err := idTok.Claims(&claims); err != nil {
		return nil, err
	}

	user := &OIDCUser{
		Username: claimString(claims, cfg.UsernameClaim),
		Email:    claimString(claims, "email"),
		Groups:   claimStrings(claims, cfg.GroupsClaim),
	}

	// Some providers only expose groups via userinfo — fall back to it.
	if len(user.Groups) == 0 && provider != nil {
		if ui, err := provider.UserInfo(ctx, oauth2.StaticTokenSource(tok)); err == nil {
			var uiClaims map[string]interface{}
			if err := ui.Claims(&uiClaims); err == nil {
				user.Groups = claimStrings(uiClaims, cfg.GroupsClaim)
				if user.Username == "" {
					user.Username = claimString(uiClaims, cfg.UsernameClaim)
				}
				if user.Email == "" {
					user.Email = claimString(uiClaims, "email")
				}
			}
		}
	}

	if user.Username == "" {
		user.Username = user.Email
	}
	if user.Username == "" {
		user.Username = idTok.Subject
	}
	if user.Username == "" {
		return nil, errors.New("could not determine username from OIDC claims")
	}
	return user, nil
}

func (p *OIDCProvider) gcStatesLocked() {
	now := time.Now()
	for k, v := range p.states {
		if now.After(v.expiresAt) {
			delete(p.states, k)
		}
	}
}

func claimString(claims map[string]interface{}, key string) string {
	if v, ok := claims[key].(string); ok {
		return v
	}
	return ""
}

func claimStrings(claims map[string]interface{}, key string) []string {
	raw, ok := claims[key]
	if !ok {
		return nil
	}
	switch v := raw.(type) {
	case []interface{}:
		out := make([]string, 0, len(v))
		for _, item := range v {
			if s, ok := item.(string); ok && s != "" {
				out = append(out, s)
			}
		}
		return out
	case []string:
		return v
	case string:
		if v == "" {
			return nil
		}
		return []string{v}
	}
	return nil
}

func randToken() string {
	b := make([]byte, 32)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}