1 files changed, 59 insertions, 2 deletions
diff --git a/tests/hazmat/primitives/test_hkdf.py b/tests/hazmat/primitives/test_hkdf.py
index e33529c9..195bfb3a 100644
--- a/tests/hazmat/primitives/test_hkdf.py
+++ b/tests/hazmat/primitives/test_hkdf.py
@@ -5,6 +5,7 @@
 from __future__ import absolute_import, division, print_function
 
 import binascii
+import os
 
 import pytest
 
@@ -15,13 +16,15 @@ from cryptography.hazmat.backends.interfaces import HMACBackend
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.kdf.hkdf import HKDF, HKDFExpand
 
-from ...utils import raises_unsupported_algorithm
+from ...utils import (
+    load_nist_vectors, load_vectors_from_file, raises_unsupported_algorithm
+)
 
 
 @pytest.mark.requires_backend_interface(interface=HMACBackend)
 class TestHKDF(object):
     def test_length_limit(self, backend):
-        big_length = 255 * (hashes.SHA256().digest_size // 8) + 1
+        big_length = 255 * hashes.SHA256().digest_size + 1
 
         with pytest.raises(ValueError):
             HKDF(
@@ -142,6 +145,47 @@ class TestHKDF(object):
 
             hkdf.verify(b"foo", u"bar")
 
+    def test_derive_short_output(self, backend):
+        hkdf = HKDF(
+            hashes.SHA256(),
+            4,
+            salt=None,
+            info=None,
+            backend=backend
+        )
+
+        assert hkdf.derive(b"\x01" * 16) == b"gJ\xfb{"
+
+    def test_derive_long_output(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 = binascii.unhexlify(vector["ikm"])
+
+        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):
@@ -158,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"