gh-87138: convert blake2b/2s types to heap types

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
 };