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

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

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
 }

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,

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
 }

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()

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`)
+}