bootnode/api-go/internal/iam/jwks.go

// Package iam — production JWT verification against Hanzo IAM's JWKS.
//
// HIP-0111: IAM-issued access tokens are validated against the issuer's
// published JWKS (RS256/ES256), resolved via OIDC discovery. No shared HMAC
// secret, no introspection round-trip per request — the JWKS public keys are
// fetched once, cached, and refreshed on an unknown `kid`.
package iam

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rsa"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"net/http"
	"sync"
	"time"

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

// ecCurve maps a JWK `crv` value to its elliptic.Curve.
func ecCurve(crv string) (elliptic.Curve, error) {
	switch crv {
	case "P-256":
		return elliptic.P256(), nil
	case "P-384":
		return elliptic.P384(), nil
	case "P-521":
		return elliptic.P521(), nil
	default:
		return nil, fmt.Errorf("unsupported EC curve %q", crv)
	}
}

// JWKSVerifier validates bearer tokens against an IAM deployment's JWKS.
//
// The issuer is the canonical OIDC issuer for the brand this deployment serves
// (iam.hanzo.ai / lux.id / zoo.id / pars.id / id.bootno.de). The jwks_uri is
// read from `{issuer}/.well-known/openid-configuration`; keys are cached and
// re-fetched at most once per minTTL when a token presents an unknown `kid`.
type JWKSVerifier struct {
	issuer    string
	audience  string // expected `aud` (clientId); "" skips the audience check
	clockSkew time.Duration
	http      *http.Client

	mu         sync.RWMutex
	keys       map[string]any // kid -> *rsa.PublicKey | *ecdsa.PublicKey
	fetched    time.Time
	minRefresh time.Duration
}

// NewJWKS constructs a JWKS verifier for the given OIDC issuer. When audience is
// non-empty (the deployment's IAM client ID), tokens whose `aud` does not match
// are rejected — so a token minted for a different client of the same issuer is
// not accepted. Empty audience skips the check (issuer + signature only).
func NewJWKS(issuer, audience string) *JWKSVerifier {
	return &JWKSVerifier{
		issuer:     issuer,
		audience:   audience,
		clockSkew:  60 * time.Second,
		http:       &http.Client{Timeout: 10 * time.Second},
		keys:       map[string]any{},
		minRefresh: 1 * time.Minute,
	}
}

// Verify implements Verifier. It parses the token, resolves the signing key by
// `kid` against the cached JWKS (refreshing on a miss), validates the signature
// and standard claims, and checks the issuer.
func (j *JWKSVerifier) Verify(ctx context.Context, token string) (*Claims, error) {
	claims := &Claims{}
	parsed, err := jwt.ParseWithClaims(token, claims, func(t *jwt.Token) (any, error) {
		switch t.Method.(type) {
		case *jwt.SigningMethodRSA, *jwt.SigningMethodECDSA:
		default:
			return nil, ErrTokenInvalid
		}
		kid, _ := t.Header["kid"].(string)
		key, err := j.keyByID(ctx, kid)
		if err != nil {
			return nil, err
		}
		return key, nil
	},
		opts(j)...,
	)
	if err != nil {
		switch {
		case errors.Is(err, jwt.ErrTokenExpired):
			return nil, ErrTokenExpired
		default:
			return nil, ErrTokenInvalid
		}
	}
	if !parsed.Valid {
		return nil, ErrTokenInvalid
	}
	return claims, nil
}

// opts builds the jwt parser options: leeway + issuer always, audience only
// when configured.
func opts(j *JWKSVerifier) []jwt.ParserOption {
	o := []jwt.ParserOption{
		jwt.WithLeeway(j.clockSkew),
		jwt.WithIssuer(j.issuer),
	}
	if j.audience != "" {
		o = append(o, jwt.WithAudience(j.audience))
	}
	return o
}

// keyByID returns the public key for kid, refreshing the JWKS once (subject to
// minRefresh) if the kid is not cached.
func (j *JWKSVerifier) keyByID(ctx context.Context, kid string) (any, error) {
	j.mu.RLock()
	key, ok := j.keys[kid]
	j.mu.RUnlock()
	if ok {
		return key, nil
	}
	if err := j.refresh(ctx); err != nil {
		return nil, err
	}
	j.mu.RLock()
	key, ok = j.keys[kid]
	j.mu.RUnlock()
	if !ok {
		return nil, ErrTokenInvalid
	}
	return key, nil
}

// refresh fetches the JWKS via discovery and replaces the key cache. Rate
// limited to one fetch per minRefresh so a flood of bad-kid tokens cannot
// hammer the IAM server.
func (j *JWKSVerifier) refresh(ctx context.Context) error {
	j.mu.Lock()
	defer j.mu.Unlock()
	if time.Since(j.fetched) < j.minRefresh && len(j.keys) > 0 {
		return nil
	}

	jwksURI, err := j.discoverJWKSURI(ctx)
	if err != nil {
		return err
	}
	keys, err := j.fetchKeys(ctx, jwksURI)
	if err != nil {
		return err
	}
	j.keys = keys
	j.fetched = time.Now()
	return nil
}

type discoveryDoc struct {
	Issuer  string `json:"issuer"`
	JWKSURI string `json:"jwks_uri"`
}

func (j *JWKSVerifier) discoverJWKSURI(ctx context.Context) (string, error) {
	url := j.issuer + "/.well-known/openid-configuration"
	var doc discoveryDoc
	if err := j.getJSON(ctx, url, &doc); err != nil {
		return "", fmt.Errorf("iam discovery: %w", err)
	}
	if doc.JWKSURI == "" {
		return "", errors.New("iam discovery: empty jwks_uri")
	}
	if doc.Issuer != "" && doc.Issuer != j.issuer {
		return "", fmt.Errorf("iam discovery: issuer mismatch %q != %q", doc.Issuer, j.issuer)
	}
	return doc.JWKSURI, nil
}

// jwk is one key from a JWKS document (RSA or EC public key, RFC 7517/7518).
type jwk struct {
	Kty string `json:"kty"`
	Kid string `json:"kid"`
	Use string `json:"use"`
	Alg string `json:"alg"`
	// RSA
	N string `json:"n"`
	E string `json:"e"`
	// EC
	Crv string `json:"crv"`
	X   string `json:"x"`
	Y   string `json:"y"`
}

func (j *JWKSVerifier) fetchKeys(ctx context.Context, jwksURI string) (map[string]any, error) {
	var doc struct {
		Keys []jwk `json:"keys"`
	}
	if err := j.getJSON(ctx, jwksURI, &doc); err != nil {
		return nil, fmt.Errorf("iam jwks: %w", err)
	}
	out := make(map[string]any, len(doc.Keys))
	for _, k := range doc.Keys {
		if k.Use != "" && k.Use != "sig" {
			continue
		}
		pub, err := k.publicKey()
		if err != nil {
			continue // skip keys we can't parse (e.g. MLDSA65) — RS/ES suffice
		}
		if k.Kid != "" {
			out[k.Kid] = pub
		}
	}
	if len(out) == 0 {
		return nil, errors.New("iam jwks: no usable signing keys")
	}
	return out, nil
}

func (k jwk) publicKey() (any, error) {
	switch k.Kty {
	case "RSA":
		nb, err := base64.RawURLEncoding.DecodeString(k.N)
		if err != nil {
			return nil, err
		}
		eb, err := base64.RawURLEncoding.DecodeString(k.E)
		if err != nil {
			return nil, err
		}
		e := 0
		for _, b := range eb {
			e = e<<8 | int(b)
		}
		if e == 0 {
			e = int(binary.BigEndian.Uint32(leftPad(eb, 4)))
		}
		return &rsa.PublicKey{N: new(big.Int).SetBytes(nb), E: e}, nil
	case "EC":
		crv, err := ecCurve(k.Crv)
		if err != nil {
			return nil, err
		}
		xb, err := base64.RawURLEncoding.DecodeString(k.X)
		if err != nil {
			return nil, err
		}
		yb, err := base64.RawURLEncoding.DecodeString(k.Y)
		if err != nil {
			return nil, err
		}
		return &ecdsa.PublicKey{
			Curve: crv,
			X:     new(big.Int).SetBytes(xb),
			Y:     new(big.Int).SetBytes(yb),
		}, nil
	default:
		return nil, fmt.Errorf("unsupported kty %q", k.Kty)
	}
}

func (j *JWKSVerifier) getJSON(ctx context.Context, url string, out any) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return err
	}
	req.Header.Set("Accept", "application/json")
	resp, err := j.http.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("GET %s: status %d", url, resp.StatusCode)
	}
	return json.NewDecoder(resp.Body).Decode(out)
}

func leftPad(b []byte, size int) []byte {
	if len(b) >= size {
		return b
	}
	out := make([]byte, size)
	copy(out[size-len(b):], b)
	return out
}