Aggressively type-check for text.

2 files changed, 66 insertions, 0 deletions

diff --git a/cryptography/hazmat/primitives/kdf/hkdf.py b/cryptography/hazmat/primitives/kdf/hkdf.py
index ce264f9b..959e8c9b 100644
--- a/cryptography/hazmat/primitives/kdf/hkdf.py
+++ b/cryptography/hazmat/primitives/kdf/hkdf.py
@@ -36,11 +36,19 @@ class HKDF(object):
 
         self._length = length
 
+        if isinstance(salt, six.text_type):
+            raise TypeError(
+                "Unicode-objects must be encoded before using them as a salt.")
+
         if salt is None:
             salt = b"\x00" * (self._algorithm.digest_size // 8)
 
         self._salt = salt
 
+        if isinstance(info, six.text_type):
+            raise TypeError(
+                "Unicode-objects must be encoded before using them as info.")
+
         if info is None:
             info = b""
 
@@ -85,6 +93,12 @@ class HKDF(object):
         return b"".join(output)[:self._length]
 
     def derive(self, key_material):
+        if isinstance(key_material, six.text_type):
+            raise TypeError(
+                "Unicode-objects must be encoded before using them as key "
+                "material."
+            )
+
         if self._used:
             raise exceptions.AlreadyFinalized
 
@@ -92,5 +106,17 @@ class HKDF(object):
         return self._expand(self._extract(key_material))
 
     def verify(self, key_material, expected_key):
+        if isinstance(key_material, six.text_type):
+            raise TypeError(
+                "Unicode-objects must be encoded before using them as key "
+                "material."
+            )
+
+        if isinstance(expected_key, six.text_type):
+            raise TypeError(
+                "Unicode-objects must be encoded before using them as "
+                "expected key material."
+            )
+
         if not constant_time.bytes_eq(self.derive(key_material), expected_key):
             raise exceptions.InvalidKey
diff --git a/tests/hazmat/primitives/test_hkdf.py b/tests/hazmat/primitives/test_hkdf.py
index 53bc098d..bd896cef 100644
--- a/tests/hazmat/primitives/test_hkdf.py
+++ b/tests/hazmat/primitives/test_hkdf.py
@@ -83,3 +83,43 @@ class TestHKDF(object):
 
         with pytest.raises(exceptions.InvalidKey):
             hkdf.verify(b'\x02' * 16, b'gJ\xfb{\xb1Oi\xc5sMC\xb7\xe4@\xf7u')
+
+    def test_unicode_typeerror(self, backend):
+        with pytest.raises(TypeError):
+            HKDF(hashes.SHA256(), 16, salt=u'foo', info=None, backend=backend)
+
+        with pytest.raises(TypeError):
+            HKDF(hashes.SHA256(), 16, salt=None, info=u'foo', backend=backend)
+
+        with pytest.raises(TypeError):
+            hkdf = HKDF(
+                hashes.SHA256(),
+                16,
+                salt=None,
+                info=None,
+                backend=backend
+            )
+
+            hkdf.derive(u'foo')
+
+        with pytest.raises(TypeError):
+            hkdf = HKDF(
+                hashes.SHA256(),
+                16,
+                salt=None,
+                info=None,
+                backend=backend
+            )
+
+            hkdf.verify(u'foo', b'bar')
+
+        with pytest.raises(TypeError):
+            hkdf = HKDF(
+                hashes.SHA256(),
+                16,
+                salt=None,
+                info=None,
+                backend=backend
+            )
+
+            hkdf.verify(b'foo', u'bar')