geniusrun/backend/internal/auth/session.go

// Package auth implements geniusrun's own login gate: an OIDC Authorization
// Code flow against an existing Keycloak realm (see oidc.go), backed by a
// signed session cookie geniusrun mints itself (this file) and a chi
// middleware that checks it (middleware.go). Keycloak's own access/refresh
// tokens are never stored or refreshed -- once HandleCallback verifies the
// ID token and role, only this package's own cookie matters for subsequent
// requests. The one exception is the raw ID token itself, carried opaquely
// inside the signed session cookie (apart from Claims -- see
// MintSessionCookie's idToken parameter and IDTokenFromSessionCookie)
// solely so a later logout can pass it back to Keycloak as
// id_token_hint -- letting Keycloak
// skip its own logout-confirmation prompt for a session it can positively
// identify, rather than leaving the user a chance to cancel out of it after
// their geniusrun account (and profile) is already gone.
package auth

import (
	"fmt"
	"net/http"
	"time"

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

const (
	// SessionCookieName holds the signed proof that a request's caller
	// completed login and passed the required-role check.
	SessionCookieName = "geniusrun_session"
	// TxnCookieName holds the short-lived OIDC login transaction (state +
	// PKCE verifier) between BeginLogin's redirect and HandleCallback.
	TxnCookieName = "geniusrun_oidc_txn"

	txnCookieDuration = 10 * time.Minute
)

// Claims identifies the authenticated user, carried in the signed session
// cookie and stashed in every request's context. Deliberately does NOT
// carry the raw Keycloak ID token: that JWT lives in the cookie payload
// separately (see MintSessionCookie/IDTokenFromSessionCookie) because it's
// only ever needed once more, at logout, and has no business riding
// through every handler's context.
type Claims struct {
	Sub   string
	Name  string
	Email string
}

type sessionClaims struct {
	Sub     string `json:"sub"`
	Name    string `json:"name"`
	Email   string `json:"email"`
	IDToken string `json:"id_token,omitempty"`
	jwt.RegisteredClaims
}

// TxnState is the OIDC login transaction round-tripped via TxnCookieName
// between BeginLogin and HandleCallback.
type TxnState struct {
	State        string
	CodeVerifier string
}

type txnClaims struct {
	State        string `json:"state"`
	CodeVerifier string `json:"code_verifier"`
	jwt.RegisteredClaims
}

// MintSessionCookie signs claims into a JWT valid for duration and wraps it
// in a cookie. secure should be true whenever the app is served over HTTPS.
// idToken is the raw Keycloak ID token to carry for the eventual logout's
// id_token_hint (empty is fine, e.g. in tests -- the hint is optional).
func MintSessionCookie(claims Claims, idToken string, secret []byte, duration time.Duration, secure bool) (*http.Cookie, error) {
	now := time.Now()
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, sessionClaims{
		Sub:     claims.Sub,
		Name:    claims.Name,
		Email:   claims.Email,
		IDToken: idToken,
		RegisteredClaims: jwt.RegisteredClaims{
			IssuedAt:  jwt.NewNumericDate(now),
			ExpiresAt: jwt.NewNumericDate(now.Add(duration)),
		},
	})
	signed, err := token.SignedString(secret)
	if err != nil {
		return nil, fmt.Errorf("sign session token: %w", err)
	}
	return &http.Cookie{
		Name:     SessionCookieName,
		Value:    signed,
		Path:     "/",
		HttpOnly: true,
		Secure:   secure,
		SameSite: http.SameSiteLaxMode,
		Expires:  now.Add(duration),
	}, nil
}

// ParseSessionCookie verifies the cookie's signature and expiry and decodes
// its Claims.
func ParseSessionCookie(cookie *http.Cookie, secret []byte) (Claims, error) {
	if cookie == nil {
		return Claims{}, fmt.Errorf("no session cookie")
	}
	var sc sessionClaims
	if _, err := jwt.ParseWithClaims(cookie.Value, &sc, keyfunc(secret), jwt.WithValidMethods([]string{"HS256"})); err != nil {
		return Claims{}, fmt.Errorf("parse session token: %w", err)
	}
	return Claims{Sub: sc.Sub, Name: sc.Name, Email: sc.Email}, nil
}

// IDTokenFromSessionCookie verifies the cookie and returns the raw Keycloak
// ID token it carries, for logout's id_token_hint. Only the logout handler
// needs this -- everything else uses ParseSessionCookie's Claims.
func IDTokenFromSessionCookie(cookie *http.Cookie, secret []byte) (string, error) {
	if cookie == nil {
		return "", fmt.Errorf("no session cookie")
	}
	var sc sessionClaims
	if _, err := jwt.ParseWithClaims(cookie.Value, &sc, keyfunc(secret), jwt.WithValidMethods([]string{"HS256"})); err != nil {
		return "", fmt.Errorf("parse session token: %w", err)
	}
	return sc.IDToken, nil
}

// MintTxnCookie signs an OIDC login transaction into a short-lived cookie.
func MintTxnCookie(txn TxnState, secret []byte, secure bool) (*http.Cookie, error) {
	now := time.Now()
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, txnClaims{
		State:        txn.State,
		CodeVerifier: txn.CodeVerifier,
		RegisteredClaims: jwt.RegisteredClaims{
			IssuedAt:  jwt.NewNumericDate(now),
			ExpiresAt: jwt.NewNumericDate(now.Add(txnCookieDuration)),
		},
	})
	signed, err := token.SignedString(secret)
	if err != nil {
		return nil, fmt.Errorf("sign txn token: %w", err)
	}
	return &http.Cookie{
		Name:     TxnCookieName,
		Value:    signed,
		Path:     "/",
		HttpOnly: true,
		Secure:   secure,
		SameSite: http.SameSiteLaxMode,
		Expires:  now.Add(txnCookieDuration),
	}, nil
}

// ParseTxnCookie verifies and decodes a login transaction cookie.
func ParseTxnCookie(cookie *http.Cookie, secret []byte) (TxnState, error) {
	if cookie == nil {
		return TxnState{}, fmt.Errorf("no txn cookie")
	}
	var tc txnClaims
	if _, err := jwt.ParseWithClaims(cookie.Value, &tc, keyfunc(secret), jwt.WithValidMethods([]string{"HS256"})); err != nil {
		return TxnState{}, fmt.Errorf("parse txn token: %w", err)
	}
	return TxnState{State: tc.State, CodeVerifier: tc.CodeVerifier}, nil
}

// ClearCookie returns a cookie that immediately expires the named cookie.
func ClearCookie(name string, secure bool) *http.Cookie {
	return &http.Cookie{
		Name:     name,
		Value:    "",
		Path:     "/",
		HttpOnly: true,
		Secure:   secure,
		SameSite: http.SameSiteLaxMode,
		MaxAge:   -1,
	}
}

func keyfunc(secret []byte) jwt.Keyfunc {
	return func(*jwt.Token) (any, error) { return secret, nil }
}