package server

import (
	"bytes"
	"context"
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/hmac"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	_ "crypto/sha512" // SHA-384/512 for RS384, ES384, RS512…
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"math/big"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"pve-confsync/internal/store"
)

// Authentication of the web UI, on the model of Logstream (auth.go, auth_local.go):
// local accounts (SQLite, PBKDF2) shown on a login page, and/or an OpenID Connect
// provider (Keycloak, Authentik, Authelia…) with the authorization code flow and
// PKCE. Both can be enabled together: the local accounts stay the way in when the
// provider is down. Sessions are HMAC signed cookies. Only the standard library.

const (
	sessionCookie = "confsync_session"
	loginCookie   = "confsync_login_" // + state: one cookie per OIDC login in progress
	loginTTL      = 10 * time.Minute
	clockSkew     = time.Minute
	loginPage     = "/login.html"
)

var loginFailDelay = time.Second // slows down password guessing

// AuthConfig is read from the environment (see README).
type AuthConfig struct {
	Local         bool // local accounts (AUTH_LOCAL, default true)
	Issuer        string
	ClientID      string
	ClientSecret  string
	RedirectURL   string
	Scopes        string
	GroupsClaim   string
	AdminGroup    string // OIDC_ADMIN_GROUP: members are admins
	OperatorGroup string // OIDC_OPERATOR_GROUP: members are operators
	DefaultRole   string // OIDC_DEFAULT_ROLE: everyone else (default viewer)
	SessionTTL    time.Duration
	DataDir       string
	Demo          bool // demo mode: the login page lists the demo accounts
}

// OIDCEnabled tells whether an OpenID Connect provider is configured.
func (c AuthConfig) OIDCEnabled() bool { return c.Issuer != "" }

// Identity is the logged in user.
type Identity struct {
	User string `json:"user"`
	Role string `json:"role"`
	Mode string `json:"mode"` // local or oidc
}

type ctxKey struct{}

// identity returns the user of a request (set by the middleware).
func identity(r *http.Request) Identity {
	id, _ := r.Context().Value(ctxKey{}).(Identity)
	return id
}

type session struct {
	User  string `json:"u"`
	Role  string `json:"r"`
	Mode  string `json:"m"`
	Epoch int64  `json:"p,omitempty"` // local: password epoch, a new password ends the sessions
	Exp   int64  `json:"e"`
}

// Auth is the middleware protecting the UI and the API.
type Auth struct {
	cfg    AuthConfig
	store  *store.Store
	key    []byte
	secure bool // cookies HTTPS only (OIDC_REDIRECT_URL in https)
	audit  func(r *http.Request, id Identity, action, target string, ok bool, detail string)
	next   http.Handler

	callback string
	client   *http.Client
	mu       sync.Mutex
	meta     *oidcMeta
	keys     map[string]crypto.PublicKey
	keysAt   time.Time
}

// NewAuth checks the configuration and loads the session key.
func NewAuth(c AuthConfig, st *store.Store) (*Auth, error) {
	if c.SessionTTL <= 0 {
		c.SessionTTL = 12 * time.Hour
	}
	if c.DefaultRole == "" {
		c.DefaultRole = store.RoleViewer
	}
	if !store.ValidRole(c.DefaultRole) {
		return nil, fmt.Errorf("OIDC_DEFAULT_ROLE=%q: expected admin, operator or viewer", c.DefaultRole)
	}
	a := &Auth{cfg: c, store: st, key: sessionKey(c.DataDir), client: &http.Client{Timeout: 10 * time.Second}}
	if !c.Local && !c.OIDCEnabled() {
		return nil, errors.New("AUTH_LOCAL=false needs an OIDC provider (OIDC_ISSUER)")
	}
	if c.OIDCEnabled() {
		var missing []string
		for _, v := range [][2]string{{"OIDC_CLIENT_ID", c.ClientID}, {"OIDC_CLIENT_SECRET", c.ClientSecret}, {"OIDC_REDIRECT_URL", c.RedirectURL}} {
			if v[1] == "" {
				missing = append(missing, v[0])
			}
		}
		if len(missing) > 0 {
			return nil, fmt.Errorf("OIDC: missing %s", strings.Join(missing, ", "))
		}
		ru, err := url.Parse(c.RedirectURL)
		if err != nil || ru.Host == "" || ru.Path == "" || ru.Path == "/" {
			return nil, fmt.Errorf("OIDC_REDIRECT_URL=%q: expected a full URL such as https://confsync.example.org/auth/callback", c.RedirectURL)
		}
		a.callback, a.secure = ru.Path, ru.Scheme == "https"
		if a.cfg.Scopes == "" {
			a.cfg.Scopes = "openid profile email"
		}
		if !strings.Contains(" "+a.cfg.Scopes+" ", " openid ") {
			a.cfg.Scopes = "openid " + a.cfg.Scopes
		}
		if a.cfg.GroupsClaim == "" {
			a.cfg.GroupsClaim = "groups"
		}
		if c.AdminGroup == "" && c.OperatorGroup == "" {
			log.Printf("oidc: no OIDC_ADMIN_GROUP / OIDC_OPERATOR_GROUP, every OIDC user gets the role %s", c.DefaultRole)
		}
	}
	return a, nil
}

// CheckProvider reads the provider configuration at startup so a mistake shows in the logs.
func (a *Auth) CheckProvider() {
	if a.cfg.OIDCEnabled() {
		if _, err := a.discover(); err != nil {
			log.Printf("oidc: %v", err)
		} else {
			log.Printf("oidc authentication enabled (issuer %s)", a.cfg.Issuer)
		}
	}
}

// sessionKey is kept in the data directory so sessions survive a restart.
func sessionKey(dir string) []byte {
	path := filepath.Join(dir, "session.key")
	if k, err := os.ReadFile(path); err == nil && len(k) >= 32 {
		return k
	}
	k := make([]byte, 32)
	if _, err := rand.Read(k); err != nil {
		log.Fatalf("session key: %v", err)
	}
	if err := os.WriteFile(path, k, 0o600); err != nil {
		log.Printf("auth: cannot save %s (%v): sessions end when the server restarts", path, err)
	}
	return k
}

// public paths: reachable without a session.
func public(p string) bool {
	switch p {
	case "/healthz", loginPage, "/style.css", "/boot.js", "/login.js", "/favicon.svg", "/auth/config", "/auth/login", "/auth/logout", "/auth/oidc", "/auth/setup":
		return true
	}
	return false
}

func (a *Auth) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.URL.Path == "/auth/login":
		a.handleLogin(w, r)
		return
	case r.URL.Path == "/auth/logout":
		a.handleLogout(w, r)
		return
	case r.URL.Path == "/auth/setup":
		a.handleSetup(w, r)
		return
	case r.URL.Path == "/auth/oidc" && a.cfg.OIDCEnabled():
		a.startLogin(w, r)
		return
	case a.callback != "" && r.URL.Path == a.callback:
		a.handleCallback(w, r)
		return
	case r.URL.Path == "/auth/config":
		writeJSON(w, http.StatusOK, map[string]any{"local": a.cfg.Local, "oidc": a.cfg.OIDCEnabled(), "demo": a.cfg.Demo,
			"setup": a.needSetup(r.Context())})
		return
	}
	id, ok := a.sessionIdentity(r)
	if ok {
		if r.URL.Path == loginPage {
			http.Redirect(w, r, safeReturn(r.URL.Query().Get("r")), http.StatusFound)
			return
		}
		a.next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), ctxKey{}, id)))
		return
	}
	switch {
	case public(r.URL.Path):
		if r.URL.Path == loginPage {
			w.Header().Set("Cache-Control", "no-store")
		}
		a.next.ServeHTTP(w, r)
	case r.Method == http.MethodGet && !strings.HasPrefix(r.URL.Path, "/api/") && r.URL.Path != "/auth/me":
		target := loginPage
		if ret := r.URL.RequestURI(); ret != "/" {
			target += "?" + url.Values{"r": {ret}}.Encode()
		}
		http.Redirect(w, r, target, http.StatusFound)
	default:
		writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "authentication required", "code": "auth"})
	}
}

// sessionIdentity checks the session cookie. Local accounts are read again from
// the database: a deleted user, a new role or a new password apply at once.
func (a *Auth) sessionIdentity(r *http.Request) (Identity, bool) {
	var s session
	c, err := r.Cookie(sessionCookie)
	if err != nil || !verifyCookie(a.key, c.Value, &s) || time.Now().Unix() >= s.Exp {
		return Identity{}, false
	}
	switch s.Mode {
	case "local":
		if !a.cfg.Local {
			return Identity{}, false
		}
		u, err := a.store.User(r.Context(), s.User)
		if err != nil || u.Epoch != s.Epoch {
			return Identity{}, false
		}
		return Identity{User: u.Name, Role: u.Role, Mode: "local"}, true
	case "oidc":
		if !a.cfg.OIDCEnabled() || !store.ValidRole(s.Role) {
			return Identity{}, false
		}
		return Identity{User: s.User, Role: s.Role, Mode: "oidc"}, true
	}
	return Identity{}, false
}

func (a *Auth) setSession(w http.ResponseWriter, r *http.Request, s session) {
	s.Exp = time.Now().Add(a.cfg.SessionTTL).Unix()
	http.SetCookie(w, &http.Cookie{
		Name:     sessionCookie,
		Value:    signCookie(a.key, s),
		Path:     "/",
		MaxAge:   int(a.cfg.SessionTTL.Seconds()),
		HttpOnly: true,
		Secure:   a.secure || isHTTPS(r),
		SameSite: http.SameSiteLaxMode,
	})
}

func (a *Auth) handleLogin(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost || !a.cfg.Local {
		http.Redirect(w, r, loginPage, http.StatusFound)
		return
	}
	user, pass := strings.TrimSpace(r.PostFormValue("user")), r.PostFormValue("pass")
	ret := safeReturn(r.PostFormValue("r"))
	u, ok := a.store.Authenticate(r.Context(), user, pass)
	if !ok {
		log.Printf("auth: failed login for %q from %s", user, clientIP(r))
		a.audit(r, Identity{User: user, Mode: "local"}, "login", "", false, "mot de passe ou utilisateur incorrect")
		time.Sleep(loginFailDelay)
		q := url.Values{"e": {"1"}}
		if ret != "/" {
			q.Set("r", ret)
		}
		http.Redirect(w, r, loginPage+"?"+q.Encode(), http.StatusSeeOther)
		return
	}
	id := Identity{User: u.Name, Role: u.Role, Mode: "local"}
	log.Printf("auth: %s logged in from %s (%s)", u.Name, clientIP(r), u.Role)
	a.audit(r, id, "login", "", true, "")
	a.setSession(w, r, session{User: u.Name, Role: u.Role, Mode: "local", Epoch: u.Epoch})
	http.Redirect(w, r, ret, http.StatusSeeOther)
}

// needSetup: local accounts enabled and none yet, the login page asks for the
// first administrator.
func (a *Auth) needSetup(ctx context.Context) bool {
	if !a.cfg.Local {
		return false
	}
	n, err := a.store.CountUsers(ctx)
	return err == nil && n == 0
}

// handleSetup creates the first administrator on an empty database, then opens
// its session. Refused as soon as an account exists.
func (a *Auth) handleSetup(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost || !a.needSetup(r.Context()) {
		http.Redirect(w, r, loginPage, http.StatusFound)
		return
	}
	user, pass := strings.TrimSpace(r.PostFormValue("user")), r.PostFormValue("pass")
	fail := func(code string) {
		http.Redirect(w, r, loginPage+"?"+url.Values{"setup": {code}}.Encode(), http.StatusSeeOther)
	}
	if pass != r.PostFormValue("pass2") {
		fail("mismatch")
		return
	}
	if err := a.store.CreateFirstAdmin(r.Context(), user, pass); err != nil {
		log.Printf("auth: first administrator: %v", err)
		if errors.Is(err, store.ErrSetupDone) {
			http.Redirect(w, r, loginPage, http.StatusSeeOther)
			return
		}
		fail("invalid")
		return
	}
	u, err := a.store.User(r.Context(), user)
	if err != nil {
		fail("invalid")
		return
	}
	id := Identity{User: u.Name, Role: u.Role, Mode: "local"}
	log.Printf("auth: first administrator %q created from %s", u.Name, clientIP(r))
	a.audit(r, id, "setup", u.Name, true, "first administrator created")
	a.setSession(w, r, session{User: u.Name, Role: u.Role, Mode: "local", Epoch: u.Epoch})
	http.Redirect(w, r, "/#clusters", http.StatusSeeOther)
}

// handleLogout ends the session; with OIDC the provider's session ends on its logout page if it has one.
func (a *Auth) handleLogout(w http.ResponseWriter, r *http.Request) {
	id, ok := a.sessionIdentity(r)
	http.SetCookie(w, &http.Cookie{Name: sessionCookie, Path: "/", MaxAge: -1, HttpOnly: true, Secure: a.secure || isHTTPS(r)})
	if ok {
		a.audit(r, id, "logout", "", true, "")
	}
	if ok && id.Mode == "oidc" {
		if meta, err := a.discover(); err == nil && meta.EndSession != "" {
			http.Redirect(w, r, addQuery(meta.EndSession, url.Values{"client_id": {a.cfg.ClientID}}), http.StatusFound)
			return
		}
	}
	http.Redirect(w, r, loginPage, http.StatusFound)
}

// safeReturn keeps the page to open after login inside the application.
func safeReturn(ret string) string {
	if !strings.HasPrefix(ret, "/") || strings.HasPrefix(ret, "//") || strings.HasPrefix(ret, "/\\") || ret == loginPage {
		return "/"
	}
	return ret
}

// isHTTPS is true when the browser talks HTTPS, directly or through a reverse proxy.
func isHTTPS(r *http.Request) bool {
	return r.TLS != nil || strings.EqualFold(r.Header.Get("X-Forwarded-Proto"), "https")
}

func clientIP(r *http.Request) string {
	if f := r.Header.Get("X-Forwarded-For"); f != "" {
		return strings.TrimSpace(strings.Split(f, ",")[0])
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

// ---- OpenID Connect ----

type oidcMeta struct {
	Issuer           string   `json:"issuer"`
	AuthEndpoint     string   `json:"authorization_endpoint"`
	TokenEndpoint    string   `json:"token_endpoint"`
	JWKSURI          string   `json:"jwks_uri"`
	EndSession       string   `json:"end_session_endpoint"`
	TokenAuthMethods []string `json:"token_endpoint_auth_methods_supported"`
}

type loginState struct {
	Nonce    string `json:"n"`
	Verifier string `json:"v"`
	Return   string `json:"r"`
	Exp      int64  `json:"e"`
}

func (a *Auth) startLogin(w http.ResponseWriter, r *http.Request) {
	meta, err := a.discover()
	if err != nil {
		log.Printf("oidc: %v", err)
		http.Error(w, "identity provider unreachable, try again later", http.StatusBadGateway)
		return
	}
	state, nonce, verifier := randomString(), randomString(), randomString()+randomString()
	http.SetCookie(w, &http.Cookie{
		Name:     loginCookie + state,
		Value:    signCookie(a.key, loginState{Nonce: nonce, Verifier: verifier, Return: safeReturn(r.URL.Query().Get("r")), Exp: time.Now().Add(loginTTL).Unix()}),
		Path:     "/",
		MaxAge:   int(loginTTL.Seconds()),
		HttpOnly: true,
		Secure:   a.secure,
		SameSite: http.SameSiteLaxMode, // sent back on the redirect from the provider
	})
	challenge := sha256.Sum256([]byte(verifier))
	q := url.Values{
		"response_type":         {"code"},
		"client_id":             {a.cfg.ClientID},
		"redirect_uri":          {a.cfg.RedirectURL},
		"scope":                 {a.cfg.Scopes},
		"state":                 {state},
		"nonce":                 {nonce},
		"code_challenge":        {base64.RawURLEncoding.EncodeToString(challenge[:])},
		"code_challenge_method": {"S256"},
	}
	http.Redirect(w, r, addQuery(meta.AuthEndpoint, q), http.StatusFound)
}

func (a *Auth) handleCallback(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()
	if e := q.Get("error"); e != "" {
		log.Printf("oidc: login refused by the provider: %s %s", e, q.Get("error_description"))
		http.Error(w, "login refused by the identity provider: "+e, http.StatusForbidden)
		return
	}
	state := q.Get("state")
	var ls loginState
	c, err := r.Cookie(loginCookie + state)
	if state == "" || err != nil || !verifyCookie(a.key, c.Value, &ls) || time.Now().Unix() > ls.Exp {
		http.Error(w, "login expired or started in another browser: open pve-confsync again", http.StatusBadRequest)
		return
	}
	http.SetCookie(w, &http.Cookie{Name: loginCookie + state, Path: "/", MaxAge: -1, HttpOnly: true, Secure: a.secure})

	user, role, err := a.exchange(r, q.Get("code"), ls)
	if err != nil {
		log.Printf("oidc: login failed: %v", err)
		a.audit(r, Identity{Mode: "oidc"}, "login", "", false, err.Error())
		http.Error(w, "login failed, see the pve-confsync logs", http.StatusForbidden)
		return
	}
	id := Identity{User: user, Role: role, Mode: "oidc"}
	log.Printf("oidc: %s logged in (%s)", user, role)
	a.audit(r, id, "login", "", true, "")
	a.setSession(w, r, session{User: user, Role: role, Mode: "oidc"})
	http.Redirect(w, r, ls.Return, http.StatusFound)
}

// exchange trades the code for tokens and returns the user name from the verified
// ID token and the role given by its groups.
func (a *Auth) exchange(r *http.Request, code string, ls loginState) (string, string, error) {
	if code == "" {
		return "", "", errors.New("no code in the callback")
	}
	meta, err := a.discover()
	if err != nil {
		return "", "", err
	}
	form := url.Values{
		"grant_type":    {"authorization_code"},
		"code":          {code},
		"redirect_uri":  {a.cfg.RedirectURL},
		"code_verifier": {ls.Verifier},
	}
	// client_secret_basic is the default; some providers only accept client_secret_post.
	post := len(meta.TokenAuthMethods) > 0 && !contains(meta.TokenAuthMethods, "client_secret_basic") && contains(meta.TokenAuthMethods, "client_secret_post")
	if post {
		form.Set("client_id", a.cfg.ClientID)
		form.Set("client_secret", a.cfg.ClientSecret)
	}
	req, err := http.NewRequestWithContext(r.Context(), http.MethodPost, meta.TokenEndpoint, strings.NewReader(form.Encode()))
	if err != nil {
		return "", "", err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("Accept", "application/json")
	if !post {
		req.SetBasicAuth(url.QueryEscape(a.cfg.ClientID), url.QueryEscape(a.cfg.ClientSecret))
	}
	res, err := a.client.Do(req)
	if err != nil {
		return "", "", fmt.Errorf("token endpoint: %w", err)
	}
	defer res.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(res.Body, 1<<20))
	if res.StatusCode != http.StatusOK {
		return "", "", fmt.Errorf("token endpoint: %s: %s", res.Status, bytes.TrimSpace(body))
	}
	var tok struct {
		IDToken string `json:"id_token"`
	}
	if err := json.Unmarshal(body, &tok); err != nil || tok.IDToken == "" {
		return "", "", errors.New("token endpoint: no id_token in the response")
	}
	claims, err := a.verifyIDToken(tok.IDToken, ls.Nonce)
	if err != nil {
		return "", "", err
	}
	role := a.cfg.DefaultRole
	switch {
	case a.cfg.AdminGroup != "" && hasGroup(claims[a.cfg.GroupsClaim], a.cfg.AdminGroup):
		role = store.RoleAdmin
	case a.cfg.OperatorGroup != "" && hasGroup(claims[a.cfg.GroupsClaim], a.cfg.OperatorGroup):
		role = store.RoleOperator
	}
	for _, k := range []string{"preferred_username", "email", "name", "sub"} {
		if v, _ := claims[k].(string); v != "" {
			return v, role, nil
		}
	}
	return "", "", errors.New("id_token: no sub")
}

// hasGroup tells whether the groups claim (a list, or a single string) holds
// group; a leading "/" (Keycloak group paths) is ignored.
func hasGroup(claim any, group string) bool {
	group = strings.TrimPrefix(group, "/")
	var groups []string
	switch v := claim.(type) {
	case string:
		groups = strings.Fields(strings.ReplaceAll(v, ",", " "))
	case []any:
		for _, g := range v {
			if s, ok := g.(string); ok {
				groups = append(groups, s)
			}
		}
	}
	for _, g := range groups {
		if strings.TrimPrefix(g, "/") == group {
			return true
		}
	}
	return false
}

// verifyIDToken checks the signature (keys from jwks_uri) and the claims of an ID token.
func (a *Auth) verifyIDToken(raw, nonce string) (map[string]any, error) {
	parts := strings.Split(raw, ".")
	if len(parts) != 3 {
		return nil, errors.New("id_token: not a JWT")
	}
	var hdr struct {
		Alg string `json:"alg"`
		Kid string `json:"kid"`
	}
	if err := decodeSegment(parts[0], &hdr); err != nil {
		return nil, fmt.Errorf("id_token header: %w", err)
	}
	sig, err := base64.RawURLEncoding.DecodeString(parts[2])
	if err != nil {
		return nil, errors.New("id_token: bad signature encoding")
	}
	key, err := a.keyFor(hdr.Kid)
	if err != nil {
		return nil, err
	}
	if err := verifySignature(hdr.Alg, key, []byte(parts[0]+"."+parts[1]), sig); err != nil {
		return nil, fmt.Errorf("id_token: %w", err)
	}
	var claims map[string]any
	if err := decodeSegment(parts[1], &claims); err != nil {
		return nil, fmt.Errorf("id_token claims: %w", err)
	}
	if iss, _ := claims["iss"].(string); iss != a.cfg.Issuer {
		return nil, fmt.Errorf("id_token: issuer %q, expected %q", iss, a.cfg.Issuer)
	}
	var aud []string
	switch v := claims["aud"].(type) {
	case string:
		aud = []string{v}
	case []any:
		for _, x := range v {
			if s, ok := x.(string); ok {
				aud = append(aud, s)
			}
		}
	}
	if !contains(aud, a.cfg.ClientID) {
		return nil, fmt.Errorf("id_token: audience %v does not include %q", aud, a.cfg.ClientID)
	}
	if azp, ok := claims["azp"].(string); ok && len(aud) > 1 && azp != a.cfg.ClientID {
		return nil, fmt.Errorf("id_token: azp %q", azp)
	}
	now := time.Now()
	exp, _ := claims["exp"].(float64)
	if exp == 0 || now.After(time.Unix(int64(exp), 0).Add(clockSkew)) {
		return nil, errors.New("id_token: expired (check the clocks)")
	}
	if iat, ok := claims["iat"].(float64); ok && time.Unix(int64(iat), 0).After(now.Add(clockSkew)) {
		return nil, errors.New("id_token: issued in the future (check the clocks)")
	}
	if n, _ := claims["nonce"].(string); subtle.ConstantTimeCompare([]byte(n), []byte(nonce)) != 1 {
		return nil, errors.New("id_token: wrong nonce")
	}
	return claims, nil
}

func verifySignature(alg string, key crypto.PublicKey, signed, sig []byte) error {
	if len(alg) != 5 {
		return fmt.Errorf("unsupported algorithm %q", alg)
	}
	var h crypto.Hash
	switch alg[2:] {
	case "256":
		h = crypto.SHA256
	case "384":
		h = crypto.SHA384
	case "512":
		h = crypto.SHA512
	}
	if h == 0 {
		return fmt.Errorf("unsupported algorithm %q", alg)
	}
	hh := h.New()
	hh.Write(signed)
	digest := hh.Sum(nil)
	switch k := key.(type) {
	case *rsa.PublicKey:
		switch alg[:2] {
		case "RS":
			return rsa.VerifyPKCS1v15(k, h, digest, sig)
		case "PS":
			return rsa.VerifyPSS(k, h, digest, sig, &rsa.PSSOptions{SaltLength: rsa.PSSSaltLengthEqualsHash})
		}
	case *ecdsa.PublicKey:
		size := (k.Curve.Params().BitSize + 7) / 8
		if alg[:2] != "ES" || len(sig) != 2*size {
			break
		}
		r, s := new(big.Int).SetBytes(sig[:size]), new(big.Int).SetBytes(sig[size:])
		if ecdsa.Verify(k, digest, r, s) {
			return nil
		}
		return errors.New("bad signature")
	}
	return fmt.Errorf("algorithm %q does not match the key", alg)
}

// discover reads the provider configuration once (and again after a failure).
func (a *Auth) discover() (*oidcMeta, error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	if a.meta != nil {
		return a.meta, nil
	}
	u := strings.TrimSuffix(a.cfg.Issuer, "/") + "/.well-known/openid-configuration"
	var m oidcMeta
	if err := a.getJSON(u, &m); err != nil {
		return nil, fmt.Errorf("discovery: %w", err)
	}
	if m.Issuer != a.cfg.Issuer {
		return nil, fmt.Errorf("discovery: the provider says its issuer is %q, set OIDC_ISSUER to that exact value", m.Issuer)
	}
	if m.AuthEndpoint == "" || m.TokenEndpoint == "" || m.JWKSURI == "" {
		return nil, errors.New("discovery: incomplete provider configuration")
	}
	a.meta = &m
	return a.meta, nil
}

// keyFor returns the signing key kid; the key set is reloaded when the provider rotates its keys.
func (a *Auth) keyFor(kid string) (crypto.PublicKey, error) {
	meta, err := a.discover()
	if err != nil {
		return nil, err
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	pick := func() crypto.PublicKey {
		if k, ok := a.keys[kid]; ok {
			return k
		}
		if kid == "" && len(a.keys) == 1 {
			for _, k := range a.keys {
				return k
			}
		}
		return nil
	}
	if k := pick(); k != nil {
		return k, nil
	}
	if time.Since(a.keysAt) < 10*time.Second {
		return nil, fmt.Errorf("id_token: unknown key %q", kid)
	}
	var set struct {
		Keys []struct {
			Kty string `json:"kty"`
			Kid string `json:"kid"`
			Use string `json:"use"`
			N   string `json:"n"`
			E   string `json:"e"`
			Crv string `json:"crv"`
			X   string `json:"x"`
			Y   string `json:"y"`
		} `json:"keys"`
	}
	if err := a.getJSON(meta.JWKSURI, &set); err != nil {
		return nil, fmt.Errorf("jwks: %w", err)
	}
	keys := map[string]crypto.PublicKey{}
	for _, k := range set.Keys {
		if k.Use != "" && k.Use != "sig" {
			continue
		}
		switch k.Kty {
		case "RSA":
			n, e := decodeBig(k.N), decodeBig(k.E)
			if n != nil && e != nil && e.IsInt64() {
				keys[k.Kid] = &rsa.PublicKey{N: n, E: int(e.Int64())}
			}
		case "EC":
			var c elliptic.Curve
			switch k.Crv {
			case "P-256":
				c = elliptic.P256()
			case "P-384":
				c = elliptic.P384()
			case "P-521":
				c = elliptic.P521()
			}
			x, y := decodeBig(k.X), decodeBig(k.Y)
			if c != nil && x != nil && y != nil && c.IsOnCurve(x, y) {
				keys[k.Kid] = &ecdsa.PublicKey{Curve: c, X: x, Y: y}
			}
		}
	}
	a.keys, a.keysAt = keys, time.Now()
	if k := pick(); k != nil {
		return k, nil
	}
	return nil, fmt.Errorf("id_token: unknown key %q", kid)
}

func (a *Auth) getJSON(u string, v any) error {
	res, err := a.client.Get(u)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("%s: %s", u, res.Status)
	}
	return json.NewDecoder(io.LimitReader(res.Body, 1<<20)).Decode(v)
}

// Cookies are base64url(JSON) + "." + base64url(HMAC-SHA256).
func signCookie(key []byte, v any) string {
	b, _ := json.Marshal(v)
	p := base64.RawURLEncoding.EncodeToString(b)
	m := hmac.New(sha256.New, key)
	m.Write([]byte(p))
	return p + "." + base64.RawURLEncoding.EncodeToString(m.Sum(nil))
}

func verifyCookie(key []byte, s string, v any) bool {
	p, sig, ok := strings.Cut(s, ".")
	if !ok {
		return false
	}
	got, err := base64.RawURLEncoding.DecodeString(sig)
	if err != nil {
		return false
	}
	m := hmac.New(sha256.New, key)
	m.Write([]byte(p))
	if !hmac.Equal(got, m.Sum(nil)) {
		return false
	}
	return decodeSegment(p, v) == nil
}

func decodeSegment(s string, v any) error {
	b, err := base64.RawURLEncoding.DecodeString(s)
	if err != nil {
		return err
	}
	return json.Unmarshal(b, v)
}

func decodeBig(s string) *big.Int {
	b, err := base64.RawURLEncoding.DecodeString(s)
	if err != nil || len(b) == 0 {
		return nil
	}
	return new(big.Int).SetBytes(b)
}

func randomString() string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		panic(err)
	}
	return base64.RawURLEncoding.EncodeToString(b)
}

func addQuery(endpoint string, q url.Values) string {
	sep := "?"
	if strings.Contains(endpoint, "?") {
		sep = "&"
	}
	return endpoint + sep + q.Encode()
}

func contains(list []string, s string) bool {
	for _, v := range list {
		if v == s {
			return true
		}
	}
	return false
}