From ec4cb3ee553e8be610c08a75407d7da71241bde6 Mon Sep 17 00:00:00 2001
From: Shane Harvey <shnhrv@gmail.com>
Date: Mon, 25 Mar 2024 13:48:45 -0700
Subject: [PATCH] PYTHON-4285 More consistent PyModule_GetState checks (#1560)

---
 bson/_cbsonmodule.c       | 24 ++++++++++++++++++++++++
 pymongo/_cmessagemodule.c | 24 ++++++++++++++++++++----
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/bson/_cbsonmodule.c b/bson/_cbsonmodule.c
index 4f6db2507..8d684def7 100644
--- a/bson/_cbsonmodule.c
+++ b/bson/_cbsonmodule.c
@@ -342,6 +342,9 @@ static long long millis_from_datetime(PyObject* datetime) {
 static PyObject* datetime_ms_from_millis(PyObject* self, long long millis){
     // Allocate a new DatetimeMS object.
     struct module_state *state = GETSTATE(self);
+    if (!state) {
+        return NULL;
+    }
 
     PyObject* dt;
     PyObject* ll_millis;
@@ -480,6 +483,9 @@ static int _load_python_objects(PyObject* module) {
     PyObject* re_compile = NULL;
     PyObject* compiled = NULL;
     struct module_state *state = GETSTATE(module);
+    if (!state) {
+        return 1;
+    }
 
     /* Cache commonly used attribute names to improve performance. */
     if (!((state->_type_marker_str = PyUnicode_FromString("_type_marker")) &&
@@ -638,6 +644,9 @@ int convert_codec_options(PyObject* self, PyObject* options_obj, codec_options_t
     PyObject* type_registry_obj = NULL;
     struct module_state *state = GETSTATE(self);
     long type_marker;
+    if (!state) {
+        return 0;
+    }
 
     options->unicode_decode_error_handler = NULL;
 
@@ -864,6 +873,9 @@ static int _write_element_to_buffer(PyObject* self, buffer_t buffer,
     if (type < 0) {
         return 0;
     }
+    if (!state) {
+        return 0;
+    }
 
     switch (type) {
     case 5:
@@ -1501,6 +1513,9 @@ int write_dict(PyObject* self, buffer_t buffer,
     struct module_state *state = GETSTATE(self);
     long type_marker;
     int is_dict = PyDict_Check(dict);
+    if (!state) {
+        return 0;
+    }
 
     if (!is_dict) {
         /* check for RawBSONDocument */
@@ -1638,6 +1653,9 @@ static PyObject* _cbson_dict_to_bson(PyObject* self, PyObject* args) {
     PyObject* raw_bson_document_bytes_obj;
     long type_marker;
     struct module_state *state = GETSTATE(self);
+    if (!state) {
+        return NULL;
+    }
 
     if (!(PyArg_ParseTuple(args, "ObO|b", &dict, &check_keys,
                           &options_obj, &top_level) &&
@@ -1689,6 +1707,9 @@ static PyObject *_dbref_hook(PyObject* self, PyObject* value) {
     PyObject* database = NULL;
     PyObject* ret = NULL;
     int db_present = 0;
+    if (!state) {
+        return NULL;
+    }
 
     /* Decoding for DBRefs */
     if (PyMapping_HasKey(value, state->_dollar_ref_str) && PyMapping_HasKey(value, state->_dollar_id_str)) { /* DBRef */
@@ -1743,6 +1764,9 @@ static PyObject* get_value(PyObject* self, PyObject* name, const char* buffer,
                            unsigned max, const codec_options_t* options, int raw_array) {
     struct module_state *state = GETSTATE(self);
     PyObject* value = NULL;
+    if (!state) {
+        return NULL;
+    }
     switch (type) {
     case 1:
         {
diff --git a/pymongo/_cmessagemodule.c b/pymongo/_cmessagemodule.c
index 3b3fde8ad..f95b94938 100644
--- a/pymongo/_cmessagemodule.c
+++ b/pymongo/_cmessagemodule.c
@@ -68,8 +68,6 @@ static int buffer_write_bytes_ssize_t(buffer_t buffer, const char* data, Py_ssiz
 
 static PyObject* _cbson_query_message(PyObject* self, PyObject* args) {
     /* NOTE just using a random number as the request_id */
-    struct module_state *state = GETSTATE(self);
-
     int request_id = rand();
     unsigned int flags;
     char* collection_name = NULL;
@@ -84,6 +82,10 @@ static PyObject* _cbson_query_message(PyObject* self, PyObject* args) {
     buffer_t buffer = NULL;
     int length_location, message_length;
     PyObject* result = NULL;
+    struct module_state *state = GETSTATE(self);
+    if (!state) {
+        return NULL;
+    }
 
     if (!(PyArg_ParseTuple(args, "Iet#iiOOO",
                           &flags,
@@ -216,8 +218,6 @@ fail:
  * only checked *after* generating the entire message.
  */
 static PyObject* _cbson_op_msg(PyObject* self, PyObject* args) {
-    struct module_state *state = GETSTATE(self);
-
     /* NOTE just using a random number as the request_id */
     int request_id = rand();
     unsigned int flags;
@@ -234,6 +234,10 @@ static PyObject* _cbson_op_msg(PyObject* self, PyObject* args) {
     int max_doc_size = 0;
     PyObject* result = NULL;
     PyObject* iterator = NULL;
+    struct module_state *state = GETSTATE(self);
+    if (!state) {
+        return NULL;
+    }
 
     /*flags, command, identifier, docs, opts*/
     if (!(PyArg_ParseTuple(args, "IOet#OO",
@@ -540,6 +544,9 @@ _cbson_encode_batched_op_msg(PyObject* self, PyObject* args) {
     codec_options_t options;
     buffer_t buffer;
     struct module_state *state = GETSTATE(self);
+    if (!state) {
+        return NULL;
+    }
 
     if (!(PyArg_ParseTuple(args, "bOObOO",
                           &op, &command, &docs, &ack,
@@ -594,6 +601,9 @@ _cbson_batched_op_msg(PyObject* self, PyObject* args) {
     codec_options_t options;
     buffer_t buffer;
     struct module_state *state = GETSTATE(self);
+    if (!state) {
+        return NULL;
+    }
 
     if (!(PyArg_ParseTuple(args, "bOObOO",
                           &op, &command, &docs, &ack,
@@ -867,6 +877,9 @@ _cbson_encode_batched_write_command(PyObject* self, PyObject* args) {
     codec_options_t options;
     buffer_t buffer;
     struct module_state *state = GETSTATE(self);
+    if (!state) {
+        return NULL;
+    }
 
     if (!(PyArg_ParseTuple(args, "et#bOOOO", "utf-8",
                           &ns, &ns_len, &op, &command, &docs,
@@ -983,6 +996,9 @@ _cmessage_exec(PyObject *m)
     }
 
     state = GETSTATE(m);
+    if (state == NULL) {
+        goto fail;
+    }
     state->_cbson = _cbson;
     if (!((state->_max_bson_size_str = PyUnicode_FromString("max_bson_size")) &&
         (state->_max_message_size_str = PyUnicode_FromString("max_message_size")) &&