implement HKDF and HKDFExpand derive_into

CHANGELOG.rst

@@ -31,6 +31,10 @@ Changelog
 * Moved :class:`~cryptography.hazmat.primitives.ciphers.algorithms.Camellia`
   into  :doc:`/hazmat/decrepit/index` and deprecated it in the ``cipher`` module.
   It will be removed from the ``cipher`` module 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

@@ -646,6 +646,27 @@ HKDF
         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_into`
+                                                          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
@@ -729,6 +750,25 @@ HKDF
         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_into`
+                                                          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

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

@@ -62,6 +62,7 @@ class HKDF:
         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: ...
 
 class HKDFExpand:
@@ -73,4 +74,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/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;
 
@@ -526,6 +526,40 @@ struct Hkdf {
     used: bool,
 }
 
+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 prk = self._extract(py, 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 prk_bytes = prk.as_bytes();
+        hkdf_expand.derive_into_buffer(py, prk_bytes, output)
+    }
+}
+
 #[pyo3::pymethods]
 impl Hkdf {
     #[new]
@@ -584,27 +618,24 @@ impl Hkdf {
         hmac.finalize(py)
     }
 
+    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 = self._extract(py, key_material.as_bytes())?;
-        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 prk_bytes = prk.as_bytes();
-        let cffi_buf = CffiBuf::from_bytes(py, prk_bytes);
-        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(
@@ -640,6 +671,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]
@@ -685,44 +768,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(())
         })?)
     }

tests/hazmat/primitives/test_hkdf.py

@@ -142,6 +142,29 @@ def test_buffer_protocol(self, backend):
 
         assert hkdf.derive(ikm) == binascii.unhexlify(vector["okm"])
 
+    def test_derive_into(self):
+        hkdf = HKDF(hashes.SHA256(), 16, salt=None, info=None)
+        buf = bytearray(16)
+        n = hkdf.derive_into(b"\x01" * 16, buf)
+        assert n == 16
+        assert buf == b"gJ\xfb{\xb1Oi\xc5sMC\xb7\xe4@\xf7u"
+
+    @pytest.mark.parametrize(
+        ("buflen", "outlen"), [(15, 16), (17, 16), (22, 23), (24, 23)]
+    )
+    def test_derive_into_buffer_incorrect_size(self, buflen, outlen):
+        hkdf = HKDF(hashes.SHA256(), outlen, salt=None, info=None)
+        buf = bytearray(buflen)
+        with pytest.raises(ValueError, match="buffer must be"):
+            hkdf.derive_into(b"\x01" * 16, buf)
+
+    def test_derive_into_already_finalized(self):
+        hkdf = HKDF(hashes.SHA256(), 16, salt=None, info=None)
+        buf = bytearray(16)
+        hkdf.derive_into(b"\x01" * 16, buf)
+        with pytest.raises(AlreadyFinalized):
+            hkdf.derive_into(b"\x02" * 16, buf)
+
 
 class TestHKDFExpand:
     def test_derive(self, backend):
@@ -237,3 +260,39 @@ def test_length_limit(self):
                 big_length,
                 info=None,
             )
+
+    def test_derive_into(self):
+        prk = binascii.unhexlify(
+            b"077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5"
+        )
+
+        okm = binascii.unhexlify(
+            b"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c"
+            b"5bf34007208d5b887185865"
+        )
+
+        info = binascii.unhexlify(b"f0f1f2f3f4f5f6f7f8f9")
+        hkdf = HKDFExpand(hashes.SHA256(), 42, info)
+
+        buf = bytearray(42)
+        n = hkdf.derive_into(prk, buf)
+        assert n == 42
+        assert buf == okm
+
+    @pytest.mark.parametrize(
+        ("buflen", "outlen"), [(15, 16), (17, 16), (22, 23), (24, 23)]
+    )
+    def test_derive_into_buffer_incorrect_size(self, buflen, outlen):
+        hkdf = HKDFExpand(hashes.SHA256(), outlen, info=None)
+
+        buf = bytearray(buflen)
+        with pytest.raises(ValueError, match="buffer must be"):
+            hkdf.derive_into(b"\x00" * 16, buf)
+
+    def test_derive_into_already_finalized(self):
+        hkdf = HKDFExpand(hashes.SHA256(), 42, info=None)
+
+        buf = bytearray(42)
+        hkdf.derive_into(b"0" * 16, buf)
+        with pytest.raises(AlreadyFinalized):
+            hkdf.derive_into(b"0" * 16, buf)