Catch buffer overflows in inserttable function

These could happen if the columns specification was extremely large.
PyGreSQL · Jan 28, 2022 · a0025f8 · a0025f8
1 parent fef206b
commit a0025f8
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git a/pgconn.c b/pgconn.c
@@ -690,7 +690,7 @@ static PyObject *
 conn_inserttable(connObject *self, PyObject *args)
 {
     PGresult *result;
-    char *table, *buffer, *bufpt;
+    char *table, *buffer, *bufpt, *bufmax;
     int encoding;
     size_t bufsiz;
     PyObject *list, *sublist, *item, *columns = NULL;
@@ -761,12 +761,14 @@ conn_inserttable(connObject *self, PyObject *args)
 
     /* starts query */
     bufpt = buffer;
+    bufmax = bufpt + MAX_BUFFER_SIZE;
     table = PQescapeIdentifier(self->cnx, table, strlen(table));
-    bufpt += sprintf(bufpt, "copy %s", table);
+    bufpt += snprintf(bufpt, (size_t) (bufmax - bufpt), "copy %s", table);
     PQfreemem(table);
     if (columns) {
         /* adds a string like f" ({','.join(columns)})" */
-        bufpt += sprintf(bufpt, " (");
+        if (bufpt < bufmax)
+            bufpt += snprintf(bufpt, (size_t) (bufmax - bufpt), " (");
         for (j = 0; j < n; ++j) {
             PyObject *obj = getcolumn(columns, j);
             Py_ssize_t slen;
@@ -786,11 +788,17 @@ conn_inserttable(connObject *self, PyObject *args)
                     "The third argument must contain only strings");
             }
             col = PQescapeIdentifier(self->cnx, col, (size_t) slen);
-            bufpt += sprintf(bufpt, "%s%s", col, j == n - 1 ? ")" : ",");
+            if (bufpt < bufmax)
+                bufpt += snprintf(bufpt, (size_t) (bufmax - bufpt),
+                    "%s%s", col, j == n - 1 ? ")" : ",");
             PQfreemem(col);
         }
     }
-    sprintf(bufpt, " from stdin");
+    if (bufpt < bufmax)
+        snprintf(bufpt, (size_t) (bufmax - bufpt), " from stdin");
+    if (bufpt >= bufmax)  {
+        PyMem_Free(buffer); return PyErr_NoMemory();
+    }
 
     Py_BEGIN_ALLOW_THREADS
     result = PQexec(self->cnx, buffer);

diff --git a/tests/test_classic_connection.py b/tests/test_classic_connection.py
@@ -2075,6 +2075,11 @@ def testInserttableNoEncoding(self):
         # cannot encode non-ascii unicode without a specific encoding
         self.assertRaises(UnicodeEncodeError, self.c.inserttable, 'test', data)
 
+    def testInserttableTooLargeColumnSpecification(self):
+        # should catch buffer overflow when building the column specification
+        self.assertRaises(MemoryError, self.c.inserttable,
+                          'test', [], ['very_long_column_name'] * 1000)
+
 
 class TestDirectSocketAccess(unittest.TestCase):
     """Test copy command with direct socket access."""