More careful memory management

- Check return value of PyErr_Malloc and set an exception in case of error
- Avoid exposing variables with refcount 0 as connection attributes.
- PyErr_Free guards itself for NULL input

7 files changed, 88 insertions, 51 deletions

diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c
index 882b0ef..37fc9b1 100644
--- a/psycopg/connection_int.c
+++ b/psycopg/connection_int.c
@@ -984,14 +984,22 @@ conn_set_client_encoding(connectionObject *self, const char *enc)
     }
 
     /* no error, we can proceeed and store the new encoding */
-    PyMem_Free(self->encoding);
+    {
+        char *tmp = self->encoding;
+        self->encoding = NULL;
+        PyMem_Free(tmp);
+    }
     if (!(self->encoding = psycopg_strdup(enc, 0))) {
         res = 1;  /* don't call pq_complete_error below */
         goto endlock;
     }
 
     /* Store the python codec too. */
-    PyMem_Free(self->codec);
+    {
+        char *tmp = self->codec;
+        self->codec = NULL;
+        PyMem_Free(tmp);
+    }
     self->codec = codec;
 
     Dprintf("conn_set_client_encoding: set encoding to %s (codec: %s)",
diff --git a/psycopg/cursor_type.c b/psycopg/cursor_type.c
index 5bb9095..4eee0d5 100644
--- a/psycopg/cursor_type.c
+++ b/psycopg/cursor_type.c
@@ -1011,7 +1011,9 @@ psyco_curs_callproc(cursorObject *self, PyObject *args, PyObject *kwargs)
     /* allocate some memory, build the SQL and create a PyString from it */
     sl = procname_len + 17 + nparameters*3 - (nparameters ? 1 : 0);
     sql = (char*)PyMem_Malloc(sl);
-    if (sql == NULL) return NULL;
+    if (sql == NULL) {
+        return PyErr_NoMemory();
+    }
 
     sprintf(sql, "SELECT * FROM %s(", procname);
     for(i=0; i<nparameters; i++) {
@@ -1233,15 +1235,16 @@ _psyco_curs_has_read_check(PyObject* o, void* var)
 static PyObject *
 psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
 {
+    char *query = NULL;
     char query_buffer[DEFAULT_COPYBUFF];
     Py_ssize_t query_size;
-    char *query;
     const char *table_name;
     const char *sep = "\t", *null = NULL;
     Py_ssize_t bufsize = DEFAULT_COPYBUFF;
     PyObject *file, *columns = NULL, *res = NULL;
     char columnlist[DEFAULT_COPYBUFF];
-    char *quoted_delimiter;
+    char *quoted_delimiter = NULL;
+    char *quoted_null = NULL;
 
     static char *kwlist[] = {
             "file", "table", "sep", "null", "size", "columns", NULL};
@@ -1262,32 +1265,32 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
     EXC_IF_GREEN(copy_from);
     EXC_IF_TPC_PREPARED(self->conn, copy_from);
 
-
-    quoted_delimiter = psycopg_escape_string((PyObject*)self->conn, sep, 0, NULL, NULL);
-    if (quoted_delimiter == NULL) {
+    if (!(quoted_delimiter = psycopg_escape_string(
+            (PyObject*)self->conn, sep, 0, NULL, NULL))) {
         PyErr_NoMemory();
-        return NULL;
+        goto exit;
     }
-    
+
     query = query_buffer;
     if (null) {
-        char *quoted_null = psycopg_escape_string((PyObject*)self->conn, null, 0, NULL, NULL);
-        if (quoted_null == NULL) {
-            PyMem_Free(quoted_delimiter);
+        if (!(quoted_null = psycopg_escape_string(
+                (PyObject*)self->conn, null, 0, NULL, NULL))) {
             PyErr_NoMemory();
-            return NULL;
+            goto exit;
         }
         query_size = PyOS_snprintf(query, DEFAULT_COPYBUFF,
             "COPY %s%s FROM stdin WITH DELIMITER AS %s NULL AS %s",
             table_name, columnlist, quoted_delimiter, quoted_null);
         if (query_size >= DEFAULT_COPYBUFF) {
             /* Got truncated, allocate dynamically */
-            query = (char *)PyMem_Malloc((query_size + 1) * sizeof(char));
+            if (!(query = PyMem_New(char, query_size + 1))) {
+                PyErr_NoMemory();
+                goto exit;
+            }
             PyOS_snprintf(query, query_size + 1,
                 "COPY %s%s FROM stdin WITH DELIMITER AS %s NULL AS %s",
                 table_name, columnlist, quoted_delimiter, quoted_null);
         }
-        PyMem_Free(quoted_null);
     }
     else {
         query_size = PyOS_snprintf(query, DEFAULT_COPYBUFF,
@@ -1295,14 +1298,16 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
            table_name, columnlist, quoted_delimiter);
         if (query_size >= DEFAULT_COPYBUFF) {
             /* Got truncated, allocate dynamically */
-            query = (char *)PyMem_Malloc((query_size + 1) * sizeof(char));
+            if (!(query = PyMem_New(char, query_size + 1))) {
+                PyErr_NoMemory();
+                goto exit;
+            }
             PyOS_snprintf(query, query_size + 1,
                 "COPY %s%s FROM stdin WITH DELIMITER AS %s",
                 table_name, columnlist, quoted_delimiter);
         }
-    }    
-    PyMem_Free(quoted_delimiter);
-    
+    }
+
     Dprintf("psyco_curs_copy_from: query = %s", query);
 
     self->copysize = bufsize;
@@ -1313,11 +1318,13 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
         Py_INCREF(Py_None);
     }
 
-    if (query && (query != query_buffer)) {
-        PyMem_Free(query);
-    }
     self->copyfile = NULL;
 
+exit:
+    PyMem_Free(quoted_delimiter);
+    PyMem_Free(quoted_null);
+    if (query != query_buffer) { PyMem_Free(query); }
+
     return res;
 }
 
@@ -1352,7 +1359,8 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
     const char *table_name;
     const char *sep = "\t", *null = NULL;
     PyObject *file, *columns = NULL, *res = NULL;
-    char *quoted_delimiter;
+    char *quoted_delimiter = NULL;
+    char *quoted_null = NULL;
 
     static char *kwlist[] = {"file", "table", "sep", "null", "columns", NULL};
 
@@ -1370,31 +1378,32 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
     EXC_IF_GREEN(copy_to);
     EXC_IF_TPC_PREPARED(self->conn, copy_to);
 
-    quoted_delimiter = psycopg_escape_string((PyObject*)self->conn, sep, 0, NULL, NULL);
-    if (quoted_delimiter == NULL) {
+    if (!(quoted_delimiter = psycopg_escape_string(
+            (PyObject*)self->conn, sep, 0, NULL, NULL))) {
         PyErr_NoMemory();
-        return NULL;
+        goto exit;
     }
-    
+
     query = query_buffer;
     if (null) {
-        char *quoted_null = psycopg_escape_string((PyObject*)self->conn, null, 0, NULL, NULL);
-        if (NULL == quoted_null) {
-            PyMem_Free(quoted_delimiter);
+        if (!(quoted_null = psycopg_escape_string(
+                (PyObject*)self->conn, null, 0, NULL, NULL))) {
             PyErr_NoMemory();
-            return NULL;
+            goto exit;
         }
         query_size = PyOS_snprintf(query, DEFAULT_COPYBUFF,
             "COPY %s%s TO stdout WITH DELIMITER AS %s"
             " NULL AS %s", table_name, columnlist, quoted_delimiter, quoted_null);
         if (query_size >= DEFAULT_COPYBUFF) {
             /* Got truncated, allocate dynamically */
-            query = (char *)PyMem_Malloc((query_size + 1) * sizeof(char));
+            if (!(query = PyMem_New(char, query_size + 1))) {
+                PyErr_NoMemory();
+                goto exit;
+            }
             PyOS_snprintf(query, query_size + 1,
                 "COPY %s%s TO stdout WITH DELIMITER AS %s"
                 " NULL AS %s", table_name, columnlist, quoted_delimiter, quoted_null);
         }
-        PyMem_Free(quoted_null);
     }
     else {
         query_size = PyOS_snprintf(query, DEFAULT_COPYBUFF,
@@ -1402,14 +1411,16 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
             table_name, columnlist, quoted_delimiter);
         if (query_size >= DEFAULT_COPYBUFF) {
             /* Got truncated, allocate dynamically */
-            query = (char *)PyMem_Malloc((query_size + 1) * sizeof(char));
+            if (!(query = PyMem_New(char, query_size + 1))) {
+                PyErr_NoMemory();
+                goto exit;
+            }
             PyOS_snprintf(query, query_size + 1,
                 "COPY %s%s TO stdout WITH DELIMITER AS %s",
                 table_name, columnlist, quoted_delimiter);
         }
     }
-    PyMem_Free(quoted_delimiter);
-    
+
     Dprintf("psyco_curs_copy_to: query = %s", query);
 
     self->copysize = 0;
@@ -1419,11 +1430,13 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
         res = Py_None;
         Py_INCREF(Py_None);
     }
-    if (query && (query != query_buffer)) {
-        PyMem_Free(query);
-    }
     self->copyfile = NULL;
 
+exit:
+    PyMem_Free(quoted_delimiter);
+    PyMem_Free(quoted_null);
+    if (query != query_buffer) { PyMem_Free(query); }
+
     return res;
 }
 
@@ -1653,8 +1666,10 @@ cursor_setup(cursorObject *self, connectionObject *conn, const char *name)
     Dprintf("cursor_setup: parameters: name = %s, conn = %p", name, conn);
 
     if (name) {
-        self->name = PyMem_Malloc(strlen(name)+1);
-        if (self->name == NULL) return 1;
+        if (!(self->name = PyMem_Malloc(strlen(name)+1))) {
+            PyErr_NoMemory();
+            return 1;
+        }
         strncpy(self->name, name, strlen(name)+1);
     }
 
@@ -1715,7 +1730,7 @@ cursor_dealloc(PyObject* obj)
 
     PyObject_GC_UnTrack(self);
 
-    if (self->name) PyMem_Free(self->name);
+    PyMem_Free(self->name);
 
     Py_CLEAR(self->conn);
     Py_CLEAR(self->casts);
diff --git a/psycopg/lobject_int.c b/psycopg/lobject_int.c
index 252d1c9..3fe1f86 100644
--- a/psycopg/lobject_int.c
+++ b/psycopg/lobject_int.c
@@ -110,6 +110,8 @@ _lobject_parse_mode(const char *mode)
 /* Return a string representing the lobject mode.
  *
  * The return value is a new string allocated on the Python heap.
+ *
+ * The function must be called holding the GIL.
  */
 static char *
 _lobject_unparse_mode(int mode)
@@ -118,7 +120,10 @@ _lobject_unparse_mode(int mode)
     char *c;
 
     /* the longest is 'rwt' */
-    c = buf = PyMem_Malloc(4);
+    if (!(c = buf = PyMem_Malloc(4))) {
+        PyErr_NoMemory();
+        return NULL;
+    }
 
     if (mode & LOBJECT_READ) { *c++ = 'r'; }
     if (mode & LOBJECT_WRITE) { *c++ = 'w'; }
@@ -204,7 +209,14 @@ lobject_open(lobjectObject *self, connectionObject *conn,
 
     /* set the mode for future reference */
     self->mode = mode;
+    Py_BLOCK_THREADS;
     self->smode = _lobject_unparse_mode(mode);
+    Py_UNBLOCK_THREADS;
+    if (NULL == self->smode) {
+        retvalue = 1;  /* exception already set */
+        goto end;
+    }
+
     retvalue = 0;
 
  end:
@@ -213,6 +225,8 @@ lobject_open(lobjectObject *self, connectionObject *conn,
 
     if (retvalue < 0)
         pq_complete_error(self->conn, &pgres, &error);
+    /* if retvalue > 0, an exception is already set */
+
     return retvalue;
 }
 
diff --git a/psycopg/lobject_type.c b/psycopg/lobject_type.c
index 39ff6ea..51c41d3 100644
--- a/psycopg/lobject_type.c
+++ b/psycopg/lobject_type.c
@@ -344,7 +344,7 @@ lobject_setup(lobjectObject *self, connectionObject *conn,
     self->fd = -1;
     self->oid = InvalidOid;
 
-    if (lobject_open(self, conn, oid, smode, new_oid, new_file) == -1)
+    if (0 != lobject_open(self, conn, oid, smode, new_oid, new_file))
         return -1;
 
    Dprintf("lobject_setup: good lobject object at %p, refcnt = "
diff --git a/psycopg/typecast_array.c b/psycopg/typecast_array.c
index 8349a1d..6bb256f 100644
--- a/psycopg/typecast_array.c
+++ b/psycopg/typecast_array.c
@@ -135,7 +135,10 @@ typecast_array_tokenize(const char *str, Py_ssize_t strlength,
     if (res == ASCAN_QUOTED) {
         Py_ssize_t j;
         char *buffer = PyMem_Malloc(l+1);
-        if (buffer == NULL) return ASCAN_ERROR;
+        if (buffer == NULL) {
+            PyErr_NoMemory();
+            return ASCAN_ERROR;
+        }
 
         *token = buffer;
 
diff --git a/psycopg/typecast_binary.c b/psycopg/typecast_binary.c
index e93aad5..472b823 100644
--- a/psycopg/typecast_binary.c
+++ b/psycopg/typecast_binary.c
@@ -201,10 +201,8 @@ typecast_BINARY_cast(const char *s, Py_ssize_t l, PyObject *curs)
           /* str's mem was allocated by PQunescapeBytea; must use PQfreemem: */
           PQfreemem(str);
       }
-      if (buffer != NULL) {
-          /* We allocated buffer with PyMem_Malloc; must use PyMem_Free: */
-          PyMem_Free(buffer);
-      }
+      /* We allocated buffer with PyMem_Malloc; must use PyMem_Free: */
+      PyMem_Free(buffer);
 
       return res;
 }
diff --git a/psycopg/xid_type.c b/psycopg/xid_type.c
index 9027b34..82d6d89 100644
--- a/psycopg/xid_type.c
+++ b/psycopg/xid_type.c
@@ -436,7 +436,6 @@ _xid_decode64(PyObject *s)
  * in order to allow some form of interoperation.
  *
  * The function must be called while holding the GIL.
- * Return a buffer allocated with PyMem_Malloc. Use PyMem_Free to free it.
  *
  * see also: the pgjdbc implementation
  *   http://cvs.pgfoundry.org/cgi-bin/cvsweb.cgi/jdbc/pgjdbc/org/postgresql/xa/RecoveredXid.java?rev=1.2