From cd2c97facc97dc53020fba1f84d0c001bdaf528b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?B=C3=A9n=C3=A9dikt=20Tran?=
 <10796600+picnixz@users.noreply.github.com>
Date: Wed, 4 Dec 2024 16:53:19 +0100
Subject: [PATCH 1/4] remove redundant casts

---
 Modules/md5module.c  |  2 +-
 Modules/sha1module.c |  2 +-
 Modules/sha2module.c | 12 ++++--------
 Modules/sha3module.c |  3 +--
 4 files changed, 7 insertions(+), 12 deletions(-)

diff --git a/Modules/md5module.c b/Modules/md5module.c
index ef9163e8be5b6c..3f64d342498f33 100644
--- a/Modules/md5module.c
+++ b/Modules/md5module.c
@@ -72,7 +72,7 @@ md5_get_state(PyObject *module)
 static MD5object *
 newMD5object(MD5State * st)
 {
-    MD5object *md5 = (MD5object *)PyObject_GC_New(MD5object, st->md5_type);
+    MD5object *md5 = PyObject_GC_New(MD5object, st->md5_type);
     if (!md5) {
         return NULL;
     }
diff --git a/Modules/sha1module.c b/Modules/sha1module.c
index 34a427a39b5cf8..42861ffded12eb 100644
--- a/Modules/sha1module.c
+++ b/Modules/sha1module.c
@@ -73,7 +73,7 @@ sha1_get_state(PyObject *module)
 static SHA1object *
 newSHA1object(SHA1State *st)
 {
-    SHA1object *sha = (SHA1object *)PyObject_GC_New(SHA1object, st->sha1_type);
+    SHA1object *sha = PyObject_GC_New(SHA1object, st->sha1_type);
     if (sha == NULL) {
         return NULL;
     }
diff --git a/Modules/sha2module.c b/Modules/sha2module.c
index 7d6a1e40243f9d..ee205fcade9d41 100644
--- a/Modules/sha2module.c
+++ b/Modules/sha2module.c
@@ -101,8 +101,7 @@ static void SHA512copy(SHA512object *src, SHA512object *dest)
 static SHA256object *
 newSHA224object(sha2_state *state)
 {
-    SHA256object *sha = (SHA256object *)PyObject_GC_New(
-        SHA256object, state->sha224_type);
+    SHA256object *sha = PyObject_GC_New(SHA256object, state->sha224_type);
     if (!sha) {
         return NULL;
     }
@@ -115,8 +114,7 @@ newSHA224object(sha2_state *state)
 static SHA256object *
 newSHA256object(sha2_state *state)
 {
-    SHA256object *sha = (SHA256object *)PyObject_GC_New(
-        SHA256object, state->sha256_type);
+    SHA256object *sha = PyObject_GC_New(SHA256object, state->sha256_type);
     if (!sha) {
         return NULL;
     }
@@ -129,8 +127,7 @@ newSHA256object(sha2_state *state)
 static SHA512object *
 newSHA384object(sha2_state *state)
 {
-    SHA512object *sha = (SHA512object *)PyObject_GC_New(
-        SHA512object, state->sha384_type);
+    SHA512object *sha = PyObject_GC_New(SHA512object, state->sha384_type);
     if (!sha) {
         return NULL;
     }
@@ -143,8 +140,7 @@ newSHA384object(sha2_state *state)
 static SHA512object *
 newSHA512object(sha2_state *state)
 {
-    SHA512object *sha = (SHA512object *)PyObject_GC_New(
-        SHA512object, state->sha512_type);
+    SHA512object *sha = PyObject_GC_New(SHA512object, state->sha512_type);
     if (!sha) {
         return NULL;
     }
diff --git a/Modules/sha3module.c b/Modules/sha3module.c
index ca839dc55e0519..ce1427662348fa 100644
--- a/Modules/sha3module.c
+++ b/Modules/sha3module.c
@@ -71,8 +71,7 @@ typedef struct {
 static SHA3object *
 newSHA3object(PyTypeObject *type)
 {
-    SHA3object *newobj;
-    newobj = (SHA3object *)PyObject_New(SHA3object, type);
+    SHA3object *newobj = PyObject_New(SHA3object, type);
     if (newobj == NULL) {
         return NULL;
     }

From a8975ed3aaa07ccd5a2ed2ea09bbc80e330937af Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?B=C3=A9n=C3=A9dikt=20Tran?=
 <10796600+picnixz@users.noreply.github.com>
Date: Wed, 4 Dec 2024 17:09:32 +0100
Subject: [PATCH 2/4] convert SHA3 and SHAKE classes to heap types

---
 Modules/sha3module.c | 40 ++++++++++++++++++++++++++++++----------
 1 file changed, 30 insertions(+), 10 deletions(-)

diff --git a/Modules/sha3module.c b/Modules/sha3module.c
index ce1427662348fa..93296c2761cc6e 100644
--- a/Modules/sha3module.c
+++ b/Modules/sha3module.c
@@ -71,12 +71,13 @@ typedef struct {
 static SHA3object *
 newSHA3object(PyTypeObject *type)
 {
-    SHA3object *newobj = PyObject_New(SHA3object, type);
+    SHA3object *newobj = PyObject_GC_New(SHA3object, type);
     if (newobj == NULL) {
         return NULL;
     }
     HASHLIB_INIT_MUTEX(newobj);
 
+    PyObject_GC_Track(newobj);
     return newobj;
 }
 
@@ -165,15 +166,32 @@ py_sha3_new_impl(PyTypeObject *type, PyObject *data, int usedforsecurity)
 
 /* Internal methods for a hash object */
 
+static int
+SHA3_clear(SHA3object *self)
+{
+    if (self->hash_state != NULL) {
+        Hacl_Hash_SHA3_free(self->hash_state);
+        self->hash_state = NULL;
+    }
+    return 0;
+}
+
 static void
 SHA3_dealloc(SHA3object *self)
 {
-    Hacl_Hash_SHA3_free(self->hash_state);
     PyTypeObject *tp = Py_TYPE(self);
-    PyObject_Free(self);
+    PyObject_GC_UnTrack(self);
+    (void)SHA3_clear(self);
+    tp->tp_free(self);
     Py_DECREF(tp);
 }
 
+static int
+SHA3_traverse(PyObject *self, visitproc visit, void *arg)
+{
+    Py_VISIT(Py_TYPE(self));
+    return 0;
+}
 
 /* External methods for a hash object */
 
@@ -334,6 +352,7 @@ static PyObject *
 SHA3_get_capacity_bits(SHA3object *self, void *closure)
 {
     uint32_t rate = Hacl_Hash_SHA3_block_len(self->hash_state) * 8;
+    assert(rate <= 1600);
     int capacity = 1600 - rate;
     return PyLong_FromLong(capacity);
 }
@@ -360,17 +379,19 @@ static PyGetSetDef SHA3_getseters[] = {
     {"_capacity_bits", (getter)SHA3_get_capacity_bits, NULL, NULL, NULL},
     {"_rate_bits", (getter)SHA3_get_rate_bits, NULL, NULL, NULL},
     {"_suffix", (getter)SHA3_get_suffix, NULL, NULL, NULL},
-    {NULL}  /* Sentinel */
+    {NULL, NULL, NULL, NULL, NULL}  /* Sentinel */
 };
 
 #define SHA3_TYPE_SLOTS(type_slots_obj, type_doc, type_methods, type_getseters) \
     static PyType_Slot type_slots_obj[] = { \
+        {Py_tp_clear, SHA3_clear}, \
         {Py_tp_dealloc, SHA3_dealloc}, \
+        {Py_tp_traverse, SHA3_traverse}, \
         {Py_tp_doc, (char*)type_doc}, \
         {Py_tp_methods, type_methods}, \
         {Py_tp_getset, type_getseters}, \
         {Py_tp_new, py_sha3_new}, \
-        {0,0} \
+        {0, NULL} \
     }
 
 // Using _PyType_GetModuleState() on these types is safe since they
@@ -379,7 +400,8 @@ static PyGetSetDef SHA3_getseters[] = {
     static PyType_Spec type_spec_obj = { \
         .name = "_sha3." type_name, \
         .basicsize = sizeof(SHA3object), \
-        .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE, \
+        .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE \
+                 | Py_TPFLAGS_HAVE_GC, \
         .slots = type_slots \
     }
 
@@ -443,9 +465,7 @@ _SHAKE_digest(SHA3object *self, unsigned long digestlen, int hex)
         result = PyBytes_FromStringAndSize((const char *)digest,
                                            digestlen);
     }
-    if (digest != NULL) {
-        PyMem_Free(digest);
-    }
+    PyMem_Free(digest);
     return result;
 }
 
@@ -562,7 +582,7 @@ _sha3_clear(PyObject *module)
 static void
 _sha3_free(void *module)
 {
-    _sha3_clear((PyObject *)module);
+    (void)_sha3_clear((PyObject *)module);
 }
 
 static int

From 94a6460d74534965207d9d4091d619e493258367 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?B=C3=A9n=C3=A9dikt=20Tran?=
 <10796600+picnixz@users.noreply.github.com>
Date: Sun, 8 Dec 2024 14:13:23 +0100
Subject: [PATCH 3/4] Revert "remove redundant casts"

This reverts commit cd2c97facc97dc53020fba1f84d0c001bdaf528b.
---
 Modules/md5module.c  |  2 +-
 Modules/sha1module.c |  2 +-
 Modules/sha2module.c | 12 ++++++++----
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/Modules/md5module.c b/Modules/md5module.c
index 3f64d342498f33..ef9163e8be5b6c 100644
--- a/Modules/md5module.c
+++ b/Modules/md5module.c
@@ -72,7 +72,7 @@ md5_get_state(PyObject *module)
 static MD5object *
 newMD5object(MD5State * st)
 {
-    MD5object *md5 = PyObject_GC_New(MD5object, st->md5_type);
+    MD5object *md5 = (MD5object *)PyObject_GC_New(MD5object, st->md5_type);
     if (!md5) {
         return NULL;
     }
diff --git a/Modules/sha1module.c b/Modules/sha1module.c
index 42861ffded12eb..34a427a39b5cf8 100644
--- a/Modules/sha1module.c
+++ b/Modules/sha1module.c
@@ -73,7 +73,7 @@ sha1_get_state(PyObject *module)
 static SHA1object *
 newSHA1object(SHA1State *st)
 {
-    SHA1object *sha = PyObject_GC_New(SHA1object, st->sha1_type);
+    SHA1object *sha = (SHA1object *)PyObject_GC_New(SHA1object, st->sha1_type);
     if (sha == NULL) {
         return NULL;
     }
diff --git a/Modules/sha2module.c b/Modules/sha2module.c
index ee205fcade9d41..7d6a1e40243f9d 100644
--- a/Modules/sha2module.c
+++ b/Modules/sha2module.c
@@ -101,7 +101,8 @@ static void SHA512copy(SHA512object *src, SHA512object *dest)
 static SHA256object *
 newSHA224object(sha2_state *state)
 {
-    SHA256object *sha = PyObject_GC_New(SHA256object, state->sha224_type);
+    SHA256object *sha = (SHA256object *)PyObject_GC_New(
+        SHA256object, state->sha224_type);
     if (!sha) {
         return NULL;
     }
@@ -114,7 +115,8 @@ newSHA224object(sha2_state *state)
 static SHA256object *
 newSHA256object(sha2_state *state)
 {
-    SHA256object *sha = PyObject_GC_New(SHA256object, state->sha256_type);
+    SHA256object *sha = (SHA256object *)PyObject_GC_New(
+        SHA256object, state->sha256_type);
     if (!sha) {
         return NULL;
     }
@@ -127,7 +129,8 @@ newSHA256object(sha2_state *state)
 static SHA512object *
 newSHA384object(sha2_state *state)
 {
-    SHA512object *sha = PyObject_GC_New(SHA512object, state->sha384_type);
+    SHA512object *sha = (SHA512object *)PyObject_GC_New(
+        SHA512object, state->sha384_type);
     if (!sha) {
         return NULL;
     }
@@ -140,7 +143,8 @@ newSHA384object(sha2_state *state)
 static SHA512object *
 newSHA512object(sha2_state *state)
 {
-    SHA512object *sha = PyObject_GC_New(SHA512object, state->sha512_type);
+    SHA512object *sha = (SHA512object *)PyObject_GC_New(
+        SHA512object, state->sha512_type);
     if (!sha) {
         return NULL;
     }

From e23eeeeea99fc4e777508183b3270da4d89b3b35 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?B=C3=A9n=C3=A9dikt=20Tran?=
 <10796600+picnixz@users.noreply.github.com>
Date: Sun, 8 Dec 2024 14:17:32 +0100
Subject: [PATCH 4/4] remove cosmetics

---
 Modules/sha3module.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Modules/sha3module.c b/Modules/sha3module.c
index 93296c2761cc6e..b13e6a9de10114 100644
--- a/Modules/sha3module.c
+++ b/Modules/sha3module.c
@@ -379,7 +379,7 @@ static PyGetSetDef SHA3_getseters[] = {
     {"_capacity_bits", (getter)SHA3_get_capacity_bits, NULL, NULL, NULL},
     {"_rate_bits", (getter)SHA3_get_rate_bits, NULL, NULL, NULL},
     {"_suffix", (getter)SHA3_get_suffix, NULL, NULL, NULL},
-    {NULL, NULL, NULL, NULL, NULL}  /* Sentinel */
+    {NULL}  /* Sentinel */
 };
 
 #define SHA3_TYPE_SLOTS(type_slots_obj, type_doc, type_methods, type_getseters) \