python · picnixz · Jun 23, 2025 · Jun 23, 2025 · Jun 24, 2025 · Jun 26, 2025
@@ -15,10 +15,11 @@
 #endif
 
 #include "Python.h"
+#include "pycore_moduleobject.h"    // _PyModule_GetState()
+#include "pycore_strhex.h"          // _Py_strhex()
+#include "pycore_typeobject.h"      // _PyType_GetModuleState()
+
 #include "hashlib.h"
-#include "pycore_strhex.h"       // _Py_strhex()
-#include "pycore_typeobject.h"
-#include "pycore_moduleobject.h"
 
 // QUICK CPU AUTODETECTION
 //
@@ -67,6 +68,7 @@
 
 // MODULE TYPE SLOTS
 
+static struct PyModuleDef blake2module_def;
 static PyType_Spec blake2b_type_spec;
 static PyType_Spec blake2s_type_spec;
 
@@ -78,23 +80,24 @@ typedef struct {
     PyTypeObject *blake2s_type;
     bool can_run_simd128;
     bool can_run_simd256;
-} Blake2State;
+} blake2module_state;
 
-static inline Blake2State *
-blake2_get_state(PyObject *module)
+static inline blake2module_state *
+get_blake2module_state(PyObject *module)
 {
     void *state = _PyModule_GetState(module);
     assert(state != NULL);
-    return (Blake2State *)state;
+    return (blake2module_state *)state;
 }
 
 #if defined(HACL_CAN_COMPILE_SIMD128) || defined(HACL_CAN_COMPILE_SIMD256)
-static inline Blake2State *
-blake2_get_state_from_type(PyTypeObject *module)
+static inline blake2module_state *
+get_blake2module_state_by_cls(PyTypeObject *cls)
 {
-    void *state = _PyType_GetModuleState(module);
+    _Py_hashlib_check_exported_type(cls, &blake2module_def);
+    void *state = _PyType_GetModuleState(cls);
     assert(state != NULL);
-    return (Blake2State *)state;
+    return (blake2module_state *)state;
 }
 #endif
 
@@ -105,7 +108,7 @@ static struct PyMethodDef blake2mod_functions[] = {
 static int
 _blake2_traverse(PyObject *module, visitproc visit, void *arg)
 {
-    Blake2State *state = blake2_get_state(module);
+    blake2module_state *state = get_blake2module_state(module);
     Py_VISIT(state->blake2b_type);
     Py_VISIT(state->blake2s_type);
     return 0;
@@ -114,7 +117,7 @@ _blake2_traverse(PyObject *module, visitproc visit, void *arg)
 static int
 _blake2_clear(PyObject *module)
 {
-    Blake2State *state = blake2_get_state(module);
+    blake2module_state *state = get_blake2module_state(module);
     Py_CLEAR(state->blake2b_type);
     Py_CLEAR(state->blake2s_type);
     return 0;
@@ -127,7 +130,7 @@ _blake2_free(void *module)
 }
 
 static void
-blake2module_init_cpu_features(Blake2State *state)
+blake2module_init_cpu_features(blake2module_state *state)
 {
     /* This must be kept in sync with hmacmodule_init_cpu_features()
      * in hmacmodule.c */
@@ -205,7 +208,7 @@ blake2module_init_cpu_features(Blake2State *state)
 static int
 blake2_exec(PyObject *m)
 {
-    Blake2State *st = blake2_get_state(m);
+    blake2module_state *st = get_blake2module_state(m);
     blake2module_init_cpu_features(st);
 
 #define ADD_INT(DICT, NAME, VALUE)                      \
@@ -285,11 +288,11 @@ static PyModuleDef_Slot _blake2_slots[] = {
     {0, NULL}
 };
 
-static struct PyModuleDef blake2_module = {
+static struct PyModuleDef blake2module_def = {
     .m_base = PyModuleDef_HEAD_INIT,
     .m_name = "_blake2",
     .m_doc = blake2mod__doc__,
-    .m_size = sizeof(Blake2State),
+    .m_size = sizeof(blake2module_state),
     .m_methods = blake2mod_functions,
     .m_slots = _blake2_slots,
     .m_traverse = _blake2_traverse,
@@ -300,7 +303,7 @@ static struct PyModuleDef blake2_module = {
 PyMODINIT_FUNC
 PyInit__blake2(void)
 {
-    return PyModuleDef_Init(&blake2_module);
+    return PyModuleDef_Init(&blake2module_def);
 }
 
 // IMPLEMENTATION OF METHODS
@@ -333,7 +336,7 @@ static inline blake2_impl
 type_to_impl(PyTypeObject *type)
 {
 #if defined(HACL_CAN_COMPILE_SIMD128) || defined(HACL_CAN_COMPILE_SIMD256)
-    Blake2State *st = blake2_get_state_from_type(type);
+    blake2module_state *st = get_blake2module_state_by_cls(type);
 #endif
     if (!strcmp(type->tp_name, blake2b_type_spec.name)) {
 #if HACL_CAN_COMPILE_SIMD256
@@ -385,6 +388,7 @@ class _blake2.blake2s "Blake2Object *" "&PyType_Type"
 static Blake2Object *
 new_Blake2Object(PyTypeObject *type)
 {
+    _Py_hashlib_check_exported_type(type, &blake2module_def);
     Blake2Object *self = PyObject_GC_New(Blake2Object, type);
     if (self == NULL) {
         return NULL;

@@ -1,6 +1,34 @@
 /* Common code for use by all hashlib related modules. */
 
-#include "pycore_lock.h"        // PyMutex
+#include "pycore_lock.h"            // PyMutex
+#include "pycore_moduleobject.h"    // _PyModule_GetDef()
+
+#ifndef NDEBUG
+/*
+ * Assert that a type cannot be subclassed and that
+ * its associated module definition matches 'moddef'.
+ *
+ * Use this helper to ensure that _PyType_GetModuleState() can be safely used.
+ */
+static inline void
+_Py_hashlib_check_exported_type(PyTypeObject *type, PyModuleDef *moddef)
+{
+    assert(type != NULL);
+    assert(moddef != NULL);
+    /* ensure that the type is a final heap type */
+    assert(PyType_Check(type));
+    assert(type->tp_flags & Py_TPFLAGS_HEAPTYPE);
+    assert(!(type->tp_flags & Py_TPFLAGS_BASETYPE));
+    /* ensure that the associated module definition matches 'moddef' */
+    PyHeapTypeObject *ht = (PyHeapTypeObject *)type;
+    assert(ht->ht_module != NULL);
+    PyModuleDef *ht_moddef = _PyModule_GetDef(ht->ht_module);
+    assert(ht_moddef != NULL);
+    assert(ht_moddef == moddef);
+}
+#else
+#define _Py_hashlib_check_exported_type(_TYPE, _MODDEF)
+#endif
 
 /*
  * Given a PyObject* obj, fill in the Py_buffer* viewp with the result

@@ -18,7 +18,9 @@
 
 #include "Python.h"
 #include "pycore_hashtable.h"
+#include "pycore_moduleobject.h"        // _PyModule_GetState()
 #include "pycore_strhex.h"              // _Py_strhex()
+#include "pycore_typeobject.h"          // _PyType_GetModuleState()
 
 /*
  * Taken from blake2module.c. In the future, detection of SIMD support
@@ -250,6 +252,8 @@ typedef struct py_hmac_hinfo {
 
 // --- HMAC module state ------------------------------------------------------
 
+static struct PyModuleDef hmacmodule_def;
+
 typedef struct hmacmodule_state {
     _Py_hashtable_t *hinfo_table;
     PyObject *unknown_hash_error;
@@ -265,15 +269,16 @@ typedef struct hmacmodule_state {
 static inline hmacmodule_state *
 get_hmacmodule_state(PyObject *module)
 {
-    void *state = PyModule_GetState(module);
+    void *state = _PyModule_GetState(module);
     assert(state != NULL);
     return (hmacmodule_state *)state;
 }
 
 static inline hmacmodule_state *
 get_hmacmodule_state_by_cls(PyTypeObject *cls)
 {
-    void *state = PyType_GetModuleState(cls);
+    _Py_hashlib_check_exported_type(cls, &hmacmodule_def);
+    void *state = _PyType_GetModuleState(cls);
     assert(state != NULL);
     return (hmacmodule_state *)state;
 }
@@ -301,13 +306,11 @@ typedef struct HMACObject {
 
 /*[clinic input]
 module _hmac
-class _hmac.HMAC "HMACObject *" "clinic_state()->hmac_type"
+class _hmac.HMAC "HMACObject *" "&PyType_Type"
 [clinic start generated code]*/
-/*[clinic end generated code: output=da39a3ee5e6b4b0d input=c8bab73fde49ba8a]*/
+/*[clinic end generated code: output=da39a3ee5e6b4b0d input=72bc06d6dc634770]*/
 
-#define clinic_state()  (get_hmacmodule_state_by_cls(Py_TYPE(self)))
 #include "clinic/hmacmodule.c.h"
-#undef clinic_state
 
 // --- Helpers ----------------------------------------------------------------
 //
@@ -1683,7 +1686,7 @@ static struct PyModuleDef_Slot hmacmodule_slots[] = {
     {0, NULL} /* sentinel */
 };
 
-static struct PyModuleDef _hmacmodule = {
+static struct PyModuleDef hmacmodule_def = {
     PyModuleDef_HEAD_INIT,
     .m_name = "_hmac",
     .m_size = sizeof(hmacmodule_state),
@@ -1697,5 +1700,5 @@ static struct PyModuleDef _hmacmodule = {
 PyMODINIT_FUNC
 PyInit__hmac(void)
 {
-    return PyModuleDef_Init(&_hmacmodule);
+    return PyModuleDef_Init(&hmacmodule_def);
 }
@@ -22,7 +22,9 @@
 #endif
 
 #include "Python.h"
-#include "pycore_strhex.h" // _Py_strhex()
+#include "pycore_moduleobject.h"    // _PyModule_GetState()
+#include "pycore_strhex.h"          // _Py_strhex()
+#include "pycore_typeobject.h"      // _PyType_GetModuleState()
 
 #include "hashlib.h"
 
@@ -44,16 +46,27 @@ typedef struct {
 
 // --- Module state -----------------------------------------------------------
 
+static struct PyModuleDef md5module_def;
+
 typedef struct {
-    PyTypeObject* md5_type;
-} MD5State;
+    PyTypeObject *md5_type;
+} md5module_state;
+
+static inline md5module_state *
+get_md5module_state(PyObject *module)
+{
+    void *state = _PyModule_GetState(module);
+    assert(state != NULL);
+    return (md5module_state *)state;
+}
 
-static inline MD5State*
-md5_get_state(PyObject *module)
+static inline md5module_state *
+get_md5module_state_by_cls(PyTypeObject *cls)
 {
-    void *state = PyModule_GetState(module);
+    _Py_hashlib_check_exported_type(cls, &md5module_def);
+    void *state = _PyType_GetModuleState(cls);
     assert(state != NULL);
-    return (MD5State *)state;
+    return (md5module_state *)state;
 }
 
 // --- Module clinic configuration --------------------------------------------
@@ -69,7 +82,7 @@ class MD5Type "MD5object *" "&PyType_Type"
 // --- MD5 object interface ---------------------------------------------------
 
 static MD5object *
-newMD5object(MD5State * st)
+newMD5object(md5module_state *st)
 {
     MD5object *md5 = PyObject_GC_New(MD5object, st->md5_type);
     if (!md5) {
@@ -115,7 +128,7 @@ static PyObject *
 MD5Type_copy_impl(MD5object *self, PyTypeObject *cls)
 /*[clinic end generated code: output=bf055e08244bf5ee input=d89087dcfb2a8620]*/
 {
-    MD5State *st = PyType_GetModuleState(cls);
+    md5module_state *st = get_md5module_state_by_cls(cls);
 
     MD5object *newobj;
     if ((newobj = newMD5object(st)) == NULL) {
@@ -288,7 +301,7 @@ _md5_md5_impl(PyObject *module, PyObject *data, int usedforsecurity,
         GET_BUFFER_VIEW_OR_ERROUT(string, &buf);
     }
 
-    MD5State *st = md5_get_state(module);
+    md5module_state *st = get_md5module_state(module);
     if ((new = newMD5object(st)) == NULL) {
         if (string) {
             PyBuffer_Release(&buf);
@@ -329,15 +342,15 @@ static struct PyMethodDef MD5_functions[] = {
 static int
 _md5_traverse(PyObject *module, visitproc visit, void *arg)
 {
-    MD5State *state = md5_get_state(module);
+    md5module_state *state = get_md5module_state(module);
     Py_VISIT(state->md5_type);
     return 0;
 }
 
 static int
 _md5_clear(PyObject *module)
 {
-    MD5State *state = md5_get_state(module);
+    md5module_state *state = get_md5module_state(module);
     Py_CLEAR(state->md5_type);
     return 0;
 }
@@ -352,7 +365,7 @@ _md5_free(void *module)
 static int
 md5_exec(PyObject *m)
 {
-    MD5State *st = md5_get_state(m);
+    md5module_state *st = get_md5module_state(m);
 
     st->md5_type = (PyTypeObject *)PyType_FromModuleAndSpec(
         m, &md5_type_spec, NULL);
@@ -375,10 +388,10 @@ static PyModuleDef_Slot _md5_slots[] = {
 };
 
 
-static struct PyModuleDef _md5module = {
+static struct PyModuleDef md5module_def = {
     PyModuleDef_HEAD_INIT,
     .m_name = "_md5",
-    .m_size = sizeof(MD5State),
+    .m_size = sizeof(md5module_state),
     .m_methods = MD5_functions,
     .m_slots = _md5_slots,
     .m_traverse = _md5_traverse,
@@ -389,5 +402,5 @@ static struct PyModuleDef _md5module = {
 PyMODINIT_FUNC
 PyInit__md5(void)
 {
-    return PyModuleDef_Init(&_md5module);
+    return PyModuleDef_Init(&md5module_def);
 }