Implement WithMessage option for jws (lestrrat-go#408)

Changes

@@ -26,6 +26,11 @@ v1.2.2
     process. The hook is called right after the JWE message has been parsed,
     but before the actual decryption has taken place.
 
+  JWS
+  * Option `jwe.WithMessage()` has been added. This allows the user to
+    obtain both the verified payload _and_ the raw `*jws.Message` in one
+    go when `jws.Verify()` is called
+
 v1.2.1 02 Jun 2021
 [New features]
   * Option `jwt.WithTypedClaim()` and `jwk.WithTypedField()` have been added.

jws/jws.go

@@ -161,16 +161,25 @@ func SignMulti(payload []byte, options ...Option) ([]byte, error) {
 // `Verifier` in `verify` subpackage, and call `Verify` method on it.
 // If you need to access signatures and JOSE headers in a JWS message,
 // use `Parse` function to get `Message` object.
-func Verify(buf []byte, alg jwa.SignatureAlgorithm, key interface{}) ([]byte, error) {
+func Verify(buf []byte, alg jwa.SignatureAlgorithm, key interface{}, options ...VerifyOption) ([]byte, error) {
+	var dst *Message
+	//nolint:forcetypeassert
+	for _, option := range options {
+		switch option.Ident() {
+		case identMessage{}:
+			dst = option.Value().(*Message)
+		}
+	}
+
 	buf = bytes.TrimSpace(buf)
 	if len(buf) == 0 {
 		return nil, errors.New(`attempt to verify empty buffer`)
 	}
 
 	if buf[0] == '{' {
-		return verifyJSON(buf, alg, key)
+		return verifyJSON(buf, alg, key, dst)
 	}
-	return verifyCompact(buf, alg, key)
+	return verifyCompact(buf, alg, key, dst)
 }
 
 // VerifySet uses keys store in a jwk.Set to verify the payload in `buf`.
@@ -208,7 +217,7 @@ func VerifySet(buf []byte, set jwk.Set) ([]byte, error) {
 	return nil, errors.New(`failed to verify message with any of the keys in the jwk.Set object`)
 }
 
-func verifyJSON(signed []byte, alg jwa.SignatureAlgorithm, key interface{}) ([]byte, error) {
+func verifyJSON(signed []byte, alg jwa.SignatureAlgorithm, key interface{}, dst *Message) ([]byte, error) {
 	verifier, err := NewVerifier(alg)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create verifier")
@@ -245,13 +254,16 @@ func verifyJSON(signed []byte, alg jwa.SignatureAlgorithm, key interface{}) ([]b
 		buf.WriteString(payload)
 
 		if err := verifier.Verify(buf.Bytes(), sig.signature, key); err == nil {
+			if dst != nil {
+				*dst = m
+			}
 			return m.payload, nil
 		}
 	}
 	return nil, errors.New(`could not verify with any of the signatures`)
 }
 
-func verifyCompact(signed []byte, alg jwa.SignatureAlgorithm, key interface{}) ([]byte, error) {
+func verifyCompact(signed []byte, alg jwa.SignatureAlgorithm, key interface{}, dst *Message) ([]byte, error) {
 	protected, payload, signature, err := SplitCompact(signed)
 	if err != nil {
 		return nil, errors.Wrap(err, `failed extract from compact serialization format`)
@@ -299,6 +311,18 @@ func verifyCompact(signed []byte, alg jwa.SignatureAlgorithm, key interface{}) (
 	if err != nil {
 		return nil, errors.Wrap(err, `message verified, failed to decode payload`)
 	}
+
+	if dst != nil {
+		// Construct a new Message object
+		m := NewMessage()
+		m.SetPayload(decodedPayload)
+		sig := NewSignature()
+		sig.SetProtectedHeaders(hdr)
+		sig.SetSignature(decodedSignature)
+		m.AppendSignature(sig)
+
+		*dst = *m
+	}
 	return decodedPayload, nil
 }
 

jws/jws_test.go

@@ -1102,3 +1102,38 @@ func TestCustomField(t *testing.T) {
 		}
 	})
 }
+
+func TestWithMessage(t *testing.T) {
+	key, err := jwxtest.GenerateRsaKey()
+	if !assert.NoError(t, err, "jwxtest.Generate should succeed") {
+		return
+	}
+
+	const text = "hello, world"
+	signed, err := jws.Sign([]byte(text), jwa.RS256, key)
+	if !assert.NoError(t, err, `jws.Sign should succeed`) {
+		return
+	}
+
+	m := jws.NewMessage()
+	payload, err := jws.Verify(signed, jwa.RS256, key.PublicKey, jws.WithMessage(m))
+	if !assert.NoError(t, err, `jws.Verify should succeed`) {
+		return
+	}
+	if !assert.Equal(t, payload, []byte(text), `jws.Verify should produce the correct payload`) {
+		return
+	}
+
+	parsed, err := jws.Parse(signed)
+	if !assert.NoError(t, err, `jws.Parse should succeed`) {
+		return
+	}
+
+	// The result of using jws.WithMessage should match the result of jws.Parse
+	buf1, _ := json.Marshal(m)
+	buf2, _ := json.Marshal(parsed)
+
+	if !assert.Equal(t, buf1, buf2, `result of jws.PArse and jws.Verify(..., jws.WithMessage()) should match`) {
+		return
+	}
+}

jws/option.go

@@ -8,6 +8,7 @@ type Option = option.Interface
 
 type identPayloadSigner struct{}
 type identHeaders struct{}
+type identMessage struct{}
 
 func WithSigner(signer Signer, key interface{}, public, protected Headers) Option {
 	return option.New(identPayloadSigner{}, &payloadSigner{
@@ -18,6 +19,26 @@ func WithSigner(signer Signer, key interface{}, public, protected Headers) Optio
 	})
 }
 
+// WithHeaders allows you to specify extra header values to include in the
+// final JWS message
 func WithHeaders(h Headers) Option {
 	return option.New(identHeaders{}, h)
 }
+
+// VerifyOption describes an option that can be passed to the jws.Verify function
+type VerifyOption interface {
+	Option
+	verifyOption()
+}
+
+type verifyOption struct {
+	Option
+}
+
+func (*verifyOption) verifyOption() {}
+
+// WithMessage can be passed to Verify() to obtain the jws.Message upon
+// a successful verification.
+func WithMessage(m *Message) VerifyOption {
+	return &verifyOption{option.New(identMessage{}, m)}
+}