package main

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/hmac"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"math/big"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

// OIDC is generic: discovery, authorization code with PKCE, and JWKS
// verification. Identity is iss plus sub; the name is derived once, since
// preferred_username promises neither uniqueness nor stability.

const (
	oidcCookie  = "oidc"
	oidcLife    = 10 * time.Minute // how long a login may take
	oidcRefresh = time.Hour        // how long discovery and keys are kept
	oidcTimeout = 15 * time.Second
)

// discovery is the part of a provider's metadata that gets used.
type discovery struct {
	Issuer        string `json:"issuer"`
	Authorization string `json:"authorization_endpoint"`
	Token         string `json:"token_endpoint"`
	JWKS          string `json:"jwks_uri"`
}

type providerCache struct {
	meta   discovery
	keys   []jwk
	metaAt time.Time
	keysAt time.Time
}

// OIDC is a cache; no session state lives here.
type OIDC struct {
	client *http.Client
	mu     sync.Mutex
	cache  map[string]*providerCache
}

func newOIDC() *OIDC {
	return &OIDC{client: &http.Client{Timeout: oidcTimeout}, cache: map[string]*providerCache{}}
}

func (o *OIDC) entry(issuer string) *providerCache {
	o.mu.Lock()
	defer o.mu.Unlock()
	if o.cache[issuer] == nil {
		o.cache[issuer] = &providerCache{}
	}
	return o.cache[issuer]
}

// meta refuses a document that does not name its own issuer.
func (o *OIDC) meta(issuer string) (discovery, error) {
	entry := o.entry(issuer)
	o.mu.Lock()
	if time.Since(entry.metaAt) < oidcRefresh && entry.meta.Issuer != "" {
		meta := entry.meta
		o.mu.Unlock()
		return meta, nil
	}
	o.mu.Unlock()

	var meta discovery
	if err := o.fetch(issuer+"/.well-known/openid-configuration", &meta); err != nil {
		return meta, err
	}
	switch {
	case meta.Issuer != issuer:
		return meta, fmt.Errorf("oidc: %s says it is %q", issuer, meta.Issuer)
	case meta.Authorization == "" || meta.Token == "" || meta.JWKS == "":
		return meta, fmt.Errorf("oidc: %s is missing an endpoint", issuer)
	}

	o.mu.Lock()
	entry.meta, entry.metaAt = meta, time.Now()
	o.mu.Unlock()
	return meta, nil
}

// signingKeys fetches the key set, again when a token names a key not held: a
// rotation.
func (o *OIDC) signingKeys(issuer, jwksURL string, force bool) ([]jwk, error) {
	entry := o.entry(issuer)
	o.mu.Lock()
	fresh := time.Since(entry.keysAt) < oidcRefresh && len(entry.keys) > 0
	keys := entry.keys
	o.mu.Unlock()
	if fresh && !force {
		return keys, nil
	}

	var set struct {
		Keys []jwk `json:"keys"`
	}
	if err := o.fetch(jwksURL, &set); err != nil {
		return nil, err
	}
	o.mu.Lock()
	entry.keys, entry.keysAt = set.Keys, time.Now()
	o.mu.Unlock()
	return set.Keys, nil
}

func (o *OIDC) fetch(url string, into any) error {
	resp, err := o.client.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("%s: %s", url, resp.Status)
	}
	body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return err
	}
	return json.Unmarshal(body, into)
}

// The state and PKCE verifier go into a short lived signed cookie, which
// survives a restart.
func (f *Forge) handleOIDCStart(w http.ResponseWriter, r *http.Request) {
	provider, ok := f.provider(r.PathValue("key"))
	if !ok {
		f.fail(w, r, http.StatusNotFound, "no such provider")
		return
	}
	meta, err := f.oidc.meta(provider.Issuer)
	if err != nil {
		log.Printf("oidc: %v", err)
		f.fail(w, r, http.StatusBadGateway, "that provider is not answering")
		return
	}

	state, verifier := randomToken(16), randomToken(32)
	f.setOIDCCookie(w, provider.Key, state, verifier)

	sum := sha256.Sum256([]byte(verifier))
	q := url.Values{
		"response_type":         {"code"},
		"client_id":             {provider.ClientID},
		"redirect_uri":          {f.oidcRedirect(provider.Key)},
		"scope":                 {"openid email profile"},
		"state":                 {state},
		"code_challenge":        {base64.RawURLEncoding.EncodeToString(sum[:])},
		"code_challenge_method": {"S256"},
	}
	http.Redirect(w, r, meta.Authorization+"?"+q.Encode(), http.StatusSeeOther)
}

func (f *Forge) handleOIDCCallback(w http.ResponseWriter, r *http.Request) {
	provider, ok := f.provider(r.PathValue("key"))
	if !ok {
		f.fail(w, r, http.StatusNotFound, "no such provider")
		return
	}
	key, state, verifier, ok := f.oidcCookie(r)
	f.clearOIDCCookie(w)
	if !ok || key != provider.Key {
		f.fail(w, r, http.StatusForbidden, "that login took too long, or did not start here. try again")
		return
	}
	q := r.URL.Query()
	if !hmac.Equal([]byte(q.Get("state")), []byte(state)) {
		f.fail(w, r, http.StatusForbidden, "that login did not start here. try again")
		return
	}
	if reason := q.Get("error"); reason != "" {
		f.fail(w, r, http.StatusForbidden, "the provider refused: "+reason)
		return
	}

	claims, err := f.oidcExchange(provider, q.Get("code"), verifier)
	if err != nil {
		log.Printf("oidc: %v", err)
		f.fail(w, r, http.StatusForbidden, "that login could not be verified")
		return
	}
	// Without a list, every address the provider will vouch for mints an
	// account.
	if !provider.allows(claims.Email) {
		log.Printf("oidc: %s refused %q: not in the allow list", provider.Key, claims.Email)
		f.fail(w, r, http.StatusForbidden, "that address may not sign in here")
		return
	}

	name, u, err := f.oidcAccount(provider, claims)
	if err != nil {
		f.problem(w, r, http.StatusInternalServerError, err)
		return
	}
	// The provider is the first factor, not the second.
	if u.TOTP != nil {
		f.askForCode(w, r, name, u, "")
		return
	}
	f.setSession(w, name, u.Epoch)
	http.Redirect(w, r, "/", http.StatusSeeOther)
}

// oidcExchange trades the code for an id token and verifies it.
func (f *Forge) oidcExchange(provider OIDCProvider, code, verifier string) (*claims, error) {
	if code == "" {
		return nil, errors.New("no code")
	}
	meta, err := f.oidc.meta(provider.Issuer)
	if err != nil {
		return nil, err
	}

	form := url.Values{
		"grant_type":    {"authorization_code"},
		"code":          {code},
		"redirect_uri":  {f.oidcRedirect(provider.Key)},
		"client_id":     {provider.ClientID},
		"code_verifier": {verifier},
	}
	req, err := http.NewRequest(http.MethodPost, meta.Token, strings.NewReader(form.Encode()))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("Accept", "application/json")
	if provider.ClientSecret != "" {
		req.SetBasicAuth(url.QueryEscape(provider.ClientID), url.QueryEscape(provider.ClientSecret))
	}
	resp, err := f.oidc.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("token endpoint: %s: %s", resp.Status, strings.TrimSpace(string(body)))
	}
	var token struct {
		IDToken string `json:"id_token"`
	}
	if err := json.Unmarshal(body, &token); err != nil {
		return nil, err
	}
	if token.IDToken == "" {
		return nil, errors.New("no id_token in the token response")
	}
	return f.verifyIDToken(provider, token.IDToken)
}

func (f *Forge) verifyIDToken(provider OIDCProvider, token string) (*claims, error) {
	signed, signature, kid, alg, payload, err := splitJWT(token)
	if err != nil {
		return nil, err
	}
	meta, err := f.oidc.meta(provider.Issuer)
	if err != nil {
		return nil, err
	}
	keys, err := f.oidc.signingKeys(provider.Issuer, meta.JWKS, false)
	if err != nil {
		return nil, err
	}
	key, ok := findKey(keys, kid)
	if !ok {
		if keys, err = f.oidc.signingKeys(provider.Issuer, meta.JWKS, true); err != nil {
			return nil, err
		}
		if key, ok = findKey(keys, kid); !ok {
			return nil, fmt.Errorf("id token signed by unknown key %q", kid)
		}
	}
	if err := verifySignature(key, alg, signed, signature); err != nil {
		return nil, err
	}

	var c claims
	if err := json.Unmarshal(payload, &c); err != nil {
		return nil, err
	}
	now := time.Now()
	switch {
	case c.Iss != provider.Issuer:
		return nil, fmt.Errorf("id token issued by %q, not %q", c.Iss, provider.Issuer)
	case c.Sub == "":
		return nil, errors.New("id token has no subject")
	case !c.Aud.has(provider.ClientID):
		return nil, fmt.Errorf("id token is for %v, not for %s", c.Aud, provider.ClientID)
	case c.Exp == 0 || now.After(time.Unix(c.Exp, 0).Add(time.Minute)):
		return nil, errors.New("id token has expired")
	case c.Iat != 0 && time.Unix(c.Iat, 0).After(now.Add(5*time.Minute)):
		return nil, errors.New("id token was issued in the future")
	case c.EmailVerified == nil || !*c.EmailVerified:
		return nil, errors.New("the provider did not say the address is verified")
	case c.Email == "" || !validEmail(c.Email):
		// This reaches git's author field and a mail envelope.
		return nil, fmt.Errorf("id token carries %q, which is not an address", c.Email)
	}
	return &c, nil
}

// oidcAccount finds the account this identity belongs to, or provisions one.
func (f *Forge) oidcAccount(provider OIDCProvider, c *claims) (string, User, error) {
	for name, u := range f.store.Users() {
		if u.Origin == originOIDC && u.Iss == c.Iss && u.Sub == c.Sub {
			// The provider owns the address; the name is what grants point at
			// and never changes.
			if u.Email != c.Email || !u.EmailConfirmed {
				err := f.store.UpdateUser(name, "users: refresh the address of "+name, Author{Name: name, Email: c.Email},
					func(acct *User) error {
						acct.Email, acct.EmailConfirmed = c.Email, true
						return nil
					})
				if err != nil {
					return "", User{}, err
				}
				u, _ = f.store.User(name)
			}
			return name, u, nil
		}
	}

	name := f.deriveName(provider.Key, c)
	acct := User{
		Origin: originOIDC, Iss: c.Iss, Sub: c.Sub,
		Email: c.Email, EmailConfirmed: true, Notify: true,
	}
	err := f.store.UpdateUsers("users: provision "+name+" from "+provider.Label, Author{Name: name, Email: c.Email},
		func(m map[string]User) error {
			if _, taken := m[name]; taken {
				return fmt.Errorf("%s was taken while provisioning it", name)
			}
			m[name] = acct
			return nil
		})
	if err != nil {
		return "", User{}, err
	}
	log.Printf("oidc: provisioned %s from %s with no grants", name, provider.Issuer)
	f.mailAdmins("", "new account "+name+" is waiting for a grant",
		name+" signed in through "+provider.Label+" for the first time and holds no grants yet.\n\n"+
			f.url("/admin")+"\n")
	return name, acct, nil
}

// deriveName is run once, at first login.
func (f *Forge) deriveName(key string, c *claims) string {
	base := sanitizeName(c.Username)
	if base == "" {
		local, _, _ := strings.Cut(c.Email, "@")
		base = sanitizeName(local)
	}
	if base == "" {
		base = sanitizeName(c.Sub)
	}
	if base == "" {
		base = "user"
	}
	if len(base) > 24 {
		base = base[:24]
	}

	users := f.store.Users()
	name := key + ":" + base
	for n := 2; ; n++ {
		if _, taken := users[name]; !taken {
			return name
		}
		name = key + ":" + base + strconv.Itoa(n)
	}
}

var unsafeName = regexp.MustCompile(`[^a-z0-9._-]+`)

// sanitizeName is blunt: the name becomes a key in repos.json and part of a
// URL, out of a claim from elsewhere.
func sanitizeName(s string) string {
	s = unsafeName.ReplaceAllString(strings.ToLower(s), "-")
	s = strings.Trim(s, "-._")
	if s == "" || !localNameRE.MatchString(s) {
		return ""
	}
	return s
}

func (f *Forge) provider(key string) (OIDCProvider, bool) {
	for _, p := range f.cfg.OIDC {
		if p.Key == key {
			return p, true
		}
	}
	return OIDCProvider{}, false
}

func (f *Forge) oidcRedirect(key string) string { return f.cfg.BaseURL + "/oidc/" + key + "/callback" }

// The login cookie: the provider, the state, the verifier, and an expiry.
func (f *Forge) setOIDCCookie(w http.ResponseWriter, key, state, verifier string) {
	expires := strconv.FormatInt(time.Now().Add(oidcLife).Unix(), 10)
	value := strings.Join([]string{key, state, verifier, expires}, ".")
	http.SetCookie(w, &http.Cookie{
		Name:     oidcCookie,
		Value:    value + "." + f.mac("oidc", key, state, verifier, expires),
		Path:     "/oidc/",
		MaxAge:   int(oidcLife / time.Second),
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode, // the callback arrives from the provider
		Secure:   strings.HasPrefix(f.cfg.BaseURL, "https://"),
	})
}

func (f *Forge) oidcCookie(r *http.Request) (key, state, verifier string, ok bool) {
	c, err := r.Cookie(oidcCookie)
	if err != nil {
		return "", "", "", false
	}
	parts := strings.Split(c.Value, ".")
	if len(parts) != 5 {
		return "", "", "", false
	}
	key, state, verifier, expires, mac := parts[0], parts[1], parts[2], parts[3], parts[4]
	unix, err := strconv.ParseInt(expires, 10, 64)
	if err != nil || time.Now().After(time.Unix(unix, 0)) {
		return "", "", "", false
	}
	if !hmac.Equal([]byte(mac), []byte(f.mac("oidc", key, state, verifier, expires))) {
		return "", "", "", false
	}
	return key, state, verifier, true
}

func (f *Forge) clearOIDCCookie(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{Name: oidcCookie, Path: "/oidc/", MaxAge: -1, HttpOnly: true})
}

// claims is the part of an id token that gets used.
type claims struct {
	Iss           string   `json:"iss"`
	Sub           string   `json:"sub"`
	Aud           audience `json:"aud"`
	Exp           int64    `json:"exp"`
	Iat           int64    `json:"iat"`
	Email         string   `json:"email"`
	EmailVerified *bool    `json:"email_verified"`
	Username      string   `json:"preferred_username"`
}

// audience is one string or several, as providers send it.
type audience []string

func (a *audience) UnmarshalJSON(data []byte) error {
	var one string
	if err := json.Unmarshal(data, &one); err == nil {
		*a = audience{one}
		return nil
	}
	var many []string
	if err := json.Unmarshal(data, &many); err != nil {
		return err
	}
	*a = many
	return nil
}

func (a audience) has(want string) bool {
	for _, s := range a {
		if s == want {
			return true
		}
	}
	return false
}

// jwk is one signing key from the provider's key set.
type jwk 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"`
}

// An exact kid wins over a key naming none, or a rotation would hide behind a
// failed verification.
func findKey(keys []jwk, kid string) (jwk, bool) {
	for _, k := range keys {
		if signingKey(k) && kid != "" && k.Kid == kid {
			return k, true
		}
	}
	for _, k := range keys {
		if signingKey(k) && (kid == "" || k.Kid == "") {
			return k, true
		}
	}
	return jwk{}, false
}

// An absent use is unrestricted.
func signingKey(k jwk) bool { return k.Use == "" || k.Use == "sig" }

// splitJWT takes a token apart without trusting any of it.
func splitJWT(token string) (signed, signature []byte, kid, alg string, payload []byte, err error) {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return nil, nil, "", "", nil, errors.New("id token is not a jwt")
	}
	head, err := base64.RawURLEncoding.DecodeString(parts[0])
	if err != nil {
		return nil, nil, "", "", nil, err
	}
	var header struct {
		Alg string `json:"alg"`
		Kid string `json:"kid"`
	}
	if err := json.Unmarshal(head, &header); err != nil {
		return nil, nil, "", "", nil, err
	}
	payload, err = base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return nil, nil, "", "", nil, err
	}
	signature, err = base64.RawURLEncoding.DecodeString(parts[2])
	if err != nil {
		return nil, nil, "", "", nil, err
	}
	return []byte(parts[0] + "." + parts[1]), signature, header.Kid, header.Alg, payload, nil
}

// RS256 and ES256 only.
func verifySignature(key jwk, alg string, signed, signature []byte) error {
	sum := sha256.Sum256(signed)
	switch {
	case alg == "RS256" && key.Kty == "RSA":
		n, err := bigFromJWK(key.N)
		if err != nil {
			return err
		}
		e, err := base64.RawURLEncoding.DecodeString(key.E)
		if err != nil {
			return err
		}
		pub := &rsa.PublicKey{N: n, E: int(new(big.Int).SetBytes(e).Int64())}
		return rsa.VerifyPKCS1v15(pub, crypto.SHA256, sum[:], signature)

	case alg == "ES256" && key.Kty == "EC" && key.Crv == "P-256":
		x, err := bigFromJWK(key.X)
		if err != nil {
			return err
		}
		y, err := bigFromJWK(key.Y)
		if err != nil {
			return err
		}
		if len(signature) != 64 {
			return errors.New("es256 signature is the wrong size")
		}
		pub := &ecdsa.PublicKey{Curve: elliptic.P256(), X: x, Y: y}
		r := new(big.Int).SetBytes(signature[:32])
		s := new(big.Int).SetBytes(signature[32:])
		if !ecdsa.Verify(pub, sum[:], r, s) {
			return errors.New("id token signature does not verify")
		}
		return nil
	}
	return fmt.Errorf("id token algorithm %q with key type %q is not supported", alg, key.Kty)
}

func bigFromJWK(field string) (*big.Int, error) {
	raw, err := base64.RawURLEncoding.DecodeString(field)
	if err != nil {
		return nil, err
	}
	return new(big.Int).SetBytes(raw), nil
}