From c1715f66fe483ceacda59f8f763d5128bc0f89e7 Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Wed, 23 Feb 2011 00:28:11 +0000 Subject: [PATCH] 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 --- psycopg/connection_int.c | 12 ++++- psycopg/cursor_type.c | 97 ++++++++++++++++++++++----------------- psycopg/lobject_int.c | 16 ++++++- psycopg/lobject_type.c | 2 +- psycopg/typecast_array.c | 5 +- psycopg/typecast_binary.c | 6 +-- psycopg/xid_type.c | 1 - 7 files changed, 88 insertions(+), 51 deletions(-) diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c index 882b0eff..37fc9b13 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 5bb9095a..4eee0d5d 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; iconn, 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 252d1c93..3fe1f86e 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 39ff6ead..51c41d35 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 8349a1d9..6bb256f1 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 e93aad57..472b823f 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 9027b349..82d6d899 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