[jwk] Add a Set() method to jwk.Set (#503)

* [jwk] Add a Set() method to jwk.Set

* Appease linter

Author: lestrrat

jwk/interface.go

@@ -57,10 +57,13 @@ const (
 type Set interface {
 	// Add adds the specified key. If the key already exists in the set, it is
 	// not added.
+	// This method will be renamed to `AddKey(Key)` in a future major release.
 	Add(Key) bool
 
 	// Clear resets the list of keys associated with this set, emptying the
 	// internal list of `jwk.Key`s
+	// This method will be changed in the future to clear all contents in the
+	// `jwk.Set` instead of just the keys.
 	Clear()
 
 	// Get returns the key at index `idx`. If the index is out of range,
@@ -69,11 +72,23 @@ type Set interface {
 	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.
+	// considered to be a private field. In other words, you cannot use this
+	// method to directly access the list of keys in the set
+	//
 	// This method will be renamed to `Get(string)` in a future major release.
 	Field(string) (interface{}, bool)
 
+	// Set sets the value of a single field.
+	//
+	// This method, which takes an `interface{}`, exists because
+	// these objects can contain extra _arbitrary_ fields that users can
+	// specify, and there is no way of knowing what type they could be.
+	Set(string, interface{}) error
+
+	// Remove removes the field associated with the specified key.
+	// There is no way to remove the `kty` (key type). You will ALWAYS be left with one field in a jwk.Key.
 	// Index returns the index where the given key exists, -1 otherwise
 	Index(Key) int
 

jwk/jwk_test.go

@@ -1769,4 +1769,37 @@ func TestSetWithPrivateParams(t *testing.T) {
 			check(t, js)
 		})
 	})
+	t.Run("Set private parameters", func(t *testing.T) {
+		set := jwk.NewSet()
+		if !assert.NoError(t, set.Set(`renewal_kid`, `foo`), `set.Set should succeed`) {
+			return
+		}
+
+		v, ok := set.Field(`renewal_kid`)
+		if !assert.True(t, ok, `set.Get("renewal_kid") should succeed`) {
+			return
+		}
+
+		if !assert.Equal(t, `foo`, v, `set.Get("renewal_kid") should return "foo"`) {
+			return
+		}
+
+		if !assert.Error(t, set.Set(`keys`, []string{"foo"}), `set.Set should fail`) {
+			return
+		}
+
+		k, err := jwk.New([]byte("foobar"))
+		if !assert.NoError(t, err, `jwk.New should succeed`) {
+			return
+		}
+		keys := []jwk.Key{k}
+
+		if !assert.NoError(t, set.Set(`keys`, keys), `set.Set should succeed`) {
+			return
+		}
+
+		if !assert.Equal(t, set.Len(), 1, `set should have 1 key`) {
+			return
+		}
+	})
 }

jwk/set.go

@@ -12,13 +12,32 @@ import (
 	"github.com/pkg/errors"
 )
 
+const keysKey = `keys` // appease linter
+
 // NewSet creates and empty `jwk.Set` object
 func NewSet() Set {
 	return &set{
 		privateParams: make(map[string]interface{}),
 	}
 }
 
+func (s *set) Set(n string, v interface{}) error {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	if n == keysKey {
+		vl, ok := v.([]Key)
+		if !ok {
+			return errors.Errorf(`value for field "keys" must be []jwk.Key`)
+		}
+		s.keys = vl
+		return nil
+	}
+
+	s.privateParams[n] = v
+	return nil
+}
+
 func (s *set) Field(n string) (interface{}, bool) {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
@@ -126,7 +145,7 @@ func (s *set) MarshalJSON() ([]byte, error) {
 	defer pool.ReleaseBytesBuffer(buf)
 	enc := json.NewEncoder(buf)
 
-	fields := []string{`keys`}
+	fields := []string{keysKey}
 	for k := range s.privateParams {
 		fields = append(fields, k)
 	}
@@ -138,7 +157,7 @@ func (s *set) MarshalJSON() ([]byte, error) {
 			buf.WriteByte(',')
 		}
 		fmt.Fprintf(buf, `%q:`, field)
-		if field != `keys` {
+		if field != keysKey {
 			if err := enc.Encode(s.privateParams[field]); err != nil {
 				return nil, errors.Wrapf(err, `failed to marshal field %q`, field)
 			}