Issue #26798: Add BLAKE2 (blake2b and blake2s) to hashlib.

2 files changed, 217 insertions, 16 deletions

diff --git a/Lib/hashlib.py b/Lib/hashlib.py
index 348ea14..40ccdec 100644
--- a/Lib/hashlib.py
+++ b/Lib/hashlib.py
@@ -4,14 +4,14 @@
 
 __doc__ = """hashlib module - A common interface to many hash functions.
 
-new(name, data=b'') - returns a new hash object implementing the
-                      given hash function; initializing the hash
-                      using the given binary data.
+new(name, data=b'', **kwargs) - returns a new hash object implementing the
+                                given hash function; initializing the hash
+                                using the given binary data.
 
 Named constructor functions are also available, these are faster
 than using new(name):
 
-md5(), sha1(), sha224(), sha256(), sha384(), and sha512()
+md5(), sha1(), sha224(), sha256(), sha384(), sha512(), blake2b(), and blake2s()
 
 More algorithms may be available on your platform but the above are guaranteed
 to exist.  See the algorithms_guaranteed and algorithms_available attributes
@@ -54,7 +54,8 @@ More condensed:
 
 # This tuple and __get_builtin_constructor() must be modified if a new
 # always available algorithm is added.
-__always_supported = ('md5', 'sha1', 'sha224', 'sha256', 'sha384', 'sha512')
+__always_supported = ('md5', 'sha1', 'sha224', 'sha256', 'sha384', 'sha512',
+                      'blake2b', 'blake2s')
 
 algorithms_guaranteed = set(__always_supported)
 algorithms_available = set(__always_supported)
@@ -85,6 +86,10 @@ def __get_builtin_constructor(name):
             import _sha512
             cache['SHA384'] = cache['sha384'] = _sha512.sha384
             cache['SHA512'] = cache['sha512'] = _sha512.sha512
+        elif name in ('blake2b', 'blake2s'):
+            import _blake2
+            cache['blake2b'] = _blake2.blake2b
+            cache['blake2s'] = _blake2.blake2s
     except ImportError:
         pass  # no extension module, this hash is unsupported.
 
@@ -107,17 +112,23 @@ def __get_openssl_constructor(name):
         return __get_builtin_constructor(name)
 
 
-def __py_new(name, data=b''):
-    """new(name, data=b'') - Return a new hashing object using the named algorithm;
-    optionally initialized with data (which must be bytes).
+def __py_new(name, data=b'', **kwargs):
+    """new(name, data=b'', **kwargs) - Return a new hashing object using the
+    named algorithm; optionally initialized with data (which must be bytes).
     """
-    return __get_builtin_constructor(name)(data)
+    return __get_builtin_constructor(name)(data, **kwargs)
 
 
-def __hash_new(name, data=b''):
+def __hash_new(name, data=b'', **kwargs):
     """new(name, data=b'') - Return a new hashing object using the named algorithm;
     optionally initialized with data (which must be bytes).
     """
+    if name in {'blake2b', 'blake2s'}:
+        # Prefer our blake2 implementation.
+        # OpenSSL 1.1.0 comes with a limited implementation of blake2b/s.
+        # It does neither support keyed blake2 nor advanced features like
+        # salt, personal, tree hashing or SSE.
+        return __get_builtin_constructor(name)(data, **kwargs)
     try:
         return _hashlib.new(name, data)
     except ValueError:
@@ -218,6 +229,7 @@ for __func_name in __always_supported:
         import logging
         logging.exception('code for hash %s was not found.', __func_name)
 
+
 # Cleanup locals()
 del __always_supported, __func_name, __get_hash
 del __py_new, __hash_new, __get_openssl_constructor
diff --git a/Lib/test/test_hashlib.py b/Lib/test/test_hashlib.py
index b010a74..e6df09c 100644
--- a/Lib/test/test_hashlib.py
+++ b/Lib/test/test_hashlib.py
@@ -27,6 +27,14 @@ COMPILED_WITH_PYDEBUG = hasattr(sys, 'gettotalrefcount')
 c_hashlib = import_fresh_module('hashlib', fresh=['_hashlib'])
 py_hashlib = import_fresh_module('hashlib', blocked=['_hashlib'])
 
+try:
+    import _blake2
+except ImportError:
+    _blake2 = None
+
+requires_blake2 = unittest.skipUnless(_blake2, 'requires _blake2')
+
+
 def hexstr(s):
     assert isinstance(s, bytes), repr(s)
     h = "0123456789abcdef"
@@ -36,10 +44,24 @@ def hexstr(s):
     return r
 
 
+URL = "https://raw.githubusercontent.com/tiran/python_vectors/master/{}.txt"
+
+def read_vectors(hash_name):
+    with support.open_urlresource(URL.format(hash_name)) as f:
+        for line in f:
+            line = line.strip()
+            if line.startswith('#') or not line:
+                continue
+            parts = line.split(',')
+            parts[0] = bytes.fromhex(parts[0])
+            yield parts
+
+
 class HashLibTestCase(unittest.TestCase):
     supported_hash_names = ( 'md5', 'MD5', 'sha1', 'SHA1',
                              'sha224', 'SHA224', 'sha256', 'SHA256',
-                             'sha384', 'SHA384', 'sha512', 'SHA512')
+                             'sha384', 'SHA384', 'sha512', 'SHA512',
+                             'blake2b', 'blake2s')
 
     # Issue #14693: fallback modules are always compiled under POSIX
     _warn_on_extension_import = os.name == 'posix' or COMPILED_WITH_PYDEBUG
@@ -57,6 +79,11 @@ class HashLibTestCase(unittest.TestCase):
         algorithms = set()
         for algorithm in self.supported_hash_names:
             algorithms.add(algorithm.lower())
+
+        _blake2 = self._conditional_import_module('_blake2')
+        if _blake2:
+            algorithms.update({'blake2b', 'blake2s'})
+
         self.constructors_to_test = {}
         for algorithm in algorithms:
             self.constructors_to_test[algorithm] = set()
@@ -65,10 +92,10 @@ class HashLibTestCase(unittest.TestCase):
         # of hashlib.new given the algorithm name.
         for algorithm, constructors in self.constructors_to_test.items():
             constructors.add(getattr(hashlib, algorithm))
-            def _test_algorithm_via_hashlib_new(data=None, _alg=algorithm):
+            def _test_algorithm_via_hashlib_new(data=None, _alg=algorithm, **kwargs):
                 if data is None:
-                    return hashlib.new(_alg)
-                return hashlib.new(_alg, data)
+                    return hashlib.new(_alg, **kwargs)
+                return hashlib.new(_alg, data, **kwargs)
             constructors.add(_test_algorithm_via_hashlib_new)
 
         _hashlib = self._conditional_import_module('_hashlib')
@@ -100,6 +127,9 @@ class HashLibTestCase(unittest.TestCase):
         if _sha512:
             add_builtin_constructor('sha384')
             add_builtin_constructor('sha512')
+        if _blake2:
+            add_builtin_constructor('blake2s')
+            add_builtin_constructor('blake2b')
 
         super(HashLibTestCase, self).__init__(*args, **kwargs)
 
@@ -194,13 +224,13 @@ class HashLibTestCase(unittest.TestCase):
             self.assertEqual(m1.digest(), m4_copy.digest())
             self.assertEqual(m4.digest(), m4_digest)
 
-    def check(self, name, data, hexdigest):
+    def check(self, name, data, hexdigest, **kwargs):
         hexdigest = hexdigest.lower()
         constructors = self.constructors_to_test[name]
         # 2 is for hashlib.name(...) and hashlib.new(name, ...)
         self.assertGreaterEqual(len(constructors), 2)
         for hash_object_constructor in constructors:
-            m = hash_object_constructor(data)
+            m = hash_object_constructor(data, **kwargs)
             computed = m.hexdigest()
             self.assertEqual(
                     computed, hexdigest,
@@ -227,6 +257,11 @@ class HashLibTestCase(unittest.TestCase):
         self.check_no_unicode('sha384')
         self.check_no_unicode('sha512')
 
+    @requires_blake2
+    def test_no_unicode_blake2(self):
+        self.check_no_unicode('blake2b')
+        self.check_no_unicode('blake2s')
+
     def check_blocksize_name(self, name, block_size=0, digest_size=0):
         constructors = self.constructors_to_test[name]
         for hash_object_constructor in constructors:
@@ -246,6 +281,11 @@ class HashLibTestCase(unittest.TestCase):
         self.check_blocksize_name('sha384', 128, 48)
         self.check_blocksize_name('sha512', 128, 64)
 
+    @requires_blake2
+    def test_blocksize_name_blake2(self):
+        self.check_blocksize_name('blake2b', 128, 64)
+        self.check_blocksize_name('blake2s', 64, 32)
+
     def test_case_md5_0(self):
         self.check('md5', b'', 'd41d8cd98f00b204e9800998ecf8427e')
 
@@ -374,6 +414,155 @@ class HashLibTestCase(unittest.TestCase):
           "e718483d0ce769644e2e42c7bc15b4638e1f98b13b2044285632a803afa973eb"+
           "de0ff244877ea60a4cb0432ce577c31beb009c5c2c49aa2e4eadb217ad8cc09b")
 
+    def check_blake2(self, constructor, salt_size, person_size, key_size,
+                     digest_size, max_offset):
+        self.assertEqual(constructor.SALT_SIZE, salt_size)
+        for i in range(salt_size + 1):
+            constructor(salt=b'a' * i)
+        salt = b'a' * (salt_size + 1)
+        self.assertRaises(ValueError, constructor, salt=salt)
+
+        self.assertEqual(constructor.PERSON_SIZE, person_size)
+        for i in range(person_size+1):
+            constructor(person=b'a' * i)
+        person = b'a' * (person_size + 1)
+        self.assertRaises(ValueError, constructor, person=person)
+
+        self.assertEqual(constructor.MAX_DIGEST_SIZE, digest_size)
+        for i in range(1, digest_size + 1):
+            constructor(digest_size=i)
+        self.assertRaises(ValueError, constructor, digest_size=-1)
+        self.assertRaises(ValueError, constructor, digest_size=0)
+        self.assertRaises(ValueError, constructor, digest_size=digest_size+1)
+
+        self.assertEqual(constructor.MAX_KEY_SIZE, key_size)
+        for i in range(key_size+1):
+            constructor(key=b'a' * i)
+        key = b'a' * (key_size + 1)
+        self.assertRaises(ValueError, constructor, key=key)
+        self.assertEqual(constructor().hexdigest(),
+                         constructor(key=b'').hexdigest())
+
+        for i in range(0, 256):
+            constructor(fanout=i)
+        self.assertRaises(ValueError, constructor, fanout=-1)
+        self.assertRaises(ValueError, constructor, fanout=256)
+
+        for i in range(1, 256):
+            constructor(depth=i)
+        self.assertRaises(ValueError, constructor, depth=-1)
+        self.assertRaises(ValueError, constructor, depth=0)
+        self.assertRaises(ValueError, constructor, depth=256)
+
+        for i in range(0, 256):
+            constructor(node_depth=i)
+        self.assertRaises(ValueError, constructor, node_depth=-1)
+        self.assertRaises(ValueError, constructor, node_depth=256)
+
+        for i in range(0, digest_size + 1):
+            constructor(inner_size=i)
+        self.assertRaises(ValueError, constructor, inner_size=-1)
+        self.assertRaises(ValueError, constructor, inner_size=digest_size+1)
+
+        constructor(leaf_size=0)
+        constructor(leaf_size=(1<<32)-1)
+        self.assertRaises(OverflowError, constructor, leaf_size=-1)
+        self.assertRaises(OverflowError, constructor, leaf_size=1<<32)
+
+        constructor(node_offset=0)
+        constructor(node_offset=max_offset)
+        self.assertRaises(OverflowError, constructor, node_offset=-1)
+        self.assertRaises(OverflowError, constructor, node_offset=max_offset+1)
+
+        constructor(
+            string=b'',
+            key=b'',
+            salt=b'',
+            person=b'',
+            digest_size=17,
+            fanout=1,
+            depth=1,
+            leaf_size=256,
+            node_offset=512,
+            node_depth=1,
+            inner_size=7,
+            last_node=True
+        )
+
+    def blake2_rfc7693(self, constructor, md_len, in_len):
+        def selftest_seq(length, seed):
+            mask = (1<<32)-1
+            a = (0xDEAD4BAD * seed) & mask
+            b = 1
+            out = bytearray(length)
+            for i in range(length):
+                t = (a + b) & mask
+                a, b = b, t
+                out[i] = (t >> 24) & 0xFF
+            return out
+        outer = constructor(digest_size=32)
+        for outlen in md_len:
+            for inlen in in_len:
+                indata = selftest_seq(inlen, inlen)
+                key = selftest_seq(outlen, outlen)
+                unkeyed = constructor(indata, digest_size=outlen)
+                outer.update(unkeyed.digest())
+                keyed = constructor(indata, key=key, digest_size=outlen)
+                outer.update(keyed.digest())
+        return outer.hexdigest()
+
+    @requires_blake2
+    def test_blake2b(self):
+        self.check_blake2(hashlib.blake2b, 16, 16, 64, 64, (1<<64)-1)
+        b2b_md_len = [20, 32, 48, 64]
+        b2b_in_len = [0, 3, 128, 129, 255, 1024]
+        self.assertEqual(
+            self.blake2_rfc7693(hashlib.blake2b, b2b_md_len, b2b_in_len),
+            "c23a7800d98123bd10f506c61e29da5603d763b8bbad2e737f5e765a7bccd475")
+
+    @requires_blake2
+    def test_case_blake2b_0(self):
+        self.check('blake2b', b"",
+          "786a02f742015903c6c6fd852552d272912f4740e15847618a86e217f71f5419"+
+          "d25e1031afee585313896444934eb04b903a685b1448b755d56f701afe9be2ce")
+
+    @requires_blake2
+    def test_case_blake2b_1(self):
+        self.check('blake2b', b"abc",
+          "ba80a53f981c4d0d6a2797b69f12f6e94c212f14685ac4b74b12bb6fdbffa2d1"+
+          "7d87c5392aab792dc252d5de4533cc9518d38aa8dbf1925ab92386edd4009923")
+
+    @requires_blake2
+    def test_blake2b_vectors(self):
+        for msg, key, md in read_vectors('blake2b'):
+            key = bytes.fromhex(key)
+            self.check('blake2b', msg, md, key=key)
+
+    @requires_blake2
+    def test_blake2s(self):
+        self.check_blake2(hashlib.blake2s, 8, 8, 32, 32, (1<<48)-1)
+        b2s_md_len = [16, 20, 28, 32]
+        b2s_in_len = [0, 3, 64, 65, 255, 1024]
+        self.assertEqual(
+            self.blake2_rfc7693(hashlib.blake2s, b2s_md_len, b2s_in_len),
+            "6a411f08ce25adcdfb02aba641451cec53c598b24f4fc787fbdc88797f4c1dfe")
+
+    @requires_blake2
+    def test_case_blake2s_0(self):
+        self.check('blake2s', b"",
+          "69217a3079908094e11121d042354a7c1f55b6482ca1a51e1b250dfd1ed0eef9")
+
+    @requires_blake2
+    def test_case_blake2s_1(self):
+        self.check('blake2s', b"abc",
+          "508c5e8c327c14e2e1a72ba34eeb452f37458b209ed63a294d999b4c86675982")
+
+    @requires_blake2
+    def test_blake2s_vectors(self):
+        for msg, key, md in read_vectors('blake2s'):
+            key = bytes.fromhex(key)
+            self.check('blake2s', msg, md, key=key)
+
     def test_gil(self):
         # Check things work fine with an input larger than the size required
         # for multithreaded operation (which is hardwired to 2048).