Fixes #3211 -- fixed hkdf's output with short length (#3215)

pyca · Nov 6, 2016 · b924696 · b924696
1 parent 21ac453
commit b924696
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/src/cryptography/hazmat/primitives/kdf/hkdf.py b/src/cryptography/hazmat/primitives/kdf/hkdf.py
@@ -91,7 +91,7 @@ def _expand(self, key_material):
         output = [b""]
         counter = 1
 
-        while (self._algorithm.digest_size // 8) * len(output) < self._length:
+        while self._algorithm.digest_size * (len(output) - 1) < self._length:
             h = hmac.HMAC(key_material, self._algorithm, backend=self._backend)
             h.update(output[-1])
             h.update(self._info)

diff --git a/tests/hazmat/primitives/test_hkdf.py b/tests/hazmat/primitives/test_hkdf.py
@@ -142,6 +142,17 @@ def test_unicode_typeerror(self, backend):
 
             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{"
+
 
 @pytest.mark.requires_backend_interface(interface=HMACBackend)
 class TestHKDFExpand(object):