llmgw/internal/web/oidc.go

package web

import (
	"context"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"net/http"
	"slices"
	"strings"
	"sync"
	"time"

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

	"gitea.mynodes.xyz/ilker/llmgw/internal/audit"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/secrets"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// OIDCStore is what signing in through a provider needs of the store.
type OIDCStore interface {
	UserByEmail(ctx context.Context, email string) (store.User, error)
	CreateUser(ctx context.Context, u store.User) error
	SetStaffRole(ctx context.Context, id, role string) error
	AddMembership(ctx context.Context, user string, scope store.Scope, id, role, by string, at time.Time) error
	AllOrgs(ctx context.Context) ([]store.Org, error)
	AllTeams(ctx context.Context) ([]store.TeamPath, error)
	AllProjects(ctx context.Context) ([]store.ProjectPath, error)
	// ClusterSecret is a secret every control plane reads alike.
	ClusterSecret(ctx context.Context, name string) (string, error)
}

// oidcLogin is sign-in through an identity provider.
type oidcLogin struct {
	cfg    config.OIDC
	secret string
	st     OIDCStore
	key    []byte // signs the cookie that carries a sign-in across the provider

	mu       sync.Mutex
	provider *oidc.Provider
}

const oidcCookie = "llmgw_oidc"

// RegisterOIDC mounts /login/oidc and its callback, and puts the button on
// the sign-in page.
func (p *AccountPages) RegisterOIDC(mux *http.ServeMux, cfg config.OIDC, st OIDCStore) error {
	secret := ""
	if cfg.ClientSecretFile != "" {
		b, err := secrets.Read(context.Background(), cfg.ClientSecretFile)
		if err != nil {
			return fmt.Errorf("auth.oidc.client_secret_file: %w", err)
		}
		secret = strings.TrimSpace(string(b))
	}
	// From the store, not made here: a sign-in begun on one control plane
	// comes back from the provider to whichever one the balancer picks, and
	// a key of this process's own made that one refuse it — with two behind
	// a balancer, half of all sign-ins.
	key, err := st.ClusterSecret(context.Background(), "oidc-state")
	if err != nil {
		return fmt.Errorf("the key that carries a sign-in across the provider: %w", err)
	}
	o := &oidcLogin{cfg: cfg, secret: secret, st: st, key: []byte(key)}
	p.oidc = o
	mux.HandleFunc("GET /login/oidc", p.oidcStart)
	mux.HandleFunc("GET /login/oidc/callback", p.oidcCallback)
	return nil
}

// providerFor discovers the provider the first time it is needed, and again
// after a failure: a provider down when the gateway started is no reason for
// sign-in never to work.
func (o *oidcLogin) providerFor(ctx context.Context) (*oidc.Provider, error) {
	o.mu.Lock()
	defer o.mu.Unlock()
	if o.provider != nil {
		return o.provider, nil
	}
	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	pr, err := oidc.NewProvider(ctx, o.cfg.Issuer)
	if err != nil {
		return nil, err
	}
	o.provider = pr
	return pr, nil
}

func (o *oidcLogin) oauth(pr *oidc.Provider, r *http.Request) *oauth2.Config {
	redirect := o.cfg.RedirectURL
	if redirect == "" {
		redirect = baseURL(r) + "/login/oidc/callback"
	}
	scopes := []string{oidc.ScopeOpenID}
	if len(o.cfg.Scopes) > 0 {
		scopes = append(scopes, o.cfg.Scopes...)
	} else {
		scopes = append(scopes, "email", "profile")
		if len(o.cfg.Roles) > 0 || len(o.cfg.Memberships) > 0 {
			scopes = append(scopes, "groups")
		}
	}
	return &oauth2.Config{ClientID: o.cfg.ClientID, ClientSecret: o.secret, Endpoint: pr.Endpoint(),
		RedirectURL: redirect, Scopes: scopes}
}

// The cookie: state, nonce and PKCE verifier, and when, signed.
func (o *oidcLogin) seal(parts ...string) string {
	v := strings.Join(parts, ".")
	m := hmac.New(sha256.New, o.key)
	m.Write([]byte(v))
	return v + "." + hex.EncodeToString(m.Sum(nil))
}

func (o *oidcLogin) open(c string) ([]string, bool) {
	i := strings.LastIndex(c, ".")
	if i < 0 {
		return nil, false
	}
	v, sig := c[:i], c[i+1:]
	m := hmac.New(sha256.New, o.key)
	m.Write([]byte(v))
	want := hex.EncodeToString(m.Sum(nil))
	if !hmac.Equal([]byte(sig), []byte(want)) {
		return nil, false
	}
	return strings.Split(v, "."), true
}

func randomToken() string {
	b := make([]byte, 24)
	_, _ = rand.Read(b)
	return base64.RawURLEncoding.EncodeToString(b)
}

func (p *AccountPages) oidcStart(w http.ResponseWriter, r *http.Request) {
	noStore(w)
	o := p.oidc
	pr, err := o.providerFor(r.Context())
	if err != nil {
		p.log.Warn("the identity provider did not answer", "issuer", o.cfg.Issuer, "err", err)
		renderStatus(w, r, http.StatusBadGateway, "login", loginData{SignupOpen: p.accounts.SignupOpen(), Version: p.version,
			CSRF: setLoginToken(w, r), SSO: o.cfg.Label, Error: langOf(r).T("sso.err.no_answer")})
		return
	}
	state, nonce, verifier := randomToken(), randomToken(), oauth2.GenerateVerifier()
	http.SetCookie(w, &http.Cookie{Name: oidcCookie, Value: o.seal(state, nonce, verifier,
		fmt.Sprint(time.Now().Unix())), Path: "/login/oidc", MaxAge: 600, HttpOnly: true,
		Secure: isHTTPS(r), SameSite: http.SameSiteLaxMode})
	http.Redirect(w, r, o.oauth(pr, r).AuthCodeURL(state, oidc.Nonce(nonce), oauth2.S256ChallengeOption(verifier)),
		http.StatusFound)
}

func (p *AccountPages) oidcCallback(w http.ResponseWriter, r *http.Request) {
	noStore(w)
	o := p.oidc
	fail := func(why string, err error) {
		p.log.Warn("an OIDC sign-in was refused", "why", why, "err", err)
		http.SetCookie(w, &http.Cookie{Name: oidcCookie, Path: "/login/oidc", MaxAge: -1})
		l := langOf(r)
		msg := l.T("sso.err.failed")
		var ae accountError
		if errors.As(err, &ae) {
			msg = ae.in(l)
		}
		renderStatus(w, r, http.StatusForbidden, "login", loginData{SignupOpen: p.accounts.SignupOpen(),
			Version: p.version, CSRF: setLoginToken(w, r), SSO: o.cfg.Label, Error: msg})
	}
	c, err := r.Cookie(oidcCookie)
	if err != nil {
		fail("no sign-in under way", err)
		return
	}
	parts, ok := o.open(c.Value)
	if !ok || len(parts) != 4 || r.URL.Query().Get("state") != parts[0] {
		fail("state", nil)
		return
	}
	if e := r.URL.Query().Get("error"); e != "" {
		fail("the provider said "+e, nil)
		return
	}
	pr, err := o.providerFor(r.Context())
	if err != nil {
		fail("provider", err)
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), 15*time.Second)
	defer cancel()
	tok, err := o.oauth(pr, r).Exchange(ctx, r.URL.Query().Get("code"), oauth2.VerifierOption(parts[2]))
	if err != nil {
		fail("code exchange", err)
		return
	}
	raw, _ := tok.Extra("id_token").(string)
	idt, err := pr.Verifier(&oidc.Config{ClientID: o.cfg.ClientID}).Verify(ctx, raw)
	if err != nil {
		fail("id token", err)
		return
	}
	if idt.Nonce != parts[1] {
		fail("nonce", nil)
		return
	}
	var claims map[string]any
	if err := idt.Claims(&claims); err != nil {
		fail("claims", err)
		return
	}
	user, err := p.oidcUser(ctx, claims)
	if err != nil {
		fail("account", err)
		return
	}
	token, csrf, err := p.accounts.OpenSession(ctx, user.ID)
	if err != nil {
		fail("session", err)
		return
	}
	http.SetCookie(w, &http.Cookie{Name: oidcCookie, Path: "/login/oidc", MaxAge: -1})
	setSessionCookie(w, r, token)
	setCSRFCookie(w, r, csrf)
	audit.Write(r.Context(), p.audit, p.log, store.AuditEvent{At: time.Now(), Actor: user.Email, Source: "web",
		Action: "sign in through " + o.cfg.Issuer, Status: http.StatusSeeOther,
		Detail: audit.Redact(map[string]any{"role": user.StaffRole()})})
	http.Redirect(w, r, "/profile", http.StatusSeeOther)
}

// errOIDCAccount is a provider's person the gateway will not let in, with why.
var errOIDCAccount = errors.New("")

// accountError is the sentence for the person, by its key and the address it
// names, so the page can say it in their language; Error is the English, for
// the log.
type accountError struct {
	key  string
	args []any
}

func (e accountError) Error() string     { return english.T(e.key, e.args...) }
func (e accountError) Is(t error) bool   { return t == errOIDCAccount }
func (e accountError) in(l *Lang) string { return l.T(e.key, e.args...) }

func init() {
	registerKeys("sso.err.no_email", "sso.err.unverified", "sso.err.no_account",
		"sso.err.pending", "sso.err.suspended", "sso.err.state")
}

// oidcUser is the account a provider's person signs in as: found by their
// verified email, made when the gateway makes accounts this way, with the
// staff role and the memberships their groups give.
func (p *AccountPages) oidcUser(ctx context.Context, claims map[string]any) (store.User, error) {
	o := p.oidc
	email, _ := claims["email"].(string)
	email = store.NormalizeEmail(email)
	if email == "" {
		return store.User{}, accountError{key: "sso.err.no_email"}
	}
	if v, ok := claims["email_verified"].(bool); ok && !v {
		return store.User{}, accountError{"sso.err.unverified", []any{email}}
	}
	u, err := o.st.UserByEmail(ctx, email)
	switch {
	case errors.Is(err, store.ErrNotFound) && o.cfg.CreateAccounts:
		name, _ := claims["name"].(string)
		if name == "" {
			name = email
		}
		u = store.User{ID: "usr-" + strings.ToLower(randomToken()[:16]), Email: email, Name: name,
			// No password works: this account signs in through the provider.
			PasswordHash: "!oidc", State: store.UserActive, CreatedAt: time.Now()}
		if err := o.st.CreateUser(ctx, u); err != nil {
			return store.User{}, err
		}
		p.log.Info("an account was made for a person the identity provider vouched for", "email", email)
	case errors.Is(err, store.ErrNotFound):
		return store.User{}, accountError{"sso.err.no_account", []any{email}}
	case err != nil:
		return store.User{}, err
	}
	if !u.CanSignIn() {
		switch u.State {
		case store.UserPending:
			return store.User{}, accountError{"sso.err.pending", []any{email}}
		case store.UserSuspended:
			return store.User{}, accountError{"sso.err.suspended", []any{email}}
		}
		return store.User{}, accountError{"sso.err.state", []any{email, string(u.State)}}
	}

	groups := claimStrings(claims[or(o.cfg.GroupsClaim, "groups")])
	if len(o.cfg.Roles) > 0 {
		role := ""
		for _, r := range []string{store.RoleAdmin, store.RoleOperator, store.RoleBilling, store.RoleViewer} {
			if slices.ContainsFunc(o.cfg.Roles[r], func(g string) bool { return slices.Contains(groups, g) }) {
				role = r
				break
			}
		}
		if role != u.StaffRole() {
			if err := o.st.SetStaffRole(ctx, u.ID, role); err != nil {
				return store.User{}, err
			}
			u.IsAdmin, u.Role = role == store.RoleAdmin, role
			if role == store.RoleAdmin {
				u.Role = ""
			}
		}
	}
	for _, g := range groups {
		where, ok := o.cfg.Memberships[g]
		if !ok {
			continue
		}
		scope, id, err := p.oidcPlace(ctx, where)
		if err != nil {
			p.log.Warn("a group maps to a place the tree does not have", "group", g, "place", where, "err", err)
			continue
		}
		if err := o.st.AddMembership(ctx, u.ID, scope, id, store.RoleMember, "oidc", time.Now()); err != nil {
			p.log.Warn("could not add a membership from a group", "group", g, "err", err)
		}
	}
	return u, nil
}

// oidcPlace finds org:NAME, team:ORG/TEAM, project:ORG/TEAM/PROJECT.
func (p *AccountPages) oidcPlace(ctx context.Context, where string) (store.Scope, string, error) {
	kind, path, _ := strings.Cut(where, ":")
	// The store writes paths "Acme / ml / train"; a mapping writes them
	// either way.
	norm := func(s string) string {
		parts := strings.Split(s, "/")
		for i := range parts {
			parts[i] = strings.TrimSpace(parts[i])
		}
		return strings.Join(parts, "/")
	}
	same := func(a string) bool { return strings.EqualFold(norm(a), norm(path)) }
	switch kind {
	case "org":
		orgs, err := p.oidc.st.AllOrgs(ctx)
		for _, o := range orgs {
			if same(o.Name) {
				return store.ScopeOrg, o.ID, nil
			}
		}
		return "", "", errors.Join(err, errors.New("no such organization"))
	case "team":
		teams, err := p.oidc.st.AllTeams(ctx)
		for _, t := range teams {
			if same(t.OrgName + "/" + t.Name) {
				return store.ScopeTeam, t.ID, nil
			}
		}
		return "", "", errors.Join(err, errors.New("no such team"))
	default:
		projects, err := p.oidc.st.AllProjects(ctx)
		for _, pp := range projects {
			if same(pp.Path) {
				return store.ScopeProject, pp.ID, nil
			}
		}
		return "", "", errors.Join(err, errors.New("no such project"))
	}
}

func claimStrings(v any) []string {
	switch t := v.(type) {
	case []any:
		var out []string
		for _, x := range t {
			if s, ok := x.(string); ok {
				out = append(out, s)
			}
		}
		return out
	case string:
		return []string{t}
	}
	return nil
}

func or(s, def string) string {
	if s == "" {
		return def
	}
	return s
}