bpo-32433: Optimized HMAC digest (#5023)

The hmac module now has hmac.digest(), which provides an optimized HMAC
digest for short messages. hmac.digest() is up to three times faster
than hmac.HMAC().digest().

Signed-off-by: Christian Heimes <christian@python.org>

2 files changed, 80 insertions, 0 deletions

diff --git a/Lib/hmac.py b/Lib/hmac.py
index 121029a..93c084e 100644
--- a/Lib/hmac.py
+++ b/Lib/hmac.py
@@ -5,6 +5,13 @@ Implements the HMAC algorithm as described by RFC 2104.
 
 import warnings as _warnings
 from _operator import _compare_digest as compare_digest
+try:
+    import _hashlib as _hashopenssl
+except ImportError:
+    _hashopenssl = None
+    _openssl_md_meths = None
+else:
+    _openssl_md_meths = frozenset(_hashopenssl.openssl_md_meth_names)
 import hashlib as _hashlib
 
 trans_5C = bytes((x ^ 0x5C) for x in range(256))
@@ -142,3 +149,38 @@ def new(key, msg = None, digestmod = None):
     method.
     """
     return HMAC(key, msg, digestmod)
+
+
+def digest(key, msg, digest):
+    """Fast inline implementation of HMAC
+
+    key:    key for the keyed hash object.
+    msg:    input message
+    digest: A hash name suitable for hashlib.new() for best performance. *OR*
+            A hashlib constructor returning a new hash object. *OR*
+            A module supporting PEP 247.
+
+    Note: key and msg must be a bytes or bytearray objects.
+    """
+    if (_hashopenssl is not None and
+            isinstance(digest, str) and digest in _openssl_md_meths):
+        return _hashopenssl.hmac_digest(key, msg, digest)
+
+    if callable(digest):
+        digest_cons = digest
+    elif isinstance(digest, str):
+        digest_cons = lambda d=b'': _hashlib.new(digest, d)
+    else:
+        digest_cons = lambda d=b'': digest.new(d)
+
+    inner = digest_cons()
+    outer = digest_cons()
+    blocksize = getattr(inner, 'block_size', 64)
+    if len(key) > blocksize:
+        key = digest_cons(key).digest()
+    key = key + b'\x00' * (blocksize - len(key))
+    inner.update(key.translate(trans_36))
+    outer.update(key.translate(trans_5C))
+    inner.update(msg)
+    outer.update(inner.digest())
+    return outer.digest()
diff --git a/Lib/test/test_hmac.py b/Lib/test/test_hmac.py
index 067e13f..4e4ef0e 100644
--- a/Lib/test/test_hmac.py
+++ b/Lib/test/test_hmac.py
@@ -1,7 +1,9 @@
+import binascii
 import functools
 import hmac
 import hashlib
 import unittest
+import unittest.mock
 import warnings
 
 
@@ -23,16 +25,27 @@ class TestVectorsTestCase(unittest.TestCase):
         def md5test(key, data, digest):
             h = hmac.HMAC(key, data, digestmod=hashlib.md5)
             self.assertEqual(h.hexdigest().upper(), digest.upper())
+            self.assertEqual(h.digest(), binascii.unhexlify(digest))
             self.assertEqual(h.name, "hmac-md5")
             self.assertEqual(h.digest_size, 16)
             self.assertEqual(h.block_size, 64)
 
             h = hmac.HMAC(key, data, digestmod='md5')
             self.assertEqual(h.hexdigest().upper(), digest.upper())
+            self.assertEqual(h.digest(), binascii.unhexlify(digest))
             self.assertEqual(h.name, "hmac-md5")
             self.assertEqual(h.digest_size, 16)
             self.assertEqual(h.block_size, 64)
 
+            self.assertEqual(
+                hmac.digest(key, data, digest='md5'),
+                binascii.unhexlify(digest)
+            )
+            with unittest.mock.patch('hmac._openssl_md_meths', {}):
+                self.assertEqual(
+                    hmac.digest(key, data, digest='md5'),
+                    binascii.unhexlify(digest)
+                )
 
         md5test(b"\x0b" * 16,
                 b"Hi There",
@@ -67,16 +80,23 @@ class TestVectorsTestCase(unittest.TestCase):
         def shatest(key, data, digest):
             h = hmac.HMAC(key, data, digestmod=hashlib.sha1)
             self.assertEqual(h.hexdigest().upper(), digest.upper())
+            self.assertEqual(h.digest(), binascii.unhexlify(digest))
             self.assertEqual(h.name, "hmac-sha1")
             self.assertEqual(h.digest_size, 20)
             self.assertEqual(h.block_size, 64)
 
             h = hmac.HMAC(key, data, digestmod='sha1')
             self.assertEqual(h.hexdigest().upper(), digest.upper())
+            self.assertEqual(h.digest(), binascii.unhexlify(digest))
             self.assertEqual(h.name, "hmac-sha1")
             self.assertEqual(h.digest_size, 20)
             self.assertEqual(h.block_size, 64)
 
+            self.assertEqual(
+                hmac.digest(key, data, digest='sha1'),
+                binascii.unhexlify(digest)
+            )
+
 
         shatest(b"\x0b" * 20,
                 b"Hi There",
@@ -122,6 +142,24 @@ class TestVectorsTestCase(unittest.TestCase):
             self.assertEqual(h.digest_size, digest_size)
             self.assertEqual(h.block_size, block_size)
 
+            self.assertEqual(
+                hmac.digest(key, data, digest=hashfunc),
+                binascii.unhexlify(hexdigests[hashfunc])
+            )
+            self.assertEqual(
+                hmac.digest(key, data, digest=hash_name),
+                binascii.unhexlify(hexdigests[hashfunc])
+            )
+
+            with unittest.mock.patch('hmac._openssl_md_meths', {}):
+                self.assertEqual(
+                    hmac.digest(key, data, digest=hashfunc),
+                    binascii.unhexlify(hexdigests[hashfunc])
+                )
+                self.assertEqual(
+                    hmac.digest(key, data, digest=hash_name),
+                    binascii.unhexlify(hexdigests[hashfunc])
+                )
 
         # 4.2.  Test Case 1
         hmactest(key = b'\x0b'*20,