From aebfba28a42b608a2c5f8bf7d45402c415b893a4 Mon Sep 17 00:00:00 2001
From: Alex Gaynor <alex.gaynor@gmail.com>
Date: Sat, 5 Nov 2016 23:01:15 -0400
Subject: [PATCH] Backport hkdf short output (#3216)

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

* added a changelog
---
 CHANGELOG.rst                                  |  7 +++++++
 src/cryptography/hazmat/primitives/kdf/hkdf.py |  2 +-
 tests/hazmat/primitives/test_hkdf.py           | 11 +++++++++++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 0bfd3281faee..9b0bf29de519 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -1,6 +1,13 @@
 Changelog
 =========
 
+1.5.3 - 2016-11-05
+~~~~~~~~~~~~~~~~~~
+
+* **SECURITY ISSUE**: Fixed a bug where ``HKDF`` would return an empty
+  byte-string if used with a ``length`` less than ``algorithm.digest_size``.
+  Credit to **Markus Döring** for reporting the issue.
+
 1.5.2 - 2016-09-26
 ~~~~~~~~~~~~~~~~~~
 
diff --git a/src/cryptography/hazmat/primitives/kdf/hkdf.py b/src/cryptography/hazmat/primitives/kdf/hkdf.py
index f738bbdc72e5..82ed9b1c85e9 100644
--- 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
index e33529c9d596..a05fd752ff55 100644
--- 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):