optimize calls to PyMem_Malloc in SHAKE digest computation

Author: picnixz

Lib/test/test_hashlib.py

@@ -207,6 +207,11 @@ def hash_constructors(self):
         constructors = self.constructors_to_test.values()
         return itertools.chain.from_iterable(constructors)
 
+    @property
+    def shake_constructors(self):
+        for shake_name in self.shakes:
+            yield from self.constructors_to_test.get(shake_name, ())
+
     @property
     def is_fips_mode(self):
         return get_fips_mode()
@@ -376,6 +381,13 @@ def test_hexdigest(self):
                 self.assertIsInstance(h.digest(), bytes)
                 self.assertEqual(hexstr(h.digest()), h.hexdigest())
 
+    def test_shakes_digest_zero_length(self):
+        for constructor in self.shake_constructors:
+            with self.subTest(constructor=constructor):
+                h = constructor(b'abcdef', usedforsecurity=False)
+                self.assertEqual(h.digest(0), b'')
+                self.assertEqual(h.hexdigest(0), '')
+
     def test_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)

Misc/NEWS.d/next/Library/2025-06-20-15-10-13.gh-issue-135532.k6r0vA.rst

@@ -0,0 +1,3 @@
+:mod:`hashlib`: ensure that exceptions raised by SHAKE objects when a digest
+length exceeding ``2 ** 29`` is given are consistent across implementations.
+Patch by Bénédikt Tran.

Modules/_hashopenssl.c

@@ -938,8 +938,13 @@ _hashlib_HASHXOF_digest_impl(HASHobject *self, Py_ssize_t length)
 /*[clinic end generated code: output=dcb09335dd2fe908 input=3eb034ce03c55b21]*/
 {
     EVP_MD_CTX *temp_ctx;
-    PyObject *retval = PyBytes_FromStringAndSize(NULL, length);
+    PyObject *retval;
+
+    if (length == 0) {
+        return Py_GetConstant(Py_CONSTANT_EMPTY_BYTES);
+    }
 
+    retval =  PyBytes_FromStringAndSize(NULL, length);
     if (retval == NULL) {
         return NULL;
     }
@@ -986,6 +991,10 @@ _hashlib_HASHXOF_hexdigest_impl(HASHobject *self, Py_ssize_t length)
     EVP_MD_CTX *temp_ctx;
     PyObject *retval;
 
+    if (length == 0) {
+        return Py_GetConstant(Py_CONSTANT_EMPTY_STR);
+    }
+
     digest = (unsigned char*)PyMem_Malloc(length);
     if (digest == NULL) {
         PyErr_NoMemory();

Modules/sha3module.c

@@ -472,40 +472,6 @@ 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);
 
-static PyObject *
-_SHAKE_digest(PyObject *op, unsigned long digestlen, int hex)
-{
-    unsigned char *digest = NULL;
-    PyObject *result = NULL;
-    SHA3object *self = _SHA3object_CAST(op);
-
-    if (digestlen >= (1 << 29)) {
-        PyErr_SetString(PyExc_ValueError, "length is too large");
-        return NULL;
-    }
-    digest = (unsigned char*)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);
-    }
-    if (hex) {
-        result = _Py_strhex((const char *)digest, digestlen);
-    }
-    else {
-        result = PyBytes_FromStringAndSize((const char *)digest, digestlen);
-    }
-    PyMem_Free(digest);
-    return result;
-}
-
-
 /*[clinic input]
 _sha3.shake_128.digest
 
@@ -518,7 +484,21 @@ static PyObject *
 _sha3_shake_128_digest_impl(SHA3object *self, unsigned long length)
 /*[clinic end generated code: output=2313605e2f87bb8f input=93d6d6ff32904f18]*/
 {
-    return _SHAKE_digest((PyObject *)self, length, 0);
+    /*
+     * Hacl_Hash_SHA3_squeeze() fails if the algorithm is not SHAKE,
+     * or if the length is 0. In the latter case, we follow OpenSSL's
+     * behavior and return an empty digest, without raising an error.
+     */
+    if (length == 0) {
+        return Py_GetConstant(Py_CONSTANT_EMPTY_BYTES);
+    }
+
+    PyObject *digest = PyBytes_FromStringAndSize(NULL, length);
+    uint8_t *buffer = (uint8_t *)PyBytes_AS_STRING(digest);
+    ENTER_HASHLIB(self);
+    (void)Hacl_Hash_SHA3_squeeze(self->hash_state, buffer, length);
+    LEAVE_HASHLIB(self);
+    return digest;
 }
 
 
@@ -534,7 +514,21 @@ static PyObject *
 _sha3_shake_128_hexdigest_impl(SHA3object *self, unsigned long length)
 /*[clinic end generated code: output=bf8e2f1e490944a8 input=562d74e7060b56ab]*/
 {
-    return _SHAKE_digest((PyObject *)self, length, 1);
+    /* See _sha3_shake_128_digest_impl() for the fast path rationale. */
+    if (length == 0) {
+        return Py_GetConstant(Py_CONSTANT_EMPTY_STR);
+    }
+
+    uint8_t *buffer = PyMem_Malloc(length);
+    if (buffer == NULL) {
+        return PyErr_NoMemory();
+    }
+    ENTER_HASHLIB(self);
+    (void)Hacl_Hash_SHA3_squeeze(self->hash_state, buffer, length);
+    LEAVE_HASHLIB(self);
+    PyObject *digest = _Py_strhex((const char *)buffer, length);
+    PyMem_Free(buffer);
+    return digest;
 }
 
 static PyObject *