jwe: WithKey validates alg-vs-key shape at option-time (#2120)

Mirror jws.validateAlgorithmForKey for the JWE side. Mismatched
(alg, key) pairs that cannot succeed downstream (RSA-OAEP + []byte,
A128KW + RSA key, ECDH-ES + bytes, etc.) now produce a crisp
`jwe.WithKey: ...` error at option-time instead of a deep nested
error from requireByteKey or keyconv inside the dispatcher.

Permissive carve-outs (defer to dispatch-time validation):
- jwk.Key wrappers — kty-vs-alg check happens at jwk.Export time
- caller-supplied KeyEncrypter / KeyDecrypter implementations
  (HSM adapters, custom-decrypter test doubles) — caller owns the
  shape contract
- HPKEKeyEncrypter / HPKEKeyDecrypter raw types (extension-pluggable)
- HPKE / ML-KEM algorithm families (extension-registered, key shape
  is the extension's contract)

Bundled doc improvement: KeyDecrypter interface godoc now spells out
the library contract (what's already validated, what the implementer
owns, what nil-vs-error means, what constant-time considerations
remain the implementer's responsibility).

Regression test covers reject paths (RSA-OAEP + bytes, A128KW + RSA,
ECDH-ES + bytes) and accept paths (jwk.Key wrapper, AES-KW + bytes).

Author: lestrrat

jwe/encrypt.go

@@ -1,6 +1,8 @@
 package jwe
 
 import (
+	"crypto/ecdh"
+	"crypto/ecdsa"
 	"crypto/rsa"
 	"fmt"
 
@@ -101,6 +103,88 @@ func requireByteKey(key any, alg string) ([]byte, error) {
 	return b, nil
 }
 
+// validateAlgorithmForKey checks that alg is family-compatible with
+// key at the WithKey option boundary, surfacing wrong-shape mismatches
+// as crisp `jwe.WithKey: ...` errors instead of nested errors deep in
+// the dispatcher (e.g. requireByteKey inside the AESKW path). Mirrors
+// jws.validateAlgorithmForKey in spirit but is shaped to JWE's
+// per-family key conventions.
+//
+// Permissive carve-outs (return nil, deferring validation):
+//
+//   - Untyped/extension algorithms: HPKE and ML-KEM. These are
+//     extension-pluggable via Register{HPKE,MLKEM}Algorithm; an
+//     extension may accept arbitrary key shapes (e.g. an
+//     HPKEKeyEncrypter raw type), so the option-time gate cannot
+//     enforce a closed-set rule. The downstream dispatch handles
+//     unsupported shapes via a typed error.
+//   - jwk.Key: the dispatcher unwraps via jwk.Export to a raw key,
+//     so the kty-vs-alg check happens then.
+//   - Nil key: legitimate for `dir` (caller provides CEK separately
+//     via WithCEK) and for callers exploring the API.
+//
+// All other built-in algorithm families enforce a concrete key-shape
+// expectation here. The error is wrapped by the WithKey site so the
+// caller sees `jwe.WithKey: ...` consistently.
+func validateAlgorithmForKey(alg jwa.KeyEncryptionAlgorithm, key any) error {
+	if key == nil {
+		return nil
+	}
+	// jwk.Key wrappers: defer to dispatch-time kty validation.
+	if _, ok := key.(jwk.Key); ok {
+		return nil
+	}
+	// Caller-supplied KeyEncrypter / KeyDecrypter implementations
+	// take responsibility for their own key-shape validation. Defer.
+	if _, ok := key.(KeyEncrypter); ok {
+		return nil
+	}
+	if _, ok := key.(KeyDecrypter); ok {
+		return nil
+	}
+	// HPKE raw key types implementing the HPKE key interfaces are
+	// extension-pluggable; defer.
+	if _, ok := key.(jwebb.HPKEKeyEncrypter); ok {
+		return nil
+	}
+	if _, ok := key.(jwebb.HPKEKeyDecrypter); ok {
+		return nil
+	}
+
+	algStr := alg.String()
+	switch {
+	case jwebb.IsHPKE(algStr) || jwebb.IsMLKEM(algStr) || jwebb.IsMLKEMDirect(algStr):
+		// Extension-pluggable; key shape is the extension's contract.
+		return nil
+	case jwebb.IsDirect(algStr):
+		// "dir" requires a byte slice (the CEK) or a symmetric jwk.
+		if _, ok := key.([]byte); !ok {
+			return fmt.Errorf(`algorithm %q requires a []byte key (got %T)`, algStr, key)
+		}
+	case jwebb.IsAESKW(algStr) || jwebb.IsAESGCMKW(algStr) || jwebb.IsPBES2(algStr):
+		if _, ok := key.([]byte); !ok {
+			return fmt.Errorf(`algorithm %q requires a []byte key (got %T)`, algStr, key)
+		}
+	case jwebb.IsRSA15(algStr) || jwebb.IsRSAOAEP(algStr):
+		switch key.(type) {
+		case *rsa.PublicKey, rsa.PublicKey, *rsa.PrivateKey, rsa.PrivateKey:
+		default:
+			return fmt.Errorf(`algorithm %q requires an RSA key (got %T)`, algStr, key)
+		}
+	case jwebb.IsECDHES(algStr):
+		switch key.(type) {
+		case *ecdsa.PublicKey, ecdsa.PublicKey, *ecdsa.PrivateKey, ecdsa.PrivateKey,
+			*ecdh.PublicKey, ecdh.PublicKey, *ecdh.PrivateKey, ecdh.PrivateKey:
+		default:
+			return fmt.Errorf(`algorithm %q requires an ECDSA or ECDH key (got %T)`, algStr, key)
+		}
+	default:
+		// Unknown algorithm family: defer to dispatch.
+		return nil
+	}
+	return nil
+}
+
 func encryptKeyRSA(cek []byte, alg string, key any, encryptFn func([]byte, string, *rsa.PublicKey) (keygen.ByteSource, error)) (keygen.ByteSource, error) {
 	// Handle rsa.PublicKey by value - convert to pointer
 	if pk, ok := key.(rsa.PublicKey); ok {

jwe/interface.go

@@ -54,6 +54,28 @@ type KeyIDer interface {
 // expose the secret key in memory, for example, when you want to use
 // hardware security modules (HSMs) to decrypt the key.
 //
+// Library contract for implementers (read carefully):
+//
+//   - The library has already verified that the wire-level `alg` is
+//     consistent across the protected header and per-recipient header
+//     (RFC 7516 §7.2.1 disjointness). Your DecryptKey is invoked with
+//     the alg the library has decided to use for this attempt.
+//   - The library has NOT validated key-shape-vs-alg compatibility for
+//     your custom decrypter. You receive the raw recipient and message;
+//     headers are split between protected (signed/integrity-protected)
+//     and per-recipient (unprotected). If you read a value from the
+//     unprotected per-recipient header for a security decision, you
+//     must enforce its consistency with the protected header yourself
+//     — the standard built-in decrypters route through a merged-headers
+//     helper that performs this check, but DecryptKey receives the
+//     unmerged inputs.
+//   - Returning a non-nil error short-circuits this recipient. Returning
+//     nil bytes with nil error is treated as "decryption failed" by the
+//     dispatcher (use a non-nil error for clarity).
+//   - You are responsible for any constant-time considerations relevant
+//     to your decryption primitive (e.g. RFC 3218 random-CEK fallback
+//     for RSA-PKCS1v1.5; the library does this for the built-in path).
+//
 // This API is experimental and may change without notice, even
 // in minor releases.
 type KeyDecrypter interface {

jwe/jwe.go

@@ -371,6 +371,9 @@ func (dc *decryptContext) ProcessOptions(options []DecryptOption) error {
 			if !ok {
 				return fmt.Errorf("jwe.decrypt: WithKey() option must be specified using jwa.KeyEncryptionAlgorithm (got %T)", pair.alg)
 			}
+			if err := validateAlgorithmForKey(alg, pair.key); err != nil {
+				return fmt.Errorf("jwe.WithKey: %w", err)
+			}
 			dc.keyProviders = append(dc.keyProviders, &staticKeyProvider{alg: alg, key: pair.key})
 		case identCEK{}:
 			dc.cek = option.MustGet[*[]byte](opt)
@@ -786,6 +789,9 @@ func (ec *encryptContext) ProcessOptions(options []EncryptOption) error {
 			if !ok {
 				return fmt.Errorf("jwe.encrypt: WithKey() option must be specified using jwa.KeyEncryptionAlgorithm (got %T)", wk.alg)
 			}
+			if err := validateAlgorithmForKey(v, wk.key); err != nil {
+				return fmt.Errorf("jwe.WithKey: %w", err)
+			}
 			if jwebb.IsDirectCEK(v.String()) {
 				useRawCEK = true
 			}

jwe/jwe_test.go

@@ -2195,3 +2195,62 @@ func TestDecryptHonorsContextCancellation(t *testing.T) {
 	require.ErrorIs(t, err, context.Canceled,
 		`cancellation must propagate as context.Canceled`)
 }
+
+// TestWithKeyValidatesAlgKeyShape locks the option-time alg-vs-key
+// shape check on jwe.WithKey. Mismatched (alg, key) pairs that cannot
+// succeed downstream produce a crisp `jwe.WithKey: ...` error at
+// option-time instead of a deep nested error from the dispatcher's
+// requireByteKey / keyconv gate.
+func TestWithKeyValidatesAlgKeyShape(t *testing.T) {
+	rsaKey, err := jwxtest.GenerateRsaJwk()
+	require.NoError(t, err)
+	rsaPub, err := jwk.PublicKeyOf(rsaKey)
+	require.NoError(t, err)
+	rawRSAPub, err := jwk.Export[*rsa.PublicKey](rsaPub)
+	require.NoError(t, err)
+
+	t.Run("rejects RSA-OAEP + []byte at option-time on Encrypt", func(t *testing.T) {
+		_, err := jwe.Encrypt([]byte("payload"),
+			jwe.WithKey(jwa.RSA_OAEP(), []byte("not-an-rsa-key")),
+		)
+		require.Error(t, err)
+		require.Contains(t, err.Error(), `jwe.WithKey:`,
+			`error must originate at the option boundary`)
+		require.Contains(t, err.Error(), `requires an RSA key`,
+			`error must name the family expected`)
+	})
+
+	t.Run("rejects A128KW + RSA key at option-time on Encrypt", func(t *testing.T) {
+		_, err := jwe.Encrypt([]byte("payload"),
+			jwe.WithKey(jwa.A128KW(), rawRSAPub),
+		)
+		require.Error(t, err)
+		require.Contains(t, err.Error(), `jwe.WithKey:`)
+		require.Contains(t, err.Error(), `requires a []byte key`)
+	})
+
+	t.Run("rejects ECDH-ES + []byte at option-time", func(t *testing.T) {
+		_, err := jwe.Encrypt([]byte("payload"),
+			jwe.WithKey(jwa.ECDH_ES(), []byte("not-an-ec-key")),
+		)
+		require.Error(t, err)
+		require.Contains(t, err.Error(), `jwe.WithKey:`)
+		require.Contains(t, err.Error(), `requires an ECDSA or ECDH key`)
+	})
+
+	t.Run("accepts jwk.Key wrapper (defers to dispatch)", func(t *testing.T) {
+		// rsaPub wrapped as jwk.Key with RSA-OAEP — legitimate.
+		encrypted, err := jwe.Encrypt([]byte("payload"),
+			jwe.WithKey(jwa.RSA_OAEP(), rsaPub),
+		)
+		require.NoError(t, err)
+		require.NotEmpty(t, encrypted)
+	})
+
+	t.Run("accepts symmetric byte key for AES-KW", func(t *testing.T) {
+		_, err := jwe.Encrypt([]byte("payload"),
+			jwe.WithKey(jwa.A128KW(), make([]byte, 16)),
+		)
+		require.NoError(t, err)
+	})
+}