From fb0b94223d481ca58b7c77332348d1ab2c9ab272 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?B=C3=A9n=C3=A9dikt=20Tran?=
 <10796600+picnixz@users.noreply.github.com>
Date: Thu, 26 Dec 2024 16:03:57 +0100
Subject: [PATCH] gh-87138: convert blake2b/2s types to heap types (#127669)

---
 Modules/blake2module.c | 89 ++++++++++++++++++++++++++++++++----------
 1 file changed, 69 insertions(+), 20 deletions(-)

diff --git a/Modules/blake2module.c b/Modules/blake2module.c
index 94cdfe7fd2e962..6723e7de4675a5 100644
--- a/Modules/blake2module.c
+++ b/Modules/blake2module.c
@@ -379,13 +379,13 @@ class _blake2.blake2s "Blake2Object *" "&PyBlake2_BLAKE2sType"
 static Blake2Object *
 new_Blake2Object(PyTypeObject *type)
 {
-    Blake2Object *self;
-    self = (Blake2Object *)type->tp_alloc(type, 0);
+    Blake2Object *self = PyObject_GC_New(Blake2Object, type);
     if (self == NULL) {
         return NULL;
     }
     HASHLIB_INIT_MUTEX(self);
 
+    PyObject_GC_Track(self);
     return self;
 }
 
@@ -454,7 +454,28 @@ py_blake2b_or_s_new(PyTypeObject *type, PyObject *data, int digest_size,
     }
 
     self->impl = type_to_impl(type);
-
+    // Ensure that the states are NULL-initialized in case of an error.
+    // See: py_blake2_clear() for more details.
+    switch (self->impl) {
+#if HACL_CAN_COMPILE_SIMD256
+        case Blake2b_256:
+            self->blake2b_256_state = NULL;
+            break;
+#endif
+#if HACL_CAN_COMPILE_SIMD128
+        case Blake2s_128:
+            self->blake2s_128_state = NULL;
+            break;
+#endif
+        case Blake2b:
+            self->blake2b_state = NULL;
+            break;
+        case Blake2s:
+            self->blake2s_state = NULL;
+            break;
+        default:
+            Py_UNREACHABLE();
+    }
     // Using Blake2b because we statically know that these are greater than the
     // Blake2s sizes -- this avoids a VLA.
     uint8_t salt_[HACL_HASH_BLAKE2B_SALT_BYTES] = { 0 };
@@ -595,7 +616,7 @@ py_blake2b_or_s_new(PyTypeObject *type, PyObject *data, int digest_size,
 
     return (PyObject *)self;
 error:
-	Py_XDECREF(self);
+    Py_XDECREF(self);
     return NULL;
 }
 
@@ -875,46 +896,70 @@ static PyGetSetDef py_blake2b_getsetters[] = {
     {NULL}
 };
 
-
-static void
-py_blake2b_dealloc(Blake2Object *self)
+static int
+py_blake2_clear(PyObject *op)
 {
+    Blake2Object *self = (Blake2Object *)op;
+    // The initialization function uses PyObject_GC_New() but explicitly
+    // initializes the HACL* internal state to NULL before allocating
+    // it. If an error occurs in the constructor, we should only free
+    // states that were allocated (i.e. that are not NULL).
     switch (self->impl) {
 #if HACL_CAN_COMPILE_SIMD256
         case Blake2b_256:
-            if (self->blake2b_256_state != NULL)
+            if (self->blake2b_256_state != NULL) {
                 Hacl_Hash_Blake2b_Simd256_free(self->blake2b_256_state);
+                self->blake2b_256_state = NULL;
+            }
             break;
 #endif
 #if HACL_CAN_COMPILE_SIMD128
         case Blake2s_128:
-            if (self->blake2s_128_state != NULL)
+            if (self->blake2s_128_state != NULL) {
                 Hacl_Hash_Blake2s_Simd128_free(self->blake2s_128_state);
+                self->blake2s_128_state = NULL;
+            }
             break;
 #endif
         case Blake2b:
-            // This happens if we hit "goto error" in the middle of the
-            // initialization function. We leverage the fact that tp_alloc
-            // guarantees that the contents of the object are NULL-initialized
-            // (see documentation for PyType_GenericAlloc) to detect this case.
-            if (self->blake2b_state != NULL)
+            if (self->blake2b_state != NULL) {
                 Hacl_Hash_Blake2b_free(self->blake2b_state);
+                self->blake2b_state = NULL;
+            }
             break;
         case Blake2s:
-            if (self->blake2s_state != NULL)
+            if (self->blake2s_state != NULL) {
                 Hacl_Hash_Blake2s_free(self->blake2s_state);
+                self->blake2s_state = NULL;
+            }
             break;
         default:
             Py_UNREACHABLE();
     }
+    return 0;
+}
 
+static void
+py_blake2_dealloc(PyObject *self)
+{
     PyTypeObject *type = Py_TYPE(self);
-    PyObject_Free(self);
+    PyObject_GC_UnTrack(self);
+    (void)py_blake2_clear(self);
+    type->tp_free(self);
     Py_DECREF(type);
 }
 
+static int
+py_blake2_traverse(PyObject *self, visitproc visit, void *arg)
+{
+    Py_VISIT(Py_TYPE(self));
+    return 0;
+}
+
 static PyType_Slot blake2b_type_slots[] = {
-    {Py_tp_dealloc, py_blake2b_dealloc},
+    {Py_tp_clear, py_blake2_clear},
+    {Py_tp_dealloc, py_blake2_dealloc},
+    {Py_tp_traverse, py_blake2_traverse},
     {Py_tp_doc, (char *)py_blake2b_new__doc__},
     {Py_tp_methods, py_blake2b_methods},
     {Py_tp_getset, py_blake2b_getsetters},
@@ -923,7 +968,9 @@ static PyType_Slot blake2b_type_slots[] = {
 };
 
 static PyType_Slot blake2s_type_slots[] = {
-    {Py_tp_dealloc, py_blake2b_dealloc},
+    {Py_tp_clear, py_blake2_clear},
+    {Py_tp_dealloc, py_blake2_dealloc},
+    {Py_tp_traverse, py_blake2_traverse},
     {Py_tp_doc, (char *)py_blake2s_new__doc__},
     {Py_tp_methods, py_blake2b_methods},
     {Py_tp_getset, py_blake2b_getsetters},
@@ -936,13 +983,15 @@ static PyType_Slot blake2s_type_slots[] = {
 static PyType_Spec blake2b_type_spec = {
     .name = "_blake2.blake2b",
     .basicsize =  sizeof(Blake2Object),
-    .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE,
+    .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE
+             | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_HEAPTYPE,
     .slots = blake2b_type_slots
 };
 
 static PyType_Spec blake2s_type_spec = {
     .name = "_blake2.blake2s",
     .basicsize =  sizeof(Blake2Object),
-    .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE,
+    .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE
+             | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_HEAPTYPE,
     .slots = blake2s_type_slots
 };