Repository: lestrrat-go/jwx
Pull request #165 (2020-04-17 09:59:13)

Re-instate jwe.Recipient json Marshaling


Patch:
diff --git a/jwe/encrypt.go b/jwe/encrypt.go
index 7b988a68..bf998942 100644
--- a/jwe/encrypt.go
+++ b/jwe/encrypt.go
@@ -49,9 +49,9 @@ func (e encryptCtx) Encrypt(plaintext []byte) (*Message, error) {
 	recipients := make([]Recipient, len(e.keyEncrypters))
 	for i, enc := range e.keyEncrypters {
 		r := NewRecipient()
-		r.headers.Set("alg", enc.Algorithm())
+		r.Headers().Set(AlgorithmKey, enc.Algorithm())
 		if v := enc.KeyID(); v != "" {
-			r.headers.Set("kid", v)
+			r.Headers().Set(KeyIDKey, v)
 		}
 		enckey, err := enc.Encrypt(cek)
 		if err != nil {
@@ -60,20 +60,20 @@ func (e encryptCtx) Encrypt(plaintext []byte) (*Message, error) {
 			}
 			return nil, errors.Wrap(err, `failed to encrypt key`)
 		}
-		r.encryptedKey = enckey.Bytes()
+		r.SetEncryptedKey(enckey.Bytes())
 		if hp, ok := enckey.(populater); ok {
-			hp.Populate(r.headers)
+			hp.Populate(r.Headers())
 		}
 		if debug.Enabled {
 			debug.Printf("Encrypt: encrypted_key = %x (%d)", enckey.Bytes(), len(enckey.Bytes()))
 		}
-		recipients[i] = *r
+		recipients[i] = r
 	}
 
 	// If there's only one recipient, you want to include that in the
 	// protected header
 	if len(recipients) == 1 {
-		h, err := mergeHeaders(context.TODO(), protected, recipients[0].headers)
+		h, err := mergeHeaders(context.TODO(), protected, recipients[0].Headers())
 		if err != nil {
 			return nil, errors.Wrap(err, "failed to merge protected headers")
 		}
diff --git a/jwe/interface.go b/jwe/interface.go
index 90d5b1e0..679a3cb8 100644
--- a/jwe/interface.go
+++ b/jwe/interface.go
@@ -15,9 +15,16 @@ const (
 )
 
 // Recipient holds the encrypted key and hints to decrypt the key
-type Recipient struct {
-	headers      Headers
-	encryptedKey buffer.Buffer
+type Recipient interface {
+	Headers() Headers
+	EncryptedKey() buffer.Buffer
+	SetHeaders(Headers) error
+	SetEncryptedKey(interface{}) error
+}
+
+type stdRecipient struct {
+	headers      Headers       `json:"header"`
+	encryptedKey buffer.Buffer `json:"encrypted_key"`
 }
 
 // Message contains the entire encrypted JWE message
diff --git a/jwe/jwe.go b/jwe/jwe.go
index 6a1489c2..009a9b53 100644
--- a/jwe/jwe.go
+++ b/jwe/jwe.go
@@ -141,8 +141,10 @@ func ParseString(s string) (*Message, error) {
 func parseJSON(buf []byte) (*Message, error) {
 	m := struct {
 		*Message
-		*Recipient
-	}{}
+		Recipient
+	}{
+		Recipient: NewRecipient(),
+	}
 
 	if err := json.Unmarshal(buf, &m); err != nil {
 		return nil, errors.Wrap(err, "failed to parse JSON")
@@ -154,7 +156,7 @@ func parseJSON(buf []byte) (*Message, error) {
 			return nil, errors.New("invalid message: mixed flattened/full json serialization")
 		}
 
-		m.Message.recipients = []Recipient{*m.Recipient}
+		m.Message.recipients = []Recipient{m.Recipient}
 	}
 
 	return m.Message, nil
@@ -215,7 +217,7 @@ func parseCompact(buf []byte) (*Message, error) {
 	m.cipherText = ctbuf
 	m.initializationVector = ivbuf
 	m.recipients = []Recipient{
-		{
+		&stdRecipient{
 			headers:      hdr,
 			encryptedKey: enckeybuf,
 		},
diff --git a/jwe/message.go b/jwe/message.go
index a9db2f5f..be2fca62 100644
--- a/jwe/message.go
+++ b/jwe/message.go
@@ -15,20 +15,54 @@ import (
 )
 
 // NewRecipient creates a Recipient object
-func NewRecipient() *Recipient {
-	return &Recipient{
+func NewRecipient() Recipient {
+	return &stdRecipient{
 		headers: NewHeaders(),
 	}
 }
 
-func (r *Recipient) Headers() Headers {
+func (r *stdRecipient) SetHeaders(h Headers) error {
+	r.headers = h
+	return nil
+}
+
+func (r *stdRecipient) SetEncryptedKey(v interface{}) error {
+	return r.encryptedKey.Accept(v)
+}
+
+func (r *stdRecipient) Headers() Headers {
 	return r.headers
 }
 
-func (r *Recipient) EncryptedKey() buffer.Buffer {
+func (r *stdRecipient) EncryptedKey() buffer.Buffer {
 	return r.encryptedKey
 }
 
+type recipientMarshalProxy struct {
+	Headers      Headers       `json:"header"`
+	EncryptedKey buffer.Buffer `json:"encrypted_key"`
+}
+
+func (r *stdRecipient) UnmarshalJSON(buf []byte) error {
+	var proxy recipientMarshalProxy
+	proxy.Headers = NewHeaders()
+	if err := json.Unmarshal(buf, &proxy); err != nil {
+		return errors.Wrap(err, `failed to unmarshal json into recipient`)
+	}
+
+	r.headers = proxy.Headers
+	r.encryptedKey = proxy.EncryptedKey
+	return nil
+}
+
+func (r *stdRecipient) MarshalJSON() ([]byte, error) {
+	var proxy recipientMarshalProxy
+	proxy.Headers = r.headers
+	proxy.EncryptedKey = r.encryptedKey
+
+	return json.Marshal(proxy)
+}
+
 func mergeHeaders(ctx context.Context, h1, h2 Headers) (Headers, error) {
 	h3 := NewHeaders()
 
@@ -242,9 +276,9 @@ func (m *Message) Decrypt(alg jwa.KeyEncryptionAlgorithm, key interface{}) ([]by
 	var lastError error
 	for _, recipient := range m.recipients {
 		if debug.Enabled {
-			debug.Printf("Attempting to check if we can decode for recipient (alg = %s)", recipient.headers.Algorithm())
+			debug.Printf("Attempting to check if we can decode for recipient (alg = %s)", recipient.Headers().Algorithm())
 		}
-		if recipient.headers.Algorithm() != alg {
+		if recipient.Headers().Algorithm() != alg {
 			continue
 		}
 
@@ -257,7 +291,7 @@ func (m *Message) Decrypt(alg jwa.KeyEncryptionAlgorithm, key interface{}) ([]by
 			continue
 		}
 
-		h2, err = mergeHeaders(context.TODO(), h2, recipient.headers)
+		h2, err = mergeHeaders(context.TODO(), h2, recipient.Headers())
 		if err != nil {
 			if debug.Enabled {
 				debug.Printf("Failed to merge! %s", err)
@@ -275,7 +309,7 @@ func (m *Message) Decrypt(alg jwa.KeyEncryptionAlgorithm, key interface{}) ([]by
 			continue
 		}
 
-		cek, err := k.Decrypt(recipient.encryptedKey.Bytes())
+		cek, err := k.Decrypt(recipient.EncryptedKey().Bytes())
 		if err != nil {
 			if debug.Enabled {
 				debug.Printf("failed to decrypt key: %s", err)
diff --git a/jwe/serializer.go b/jwe/serializer.go
index 91e3107b..61bef119 100644
--- a/jwe/serializer.go
+++ b/jwe/serializer.go
@@ -50,7 +50,7 @@ func Compact(m *Message, _ ...Option) ([]byte, error) {
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to merge unprotected header")
 	}
-	hcopy, err = mergeHeaders(context.TODO(), hcopy, recipient.headers)
+	hcopy, err = mergeHeaders(context.TODO(), hcopy, recipient.Headers())
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to merge recipient header")
 	}
@@ -60,7 +60,7 @@ func Compact(m *Message, _ ...Option) ([]byte, error) {
 		return nil, errors.Wrap(err, "failed to encode header")
 	}
 
-	encryptedKey, err := recipient.encryptedKey.Base64Encode()
+	encryptedKey, err := recipient.EncryptedKey().Base64Encode()
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to encode encryption key")
 	}


Test patch:
diff --git a/jwe/lowlevel_test.go b/jwe/lowlevel_test.go
index 1a3e2128..a47e1036 100644
--- a/jwe/lowlevel_test.go
+++ b/jwe/lowlevel_test.go
@@ -96,8 +96,8 @@ func TestLowLevelParts_A128KW_A128CBCHS256(t *testing.T) {
 	}
 
 	r := NewRecipient()
-	r.headers.Set(AlgorithmKey, jwa.A128KW)
-	r.encryptedKey = enckey.Bytes()
+	r.Headers().Set(AlgorithmKey, jwa.A128KW)
+	r.SetEncryptedKey(enckey.Bytes())
 
 	protected := NewHeaders()
 	protected.Set(ContentEncryptionKey, jwa.A128CBC_HS256)
@@ -108,7 +108,7 @@ func TestLowLevelParts_A128KW_A128CBCHS256(t *testing.T) {
 	msg.cipherText = ciphertext
 	msg.initializationVector = iv
 	msg.tag = tag
-	msg.recipients = []Recipient{*r}
+	msg.recipients = []Recipient{r}
 
 	serialized, err := Compact(msg)
 	if !assert.NoError(t, err, "compact serialization is successful") {
diff --git a/jwe/message_test.go b/jwe/message_test.go
new file mode 100644
index 00000000..af3e64f8
--- /dev/null
+++ b/jwe/message_test.go
@@ -0,0 +1,29 @@
+package jwe_test
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/lestrrat-go/jwx/jwe"
+)
+
+func TestRecipient(t *testing.T) {
+	t.Run("JSON Marshaling", func(t *testing.T) {
+		const src = `{"header":{"foo":"bar"},"encrypted_key":"Zm9vYmFyYmF6"}`
+		r1 := jwe.NewRecipient()
+
+		if !assert.NoError(t, json.Unmarshal([]byte(src), r1), `json.Unmarshal should succeed`) {
+			return
+		}
+		
+		buf, err := json.Marshal(r1)
+		if !assert.NoError(t, err, `json.Marshal should succeed`) {
+			return
+		}
+
+		if !assert.Equal(t, []byte(src), buf, `json representation should match`) {
+			return
+		}
+	})
+}