From 929907364962c52f5652cf748f632850d97ad6f0 Mon Sep 17 00:00:00 2001
From: Federico Di Gregorio <fog@initd.org>
Date: Wed, 24 May 2006 09:43:55 +0000
Subject: [PATCH] First half of 8.1.4 securiy patch.

---
 lib/extras.py             |  9 ++++++++-
 psycopg/adapter_qstring.c | 29 ++++++++++++++++++++++++-----
 psycopg/adapter_qstring.h |  2 ++
 3 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/lib/extras.py b/lib/extras.py
index 7a57c710..847bdb70 100644
--- a/lib/extras.py
+++ b/lib/extras.py
@@ -125,10 +125,17 @@ class SQL_IN(object):
     def __init__(self, seq):
         self._seq = seq
 
+    def prepare(self, conn):
+        self._conn = conn
+    
     def getquoted(self):
         # this is the important line: note how every object in the
         # list is adapted and then how getquoted() is called on it
-        qobjs = [str(_A(o).getquoted()) for o in self._seq]
+        pobjs = [_A(o) for o in self._seq]
+        for obj in pobjs:
+            if hasattr(obj, 'prepare'):
+                obj.prepare(self._conn)
+        qobjs = [str(o.getquoted()) for o in pobjs]
         return '(' + ', '.join(qobjs) + ')'
 
     __str__ = getquoted
diff --git a/psycopg/adapter_qstring.c b/psycopg/adapter_qstring.c
index 788d9c2f..50c62c89 100644
--- a/psycopg/adapter_qstring.c
+++ b/psycopg/adapter_qstring.c
@@ -38,10 +38,19 @@
 /** the quoting code */
 
 #ifndef PSYCOPG_OWN_QUOTING
-#define qstring_escape PQescapeString
+static size_t
+qstring_escape(char *to, char *from, size_t len, PGconn *conn)
+{
+    int err = 0;
+
+    if (conn)
+        return PQescapeStringConn(conn, to, from, len, &err);
+    else
+        return PQescapeString(to, from, len);
+}
 #else
 static size_t
-qstring_escape(char *to, char *from, size_t len)
+qstring_escape(char *to, char *from, size_t len, PGconn *conn)
 {
     int i, j;
 
@@ -134,7 +143,8 @@ qstring_quote(qstringObject *self)
     }
 
     Py_BEGIN_ALLOW_THREADS;
-    len = qstring_escape(buffer+1, s, len);
+    len = qstring_escape(buffer+1, s, len,
+                         ((connectionObject*)self->conn)->pgconn);
     buffer[0] = '\'' ; buffer[len+1] = '\'';
     Py_END_ALLOW_THREADS;
     
@@ -179,7 +189,13 @@ qstring_prepare(qstringObject *self, PyObject *args)
         self->encoding = strdup(conn->encoding);
         Dprintf("qstring_prepare: set encoding to %s", conn->encoding);
     }
-    
+
+    Py_XDECREF(self->conn);
+    if (conn) {
+        self->conn = (PyObject*)conn;
+        Py_INCREF(self->conn);
+    }
+
     Py_INCREF(Py_None);
     return Py_None;
 }
@@ -217,7 +233,7 @@ static PyMethodDef qstringObject_methods[] = {
     {"getquoted", (PyCFunction)qstring_getquoted, METH_VARARGS,
      "getquoted() -> wrapped object value as SQL-quoted string"},
     {"prepare", (PyCFunction)qstring_prepare, METH_VARARGS,
-     "prepare(conn) -> set encoding to conn->encoding"},
+     "prepare(conn) -> set encoding to conn->encoding and store conn"},
     {"__conform__", (PyCFunction)qstring_conform, METH_VARARGS, NULL},
     {NULL}  /* Sentinel */
 };
@@ -231,6 +247,7 @@ qstring_setup(qstringObject *self, PyObject *str, char *enc)
             self, ((PyObject *)self)->ob_refcnt);
 
     self->buffer = NULL;
+    self->conn = NULL;
 
     /* FIXME: remove this orrible strdup */
     if (enc) self->encoding = strdup(enc);
@@ -250,6 +267,8 @@ qstring_dealloc(PyObject* obj)
 
     Py_XDECREF(self->wrapped);
     Py_XDECREF(self->buffer);
+    Py_XDECREF(self->conn);
+
     if (self->encoding) free(self->encoding);
     
     Dprintf("qstring_dealloc: deleted qstring object at %p, refcnt = %d",
diff --git a/psycopg/adapter_qstring.h b/psycopg/adapter_qstring.h
index 9e6c33dd..544b32b6 100644
--- a/psycopg/adapter_qstring.h
+++ b/psycopg/adapter_qstring.h
@@ -36,6 +36,8 @@ typedef struct {
     PyObject *wrapped;
     PyObject *buffer;
     char     *encoding;
+
+    PyObject *conn;
 } qstringObject;
     
 /* functions exported to psycopgmodule.c */