llmgw/internal/web/account.go

package web

import (
	"context"
	"crypto/rand"
	"encoding/base32"
	"errors"
	"log/slog"
	"net/http"
	"strconv"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/account"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// SessionCookie is the name of the cookie a signed-in browser carries.
const SessionCookie = "llmgw_session"

// Accounts is what the pages need from the account service.
type Accounts interface {
	SignUp(ctx context.Context, email, name, password string) error
	SignIn(ctx context.Context, email, password string) (token, csrf string, err error)
	Resolve(ctx context.Context, token string) (store.Session, store.User, error)
	SignOut(ctx context.Context, token string) error
	CheckCSRF(sess store.Session, submitted string) bool
	SignupOpen() bool
	ChangePassword(ctx context.Context, pw account.PasswordChanger,
		u store.User, current, next string) error
	OpenSession(ctx context.Context, userID string) (token, csrf string, err error)
}

// KeyOwner is the store half of the profile page.
type KeyOwner interface {
	KeysForUser(ctx context.Context, userID string) ([]store.Key, error)
	CreateUserKey(ctx context.Context, userID string, k store.Key) error
	RevokeUserKey(ctx context.Context, userID, keyID string, at time.Time) error
	Resolve(ctx context.Context, keyID string) (store.Effective, error)
	LimitsFor(ctx context.Context, scope store.Scope, id string) (store.Limits, error)
}

// AccountPages serves signup, sign-in and the profile.
type AccountPages struct {
	accounts  Accounts
	keys      KeyOwner
	passwords account.PasswordChanger
	version   string
	log       *slog.Logger
	// served, when set, names the models an endpoint is currently serving. The
	// curl example on this page fills its model field from it: every other
	// field in that snippet is real, so the one placeholder left in it was
	// copied verbatim and sent as the literal string MODEL.
	served func() []string
}

// ServedModels supplies the names the profile page's example should use.
//
// A setter rather than a constructor argument because only the router knows
// what is being served, and the page is useful without it — with no source the
// example falls back to a placeholder, which is what it always was.
func (p *AccountPages) ServedModels(f func() []string) { p.served = f }

// exampleModel picks the model the copy-pasteable example should name.
//
// A real model over a simulated one is not expressible here — this package does
// not know which is which — so it takes the first served name and relies on the
// router listing real models first.
func (p *AccountPages) exampleModel() string {
	if p.served == nil {
		return "<model-name>"
	}
	if names := p.served(); len(names) > 0 {
		return names[0]
	}
	return "<model-name>"
}

// NewAccountPages builds the handlers.
func NewAccountPages(a Accounts, k KeyOwner, pw account.PasswordChanger,
	version string, log *slog.Logger,
) *AccountPages {
	if log == nil {
		log = slog.Default()
	}
	return &AccountPages{accounts: a, keys: k, passwords: pw, version: version, log: log}
}

// Register mounts the pages. They are all browser surfaces and none of them
// takes an API key, so the caller must exempt these paths from key
// authentication.
func (p *AccountPages) Register(mux *http.ServeMux) {
	if p.accounts.SignupOpen() {
		mux.HandleFunc("GET /signup", p.signupForm)
		mux.HandleFunc("POST /signup", p.signup)
	}
	mux.HandleFunc("GET /login", p.loginForm)
	mux.HandleFunc("POST /login", p.login)
	mux.HandleFunc("POST /logout", p.logout)
	mux.HandleFunc("GET /profile", p.profile)
	mux.HandleFunc("POST /profile/keys", p.createKey)
	mux.HandleFunc("POST /profile/keys/{id}/revoke", p.revokeKey)
	mux.HandleFunc("POST /profile/password", p.changePassword)
}

// --- signup ---

type signupData struct {
	Email, Name, Error string
	Done               bool
	MinPassword        int
	Version            string
}

func (p *AccountPages) signupForm(w http.ResponseWriter, _ *http.Request) {
	noStore(w)
	render(w, "signup", signupData{MinPassword: account.MinPasswordLen, Version: p.version})
}

func (p *AccountPages) signup(w http.ResponseWriter, r *http.Request) {
	noStore(w)
	d := signupData{
		Email:       strings.TrimSpace(r.FormValue("email")),
		Name:        strings.TrimSpace(r.FormValue("name")),
		MinPassword: account.MinPasswordLen,
		Version:     p.version,
	}

	err := p.accounts.SignUp(r.Context(), d.Email, d.Name, r.FormValue("password"))
	switch {
	case err == nil:
		// The same page whether or not the address was already taken. Anything
		// else is an account enumerator that anyone can run.
		d.Done = true
	case errors.Is(err, account.ErrWeakPassword), errors.Is(err, account.ErrInvalidEmail):
		d.Error = err.Error()
	default:
		p.log.Error("signup failed", "err", err)
		d.Error = "something went wrong; try again"
	}
	render(w, "signup", d)
}

// --- sign in and out ---

type loginData struct {
	Email, Error string
	SignupOpen   bool
	Version      string
}

func (p *AccountPages) loginForm(w http.ResponseWriter, _ *http.Request) {
	noStore(w)
	render(w, "login", loginData{SignupOpen: p.accounts.SignupOpen(), Version: p.version})
}

func (p *AccountPages) login(w http.ResponseWriter, r *http.Request) {
	noStore(w)
	email := strings.TrimSpace(r.FormValue("email"))

	token, csrf, err := p.accounts.SignIn(r.Context(), email, r.FormValue("password"))
	if err != nil {
		d := loginData{Email: email, SignupOpen: p.accounts.SignupOpen(), Version: p.version}
		switch {
		case errors.Is(err, account.ErrBadCredentials),
			errors.Is(err, account.ErrNotApproved),
			errors.Is(err, account.ErrSuspended),
			errors.Is(err, account.ErrTooManyTries):
			d.Error = err.Error()
		default:
			p.log.Error("sign-in failed", "err", err)
			d.Error = "something went wrong; try again"
		}
		w.WriteHeader(http.StatusUnauthorized)
		render(w, "login", d)
		return
	}

	setSessionCookie(w, r, token)
	setCSRFCookie(w, r, csrf)
	http.Redirect(w, r, "/profile", http.StatusSeeOther)
}

func (p *AccountPages) logout(w http.ResponseWriter, r *http.Request) {
	noStore(w)
	// Even signing out is a state change, so it needs its token: otherwise any
	// page can sign a visitor out of this one.
	sess, _, ok := p.session(r)
	if ok && !p.accounts.CheckCSRF(sess, r.FormValue("csrf")) {
		http.Error(w, "bad request", http.StatusBadRequest)
		return
	}
	if c, err := r.Cookie(SessionCookie); err == nil {
		if err := p.accounts.SignOut(r.Context(), c.Value); err != nil {
			p.log.Warn("could not close a session", "err", err)
		}
	}
	clearSessionCookie(w, r)
	http.Redirect(w, r, "/login", http.StatusSeeOther)
}

// --- profile ---

type limitRow struct{ Name, Value string }

type profileData struct {
	User   store.User
	Keys   []store.Key
	Limits []limitRow
	NewKey string
	Error  string

	PasswordChanged bool
	PasswordError   string
	MinPassword     int

	CSRF    string
	BaseURL string
	// ExampleModel is a model that is actually being served, so the example on
	// the page can be copied and run as it stands.
	ExampleModel string
	Version      string
	Now          string
}

func (p *AccountPages) profile(w http.ResponseWriter, r *http.Request) {
	noStore(w)
	sess, user, ok := p.session(r)
	if !ok {
		http.Redirect(w, r, "/login", http.StatusSeeOther)
		return
	}
	p.renderProfile(w, r, sess, user, "", "")
}

func (p *AccountPages) createKey(w http.ResponseWriter, r *http.Request) {
	noStore(w)
	sess, user, ok := p.session(r)
	if !ok {
		http.Redirect(w, r, "/login", http.StatusSeeOther)
		return
	}
	if !p.accounts.CheckCSRF(sess, r.FormValue("csrf")) {
		http.Error(w, "bad request", http.StatusBadRequest)
		return
	}

	name := strings.TrimSpace(r.FormValue("name"))
	if name == "" {
		name = "untitled"
	}
	secret, k, err := store.NewKey(randomKeyID(), name, user.ProjectID)
	if err != nil {
		p.log.Error("could not mint a key", "err", err)
		p.renderProfile(w, r, sess, user, "", "something went wrong; try again")
		return
	}
	k.CreatedAt = time.Now()
	if err := p.keys.CreateUserKey(r.Context(), user.ID, k); err != nil {
		p.log.Error("could not store a key", "err", err)
		p.renderProfile(w, r, sess, user, "", "something went wrong; try again")
		return
	}
	p.log.Info("a user created a key", "user", user.ID, "key", k.ID)

	// Shown on this response rather than after a redirect: a redirect would
	// have to carry the key in a URL or a flash, and both put it somewhere it
	// does not belong.
	p.renderProfile(w, r, sess, user, secret, "")
}

func (p *AccountPages) revokeKey(w http.ResponseWriter, r *http.Request) {
	noStore(w)
	sess, user, ok := p.session(r)
	if !ok {
		http.Redirect(w, r, "/login", http.StatusSeeOther)
		return
	}
	if !p.accounts.CheckCSRF(sess, r.FormValue("csrf")) {
		http.Error(w, "bad request", http.StatusBadRequest)
		return
	}

	// Ownership is enforced in the statement, so a guessed key id from another
	// account changes nothing.
	err := p.keys.RevokeUserKey(r.Context(), user.ID, r.PathValue("id"), time.Now())
	if err != nil && !errors.Is(err, store.ErrNotFound) {
		p.log.Error("could not revoke a key", "err", err)
	}
	http.Redirect(w, r, "/profile", http.StatusSeeOther)
}

func (p *AccountPages) renderProfile(w http.ResponseWriter, r *http.Request,
	sess store.Session, user store.User, newKey, errMsg string,
) {
	p.render(w, r, sess, user, profileData{NewKey: newKey, Error: errMsg})
}

func (p *AccountPages) render(w http.ResponseWriter, r *http.Request,
	sess store.Session, user store.User, d profileData,
) {
	d.User = user
	d.MinPassword = account.MinPasswordLen
	d.BaseURL = baseURL(r)
	d.ExampleModel = p.exampleModel()
	d.Version = p.version
	d.Now = time.Now().UTC().Format(time.RFC3339)
	// The form token comes from the cookie, never from whatever the request
	// happened to carry: echoing a submitted value back into the next form
	// would let an attacker choose both halves of the comparison.
	d.CSRF = csrfFor(r)
	_ = sess

	keys, err := p.keys.KeysForUser(r.Context(), user.ID)
	if err != nil {
		p.log.Error("could not list keys", "err", err)
	}
	d.Keys = keys
	d.Limits = p.limitsFor(r.Context(), user)
	render(w, "profile", d)
}

// limitsFor renders the limits a person's keys inherit from their project.
//
// From the project rather than from a key, because a person with no keys yet
// still wants to know what they are getting.
func (p *AccountPages) limitsFor(ctx context.Context, user store.User) []limitRow {
	l, err := p.keys.LimitsFor(ctx, store.ScopeProject, user.ProjectID)
	if err != nil {
		return nil
	}
	var out []limitRow
	if l.RPM != nil {
		out = append(out, limitRow{"Requests per minute", itoa(*l.RPM)})
	}
	if l.TPM != nil {
		out = append(out, limitRow{"Tokens per minute", itoa(*l.TPM)})
	}
	if l.Budget != nil {
		out = append(out, limitRow{"Budget", units(*l.Budget) + " per " + string(l.BudgetPeriod)})
	}
	if l.AllowedModels != nil {
		if len(l.AllowedModels) == 0 {
			out = append(out, limitRow{"Models", "none"})
		} else {
			out = append(out, limitRow{"Models", strings.Join(l.AllowedModels, ", ")})
		}
	}
	return out
}

// --- helpers ---

func (p *AccountPages) session(r *http.Request) (store.Session, store.User, bool) {
	c, err := r.Cookie(SessionCookie)
	if err != nil {
		return store.Session{}, store.User{}, false
	}
	sess, user, err := p.accounts.Resolve(r.Context(), c.Value)
	if err != nil {
		return store.Session{}, store.User{}, false
	}
	return sess, user, true
}

// csrfFor returns the token a form in this session must carry.
//
// The session stores only a hash, so the value itself has to travel with the
// browser. It rides the same cookie jar as the session under its own name,
// which is the double-submit pattern: an attacker's page can make the browser
// send our cookies but cannot read them to put one in a form field.
func csrfFor(r *http.Request) string {
	if c, err := r.Cookie(csrfCookie); err == nil {
		return c.Value
	}
	return ""
}

func setCSRFCookie(w http.ResponseWriter, r *http.Request, csrf string) {
	// HttpOnly as well: the server renders it into forms, so no script needs to
	// read it, and the attack this defends against works by making the browser
	// send cookies it cannot read.
	http.SetCookie(w, &http.Cookie{
		Name: csrfCookie, Value: csrf, Path: "/",
		HttpOnly: true, Secure: isHTTPS(r), SameSite: http.SameSiteLaxMode,
		MaxAge: int(store.SessionTTL.Seconds()),
	})
}

func randomKeyID() string {
	var buf [10]byte
	_, _ = rand.Read(buf[:])
	return "key-" + strings.ToLower(
		base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(buf[:]))
}

func itoa(v int) string { return strconv.Itoa(v) }

func units(micros int64) string {
	return strconv.FormatFloat(float64(micros)/1e6, 'f', 2, 64)
}

const csrfCookie = "llmgw_csrf"

func setSessionCookie(w http.ResponseWriter, r *http.Request, token string) {
	http.SetCookie(w, &http.Cookie{
		Name: SessionCookie, Value: token, Path: "/",
		HttpOnly: true, Secure: isHTTPS(r), SameSite: http.SameSiteLaxMode,
		MaxAge: int(store.SessionTTL.Seconds()),
	})
}

func clearSessionCookie(w http.ResponseWriter, r *http.Request) {
	for _, name := range []string{SessionCookie, csrfCookie} {
		http.SetCookie(w, &http.Cookie{
			Name: name, Value: "", Path: "/",
			HttpOnly: name == SessionCookie, Secure: isHTTPS(r),
			SameSite: http.SameSiteLaxMode, MaxAge: -1,
		})
	}
}

// isHTTPS decides whether to mark cookies Secure.
//
// Behind a reverse proxy this process sees plain HTTP, so the forwarded header
// is the only thing that knows. Getting it wrong in the safe direction — marking
// Secure when the connection is plain — would break sign-in over http://, which
// is why it is not simply hardcoded.
func isHTTPS(r *http.Request) bool {
	if r.TLS != nil {
		return true
	}
	return strings.EqualFold(firstValue(r.Header.Get("X-Forwarded-Proto")), "https")
}

func noStore(w http.ResponseWriter) {
	w.Header().Set("Cache-Control", "no-store, no-cache, must-revalidate, private")
	w.Header().Set("Referrer-Policy", "no-referrer")
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.Header().Set("X-Frame-Options", "DENY")
}

func (p *AccountPages) changePassword(w http.ResponseWriter, r *http.Request) {
	noStore(w)
	sess, user, ok := p.session(r)
	if !ok {
		http.Redirect(w, r, "/login", http.StatusSeeOther)
		return
	}
	if !p.accounts.CheckCSRF(sess, r.FormValue("csrf")) {
		http.Error(w, "bad request", http.StatusBadRequest)
		return
	}

	err := p.accounts.ChangePassword(r.Context(), p.passwords, user,
		r.FormValue("current"), r.FormValue("new"))
	if err != nil {
		d := profileData{}
		switch {
		case errors.Is(err, account.ErrBadCredentials):
			d.PasswordError = "that is not your current password"
		case errors.Is(err, account.ErrWeakPassword),
			errors.Is(err, account.ErrSamePassword),
			errors.Is(err, account.ErrTooManyTries):
			d.PasswordError = err.Error()
		default:
			p.log.Error("could not change a password", "err", err)
			d.PasswordError = "something went wrong; try again"
		}
		p.render(w, r, sess, user, d)
		return
	}

	// The change closed every session including this one, which is the point.
	// Open a new one for the browser in front of us, so the person who just
	// proved their password is not thrown back to the sign-in page.
	token, csrf, err := p.accounts.OpenSession(r.Context(), user.ID)
	if err != nil {
		p.log.Error("could not reopen a session", "err", err)
		clearSessionCookie(w, r)
		http.Redirect(w, r, "/login", http.StatusSeeOther)
		return
	}
	setSessionCookie(w, r, token)
	setCSRFCookie(w, r, csrf)

	// The new CSRF cookie is set on this response, so the forms rendered now
	// must carry the same value rather than the one the request arrived with.
	newSess, newUser, rerr := p.accounts.Resolve(r.Context(), token)
	if rerr != nil {
		http.Redirect(w, r, "/profile", http.StatusSeeOther)
		return
	}
	d := profileData{PasswordChanged: true}
	d.CSRF = csrf
	p.renderWithCSRF(w, r, newSess, newUser, d)
}

// renderWithCSRF is render for the one case where the token in the cookie jar
// is not yet the one the browser has: the response that sets a new one.
func (p *AccountPages) renderWithCSRF(w http.ResponseWriter, r *http.Request,
	sess store.Session, user store.User, d profileData,
) {
	csrf := d.CSRF
	d.User = user
	d.MinPassword = account.MinPasswordLen
	d.BaseURL = baseURL(r)
	d.ExampleModel = p.exampleModel()
	d.Version = p.version
	d.Now = time.Now().UTC().Format(time.RFC3339)
	d.CSRF = csrf
	_ = sess

	keys, err := p.keys.KeysForUser(r.Context(), user.ID)
	if err != nil {
		p.log.Error("could not list keys", "err", err)
	}
	d.Keys = keys
	d.Limits = p.limitsFor(r.Context(), user)
	render(w, "profile", d)
}