fix(hpke): require psk_id in protected header and update test vectors

Author: dajiaji

cwt/cose.py

@@ -699,8 +699,7 @@ def _validate_cose_message(
                 if not isinstance(v, (bytes, bytearray)):
                     raise ValueError("ek (-4) must be bstr.")
             if k == -5:  # psk_id
-                if not isinstance(v, (bytes, bytearray)):
-                    raise ValueError("psk_id (-5) must be bstr.")
+                raise ValueError("psk_id (-5) must be placed only in the protected header.")
             h[k] = v
         if len(h) != len(p) + len(u):
             raise ValueError("The same keys are both in protected and unprotected headers.")

cwt/recipient_algs/hpke.py

@@ -102,7 +102,8 @@ def encode(self, plaintext: bytes = b"", aad: bytes = b"") -> Tuple[List[Any], O
             raise ValueError("recipient_key should be set in advance.")
         self._kem_key = self._to_kem_key(self._recipient_key)
         try:
-            psk_id = self._unprotected.get(-5, None)
+            # psk_id MUST be in the protected header (draft-ietf-cose-hpke)
+            psk_id = self._protected.get(-5, None) if isinstance(self._protected, dict) else None
             if psk_id is not None and not isinstance(psk_id, (bytes, bytearray)):
                 raise EncodeError("psk_id (-5) must be bstr.")
             if self._psk is not None and psk_id is None:
@@ -147,7 +148,8 @@ def decode(
         if not isinstance(ek, (bytes, bytearray)):
             raise DecodeError("ek (-4) must be bstr.")
         try:
-            psk_id = self._unprotected.get(-5, None)
+            # psk_id MUST be in the protected header (draft-ietf-cose-hpke)
+            psk_id = self._protected.get(-5, None) if isinstance(self._protected, dict) else None
             if psk_id is not None and not isinstance(psk_id, (bytes, bytearray)):
                 raise DecodeError("psk_id (-5) must be bstr.")
             if self._psk is not None and psk_id is None:

tests/test_cose_hpke.py

@@ -186,8 +186,8 @@ def test_cose_hpke_encrypt0_psk_id_wrong_type_header_validation(self):
             sender.encode_and_encrypt(
                 b"This is the content.",
                 rpk,
-                protected={COSEHeaders.ALG: COSEAlgs.HPKE_0},
-                unprotected={COSEHeaders.KID: b"01", COSEHeaders.PSK_ID: 123},
+                protected={COSEHeaders.ALG: COSEAlgs.HPKE_0, COSEHeaders.PSK_ID: 123},
+                unprotected={COSEHeaders.KID: b"01"},
             )
         assert "psk_id (-5) must be bstr." in str(err.value)
 
@@ -205,8 +205,8 @@ def test_cose_hpke_encrypt0_with_psk_id_roundtrip(self):
         encoded = sender.encode_and_encrypt(
             b"This is the content.",
             rpk,
-            protected={COSEHeaders.ALG: COSEAlgs.HPKE_0},
-            unprotected={COSEHeaders.KID: b"01", COSEHeaders.PSK_ID: b"psk-01"},
+            protected={COSEHeaders.ALG: COSEAlgs.HPKE_0, COSEHeaders.PSK_ID: b"psk-01"},
+            unprotected={COSEHeaders.KID: b"01"},
             hpke_psk=b"secret-psk",
         )
 
@@ -279,8 +279,8 @@ def test_cose_hpke_encrypt0_psk_id_without_psk_should_error_on_encode(self):
             sender.encode_and_encrypt(
                 b"This is the content.",
                 rpk,
-                protected={COSEHeaders.ALG: COSEAlgs.HPKE_0},
-                unprotected={COSEHeaders.KID: b"01", COSEHeaders.PSK_ID: b"psk-01"},
+                protected={COSEHeaders.ALG: COSEAlgs.HPKE_0, COSEHeaders.PSK_ID: b"psk-01"},
+                unprotected={COSEHeaders.KID: b"01"},
             )
         assert "hpke_psk is required when psk_id (-5) is provided." in str(err.value)
 
@@ -295,7 +295,8 @@ def test_cose_hpke_encrypt0_psk_id_without_psk_should_error_on_decode(self):
             }
         )
         sender = COSE.new()
-        # First, produce a base-mode (no psk_id) and then inject psk_id to simulate peer mismatch
+        # First, produce a base-mode (no psk_id) and then inject psk_id into the
+        # protected header to simulate peer mismatch
         encoded = sender.encode_and_encrypt(
             b"This is the content.",
             rpk,
@@ -304,8 +305,10 @@ def test_cose_hpke_encrypt0_psk_id_without_psk_should_error_on_decode(self):
         )
         tag = cbor2.loads(encoded)
         p, u, c = tag.value
-        u[-5] = b"psk-01"
-        tampered = cbor2.dumps(cbor2.CBORTag(16, [p, u, c]))
+        p_map = cbor2.loads(p)
+        p_map[-5] = b"psk-01"
+        tampered_p = cbor2.dumps(p_map)
+        tampered = cbor2.dumps(cbor2.CBORTag(16, [tampered_p, u, c]))
 
         rsk = COSEKey.from_jwk(
             {
@@ -676,8 +679,8 @@ def test_cose_hpke_ke_with_psk_roundtrip(self):
         )
 
         r = Recipient.new(
-            protected={COSEHeaders.ALG: COSEAlgs.HPKE_0_KE},
-            unprotected={COSEHeaders.KID: b"01", COSEHeaders.PSK_ID: b"psk-01"},
+            protected={COSEHeaders.ALG: COSEAlgs.HPKE_0_KE, COSEHeaders.PSK_ID: b"psk-01"},
+            unprotected={COSEHeaders.KID: b"01"},
             recipient_key=rpk,
             hpke_psk=b"secret-psk",
         )