Repository: lestrrat-go/jwx
Pull request #249 (2020-12-16 05:44:12)

DIRECT support
Not sure if I missed something, this was a lot easier than I anticipated.


Patch:
diff --git a/README.md b/README.md
index bc1979c8..4f0e70f7 100644
--- a/README.md
+++ b/README.md
@@ -20,10 +20,6 @@ PR/issues welcome.
 | [jws](https://github.com/lestrrat-go/jwx/tree/master/jws) | [RFC 7515](https://tools.ietf.org/html/rfc7515) |
 | [jwe](https://github.com/lestrrat-go/jwx/tree/master/jwe) | [RFC 7516](https://tools.ietf.org/html/rfc7516) |
 
-### In progress:
-
-* jwe - more algorithms
-
 ## Why?
 
 My goal was to write a server that heavily uses JWK and JWT. At first glance
@@ -38,8 +34,6 @@ Because I was writing the server side (and the client side for testing), I neede
 
 So here's go-jwx. This library is extensible, customizable, and hopefully well organized to the point that it is easy for you to slice and dice it.
 
-As of this writing (Nov 2015), it's still lacking a few of the algorithms for JWE that are described in JWA (which I believe to be less frequently used), but in general you should be able to do pretty much everything allowed in the specifications.
-
 ## Notes for users of pre-1.0.0 release
 
 The API has been reworked quite substantially between pre- and post 1.0.0 releases. Please check out the [Changes](./Changes) file (or the [diff](https://github.com/lestrrat-go/jwx/compare/v0.9.2...v1.0.0), if you are into that sort of thing)
@@ -373,7 +367,7 @@ Supported key encryption algorithm:
 | AES key wrap (128)                       | YES        | jwa.A128KW             |
 | AES key wrap (192)                       | YES        | jwa.A192KW             |
 | AES key wrap (256)                       | YES        | jwa.A256KW             |
-| Direct encryption                        | NO         | jwa.DIRECT             |
+| Direct encryption                        | YES (1)    | jwa.DIRECT             |
 | ECDH-ES                                  | YES (1)    | jwa.ECDH_ES            |
 | ECDH-ES + AES key wrap (128)             | YES        | jwa.ECDH_ES_A128KW     |
 | ECDH-ES + AES key wrap (192)             | YES        | jwa.ECDH_ES_A192KW     |
diff --git a/internal/jose/jose.go b/internal/jose/jose.go
index e037f464..7970b7e7 100644
--- a/internal/jose/jose.go
+++ b/internal/jose/jose.go
@@ -111,10 +111,15 @@ func GenerateJwk(ctx context.Context, t *testing.T, template string) (string, fu
 // a cleanup function.
 // The caller is responsible for calling the cleanup
 // function and make sure all resources are released
-func EncryptJwe(ctx context.Context, t *testing.T, payload []byte, keyfile string, enc string, compact bool) (string, func(), error) {
+func EncryptJwe(ctx context.Context, t *testing.T, payload []byte, alg string, keyfile string, enc string, compact bool) (string, func(), error) {
 	t.Helper()
 
-	arg := fmt.Sprintf(`{"protected":{"enc":"%s"}}`, enc)
+	var arg string
+	if alg == "dir" {
+		arg = fmt.Sprintf(`{"protected":{"alg":"dir","enc":"%s"}}`, enc)
+	} else {
+		arg = fmt.Sprintf(`{"protected":{"enc":"%s"}}`, enc)
+	}
 
 	cmdargs := []string{"jwe", "enc", "-k", keyfile, "-i", arg}
 	if compact {
diff --git a/jwe/encrypt.go b/jwe/encrypt.go
index f8c95a9e..69a4e736 100644
--- a/jwe/encrypt.go
+++ b/jwe/encrypt.go
@@ -82,7 +82,7 @@ func (e encryptCtx) Encrypt(plaintext []byte) (*Message, error) {
 			}
 			return nil, errors.Wrap(err, `failed to encrypt key`)
 		}
-		if enc.Algorithm() == jwa.ECDH_ES {
+		if enc.Algorithm() == jwa.ECDH_ES || enc.Algorithm() == jwa.DIRECT {
 			if len(e.keyEncrypters) > 1 {
 				return nil, errors.Errorf("unable to support multiple recipients for ECDH-ES")
 			}
diff --git a/jwe/internal/keyenc/interface.go b/jwe/internal/keyenc/interface.go
index 18da3984..5ec86025 100644
--- a/jwe/internal/keyenc/interface.go
+++ b/jwe/internal/keyenc/interface.go
@@ -25,6 +25,12 @@ type Decrypter interface {
 	Decrypt([]byte) ([]byte, error)
 }
 
+type Noop struct {
+	alg       jwa.KeyEncryptionAlgorithm
+	sharedkey []byte
+	keyID     string
+}
+
 // AES encrypts content encryption keys using AES key wrap.
 // Contrary to what the name implies, it also decrypt encrypted keys
 type AES struct {
diff --git a/jwe/internal/keyenc/keyenc.go b/jwe/internal/keyenc/keyenc.go
index 677e73cf..067ac5ef 100644
--- a/jwe/internal/keyenc/keyenc.go
+++ b/jwe/internal/keyenc/keyenc.go
@@ -27,6 +27,25 @@ import (
 	"github.com/pkg/errors"
 )
 
+func NewNoop(alg jwa.KeyEncryptionAlgorithm, sharedkey []byte) (*Noop, error) {
+	return &Noop{
+		alg:       alg,
+		sharedkey: sharedkey,
+	}, nil
+}
+
+func (kw *Noop) Algorithm() jwa.KeyEncryptionAlgorithm {
+	return kw.alg
+}
+
+func (kw *Noop) KeyID() string {
+	return kw.keyID
+}
+
+func (kw *Noop) Encrypt(cek []byte) (keygen.ByteSource, error) {
+	return keygen.ByteKey(kw.sharedkey), nil
+}
+
 // NewAES creates a key-wrap encrypter using AES.
 // Although the name suggests otherwise, this does the decryption as well.
 func NewAES(alg jwa.KeyEncryptionAlgorithm, sharedkey []byte) (*AES, error) {
diff --git a/jwe/jwe.go b/jwe/jwe.go
index 90c44236..380e1b1f 100644
--- a/jwe/jwe.go
+++ b/jwe/jwe.go
@@ -101,6 +101,12 @@ func Encrypt(payload []byte, keyalg jwa.KeyEncryptionAlgorithm, key interface{},
 		if err != nil {
 			return nil, errors.Wrap(err, "failed to create ECDHS key wrap encrypter")
 		}
+	case jwa.DIRECT:
+		sharedkey, ok := key.([]byte)
+		if !ok {
+			return nil, errors.New("invalid key: []byte required")
+		}
+		enc, _ = keyenc.NewNoop(keyalg, sharedkey)
 	default:
 		if pdebug.Enabled {
 			pdebug.Printf("Encrypt: unknown key encryption algorithm: %s", keyalg)


Test patch:
diff --git a/jwe/jwe_test.go b/jwe/jwe_test.go
index c53543c9..53447f34 100644
--- a/jwe/jwe_test.go
+++ b/jwe/jwe_test.go
@@ -483,43 +483,41 @@ func Test_GHIssue207(t *testing.T) {
 
 // tests direct key encryption by encrypting-decrypting a plaintext
 func TestEncode_Direct(t *testing.T) {
-	// TODO: uncomment once jwe.Encrypt() supports jwa.DIRECT
-	/*	var testcases = []struct {
-			Algorithm jwa.ContentEncryptionAlgorithm
-			KeySize   int // in bytes
-		}{
-			{jwa.A128CBC_HS256, 16},
-			{jwa.A128GCM, 16},
-			{jwa.A192CBC_HS384, 24},
-			{jwa.A192GCM, 24},
-			{jwa.A256CBC_HS512, 32},
-			{jwa.A256GCM, 32},
-		}
-		plaintext := []byte("Lorem ipsum")
-
-		for _, tc := range testcases {
-			tc := tc
-			t.Run(tc.Algorithm.String(), func(t *testing.T) {
-				key := make([]byte, tc.KeySize)
-				_, err := rand.Read(key)
-				if !assert.NoError(t, err, "Key generation succeeds") {
-					return
-				}
+	var testcases = []struct {
+		Algorithm jwa.ContentEncryptionAlgorithm
+		KeySize   int // in bytes
+	}{
+		{jwa.A128CBC_HS256, 32},
+		{jwa.A128GCM, 16},
+		{jwa.A192CBC_HS384, 48},
+		{jwa.A192GCM, 24},
+		{jwa.A256CBC_HS512, 64},
+		{jwa.A256GCM, 32},
+	}
+	plaintext := []byte("Lorem ipsum")
 
-				encrypted, err := jwe.Encrypt(plaintext, jwa.DIRECT, key, tc.Algorithm, jwa.NoCompress)
-				if !assert.NoError(t, err, "Encrypt succeeds") {
-					return
-				}
+	for _, tc := range testcases {
+		tc := tc
+		t.Run(tc.Algorithm.String(), func(t *testing.T) {
+			key := make([]byte, tc.KeySize)
+			_, err := rand.Read(key)
+			if !assert.NoError(t, err, "Key generation succeeds") {
+				return
+			}
 
-				decrypted, err := jwe.Decrypt(encrypted, jwa.DIRECT, key)
-				if !assert.NoError(t, err, `jwe.Decrypt should succeed`) {
-					return
-				}
+			encrypted, err := jwe.Encrypt(plaintext, jwa.DIRECT, key, tc.Algorithm, jwa.NoCompress)
+			if !assert.NoError(t, err, "Encrypt succeeds") {
+				return
+			}
 
-				assert.Equal(t, plaintext, decrypted, `jwe.Decrypt should match input plaintext`)
-			})
-		}
-	*/
+			decrypted, err := jwe.Decrypt(encrypted, jwa.DIRECT, key)
+			if !assert.NoError(t, err, `jwe.Decrypt should succeed`) {
+				return
+			}
+
+			assert.Equal(t, plaintext, decrypted, `jwe.Decrypt should match input plaintext`)
+		})
+	}
 }
 
 // Decrypts messages generated by `jose` tool. It helps check compatibility with other jwx implementations.
diff --git a/jwx_test.go b/jwx_test.go
index 2cd82b26..5144d1d2 100644
--- a/jwx_test.go
+++ b/jwx_test.go
@@ -189,6 +189,10 @@ func TestJoseCompatibility(t *testing.T) {
 			{jwa.PBES2_HS256_A128KW, jwa.A128CBC_HS256},
 			{jwa.PBES2_HS512_A256KW, jwa.A256GCM},
 			{jwa.PBES2_HS512_A256KW, jwa.A256CBC_HS512},
+			{jwa.DIRECT, jwa.A128GCM},
+			{jwa.DIRECT, jwa.A128CBC_HS256},
+			{jwa.DIRECT, jwa.A256GCM},
+			{jwa.DIRECT, jwa.A256CBC_HS512},
 		}
 		for _, test := range tests {
 			test := test
@@ -209,7 +213,11 @@ func joseInteropTest(ctx context.Context, spec interopTest, t *testing.T) {
 	expected := []byte("Lorem ipsum")
 
 	// let jose generate a key file
-	joseJwkFile, joseJwkCleanup, err := jose.GenerateJwk(ctx, t, fmt.Sprintf(`{"alg": "%s"}`, spec.alg))
+	alg := spec.alg.String()
+	if spec.alg == jwa.DIRECT {
+		alg = spec.enc.String()
+	}
+	joseJwkFile, joseJwkCleanup, err := jose.GenerateJwk(ctx, t, fmt.Sprintf(`{"alg": "%s"}`, alg))
 	if !assert.NoError(t, err, `jose.GenerateJwk should succeed`) {
 		return
 	}
@@ -242,7 +250,7 @@ func joseInteropTest(ctx context.Context, spec interopTest, t *testing.T) {
 	})
 	t.Run("Encrypt with jose, Decrypt with jwx", func(t *testing.T) {
 		// let jose encrypt payload using the key file
-		joseCryptFile, joseCryptCleanup, err := jose.EncryptJwe(ctx, t, expected, joseJwkFile, spec.enc.String(), true)
+		joseCryptFile, joseCryptCleanup, err := jose.EncryptJwe(ctx, t, expected, spec.alg.String(), joseJwkFile, spec.enc.String(), true)
 		if !assert.NoError(t, err, `jose.EncryptJwe should succeed`) {
 			return
 		}