collab/pkg/auth/auth.go

// Package auth verifies Hanzo IAM JWTs via the published JWKS.
//
// We do NOT share code with team-go/extract because the WS upgrade
// path needs to extract the token from either the Authorization header
// or the offered WebSocket subprotocols (browsers cannot set headers on
// the WS handshake), and we want a tiny self-contained dependency surface.
package auth

import (
	"context"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

// Claims is the subset of IAM token claims the collab service uses.
//
// `owner` is the org slug (Hanzo IAM convention — see HIP-0026).
type Claims struct {
	Sub   string `json:"sub"`
	Email string `json:"email"`
	Owner string `json:"owner"`
	jwt.RegisteredClaims
}

// Verifier resolves JWKS keys and validates tokens. Safe for concurrent use.
type Verifier struct {
	jwksURL string
	client  *http.Client

	mu      sync.RWMutex
	keys    map[string]*rsa.PublicKey
	fetched time.Time
	ttl     time.Duration

	// A token whose kid the cached set lacks forces a fetch, and anyone can
	// send one, so fetches are spaced at least minRefresh apart: a caller
	// inside the window gets the last fetch's outcome instead of a new fetch.
	refreshMu  sync.Mutex
	attempted  time.Time
	lastErr    error
	minRefresh time.Duration
}

// New constructs a Verifier against the IAM JWKS URL.
func New(jwksURL string) *Verifier {
	return &Verifier{
		jwksURL:    jwksURL,
		client:     &http.Client{Timeout: 10 * time.Second},
		keys:       map[string]*rsa.PublicKey{},
		ttl:        10 * time.Minute,
		minRefresh: 30 * time.Second,
	}
}

// Verify parses + validates a raw bearer token, returning the typed claims.
func (v *Verifier) Verify(ctx context.Context, raw string) (*Claims, error) {
	claims := &Claims{}
	tok, err := jwt.ParseWithClaims(raw, claims, func(t *jwt.Token) (any, error) {
		if _, ok := t.Method.(*jwt.SigningMethodRSA); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"])
		}
		kid, _ := t.Header["kid"].(string)
		return v.lookupKey(ctx, kid)
	})
	if err != nil {
		return nil, err
	}
	if !tok.Valid {
		return nil, errors.New("token invalid")
	}
	return claims, nil
}

// ExtractToken pulls a bearer token from the Authorization header or, for a
// browser, whose WebSocket API cannot set headers, from a `bearer.<token>`
// entry in the Sec-WebSocket-Protocol list it offers alongside `yjs`:
//
//	new WebSocket(url, ["yjs", "bearer." + token])
//
// A token is never read from the URL, because the edge writes query strings
// to its access log.
func ExtractToken(r *http.Request) string {
	if h := r.Header.Get("Authorization"); h != "" {
		if strings.HasPrefix(strings.ToLower(h), "bearer ") {
			return strings.TrimSpace(h[7:])
		}
	}
	for _, v := range r.Header.Values("Sec-WebSocket-Protocol") {
		for p := range strings.SplitSeq(v, ",") {
			if tok, ok := strings.CutPrefix(strings.TrimSpace(p), "bearer."); ok {
				return tok
			}
		}
	}
	return ""
}

// lookupKey returns the key for kid, refreshing the JWKS when the cached set
// is stale or lacks kid. A key the last good fetch published still verifies
// when a refresh fails, so an IAM outage does not refuse the tokens IAM has
// already vouched for; a successful refresh that drops a key revokes it.
func (v *Verifier) lookupKey(ctx context.Context, kid string) (*rsa.PublicKey, error) {
	v.mu.RLock()
	k, ok := v.keys[kid]
	fresh := time.Since(v.fetched) < v.ttl
	v.mu.RUnlock()
	if ok && fresh {
		return k, nil
	}

	err := v.refresh(ctx)
	v.mu.RLock()
	defer v.mu.RUnlock()
	if k, ok := v.keys[kid]; ok {
		return k, nil
	}
	if err != nil {
		return nil, err
	}
	return nil, fmt.Errorf("kid %q not found", kid)
}

type jwk struct {
	Kty string `json:"kty"`
	Kid string `json:"kid"`
	N   string `json:"n"`
	E   string `json:"e"`
}

type jwks struct {
	Keys []jwk `json:"keys"`
}

// refresh fetches the JWKS at most once per minRefresh and returns that
// fetch's result to every caller inside the window. The fetch ignores the
// caller's cancellation, so one aborted handshake cannot record a failure
// for everyone; the client timeout still bounds it.
func (v *Verifier) refresh(ctx context.Context) error {
	v.refreshMu.Lock()
	defer v.refreshMu.Unlock()
	if time.Since(v.attempted) < v.minRefresh {
		return v.lastErr
	}
	v.attempted = time.Now()
	v.lastErr = v.fetch(context.WithoutCancel(ctx))
	return v.lastErr
}

func (v *Verifier) fetch(ctx context.Context) error {
	u, err := url.Parse(v.jwksURL)
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
	if err != nil {
		return err
	}
	resp, err := v.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("jwks fetch: %s", resp.Status)
	}
	var doc jwks
	if err := json.NewDecoder(resp.Body).Decode(&doc); err != nil {
		return err
	}
	next := map[string]*rsa.PublicKey{}
	for _, k := range doc.Keys {
		if k.Kty != "RSA" {
			continue
		}
		pk, err := decodeRSA(k.N, k.E)
		if err != nil {
			continue
		}
		next[k.Kid] = pk
	}
	v.mu.Lock()
	v.keys = next
	v.fetched = time.Now()
	v.mu.Unlock()
	return nil
}

func decodeRSA(nB64, eB64 string) (*rsa.PublicKey, error) {
	n, err := base64.RawURLEncoding.DecodeString(nB64)
	if err != nil {
		return nil, err
	}
	e, err := base64.RawURLEncoding.DecodeString(eB64)
	if err != nil {
		return nil, err
	}
	ei := new(big.Int).SetBytes(e).Int64()
	return &rsa.PublicKey{
		N: new(big.Int).SetBytes(n),
		E: int(ei),
	}, nil
}