optimize calls to PyMem_Malloc in SHAKE digest computation

Author: picnixz

Lib/test/test_hashlib.py

@@ -376,7 +376,15 @@ def test_hexdigest(self):
                 self.assertIsInstance(h.digest(), bytes)
                 self.assertEqual(hexstr(h.digest()), h.hexdigest())
 
-    def test_digest_length_overflow(self):
+    def test_shakes_digest_zero_length(self):
+        for shake_name in self.shakes:
+            for constructor in self.hash_constructors[shake_name]:
+                with self.subTest(constructor=constructor):
+                    h = constructor(b'abcdef', usedforsecurity=False)
+                    self.assertEqual(h.digest(0), b'')
+                    self.assertEqual(h.hexdigest(0), '')
+
+    def test_shakes_digest_length_overflow(self):
         # See issue #34922
         large_sizes = (2**29, 2**32-10, 2**32+10, 2**61, 2**64-10, 2**64+10)
         for cons in self.hash_constructors:

Modules/clinic/sha3module.c.h

@@ -472,29 +472,48 @@ SHA3_TYPE_SPEC(sha3_384_spec, "sha3_384", sha3_384_slots);
 SHA3_TYPE_SLOTS(sha3_512_slots, sha3_512__doc__, SHA3_methods, SHA3_getseters);
 SHA3_TYPE_SPEC(sha3_512_spec, "sha3_512", sha3_512_slots);
 
+/*
+ * Compute the digest of SHAKE object.
+ *
+ * If 'hex' is false, the output is the raw binary digest as bytes.
+ * Otherwise, the output is its hexadecimal (string) representation.
+ */
 static PyObject *
-_SHAKE_digest(PyObject *op, unsigned long digestlen, int hex)
+_sha3_shake_compute_digest(SHA3object *self, uint32_t digestlen, bool hex)
 {
-    unsigned char *digest = NULL;
+    uint8_t *digest = NULL;
     PyObject *result = NULL;
-    SHA3object *self = _SHA3object_CAST(op);
-
-    if (digestlen >= (1 << 29)) {
-        PyErr_SetString(PyExc_ValueError, "length is too large");
+#define MAX_SHAKE_DIGESTSIZE (1 << 29)
+    if (digestlen >= MAX_SHAKE_DIGESTSIZE) {
+        PyErr_Format(PyExc_ValueError,
+                     "SHAKE digest length must be < %" PRIu32,
+                     MAX_SHAKE_DIGESTSIZE);
         return NULL;
     }
-    digest = (unsigned char*)PyMem_Malloc(digestlen);
+#undef MAX_SHAKE_DIGESTSIZE
+
+    /*
+     * The HACL* function fails if the algorithm is not SHAKE, which is not
+     * the case as the caller is _sha3_shake_{128,256}[hex]digest_impl(),
+     * or if the output length is zero.
+     *
+     * In the latter case, we follow the existing behavior
+     * and return an empty digest, without raising an error.
+     */
+    if (digestlen == 0) {
+        return Py_GetConstant(
+            hex ? Py_CONSTANT_EMPTY_STR : Py_CONSTANT_EMPTY_BYTES
+        );
+    }
+
+    digest = PyMem_Malloc(digestlen);
     if (digest == NULL) {
         return PyErr_NoMemory();
     }
 
-    /* Get the raw (binary) digest value. The HACL functions errors out if:
-     * - the algorithm is not shake -- not the case here
-     * - the output length is zero -- we follow the existing behavior and return
-     *   an empty digest, without raising an error */
-    if (digestlen > 0) {
-        (void)Hacl_Hash_SHA3_squeeze(self->hash_state, digest, digestlen);
-    }
+    ENTER_HASHLIB(self);
+    (void)Hacl_Hash_SHA3_squeeze(self->hash_state, digest, digestlen);
+    LEAVE_HASHLIB(self);
     if (hex) {
         result = _Py_strhex((const char *)digest, digestlen);
     }
@@ -509,32 +528,32 @@ _SHAKE_digest(PyObject *op, unsigned long digestlen, int hex)
 /*[clinic input]
 _sha3.shake_128.digest
 
-    length: unsigned_long
+    length: uint32
 
 Return the digest value as a bytes object.
 [clinic start generated code]*/
 
 static PyObject *
-_sha3_shake_128_digest_impl(SHA3object *self, unsigned long length)
-/*[clinic end generated code: output=2313605e2f87bb8f input=93d6d6ff32904f18]*/
+_sha3_shake_128_digest_impl(SHA3object *self, uint32_t length)
+/*[clinic end generated code: output=b8be6cd55400959c input=8a6901d3d699ee26]*/
 {
-    return _SHAKE_digest((PyObject *)self, length, 0);
+    return _sha3_shake_compute_digest(self, length, false);
 }
 
 
 /*[clinic input]
 _sha3.shake_128.hexdigest
 
-    length: unsigned_long
+    length: uint32
 
 Return the digest value as a string of hexadecimal digits.
 [clinic start generated code]*/
 
 static PyObject *
-_sha3_shake_128_hexdigest_impl(SHA3object *self, unsigned long length)
-/*[clinic end generated code: output=bf8e2f1e490944a8 input=562d74e7060b56ab]*/
+_sha3_shake_128_hexdigest_impl(SHA3object *self, uint32_t length)
+/*[clinic end generated code: output=cc1abb908d8ac614 input=acab18f97f8af822]*/
 {
-    return _SHAKE_digest((PyObject *)self, length, 1);
+    return _sha3_shake_compute_digest(self, length, true);
 }
 
 static PyObject *