[jwk] Add the ability to work with private fields in jwk.Set (#502)

* [jwk] Add the ability to work with private fields in jwk.Set

* appease linter

* check case when JWKS == JWK

Author: lestrrat

Changes

@@ -1,6 +1,12 @@
 Changes
 =======
 
+v1.2.12 UNRELEASED
+[New Features]
+  * `jwk.Set` can now parse private parameters. For example, after parsing
+    `{"foo": "bar", "keys": [...]}`, users can get to the value of `"foo"`
+     by calling `set.Field("foo")`
+
 v1.2.11 14 Nov 2021
 [Security Fix]
   * It was reported that since v1.2.6, it was possible to craft

jwk/interface.go

@@ -48,6 +48,12 @@ const (
 // `"encoding/json".Marshal` and `"encoding/json".Unmarshal`. However,
 // if you do not know if the payload contains a single JWK or a JWK set,
 // consider using `jwk.Parse()` to always get a `jwk.Set` out of it.
+//
+// Since v1.2.12, JWK sets with private parameters can be parsed as well.
+// Such private parameters can be accessed via the `Field()` method.
+// If a resource contains a single JWK instead of a JWK set, private parameters
+// are stored in _both_ the resulting `jwk.Set` object and the `jwk.Key` object .
+//
 type Set interface {
 	// Add adds the specified key. If the key already exists in the set, it is
 	// not added.
@@ -59,8 +65,15 @@ type Set interface {
 
 	// Get returns the key at index `idx`. If the index is out of range,
 	// then the second return value is false.
+	// This method will be renamed to `Key(int)` in a future major release.
 	Get(int) (Key, bool)
 
+	// Field returns the value of a private field in the key set.
+	// For the purposes of a key set, any field other than the "keys" field is
+	// considered to be a private field.
+	// This method will be renamed to `Get(string)` in a future major release.
+	Field(string) (interface{}, bool)
+
 	// Index returns the index where the given key exists, -1 otherwise
 	Index(Key) int
 
@@ -84,9 +97,10 @@ type Set interface {
 }
 
 type set struct {
-	keys []Key
-	mu   sync.RWMutex
-	dc   DecodeCtx
+	keys          []Key
+	mu            sync.RWMutex
+	dc            DecodeCtx
+	privateParams map[string]interface{}
 }
 
 type HeaderVisitor = iter.MapVisitor

jwk/jwk.go

@@ -466,7 +466,7 @@ func ParseKey(data []byte, options ...ParseOption) (Key, error) {
 // call `json.Unmarshal` against an empty set created by `jwk.NewSet()`
 // to parse a JSON buffer into a `jwk.Set`.
 //
-// This method exists because many times the user does not know before hand
+// This function exists because many times the user does not know before hand
 // if a JWK(s) resource at a remote location contains a single JWK key or
 // a JWK set, and `jwk.Parse()` can handle either case, returning a JWK Set
 // even if the data only contains a single JWK key

jwk/jwk_test.go

@@ -1,6 +1,7 @@
 package jwk_test
 
 import (
+	"bytes"
 	"context"
 	"crypto"
 	"crypto/ecdsa"
@@ -1653,3 +1654,119 @@ func TestGH491(t *testing.T) {
 		return
 	}
 }
+
+func TestSetWithPrivateParams(t *testing.T) {
+	k1, err := jwxtest.GenerateRsaJwk()
+	if !assert.NoError(t, err, `jwxtest.GenerateRsaJwk should succeed`) {
+		return
+	}
+	k2, err := jwxtest.GenerateEcdsaJwk()
+	if !assert.NoError(t, err, `jwxtest.GenerateEcdsaJwk should succeed`) {
+		return
+	}
+	k3, err := jwxtest.GenerateSymmetricJwk()
+	if !assert.NoError(t, err, `jwxtest.GenerateSymmetricJwk should succeed`) {
+		return
+	}
+
+	t.Run("JWK instead of JWKS", func(t *testing.T) {
+		var buf bytes.Buffer
+		_ = k1.Set(`renewal_kid`, "foo")
+		_ = json.NewEncoder(&buf).Encode(k1)
+
+		var check = func(t *testing.T, buf []byte) {
+			set, err := jwk.Parse(buf)
+			if !assert.NoError(t, err, `jwk.Parse should succeed`) {
+				return
+			}
+
+			if !assert.Equal(t, 1, set.Len(), `set.Len() should be 1`) {
+				return
+			}
+
+			v, ok := set.Field(`renewal_kid`)
+			if !assert.True(t, ok, `set.Field("renewal_kid") should return ok = true`) {
+				return
+			}
+
+			if !assert.Equal(t, `foo`, v, `set.Field("renewal_kid") should return "foo"`) {
+				return
+			}
+
+			key, ok := set.Get(0)
+			if !assert.True(t, ok, `set.Get(0) should return ok = true`) {
+				return
+			}
+
+			v, ok = key.Get(`renewal_kid`)
+			if !assert.True(t, ok, `key.Get("renewal_kid") should return ok = true`) {
+				return
+			}
+
+			if !assert.Equal(t, `foo`, v, `key.Get("renewal_kid") should return "foo"`) {
+				return
+			}
+		}
+
+		t.Run("Check original buffer", func(t *testing.T) {
+			check(t, buf.Bytes())
+		})
+		t.Run("Check serialized", func(t *testing.T) {
+			set, err := jwk.Parse(buf.Bytes())
+			if !assert.NoError(t, err, `jwk.Parse should succeed`) {
+				return
+			}
+			js, err := json.MarshalIndent(set, "", "  ")
+			if !assert.NoError(t, err, `json.MarshalIndent should succeed`) {
+				return
+			}
+			check(t, js)
+		})
+	})
+	t.Run("JWKS with multiple keys", func(t *testing.T) {
+		var buf bytes.Buffer
+		buf.WriteString(`{"renewal_kid":"foo","keys":[`)
+		enc := json.NewEncoder(&buf)
+		_ = enc.Encode(k1)
+		buf.WriteByte(',')
+		_ = enc.Encode(k2)
+		buf.WriteByte(',')
+		_ = enc.Encode(k3)
+		buf.WriteString(`]}`)
+
+		var check = func(t *testing.T, buf []byte) {
+			set, err := jwk.Parse(buf)
+			if !assert.NoError(t, err, `jwk.Parse should succeed`) {
+				return
+			}
+
+			if !assert.Equal(t, 3, set.Len(), `set.Len() should be 3`) {
+				return
+			}
+
+			v, ok := set.Field(`renewal_kid`)
+			if !assert.True(t, ok, `set.Field("renewal_kid") should return ok = true`) {
+				return
+			}
+
+			if !assert.Equal(t, `foo`, v, `set.Field("renewal_kid") should return "foo"`) {
+				return
+			}
+		}
+
+		t.Run("Check original buffer", func(t *testing.T) {
+			check(t, buf.Bytes())
+		})
+		t.Run("Check serialized", func(t *testing.T) {
+			set, err := jwk.Parse(buf.Bytes())
+			if !assert.NoError(t, err, `jwk.Parse should succeed`) {
+				return
+			}
+			js, err := json.MarshalIndent(set, "", "  ")
+			if !assert.NoError(t, err, `json.MarshalIndent should succeed`) {
+				return
+			}
+			check(t, js)
+		})
+	})
+}

jwk/set.go

@@ -1,7 +1,10 @@
 package jwk
 
 import (
+	"bytes"
 	"context"
+	"fmt"
+	"sort"
 
 	"github.com/lestrrat-go/iter/arrayiter"
 	"github.com/lestrrat-go/jwx/internal/json"
@@ -11,7 +14,17 @@ import (
 
 // NewSet creates and empty `jwk.Set` object
 func NewSet() Set {
-	return &set{}
+	return &set{
+		privateParams: make(map[string]interface{}),
+	}
+}
+
+func (s *set) Field(n string) (interface{}, bool) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	v, ok := s.privateParams[n]
+	return v, ok
 }
 
 func (s *set) Get(idx int) (Key, bool) {
@@ -105,10 +118,6 @@ func iterate(ctx context.Context, keys []Key, ch chan *KeyPair) {
 	}
 }
 
-type keySetMarshalProxy struct {
-	Keys []json.RawMessage `json:"keys"`
-}
-
 func (s *set) MarshalJSON() ([]byte, error) {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
@@ -117,16 +126,36 @@ func (s *set) MarshalJSON() ([]byte, error) {
 	defer pool.ReleaseBytesBuffer(buf)
 	enc := json.NewEncoder(buf)
 
-	buf.WriteString(`{"keys":[`)
-	for i, k := range s.keys {
+	fields := []string{`keys`}
+	for k := range s.privateParams {
+		fields = append(fields, k)
+	}
+	sort.Strings(fields)
+
+	buf.WriteByte('{')
+	for i, field := range fields {
 		if i > 0 {
 			buf.WriteByte(',')
 		}
-		if err := enc.Encode(k); err != nil {
-			return nil, errors.Wrapf(err, `failed to marshal key #%d`, i)
+		fmt.Fprintf(buf, `%q:`, field)
+		if field != `keys` {
+			if err := enc.Encode(s.privateParams[field]); err != nil {
+				return nil, errors.Wrapf(err, `failed to marshal field %q`, field)
+			}
+		} else {
+			buf.WriteByte('[')
+			for j, k := range s.keys {
+				if j > 0 {
+					buf.WriteByte(',')
+				}
+				if err := enc.Encode(k); err != nil {
+					return nil, errors.Wrapf(err, `failed to marshal key #%d`, i)
+				}
+			}
+			buf.WriteByte(']')
 		}
 	}
-	buf.WriteString("]}")
+	buf.WriteByte('}')
 
 	ret := make([]byte, buf.Len())
 	copy(ret, buf.Bytes())
@@ -137,10 +166,8 @@ func (s *set) UnmarshalJSON(data []byte) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	var proxy keySetMarshalProxy
-	if err := json.Unmarshal(data, &proxy); err != nil {
-		return errors.Wrap(err, `failed to unmarshal into Key (proxy)`)
-	}
+	s.privateParams = make(map[string]interface{})
+	s.keys = nil
 
 	var options []ParseOption
 	if dc := s.dc; dc != nil {
@@ -149,20 +176,61 @@ func (s *set) UnmarshalJSON(data []byte) error {
 		}
 	}
 
-	if len(proxy.Keys) == 0 {
-		k, err := ParseKey(data, options...)
+	var sawKeysField bool
+	dec := json.NewDecoder(bytes.NewReader(data))
+LOOP:
+	for {
+		tok, err := dec.Token()
 		if err != nil {
-			return errors.Wrap(err, `failed to unmarshal key from JSON headers`)
+			return errors.Wrap(err, `error reading token`)
 		}
-		s.keys = append(s.keys, k)
-	} else {
-		for i, buf := range proxy.Keys {
-			k, err := ParseKey([]byte(buf), options...)
-			if err != nil {
-				return errors.Wrapf(err, `failed to unmarshal key #%d (total %d) from multi-key JWK set`, i+1, len(proxy.Keys))
+
+		switch tok := tok.(type) {
+		case json.Delim:
+			// Assuming we're doing everything correctly, we should ONLY
+			// get either '{' or '}' here.
+			if tok == '}' { // End of object
+				break LOOP
+			} else if tok != '{' {
+				return errors.Errorf(`expected '{', but got '%c'`, tok)
 			}
-			s.keys = append(s.keys, k)
+		case string:
+			switch tok {
+			case "keys":
+				sawKeysField = true
+				var list []json.RawMessage
+				if err := dec.Decode(&list); err != nil {
+					return errors.Wrap(err, `failed to decode "keys"`)
+				}
+
+				for i, keysrc := range list {
+					key, err := ParseKey(keysrc, options...)
+					if err != nil {
+						return errors.Wrapf(err, `failed to code key #%d in "keys"`, i)
+					}
+					s.keys = append(s.keys, key)
+				}
+			default:
+				var v interface{}
+				if err := dec.Decode(&v); err != nil {
+					return errors.Wrapf(err, `failed to decode value for key %q`, tok)
+				}
+				s.privateParams[tok] = v
+			}
+		}
+	}
+
+	// This is really silly, but we can only detect the
+	// lack of the "keys" field after going through the
+	// entire object once
+	// Not checking for len(s.keys) == 0, because it could be
+	// an empty key set
+	if !sawKeysField {
+		key, err := ParseKey(data, options...)
+		if err != nil {
+			return errors.Wrapf(err, `failed to parse sole key in key set`)
 		}
+		s.keys = append(s.keys, key)
 	}
 	return nil
 }