From 5fd1fa53b33a0f6274d76e7c66e36884a83ad09f Mon Sep 17 00:00:00 2001
From: lestrrat <49281+lestrrat@users.noreply.github.com>
Date: Tue, 18 Jun 2024 07:31:04 +0900
Subject: [PATCH] Fix conversion in JWE when WithUseNumber global settings is
 used (#1140) (#1141)

* Fix conversion in JWE when WithUseNumber global settings is used (#1140)

* Add Changes
---
 Changes                 |  5 +++++
 internal/json/goccy.go  |  4 +---
 internal/json/json.go   | 17 +++++++++++------
 internal/json/stdlib.go |  4 +---
 jwe/jwe.go              | 24 +++++++++++++++++++++---
 jwx_test.go             | 21 +++++++++++++++++++++
 6 files changed, 60 insertions(+), 15 deletions(-)

diff --git a/Changes b/Changes
index 35d09f910..c7233a6a2 100644
--- a/Changes
+++ b/Changes
@@ -78,6 +78,11 @@ v2.0.22 UNRELEASED
 
     5. We have no plans to include more curves this way. One is already one too many.
 
+  * [jwe] Fixed a bug when using encryption algorithms involving PBES2 along with the
+    jwx.WithUseNumber() global option. Enabling this option would turn all values
+    stored in the JSON content to be of type `json.Number`, but we did not account for
+    it when checking for the value of `p2c` header, resulting in a conversion error.
+
 v2.0.21 07 Mar 2024
 [Security]
   * [jwe] Added `jwe.Settings(jwe.WithMaxDecompressBufferSize(int64))` to specify the
diff --git a/internal/json/goccy.go b/internal/json/goccy.go
index 59682104b..b63c015a2 100644
--- a/internal/json/goccy.go
+++ b/internal/json/goccy.go
@@ -26,11 +26,9 @@ func Engine() string {
 func NewDecoder(r io.Reader) *json.Decoder {
 	dec := json.NewDecoder(r)
 
-	muGlobalConfig.RLock()
-	if useNumber {
+	if UseNumber() {
 		dec.UseNumber()
 	}
-	muGlobalConfig.RUnlock()
 
 	return dec
 }
diff --git a/internal/json/json.go b/internal/json/json.go
index a4f1026a5..89209721b 100644
--- a/internal/json/json.go
+++ b/internal/json/json.go
@@ -4,19 +4,24 @@ import (
 	"bytes"
 	"fmt"
 	"os"
-	"sync"
+	"sync/atomic"
 
 	"github.com/lestrrat-go/jwx/v2/internal/base64"
 )
 
-var muGlobalConfig sync.RWMutex
-var useNumber bool
+var useNumber uint32 // TODO: at some point, change to atomic.Bool
+
+func UseNumber() bool {
+	return atomic.LoadUint32(&useNumber) == 1
+}
 
 // Sets the global configuration for json decoding
 func DecoderSettings(inUseNumber bool) {
-	muGlobalConfig.Lock()
-	useNumber = inUseNumber
-	muGlobalConfig.Unlock()
+	var val uint32
+	if inUseNumber {
+		val = 1
+	}
+	atomic.StoreUint32(&useNumber, val)
 }
 
 // Unmarshal respects the values specified in DecoderSettings,
diff --git a/internal/json/stdlib.go b/internal/json/stdlib.go
index 62b1a5ff5..759306260 100644
--- a/internal/json/stdlib.go
+++ b/internal/json/stdlib.go
@@ -25,11 +25,9 @@ func Engine() string {
 func NewDecoder(r io.Reader) *json.Decoder {
 	dec := json.NewDecoder(r)
 
-	muGlobalConfig.RLock()
-	if useNumber {
+	if UseNumber() {
 		dec.UseNumber()
 	}
-	muGlobalConfig.RUnlock()
 
 	return dec
 }
diff --git a/jwe/jwe.go b/jwe/jwe.go
index 44909a266..1e6ecda5a 100644
--- a/jwe/jwe.go
+++ b/jwe/jwe.go
@@ -745,10 +745,28 @@ func (dctx *decryptCtx) decryptContent(ctx context.Context, alg jwa.KeyEncryptio
 		if !ok {
 			return nil, fmt.Errorf(`failed to get 'p2c' field`)
 		}
-		countFlt, ok := count.(float64)
-		if !ok {
-			return nil, fmt.Errorf("unexpected type for 'p2c': %T", count)
+
+		// check if WithUseNumber is effective, because it will change the
+		// type of the underlying value (#1140)
+		var countFlt float64
+		if json.UseNumber() {
+			num, ok := count.(json.Number)
+			if !ok {
+				return nil, fmt.Errorf("unexpected type for 'p2c': %T", count)
+			}
+			v, err := num.Float64()
+			if err != nil {
+				return nil, fmt.Errorf("failed to convert 'p2c' to float64: %w", err)
+			}
+			countFlt = v
+		} else {
+			v, ok := count.(float64)
+			if !ok {
+				return nil, fmt.Errorf("unexpected type for 'p2c': %T", count)
+			}
+			countFlt = v
 		}
+
 		muSettings.RLock()
 		maxCount := maxPBES2Count
 		muSettings.RUnlock()
diff --git a/jwx_test.go b/jwx_test.go
index 40011ff7b..b21e15186 100644
--- a/jwx_test.go
+++ b/jwx_test.go
@@ -635,3 +635,24 @@ func TestGH996(t *testing.T) {
 		})
 	}
 }
+
+func TestGH1140(t *testing.T) {
+	// Using WithUseNumber changes the type of value obtained from the
+	// source JSON, which may cause issues
+	jwx.DecoderSettings(jwx.WithUseNumber(true))
+	t.Cleanup(func() {
+		jwx.DecoderSettings(jwx.WithUseNumber(false))
+	})
+	key, err := jwk.FromRaw([]byte("secure-key"))
+	require.NoError(t, err, `jwk.FromRaw should succeed`)
+
+	var encrypted []byte
+	encrypted, err = jwe.Encrypt(
+		[]byte("test-encryption-payload"),
+		jwe.WithKey(jwa.PBES2_HS256_A128KW, key),
+	)
+	require.NoError(t, err, `jwe.Encrypt should succeed`)
+
+	_, err = jwe.Decrypt(encrypted, jwe.WithKey(jwa.PBES2_HS256_A128KW, key))
+	require.NoError(t, err, `jwe.Decrypt should succeed`)
+}