Support byteslike in HKDF and PBKDF2HMAC (#4707)

* support byteslike in HKDF

* support byteslike in PBKDF2HMAC

* add missing docs

2 files changed, 33 insertions, 0 deletions

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()