From 4b2101afac6fdbfef94eaf7a352cc9b07f70f363 Mon Sep 17 00:00:00 2001
From: Bojidar Marinov <bojidar.marinov.bg@gmail.com>
Date: Tue, 12 Dec 2023 18:48:17 +0200
Subject: [PATCH] jwe: Support overriding the algorithm when supplying a JWK

Now, passing a JWK (via EncryptWithJwe / JSONWebKey.MarshalJSON) will
allow for ECDSA keys and for customizing the algorithm used with a
particular key.

Previously, the code made it impossible to supply a JWK-encoded ECDSA
public key in the encryption config, as all keys passed as JSONWebKey-s
were treated as RSA_OAEP keys, since utils.ParsePublicKey delegates to
parseJWKPublicKey which returns the JWK itself; and hence the switch in
the JWE keywrap failed to detect those as an ecdsa public key.
A simpler patch here would have been to change parseJWKPublicKey to return
the key contained inside the JWK directly, however, as pointed out by
stefanberger, this would have broken backwards compatibility of the public
API. Plus, using the algorithm encoded in the JWK allows us to more easily
extend the JWE encoder to new algorithms.

Risks: JWK-s containing RSA keys but with .Algorithm not set to "" (the
default value) or string(jose.RSA_OLAP) will end up erroring or producing
different encryptions than before. However, such keys would have failed to
decrypt the contents regardless, so it should be fine to consider this a
correction rather than breakage of old behavior. (Hyrum's law
notwithstanding)

Signed-off-by: Bojidar Marinov <bojidar.marinov.bg@gmail.com>
---
 keywrap/jwe/keywrapper_jwe.go      | 17 +++++++++++++-
 keywrap/jwe/keywrapper_jwe_test.go | 36 ++++++++++++++++++++++++++++++
 utils/testing.go                   | 13 +++++++++--
 3 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/keywrap/jwe/keywrapper_jwe.go b/keywrap/jwe/keywrapper_jwe.go
index cd2241c..24e1d61 100644
--- a/keywrap/jwe/keywrapper_jwe.go
+++ b/keywrap/jwe/keywrapper_jwe.go
@@ -123,9 +123,24 @@ func addPubKeys(joseRecipients *[]jose.Recipient, pubKeys [][]byte) error {
 		}
 
 		alg := jose.RSA_OAEP
-		switch key.(type) {
+		switch key := key.(type) {
 		case *ecdsa.PublicKey:
 			alg = jose.ECDH_ES_A256KW
+		case *jose.JSONWebKey:
+			if key.Algorithm != "" {
+				alg = jose.KeyAlgorithm(key.Algorithm)
+				switch alg {
+				/* accepted algorithms */
+				case jose.RSA_OAEP:
+				case jose.RSA_OAEP_256:
+				case jose.ECDH_ES_A128KW:
+				case jose.ECDH_ES_A192KW:
+				case jose.ECDH_ES_A256KW:
+				/* all others are rejected */
+				default:
+					return fmt.Errorf("%s is an unsupported JWE key algorithm", alg)
+				}
+			}
 		}
 
 		*joseRecipients = append(*joseRecipients, jose.Recipient{
diff --git a/keywrap/jwe/keywrapper_jwe_test.go b/keywrap/jwe/keywrapper_jwe_test.go
index 3beea39..3d5e053 100644
--- a/keywrap/jwe/keywrapper_jwe_test.go
+++ b/keywrap/jwe/keywrapper_jwe_test.go
@@ -70,6 +70,21 @@ func createValidJweCcs() ([]*config.CryptoConfig, error) {
 		return nil, err
 	}
 
+	ecKey, err := utils.CreateECDSAKey(elliptic.P521())
+	if err != nil {
+		return nil, err
+	}
+
+	jweEcPrivKeyJwk, err := jose.JSONWebKey{Key: ecKey, Algorithm: string(jose.ECDH_ES_A256KW)}.MarshalJSON()
+	if err != nil {
+		return nil, err
+	}
+
+	jweEcPubKeyJwk, err := jose.JSONWebKey{Key: &ecKey.PublicKey, Algorithm: string(jose.ECDH_ES_A256KW)}.MarshalJSON()
+	if err != nil {
+		return nil, err
+	}
+
 	validJweCcs := []*config.CryptoConfig{
 		// Key 1
 		{
@@ -226,6 +241,27 @@ func createValidJweCcs() ([]*config.CryptoConfig, error) {
 				},
 			},
 		},
+		// EC Key (JWK format)
+		{
+			EncryptConfig: &config.EncryptConfig{
+				Parameters: map[string][][]byte{
+					"pubkeys": {jweEcPubKeyJwk},
+				},
+				DecryptConfig: config.DecryptConfig{
+					Parameters: map[string][][]byte{
+						"privkeys":           {jweEcPrivKeyJwk},
+						"privkeys-passwords": {oneEmpty},
+					},
+				},
+			},
+
+			DecryptConfig: &config.DecryptConfig{
+				Parameters: map[string][][]byte{
+					"privkeys":           {jweEcPrivKeyJwk},
+					"privkeys-passwords": {oneEmpty},
+				},
+			},
+		},
 	}
 	return validJweCcs, nil
 }
diff --git a/utils/testing.go b/utils/testing.go
index 69bb9d1..050aa88 100644
--- a/utils/testing.go
+++ b/utils/testing.go
@@ -38,6 +38,15 @@ func CreateRSAKey(bits int) (*rsa.PrivateKey, error) {
 	return key, nil
 }
 
+// CreateECDSAKey creates an elliptic curve key for the given curve
+func CreateECDSAKey(curve elliptic.Curve) (*ecdsa.PrivateKey, error) {
+	key, err := ecdsa.GenerateKey(curve, rand.Reader)
+	if err != nil {
+		return nil, fmt.Errorf("ecdsa.GenerateKey failed: %w", err)
+	}
+	return key, nil
+}
+
 // CreateRSATestKey creates an RSA key of the given size and returns
 // the public and private key in PEM or DER format
 func CreateRSATestKey(bits int, password []byte, pemencode bool) ([]byte, []byte, error) {
@@ -85,9 +94,9 @@ func CreateRSATestKey(bits int, password []byte, pemencode bool) ([]byte, []byte
 // CreateECDSATestKey creates and elliptic curve key for the given curve and returns
 // the public and private key in DER format
 func CreateECDSATestKey(curve elliptic.Curve) ([]byte, []byte, error) {
-	key, err := ecdsa.GenerateKey(curve, rand.Reader)
+	key, err := CreateECDSAKey(curve)
 	if err != nil {
-		return nil, nil, fmt.Errorf("ecdsa.GenerateKey failed: %w", err)
+		return nil, nil, err
 	}
 
 	pubData, err := x509.MarshalPKIXPublicKey(&key.PublicKey)