dns/plugin/hanzodns/oidc.go

package hanzodns

import (
	"context"
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math/big"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"
)

type contextKey string

const (
	orgIDKey  contextKey = "org_id"
	bearerKey contextKey = "bearer"
)

// OrgIDFromContext returns the caller's org (the IAM `owner` claim) extracted
// from the validated JWT. It is the tenant key used to scope every per-org
// operation, including the KMS secret read.
func OrgIDFromContext(ctx context.Context) string {
	v, _ := ctx.Value(orgIDKey).(string)
	return v
}

// BearerFromContext returns the caller's raw, validated bearer token so a
// downstream org-scoped call (the KMS secret read) can be authorized as the
// SAME principal. hanzodns holds no standing cross-tenant credential of its own:
// it relays the caller's authority and lets the target enforce org isolation.
func BearerFromContext(ctx context.Context) string {
	v, _ := ctx.Value(bearerKey).(string)
	return v
}

// oidcMiddleware validates OIDC JWT tokens when HANZO_DNS_OIDC_ISSUER is set.
func oidcMiddleware(next http.Handler) http.Handler {
	issuer := os.Getenv("HANZO_DNS_OIDC_ISSUER")
	if issuer == "" {
		return next
	}

	audience := os.Getenv("HANZO_DNS_OIDC_AUDIENCE")
	cache := &jwksCache{issuer: issuer}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		auth := r.Header.Get("Authorization")
		if auth == "" {
			writeError(w, http.StatusUnauthorized, "unauthorized", "missing Authorization header")
			return
		}

		parts := strings.SplitN(auth, " ", 2)
		if len(parts) != 2 || !strings.EqualFold(parts[0], "bearer") {
			writeError(w, http.StatusUnauthorized, "unauthorized", "invalid Authorization header")
			return
		}

		keys, err := cache.getKeys()
		if err != nil {
			writeError(w, http.StatusInternalServerError, "internal_error", "failed to fetch JWKS")
			return
		}

		claims, err := validateToken(parts[1], keys, issuer, audience)
		if err != nil {
			writeError(w, http.StatusUnauthorized, "unauthorized", err.Error())
			return
		}

		ctx := context.WithValue(r.Context(), orgIDKey, claims.org())
		ctx = context.WithValue(ctx, bearerKey, parts[1])
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

type jwtClaims struct {
	Issuer   string        `json:"iss"`
	Subject  string        `json:"sub"`
	Audience audienceClaim `json:"aud"`
	Expiry   int64         `json:"exp"`
	IssuedAt int64         `json:"iat"`
	// Owner is the canonical Hanzo IAM org claim (the org slug). It is the same
	// value cloud/gateway derive tenancy from, so using it here keeps the KMS
	// org path in lockstep with KMS's org-scope guard.
	Owner string `json:"owner"`
	OrgID string `json:"org_id"`
}

// org returns the caller's org, preferring the canonical `owner` claim and
// falling back to `org_id` for tokens that only carry the legacy field.
func (c *jwtClaims) org() string {
	if c.Owner != "" {
		return c.Owner
	}
	return c.OrgID
}

type jwk struct {
	Kty string `json:"kty"`
	Kid string `json:"kid"`
	Alg string `json:"alg"`
	Use string `json:"use"`
	N   string `json:"n"`
	E   string `json:"e"`
	Crv string `json:"crv"`
	X   string `json:"x"`
	Y   string `json:"y"`
}

func (k jwk) toPublicKey() (crypto.PublicKey, error) {
	switch k.Kty {
	case "RSA":
		nb, err := base64URLDecode(k.N)
		if err != nil {
			return nil, err
		}
		eb, err := base64URLDecode(k.E)
		if err != nil {
			return nil, err
		}
		return &rsa.PublicKey{
			N: new(big.Int).SetBytes(nb),
			E: int(new(big.Int).SetBytes(eb).Int64()),
		}, nil
	case "EC":
		xb, err := base64URLDecode(k.X)
		if err != nil {
			return nil, err
		}
		yb, err := base64URLDecode(k.Y)
		if err != nil {
			return nil, err
		}
		var curve elliptic.Curve
		switch k.Crv {
		case "P-256":
			curve = elliptic.P256()
		case "P-384":
			curve = elliptic.P384()
		default:
			return nil, fmt.Errorf("unsupported curve: %s", k.Crv)
		}
		return &ecdsa.PublicKey{Curve: curve, X: new(big.Int).SetBytes(xb), Y: new(big.Int).SetBytes(yb)}, nil
	default:
		return nil, fmt.Errorf("unsupported key type: %s", k.Kty)
	}
}

type jwksCache struct {
	issuer    string
	mu        sync.RWMutex
	keys      []jwk
	fetchedAt time.Time
}

func (c *jwksCache) getKeys() ([]jwk, error) {
	c.mu.RLock()
	if time.Since(c.fetchedAt) < time.Hour && len(c.keys) > 0 {
		keys := c.keys
		c.mu.RUnlock()
		return keys, nil
	}
	c.mu.RUnlock()

	c.mu.Lock()
	defer c.mu.Unlock()

	if time.Since(c.fetchedAt) < time.Hour && len(c.keys) > 0 {
		return c.keys, nil
	}

	keys, err := fetchJWKS(c.issuer)
	if err != nil {
		if len(c.keys) > 0 {
			return c.keys, nil
		}
		return nil, err
	}
	c.keys = keys
	c.fetchedAt = time.Now()
	return keys, nil
}

func fetchJWKS(issuerURL string) ([]jwk, error) {
	issuerURL = strings.TrimSuffix(issuerURL, "/")
	client := &http.Client{Timeout: 10 * time.Second}

	resp, err := client.Get(issuerURL + "/.well-known/openid-configuration")
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	var config struct {
		JWKSURI string `json:"jwks_uri"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&config); err != nil {
		return nil, err
	}
	if config.JWKSURI == "" {
		return nil, fmt.Errorf("no jwks_uri in oidc config")
	}

	resp2, err := client.Get(config.JWKSURI)
	if err != nil {
		return nil, err
	}
	defer resp2.Body.Close()

	var jwks struct {
		Keys []jwk `json:"keys"`
	}
	if err := json.NewDecoder(resp2.Body).Decode(&jwks); err != nil {
		return nil, err
	}
	return jwks.Keys, nil
}

func validateToken(tokenStr string, keys []jwk, issuer, audience string) (*jwtClaims, error) {
	parts := strings.Split(tokenStr, ".")
	if len(parts) != 3 {
		return nil, fmt.Errorf("invalid token format")
	}

	headerBytes, err := base64URLDecode(parts[0])
	if err != nil {
		return nil, fmt.Errorf("decode header: %w", err)
	}
	var header struct {
		Alg string `json:"alg"`
		Kid string `json:"kid"`
	}
	if err := json.Unmarshal(headerBytes, &header); err != nil {
		return nil, err
	}

	var matchedKey *jwk
	for _, k := range keys {
		if k.Kid == header.Kid || (header.Kid == "" && k.Use == "sig") {
			matchedKey = &k
			break
		}
	}
	if matchedKey == nil {
		return nil, fmt.Errorf("no matching key for kid %q", header.Kid)
	}

	signingInput := parts[0] + "." + parts[1]
	signature, err := base64URLDecode(parts[2])
	if err != nil {
		return nil, fmt.Errorf("decode signature: %w", err)
	}

	pubKey, err := matchedKey.toPublicKey()
	if err != nil {
		return nil, err
	}

	hash := sha256.Sum256([]byte(signingInput))
	switch key := pubKey.(type) {
	case *rsa.PublicKey:
		if err := rsa.VerifyPKCS1v15(key, crypto.SHA256, hash[:], signature); err != nil {
			return nil, fmt.Errorf("invalid signature")
		}
	case *ecdsa.PublicKey:
		if !ecdsa.VerifyASN1(key, hash[:], signature) {
			return nil, fmt.Errorf("invalid signature")
		}
	default:
		return nil, fmt.Errorf("unsupported key type")
	}

	claimsBytes, err := base64URLDecode(parts[1])
	if err != nil {
		return nil, err
	}
	var c jwtClaims
	if err := json.Unmarshal(claimsBytes, &c); err != nil {
		return nil, err
	}

	// A token with no expiry is rejected: an unbounded token that leaks never
	// stops being valid.
	if c.Expiry == 0 {
		return nil, fmt.Errorf("token has no expiry")
	}
	if time.Now().Unix() >= c.Expiry {
		return nil, fmt.Errorf("token expired")
	}
	if c.Issuer != issuer {
		return nil, fmt.Errorf("invalid issuer: got %q, want %q", c.Issuer, issuer)
	}
	// When an audience allowlist is configured, the token must carry at least one
	// of the allowed audiences — so a token minted for another service is refused.
	if allow := splitAudience(audience); len(allow) > 0 && !c.Audience.matchesAny(allow) {
		return nil, fmt.Errorf("invalid audience")
	}

	return &c, nil
}

// audienceClaim is a JWT `aud`, which per RFC 7519 may be a single string or an
// array of strings.
type audienceClaim []string

func (a *audienceClaim) UnmarshalJSON(b []byte) error {
	var s string
	if err := json.Unmarshal(b, &s); err == nil {
		*a = audienceClaim{s}
		return nil
	}
	var arr []string
	if err := json.Unmarshal(b, &arr); err != nil {
		return err
	}
	*a = arr
	return nil
}

// matchesAny reports whether the token carries any of the allowed audiences.
func (a audienceClaim) matchesAny(allow []string) bool {
	for _, want := range allow {
		for _, got := range a {
			if got == want {
				return true
			}
		}
	}
	return false
}

// splitAudience parses a comma-separated audience allowlist, trimming blanks.
func splitAudience(s string) []string {
	var out []string
	for _, p := range strings.Split(s, ",") {
		if p = strings.TrimSpace(p); p != "" {
			out = append(out, p)
		}
	}
	return out
}

func base64URLDecode(s string) ([]byte, error) {
	switch len(s) % 4 {
	case 2:
		s += "=="
	case 3:
		s += "="
	}
	return base64.URLEncoding.DecodeString(s)
}