Repository: lestrrat-go/jwx
Pull request #674 (2022-03-29 23:51:56)

Add jwe.WithKeyUsed


Patch:
diff --git a/Changes.v2 b/Changes.v2
index 573b7108..082a9261 100644
--- a/Changes.v2
+++ b/Changes.v2
@@ -109,6 +109,10 @@ Changes
   * [jwe] `jwe.ReadFile` now supports the option `jwe.WithFS` which allows you to
     read data from arbitrary `fs.FS` objects
 
+  * [jwe] jwe.WithKeyUsed has been added to allow users to retrieve
+    the key used for decryption. This is useful in cases you provided
+    multiple keys and you want to know which one was successful
+
   # JWK
 
   * [jwk] `jwk.New()` has been renamed to `jwk.FromRaw()`, which hopefully will
diff --git a/jwe/jwe.go b/jwe/jwe.go
index 7a14fcbb..4993cb40 100644
--- a/jwe/jwe.go
+++ b/jwe/jwe.go
@@ -12,6 +12,7 @@ import (
 	"io"
 	"io/ioutil"
 
+	"github.com/lestrrat-go/blackmagic"
 	"github.com/lestrrat-go/jwx/v2/internal/base64"
 	"github.com/lestrrat-go/jwx/v2/internal/json"
 	"github.com/lestrrat-go/jwx/v2/internal/keyconv"
@@ -403,6 +404,7 @@ type decryptCtx struct {
 // `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, options ...DecryptOption) ([]byte, error) {
 	var keyProviders []KeyProvider
+	var keyUsed interface{}
 
 	var dst *Message
 	//nolint:forcetypeassert
@@ -412,6 +414,8 @@ func Decrypt(buf []byte, options ...DecryptOption) ([]byte, error) {
 			dst = option.Value().(*Message)
 		case identKeyProvider{}:
 			keyProviders = append(keyProviders, option.Value().(KeyProvider))
+		case identKeyUsed{}:
+			keyUsed = option.Value()
 		case identKey{}:
 			pair := option.Value().(*withKey)
 			alg, ok := pair.alg.(jwa.KeyEncryptionAlgorithm)
@@ -483,7 +487,7 @@ func Decrypt(buf []byte, options ...DecryptOption) ([]byte, error) {
 
 	var lastError error
 	for _, recipient := range recipients {
-		decrypted, err := dctx.try(ctx, recipient)
+		decrypted, err := dctx.try(ctx, recipient, keyUsed)
 		if err != nil {
 			lastError = err
 			continue
@@ -498,7 +502,7 @@ func Decrypt(buf []byte, options ...DecryptOption) ([]byte, error) {
 	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) {
+func (dctx *decryptCtx) try(ctx context.Context, recipient Recipient, keyUsed interface{}) ([]byte, error) {
 	var tried int
 	var lastError error
 	for i, kp := range dctx.keyProviders {
@@ -522,6 +526,11 @@ func (dctx *decryptCtx) try(ctx context.Context, recipient Recipient) ([]byte, e
 				continue
 			}
 
+			if keyUsed != nil {
+				if err := blackmagic.AssignIfCompatible(keyUsed, key); err != nil {
+					return nil, fmt.Errorf(`failed to assign used key (%T) to %T: %w`, key, keyUsed, err)
+				}
+			}
 			return decrypted, nil
 		}
 	}
diff --git a/jwe/options.yaml b/jwe/options.yaml
index 1c9df790..84f89666 100644
--- a/jwe/options.yaml
+++ b/jwe/options.yaml
@@ -105,4 +105,18 @@ options:
     argument_type: fs.FS
     comment: |
       WithFS specifies the source `fs.FS` object to read the file from.
+  - ident: KeyUsed
+    interface: DecryptOption
+    argument_type: 'interface{}'
+    comment: |
+      WithKeyUsed allows you to specify the `jwe.Decrypt()` function to
+      return the key used for decryption. This may be useful when
+      you specify multiple key sources or if you pass a `jwk.Set`
+      and you want to know which key was successful at decrypting the
+      signature.
+      
+      `v` must be a pointer to an empty `interface{}`. Do not use
+      `jwk.Key` here unless you are 100% sure that all keys that you
+      have provided are instances of `jwk.Key` (remember that the
+      jwx API allows users to specify a raw key such as *rsa.PublicKey)
 
diff --git a/jwe/options_gen.go b/jwe/options_gen.go
index ded925c8..41b43895 100644
--- a/jwe/options_gen.go
+++ b/jwe/options_gen.go
@@ -115,6 +115,7 @@ type identContentEncryptionAlgorithm struct{}
 type identFS struct{}
 type identKey struct{}
 type identKeyProvider struct{}
+type identKeyUsed struct{}
 type identMergeProtectedHeaders struct{}
 type identMessage struct{}
 type identPerRecipientHeaders struct{}
@@ -143,6 +144,10 @@ func (identKeyProvider) String() string {
 	return "WithKeyProvider"
 }
 
+func (identKeyUsed) String() string {
+	return "WithKeyUsed"
+}
+
 func (identMergeProtectedHeaders) String() string {
 	return "WithMergeProtectedHeaders"
 }
@@ -194,6 +199,20 @@ func WithKeyProvider(v KeyProvider) DecryptOption {
 	return &decryptOption{option.New(identKeyProvider{}, v)}
 }
 
+// WithKeyUsed allows you to specify the `jwe.Decrypt()` function to
+// return the key used for decryption. This may be useful when
+// you specify multiple key sources or if you pass a `jwk.Set`
+// and you want to know which key was successful at decrypting the
+// signature.
+//
+// `v` must be a pointer to an empty `interface{}`. Do not use
+// `jwk.Key` here unless you are 100% sure that all keys that you
+// have provided are instances of `jwk.Key` (remember that the
+// jwx API allows users to specify a raw key such as *rsa.PublicKey)
+func WithKeyUsed(v interface{}) DecryptOption {
+	return &decryptOption{option.New(identKeyUsed{}, v)}
+}
+
 // WithMergeProtectedHeaders specify that when given multiple headers
 // as options to `jwe.Encrypt`, these headers should be merged instead
 // of overwritten


Test patch:
diff --git a/jwe/jwe_test.go b/jwe/jwe_test.go
index fc7c659f..468c1c36 100644
--- a/jwe/jwe_test.go
+++ b/jwe/jwe_test.go
@@ -269,7 +269,9 @@ func TestParse_RSAES_OAEP_AES_GCM(t *testing.T) {
 				_ = pkJwk.Set(jwe.AlgorithmKey, jwa.RSA_OAEP)
 				set := jwk.NewSet()
 				set.Add(pkJwk)
-				plaintext, err = jwe.Decrypt(encrypted, jwe.WithKeySet(set, jwe.WithRequireKid(false)))
+
+				var used interface{}
+				plaintext, err = jwe.Decrypt(encrypted, jwe.WithKeySet(set, jwe.WithRequireKid(false)), jwe.WithKeyUsed(&used))
 				if !assert.NoError(t, err, "jwe.Decrypt should succeed") {
 					return
 				}
@@ -277,6 +279,10 @@ func TestParse_RSAES_OAEP_AES_GCM(t *testing.T) {
 				if !assert.Equal(t, payload, string(plaintext), "jwe.Decrypt should produce the same plaintext") {
 					return
 				}
+
+				if !assert.Equal(t, pkJwk, used) {
+					return
+				}
 			})
 		})
 	}
diff --git a/jwe/options_gen_test.go b/jwe/options_gen_test.go
index 27964b62..a1d2590f 100644
--- a/jwe/options_gen_test.go
+++ b/jwe/options_gen_test.go
@@ -14,6 +14,7 @@ func TestOptionIdent(t *testing.T) {
 	require.Equal(t, "WithFS", identFS{}.String())
 	require.Equal(t, "WithKey", identKey{}.String())
 	require.Equal(t, "WithKeyProvider", identKeyProvider{}.String())
+	require.Equal(t, "WithKeyUsed", identKeyUsed{}.String())
 	require.Equal(t, "WithMergeProtectedHeaders", identMergeProtectedHeaders{}.String())
 	require.Equal(t, "WithMessage", identMessage{}.String())
 	require.Equal(t, "WithPerRecipientHeaders", identPerRecipientHeaders{}.String())