petal/internal/auth/oidc.go

package auth

import (
	"context"
	"crypto/rand"
	"crypto/subtle"
	"encoding/base64"
	"errors"
	"fmt"
	"log"
	"net/http"
	"strings"
	"sync"

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

// Temporary cookies that carry one login attempt from /auth/login to
// /auth/callback. They live for ten minutes and are cleared the moment the
// callback runs.
const (
	stateCookie = "petal_oidc_state"
	nonceCookie = "petal_oidc_nonce"
	pkceCookie  = "petal_oidc_pkce"

	loginAttemptTTL = 600 // seconds
)

// Options configures the OIDC client.
type Options struct {
	IssuerURL    string // Authentik's issuer, e.g. https://auth.example.com/application/o/petal/
	ClientID     string
	ClientSecret string
	BaseURL      string // Petal's public base URL; the redirect URI is derived from it
	Allowed      Allowlist
}

// OIDC implements Petal's half of an authorization-code login against
// Authentik: /auth/login starts it, /auth/callback finishes it by provisioning
// the account and issuing a session, /auth/logout ends it.
//
// Petal is the OIDC client itself rather than trusting a proxy-injected header.
// The header approach is far less code, but it is only safe while the container
// is unreachable except through that proxy — an invariant enforced by network
// configuration, not by anything in the repository, on a public host that also
// runs half a dozen other services. Petal holds someone's private journals; it
// should be safe to expose directly.
type OIDC struct {
	opts     Options
	sessions *SessionStore
	users    *UserStore
	secure   bool

	// The provider is discovered over the network, which means it can fail at
	// startup for reasons that have nothing to do with Petal. Discovery is
	// therefore lazy and retried: an Authentik outage blocks new logins but
	// leaves every existing session working, since those only need the database.
	mu       sync.Mutex
	provider *oidc.Provider
	oauth    *oauth2.Config
	verifier *oidc.IDTokenVerifier
}

// NewOIDC builds the login flow. It attempts discovery once so a misconfigured
// issuer shows up in the startup log rather than on the writer's first login,
// but a failure here is not fatal.
func NewOIDC(ctx context.Context, opts Options, sessions *SessionStore, users *UserStore) *OIDC {
	o := &OIDC{
		opts:     opts,
		sessions: sessions,
		users:    users,
		secure:   strings.HasPrefix(strings.ToLower(opts.BaseURL), "https://"),
	}
	if err := o.discover(ctx); err != nil {
		log.Printf("auth: OIDC discovery failed (%v) — login will retry on demand", err)
	}
	return o
}

// RedirectURI is the callback Authentik must have registered for this client.
func (o *OIDC) RedirectURI() string {
	return strings.TrimSuffix(o.opts.BaseURL, "/") + "/auth/callback"
}

// discover resolves the provider metadata and builds the oauth2 config.
func (o *OIDC) discover(ctx context.Context) error {
	o.mu.Lock()
	defer o.mu.Unlock()
	if o.provider != nil {
		return nil
	}
	// The issuer is passed through exactly as configured, trailing slash and
	// all: OIDC requires the discovered issuer to match the requested one
	// byte-for-byte, and Authentik's ends in a slash. (go-oidc trims it itself
	// when building the .well-known URL, so a slash here costs nothing.)
	provider, err := oidc.NewProvider(ctx, o.opts.IssuerURL)
	if err != nil {
		return err
	}
	o.provider = provider
	o.verifier = provider.Verifier(&oidc.Config{ClientID: o.opts.ClientID})
	o.oauth = &oauth2.Config{
		ClientID:     o.opts.ClientID,
		ClientSecret: o.opts.ClientSecret,
		Endpoint:     provider.Endpoint(),
		RedirectURL:  o.RedirectURI(),
		Scopes:       []string{oidc.ScopeOpenID, "profile", "email"},
	}
	return nil
}

// ready returns the discovered client, discovering it first if an earlier
// attempt failed.
func (o *OIDC) ready(ctx context.Context) (*oauth2.Config, *oidc.IDTokenVerifier, error) {
	if err := o.discover(ctx); err != nil {
		return nil, nil, err
	}
	o.mu.Lock()
	defer o.mu.Unlock()
	return o.oauth, o.verifier, nil
}

// Routes mounts the login endpoints. Mount at "/auth", outside /api: these are
// browser navigations, not API calls, and they must be reachable without a
// session — that is their whole purpose.
func (o *OIDC) Routes() chi.Router {
	r := chi.NewRouter()
	r.Get("/login", o.login)
	r.Get("/callback", o.callback)
	// POST only. Signing out is a state change, and SameSite=Lax deliberately
	// *does* send the session cookie on a top-level cross-site GET — so a GET
	// route here means any page on the internet can sign her out mid-draft by
	// linking to it, or embedding it as an image. Small harm, free to remove.
	r.Post("/logout", o.logout)
	return r
}

// login starts an authorization-code flow with PKCE.
func (o *OIDC) login(w http.ResponseWriter, r *http.Request) {
	// Already signed in? Don't bounce a valid session through the IdP.
	if _, err := o.sessions.Resolve(r); err == nil {
		http.Redirect(w, r, "/", http.StatusFound)
		return
	}

	conf, _, err := o.ready(r.Context())
	if err != nil {
		o.page(w, http.StatusServiceUnavailable,
			"登录暂时不可用", "Sign-in is unavailable right now",
			"Petal 联系不上登录服务。请稍后再试。",
			"Petal can't reach the sign-in service. Please try again in a moment.")
		return
	}

	state, err := randomToken()
	if err != nil {
		o.page(w, http.StatusInternalServerError, "出了点问题", "Something went wrong", "请再试一次。", "Please try again.")
		return
	}
	nonce, err := randomToken()
	if err != nil {
		o.page(w, http.StatusInternalServerError, "出了点问题", "Something went wrong", "请再试一次。", "Please try again.")
		return
	}
	pkce := oauth2.GenerateVerifier()

	// state defends the callback against CSRF (a forged callback can't know the
	// cookie); nonce ties the returned ID token to this attempt; PKCE binds the
	// code to this client even if it leaks in transit.
	o.setTemp(w, stateCookie, state)
	o.setTemp(w, nonceCookie, nonce)
	o.setTemp(w, pkceCookie, pkce)

	http.Redirect(w, r, conf.AuthCodeURL(state,
		oidc.Nonce(nonce),
		oauth2.S256ChallengeOption(pkce),
	), http.StatusFound)
}

// callback completes the flow: verify, allowlist, provision, issue a session.
func (o *OIDC) callback(w http.ResponseWriter, r *http.Request) {
	// Expire the one-shot login cookies up front, not on the way out: every exit
	// from here writes a response, and a Set-Cookie added after the header is
	// written is silently dropped. They're read from the request below, so
	// clearing them on the response now costs nothing.
	o.clearTemp(w)

	if errParam := r.URL.Query().Get("error"); errParam != "" {
		o.page(w, http.StatusForbidden,
			"登录未完成", "Sign-in didn't finish",
			"登录服务拒绝了这次请求。你可以再试一次。",
			"The sign-in service turned that request down. You can try again.")
		return
	}

	state, err := r.Cookie(stateCookie)
	if err != nil || state.Value == "" ||
		subtle.ConstantTimeCompare([]byte(state.Value), []byte(r.URL.Query().Get("state"))) != 1 {
		o.page(w, http.StatusBadRequest,
			"这个登录链接过期了", "That sign-in link expired",
			"请回到 Petal 重新登录。",
			"Head back to Petal and sign in again.")
		return
	}

	conf, verifier, err := o.ready(r.Context())
	if err != nil {
		o.page(w, http.StatusServiceUnavailable,
			"登录暂时不可用", "Sign-in is unavailable right now",
			"Petal 联系不上登录服务。请稍后再试。",
			"Petal can't reach the sign-in service. Please try again in a moment.")
		return
	}

	pkce, err := r.Cookie(pkceCookie)
	if err != nil {
		o.page(w, http.StatusBadRequest, "这个登录链接过期了", "That sign-in link expired",
			"请回到 Petal 重新登录。", "Head back to Petal and sign in again.")
		return
	}

	token, err := conf.Exchange(r.Context(), r.URL.Query().Get("code"), oauth2.VerifierOption(pkce.Value))
	if err != nil {
		log.Printf("auth: code exchange failed: %v", err)
		o.page(w, http.StatusBadGateway, "登录没有成功", "Sign-in didn't go through",
			"请再试一次。", "Please try again.")
		return
	}

	claims, err := o.claims(r.Context(), verifier, token)
	if err != nil {
		log.Printf("auth: id token rejected: %v", err)
		o.page(w, http.StatusBadGateway, "登录没有成功", "Sign-in didn't go through",
			"请再试一次。", "Please try again.")
		return
	}

	// Nonce check: this ID token must belong to the attempt that started here.
	nonce, err := r.Cookie(nonceCookie)
	if err != nil || subtle.ConstantTimeCompare([]byte(nonce.Value), []byte(claims.nonce)) != 1 {
		o.page(w, http.StatusBadRequest, "这个登录链接过期了", "That sign-in link expired",
			"请回到 Petal 重新登录。", "Head back to Petal and sign in again.")
		return
	}

	if !o.opts.Allowed.Permits(claims.Subject, claims.Email) {
		log.Printf("auth: rejected sign-in for sub=%s email=%s (not on the allowlist)", claims.Subject, claims.Email)
		o.page(w, http.StatusForbidden,
			"这个 Petal 不是给你写的", "This Petal isn't yours to write in",
			"你的账号是有效的，但还没有被邀请到这个 Petal。如果这是个误会，找管理员说一声就好。",
			"Your account is valid, but it hasn't been invited to this Petal. If that's a mistake, a word with whoever runs it will sort it out.")
		return
	}

	if err := o.users.Upsert(claims.Subject, claims.Email, claims.displayName()); err != nil {
		log.Printf("auth: provisioning failed: %v", err)
		o.page(w, http.StatusInternalServerError, "出了点问题", "Something went wrong",
			"请再试一次。", "Please try again.")
		return
	}

	session, err := o.sessions.Create(claims.Subject, r.UserAgent())
	if err != nil {
		log.Printf("auth: session creation failed: %v", err)
		o.page(w, http.StatusInternalServerError, "出了点问题", "Something went wrong",
			"请再试一次。", "Please try again.")
		return
	}
	SetSessionCookie(w, session, o.secure)
	log.Printf("auth: signed in %s (%s)", claims.Email, claims.Subject)

	http.Redirect(w, r, "/", http.StatusFound)
}

// logout revokes the session server-side and clears the cookie. Doing both
// matters: clearing only the cookie leaves a token that still works if it was
// ever captured.
func (o *OIDC) logout(w http.ResponseWriter, r *http.Request) {
	if token := SessionToken(r); token != "" {
		if err := o.sessions.Revoke(token); err != nil {
			log.Printf("auth: revoke failed: %v", err)
		}
	}
	ClearSessionCookie(w, o.secure)
	http.Redirect(w, r, "/", http.StatusFound)
}

// idClaims is the subset of the ID token Petal cares about.
type idClaims struct {
	Subject           string `json:"sub"`
	Email             string `json:"email"`
	Name              string `json:"name"`
	PreferredUsername string `json:"preferred_username"`

	nonce string
}

func (c idClaims) displayName() string {
	if c.Name != "" {
		return c.Name
	}
	if c.PreferredUsername != "" {
		return c.PreferredUsername
	}
	return c.Email
}

// claims verifies the ID token in a token response and extracts its claims.
func (o *OIDC) claims(ctx context.Context, verifier *oidc.IDTokenVerifier, token *oauth2.Token) (idClaims, error) {
	raw, ok := token.Extra("id_token").(string)
	if !ok || raw == "" {
		return idClaims{}, errors.New("no id_token in the token response")
	}
	idToken, err := verifier.Verify(ctx, raw)
	if err != nil {
		return idClaims{}, err
	}
	var claims idClaims
	if err := idToken.Claims(&claims); err != nil {
		return idClaims{}, err
	}
	if claims.Subject == "" {
		claims.Subject = idToken.Subject
	}
	claims.nonce = idToken.Nonce
	return claims, nil
}

func (o *OIDC) setTemp(w http.ResponseWriter, name, value string) {
	http.SetCookie(w, &http.Cookie{
		Name:     name,
		Value:    value,
		Path:     "/auth",
		HttpOnly: true,
		Secure:   o.secure,
		SameSite: http.SameSiteLaxMode,
		MaxAge:   loginAttemptTTL,
	})
}

func (o *OIDC) clearTemp(w http.ResponseWriter) {
	for _, name := range []string{stateCookie, nonceCookie, pkceCookie} {
		http.SetCookie(w, &http.Cookie{
			Name: name, Value: "", Path: "/auth",
			HttpOnly: true, Secure: o.secure, SameSite: http.SameSiteLaxMode, MaxAge: -1,
		})
	}
}

func randomToken() (string, error) {
	b := make([]byte, 24)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(b), nil
}

// page renders one of the flow's dead ends. Every one of them is a full stop in
// front of someone who was just trying to write, so they read as warm bilingual
// sentences rather than as a status code — the same standard as the rest of the
// app, and the reason these aren't plain http.Error calls.
func (o *OIDC) page(w http.ResponseWriter, status int, titleZH, titleEN, bodyZH, bodyEN string) {
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.WriteHeader(status)
	fmt.Fprintf(w, `<!doctype html>
<html lang="zh"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>%s · Petal</title>
<style>
  :root { color-scheme: light dark; }
  body { margin:0; min-height:100vh; display:flex; align-items:center; justify-content:center;
         background:#fdf8f5; color:#5b4b52;
         font-family:'Nunito','PingFang SC','Microsoft YaHei','Noto Sans CJK SC',system-ui,sans-serif; }
  main { max-width:30rem; padding:2.5rem; text-align:center; }
  .mark { font-size:2.5rem; }
  h1 { font-size:1.5rem; margin:.75rem 0 .25rem; font-weight:700; }
  h2 { font-size:1rem; margin:0 0 1.25rem; font-weight:600; opacity:.65; }
  p { line-height:1.7; margin:.4rem 0; }
  p.en { opacity:.7; font-size:.95rem; }
  a { display:inline-block; margin-top:1.75rem; padding:.6rem 1.4rem; border-radius:999px;
      background:#f3c7d3; color:#5b4b52; text-decoration:none; font-weight:700; }
  @media (prefers-color-scheme: dark) { body { background:#231b28; color:#e9dfe6; } a { background:#7c5f78; color:#fdf8f5; } }
</style></head>
<body><main>
  <div class="mark">🌸</div>
  <h1>%s</h1><h2>%s</h2>
  <p>%s</p><p class="en">%s</p>
  <a href="/">回到 Petal · Back to Petal</a>
</main></body></html>
`, titleEN, titleZH, titleEN, bodyZH, bodyEN)
}