implement derive_into for kbkdfhmac

Author: reaperhulk

CHANGELOG.rst

@@ -55,6 +55,7 @@ Changelog
   :class:`~cryptography.hazmat.primitives.kdf.concatkdf.ConcatKDFHMAC`,
   :class:`~cryptography.hazmat.primitives.kdf.argon2.Argon2id`,
   :class:`~cryptography.hazmat.primitives.kdf.pbkdf2.PBKDF2HMAC`,
+  :class:`~cryptography.hazmat.primitives.kdf.kbkdf.KBKDFHMAC`,
   :class:`~cryptography.hazmat.primitives.kdf.scrypt.Scrypt`, and
   :class:`~cryptography.hazmat.primitives.kdf.x963kdf.X963KDF` to allow
   deriving keys directly into pre-allocated buffers.

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

@@ -1090,13 +1090,40 @@ KBKDF
         :raises TypeError: This exception is raised if ``key_material`` is
                             not ``bytes``.
         :raises cryptography.exceptions.AlreadyFinalized: This is raised when
-                                                          :meth:`derive` or
+                                                          :meth:`derive`,
+                                                          :meth:`derive_into`, or
                                                           :meth:`verify` is
                                                           called more than
                                                           once.
 
         Derives a new key from the input key material.
 
+    .. 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. The
+                       buffer must be equal to the length supplied in the
+                       constructor.
+        :type buffer: :term:`bytes-like`
+        :return int: the number of bytes written to the buffer.
+        :raises ValueError: This exception is raised if the buffer length does
+                           not match the specified ``length``.
+        :raises TypeError: This exception is raised if ``key_material`` or
+                           ``buffer`` is not ``bytes``.
+        :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 and writes it into
+        the provided buffer. This is useful when you want to avoid allocating
+        new memory for the derived key.
+
     .. method:: verify(key_material, expected_key)
 
         :param bytes key_material: The input key material. This is the same as
@@ -1108,7 +1135,8 @@ KBKDF
                                                     derived key does not match
                                                     the expected key.
         :raises cryptography.exceptions.AlreadyFinalized: This is raised when
-                                                          :meth:`derive` or
+                                                          :meth:`derive`,
+                                                          :meth:`derive_into`, or
                                                           :meth:`verify` is
                                                           called more than
                                                           once.

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

@@ -181,4 +181,5 @@ class KBKDFHMAC:
         break_location: int | None = 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: ...

src/rust/src/backend/kdf.rs

@@ -1971,6 +1971,15 @@ impl KbkdfHmac {
         })?)
     }
 
+    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 verify(
         &mut self,
         py: pyo3::Python<'_>,

tests/hazmat/primitives/test_kbkdf.py

@@ -113,6 +113,74 @@ def test_already_finalized(self, backend):
         with pytest.raises(AlreadyFinalized):
             kdf.verify(b"material", key)
 
+    def test_derive_into(self, backend):
+        kdf = KBKDFHMAC(
+            hashes.SHA256(),
+            Mode.CounterMode,
+            32,
+            4,
+            4,
+            CounterLocation.BeforeFixed,
+            b"label",
+            b"context",
+            None,
+            backend=backend,
+        )
+        buf = bytearray(32)
+        n = kdf.derive_into(b"material", buf)
+        assert n == 32
+        # Verify the output matches what derive would produce
+        kdf2 = KBKDFHMAC(
+            hashes.SHA256(),
+            Mode.CounterMode,
+            32,
+            4,
+            4,
+            CounterLocation.BeforeFixed,
+            b"label",
+            b"context",
+            None,
+            backend=backend,
+        )
+        expected = kdf2.derive(b"material")
+        assert buf == expected
+
+    @pytest.mark.parametrize(("buflen", "outlen"), [(31, 32), (33, 32)])
+    def test_derive_into_buffer_incorrect_size(self, buflen, outlen, backend):
+        kdf = KBKDFHMAC(
+            hashes.SHA256(),
+            Mode.CounterMode,
+            outlen,
+            4,
+            4,
+            CounterLocation.BeforeFixed,
+            b"label",
+            b"context",
+            None,
+            backend=backend,
+        )
+        buf = bytearray(buflen)
+        with pytest.raises(ValueError, match="buffer must be"):
+            kdf.derive_into(b"material", buf)
+
+    def test_derive_into_already_finalized(self, backend):
+        kdf = KBKDFHMAC(
+            hashes.SHA256(),
+            Mode.CounterMode,
+            32,
+            4,
+            4,
+            CounterLocation.BeforeFixed,
+            b"label",
+            b"context",
+            None,
+            backend=backend,
+        )
+        buf = bytearray(32)
+        kdf.derive_into(b"material", buf)
+        with pytest.raises(AlreadyFinalized):
+            kdf.derive_into(b"material2", buf)
+
     def test_key_length(self, backend):
         error = OverflowError if sys.maxsize <= 2**31 else ValueError
         with pytest.raises(error):