implement derive_into for pbkdf2hmac

Author: reaperhulk

CHANGELOG.rst

@@ -47,8 +47,9 @@ Changelog
   private implementation will be removed in 49.0.0.
 * Added ``derive_into`` methods to
   :class:`~cryptography.hazmat.primitives.kdf.hkdf.HKDF`,
-  :class:`~cryptography.hazmat.primitives.kdf.hkdf.HKDFExpand`, and
-  :class:`~cryptography.hazmat.primitives.kdf.argon2.Argon2id` to allow
+  :class:`~cryptography.hazmat.primitives.kdf.hkdf.HKDFExpand`,
+  :class:`~cryptography.hazmat.primitives.kdf.argon2.Argon2id`, and
+  :class:`~cryptography.hazmat.primitives.kdf.pbkdf2.PBKDF2HMAC` to allow
   deriving keys directly into pre-allocated buffers.
 * Added ``encrypt_into`` methods to
   :class:`~cryptography.hazmat.primitives.ciphers.aead.AESCCM`,

docs/hazmat/primitives/key-derivation-functions.rst

@@ -282,8 +282,9 @@ PBKDF2
         :type key_material: :term:`bytes-like`
         :return bytes: the derived key.
         :raises cryptography.exceptions.AlreadyFinalized: This is raised when
-                                                          :meth:`derive` or
-                                                          :meth:`verify` is
+                                                          :meth:`derive`,
+                                                          :meth:`derive_into`,
+                                                          or :meth:`verify` is
                                                           called more than
                                                           once.
 
@@ -292,6 +293,31 @@ PBKDF2
 
         This generates and returns a new key from the supplied password.
 
+    .. method:: derive_into(key_material, buffer)
+
+        .. versionadded:: 47.0.0
+
+        :param key_material: The input key material. For PBKDF2 this
+            should be a password.
+        :type key_material: :term:`bytes-like`
+        :param buffer: A writable buffer to write the derived key into. The
+                       buffer must be equal to the length supplied in the
+                       constructor.
+        :return int: The number of bytes written to the buffer.
+        :raises TypeError: This exception is raised if ``key_material`` is not
+                           ``bytes``.
+        :raises ValueError: This exception is raised if the buffer is too small
+                           for the derived key.
+        :raises cryptography.exceptions.AlreadyFinalized: This is raised when
+                                                          :meth:`derive`,
+                                                          :meth:`derive_into`,
+                                                          or :meth:`verify` is
+                                                          called more than
+                                                          once.
+
+        This generates a new key from the supplied password and writes it
+        directly into the provided buffer.
+
     .. method:: verify(key_material, expected_key)
 
         :param bytes key_material: The input key material. This is the same as
@@ -303,8 +329,9 @@ PBKDF2
                                                     derived key does not match
                                                     the expected key.
         :raises cryptography.exceptions.AlreadyFinalized: This is raised when
-                                                          :meth:`derive` or
-                                                          :meth:`verify` is
+                                                          :meth:`derive`,
+                                                          :meth:`derive_into`,
+                                                          or :meth:`verify` is
                                                           called more than
                                                           once.
 

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

@@ -17,6 +17,7 @@ class PBKDF2HMAC:
         backend: typing.Any = None,
     ) -> None: ...
     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 Scrypt:

src/rust/src/backend/kdf.rs

@@ -15,20 +15,6 @@ use crate::buf::{CffiBuf, CffiMutBuf};
 use crate::error::{CryptographyError, CryptographyResult};
 use crate::exceptions;
 
-pub(crate) fn pbkdf2_hmac_derive<'p>(
-    py: pyo3::Python<'p>,
-    key_material: &[u8],
-    md: openssl::hash::MessageDigest,
-    salt: &[u8],
-    iterations: usize,
-    length: usize,
-) -> CryptographyResult<pyo3::Bound<'p, pyo3::types::PyBytes>> {
-    Ok(pyo3::types::PyBytes::new_with(py, length, |b| {
-        openssl::pkcs5::pbkdf2_hmac(key_material, salt, iterations, md, b).unwrap();
-        Ok(())
-    })?)
-}
-
 // NO-COVERAGE-START
 #[pyo3::pyclass(
     module = "cryptography.hazmat.primitives.kdf.pbkdf2",
@@ -43,6 +29,40 @@ struct Pbkdf2Hmac {
     used: bool,
 }
 
+impl Pbkdf2Hmac {
+    fn derive_into_buffer(
+        &mut self,
+        py: pyo3::Python<'_>,
+        key_material: &[u8],
+        output: &mut [u8],
+    ) -> CryptographyResult<usize> {
+        if self.used {
+            return Err(exceptions::already_finalized_error());
+        }
+        self.used = true;
+
+        if output.len() != self.length {
+            return Err(CryptographyError::from(
+                pyo3::exceptions::PyValueError::new_err(format!(
+                    "buffer must be {} bytes",
+                    self.length
+                )),
+            ));
+        }
+
+        openssl::pkcs5::pbkdf2_hmac(
+            key_material,
+            self.salt.as_bytes(py),
+            self.iterations,
+            self.md,
+            output,
+        )
+        .unwrap();
+
+        Ok(self.length)
+    }
+}
+
 #[pyo3::pymethods]
 impl Pbkdf2Hmac {
     #[new]
@@ -67,24 +87,24 @@ impl Pbkdf2Hmac {
         })
     }
 
+    fn derive_into(
+        &mut self,
+        py: pyo3::Python<'_>,
+        key_material: CffiBuf<'_>,
+        mut buf: CffiMutBuf<'_>,
+    ) -> CryptographyResult<usize> {
+        self.derive_into_buffer(py, key_material.as_bytes(), buf.as_mut_bytes())
+    }
+
     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;
-
-        pbkdf2_hmac_derive(
-            py,
-            key_material.as_bytes(),
-            self.md,
-            self.salt.as_bytes(py),
-            self.iterations,
-            self.length,
-        )
+        Ok(pyo3::types::PyBytes::new_with(py, self.length, |output| {
+            self.derive_into_buffer(py, key_material.as_bytes(), output)?;
+            Ok(())
+        })?)
     }
 
     fn verify(

tests/hazmat/primitives/test_pbkdf2hmac.py

@@ -61,3 +61,27 @@ def test_buffer_protocol(self, backend):
         kdf = PBKDF2HMAC(hashes.SHA1(), 10, b"salt", 10, backend)
         data = bytearray(b"data")
         assert kdf.derive(data) == b"\xe9n\xaa\x81\xbbt\xa4\xf6\x08\xce"
+
+    def test_derive_into(self, backend):
+        kdf = PBKDF2HMAC(hashes.SHA1(), 20, b"salt", 10, backend)
+        buf = bytearray(20)
+        n = kdf.derive_into(b"password", buf)
+        assert n == 20
+        # Verify the output matches what derive would produce
+        kdf2 = PBKDF2HMAC(hashes.SHA1(), 20, b"salt", 10, backend)
+        expected = kdf2.derive(b"password")
+        assert buf == expected
+
+    @pytest.mark.parametrize(("buflen", "outlen"), [(19, 20), (21, 20)])
+    def test_derive_into_buffer_incorrect_size(self, buflen, outlen, backend):
+        kdf = PBKDF2HMAC(hashes.SHA1(), outlen, b"salt", 10, backend)
+        buf = bytearray(buflen)
+        with pytest.raises(ValueError, match="buffer must be"):
+            kdf.derive_into(b"password", buf)
+
+    def test_derive_into_already_finalized(self, backend):
+        kdf = PBKDF2HMAC(hashes.SHA1(), 20, b"salt", 10, backend)
+        buf = bytearray(20)
+        kdf.derive_into(b"password", buf)
+        with pytest.raises(AlreadyFinalized):
+            kdf.derive_into(b"password2", buf)