implement derive_into on concatkdfhmac (#13799)

Author: reaperhulk

CHANGELOG.rst

@@ -49,6 +49,7 @@ Changelog
   :class:`~cryptography.hazmat.primitives.kdf.hkdf.HKDF`,
   :class:`~cryptography.hazmat.primitives.kdf.hkdf.HKDFExpand`,
   :class:`~cryptography.hazmat.primitives.kdf.concatkdf.ConcatKDFHash`,
+  :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.scrypt.Scrypt`, and

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

@@ -580,7 +580,8 @@ ConcatKDF
                                                     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.
@@ -649,13 +650,40 @@ ConcatKDF
         :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
@@ -667,7 +695,8 @@ ConcatKDF
                                                     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

@@ -118,4 +118,5 @@ class ConcatKDFHMAC:
         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: ...

src/rust/src/backend/kdf.rs

@@ -1220,6 +1220,54 @@ struct ConcatKdfHmac {
     used: bool,
 }
 
+impl ConcatKdfHmac {
+    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 digest_size = algorithm_bound
+            .getattr(pyo3::intern!(py, "digest_size"))?
+            .extract::<usize>()?;
+
+        let mut pos = 0usize;
+        let mut counter = 1u32;
+
+        while pos < self.length {
+            let mut hmac = Hmac::new_bytes(py, self.salt.as_bytes(py), algorithm_bound)?;
+            hmac.update_bytes(&counter.to_be_bytes())?;
+            hmac.update_bytes(key_material)?;
+            if let Some(ref otherinfo) = self.otherinfo {
+                hmac.update_bytes(otherinfo.as_bytes(py))?;
+            }
+            let result = hmac.finalize_bytes()?;
+
+            let copy_len = (self.length - pos).min(digest_size);
+            output[pos..pos + copy_len].copy_from_slice(&result[..copy_len]);
+            pos += copy_len;
+            counter += 1;
+        }
+
+        Ok(self.length)
+    }
+}
+
 #[pyo3::pymethods]
 impl ConcatKdfHmac {
     #[new]
@@ -1278,40 +1326,22 @@ impl ConcatKdfHmac {
         })
     }
 
+    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 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 = 1u32;
-
-            while pos < self.length {
-                let mut hmac = Hmac::new_bytes(py, self.salt.as_bytes(py), algorithm_bound)?;
-                hmac.update_bytes(&counter.to_be_bytes())?;
-                hmac.update_bytes(key_material.as_bytes())?;
-                if let Some(ref otherinfo) = self.otherinfo {
-                    hmac.update_bytes(otherinfo.as_bytes(py))?;
-                }
-                let result = hmac.finalize_bytes()?;
-
-                let copy_len = (self.length - pos).min(digest_size);
-                output[pos..pos + copy_len].copy_from_slice(&result[..copy_len]);
-                pos += copy_len;
-                counter += 1;
-            }
-
+            self.derive_into_buffer(py, key_material.as_bytes(), output)?;
             Ok(())
         })?)
     }

tests/hazmat/primitives/test_concatkdf.py

@@ -324,3 +324,38 @@ def test_unsupported_hash_algorithm(self, backend):
                 otherinfo=None,
                 backend=backend,
             )
+
+    def test_derive_into(self, backend):
+        prk = binascii.unhexlify(
+            b"013951627c1dea63ea2d7702dd24e963eef5faac6b4af7e4"
+            b"b831cde499dff1ce45f6179f741c728aa733583b02409208"
+            b"8f0af7fce1d045edbc5790931e8d5ca79c73"
+        )
+        oinfo = binascii.unhexlify(
+            b"a1b2c3d4e55e600be5f367e0e8a465f4bf2704db00c9325c"
+            b"9fbd216d12b49160b2ae5157650f43415653696421e68e"
+        )
+        ckdf = ConcatKDFHMAC(hashes.SHA512(), 32, None, oinfo, backend)
+        buf = bytearray(32)
+        n = ckdf.derive_into(prk, buf)
+        assert n == 32
+        # Verify the output matches what derive would produce
+        ckdf2 = ConcatKDFHMAC(hashes.SHA512(), 32, None, oinfo, backend)
+        expected = ckdf2.derive(prk)
+        assert buf == expected
+
+    @pytest.mark.parametrize(
+        ("buflen", "outlen"), [(31, 32), (33, 32), (16, 32), (64, 32)]
+    )
+    def test_derive_into_buffer_incorrect_size(self, buflen, outlen, backend):
+        ckdf = ConcatKDFHMAC(hashes.SHA512(), outlen, None, None, backend)
+        buf = bytearray(buflen)
+        with pytest.raises(ValueError, match="buffer must be"):
+            ckdf.derive_into(b"key", buf)
+
+    def test_derive_into_already_finalized(self, backend):
+        ckdf = ConcatKDFHMAC(hashes.SHA512(), 32, None, None, backend)
+        buf = bytearray(32)
+        ckdf.derive_into(b"key", buf)
+        with pytest.raises(AlreadyFinalized):
+            ckdf.derive_into(b"key", buf)