migrate X963KDF to rust (#13789)

* migrate X963KDF to rust

* fix test

* don't skip test on 32-bit

Author: reaperhulk

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

@@ -80,3 +80,14 @@ class HKDFExpand:
     def derive(self, key_material: Buffer) -> bytes: ...
     def derive_into(self, key_material: Buffer, buffer: Buffer) -> int: ...
     def verify(self, key_material: bytes, expected_key: bytes) -> None: ...
+
+class X963KDF:
+    def __init__(
+        self,
+        algorithm: HashAlgorithm,
+        length: int,
+        sharedinfo: bytes | None,
+        backend: typing.Any = None,
+    ) -> None: ...
+    def derive(self, key_material: Buffer) -> bytes: ...
+    def verify(self, key_material: bytes, expected_key: bytes) -> None: ...

src/cryptography/hazmat/primitives/kdf/x963kdf.py

@@ -4,58 +4,10 @@
 
 from __future__ import annotations
 
-import typing
-
-from cryptography import utils
-from cryptography.exceptions import AlreadyFinalized, InvalidKey
-from cryptography.hazmat.primitives import constant_time, hashes
+from cryptography.hazmat.bindings._rust import openssl as rust_openssl
 from cryptography.hazmat.primitives.kdf import KeyDerivationFunction
 
+X963KDF = rust_openssl.kdf.X963KDF
+KeyDerivationFunction.register(X963KDF)
 
-def _int_to_u32be(n: int) -> bytes:
-    return n.to_bytes(length=4, byteorder="big")
-
-
-class X963KDF(KeyDerivationFunction):
-    def __init__(
-        self,
-        algorithm: hashes.HashAlgorithm,
-        length: int,
-        sharedinfo: bytes | None,
-        backend: typing.Any = None,
-    ):
-        max_len = algorithm.digest_size * (2**32 - 1)
-        if length > max_len:
-            raise ValueError(f"Cannot derive keys larger than {max_len} bits.")
-        if sharedinfo is not None:
-            utils._check_bytes("sharedinfo", sharedinfo)
-
-        self._algorithm = algorithm
-        self._length = length
-        self._sharedinfo = sharedinfo
-        self._used = False
-
-    def derive(self, key_material: utils.Buffer) -> bytes:
-        if self._used:
-            raise AlreadyFinalized
-        self._used = True
-        utils._check_byteslike("key_material", key_material)
-        output = [b""]
-        outlen = 0
-        counter = 1
-
-        while self._length > outlen:
-            h = hashes.Hash(self._algorithm)
-            h.update(key_material)
-            h.update(_int_to_u32be(counter))
-            if self._sharedinfo is not None:
-                h.update(self._sharedinfo)
-            output.append(h.finalize())
-            outlen += len(output[-1])
-            counter += 1
-
-        return b"".join(output)[: self._length]
-
-    def verify(self, key_material: bytes, expected_key: bytes) -> None:
-        if not constant_time.bytes_eq(self.derive(key_material), expected_key):
-            raise InvalidKey
+__all__ = ["X963KDF"]

src/rust/src/backend/kdf.rs

@@ -834,8 +834,116 @@ impl HkdfExpand {
     }
 }
 
+// NO-COVERAGE-START
+#[pyo3::pyclass(
+    module = "cryptography.hazmat.primitives.kdf.x963kdf",
+    name = "X963KDF"
+)]
+// NO-COVERAGE-END
+struct X963Kdf {
+    algorithm: pyo3::Py<pyo3::PyAny>,
+    length: usize,
+    sharedinfo: Option<pyo3::Py<pyo3::types::PyBytes>>,
+    used: bool,
+}
+
+#[pyo3::pymethods]
+impl X963Kdf {
+    #[new]
+    #[pyo3(signature = (algorithm, length, sharedinfo, backend=None))]
+    fn new(
+        py: pyo3::Python<'_>,
+        algorithm: pyo3::Py<pyo3::PyAny>,
+        length: usize,
+        sharedinfo: Option<pyo3::Py<pyo3::types::PyBytes>>,
+        backend: Option<pyo3::Bound<'_, pyo3::PyAny>>,
+    ) -> CryptographyResult<Self> {
+        _ = backend;
+
+        let digest_size = algorithm
+            .bind(py)
+            .getattr(pyo3::intern!(py, "digest_size"))?
+            .extract::<usize>()?;
+
+        let max_len = digest_size.saturating_mul(u32::MAX as usize);
+
+        if length > max_len {
+            return Err(CryptographyError::from(
+                pyo3::exceptions::PyValueError::new_err(format!(
+                    "Cannot derive keys larger than {max_len} bits."
+                )),
+            ));
+        }
+
+        Ok(X963Kdf {
+            algorithm,
+            length,
+            sharedinfo,
+            used: false,
+        })
+    }
+
+    fn derive<'p>(
+        &mut self,
+        py: pyo3::Python<'p>,
+        key_material: CffiBuf<'_>,
+    ) -> CryptographyResult<pyo3::Bound<'p, pyo3::types::PyBytes>> {
+        if self.used {
+            return Err(exceptions::already_finalized_error());
+        }
+        self.used = true;
+
+        let algorithm_bound = self.algorithm.bind(py);
+        let digest_size = algorithm_bound
+            .getattr(pyo3::intern!(py, "digest_size"))?
+            .extract::<usize>()?;
+
+        Ok(pyo3::types::PyBytes::new_with(py, self.length, |output| {
+            let mut pos = 0usize;
+            let mut counter = 1u32;
+
+            while pos < self.length {
+                let mut hash_obj = hashes::Hash::new(py, algorithm_bound, None)?;
+                hash_obj.update_bytes(key_material.as_bytes())?;
+                hash_obj.update_bytes(&counter.to_be_bytes())?;
+                if let Some(ref sharedinfo) = self.sharedinfo {
+                    hash_obj.update_bytes(sharedinfo.as_bytes(py))?;
+                }
+                let block = hash_obj.finalize(py)?;
+                let block_bytes = block.as_bytes();
+
+                let copy_len = (self.length - pos).min(digest_size);
+                output[pos..pos + copy_len].copy_from_slice(&block_bytes[..copy_len]);
+                pos += copy_len;
+                counter += 1;
+            }
+
+            Ok(())
+        })?)
+    }
+
+    fn verify(
+        &mut self,
+        py: pyo3::Python<'_>,
+        key_material: CffiBuf<'_>,
+        expected_key: CffiBuf<'_>,
+    ) -> CryptographyResult<()> {
+        let actual = self.derive(py, key_material)?;
+        let actual_bytes = actual.as_bytes();
+        let expected_bytes = expected_key.as_bytes();
+
+        if !constant_time::bytes_eq(actual_bytes, expected_bytes) {
+            return Err(CryptographyError::from(exceptions::InvalidKey::new_err(
+                "Keys do not match.",
+            )));
+        }
+
+        Ok(())
+    }
+}
+
 #[pyo3::pymodule(gil_used = false)]
 pub(crate) mod kdf {
     #[pymodule_export]
-    use super::{Argon2id, Hkdf, HkdfExpand, Pbkdf2Hmac, Scrypt};
+    use super::{Argon2id, Hkdf, HkdfExpand, Pbkdf2Hmac, Scrypt, X963Kdf};
 }

tests/hazmat/primitives/test_x963kdf.py

@@ -4,6 +4,7 @@
 
 
 import binascii
+import sys
 
 import pytest
 
@@ -15,8 +16,9 @@
 class TestX963KDF:
     def test_length_limit(self, backend):
         big_length = hashes.SHA256().digest_size * (2**32 - 1) + 1
+        error = OverflowError if sys.maxsize <= 2**31 else ValueError
 
-        with pytest.raises(ValueError):
+        with pytest.raises(error):
             X963KDF(hashes.SHA256(), big_length, None, backend)
 
     def test_already_finalized(self, backend):