implement pbkdf2hmac in rust

src/cryptography/hazmat/backends/openssl/backend.py

@@ -129,23 +129,7 @@ def argon2_supported(self) -> bool:
             return hasattr(rust_openssl.kdf.Argon2id, "derive")
 
     def hmac_supported(self, algorithm: hashes.HashAlgorithm) -> bool:
-        # FIPS mode still allows SHA1 for HMAC
-        if self._fips_enabled and isinstance(algorithm, hashes.SHA1):
-            return True
-        if rust_openssl.CRYPTOGRAPHY_IS_AWSLC:
-            return isinstance(
-                algorithm,
-                (
-                    hashes.SHA1,
-                    hashes.SHA224,
-                    hashes.SHA256,
-                    hashes.SHA384,
-                    hashes.SHA512,
-                    hashes.SHA512_224,
-                    hashes.SHA512_256,
-                ),
-            )
-        return self.hash_supported(algorithm)
+        return rust_openssl.kdf._hmac_supported(algorithm)
 
     def cipher_supported(self, cipher: CipherAlgorithm, mode: Mode) -> bool:
         if self._fips_enabled:

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

@@ -7,13 +7,19 @@ import typing
 from cryptography.hazmat.primitives.hashes import HashAlgorithm
 from cryptography.utils import Buffer
 
-def derive_pbkdf2_hmac(
-    key_material: Buffer,
-    algorithm: HashAlgorithm,
-    salt: bytes,
-    iterations: int,
-    length: int,
-) -> bytes: ...
+def _hmac_supported(algorithm: HashAlgorithm) -> bool: ...
+
+class PBKDF2HMAC:
+    def __init__(
+        self,
+        algorithm: HashAlgorithm,
+        length: int,
+        salt: bytes,
+        iterations: int,
+        backend: typing.Any = None,
+    ) -> None: ...
+    def derive(self, key_material: Buffer) -> bytes: ...
+    def verify(self, key_material: bytes, expected_key: bytes) -> None: ...
 
 class Scrypt:
     def __init__(

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

@@ -4,59 +4,10 @@
 
 from __future__ import annotations
 
-import typing
-
-from cryptography import utils
-from cryptography.exceptions import (
-    AlreadyFinalized,
-    InvalidKey,
-    UnsupportedAlgorithm,
-    _Reasons,
-)
 from cryptography.hazmat.bindings._rust import openssl as rust_openssl
-from cryptography.hazmat.primitives import constant_time, hashes
 from cryptography.hazmat.primitives.kdf import KeyDerivationFunction
 
+PBKDF2HMAC = rust_openssl.kdf.PBKDF2HMAC
+KeyDerivationFunction.register(PBKDF2HMAC)
 
-class PBKDF2HMAC(KeyDerivationFunction):
-    def __init__(
-        self,
-        algorithm: hashes.HashAlgorithm,
-        length: int,
-        salt: bytes,
-        iterations: int,
-        backend: typing.Any = None,
-    ):
-        from cryptography.hazmat.backends.openssl.backend import (
-            backend as ossl,
-        )
-
-        if not ossl.pbkdf2_hmac_supported(algorithm):
-            raise UnsupportedAlgorithm(
-                f"{algorithm.name} is not supported for PBKDF2.",
-                _Reasons.UNSUPPORTED_HASH,
-            )
-        self._used = False
-        self._algorithm = algorithm
-        self._length = length
-        utils._check_bytes("salt", salt)
-        self._salt = salt
-        self._iterations = iterations
-
-    def derive(self, key_material: utils.Buffer) -> bytes:
-        if self._used:
-            raise AlreadyFinalized("PBKDF2 instances can only be used once.")
-        self._used = True
-
-        return rust_openssl.kdf.derive_pbkdf2_hmac(
-            key_material,
-            self._algorithm,
-            self._salt,
-            self._iterations,
-            self._length,
-        )
-
-    def verify(self, key_material: bytes, expected_key: bytes) -> None:
-        derived_key = self.derive(key_material)
-        if not constant_time.bytes_eq(derived_key, expected_key):
-            raise InvalidKey("Keys do not match.")
+__all__ = ["PBKDF2HMAC"]

src/rust/src/backend/kdf.rs

@@ -15,7 +15,7 @@ use crate::buf::CffiBuf;
 use crate::error::{CryptographyError, CryptographyResult};
 use crate::exceptions;
 
-#[pyo3::pyfunction]
+// TODO: remove this function
 pub(crate) fn derive_pbkdf2_hmac<'p>(
     py: pyo3::Python<'p>,
     key_material: CffiBuf<'_>,
@@ -32,6 +32,130 @@ pub(crate) fn derive_pbkdf2_hmac<'p>(
     })?)
 }
 
+#[pyo3::pyfunction]
+pub(crate) fn _hmac_supported(
+    py: pyo3::Python<'_>,
+    algorithm: &pyo3::Bound<'_, pyo3::PyAny>,
+) -> CryptographyResult<bool> {
+    let fips_enabled = cryptography_openssl::fips::is_enabled();
+
+    // Get algorithm name
+    let name = algorithm
+        .getattr(pyo3::intern!(py, "name"))?
+        .extract::<pyo3::pybacked::PyBackedStr>()?;
+
+    // FIPS mode still allows SHA1 for HMAC
+    if fips_enabled && name == "sha1" {
+        return Ok(true);
+    }
+
+    cfg_if::cfg_if! {
+        if #[cfg(CRYPTOGRAPHY_IS_AWSLC)] {
+            // AWS-LC only supports specific hash algorithms
+            Ok(matches!(
+                name.as_ref(),
+                "sha1" | "sha224" | "sha256" | "sha384" | "sha512" | "sha512-224" | "sha512-256"
+            ))
+        } else {
+            Ok(hashes::message_digest_from_algorithm(py, algorithm).is_ok())
+        }
+    }
+}
+
+#[pyo3::pyclass(
+    module = "cryptography.hazmat.primitives.kdf.pbkdf2",
+    name = "PBKDF2HMAC"
+)]
+struct Pbkdf2Hmac {
+    algorithm: pyo3::Py<pyo3::PyAny>,
+    salt: pyo3::Py<pyo3::types::PyBytes>,
+    length: usize,
+    iterations: usize,
+    used: bool,
+}
+
+#[pyo3::pymethods]
+impl Pbkdf2Hmac {
+    #[new]
+    #[pyo3(signature = (algorithm, length, salt, iterations, backend=None))]
+    fn new(
+        py: pyo3::Python<'_>,
+        algorithm: pyo3::Py<pyo3::PyAny>,
+        length: usize,
+        salt: pyo3::Py<pyo3::types::PyBytes>,
+        iterations: usize,
+        backend: Option<pyo3::Bound<'_, pyo3::PyAny>>,
+    ) -> CryptographyResult<Self> {
+        _ = backend;
+
+        let algorithm_bound = algorithm.bind(py);
+        if !_hmac_supported(py, algorithm_bound)? {
+            let name = algorithm_bound
+                .getattr(pyo3::intern!(py, "name"))?
+                .extract::<pyo3::pybacked::PyBackedStr>()?;
+            return Err(CryptographyError::from(
+                exceptions::UnsupportedAlgorithm::new_err((
+                    format!("{name} is not supported for PBKDF2."),
+                    exceptions::Reasons::UNSUPPORTED_HASH,
+                )),
+            ));
+        }
+
+        Ok(Pbkdf2Hmac {
+            algorithm,
+            salt,
+            length,
+            iterations,
+            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 md = hashes::message_digest_from_algorithm(py, algorithm_bound)?;
+
+        Ok(pyo3::types::PyBytes::new_with(py, self.length, |b| {
+            openssl::pkcs5::pbkdf2_hmac(
+                key_material.as_bytes(),
+                self.salt.as_bytes(py),
+                self.iterations,
+                md,
+                b,
+            )
+            .unwrap();
+            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::pyclass(module = "cryptography.hazmat.primitives.kdf.scrypt")]
 struct Scrypt {
     #[cfg(not(CRYPTOGRAPHY_IS_LIBRESSL))]
@@ -673,13 +797,5 @@ impl HkdfExpand {
 #[pyo3::pymodule]
 pub(crate) mod kdf {
     #[pymodule_export]
-    use super::derive_pbkdf2_hmac;
-    #[pymodule_export]
-    use super::Argon2id;
-    #[pymodule_export]
-    use super::Hkdf;
-    #[pymodule_export]
-    use super::HkdfExpand;
-    #[pymodule_export]
-    use super::Scrypt;
+    use super::{Argon2id, Hkdf, HkdfExpand, Pbkdf2Hmac, Scrypt, _hmac_supported};
 }

tests/bench/test_pbkdf2hmac.py

@@ -0,0 +1,14 @@
+# This file is dual licensed under the terms of the Apache License, Version
+# 2.0, and the BSD License. See the LICENSE file in the root of this repository
+# for complete details.
+
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+
+
+def test_pbkdf2hmac(benchmark):
+    def bench():
+        pbkdf2 = PBKDF2HMAC(hashes.SHA256(), 64, b"salt", 512)
+        pbkdf2.derive(b"0" * 64)
+
+    benchmark(bench)