Support byteslike in HKDF and PBKDF2HMAC (#4707)

* support byteslike in HKDF

* support byteslike in PBKDF2HMAC

* add missing docs

6 files changed, 42 insertions, 6 deletions

diff --git a/docs/hazmat/primitives/key-derivation-functions.rst b/docs/hazmat/primitives/key-derivation-functions.rst
index 0f3d1afe..373fdbff 100644
--- a/docs/hazmat/primitives/key-derivation-functions.rst
+++ b/docs/hazmat/primitives/key-derivation-functions.rst
@@ -92,8 +92,9 @@ Different KDFs are suitable for different tasks such as:
 
     .. method:: derive(key_material)
 
-        :param bytes key_material: The input key material. For PBKDF2 this
+        :param key_material: The input key material. For PBKDF2 this
             should be a password.
+        :type key_material: :term:`bytes-like`
         :return bytes: the derived key.
         :raises cryptography.exceptions.AlreadyFinalized: This is raised when
                                                           :meth:`derive` or
@@ -199,7 +200,8 @@ Different KDFs are suitable for different tasks such as:
 
     .. method:: derive(key_material)
 
-        :param bytes key_material: The input key material.
+        :param key_material: The input key material.
+        :type key_material: :term:`bytes-like`
         :return bytes: The derived key.
         :raises TypeError: This exception is raised if ``key_material`` is not
                            ``bytes``.
diff --git a/src/cryptography/hazmat/backends/openssl/backend.py b/src/cryptography/hazmat/backends/openssl/backend.py
index 2db63a27..f74c955a 100644
--- a/src/cryptography/hazmat/backends/openssl/backend.py
+++ b/src/cryptography/hazmat/backends/openssl/backend.py
@@ -297,8 +297,9 @@ class Backend(object):
                            key_material):
         buf = self._ffi.new("unsigned char[]", length)
         evp_md = self._evp_md_non_null_from_algorithm(algorithm)
+        key_material_ptr = self._ffi.from_buffer(key_material)
         res = self._lib.PKCS5_PBKDF2_HMAC(
-            key_material,
+            key_material_ptr,
             len(key_material),
             salt,
             len(salt),
diff --git a/src/cryptography/hazmat/primitives/kdf/hkdf.py b/src/cryptography/hazmat/primitives/kdf/hkdf.py
index 27dc9c93..307f91cc 100644
--- a/src/cryptography/hazmat/primitives/kdf/hkdf.py
+++ b/src/cryptography/hazmat/primitives/kdf/hkdf.py
@@ -43,7 +43,7 @@ class HKDF(object):
         return h.finalize()
 
     def derive(self, key_material):
-        utils._check_bytes("key_material", key_material)
+        utils._check_byteslike("key_material", key_material)
         return self._hkdf_expand.derive(self._extract(key_material))
 
     def verify(self, key_material, expected_key):
@@ -98,7 +98,7 @@ class HKDFExpand(object):
         return b"".join(output)[:self._length]
 
     def derive(self, key_material):
-        utils._check_bytes("key_material", key_material)
+        utils._check_byteslike("key_material", key_material)
         if self._used:
             raise AlreadyFinalized
 
diff --git a/src/cryptography/hazmat/primitives/kdf/pbkdf2.py b/src/cryptography/hazmat/primitives/kdf/pbkdf2.py
index fbe8964d..a47b7bbc 100644
--- a/src/cryptography/hazmat/primitives/kdf/pbkdf2.py
+++ b/src/cryptography/hazmat/primitives/kdf/pbkdf2.py
@@ -41,7 +41,7 @@ class PBKDF2HMAC(object):
             raise AlreadyFinalized("PBKDF2 instances can only be used once.")
         self._used = True
 
-        utils._check_bytes("key_material", key_material)
+        utils._check_byteslike("key_material", key_material)
         return self._backend.derive_pbkdf2_hmac(
             self._algorithm,
             self._length,
diff --git a/tests/hazmat/primitives/test_hkdf.py b/tests/hazmat/primitives/test_hkdf.py
index 5d2d1867..195bfb3a 100644
--- a/tests/hazmat/primitives/test_hkdf.py
+++ b/tests/hazmat/primitives/test_hkdf.py
@@ -171,6 +171,21 @@ class TestHKDF(object):
 
         assert hkdf.derive(ikm) == binascii.unhexlify(vector["okm"])
 
+    def test_buffer_protocol(self, backend):
+        vector = load_vectors_from_file(
+            os.path.join("KDF", "hkdf-generated.txt"), load_nist_vectors
+        )[0]
+        hkdf = HKDF(
+            hashes.SHA256(),
+            int(vector["l"]),
+            salt=vector["salt"],
+            info=vector["info"],
+            backend=backend
+        )
+        ikm = bytearray(binascii.unhexlify(vector["ikm"]))
+
+        assert hkdf.derive(ikm) == binascii.unhexlify(vector["okm"])
+
 
 @pytest.mark.requires_backend_interface(interface=HMACBackend)
 class TestHKDFExpand(object):
@@ -187,6 +202,19 @@ class TestHKDFExpand(object):
 
         assert binascii.hexlify(hkdf.derive(prk)) == okm
 
+    def test_buffer_protocol(self, backend):
+        prk = bytearray(binascii.unhexlify(
+            b"077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5"
+        ))
+
+        okm = (b"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c"
+               b"5bf34007208d5b887185865")
+
+        info = binascii.unhexlify(b"f0f1f2f3f4f5f6f7f8f9")
+        hkdf = HKDFExpand(hashes.SHA256(), 42, info, backend)
+
+        assert binascii.hexlify(hkdf.derive(prk)) == okm
+
     def test_verify(self, backend):
         prk = binascii.unhexlify(
             b"077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5"
diff --git a/tests/hazmat/primitives/test_pbkdf2hmac.py b/tests/hazmat/primitives/test_pbkdf2hmac.py
index d971ebd0..0254b216 100644
--- a/tests/hazmat/primitives/test_pbkdf2hmac.py
+++ b/tests/hazmat/primitives/test_pbkdf2hmac.py
@@ -57,6 +57,11 @@ class TestPBKDF2HMAC(object):
         with pytest.raises(TypeError):
             kdf.derive(u"unicode here")
 
+    def test_buffer_protocol(self, backend):
+        kdf = PBKDF2HMAC(hashes.SHA1(), 10, b"salt", 10, default_backend())
+        data = bytearray(b"data")
+        assert kdf.derive(data) == b"\xe9n\xaa\x81\xbbt\xa4\xf6\x08\xce"
+
 
 def test_invalid_backend():
     pretend_backend = object()