Repository: lestrrat-go/jwx
Pull request #590 (2022-03-02 04:35:57)

Rework jwe.Decrypt and its options


Patch:
diff --git a/Changes.v2 b/Changes.v2
index f46518ed..330cea18 100644
--- a/Changes.v2
+++ b/Changes.v2
@@ -5,10 +5,48 @@ Changes
 [Incompatible Changes]
   These are changes that are incompatible with the v1.x.x version.
 
+  * [jwe] `jwe.Compact()`'s signature has been changed to `Compact(*Message)`
+
+  * [jwe] `jwe.JSON()` has been removed. You can generate JSON serialization
+    using `jwe.Encrypt(jwe.WitJSON())`
+
   * [jwe] `(jwe.Message).Decrypt()` has been removed. Since formatting of the
     original serialized message matters (including whitespace), using a parsed
     object was inherently confusing.
 
+  * [jwe] `jwe.Encrypt()` can now generate JWE messages in either compact or JSON
+    forms. By default, the compact form is used. JSON format can be
+    enabled by using the `jwe.WithJSON` option.
+
+  * [jwe] `jwe.Encrypt()` can now accept multiple keys by passing multiple
+    `jwe.WithKey()` options. This can be used with `jwe.WithJSON` to
+    create JWE messages with multiple signatures.
+
+  * [jwe] `jwe.WithCompact` and `jwe.WithJSON` options have been added
+    to control the serialization format.
+
+  * [jwe] jwe.Decrypt()'s method signature has been changed to
+
+      jwt.Decrypt([]byte, ...jwe.DecryptOption) ([]byte, error)
+
+    - For static key pair, use `jwe.WithKey()`
+    - For static JWKS, use `jwe.WithKeySet()` (NOTE: InferAlgorithmFromKey like in `jws` package is NOT supported)
+    - For custom, possibly dynamic key provisioning, use `jwe.WithKeyProvider()`
+
+    These options can be stacked. Therefore, you could configure the
+    verification process to attempt a static key pair, a JWKS, and only
+    try other forms if the first two fails, for example.
+   
+  * [jwe] jwe.WithKeyProvider() has been added to specify arbitrary
+    code to specify which keys to try.
+
+  * [jwe] jwe.KeyProvider interface has been added
+
+  * [jwe] jwe.KeyProviderFunc has been added
+ 
+  * [jwe] `WithPostParser()` has been removed. You can achieve the same effect
+    by using `jwe.WithKeyProvider()`
+
   * [jwk] `(jwk.Key).Algorithm()` is now of `jwk.KeyAlgorithm` type. This field used
     to be `string` and therefore could not be passed directly to `jwt.Sign()`
     `jws.Sign()`, `jwe.Encrypt()`, et al. This is no longer the case, and
@@ -25,14 +63,14 @@ Changes
     create JWS messages with multiple signatures.
 
   * [jws] `jws.WithCompact` and `jws.WithJSON` options have been added
-    to control the signatur serialization format.
+    to control the serialization format.
 
   * [jwk] jwk.SetFetcher and jwk.SetFetchFunc has been added.
     They represent something that can fetch a jwk.Set
 
   * [jws] jws.Verify()'s method signature has been changed to
 
-      jwt.Verify([]byte, options ...jws.VerifyOption) ([]byte, error)
+      jwt.Verify([]byte, ...jws.VerifyOption) ([]byte, error)
 
     - For static key pair, use `jws.WithKey()`
     - For static JWKS, use `jws.WithKeySet()`
@@ -65,6 +103,7 @@ Changes
       )
   * [jws] `jws.VerifuAuto` has been removed in favor of using
     `jws.WithVerifyAuto` option with `jws.Verify()`
+
   * [jws] `jws.WithVerifyAuto` has been added to enable verification
     using `jku`.
 
@@ -78,6 +117,7 @@ Changes
     code to specify which keys to try.
 
   * [jws] jws.KeyProvider interface has been added
+
   * [jws] jws.KeyProviderFunc has been added
  
   * [jws] jws.WithKeyUsed has been added to allow users to retrieve
@@ -105,6 +145,7 @@ Changes
     using `jws.WithKey()`
 
   * [jwt] `jwt.Sign()` has been changed so that it works more like the new `jws.Sign()`
+
   * [jwt] `jwt.WithVerify()` has been renamed to `jwt.WithKey()`. The option can
     be used for both signing and parsing.
   * [jwt] `jwt.WithHeaders` and `jwt.WithJwsHeaders` have been removed.
diff --git a/go.mod b/go.mod
index 2c2264cd..a867bee1 100644
--- a/go.mod
+++ b/go.mod
@@ -15,4 +15,4 @@ require (
 	golang.org/x/crypto v0.0.0-20201217014255-9d1352758620
 )
 
-retract v1.2.16
+retract v1.2.16 // see https://github.com/lestrrat-go/jwx/releases/tag/v1.2.17
diff --git a/jwe/README.md b/jwe/README.md
index 216c5336..e43ae167 100644
--- a/jwe/README.md
+++ b/jwe/README.md
@@ -59,7 +59,7 @@ func ExampleEncrypt() {
 
   payload := []byte("Lorem Ipsum")
 
-  encrypted, err := jwe.Encrypt(payload, jwa.RSA1_5, &privkey.PublicKey, jwa.A128CBC_HS256, jwa.NoCompress)
+  encrypted, err := jwe.Encrypt(payload, jwe.WithKey(jwa.RSA1_5, &privkey.PublicKey), jwe.WithContentEncryption(jwa.A128CBC_HS256))
   if err != nil {
     log.Printf("failed to encrypt payload: %s", err)
     return
@@ -79,7 +79,7 @@ func ExampleDecrypt() {
     return
   }
 
-  decrypted, err := jwe.Decrypt(encrypted, jwa.RSA1_5, privkey)
+  decrypted, err := jwe.Decrypt(encrypted, jwe.WithKey(jwa.RSA1_5, privkey))
   if err != nil {
     log.Printf("failed to decrypt: %s", err)
     return
diff --git a/jwe/jwe.go b/jwe/jwe.go
index 88c6b814..8125afc5 100644
--- a/jwe/jwe.go
+++ b/jwe/jwe.go
@@ -343,7 +343,9 @@ func Encrypt(payload []byte, options ...EncryptOption) ([]byte, error) {
 	case fmtCompact:
 		return Compact(msg)
 	case fmtJSON:
-		return JSON(msg)
+		return json.Marshal(msg)
+	case fmtJSONPretty:
+		return json.MarshalIndent(msg, "", "  ")
 	default:
 		return nil, fmt.Errorf(`jwe.Encrypt: invalid serialization`)
 	}
@@ -365,33 +367,11 @@ type DecryptCtx interface {
 }
 
 type decryptCtx struct {
-	alg jwa.KeyEncryptionAlgorithm
-	key interface{}
-	msg *Message
-}
-
-func (ctx *decryptCtx) Algorithm() jwa.KeyEncryptionAlgorithm {
-	return ctx.alg
-}
-
-func (ctx *decryptCtx) SetAlgorithm(v jwa.KeyEncryptionAlgorithm) {
-	ctx.alg = v
-}
-
-func (ctx *decryptCtx) Key() interface{} {
-	return ctx.key
-}
-
-func (ctx *decryptCtx) SetKey(v interface{}) {
-	ctx.key = v
-}
-
-func (ctx *decryptCtx) Message() *Message {
-	return ctx.msg
-}
-
-func (ctx *decryptCtx) SetMessage(m *Message) {
-	ctx.msg = m
+	msg              *Message
+	aad              []byte
+	computedAad      []byte
+	keyProviders     []KeyProvider
+	protectedHeaders Headers
 }
 
 // Decrypt takes the key encryption algorithm and the corresponding
@@ -403,52 +383,259 @@ func (ctx *decryptCtx) SetMessage(m *Message) {
 // `jwa.KeyEncryptionAlgorithm` or otherwise it will cause an error.
 //
 // `key` must be a private key. It can be either in its raw format (e.g. *rsa.PrivateKey) or a jwk.Key
-func Decrypt(buf []byte, alg jwa.KeyAlgorithm, key interface{}, options ...DecryptOption) ([]byte, error) {
-	keyalg, ok := alg.(jwa.KeyEncryptionAlgorithm)
-	if !ok {
-		return nil, errors.Errorf(`expected alg to be jwa.KeyEncryptionAlgorithm, but got %T`, alg)
-	}
-
-	var ctx decryptCtx
-	ctx.key = key
-	ctx.alg = keyalg
+func Decrypt(buf []byte, options ...DecryptOption) ([]byte, error) {
+	var keyProviders []KeyProvider
 
 	var dst *Message
-	var postParse PostParser
 	//nolint:forcetypeassert
 	for _, option := range options {
 		switch option.Ident() {
 		case identMessage{}:
 			dst = option.Value().(*Message)
-		case identPostParser{}:
-			postParse = option.Value().(PostParser)
+		case identKeyProvider{}:
+			keyProviders = append(keyProviders, option.Value().(KeyProvider))
+		case identKey{}:
+			pair := option.Value().(*withKey)
+			alg, ok := pair.alg.(jwa.KeyEncryptionAlgorithm)
+			if !ok {
+				return nil, fmt.Errorf(`WithKey() option must be specified using jwa.KeyEncryptionAlgorithm (got %T)`, pair.alg)
+			}
+			keyProviders = append(keyProviders, &staticKeyProvider{
+				alg: alg,
+				key: pair.key,
+			})
 		}
 	}
 
+	if len(keyProviders) < 1 {
+		return nil, fmt.Errorf(`jwe.Decrypt: no key providers have been provided (see jwe.WithKey(), jwe.WithKeySet(), and jwe.WithKeyProvider()`)
+	}
+
 	msg, err := parseJSONOrCompact(buf, true)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to parse buffer for Decrypt")
 	}
 
-	ctx.msg = msg
-	if postParse != nil {
-		if err := postParse.PostParse(&ctx); err != nil {
-			return nil, errors.Wrap(err, `failed to execute PostParser hook`)
+	// Process things that are common to the message
+	ctx := context.TODO()
+	h, err := msg.protectedHeaders.Clone(ctx)
+	if err != nil {
+		return nil, errors.Wrap(err, `failed to copy protected headers`)
+	}
+	h, err = h.Merge(ctx, msg.unprotectedHeaders)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to merge headers for message decryption")
+	}
+
+	var aad []byte
+	if aadContainer := msg.authenticatedData; aadContainer != nil {
+		aad = base64.Encode(aadContainer)
+	}
+
+	var computedAad []byte
+	if len(msg.rawProtectedHeaders) > 0 {
+		computedAad = msg.rawProtectedHeaders
+	} else {
+		// this is probably not required once msg.Decrypt is deprecated
+		var err error
+		computedAad, err = msg.protectedHeaders.Encode()
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to encode protected headers")
+		}
+	}
+
+	// for each recipient, attempt to match the key providers
+	// if we have no recipients, pretend like we only have one
+	recipients := msg.recipients
+	if len(recipients) == 0 {
+		r := NewRecipient()
+		if err := r.SetHeaders(msg.protectedHeaders); err != nil {
+			return nil, errors.Wrap(err, `failed to set headers to recipient`)
+		}
+		recipients = append(recipients, r)
+	}
+
+	var dctx decryptCtx
+
+	dctx.aad = aad
+	dctx.computedAad = computedAad
+	dctx.msg = msg
+	dctx.keyProviders = keyProviders
+	dctx.protectedHeaders = h
+
+	var lastError error
+	for _, recipient := range recipients {
+		decrypted, err := dctx.try(ctx, recipient)
+		if err != nil {
+			lastError = err
+			continue
+		}
+		if dst != nil {
+			*dst = *msg
+			dst.rawProtectedHeaders = nil
+			dst.storeProtectedHeaders = false
 		}
+		return decrypted, nil
 	}
+	return nil, fmt.Errorf(`jwe.Decrypt: failed to decrypt any of the recipients (last error = %w)`, lastError)
+}
+
+func (dctx *decryptCtx) try(ctx context.Context, recipient Recipient) ([]byte, error) {
+	for i, kp := range dctx.keyProviders {
+		var sink algKeySink
+		if err := kp.FetchKeys(ctx, &sink, recipient, dctx.msg); err != nil {
+			return nil, fmt.Errorf(`key provider %d failed: %w`, i, err)
+		}
+
+		for _, pair := range sink.list {
+			// alg is converted here because pair.alg is of type jwa.KeyAlgorithm.
+			// this may seem ugly, but we're trying to avoid declaring separate
+			// structs for `alg jwa.KeyAlgorithm` and `alg jwa.SignatureAlgorithm`
+			//nolint:forcetypeassert
+			alg := pair.alg.(jwa.KeyEncryptionAlgorithm)
+			key := pair.key
+
+			decrypted, err := dctx.decryptKey(ctx, alg, key, recipient)
+			if err != nil {
+				continue
+			}
 
-	payload, err := doDecryptCtx(&ctx)
+			return decrypted, nil
+		}
+	}
+	return nil, fmt.Errorf(`jwe.Encrypt: failed to match any of the keys with recipient`)
+}
+
+func (dctx *decryptCtx) decryptKey(ctx context.Context, alg jwa.KeyEncryptionAlgorithm, key interface{}, recipient Recipient) ([]byte, error) {
+	if jwkKey, ok := key.(jwk.Key); ok {
+		var raw interface{}
+		if err := jwkKey.Raw(&raw); err != nil {
+			return nil, errors.Wrapf(err, `failed to retrieve raw key from %T`, key)
+		}
+		key = raw
+	}
+	dec := NewDecrypter(alg, dctx.msg.protectedHeaders.ContentEncryption(), key).
+		AuthenticatedData(dctx.aad).
+		ComputedAuthenticatedData(dctx.computedAad).
+		InitializationVector(dctx.msg.initializationVector).
+		Tag(dctx.msg.tag)
+
+	if recipient.Headers().Algorithm() != alg {
+		// algorithms don't match
+		return nil, fmt.Errorf(`jwe.Decrypt: key and recipient algorithms do not match`)
+	}
+
+	h2, err := dctx.protectedHeaders.Clone(ctx)
+	if err != nil {
+		return nil, fmt.Errorf(`jwe.Decrypt: failed to copy headers (1): %w`, err)
+	}
+
+	h2, err = h2.Merge(ctx, recipient.Headers())
+	if err != nil {
+		return nil, fmt.Errorf(`failed to copy headers (2): %w`, err)
+	}
+
+	switch alg {
+	case jwa.ECDH_ES, jwa.ECDH_ES_A128KW, jwa.ECDH_ES_A192KW, jwa.ECDH_ES_A256KW:
+		epkif, ok := h2.Get(EphemeralPublicKeyKey)
+		if !ok {
+			return nil, errors.New("failed to get 'epk' field")
+		}
+		switch epk := epkif.(type) {
+		case jwk.ECDSAPublicKey:
+			var pubkey ecdsa.PublicKey
+			if err := epk.Raw(&pubkey); err != nil {
+				return nil, errors.Wrap(err, "failed to get public key")
+			}
+			dec.PublicKey(&pubkey)
+		case jwk.OKPPublicKey:
+			var pubkey interface{}
+			if err := epk.Raw(&pubkey); err != nil {
+				return nil, errors.Wrap(err, "failed to get public key")
+			}
+			dec.PublicKey(pubkey)
+		default:
+			return nil, errors.Errorf("unexpected 'epk' type %T for alg %s", epkif, alg)
+		}
+
+		if apu := h2.AgreementPartyUInfo(); len(apu) > 0 {
+			dec.AgreementPartyUInfo(apu)
+		}
+
+		if apv := h2.AgreementPartyVInfo(); len(apv) > 0 {
+			dec.AgreementPartyVInfo(apv)
+		}
+	case jwa.A128GCMKW, jwa.A192GCMKW, jwa.A256GCMKW:
+		ivB64, ok := h2.Get(InitializationVectorKey)
+		if !ok {
+			return nil, errors.New("failed to get 'iv' field")
+		}
+		ivB64Str, ok := ivB64.(string)
+		if !ok {
+			return nil, errors.Errorf("unexpected type for 'iv': %T", ivB64)
+		}
+		tagB64, ok := h2.Get(TagKey)
+		if !ok {
+			return nil, errors.New("failed to get 'tag' field")
+		}
+		tagB64Str, ok := tagB64.(string)
+		if !ok {
+			return nil, errors.Errorf("unexpected type for 'tag': %T", tagB64)
+		}
+		iv, err := base64.DecodeString(ivB64Str)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to b64-decode 'iv'")
+		}
+		tag, err := base64.DecodeString(tagB64Str)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to b64-decode 'tag'")
+		}
+		dec.KeyInitializationVector(iv)
+		dec.KeyTag(tag)
+	case jwa.PBES2_HS256_A128KW, jwa.PBES2_HS384_A192KW, jwa.PBES2_HS512_A256KW:
+		saltB64, ok := h2.Get(SaltKey)
+		if !ok {
+			return nil, errors.New("failed to get 'p2s' field")
+		}
+		saltB64Str, ok := saltB64.(string)
+		if !ok {
+			return nil, errors.Errorf("unexpected type for 'p2s': %T", saltB64)
+		}
+
+		count, ok := h2.Get(CountKey)
+		if !ok {
+			return nil, errors.New("failed to get 'p2c' field")
+		}
+		countFlt, ok := count.(float64)
+		if !ok {
+			return nil, errors.Errorf("unexpected type for 'p2c': %T", count)
+		}
+		salt, err := base64.DecodeString(saltB64Str)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to b64-decode 'salt'")
+		}
+		dec.KeySalt(salt)
+		dec.KeyCount(int(countFlt))
+	}
+
+	plaintext, err := dec.Decrypt(recipient.EncryptedKey(), dctx.msg.cipherText)
 	if err != nil {
-		return nil, errors.Wrap(err, `failed to decrypt message`)
+		return nil, fmt.Errorf(`jwe.Decrypt: decryption failed: %w`, err)
+	}
+
+	if h2.Compression() == jwa.Deflate {
+		buf, err := uncompress(plaintext)
+		if err != nil {
+			return nil, fmt.Errorf(`jwe.Derypt: failed to uncompress payload: %w`, err)
+		}
+		plaintext = buf
 	}
 
-	if dst != nil {
-		*dst = *msg
-		dst.rawProtectedHeaders = nil
-		dst.storeProtectedHeaders = false
+	if plaintext == nil {
+		return nil, errors.New("failed to find matching recipient")
 	}
 
-	return payload, nil
+	return plaintext, nil
 }
 
 // Parse parses the JWE message into a Message object. The JWE message
diff --git a/jwe/key_provider.go b/jwe/key_provider.go
new file mode 100644
index 00000000..c1581041
--- /dev/null
+++ b/jwe/key_provider.go
@@ -0,0 +1,100 @@
+package jwe
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"github.com/lestrrat-go/jwx/jwa"
+	"github.com/lestrrat-go/jwx/jwk"
+)
+
+type KeyProvider interface {
+	FetchKeys(context.Context, KeySink, Recipient, *Message) error
+}
+
+type KeySink interface {
+	Key(jwa.KeyEncryptionAlgorithm, interface{})
+}
+
+type algKeyPair struct {
+	alg jwa.KeyAlgorithm
+	key interface{}
+}
+
+type algKeySink struct {
+	mu   sync.Mutex
+	list []algKeyPair
+}
+
+func (s *algKeySink) Key(alg jwa.KeyEncryptionAlgorithm, key interface{}) {
+	s.mu.Lock()
+	s.list = append(s.list, algKeyPair{alg, key})
+	s.mu.Unlock()
+}
+
+type staticKeyProvider struct {
+	alg jwa.KeyEncryptionAlgorithm
+	key interface{}
+}
+
+func (kp *staticKeyProvider) FetchKeys(_ context.Context, sink KeySink, _ Recipient, _ *Message) error {
+	sink.Key(kp.alg, kp.key)
+	return nil
+}
+
+type keySetProvider struct {
+	set        jwk.Set
+	requireKid bool
+}
+
+func (kp *keySetProvider) selectKey(sink KeySink, key jwk.Key, _ Recipient, _ *Message) error {
+	if usage := key.KeyUsage(); usage != "" && usage != jwk.ForEncryption.String() {
+		return nil
+	}
+
+	if v := key.Algorithm(); v.String() != "" {
+		var alg jwa.KeyEncryptionAlgorithm
+		if err := alg.Accept(v); err != nil {
+			return fmt.Errorf(`invalid key encryption algorithm %s: %w`, key.Algorithm(), err)
+		}
+
+		sink.Key(alg, key)
+		return nil
+	}
+
+	return nil
+}
+
+func (kp *keySetProvider) FetchKeys(_ context.Context, sink KeySink, r Recipient, msg *Message) error {
+	if kp.requireKid {
+		var key jwk.Key
+
+		wantedKid := r.Headers().KeyID()
+		if wantedKid == "" {
+			return fmt.Errorf(`failed to find matching key: no key ID ("kid") specified in token but multiple keys available in key set`)
+		}
+		// Otherwise we better be able to look up the key, baby.
+		v, ok := kp.set.LookupKeyID(wantedKid)
+		if !ok {
+			return fmt.Errorf(`failed to find key with key ID %q in key set`, wantedKid)
+		}
+		key = v
+
+		return kp.selectKey(sink, key, r, msg)
+	}
+
+	for i := 0; i < kp.set.Len(); i++ {
+		key, _ := kp.set.Get(i)
+		if err := kp.selectKey(sink, key, r, msg); err != nil {
+			continue
+		}
+	}
+	return nil
+}
+
+type KeyProviderFunc func(context.Context, KeySink, Recipient, *Message) error
+
+func (kp KeyProviderFunc) FetchKeys(ctx context.Context, sink KeySink, r Recipient, msg *Message) error {
+	return kp(ctx, sink, r, msg)
+}
diff --git a/jwe/message.go b/jwe/message.go
index 6e399f5f..48d7115b 100644
--- a/jwe/message.go
+++ b/jwe/message.go
@@ -2,15 +2,12 @@ package jwe
 
 import (
 	"context"
-	"crypto/ecdsa"
 	"fmt"
 
 	"github.com/lestrrat-go/jwx/internal/json"
 	"github.com/lestrrat-go/jwx/internal/pool"
-	"github.com/lestrrat-go/jwx/jwk"
 
 	"github.com/lestrrat-go/jwx/internal/base64"
-	"github.com/lestrrat-go/jwx/jwa"
 	"github.com/pkg/errors"
 )
 
@@ -436,193 +433,61 @@ func (m *Message) makeDummyRecipient(enckeybuf string, protected Headers) error
 	return nil
 }
 
-func doDecryptCtx(dctx *decryptCtx) ([]byte, error) {
-	m := dctx.msg
-	alg := dctx.alg
-	key := dctx.key
+func Compact(m *Message) ([]byte, error) {
+	if len(m.recipients) != 1 {
+		return nil, errors.New("wrong number of recipients for compact serialization")
+	}
 
-	if jwkKey, ok := key.(jwk.Key); ok {
-		var raw interface{}
-		if err := jwkKey.Raw(&raw); err != nil {
-			return nil, errors.Wrapf(err, `failed to retrieve raw key from %T`, key)
-		}
-		key = raw
+	recipient := m.recipients[0]
+
+	// The protected header must be a merge between the message-wide
+	// protected header AND the recipient header
+
+	// There's something wrong if m.protectedHeaders is nil, but
+	// it could happen
+	if m.protectedHeaders == nil {
+		return nil, errors.New("invalid protected header")
 	}
 
-	var err error
 	ctx := context.TODO()
-	h, err := m.protectedHeaders.Clone(ctx)
+	hcopy, err := m.protectedHeaders.Clone(ctx)
 	if err != nil {
-		return nil, errors.Wrap(err, `failed to copy protected headers`)
+		return nil, errors.Wrap(err, "failed to copy protected header")
 	}
-	h, err = h.Merge(ctx, m.unprotectedHeaders)
+	hcopy, err = hcopy.Merge(ctx, m.unprotectedHeaders)
 	if err != nil {
-		return nil, errors.Wrap(err, "failed to merge headers for message decryption")
+		return nil, errors.Wrap(err, "failed to merge unprotected header")
 	}
-
-	enc := m.protectedHeaders.ContentEncryption()
-	var aad []byte
-	if aadContainer := m.authenticatedData; aadContainer != nil {
-		aad = base64.Encode(aadContainer)
-	}
-
-	var computedAad []byte
-	if len(m.rawProtectedHeaders) > 0 {
-		computedAad = m.rawProtectedHeaders
-	} else {
-		// this is probably not required once msg.Decrypt is deprecated
-		var err error
-		computedAad, err = m.protectedHeaders.Encode()
-		if err != nil {
-			return nil, errors.Wrap(err, "failed to encode protected headers")
-		}
+	hcopy, err = hcopy.Merge(ctx, recipient.Headers())
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to merge recipient header")
 	}
 
-	dec := NewDecrypter(alg, enc, key).
-		AuthenticatedData(aad).
-		ComputedAuthenticatedData(computedAad).
-		InitializationVector(m.initializationVector).
-		Tag(m.tag)
-
-	var plaintext []byte
-	var lastError error
-
-	// if we have no recipients, pretend like we only have one
-	recipients := m.recipients
-	if len(recipients) == 0 {
-		r := NewRecipient()
-		if err := r.SetHeaders(m.protectedHeaders); err != nil {
-			return nil, errors.Wrap(err, `failed to set headers to recipient`)
-		}
-		recipients = append(recipients, r)
+	protected, err := hcopy.Encode()
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to encode header")
 	}
 
-	for _, recipient := range recipients {
-		// strategy: try each recipient. If we fail in one of the steps,
-		// keep looping because there might be another key with the same algo
-		if recipient.Headers().Algorithm() != alg {
-			// algorithms don't match
-			continue
-		}
-
-		h2, err := h.Clone(ctx)
-		if err != nil {
-			lastError = errors.Wrap(err, `failed to copy headers (1)`)
-			continue
-		}
+	encryptedKey := base64.Encode(recipient.EncryptedKey())
+	iv := base64.Encode(m.initializationVector)
+	cipher := base64.Encode(m.cipherText)
+	tag := base64.Encode(m.tag)
 
-		h2, err = h2.Merge(ctx, recipient.Headers())
-		if err != nil {
-			lastError = errors.Wrap(err, `failed to copy headers (2)`)
-			continue
-		}
-
-		switch alg {
-		case jwa.ECDH_ES, jwa.ECDH_ES_A128KW, jwa.ECDH_ES_A192KW, jwa.ECDH_ES_A256KW:
-			epkif, ok := h2.Get(EphemeralPublicKeyKey)
-			if !ok {
-				return nil, errors.New("failed to get 'epk' field")
-			}
-			switch epk := epkif.(type) {
-			case jwk.ECDSAPublicKey:
-				var pubkey ecdsa.PublicKey
-				if err := epk.Raw(&pubkey); err != nil {
-					return nil, errors.Wrap(err, "failed to get public key")
-				}
-				dec.PublicKey(&pubkey)
-			case jwk.OKPPublicKey:
-				var pubkey interface{}
-				if err := epk.Raw(&pubkey); err != nil {
-					return nil, errors.Wrap(err, "failed to get public key")
-				}
-				dec.PublicKey(pubkey)
-			default:
-				return nil, errors.Errorf("unexpected 'epk' type %T for alg %s", epkif, alg)
-			}
-
-			if apu := h2.AgreementPartyUInfo(); len(apu) > 0 {
-				dec.AgreementPartyUInfo(apu)
-			}
-
-			if apv := h2.AgreementPartyVInfo(); len(apv) > 0 {
-				dec.AgreementPartyVInfo(apv)
-			}
-		case jwa.A128GCMKW, jwa.A192GCMKW, jwa.A256GCMKW:
-			ivB64, ok := h2.Get(InitializationVectorKey)
-			if !ok {
-				return nil, errors.New("failed to get 'iv' field")
-			}
-			ivB64Str, ok := ivB64.(string)
-			if !ok {
-				return nil, errors.Errorf("unexpected type for 'iv': %T", ivB64)
-			}
-			tagB64, ok := h2.Get(TagKey)
-			if !ok {
-				return nil, errors.New("failed to get 'tag' field")
-			}
-			tagB64Str, ok := tagB64.(string)
-			if !ok {
-				return nil, errors.Errorf("unexpected type for 'tag': %T", tagB64)
-			}
-			iv, err := base64.DecodeString(ivB64Str)
-			if err != nil {
-				return nil, errors.Wrap(err, "failed to b64-decode 'iv'")
-			}
-			tag, err := base64.DecodeString(tagB64Str)
-			if err != nil {
-				return nil, errors.Wrap(err, "failed to b64-decode 'tag'")
-			}
-			dec.KeyInitializationVector(iv)
-			dec.KeyTag(tag)
-		case jwa.PBES2_HS256_A128KW, jwa.PBES2_HS384_A192KW, jwa.PBES2_HS512_A256KW:
-			saltB64, ok := h2.Get(SaltKey)
-			if !ok {
-				return nil, errors.New("failed to get 'p2s' field")
-			}
-			saltB64Str, ok := saltB64.(string)
-			if !ok {
-				return nil, errors.Errorf("unexpected type for 'p2s': %T", saltB64)
-			}
-
-			count, ok := h2.Get(CountKey)
-			if !ok {
-				return nil, errors.New("failed to get 'p2c' field")
-			}
-			countFlt, ok := count.(float64)
-			if !ok {
-				return nil, errors.Errorf("unexpected type for 'p2c': %T", count)
-			}
-			salt, err := base64.DecodeString(saltB64Str)
-			if err != nil {
-				return nil, errors.Wrap(err, "failed to b64-decode 'salt'")
-			}
-			dec.KeySalt(salt)
-			dec.KeyCount(int(countFlt))
-		}
-
-		plaintext, err = dec.Decrypt(recipient.EncryptedKey(), m.cipherText)
-		if err != nil {
-			lastError = errors.Wrap(err, `failed to decrypt`)
-			continue
-		}
-
-		if h2.Compression() == jwa.Deflate {
-			buf, err := uncompress(plaintext)
-			if err != nil {
-				lastError = errors.Wrap(err, `failed to uncompress payload`)
-				continue
-			}
-			plaintext = buf
-		}
-		break
-	}
-
-	if plaintext == nil {
-		if lastError != nil {
-			return nil, errors.Errorf(`failed to find matching recipient to decrypt key (last error = %s)`, lastError)
-		}
-		return nil, errors.New("failed to find matching recipient")
-	}
+	buf := pool.GetBytesBuffer()
+	defer pool.ReleaseBytesBuffer(buf)
 
-	return plaintext, nil
+	buf.Grow(len(protected) + len(encryptedKey) + len(iv) + len(cipher) + len(tag) + 4)
+	buf.Write(protected)
+	buf.WriteByte('.')
+	buf.Write(encryptedKey)
+	buf.WriteByte('.')
+	buf.Write(iv)
+	buf.WriteByte('.')
+	buf.Write(cipher)
+	buf.WriteByte('.')
+	buf.Write(tag)
+
+	result := make([]byte, buf.Len())
+	copy(result, buf.Bytes())
+	return result, nil
 }
diff --git a/jwe/options.go b/jwe/options.go
index 7621364b..5d50452e 100644
--- a/jwe/options.go
+++ b/jwe/options.go
@@ -4,26 +4,10 @@ import (
 	"context"
 
 	"github.com/lestrrat-go/jwx/jwa"
+	"github.com/lestrrat-go/jwx/jwk"
 	"github.com/lestrrat-go/option"
 )
 
-type SerializerOption interface {
-	Option
-	serializerOption()
-}
-
-type serializerOption struct {
-	Option
-}
-
-func (*serializerOption) serializerOption() {}
-
-// WithPrettyFormat specifies if the `jwe.JSON` serialization tool
-// should generate pretty-formatted output
-func WithPrettyFormat(b bool) SerializerOption {
-	return &serializerOption{option.New(identPrettyFormat{}, b)}
-}
-
 // Specify contents of the protected header. Some fields such as
 // "enc" and "zip" will be overwritten when encryption is performed.
 func WithProtectedHeaders(h Headers) EncryptOption {
@@ -31,22 +15,6 @@ func WithProtectedHeaders(h Headers) EncryptOption {
 	return &encryptOption{option.New(identProtectedHeaders{}, cloned)}
 }
 
-// WithPostParser specifies the handler to be called immediately
-// after the JWE message has been parsed, but before decryption
-// takes place during `jwe.Decrypt`.
-//
-// This option exists to allow advanced users that require the use
-// of information stored in the JWE message to determine how the
-// decryption should be handled.
-//
-// For security reasons it is highly recommended that you thoroughly
-// study how the process works before using this option. This is especially
-// true if you are trying to infer key algorithms and keys to use to
-// decrypt a message using non-standard hints.
-func WithPostParser(p PostParser) DecryptOption {
-	return &decryptOption{option.New(identPostParser{}, p)}
-}
-
 type DecryptEncryptOption interface {
 	DecryptOption
 	EncryptOption
@@ -97,19 +65,20 @@ func WithKey(alg jwa.KeyAlgorithm, key interface{}, options ...WithKeySuboption)
 	})}
 }
 
-type JSONSuboption interface {
-	Option
-	withJSONSuboption()
-}
-
-type jsonSuboption struct {
-	Option
-}
-
-func (*jsonSuboption) withJSONSuboption() {}
+func WithKeySet(set jwk.Set, options ...WithKeySetSuboption) DecryptOption {
+	var requireKid bool
+	for _, option := range options {
+		//nolint:forcetypeassert
+		switch option.Ident() {
+		case identRequireKid{}:
+			requireKid = option.Value().(bool)
+		}
+	}
 
-func WithPretty(v bool) JSONSuboption {
-	return &jsonSuboption{option.New(identPretty{}, v)}
+	return WithKeyProvider(&keySetProvider{
+		set:        set,
+		requireKid: requireKid,
+	})
 }
 
 // WithJSON specifies that the result of `jwe.Encrypt()` is serialized in
@@ -117,7 +86,7 @@ func WithPretty(v bool) JSONSuboption {
 //
 // If you pass multiple keys to `jwe.Encrypt()`, it will fail unless
 // you also pass this option.
-func WithJSON(options ...JSONSuboption) EncryptOption {
+func WithJSON(options ...WithJSONSuboption) EncryptOption {
 	var pretty bool
 	for _, option := range options {
 		//nolint:forcetypeassert
diff --git a/jwe/options.yaml b/jwe/options.yaml
index 208154ae..74715aec 100644
--- a/jwe/options.yaml
+++ b/jwe/options.yaml
@@ -7,6 +7,13 @@ interfaces:
   - name: EncryptOption
     comment: |
       EncryptOption describes options that can be passed to `jwe.Encrypt`
+  - name: WithJSONSuboption
+    concrete_type: withJSONSuboption
+    comment: |
+      JSONSuboption describes suboptions that can be passed to `jwe.WithJSON()` option
+  - name: WithKeySetSuboption
+    comment: |
+      WithKeySetSuboption is a suboption passed to the WithKeySet() option
 options:
   - ident: Key
     skip_option: true
@@ -14,13 +21,11 @@ options:
     skip_option: true
   - ident: ProtectedHeaders
     skip_option: true
-  # These should be deprecated?
-  - ident: PostParser
-    skip_option: true
-  - ident: PrettyFormat
-    skip_option: true
   - ident: RecipientHeaders
     skip_option: true
+  - ident: KeyProvider
+    interface: DecryptOption
+    argument_type: KeyProvider
   - ident: Serialization
     option_name: WithCompact
     interface: EncryptOption
@@ -58,4 +63,17 @@ options:
       than inspecting its contents. Particularly, do not expect the message
       reliable when you call `Decrypt` on it. `(jwe.Message).Decrypt` is
       slated to be deprecated in the next major version.
+  - ident: RequireKid
+    interface: WithKeySetSuboption
+    argument_type: bool
+    comment: |
+      WithrequiredKid specifies whether the keys in the jwk.Set should
+      only be matched if the target JWE message's Key ID and the Key ID
+      in the given key matches.
+  - ident: Pretty
+    interface: WithJSONSuboption
+    argument_type: bool
+    comment: |
+      WithPretty specifies whether the JSON output should be formatted and
+      indented
 
diff --git a/jwe/options_gen.go b/jwe/options_gen.go
index 4c9cfdab..e6237e8a 100644
--- a/jwe/options_gen.go
+++ b/jwe/options_gen.go
@@ -33,15 +33,39 @@ type encryptOption struct {
 
 func (*encryptOption) encryptOption() {}
 
+// JSONSuboption describes suboptions that can be passed to `jwe.WithJSON()` option
+type WithJSONSuboption interface {
+	Option
+	withJSONSuboption()
+}
+
+type withJSONSuboption struct {
+	Option
+}
+
+func (*withJSONSuboption) withJSONSuboption() {}
+
+// WithKeySetSuboption is a suboption passed to the WithKeySet() option
+type WithKeySetSuboption interface {
+	Option
+	withKeySetSuboption()
+}
+
+type withKeySetSuboption struct {
+	Option
+}
+
+func (*withKeySetSuboption) withKeySetSuboption() {}
+
 type identCompress struct{}
 type identContentEncryptionAlgorithm struct{}
 type identKey struct{}
+type identKeyProvider struct{}
 type identMessage struct{}
-type identPostParser struct{}
 type identPretty struct{}
-type identPrettyFormat struct{}
 type identProtectedHeaders struct{}
 type identRecipientHeaders struct{}
+type identRequireKid struct{}
 type identSerialization struct{}
 
 func (identCompress) String() string {
@@ -56,22 +80,18 @@ func (identKey) String() string {
 	return "WithKey"
 }
 
-func (identMessage) String() string {
-	return "WithMessage"
+func (identKeyProvider) String() string {
+	return "WithKeyProvider"
 }
 
-func (identPostParser) String() string {
-	return "WithPostParser"
+func (identMessage) String() string {
+	return "WithMessage"
 }
 
 func (identPretty) String() string {
 	return "WithPretty"
 }
 
-func (identPrettyFormat) String() string {
-	return "WithPrettyFormat"
-}
-
 func (identProtectedHeaders) String() string {
 	return "WithProtectedHeaders"
 }
@@ -80,6 +100,10 @@ func (identRecipientHeaders) String() string {
 	return "WithRecipientHeaders"
 }
 
+func (identRequireKid) String() string {
+	return "WithRequireKid"
+}
+
 func (identSerialization) String() string {
 	return "WithCompact"
 }
@@ -98,6 +122,10 @@ func WithContentEncryption(v jwa.ContentEncryptionAlgorithm) EncryptOption {
 	return &encryptOption{option.New(identContentEncryptionAlgorithm{}, v)}
 }
 
+func WithKeyProvider(v KeyProvider) DecryptOption {
+	return &decryptOption{option.New(identKeyProvider{}, v)}
+}
+
 // WithMessage provides a message object to be populated by `jwe.Decrpt`
 // Using this option allows you to decrypt AND obtain the `jwe.Message`
 // in one go.
@@ -110,6 +138,19 @@ func WithMessage(v *Message) DecryptOption {
 	return &decryptOption{option.New(identMessage{}, v)}
 }
 
+// WithPretty specifies whether the JSON output should be formatted and
+// indented
+func WithPretty(v bool) WithJSONSuboption {
+	return &withJSONSuboption{option.New(identPretty{}, v)}
+}
+
+// WithrequiredKid specifies whether the keys in the jwk.Set should
+// only be matched if the target JWE message's Key ID and the Key ID
+// in the given key matches.
+func WithRequireKid(v bool) WithKeySetSuboption {
+	return &withKeySetSuboption{option.New(identRequireKid{}, v)}
+}
+
 // WithCompact specifies that the result of `jwe.Encrypt()` is serialized in
 // compact format.
 //
diff --git a/jwe/serializer.go b/jwe/serializer.go
index 27ca9b86..1ca88c03 100644
--- a/jwe/serializer.go
+++ b/jwe/serializer.go
@@ -1,94 +1,1 @@
 package jwe
-
-import (
-	"context"
-
-	"github.com/lestrrat-go/jwx/internal/base64"
-	"github.com/lestrrat-go/jwx/internal/json"
-
-	"github.com/lestrrat-go/jwx/internal/pool"
-	"github.com/pkg/errors"
-)
-
-// Compact encodes the given message into a JWE compact serialization format.
-//
-// Currently `Compact()` does not take any options, but the API is
-// set up as such to allow future expansions
-func Compact(m *Message, _ ...SerializerOption) ([]byte, error) {
-	if len(m.recipients) != 1 {
-		return nil, errors.New("wrong number of recipients for compact serialization")
-	}
-
-	recipient := m.recipients[0]
-
-	// The protected header must be a merge between the message-wide
-	// protected header AND the recipient header
-
-	// There's something wrong if m.protectedHeaders is nil, but
-	// it could happen
-	if m.protectedHeaders == nil {
-		return nil, errors.New("invalid protected header")
-	}
-
-	ctx := context.TODO()
-	hcopy, err := m.protectedHeaders.Clone(ctx)
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to copy protected header")
-	}
-	hcopy, err = hcopy.Merge(ctx, m.unprotectedHeaders)
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to merge unprotected header")
-	}
-	hcopy, err = hcopy.Merge(ctx, recipient.Headers())
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to merge recipient header")
-	}
-
-	protected, err := hcopy.Encode()
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to encode header")
-	}
-
-	encryptedKey := base64.Encode(recipient.EncryptedKey())
-	iv := base64.Encode(m.initializationVector)
-	cipher := base64.Encode(m.cipherText)
-	tag := base64.Encode(m.tag)
-
-	buf := pool.GetBytesBuffer()
-	defer pool.ReleaseBytesBuffer(buf)
-
-	buf.Grow(len(protected) + len(encryptedKey) + len(iv) + len(cipher) + len(tag) + 4)
-	buf.Write(protected)
-	buf.WriteByte('.')
-	buf.Write(encryptedKey)
-	buf.WriteByte('.')
-	buf.Write(iv)
-	buf.WriteByte('.')
-	buf.Write(cipher)
-	buf.WriteByte('.')
-	buf.Write(tag)
-
-	result := make([]byte, buf.Len())
-	copy(result, buf.Bytes())
-	return result, nil
-}
-
-// JSON encodes the message into a JWE JSON serialization format.
-//
-// If `WithPrettyFormat(true)` is passed as an option, the returned
-// value will be formatted using `json.MarshalIndent()`
-func JSON(m *Message, options ...SerializerOption) ([]byte, error) {
-	var pretty bool
-	for _, option := range options {
-		//nolint:forcetypeassert
-		switch option.Ident() {
-		case identPrettyFormat{}:
-			pretty = option.Value().(bool)
-		}
-	}
-
-	if pretty {
-		return json.MarshalIndent(m, "", "  ")
-	}
-	return json.Marshal(m)
-}
diff --git a/jws/jws.go b/jws/jws.go
index 1dbbe6dd..8d17ed86 100644
--- a/jws/jws.go
+++ b/jws/jws.go
@@ -313,6 +313,10 @@ func Verify(buf []byte, options ...VerifyOption) ([]byte, error) {
 		}
 	}
 
+	if len(keyProviders) < 1 {
+		return nil, fmt.Errorf(`jws.Verify: no key providers have been provided (see jws.WithKey(), jws.WithKeySet(), jws.WithVerifyAuto(), and jws.WithKeyProvider()`)
+	}
+
 	msg, err := Parse(buf)
 	if err != nil {
 		return nil, fmt.Errorf(`failed to parse jws: %w`, err)
@@ -363,7 +367,7 @@ func Verify(buf []byte, options ...VerifyOption) ([]byte, error) {
 
 		for i, kp := range keyProviders {
 			var sink algKeySink
-			if err := kp.FetchKeys(ctx, &sink, sig); err != nil {
+			if err := kp.FetchKeys(ctx, &sink, sig, msg); err != nil {
 				return nil, fmt.Errorf(`key provider %d failed: %w`, i, err)
 			}
 
diff --git a/jws/key_provider.go b/jws/key_provider.go
index 4c4ba5e2..a333161e 100644
--- a/jws/key_provider.go
+++ b/jws/key_provider.go
@@ -11,7 +11,7 @@ import (
 )
 
 type KeyProvider interface {
-	FetchKeys(context.Context, KeySink, *Signature) error
+	FetchKeys(context.Context, KeySink, *Signature, *Message) error
 }
 
 type KeySink interface {
@@ -39,7 +39,7 @@ type staticKeyProvider struct {
 	key interface{}
 }
 
-func (kp *staticKeyProvider) FetchKeys(_ context.Context, sink KeySink, _ *Signature) error {
+func (kp *staticKeyProvider) FetchKeys(_ context.Context, sink KeySink, _ *Signature, _ *Message) error {
 	sink.Key(kp.alg, kp.key)
 	return nil
 }
@@ -51,7 +51,7 @@ type keySetProvider struct {
 	inferAlgorithm bool // true if the algorithm should be inferred from key type
 }
 
-func (kp *keySetProvider) selectKey(sink KeySink, key jwk.Key, sig *Signature) error {
+func (kp *keySetProvider) selectKey(sink KeySink, key jwk.Key, sig *Signature, _ *Message) error {
 	if usage := key.KeyUsage(); usage != "" && usage != jwk.ForSignature.String() {
 		return nil
 	}
@@ -92,7 +92,7 @@ func (kp *keySetProvider) selectKey(sink KeySink, key jwk.Key, sig *Signature) e
 	return nil
 }
 
-func (kp *keySetProvider) FetchKeys(_ context.Context, sink KeySink, sig *Signature) error {
+func (kp *keySetProvider) FetchKeys(_ context.Context, sink KeySink, sig *Signature, msg *Message) error {
 	if kp.requireKid {
 		var key jwk.Key
 
@@ -118,12 +118,12 @@ func (kp *keySetProvider) FetchKeys(_ context.Context, sink KeySink, sig *Signat
 			key = v
 		}
 
-		return kp.selectKey(sink, key, sig)
+		return kp.selectKey(sink, key, sig, msg)
 	}
 
 	for i := 0; i < kp.set.Len(); i++ {
 		key, _ := kp.set.Get(i)
-		if err := kp.selectKey(sink, key, sig); err != nil {
+		if err := kp.selectKey(sink, key, sig, msg); err != nil {
 			continue
 		}
 	}
@@ -135,7 +135,7 @@ type jkuProvider struct {
 	options []jwk.FetchOption
 }
 
-func (kp jkuProvider) FetchKeys(ctx context.Context, sink KeySink, sig *Signature) error {
+func (kp jkuProvider) FetchKeys(ctx context.Context, sink KeySink, sig *Signature, _ *Message) error {
 	kid := sig.ProtectedHeaders().KeyID()
 	if kid == "" {
 		return nil
@@ -185,8 +185,8 @@ func (kp jkuProvider) FetchKeys(ctx context.Context, sink KeySink, sig *Signatur
 	return nil
 }
 
-type KeyProviderFunc func(context.Context, KeySink, *Signature) error
+type KeyProviderFunc func(context.Context, KeySink, *Signature, *Message) error
 
-func (kp KeyProviderFunc) FetchKeys(ctx context.Context, sink KeySink, sig *Signature) error {
-	return kp(ctx, sink, sig)
+func (kp KeyProviderFunc) FetchKeys(ctx context.Context, sink KeySink, sig *Signature, msg *Message) error {
+	return kp(ctx, sink, sig, msg)
 }
diff --git a/jwt/jwt.go b/jwt/jwt.go
index ee1dfea3..18031cc3 100644
--- a/jwt/jwt.go
+++ b/jwt/jwt.go
@@ -259,7 +259,8 @@ OUTER:
 				decryptOpts = []jwe.DecryptOption{jwe.WithMessage(m)}
 			}
 
-			v, err := jwe.Decrypt(data, dp.Algorithm(), dp.Key(), decryptOpts...)
+			decryptOpts = append(decryptOpts, jwe.WithKey(dp.Algorithm(), dp.Key()))
+			v, err := jwe.Decrypt(data, decryptOpts...)
 			if err != nil {
 				return nil, errors.Wrap(err, `failed to decrypt payload`)
 			}


Test patch:
diff --git a/examples/jwe_example_test.go b/examples/jwe_example_test.go
index 4284e846..186d1a70 100644
--- a/examples/jwe_example_test.go
+++ b/examples/jwe_example_test.go
@@ -1,6 +1,7 @@
 package examples_test
 
 import (
+	"context"
 	"crypto/rand"
 	"crypto/rsa"
 	"fmt"
@@ -52,7 +53,7 @@ func ExampleJWE_Decrypt() {
 		return
 	}
 
-	decrypted, err := jwe.Decrypt(encrypted, jwa.RSA1_5, privkey)
+	decrypted, err := jwe.Decrypt(encrypted, jwe.WithKey(jwa.RSA1_5, privkey))
 	if err != nil {
 		log.Printf("failed to decrypt: %s", err)
 		return
@@ -99,8 +100,7 @@ func ExampleJWE_ComplexDecrypt() {
 	// Here we are using a function turned into an interface for brevity, but in real life
 	// I would personally recommend creating a real type for your specific needs
 	// instead of passing adhoc closures. YMMV.
-	pp := func(ctx jwe.DecryptCtx) error {
-		msg := ctx.Message()
+	kp := func(ctx context.Context, sink jwe.KeySink, _ jwe.Recipient, msg *jwe.Message) error {
 		rawhint, _ := msg.ProtectedHeaders().Get(`jwx-hints`)
 		//nolint:forcetypeassert
 		hint, ok := rawhint.(string)
@@ -113,7 +113,7 @@ func ExampleJWE_ComplexDecrypt() {
 			// You may opt to set both the algorithm and key here as well.
 			// BUT BE CAREFUL so that you don't accidentally create a
 			// vulnerability
-			ctx.SetKey(privkey)
+			sink.Key(jwa.RSA_OAEP, privkey)
 			return nil
 		}
 
@@ -124,7 +124,7 @@ func ExampleJWE_ComplexDecrypt() {
 	// Calling jwe.Decrypt with the extra argument of jwe.WithPostParser().
 	// Here we pass a nil key to jwe.Decrypt, because the PostParser will be
 	// determining the key to use when its PostParse() method is called
-	decrypted, err := jwe.Decrypt(encrypted, jwa.RSA_OAEP, nil, jwe.WithPostParser(jwe.PostParseFunc(pp)))
+	decrypted, err := jwe.Decrypt(encrypted, jwe.WithKeyProvider(jwe.KeyProviderFunc(kp)))
 	if err != nil {
 		fmt.Printf("failed to decrypt message: %s\n", err)
 		return
diff --git a/internal/jwxtest/jwxtest.go b/internal/jwxtest/jwxtest.go
index 65263f06..f5b19c1d 100644
--- a/internal/jwxtest/jwxtest.go
+++ b/internal/jwxtest/jwxtest.go
@@ -276,7 +276,7 @@ func DecryptJweFile(ctx context.Context, file string, alg jwa.KeyEncryptionAlgor
 		return nil, errors.Wrap(err, `failed to obtain raw key from JWK`)
 	}
 
-	return jwe.Decrypt(buf, alg, rawkey)
+	return jwe.Decrypt(buf, jwe.WithKey(alg, rawkey))
 }
 
 func EncryptJweFile(ctx context.Context, payload []byte, keyalg jwa.KeyEncryptionAlgorithm, keyfile string, contentalg jwa.ContentEncryptionAlgorithm, compressalg jwa.CompressionAlgorithm) (string, func(), error) {
diff --git a/jwe/gh402_test.go b/jwe/gh402_test.go
index c82aa2c7..19d81d7d 100644
--- a/jwe/gh402_test.go
+++ b/jwe/gh402_test.go
@@ -1,6 +1,7 @@
 package jwe_test
 
 import (
+	"context"
 	"testing"
 
 	"github.com/lestrrat-go/jwx/internal/json"
@@ -56,11 +57,6 @@ func TestGH402(t *testing.T) {
 	key := []byte{195, 170, 42, 171, 98, 176, 98, 162, 57, 170, 62, 69, 175, 209, 200, 151, 81, 135, 63, 43, 93, 20, 16, 111, 13, 26, 138, 188, 15, 19, 26, 242}
 	data := "eyJhbGciOiJkaXIiLCJjbGV2aXMiOnsicGluIjoic3NzIiwic3NzIjp7Imp3ZSI6WyJleUpoYkdjaU9pSkZRMFJJTFVWVElpd2lZMnhsZG1seklqcDdJbkJwYmlJNkluUmhibWNpTENKMFlXNW5JanA3SW1Ga2RpSTZleUpyWlhseklqcGJleUpoYkdjaU9pSkZRMDFTSWl3aVkzSjJJam9pVUMwMU1qRWlMQ0pyWlhsZmIzQnpJanBiSW1SbGNtbDJaVXRsZVNKZExDSnJkSGtpT2lKRlF5SXNJbmdpT2lKQlZqbEplVmszZDNCamNtMVpZMDkxZWsxTE9VSjJaMDFDZDNJd09FNXhURXBtUW5Ca1IzVXpNa1ZDWjJORmRITnpNRGswVEdweFdXSlFhVmhQTFhWTWMyeDVValZUVlRsdlNHbGhaRlI1VlhWb1UzUTJhbXQzSWl3aWVTSTZJa0ZIVkRseVpYcHlkV0pITFU1eFgxWkxPRUl3UXpGTE9WSTJNV05YU1U1UFNVaEZSMU5MYlZKclUxZEtjVEJRWHpkaFVHaDJlV3hhWDNaTFdHSk5TMGhrU1d0VFFrd3dSMWRRTmxaaVdqbExZVEJHVkZGSFpqTWlmU3g3SW1Gc1p5STZJa1ZUTlRFeUlpd2lZM0oySWpvaVVDMDFNakVpTENKclpYbGZiM0J6SWpwYkluWmxjbWxtZVNKZExDSnJkSGtpT2lKRlF5SXNJbmdpT2lKQlNIZFFhR2hPZDBac1ptc3RNR2d6VDNGS1JFVXpkekpVWDJGcExUbGlZWGhXY0VKa05TMWtNbXhIU0VseFgyTXhVSEZrTW1RMlYxSkRXbEZVYkc5cWQyUmhOblJUUkVndFUxTmhaVE5aTmxGVlZIVkVla1puSWl3aWVTSTZJa0ZrYVhkdVozaDVlamt4UVMxUFRGODFhekIwVEcwMlZtSTNSVkZxWVcxQ1NWUlhlVzVuUldSRVpuRlRSbm8xWlRkdVkycFdlR3BpWDFCSVMwbFNaRzFPT1hCa2RFMTRlRUpPYTFkV2RsOWZUbkZxVTI5SU5UUWlmVjE5TENKMWNtd2lPaUpvZEhSd09pOHZiRzlqWVd4b2IzTjBPak16TURJeEluMTlMQ0psYm1NaU9pSkJNalUyUjBOTklpd2laWEJySWpwN0ltTnlkaUk2SWxBdE5USXhJaXdpYTNSNUlqb2lSVU1pTENKNElqb2lRVWR1YjNrMFQyMTBhRWt3TjNKV1QxbFNTVGhxUTIxVGMxQjZXVVYwUkcwNWNqRkRaa0ptVmxCa2RscEdVVE0wY2taVmJFTTFRbGRrVlZweVMyOXNYMHh5UVhsdFlrY3pWVEZhY1VRMU9YSnBlVGQzV1hGWVV5SXNJbmtpT2lKQlNHaHNSVGx0T1hoR04xVkVVVXd4WjJSd2RGaE5Zbm94VUhNMloweDRObWx1ZGpodU4xVXpUamhpU0dSTlJIRldhalZKYVdGVGFGbDRZM1kyYzFSQ2VuUXpNMmxrUlRaZlYyUm1Oell3WVhkV1VVUkRXa1oxSW4wc0ltdHBaQ0k2SWpWQlNVaHJjblpDYzFOZmQzRjFMVGN5TUd0Q1dtMTBWblpxYkdsVGVWSmpUVXREVFU5dmJuWlFSVWtpZlEuLlR0SzJzWTE2bDJnSmx6c18uRFZqcGVkbGUxY29xZkFkLUZ3bnM5RmtRVkR0ak4zWU1NWjc4VjNLVmI3VWpFd2p1eTBXYjZ0eUxzTmZsVnJaVU85dDM1X2pfSlpobU0ycllfV1RyOHcuVkZMYjFvcTFqZzRsdnJpY3laaXF6USIsImV5SmhiR2NpT2lKRlEwUklMVVZUSWl3aVkyeGxkbWx6SWpwN0luQnBiaUk2SW5SaGJtY2lMQ0owWVc1bklqcDdJbUZrZGlJNmV5SnJaWGx6SWpwYmV5SmhiR2NpT2lKRlEwMVNJaXdpWTNKMklqb2lVQzAxTWpFaUxDSnJaWGxmYjNCeklqcGJJbVJsY21sMlpVdGxlU0pkTENKcmRIa2lPaUpGUXlJc0luZ2lPaUpCVUZWTVMxVllhMUp2TjJOdGFXOXFVRUV3ZFVZMGVIcFFXakJ0V21sQ1NFeFBhVWRyV0hkQ1pGTkVibVpNV1ZaTmJFeFZWR0UzTXkxNWJYSnVWME5mUlRZdE9DMWpUbnA0U1RaMWVtMTRhVFphTXpCUU1EaHhJaXdpZVNJNklrRlBPVTFWTVdGQmRWRnJNVEJrY1dGbVIwRlFYMlpKVGkwelRXTldNRkpGT0drNGFVdG5UbGhuWm14SlFtbDZNRGMxVmpCWVJIZEZNbkpSU2xka1VsRktlWGczT0hFeVJrTndkekZGTTFCMk1YRnZZMmhCZDJnaWZTeDdJbUZzWnlJNklrVlROVEV5SWl3aVkzSjJJam9pVUMwMU1qRWlMQ0pyWlhsZmIzQnpJanBiSW5abGNtbG1lU0pkTENKcmRIa2lPaUpGUXlJc0luZ2lPaUpCVUUxSFZYWTFjSFZUZWt0MU1HRjRZVFJRYldOeFRsQkxiMjR0TW1oYVVGSjJXbmRwV2xkc05sOVNNREJ2YlZsMldESXplRlJ4UldGUVdFMXphemxFYm1SdGNrUnpWbEZvTlRFMFEzTTBSbGszV0VaTlkyUkVJaXdpZVNJNklrRlZNa051VUV0RlJVeDBhSFZaYlZaeWJWZEJORFZLYjNocE1FRkhjM2w1T0daak5UTlJiMHA2UjA4NWNYVnRXVVpRTTFWbVoyWjZibkZsTnpWcE5uUXpkSFJPUW0xSWNGVlhTa1JuUkRkdGRGbEdVbXRzYzFjaWZWMTlMQ0oxY213aU9pSm9kSFJ3T2k4dmJHOWpZV3hvYjNOME9qUXhPRFl4SW4xOUxDSmxibU1pT2lKQk1qVTJSME5OSWl3aVpYQnJJanA3SW1OeWRpSTZJbEF0TlRJeElpd2lhM1I1SWpvaVJVTWlMQ0o0SWpvaVFWazJVMmxoUzFaRFMwUnFRMjUwTFdWNFJFZzJTa3RhYVdkaWRXOVpZVzV3Y0hWRVkxOXNOVXBxY1ROR05GWnpMV3hJYzFGbVYzQmFOVXRPWTNCUWFITmZObFpVUmt4ZlYwMVdja2xCZG5wWWVFeExZMmhKY0NJc0lua2lPaUpCU1RCVFZVSTFUMG80ZGtkR05YRndXRzVFZVRkd1pteFBNVjlsZW14WWQyTTNValpxY25RelJrSklPRTEwZG0xcFRHVkZNRzR0VEhZNFMyNWZMVWMzUlhJM1MwRklNSEoyWW5kMlVFZzBjbXhrU1VrM2RETmxJbjBzSW10cFpDSTZJa1JNT1ZaclRHcFVaVXBKTVhRd2ExUnNUbnBzVGxrM1RHcDFTWGxLZHpGUE1HTkdYMHN0YW5GdGFHOGlmUS4uSjJ4c2VjeTN3c2FGRjFfdi5uQzdPSzVrdXAwbkgxVGwwOFBRUE82QTllNjFYMXNGY1RFUVZjTjNkd0NWYTFMZ2p2STFfSWZqTU5TX2cweTI4WGM1ZThQZ2hFS01UQ0hkQk44UzFBdy5yV0s5Qi0yT2RtWUlIbmQzMzcxeHpRIiwiZXlKaGJHY2lPaUpGUTBSSUxVVlRJaXdpWTJ4bGRtbHpJanA3SW5CcGJpSTZJblJoYm1jaUxDSjBZVzVuSWpwN0ltRmtkaUk2ZXlKclpYbHpJanBiZXlKaGJHY2lPaUpGUTAxU0lpd2lZM0oySWpvaVVDMDFNakVpTENKclpYbGZiM0J6SWpwYkltUmxjbWwyWlV0bGVTSmRMQ0pyZEhraU9pSkZReUlzSW5naU9pSkJRbVZ1WVRaWVVUQkNTMmQwU1U5c05FMXZUakZEZFhKc2REVk1lVEZ6YUZNMFRFbHdkV05aUkVaWE9VNDVNRWgxY1RSUk4wcENPR3RNVDNabVdEUnBhSFJwTFdWQ1JXdGpVR3BCYzNWWU9HWm5kbVEyVEZSM0lpd2llU0k2SWtGUlRFWnVZMmhmZVRCaVdFNTRTRGt6UW10NWQwSmFhbXB6T1ROaFptOXFXWGhyV0d0dFFXZFRVbnA1ZGxaWGFrbHJRMEYwVWtoRVZuVmhlbkUxTWtSRFYxVnRlRU5YWm1OR05qWkdSaTB5Y1ZrellsTk9Na01pZlN4N0ltRnNaeUk2SWtWVE5URXlJaXdpWTNKMklqb2lVQzAxTWpFaUxDSnJaWGxmYjNCeklqcGJJblpsY21sbWVTSmRMQ0pyZEhraU9pSkZReUlzSW5naU9pSkJaR2hHZW05RE5FZG5WRk5CWlUxTFVYSlhjRWRCTWtndFEyeEVNVE52U1d4SVprczFTWEI2ZVZseWRuVm1Sbk5DZDJkYWJWYzRWR0ZXUkdSYVltbGlkM3BtTldSMk1rOHRkVkpmU2pSS2IyNDRNekJ0ZFVGa0lpd2llU0k2SWtGUlpFMVdka3RGWW5obWMwbHRiVFJtVHpZMlFsOHdVMnhPWDBaRmVEQkthVlpyVjFKQlFrRnBhRVExV1VKUWVVbEpjM1pGUjJoR1dqaFpOVXBWTXpoQ1JWVlJUWGhyVkZsT1pEVlpOR1UxVG1STVVVbDNObDhpZlYxOUxDSjFjbXdpT2lKb2RIUndPaTh2Ykc5allXeG9iM04wT2pNME1qVTVJbjE5TENKbGJtTWlPaUpCTWpVMlIwTk5JaXdpWlhCcklqcDdJbU55ZGlJNklsQXROVEl4SWl3aWEzUjVJam9pUlVNaUxDSjRJam9pUVZaU1praEpWVVJ2U0RCb1FYRnhOSFpJYVVGMmNVOVNXblJUVFZkS2VGVmtPWEozV2tWdVZFMUVWbWt4VG5KbGRWTnZWWHBNWmxkcFVVcDBiR3BxU21ZM1NqRnpUbGRVUmpKaVFtTnJNSEJrUnpkcFJHVnhVeUlzSW5raU9pSkJSak5aUms5dlNXNW1jVU5yVFRKa1pISnlkV1JyY21zMmN6Wk1SMUpNU0U5YVl6QlhhMmRhZFZsb1FsRk9ZbVZXY25sWFZTMW9aMTgwYkhCRE5EUjRaRmh3ZUc4NWFGOWhSVlZ0U0V0UGFWOWpVM2xhY1VjM0luMHNJbXRwWkNJNkluQnNablF4UW5WR2EyeGZaRGhyZFZOaWNVVlVSbXhxYjJwM01WZDBMVFpLYTJ0VGVXOWtWa1JEZDNjaWZRLi5CUUVtRXFhc1FHNHNKY1d2Lmo5RzBweVEwbTY0TU5OTEZicm4wT01BbmZaOTJ5bWNxV19ZM2ZIMWJEd2JFbnI1U1FHWF9jWjNfbUpGSS0zMFVJOEZuZzRUaXRxMU9JT2JySktETTh3LndRTmVJVzhzZnNua1V1aTUwTWpYd0EiLCJleUpoYkdjaU9pSkZRMFJJTFVWVElpd2lZMnhsZG1seklqcDdJbkJwYmlJNkluUmhibWNpTENKMFlXNW5JanA3SW1Ga2RpSTZleUpyWlhseklqcGJleUpoYkdjaU9pSkZRMDFTSWl3aVkzSjJJam9pVUMwMU1qRWlMQ0pyWlhsZmIzQnpJanBiSW1SbGNtbDJaVXRsZVNKZExDSnJkSGtpT2lKRlF5SXNJbmdpT2lKQlJVdzRjbkZ3UjE4dE9FRXRiM1UwYTJaMFJXbFNjbVZ5VldOb1ozWmlVRzFXYUV0bVJVNWtXVXd0ZDNOVlRucGZNMlpRZW5GUVVtWkhiRzltTkZkR2RFMTNNM0I1YnpaSVFUaFFTM0JHVlc5VWJ6ZHRjbWRSSWl3aWVTSTZJa0ZDUWxsVU9WOWpWRnBuT1RCWlJXWk9jWFV4V2tkSGFtaFphMWxwTTFkWlRqWmpialZJUmtjdGJ5MUdOR3BIWWpkVFpEazFTbWsxTkhkVWRUQmxOVmM1UjE5cGVFRlZYMWh0Wm1GSmQyZFRXVXN3WDJScVpEVWlmU3g3SW1Gc1p5STZJa1ZUTlRFeUlpd2lZM0oySWpvaVVDMDFNakVpTENKclpYbGZiM0J6SWpwYkluWmxjbWxtZVNKZExDSnJkSGtpT2lKRlF5SXNJbmdpT2lKQlZYSktNME5GUzFZd2FGWlNNamhLWTE5eWRXUkpUWEF0YmxCUmRtdFNSemxGVlZsVVRrTjJiVEZrY2tWUlJVNXVWMjB3VXpNelMycEtTMWx1VFhObVJrTlRPR3BYV0hJMFkxUTNTMFV4V2xkWmFGRnJkRzF3SWl3aWVTSTZJa0ZrYUVkWGFFVjJNREE1VjJwdFRUQjFZWFpFTm5sR01HczNMVXR0WkRCTmNHNTNRelZtWVVReFZuTkhWSGhLVG5SYVptSmlWbGhGVTIxeVJYZ3pYMVJoVERZMFpVeFhObFZwTFVveWJscHZhbEpUU0hWWmFtTWlmVjE5TENKMWNtd2lPaUpvZEhSd09pOHZiRzlqWVd4b2IzTjBPalF6TVRVMUluMTlMQ0psYm1NaU9pSkJNalUyUjBOTklpd2laWEJySWpwN0ltTnlkaUk2SWxBdE5USXhJaXdpYTNSNUlqb2lSVU1pTENKNElqb2lRV1JrT0Rnd05FdFlRbGhuWmpacU1EQkdVbEIyUzBsWmFGaE5ZMnRpWjE5c1NXbHdhbVF6TlhGUk9FVk9MVEZGUTNOUVZrNHRTa1I1ZW1sVmVuSkhjVTVOUkRGSmJGOUlTMUIyVGpaUmFVMU9kVXB0VVcxcWJpSXNJbmtpT2lKQlF5MW9Tak5DUTBvMU9VSnZhbWR1Ym5oZlZuaHBWekZNTVU5UlYxTXRZMmxzV2xkRGVIaGhZVUpyYjA0MVV6RnpiMnh4Y1ZwbVFuQkNiakZSWW5OTlJEZGZiVFEwVm1weGRWQlFjMjQwVmpFNFFqQlZkV2hDSW4wc0ltdHBaQ0k2SWtaUmNWZDBUWGQ0V2pWbVEyWjFSekp1WldWT1lreEpkQzF6Tm1Ka2RUQXpaMUZSU0Vvd00wY3RNMDBpZlEuLmduekZoT21WQk41UFNWT2kuZlBkQWJ2YlgtLUFXV1pPREZEdy1Yd2VvZHJQamc5UGM4SlFQQjIwZTlPQ2Y5RTZmSEtNYUJlR2xVSWRyOHViRjdRUjllM3Vnalk2ZEdLOEoyeklyYXcuTlV6T3M4a01HVVhMa3BXNWxja3BnZyIsImV5SmhiR2NpT2lKRlEwUklMVVZUSWl3aVkyeGxkbWx6SWpwN0luQnBiaUk2SW5SaGJtY2lMQ0owWVc1bklqcDdJbUZrZGlJNmV5SnJaWGx6SWpwYmV5SmhiR2NpT2lKRlEwMVNJaXdpWTNKMklqb2lVQzAxTWpFaUxDSnJaWGxmYjNCeklqcGJJbVJsY21sMlpVdGxlU0pkTENKcmRIa2lPaUpGUXlJc0luZ2lPaUpCWTNVeWFFOTVXSGRPVG1OYVpXRnJkVTFRVW01aGFqSnJXbGxJUldkMVFuRmZiR0ZMU0VZeFUwbzFjMVp3UmswMVVsZFJYMVp5U2pKVmVtMHpXa0k0T1Y5Nk5HZExVVzVETFU1TFJHTk5NMHcxUkVsdlh6UlNJaXdpZVNJNklrRmlOVXBTU1hOeU5HSm9kRXhvV1dabVVITlpaMUZqT0hocGNXcFplak5wUm1GeVpIbEVkVFZRZHpSVFRIZExWR1JuV2tSMGNtODVNSHByTVRKVFpERmhYemgxZVRkV1drMHRTRVkzUWs5bGNXeG9RV3R2VEdnaWZTeDdJbUZzWnlJNklrVlROVEV5SWl3aVkzSjJJam9pVUMwMU1qRWlMQ0pyWlhsZmIzQnpJanBiSW5abGNtbG1lU0pkTENKcmRIa2lPaUpGUXlJc0luZ2lPaUpCVm1Sa1ZFWklNRWxwTm1OWk1sQk9XVVp1WVZkT2IwaENlRzVuVFZaR1kzZFlObVYzU25vd1VVUlZjbkF5UmxsT1pVNHpaQzFxYWxwcVdYQmZkRkJ0Ukdjd05YcG1iSHAyU0VoSWJscFlUM2RLT1RSSU5VOHhJaXdpZVNJNklrRllaRkZ5VDBFdFJXWjFTM2xHU0drdE5sVlFSRGMwUnkxaVdVdFhNbFV4YVZselpFTkRSakp5WWs5T1JXeDBTVFowT0ZSZk56bHZjR1pLTkcxU1gySnliMVZqVVhBNU1GOWFZM05RY3poZlZtRndSbmhyVmpJaWZWMTlMQ0oxY213aU9pSm9kSFJ3T2k4dmJHOWpZV3hvYjNOME9qTTBPRGM1SW4xOUxDSmxibU1pT2lKQk1qVTJSME5OSWl3aVpYQnJJanA3SW1OeWRpSTZJbEF0TlRJeElpd2lhM1I1SWpvaVJVTWlMQ0o0SWpvaVFXUkNVVlJKUkdjNGNEWm5NVzh4T0d4R1dYRm1kMXBhTWxaaFJteG9VMlF0YVhkV09EQnRXalZaWm5wWU5qbFNWV3BxYzNKc1ZHZEtjbXMwY2pKaWRrSlZUbE5YY0cxbVEyeHBYek0xYm1Wb2NrZHpRbW94VUNJc0lua2lPaUpCVTJGVVYxY3RSRk5ZVVhObFgxUmxibDg0VkdkVmFrWTNVRWxRTm5wcVpUUmxSbk55YkY5WFRESkhNRGwzTFROalYzVlRlVlZmUzFvemF6VndTVlp0TlZkaU1tcGlXREptV2sxWmVFNVpRbDh5Vm10d1QxWnVJbjBzSW10cFpDSTZJbEk0YWtWSk1YQlRNVWxoY0hWVVZHTktObGhJWkROWFUzbGZRbmhXU1hWb1VVeDFVekJwYjJkU09UUWlmUS4uTEk2ZC1TNUlYSDRIY2lxeS5ESWl3M1hFeXlaMVZWZDFPcWYwS1pJeGZrUWtnZXo4bXl2N0dfZGFMOWxIdjVFVlFMQzREVG9PMWF6UDQ5SDNGUk5PRHZIcXBmaWp1a3VDX2JmLVZldy41NGpyZnZWdEYxR3ItYU9TVEFjSWV3Il0sInAiOiI0NERteG1EZVFiWTJ3eHpIYm9PODFkbTBCbENNcVJzeEtLQ0NDQWo4TVVjIiwidCI6M319LCJlbmMiOiJBMjU2R0NNIn0..zz3fUXsiaME2cSoy.LTQovHUvDP4MXT2_sHgf_cM2gicobD5kGXEl5eY.MK3Lf6IwaoVUvCTp1Q5VOA"
 
-	pp := jwe.PostParseFunc(func(ctx jwe.DecryptCtx) error {
-		ctx.SetKey(key)
-		ctx.SetAlgorithm(jwa.DIRECT)
-		return nil
-	})
 	decrypt := func(customField bool) {
 		t.Helper()
 		m := jwe.NewMessage()
@@ -68,20 +64,28 @@ func TestGH402(t *testing.T) {
 		plain, err := jwe.Decrypt([]byte(data),
 			// This is a really cheesy way of creating a jwa.KeyEncryptionAlgorithm
 			// but a bogus one.
-			jwa.KeyEncryptionAlgorithm("invalid algorithm"),
-			nil,
+			jwe.WithKey(jwa.KeyEncryptionAlgorithm("invalid algorithm"), nil),
 			jwe.WithMessage(m),
-			jwe.WithPostParser(pp),
+			jwe.WithKeyProvider(jwe.KeyProviderFunc(func(_ context.Context, sink jwe.KeySink, _ jwe.Recipient, _ *jwe.Message) error {
+				sink.Key(jwa.DIRECT, key)
+				return nil
+			})),
 		)
+
 		if !assert.NoError(t, err, `jwe.Decrypt should succeed`) {
 			return
 		}
+
 		if string(plain) != "testing Shamir Secret Sharing" {
 			t.Errorf("expected 'testing Shamir Secret Sharing', got %s", string(plain))
 			return
 		}
 
 		if customField {
+			if !assert.NotNil(t, m.ProtectedHeaders(), `m.ProtectedHeaders should be non-nil`) {
+				return
+			}
+
 			v, ok := m.ProtectedHeaders().Get("clevis")
 			if !assert.True(t, ok, `m.Get("clevis") should be true`) {
 				return
diff --git a/jwe/jwe_test.go b/jwe/jwe_test.go
index 7cdcc9c6..c00f5ae3 100644
--- a/jwe/jwe_test.go
+++ b/jwe/jwe_test.go
@@ -126,7 +126,7 @@ func TestParse(t *testing.T) {
 			return
 		}
 
-		buf, err := jwe.JSON(msg)
+		buf, err := json.Marshal(msg)
 		if !assert.NoError(t, err, "Serializing to JSON format should succeed") {
 			return
 		}
@@ -170,7 +170,7 @@ func TestParse_RSAES_OAEP_AES_GCM(t *testing.T) {
 	}
 
 	msg := jwe.NewMessage()
-	plaintext, err := jwe.Decrypt([]byte(serialized), jwa.RSA_OAEP, rawkey, jwe.WithMessage(msg))
+	plaintext, err := jwe.Decrypt([]byte(serialized), jwe.WithKey(jwa.RSA_OAEP, rawkey), jwe.WithMessage(msg))
 	if !assert.NoError(t, err, "jwe.Decrypt should be successful") {
 		return
 	}
@@ -253,7 +253,7 @@ func TestParse_RSAES_OAEP_AES_GCM(t *testing.T) {
 			}
 			t.Logf("%s", encrypted)
 
-			plaintext, err = jwe.Decrypt(encrypted, jwa.RSA_OAEP, rawkey)
+			plaintext, err = jwe.Decrypt(encrypted, jwe.WithKey(jwa.RSA_OAEP, rawkey))
 			if !assert.NoError(t, err, "jwe.Decrypt should succeed") {
 				return
 			}
@@ -304,7 +304,7 @@ func TestRoundtrip_RSAES_OAEP_AES_GCM(t *testing.T) {
 			return
 		}
 
-		decrypted, err := jwe.Decrypt(encrypted, jwa.RSA_OAEP, rsaPrivKey)
+		decrypted, err := jwe.Decrypt(encrypted, jwe.WithKey(jwa.RSA_OAEP, rsaPrivKey))
 		if !assert.NoError(t, err, "Decrypt should succeed") {
 			return
 		}
@@ -332,7 +332,7 @@ func TestRoundtrip_RSA1_5_A128CBC_HS256(t *testing.T) {
 			return
 		}
 
-		decrypted, err := jwe.Decrypt(encrypted, jwa.RSA1_5, rsaPrivKey)
+		decrypted, err := jwe.Decrypt(encrypted, jwe.WithKey(jwa.RSA1_5, rsaPrivKey))
 		if !assert.NoError(t, err, "Decrypt successful") {
 			return
 		}
@@ -365,7 +365,7 @@ func TestEncode_A128KW_A128CBC_HS256(t *testing.T) {
 			return
 		}
 
-		decrypted, err := jwe.Decrypt(encrypted, jwa.A128KW, sharedkey)
+		decrypted, err := jwe.Decrypt(encrypted, jwe.WithKey(jwa.A128KW, sharedkey))
 		if !assert.NoError(t, err, "Decrypt successful") {
 			return
 		}
@@ -400,7 +400,7 @@ func testEncodeECDHWithKey(t *testing.T, privkey interface{}, pubkey interface{}
 				return
 			}
 
-			decrypted, err := jwe.Decrypt(encrypted, alg, privkey)
+			decrypted, err := jwe.Decrypt(encrypted, jwe.WithKey(alg, privkey))
 			if !assert.NoError(t, err, "Decrypt succeeds") {
 				return
 			}
@@ -534,7 +534,7 @@ func TestEncode_Direct(t *testing.T) {
 			if !assert.NoError(t, err, `jwe.Encrypt should succeed`) {
 				return
 			}
-			decrypted, err := jwe.Decrypt(encrypted, jwa.DIRECT, key)
+			decrypted, err := jwe.Decrypt(encrypted, jwe.WithKey(jwa.DIRECT, key))
 			if !assert.NoError(t, err, `jwe.Decrypt should succeed`) {
 				return
 			}
@@ -613,7 +613,7 @@ func TestDecodePredefined_Direct(t *testing.T) {
 				return
 			}
 
-			decrypted, err := jwe.Decrypt([]byte(tc.Data), jwa.DIRECT, key)
+			decrypted, err := jwe.Decrypt([]byte(tc.Data), jwe.WithKey(jwa.DIRECT, key))
 			if !assert.NoError(t, err, `jwe.Decrypt should succeed`) {
 				return
 			}
diff --git a/jwe/lowlevel_test.go b/jwe/lowlevel_test.go
deleted file mode 100644
index 2319ee3d..00000000
--- a/jwe/lowlevel_test.go
+++ /dev/null
@@ -1,125 +0,0 @@
-package jwe
-
-import (
-	"testing"
-
-	"github.com/lestrrat-go/jwx/jwa"
-	"github.com/lestrrat-go/jwx/jwe/internal/cipher"
-	"github.com/lestrrat-go/jwx/jwe/internal/keyenc"
-	"github.com/lestrrat-go/jwx/jwe/internal/keygen"
-	"github.com/stretchr/testify/assert"
-)
-
-// This test uses Appendix 3 to verify some low level tools for
-// KeyWrap and CBC HMAC encryption.
-// This test uses a static cek so that we can validate the results
-// against the contents in the above Appendix
-func TestLowLevelParts_A128KW_A128CBCHS256(t *testing.T) {
-	var plaintext = []byte{
-		76, 105, 118, 101, 32, 108, 111, 110, 103, 32, 97, 110, 100, 32,
-		112, 114, 111, 115, 112, 101, 114, 46,
-	}
-	var cek = []byte{
-		4, 211, 31, 197, 84, 157, 252, 254, 11, 100, 157, 250, 63, 170, 106,
-		206, 107, 124, 212, 45, 111, 107, 9, 219, 200, 177, 0, 240, 143, 156,
-		44, 207,
-	}
-	var iv = []byte{
-		3, 22, 60, 12, 43, 67, 104, 105, 108, 108, 105, 99, 111, 116, 104,
-		101,
-	}
-	var sharedkey = []byte{
-		25, 172, 32, 130, 225, 114, 26, 181, 138, 106, 254, 192, 95, 133, 74, 82,
-	}
-	var encsharedkey = []byte{
-		232, 160, 123, 211, 183, 76, 245, 132, 200, 128, 123, 75, 190, 216,
-		22, 67, 201, 138, 193, 186, 9, 91, 122, 31, 246, 90, 28, 139, 57, 3,
-		76, 124, 193, 11, 98, 37, 173, 61, 104, 57,
-	}
-	var aad = []byte{
-		101, 121, 74, 104, 98, 71, 99, 105, 79, 105, 74, 66, 77, 84, 73, 52,
-		83, 49, 99, 105, 76, 67, 74, 108, 98, 109, 77, 105, 79, 105, 74, 66,
-		77, 84, 73, 52, 81, 48, 74, 68, 76, 85, 104, 84, 77, 106, 85, 50, 73,
-		110, 48,
-	}
-	var ciphertext = []byte{
-		40, 57, 83, 181, 119, 33, 133, 148, 198, 185, 243, 24, 152, 230, 6,
-		75, 129, 223, 127, 19, 210, 82, 183, 230, 168, 33, 215, 104, 143,
-		112, 56, 102,
-	}
-	var authtag = []byte{
-		83, 73, 191, 98, 104, 205, 211, 128, 201, 189, 199, 133, 32, 38,
-		194, 85,
-	}
-
-	const compactExpected = `eyJhbGciOiJBMTI4S1ciLCJlbmMiOiJBMTI4Q0JDLUhTMjU2In0.6KB707dM9YTIgHtLvtgWQ8mKwboJW3of9locizkDTHzBC2IlrT1oOQ.AxY8DCtDaGlsbGljb3RoZQ.KDlTtXchhZTGufMYmOYGS4HffxPSUrfmqCHXaI9wOGY.U0m_YmjN04DJvceFICbCVQ`
-
-	k, err := keyenc.NewAES(jwa.A128KW, sharedkey)
-	if !assert.NoError(t, err, "Create key wrap") {
-		return
-	}
-
-	enckey, err := k.Encrypt(cek)
-	if !assert.NoError(t, err, "Failed to encrypt key") {
-		return
-	}
-	if !assert.Equal(t, encsharedkey, enckey.Bytes(), "encrypted keys match") {
-		return
-	}
-
-	cipher, err := cipher.NewAES(jwa.A128CBC_HS256)
-	if !assert.NoError(t, err, "NewAesContentCipher is successful") {
-		return
-	}
-	cipher.NonceGenerator = keygen.Static(iv)
-
-	iv, encrypted, tag, err := cipher.Encrypt(cek, plaintext, aad)
-	if !assert.NoError(t, err, "encrypt() successful") {
-		return
-	}
-
-	if !assert.Equal(t, ciphertext, encrypted, "Generated cipher text does not match") {
-		return
-	}
-
-	if !assert.Equal(t, tag, authtag, "Generated tag text does not match") {
-		return
-	}
-
-	data, err := cipher.Decrypt(cek, iv, encrypted, tag, aad)
-	if !assert.NoError(t, err, "decrypt successful") {
-		return
-	}
-
-	if !assert.Equal(t, plaintext, data, "decrypt works") {
-		return
-	}
-
-	r := NewRecipient()
-	r.Headers().Set(AlgorithmKey, jwa.A128KW)
-	r.SetEncryptedKey(enckey.Bytes())
-
-	protected := NewHeaders()
-	protected.Set(ContentEncryptionKey, jwa.A128CBC_HS256)
-
-	msg := NewMessage()
-	msg.Set(ProtectedHeadersKey, protected)
-	msg.Set(AuthenticatedDataKey, aad)
-	msg.Set(CipherTextKey, ciphertext)
-	msg.Set(InitializationVectorKey, iv)
-	msg.Set(TagKey, tag)
-	msg.Set(RecipientsKey, []Recipient{r})
-
-	serialized, err := Compact(msg)
-	if !assert.NoError(t, err, "compact serialization is successful") {
-		return
-	}
-
-	if !assert.Equal(t, compactExpected, string(serialized), "compact serialization matches") {
-		serialized, err = JSON(msg, WithPrettyFormat(true))
-		if !assert.NoError(t, err, "JSON serialization is successful") {
-			return
-		}
-		t.Logf("%s", serialized)
-	}
-}
diff --git a/jwt/jwt_test.go b/jwt/jwt_test.go
index a03992f4..570303cb 100644
--- a/jwt/jwt_test.go
+++ b/jwt/jwt_test.go
@@ -1297,7 +1297,7 @@ func TestNested(t *testing.T) {
 
 	// First layer should be JWE
 	jweMessage := jwe.NewMessage()
-	decrypted, err := jwe.Decrypt(serialized, jwa.RSA_OAEP, key, jwe.WithMessage(jweMessage))
+	decrypted, err := jwe.Decrypt(serialized, jwe.WithKey(jwa.RSA_OAEP, key), jwe.WithMessage(jweMessage))
 	if !assert.NoError(t, err, `jwe.Decrypt should succeed`) {
 		return
 	}