geniusrun/backend/internal/auth/oidc.go

package auth

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net/url"
	"slices"

	"github.com/coreos/go-oidc/v3/oidc"
	"golang.org/x/oauth2"
)

// Verifier is the seam between internal/api and this package's real
// Keycloak-backed OIDC client, so tests can inject internal/auth/mock's fake
// instead of talking to a real identity provider.
type Verifier interface {
	// BeginLogin builds a Keycloak authorization URL and the transaction
	// state that must be round-tripped (via a TxnCookie) to HandleCallback.
	BeginLogin() (authURL string, txn TxnState, err error)
	// HandleCallback validates the callback query against txn, exchanges the
	// code, verifies the ID token, and reports whether the required role was
	// present.
	HandleCallback(ctx context.Context, txn TxnState, query url.Values) (LoginResult, error)
	// EndSessionURL builds the identity provider's logout URL, redirecting
	// back to postLogoutRedirectURL once Keycloak's own session is cleared.
	// idTokenHint, if non-empty, is passed as id_token_hint so Keycloak can
	// positively identify the session being ended and skip its own
	// logout-confirmation prompt (which would otherwise let the user cancel
	// out of logout after their geniusrun account is already deleted).
	EndSessionURL(postLogoutRedirectURL, idTokenHint string) string
}

// LoginResult is what a completed callback exchange resolves to.
type LoginResult struct {
	Claims     Claims
	Authorized bool
	// IDToken is the raw Keycloak ID token JWT, kept apart from Claims:
	// it's only ever needed once more, at logout (id_token_hint), so it
	// rides in the session cookie but never in the per-request Claims.
	IDToken string
}

// OIDCConfig configures NewOIDCVerifier. RequiredRole is checked against the
// ID token's realm_access.roles.
type OIDCConfig struct {
	IssuerURL    string
	ClientID     string
	ClientSecret string
	RedirectURL  string
	RequiredRole string
}

// idTokenClaims mirrors the subset of a Keycloak ID token's claims geniusrun
// cares about.
type idTokenClaims struct {
	Sub         string `json:"sub"`
	Name        string `json:"name"`
	Email       string `json:"email"`
	RealmAccess struct {
		Roles []string `json:"roles"`
	} `json:"realm_access"`
}

func (c idTokenClaims) hasRole(required string) bool {
	return slices.Contains(c.RealmAccess.Roles, required)
}

type oidcVerifier struct {
	provider     *oidc.Provider
	idTokenVerif *oidc.IDTokenVerifier
	oauth2Config oauth2.Config
	requiredRole string
}

// NewOIDCVerifier performs OIDC discovery against cfg.IssuerURL (once, at
// startup -- go-oidc caches the discovery document internally) and returns a
// Verifier backed by the real Keycloak realm.
func NewOIDCVerifier(ctx context.Context, cfg OIDCConfig) (Verifier, error) {
	provider, err := oidc.NewProvider(ctx, cfg.IssuerURL)
	if err != nil {
		return nil, fmt.Errorf("oidc discovery against %s: %w", cfg.IssuerURL, err)
	}
	return &oidcVerifier{
		provider:     provider,
		idTokenVerif: provider.Verifier(&oidc.Config{ClientID: cfg.ClientID}),
		oauth2Config: oauth2.Config{
			ClientID:     cfg.ClientID,
			ClientSecret: cfg.ClientSecret,
			RedirectURL:  cfg.RedirectURL,
			Endpoint:     provider.Endpoint(),
			Scopes:       []string{oidc.ScopeOpenID, "profile", "email"},
		},
		requiredRole: cfg.RequiredRole,
	}, nil
}

func randomString(n int) (string, error) {
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		return "", fmt.Errorf("generate random state: %w", err)
	}
	return base64.RawURLEncoding.EncodeToString(b), nil
}

func (v *oidcVerifier) BeginLogin() (string, TxnState, error) {
	state, err := randomString(24)
	if err != nil {
		return "", TxnState{}, err
	}
	verifier := oauth2.GenerateVerifier()
	authURL := v.oauth2Config.AuthCodeURL(state, oauth2.S256ChallengeOption(verifier))
	return authURL, TxnState{State: state, CodeVerifier: verifier}, nil
}

func (v *oidcVerifier) HandleCallback(ctx context.Context, txn TxnState, query url.Values) (LoginResult, error) {
	if got := query.Get("state"); got == "" || got != txn.State {
		return LoginResult{}, fmt.Errorf("state mismatch")
	}
	code := query.Get("code")
	if code == "" {
		return LoginResult{}, fmt.Errorf("callback missing code")
	}
	token, err := v.oauth2Config.Exchange(ctx, code, oauth2.VerifierOption(txn.CodeVerifier))
	if err != nil {
		return LoginResult{}, fmt.Errorf("exchange code: %w", err)
	}
	rawIDToken, ok := token.Extra("id_token").(string)
	if !ok || rawIDToken == "" {
		return LoginResult{}, fmt.Errorf("token response missing id_token")
	}
	idToken, err := v.idTokenVerif.Verify(ctx, rawIDToken)
	if err != nil {
		return LoginResult{}, fmt.Errorf("verify id_token: %w", err)
	}
	var claims idTokenClaims
	if err := idToken.Claims(&claims); err != nil {
		return LoginResult{}, fmt.Errorf("decode id_token claims: %w", err)
	}
	return LoginResult{
		Claims:     Claims{Sub: claims.Sub, Name: claims.Name, Email: claims.Email},
		Authorized: claims.hasRole(v.requiredRole),
		IDToken:    rawIDToken,
	}, nil
}

func (v *oidcVerifier) EndSessionURL(postLogoutRedirectURL, idTokenHint string) string {
	var discovery struct {
		EndSessionEndpoint string `json:"end_session_endpoint"`
	}
	if err := v.provider.Claims(&discovery); err != nil || discovery.EndSessionEndpoint == "" {
		return postLogoutRedirectURL
	}
	u, err := url.Parse(discovery.EndSessionEndpoint)
	if err != nil {
		return postLogoutRedirectURL
	}
	q := u.Query()
	q.Set("client_id", v.oauth2Config.ClientID)
	q.Set("post_logout_redirect_uri", postLogoutRedirectURL)
	if idTokenHint != "" {
		q.Set("id_token_hint", idTokenHint)
	}
	u.RawQuery = q.Encode()
	return u.String()
}