Add ChaCha20Poly1305 support to HPKE implementation (#14393)

Add CHACHA20_POLY1305 as an AEAD option for HPKE (RFC 9180 aead_id
0x0003), alongside the existing AES_128_GCM support.

Changes:
- Add CHACHA20_POLY1305 variant to the HPKE AEAD enum with parameters
  (Nk=32, Nn=12, Nt=16)
- Make Suite store the selected AEAD and dynamically dispatch to the
  correct cipher for encrypt/decrypt operations
- Make key_schedule return PyBytes directly to avoid unnecessary
  allocations when passing to AEAD constructors
- Make ChaCha20Poly1305 struct and its encrypt/decrypt methods
  pub(crate) so they can be used from the HPKE module
- Update tests to cover ChaCha20Poly1305 roundtrips and RFC 9180 test
  vector validation
- Update type stubs and documentation

https://claude.ai/code/session_019vzyAZNnzXmtyAnkc9BHZp

Co-authored-by: Claude <noreply@anthropic.com>

Author: alex

docs/hazmat/primitives/hpke.rst

@@ -10,7 +10,7 @@ Mechanism (KEM), a Key Derivation Function (KDF), and an Authenticated
 Encryption with Associated Data (AEAD) scheme. It is defined in :rfc:`9180`.
 
 This implementation supports Base mode with DHKEM(X25519, HKDF-SHA256),
-HKDF-SHA256, and AES-128-GCM.
+HKDF-SHA256, and either AES-128-GCM or ChaCha20Poly1305.
 
 HPKE provides authenticated encryption: the recipient can be certain that the
 message was encrypted by someone who knows the recipient's public key, but
@@ -99,3 +99,7 @@ specifying auxiliary authenticated information.
     .. attribute:: AES_128_GCM
 
         AES-128-GCM
+
+    .. attribute:: CHACHA20_POLY1305
+
+        ChaCha20Poly1305

src/cryptography/hazmat/bindings/_rust/openssl/hpke.pyi

@@ -13,6 +13,7 @@ class KDF:
 
 class AEAD:
     AES_128_GCM: AEAD
+    CHACHA20_POLY1305: AEAD
 
 class Suite:
     def __init__(self, kem: KEM, kdf: KDF, aead: AEAD) -> None: ...

src/rust/src/backend/aead.rs

@@ -419,7 +419,7 @@ impl EvpAead {
 }
 
 #[pyo3::pyclass(frozen, module = "cryptography.hazmat.bindings._rust.openssl.aead")]
-struct ChaCha20Poly1305 {
+pub(crate) struct ChaCha20Poly1305 {
     #[cfg(any(CRYPTOGRAPHY_IS_BORINGSSL, CRYPTOGRAPHY_IS_AWSLC))]
     ctx: EvpAead,
     #[cfg(any(CRYPTOGRAPHY_OPENSSL_320_OR_GREATER, CRYPTOGRAPHY_IS_LIBRESSL))]
@@ -436,7 +436,10 @@ struct ChaCha20Poly1305 {
 #[pyo3::pymethods]
 impl ChaCha20Poly1305 {
     #[new]
-    fn new(py: pyo3::Python<'_>, key: pyo3::Py<pyo3::PyAny>) -> CryptographyResult<Self> {
+    pub(crate) fn new(
+        py: pyo3::Python<'_>,
+        key: pyo3::Py<pyo3::PyAny>,
+    ) -> CryptographyResult<Self> {
         let key_buf = key.extract::<CffiBuf<'_>>(py)?;
         if key_buf.as_bytes().len() != 32 {
             return Err(CryptographyError::from(
@@ -495,7 +498,7 @@ impl ChaCha20Poly1305 {
     }
 
     #[pyo3(signature = (nonce, data, associated_data))]
-    fn encrypt<'p>(
+    pub(crate) fn encrypt<'p>(
         &self,
         py: pyo3::Python<'p>,
         nonce: CffiBuf<'_>,
@@ -553,7 +556,7 @@ impl ChaCha20Poly1305 {
     }
 
     #[pyo3(signature = (nonce, data, associated_data))]
-    fn decrypt<'p>(
+    pub(crate) fn decrypt<'p>(
         &self,
         py: pyo3::Python<'p>,
         nonce: CffiBuf<'_>,

src/rust/src/backend/hpke.rs

@@ -4,7 +4,7 @@
 
 use pyo3::types::{PyAnyMethods, PyBytesMethods};
 
-use crate::backend::aead::AesGcm;
+use crate::backend::aead::{AesGcm, ChaCha20Poly1305};
 use crate::backend::kdf::{hkdf_extract, HkdfExpand};
 use crate::backend::x25519;
 use crate::buf::CffiBuf;
@@ -29,6 +29,11 @@ mod aead_params {
     pub const AES_128_GCM_NK: usize = 16;
     pub const AES_128_GCM_NN: usize = 12;
     pub const AES_128_GCM_NT: usize = 16;
+
+    pub const CHACHA20_POLY1305_ID: u16 = 0x0003;
+    pub const CHACHA20_POLY1305_NK: usize = 32;
+    pub const CHACHA20_POLY1305_NN: usize = 12;
+    pub const CHACHA20_POLY1305_NT: usize = 16;
 }
 
 #[allow(clippy::upper_case_acronyms)]
@@ -70,10 +75,42 @@ pub(crate) enum KDF {
 #[derive(Clone, PartialEq, Eq, Hash)]
 pub(crate) enum AEAD {
     AES_128_GCM,
+    CHACHA20_POLY1305,
+}
+
+impl AEAD {
+    fn id(&self) -> u16 {
+        match self {
+            AEAD::AES_128_GCM => aead_params::AES_128_GCM_ID,
+            AEAD::CHACHA20_POLY1305 => aead_params::CHACHA20_POLY1305_ID,
+        }
+    }
+
+    fn key_length(&self) -> usize {
+        match self {
+            AEAD::AES_128_GCM => aead_params::AES_128_GCM_NK,
+            AEAD::CHACHA20_POLY1305 => aead_params::CHACHA20_POLY1305_NK,
+        }
+    }
+
+    fn nonce_length(&self) -> usize {
+        match self {
+            AEAD::AES_128_GCM => aead_params::AES_128_GCM_NN,
+            AEAD::CHACHA20_POLY1305 => aead_params::CHACHA20_POLY1305_NN,
+        }
+    }
+
+    fn tag_length(&self) -> usize {
+        match self {
+            AEAD::AES_128_GCM => aead_params::AES_128_GCM_NT,
+            AEAD::CHACHA20_POLY1305 => aead_params::CHACHA20_POLY1305_NT,
+        }
+    }
 }
 
 #[pyo3::pyclass(frozen, module = "cryptography.hazmat.bindings._rust.openssl.hpke")]
 pub(crate) struct Suite {
+    aead: AEAD,
     kem_suite_id: [u8; 5],
     hpke_suite_id: [u8; 10],
 }
@@ -235,6 +272,50 @@ impl Suite {
         self.hkdf_expand(py, prk, &labeled_info, length)
     }
 
+    fn aead_encrypt<'p>(
+        &self,
+        py: pyo3::Python<'p>,
+        key: &pyo3::Bound<'_, pyo3::types::PyBytes>,
+        nonce: &pyo3::Bound<'_, pyo3::types::PyBytes>,
+        plaintext: CffiBuf<'_>,
+        aad: Option<CffiBuf<'_>>,
+    ) -> CryptographyResult<pyo3::Bound<'p, pyo3::types::PyBytes>> {
+        let key_obj = key.clone().unbind().into_any();
+        let nonce_buf = CffiBuf::from_bytes(py, nonce.as_bytes());
+        match &self.aead {
+            AEAD::AES_128_GCM => {
+                let cipher = AesGcm::new(py, key_obj)?;
+                cipher.encrypt(py, nonce_buf, plaintext, aad)
+            }
+            AEAD::CHACHA20_POLY1305 => {
+                let cipher = ChaCha20Poly1305::new(py, key_obj)?;
+                cipher.encrypt(py, nonce_buf, plaintext, aad)
+            }
+        }
+    }
+
+    fn aead_decrypt<'p>(
+        &self,
+        py: pyo3::Python<'p>,
+        key: &pyo3::Bound<'_, pyo3::types::PyBytes>,
+        nonce: &pyo3::Bound<'_, pyo3::types::PyBytes>,
+        ciphertext: &[u8],
+        aad: Option<CffiBuf<'_>>,
+    ) -> CryptographyResult<pyo3::Bound<'p, pyo3::types::PyBytes>> {
+        let key_obj = key.clone().unbind().into_any();
+        let nonce_buf = CffiBuf::from_bytes(py, nonce.as_bytes());
+        match &self.aead {
+            AEAD::AES_128_GCM => {
+                let cipher = AesGcm::new(py, key_obj)?;
+                cipher.decrypt(py, nonce_buf, CffiBuf::from_bytes(py, ciphertext), aad)
+            }
+            AEAD::CHACHA20_POLY1305 => {
+                let cipher = ChaCha20Poly1305::new(py, key_obj)?;
+                cipher.decrypt(py, nonce_buf, CffiBuf::from_bytes(py, ciphertext), aad)
+            }
+        }
+    }
+
     fn encrypt_inner<'p>(
         &self,
         py: pyo3::Python<'p>,
@@ -248,8 +329,7 @@ impl Suite {
         let (shared_secret, enc) = self.encap(py, public_key)?;
         let (key, base_nonce) = self.key_schedule(py, shared_secret.as_bytes(), info_bytes)?;
 
-        let aesgcm = AesGcm::new(py, pyo3::types::PyBytes::new(py, &key).unbind().into_any())?;
-        let ct = aesgcm.encrypt(py, CffiBuf::from_bytes(py, &base_nonce), plaintext, aad)?;
+        let ct = self.aead_encrypt(py, &key, &base_nonce, plaintext, aad)?;
 
         let enc_bytes = enc.as_bytes();
         let ct_bytes = ct.as_bytes();
@@ -273,7 +353,7 @@ impl Suite {
         aad: Option<CffiBuf<'_>>,
     ) -> CryptographyResult<pyo3::Bound<'p, pyo3::types::PyBytes>> {
         let ct_bytes = ciphertext.as_bytes();
-        if ct_bytes.len() < kem_params::X25519_NENC + aead_params::AES_128_GCM_NT {
+        if ct_bytes.len() < kem_params::X25519_NENC + self.aead.tag_length() {
             return Err(CryptographyError::from(exceptions::InvalidTag::new_err(())));
         }
 
@@ -286,23 +366,17 @@ impl Suite {
             .map_err(|_| CryptographyError::from(exceptions::InvalidTag::new_err(())))?;
         let (key, base_nonce) = self.key_schedule(py, shared_secret.as_bytes(), info_bytes)?;
 
-        let aesgcm = AesGcm::new(py, pyo3::types::PyBytes::new(py, &key).unbind().into_any())?;
-        aesgcm.decrypt(
-            py,
-            CffiBuf::from_bytes(py, &base_nonce),
-            CffiBuf::from_bytes(py, ct),
-            aad,
-        )
+        self.aead_decrypt(py, &key, &base_nonce, ct, aad)
     }
 
-    fn key_schedule(
+    fn key_schedule<'p>(
         &self,
-        py: pyo3::Python<'_>,
+        py: pyo3::Python<'p>,
         shared_secret: &[u8],
         info: &[u8],
     ) -> CryptographyResult<(
-        [u8; aead_params::AES_128_GCM_NK],
-        [u8; aead_params::AES_128_GCM_NN],
+        pyo3::Bound<'p, pyo3::types::PyBytes>,
+        pyo3::Bound<'p, pyo3::types::PyBytes>,
     )> {
         let psk_id_hash = self.hpke_labeled_extract(py, None, b"psk_id_hash", b"")?;
         let info_hash = self.hpke_labeled_extract(py, None, b"info_hash", info)?;
@@ -312,34 +386,29 @@ impl Suite {
 
         let secret = self.hpke_labeled_extract(py, Some(shared_secret), b"secret", b"")?;
 
-        let key_vec = self.hpke_labeled_expand(
+        let key = self.hpke_labeled_expand(
             py,
             &secret,
             b"key",
             &key_schedule_context,
-            aead_params::AES_128_GCM_NK,
+            self.aead.key_length(),
         )?;
-        let nonce_vec = self.hpke_labeled_expand(
+        let base_nonce = self.hpke_labeled_expand(
             py,
             &secret,
             b"base_nonce",
             &key_schedule_context,
-            aead_params::AES_128_GCM_NN,
+            self.aead.nonce_length(),
         )?;
 
-        let mut key = [0u8; aead_params::AES_128_GCM_NK];
-        let mut base_nonce = [0u8; aead_params::AES_128_GCM_NN];
-        key.copy_from_slice(key_vec.as_bytes());
-        base_nonce.copy_from_slice(nonce_vec.as_bytes());
-
         Ok((key, base_nonce))
     }
 }
 
 #[pyo3::pymethods]
 impl Suite {
     #[new]
-    fn new(_kem: KEM, _kdf: KDF, _aead: AEAD) -> CryptographyResult<Suite> {
+    fn new(_kem: KEM, _kdf: KDF, aead: AEAD) -> CryptographyResult<Suite> {
         // Build suite IDs
         let mut kem_suite_id = [0u8; 5];
         kem_suite_id[..3].copy_from_slice(b"KEM");
@@ -349,9 +418,10 @@ impl Suite {
         hpke_suite_id[..4].copy_from_slice(b"HPKE");
         hpke_suite_id[4..6].copy_from_slice(&kem_params::X25519_ID.to_be_bytes());
         hpke_suite_id[6..8].copy_from_slice(&kdf_params::HKDF_SHA256_ID.to_be_bytes());
-        hpke_suite_id[8..10].copy_from_slice(&aead_params::AES_128_GCM_ID.to_be_bytes());
+        hpke_suite_id[8..10].copy_from_slice(&aead.id().to_be_bytes());
 
         Ok(Suite {
+            aead,
             kem_suite_id,
             hpke_suite_id,
         })

tests/hazmat/primitives/test_hpke.py

@@ -23,6 +23,7 @@
 
 SUPPORTED_SUITES = [
     (KEM.X25519, KDF.HKDF_SHA256, AEAD.AES_128_GCM),
+    (KEM.X25519, KDF.HKDF_SHA256, AEAD.CHACHA20_POLY1305),
 ]
 
 
@@ -224,18 +225,25 @@ def test_vector_decryption(self, subtests):
             lambda f: json.load(f),
         )
 
+        aead_map = {
+            0x0001: AEAD.AES_128_GCM,
+            0x0003: AEAD.CHACHA20_POLY1305,
+        }
+
         for vector in vectors:
-            # Currently support: mode 0 (Base), X25519, HKDF-SHA256, AES-GCM
+            # Support: mode 0 (Base), X25519, HKDF-SHA256,
+            # AES-128-GCM or ChaCha20Poly1305
             if not (
                 vector["mode"] == 0
                 and vector["kem_id"] == 0x0020
                 and vector["kdf_id"] == 0x0001
-                and vector["aead_id"] == 0x0001
+                and vector["aead_id"] in aead_map
             ):
                 continue
 
             with subtests.test():
-                suite = Suite(KEM.X25519, KDF.HKDF_SHA256, AEAD.AES_128_GCM)
+                aead = aead_map[vector["aead_id"]]
+                suite = Suite(KEM.X25519, KDF.HKDF_SHA256, aead)
 
                 sk_r_bytes = bytes.fromhex(vector["skRm"])
                 sk_r = x25519.X25519PrivateKey.from_private_bytes(sk_r_bytes)