allow passing context.Contex to jwe.Decrypt (#1222)

* allow passing context.Contex to jwe.Decrypt * really, this is OK --------- Co-authored-by: Daisuke Maki <[email protected]>
lestrrat-go · Oct 26, 2024 · c1658b3 · c1658b3
1 parent 4769e28
commit c1658b3
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 3 deletions.
diff --git a/jwe/jwe.go b/jwe/jwe.go
@@ -542,6 +542,7 @@ func decrypt(buf []byte, options ...DecryptOption) ([]byte, error) {
 	var cek *[]byte
 	var dst *Message
 	perCallMaxDecompressBufferSize := maxDecompressBufferSize
+	ctx := context.Background()
 	//nolint:forcetypeassert
 	for _, option := range options {
 		switch option.Ident() {
@@ -565,6 +566,9 @@ func decrypt(buf []byte, options ...DecryptOption) ([]byte, error) {
 			cek = option.Value().(*[]byte)
 		case identMaxDecompressBufferSize{}:
 			perCallMaxDecompressBufferSize = option.Value().(int64)
+		case identContext{}:
+			//nolint:fatcontext
+			ctx = option.Value().(context.Context)
 		}
 	}
 
@@ -625,8 +629,6 @@ func decrypt(buf []byte, options ...DecryptOption) ([]byte, error) {
 	dctx.cek = cek
 	dctx.maxDecompressBufferSize = perCallMaxDecompressBufferSize
 
-	ctx := context.TODO()
-
 	errs := make([]error, 0, len(recipients))
 	for _, recipient := range recipients {
 		decrypted, err := dctx.try(ctx, recipient, keyUsed)

diff --git a/jwe/options.yaml b/jwe/options.yaml
@@ -52,6 +52,12 @@ options:
   - ident: KeyProvider
     interface: DecryptOption
     argument_type: KeyProvider
+  - ident: Context
+    interface: DecryptOption
+    argument_type: context.Context
+    comment: |
+      WithContext specifies the context.Context object to use when decrypting a JWE message.
+      If not provided, context.Background() will be used.
   - ident: Serialization
     option_name: WithCompact
     interface: EncryptOption

diff --git a/jwe/options_gen.go b/jwe/options_gen.go
diff --git a/jwe/options_gen_test.go b/jwe/options_gen_test.go
diff --git a/tools/cmd/genoptions/main.go b/tools/cmd/genoptions/main.go
@@ -119,7 +119,7 @@ func _main() error {
 	})
 
 	if err := genOptions(&objects); err != nil {
-		return fmt.Errorf(`failed to generate %q`, objects.Output)
+		return fmt.Errorf(`failed to generate %q: %w`, objects.Output, err)
 	}
 
 	if err := genOptionTests(&objects); err != nil {