implement HKDF and HKDFExpand derive_into

CHANGELOG.rst

@@ -34,6 +34,10 @@ Changelog
 * Added :meth:`~cryptography.hazmat.primitives.kdf.hkdf.HKDF.extract`
   to :class:`~cryptography.hazmat.primitives.kdf.hkdf.HKDF`. The previous
   private implementation will be removed in 49.0.0.
+* Added ``derive_into`` methods to
+  :class:`~cryptography.hazmat.primitives.kdf.hkdf.HKDF` and
+  :class:`~cryptography.hazmat.primitives.kdf.hkdf.HKDFExpand` to allow
+  deriving keys directly into pre-allocated buffers.
 
 .. _v46-0-2:
 

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

@@ -665,14 +665,38 @@ HKDF
         :raises TypeError: This exception is raised if ``key_material`` is not
                            ``bytes``.
         :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.
 
         Derives a new key from the input key material by performing both the
         extract and expand operations.
 
+    .. method:: derive_into(key_material, buffer)
+
+        .. versionadded:: 47.0.0
+
+        :param key_material: The input key material.
+        :type key_material: :term:`bytes-like`
+        :param buffer: A writable buffer to write the derived key into.
+        :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.
+
+        Derives a new key from the input key material by performing both the
+        extract and expand operations, writing the result into the provided
+        buffer.
+
     .. method:: verify(key_material, expected_key)
 
         :param bytes key_material: The input key material. This is the same as
@@ -684,8 +708,9 @@ HKDF
                                                     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.
 
@@ -748,14 +773,36 @@ HKDF
         :raises TypeError: This exception is raised if ``key_material`` is not
                            ``bytes``.
         :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.
 
         Derives a new key from the input key material by only performing the
         expand operation.
 
+    .. method:: derive_into(key_material, buffer)
+
+        .. versionadded:: 47.0.0
+
+        :param bytes key_material: The input key material.
+        :param buffer: A writable buffer to write the derived key into.
+        :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.
+
+        Derives a new key from the input key material by only performing the
+        expand operation, writing the result into the provided buffer.
+
     .. method:: verify(key_material, expected_key)
 
         :param bytes key_material: The input key material. This is the same as
@@ -767,8 +814,9 @@ HKDF
                                                     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.
         :raises TypeError: This is raised if the provided ``key_material`` is

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

@@ -66,6 +66,7 @@ class HKDF:
         algorithm: HashAlgorithm, salt: bytes | None, key_material: Buffer
     ) -> bytes: ...
     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 HKDFExpand:
@@ -77,4 +78,5 @@ class HKDFExpand:
         backend: typing.Any = 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: ...

src/rust/src/backend/hmac.rs

@@ -3,7 +3,6 @@
 // for complete details.
 
 use cryptography_crypto::constant_time;
-use pyo3::types::PyBytesMethods;
 
 use crate::backend::hashes::message_digest_from_algorithm;
 use crate::buf::CffiBuf;
@@ -92,14 +91,12 @@ impl Hmac {
         py: pyo3::Python<'p>,
     ) -> CryptographyResult<pyo3::Bound<'p, pyo3::types::PyBytes>> {
         let data = self.finalize_bytes()?;
-        self.ctx = None;
         Ok(pyo3::types::PyBytes::new(py, &data))
     }
 
-    fn verify(&mut self, py: pyo3::Python<'_>, signature: &[u8]) -> CryptographyResult<()> {
-        let actual_bound = self.finalize(py)?;
-        let actual = actual_bound.as_bytes();
-        if !constant_time::bytes_eq(actual, signature) {
+    fn verify(&mut self, signature: &[u8]) -> CryptographyResult<()> {
+        let actual = self.finalize_bytes()?;
+        if !constant_time::bytes_eq(&actual, signature) {
             return Err(CryptographyError::from(
                 exceptions::InvalidSignature::new_err("Signature did not match digest."),
             ));

src/rust/src/backend/kdf.rs

@@ -11,7 +11,7 @@ use pyo3::types::{PyAnyMethods, PyBytesMethods};
 
 use crate::backend::hashes;
 use crate::backend::hmac::Hmac;
-use crate::buf::CffiBuf;
+use crate::buf::{CffiBuf, CffiMutBuf};
 use crate::error::{CryptographyError, CryptographyResult};
 use crate::exceptions;
 
@@ -549,6 +549,40 @@ fn hkdf_extract(
     hmac.finalize_bytes()
 }
 
+impl Hkdf {
+    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
+                )),
+            ));
+        }
+
+        let buf = CffiBuf::from_bytes(py, key_material);
+        let prk = hkdf_extract(py, &self.algorithm, self.salt.as_ref(), &buf)?;
+        let mut hkdf_expand = HkdfExpand::new(
+            py,
+            self.algorithm.clone_ref(py),
+            self.length,
+            self.info.as_ref().map(|i| i.clone_ref(py)),
+            None,
+        )?;
+        hkdf_expand.derive_into_buffer(py, &prk, output)
+    }
+}
+
 #[pyo3::pymethods]
 impl Hkdf {
     #[new]
@@ -610,26 +644,24 @@ impl Hkdf {
         Ok(pyo3::types::PyBytes::new(py, &prk))
     }
 
+    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;
-
-        let prk = hkdf_extract(py, &self.algorithm, self.salt.as_ref(), &key_material)?;
-        let mut hkdf_expand = HkdfExpand::new(
-            py,
-            self.algorithm.clone_ref(py),
-            self.length,
-            self.info.as_ref().map(|i| i.clone_ref(py)),
-            None,
-        )?;
-        let cffi_buf = CffiBuf::from_bytes(py, &prk);
-        hkdf_expand.derive(py, cffi_buf)
+        Ok(pyo3::types::PyBytes::new_with(py, self.length, |output| {
+            self.derive_into_buffer(py, key_material.as_bytes(), output)?;
+            Ok(())
+        })?)
     }
 
     fn verify(
@@ -665,6 +697,58 @@ struct HkdfExpand {
     used: bool,
 }
 
+impl HkdfExpand {
+    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
+                )),
+            ));
+        }
+
+        let algorithm_bound = self.algorithm.bind(py);
+        let h_prime = Hmac::new_bytes(py, key_material, algorithm_bound)?;
+        let digest_size = algorithm_bound
+            .getattr(pyo3::intern!(py, "digest_size"))?
+            .extract::<usize>()?;
+
+        let mut pos = 0usize;
+        let mut counter = 0u8;
+
+        while pos < self.length {
+            counter += 1;
+            let mut h = h_prime.copy(py)?;
+
+            let start = pos.saturating_sub(digest_size);
+            h.update_bytes(&output[start..pos])?;
+
+            h.update_bytes(self.info.as_bytes(py))?;
+            h.update_bytes(&[counter])?;
+
+            let block = h.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;
+        }
+
+        Ok(self.length)
+    }
+}
+
 #[pyo3::pymethods]
 impl HkdfExpand {
     #[new]
@@ -710,44 +794,22 @@ impl HkdfExpand {
         })
     }
 
+    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;
-
-        let algorithm_bound = self.algorithm.bind(py);
-        let h_prime = Hmac::new_bytes(py, key_material.as_bytes(), algorithm_bound)?;
-        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 = 0u8;
-
-            while pos < self.length {
-                counter += 1;
-                let mut h = h_prime.copy(py)?;
-
-                let start = pos.saturating_sub(digest_size);
-                h.update_bytes(&output[start..pos])?;
-
-                h.update_bytes(self.info.as_bytes(py))?;
-                h.update_bytes(&[counter])?;
-
-                let block = h.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;
-            }
-
+            self.derive_into_buffer(py, key_material.as_bytes(), output)?;
             Ok(())
         })?)
     }