foundry-core/internal/auth/auth.go

package auth

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"time"
)

var projectName = regexp.MustCompile(`^[a-z][a-z0-9-]*$`)
var tokenPattern = regexp.MustCompile(`^[A-Za-z0-9_-]{43}$`)

type service struct {
	db                            *sql.DB
	projects                      map[string]*url.URL
	redirectURI, clientID, secret string
	google                        *http.Client
}

func New(db *sql.DB, callbacks map[string]string, redirectURI, id, secret string) (http.Handler, error) {
	s := &service{db: db, projects: make(map[string]*url.URL), redirectURI: redirectURI, clientID: id, secret: secret,
		google: &http.Client{Timeout: 10 * time.Second, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}}
	for name, raw := range callbacks {
		u, err := url.Parse(raw)
		if err != nil {
			return nil, fmt.Errorf("invalid project callback for %q", name)
		}
		host := u.Hostname()
		local := u.Scheme == "http" && (host == "localhost" || host == "127.0.0.1" || host == "::1")
		if !projectName.MatchString(name) || host == "" || (u.Scheme != "https" && !local) || u.User != nil || u.Fragment != "" {
			return nil, fmt.Errorf("invalid project callback for %q", name)
		}
		s.projects[name] = u
	}
	mux := http.NewServeMux()
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) { reply(w, 200, payload{"status": "ok"}) })
	mux.HandleFunc("GET /readyz", s.wrap(func(w http.ResponseWriter, r *http.Request) error {
		if _, err := s.db.ExecContext(r.Context(), "SELECT 1 FROM project_accounts LIMIT 1"); err != nil {
			return err
		}
		reply(w, 200, payload{"status": "ready"})
		return nil
	}))
	mux.HandleFunc("GET /auth/google/callback", s.wrap(s.callback))
	mux.HandleFunc("/projects/{project}/auth/{action}", s.wrap(s.project))
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Cache-Control", "no-store")
		w.Header().Set("Referrer-Policy", "no-referrer")
		ctx, cancel := context.WithTimeout(r.Context(), 25*time.Second)
		defer cancel()
		mux.ServeHTTP(w, r.WithContext(ctx))
	}), nil
}

type payload map[string]any

func reply(w http.ResponseWriter, status int, body payload) {
	if body != nil {
		w.Header().Set("Content-Type", "application/json")
	}
	w.WriteHeader(status)
	if body != nil {
		_ = json.NewEncoder(w).Encode(body)
	}
}

func reject(w http.ResponseWriter, status int, code string) error {
	reply(w, status, payload{"error": code})
	return nil
}

func (s *service) wrap(handle func(http.ResponseWriter, *http.Request) error) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if err := handle(w, r); err != nil {
			log.Printf("request failed (%T)", err)
			_ = reject(w, 503, "service_unavailable")
		}
	}
}

func token() string {
	var value [32]byte
	rand.Read(value[:])
	return base64.RawURLEncoding.EncodeToString(value[:])
}

func digest(value string) string {
	sum := sha256.Sum256([]byte(value))
	return base64.RawURLEncoding.EncodeToString(sum[:])
}

func (s *service) project(w http.ResponseWriter, r *http.Request) error {
	project, action := r.PathValue("project"), r.PathValue("action")
	frontend := s.projects[project]
	if frontend == nil {
		return reject(w, 404, "not_found")
	}
	if r.Method == http.MethodGet && action == "google" {
		return s.start(w, r, project)
	}

	w.Header().Set("Vary", "Origin")
	if origin := r.Header.Get("Origin"); origin != "" {
		if origin != frontend.Scheme+"://"+frontend.Host {
			return reject(w, 403, "origin_denied")
		}
		w.Header().Set("Access-Control-Allow-Origin", origin)
	}
	switch {
	case r.Method == http.MethodOptions:
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST")
		w.Header().Set("Access-Control-Allow-Headers", "Authorization, X-Foundry-Code, X-Foundry-Verifier")
		reply(w, 204, nil)
	case r.Method == http.MethodPost && action == "exchange":
		return s.exchange(w, r, project)
	case r.Method == http.MethodGet && action == "me":
		return s.account(w, r, project, false)
	case r.Method == http.MethodPost && action == "logout":
		return s.account(w, r, project, true)
	default:
		return reject(w, 405, "method_not_allowed")
	}
	return nil
}

func oauthCookie(state, value string, maxAge int) *http.Cookie {
	return &http.Cookie{Name: "__Host-foundry-oauth-" + state, Value: value, Path: "/", MaxAge: maxAge,
		Secure: true, HttpOnly: true, SameSite: http.SameSiteLaxMode}
}

func (s *service) start(w http.ResponseWriter, r *http.Request, project string) error {
	challenge := r.URL.Query().Get("challenge")
	if !tokenPattern.MatchString(challenge) {
		return reject(w, 400, "invalid_challenge")
	}
	state, browser, verifier := token(), token(), token()
	_, err := s.db.ExecContext(r.Context(), `
  INSERT INTO oauth_attempts VALUES ($1, $2, $3, $4, $5, now() + interval '10 minutes')
 `, digest(state), digest(browser), project, challenge, verifier)
	if err != nil {
		return err
	}
	query := url.Values{
		"client_id": {s.clientID}, "redirect_uri": {s.redirectURI}, "response_type": {"code"},
		"scope": {"openid email"}, "state": {state}, "code_challenge": {digest(verifier)},
		"code_challenge_method": {"S256"}, "prompt": {"select_account"},
	}
	http.SetCookie(w, oauthCookie(state, browser, 600))
	http.Redirect(w, r, "https://accounts.google.com/o/oauth2/v2/auth?"+query.Encode(), 302)
	return nil
}

func (s *service) googleJSON(ctx context.Context, method, endpoint string, form url.Values, bearer string, result any) error {
	var body io.Reader
	if form != nil {
		body = strings.NewReader(form.Encode())
	}
	req, err := http.NewRequestWithContext(ctx, method, endpoint, body)
	if err != nil {
		return err
	}
	if form != nil {
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	}
	if bearer != "" {
		req.Header.Set("Authorization", "Bearer "+bearer)
	}
	res, err := s.google.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		return fmt.Errorf("Google returned HTTP %d", res.StatusCode)
	}
	return json.NewDecoder(io.LimitReader(res.Body, 1<<20)).Decode(result)
}

func (s *service) callback(w http.ResponseWriter, r *http.Request) error {
	ctx := r.Context()
	query := r.URL.Query()
	state := query.Get("state")
	if !tokenPattern.MatchString(state) {
		return reject(w, 400, "invalid_state")
	}
	browser, err := r.Cookie(oauthCookie(state, "", 0).Name)
	if err != nil || !tokenPattern.MatchString(browser.Value) {
		return reject(w, 400, "invalid_state")
	}
	var project, challenge, verifier string
	err = s.db.QueryRowContext(ctx, `
  DELETE FROM oauth_attempts
  WHERE state_hash = $1 AND browser_hash = $2 AND expires_at > now()
  RETURNING project, frontend_challenge, google_verifier
 `, digest(state), digest(browser.Value)).Scan(&project, &challenge, &verifier)
	if errors.Is(err, sql.ErrNoRows) {
		return reject(w, 400, "invalid_state")
	}
	if err != nil {
		return err
	}
	frontend := s.projects[project]
	if frontend == nil {
		return reject(w, 400, "invalid_state")
	}

	http.SetCookie(w, oauthCookie(state, "", -1))
	destination := *frontend
	code := query.Get("code")
	if code == "" || query.Has("error") {
		destination.Fragment = "error=google_login_cancelled"
		http.Redirect(w, r, destination.String(), 302)
		return nil
	}
	var tokens struct {
		AccessToken string `json:"access_token"`
	}
	err = s.googleJSON(ctx, "POST", "https://oauth2.googleapis.com/token", url.Values{
		"code": {code}, "client_id": {s.clientID}, "client_secret": {s.secret}, "redirect_uri": {s.redirectURI},
		"grant_type": {"authorization_code"}, "code_verifier": {verifier},
	}, "", &tokens)
	if err != nil {
		return err
	}
	if tokens.AccessToken == "" {
		return errors.New("missing Google access token")
	}
	var user struct {
		Subject  string  `json:"sub"`
		Email    *string `json:"email"`
		Verified bool    `json:"email_verified"`
	}
	if err := s.googleJSON(ctx, "GET", "https://openidconnect.googleapis.com/v1/userinfo", nil, tokens.AccessToken, &user); err != nil {
		return err
	}
	if user.Subject == "" {
		return errors.New("missing Google subject")
	}
	if !user.Verified {
		user.Email = nil
	}

	grant := token()
	_, err = s.db.ExecContext(ctx, `
  WITH account AS (
   INSERT INTO project_accounts (project, id, google_subject, email) VALUES ($1, gen_random_uuid(), $2, $3)
   ON CONFLICT (project, google_subject) DO UPDATE SET email = EXCLUDED.email
   RETURNING id
  )
  INSERT INTO login_grants
  SELECT $4, $1, id, $5, now() + interval '1 minute' FROM account
 `, project, user.Subject, user.Email, digest(grant), challenge)
	if err != nil {
		return err
	}
	destination.Fragment = "code=" + grant
	http.Redirect(w, r, destination.String(), 302)
	return nil
}

func (s *service) exchange(w http.ResponseWriter, r *http.Request, project string) error {
	code, verifier := r.Header.Get("X-Foundry-Code"), r.Header.Get("X-Foundry-Verifier")
	if !tokenPattern.MatchString(code) || !tokenPattern.MatchString(verifier) {
		return reject(w, 400, "invalid_exchange")
	}
	session := token()
	result, err := s.db.ExecContext(r.Context(), `
  WITH grant AS (
   DELETE FROM login_grants
   WHERE token_hash = $1 AND project = $2 AND frontend_challenge = $3 AND expires_at > now()
   RETURNING account_id
  )
  INSERT INTO project_sessions
  SELECT $4, $2, account_id, now() + interval '8 hours' FROM grant
 `, digest(code), project, digest(verifier), digest(session))
	if err != nil {
		return err
	}
	count, err := result.RowsAffected()
	if err != nil {
		return err
	}
	if count == 0 {
		return reject(w, 401, "invalid_grant")
	}
	reply(w, 200, payload{"access_token": session, "token_type": "Bearer", "expires_in": 28800})
	return nil
}

func (s *service) account(w http.ResponseWriter, r *http.Request, project string, logout bool) error {
	token, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ")
	if !ok || !tokenPattern.MatchString(token) {
		return reject(w, 401, "invalid_session")
	}
	if logout {
		_, err := s.db.ExecContext(r.Context(), "DELETE FROM project_sessions WHERE token_hash = $1 AND project = $2", digest(token), project)
		if err != nil {
			return err
		}
		reply(w, 204, nil)
		return nil
	}
	var id string
	var email *string
	err := s.db.QueryRowContext(r.Context(), `
  SELECT a.id::text, a.email FROM project_sessions s
  JOIN project_accounts a ON a.project = s.project AND a.id = s.account_id
  WHERE s.token_hash = $1 AND s.project = $2 AND s.expires_at > now()
 `, digest(token), project).Scan(&id, &email)
	if errors.Is(err, sql.ErrNoRows) {
		return reject(w, 401, "invalid_session")
	}
	if err != nil {
		return err
	}
	reply(w, 200, payload{"project": project, "id": id, "email": email})
	return nil
}