feat(jwe): allow verify only one recipient

#16

Author: lepture

src/joserfc/rfc7516/message.py

@@ -14,6 +14,7 @@
 )
 from .registry import JWERegistry
 from ..errors import (
+    JoseError,
     DecodeError,
     InvalidCEKLengthError,
     InvalidEncryptedKeyError,
@@ -99,8 +100,15 @@ def _perform_decrypt(obj: EncryptionData, registry: JWERegistry) -> None:
         # Step 6, Determine the Key Management Mode employed by the algorithm
         # specified by the "alg" (algorithm) Header Parameter.
         alg = registry.get_alg(headers["alg"])
-        cek = decrypt_recipient(alg, enc, recipient, tag)
-        cek_set.add(cek)
+        try:
+            cek = decrypt_recipient(alg, enc, recipient, tag)
+            cek_set.add(cek)
+        except (AssertionError, JoseError) as error:
+            if registry.verify_all_recipients:
+                raise error
+
+    if not cek_set:
+        raise DecodeError('Invalid recipients')
 
     if len(cek_set) > 1:  # pragma: no cover
         raise DecodeError('Multiple "cek" found')

src/joserfc/rfc7516/models.py

@@ -173,7 +173,7 @@ def generate_iv(self) -> bytes:
         return os.urandom(self.iv_size // 8)
 
     def check_iv(self, iv: bytes) -> bytes:
-        if len(iv) * 8 != self.iv_size:
+        if len(iv) * 8 != self.iv_size:  # pragma: no cover
             raise ValueError('Invalid "iv" size')
         return iv
 

src/joserfc/rfc7516/registry.py

@@ -38,12 +38,14 @@ def __init__(
             self,
             header_registry: t.Optional[HeaderRegistryDict] = None,
             algorithms: t.Optional[t.List[str]] = None,
+            verify_all_recipients: bool = True,
             strict_check_header: bool = True):
         self.header_registry: HeaderRegistryDict = {}
         self.header_registry.update(JWE_HEADER_REGISTRY)
         if header_registry is not None:
             self.header_registry.update(header_registry)
         self.allowed = algorithms
+        self.verify_all_recipients = verify_all_recipients
         self.strict_check_header = strict_check_header
 
     @classmethod

tests/jwe/test_json.py

@@ -2,7 +2,11 @@
 from joserfc import jwe
 from joserfc.jwe import GeneralJSONEncryption
 from joserfc.jwk import KeySet, RSAKey, ECKey, OctKey
-from joserfc.errors import ConflictAlgorithmError
+from joserfc.errors import (
+    DecodeError,
+    ConflictAlgorithmError,
+    InvalidKeyTypeError,
+)
 
 
 class TestJWEJSON(TestCase):
@@ -43,3 +47,27 @@ def test_with_aad(self):
         value = jwe.encrypt_json(obj, None)
         obj1 = jwe.decrypt_json(value, key1)
         self.assertEqual(obj1.aad, b"foo")
+
+    def test_decode_multiple_recipients(self):
+        key1 = RSAKey.generate_key()
+        key2 = ECKey.generate_key()
+        obj = GeneralJSONEncryption({"enc": "A128CBC-HS256"}, b"i")
+        obj.add_recipient({"alg": "RSA-OAEP"}, key1)
+        obj.add_recipient({"alg": "ECDH-ES+A128KW"}, key2)
+        value = jwe.encrypt_json(obj, None)
+        self.assertRaises(
+            InvalidKeyTypeError,
+            jwe.decrypt_json,
+            value, key1,
+        )
+        registry = jwe.JWERegistry(verify_all_recipients=False)
+        obj1 = jwe.decrypt_json(value, key1, registry=registry)
+        self.assertEqual(obj1.plaintext, b"i")
+
+        key3 = OctKey.generate_key()
+        self.assertRaises(
+            DecodeError,
+            jwe.decrypt_json,
+            value, key3,
+            registry=registry,
+        )