Merge branch 'fix-528'

7 files changed, 145 insertions, 40 deletions

diff --git a/NEWS b/NEWS
index e28202c..b2d2153 100644
--- a/NEWS
+++ b/NEWS
@@ -7,6 +7,8 @@ What's new in psycopg 2.7.2
 - Fixed inconsistent state in externally closed connections
   (:tickets:`#263, #311, #443`). Was fixed in 2.6.2 but not included in
   2.7 by mistake.
+- Don't display the password in `connection.dsn` when the connection
+  string is specified as an URI (:ticket:`#528`).
 
 
 What's new in psycopg 2.7.1
diff --git a/doc/src/connection.rst b/doc/src/connection.rst
index 53f908f..fbbc53e 100644
--- a/doc/src/connection.rst
+++ b/doc/src/connection.rst
@@ -343,6 +343,9 @@ The ``connection`` class
         Read-only string containing the connection string used by the
         connection.
 
+        If a password was specified in the connection string it will be
+        obscured.
+
 
     .. index::
         pair: Transaction; Autocommit
diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c
index 26100b2..eb4b241 100644
--- a/psycopg/connection_type.c
+++ b/psycopg/connection_type.c
@@ -1248,10 +1248,58 @@ static struct PyGetSetDef connectionObject_getsets[] = {
 
 /* initialization and finalization methods */
 
+RAISES_NEG static int
+obscure_password(connectionObject *conn)
+{
+    PQconninfoOption *options;
+    PyObject *d = NULL, *v = NULL, *dsn = NULL;
+    char *tmp;
+    int rv = -1;
+
+    if (!conn || !conn->dsn) {
+        return 0;
+    }
+
+    if (!(options = PQconninfoParse(conn->dsn, NULL))) {
+        /* unlikely: the dsn was already tested valid */
+        return 0;
+    }
+
+    if (!(d = psycopg_dict_from_conninfo_options(
+            options, /* include_password = */ 1))) {
+        goto exit;
+    }
+    if (NULL == PyDict_GetItemString(d, "password")) {
+        /* the dsn doesn't have a password */
+        rv = 0;
+        goto exit;
+    }
+
+    /* scrub the password and put back the connection string together */
+    if (!(v = Text_FromUTF8("xxx"))) { goto exit; }
+    if (0 > PyDict_SetItemString(d, "password", v)) { goto exit; }
+    if (!(dsn = psycopg_make_dsn(Py_None, d))) { goto exit; }
+    if (!(dsn = psycopg_ensure_bytes(dsn))) { goto exit; }
+
+    /* Replace the connection string on the connection object */
+    tmp = conn->dsn;
+    psycopg_strdup(&conn->dsn, Bytes_AS_STRING(dsn), -1);
+    PyMem_Free(tmp);
+
+    rv = 0;
+
+exit:
+    PQconninfoFree(options);
+    Py_XDECREF(v);
+    Py_XDECREF(d);
+    Py_XDECREF(dsn);
+
+    return rv;
+}
+
 static int
 connection_setup(connectionObject *self, const char *dsn, long int async)
 {
-    char *pos;
     int res = -1;
 
     Dprintf("connection_setup: init connection object at %p, "
@@ -1288,15 +1336,16 @@ connection_setup(connectionObject *self, const char *dsn, long int async)
 
 exit:
     /* here we obfuscate the password even if there was a connection error */
-    pos = strstr(self->dsn, "password");
-    if (pos != NULL) {
-        for (pos = pos+9 ; *pos != '\0' && *pos != ' '; pos++)
-            *pos = 'x';
+    {
+        PyObject *ptype = NULL, *pvalue = NULL, *ptb = NULL;
+        PyErr_Fetch(&ptype, &pvalue, &ptb);
+        obscure_password(self);
+        PyErr_Restore(ptype, pvalue, ptb);
     }
-
     return res;
 }
 
+
 static int
 connection_clear(connectionObject *self)
 {
diff --git a/psycopg/psycopg.h b/psycopg/psycopg.h
index 1367354..d54037f 100644
--- a/psycopg/psycopg.h
+++ b/psycopg/psycopg.h
@@ -142,6 +142,8 @@ STEALS(1) HIDDEN PyObject * psycopg_ensure_text(PyObject *obj);
 HIDDEN PyObject *psycopg_dict_from_conninfo_options(PQconninfoOption *options,
               int include_password);
 
+HIDDEN PyObject *psycopg_make_dsn(PyObject *dsn, PyObject *kwargs);
+
 /* Exceptions docstrings */
 #define Error_doc \
 "Base class for error exceptions."
diff --git a/psycopg/replication_connection_type.c b/psycopg/replication_connection_type.c
index 5e5d222..f4b5eb8 100644
--- a/psycopg/replication_connection_type.c
+++ b/psycopg/replication_connection_type.c
@@ -60,35 +60,28 @@ psyco_repl_conn_get_type(replicationConnectionObject *self)
 
 
 static int
-replicationConnection_init(PyObject *obj, PyObject *args, PyObject *kwargs)
+replicationConnection_init(replicationConnectionObject *self,
+    PyObject *args, PyObject *kwargs)
 {
-    replicationConnectionObject *self = (replicationConnectionObject *)obj;
-    PyObject *dsn = NULL, *replication_type = NULL,
-        *item = NULL, *ext = NULL, *make_dsn = NULL,
-        *extras = NULL, *cursor = NULL;
-    int async = 0;
+    PyObject *dsn = NULL, *async = Py_False, *replication_type = NULL,
+        *item = NULL, *extras = NULL, *cursor = NULL,
+        *newdsn = NULL, *newargs = NULL, *dsnopts = NULL;
     int ret = -1;
 
     /* 'replication_type' is not actually optional, but there's no
        good way to put it before 'async' in the list */
     static char *kwlist[] = {"dsn", "async", "replication_type", NULL};
-    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|iO", kwlist,
-                                     &dsn, &async, &replication_type)) { return ret; }
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|OO", kwlist,
+                                     &dsn, &async, &replication_type)) {
+        return ret;
+    }
 
     /*
       We have to call make_dsn() to add replication-specific
       connection parameters, because the DSN might be an URI (if there
       were no keyword arguments to connect() it is passed unchanged).
     */
-    /* we reuse args and kwargs to call make_dsn() and parent type's tp_init() */
-    if (!(kwargs = PyDict_New())) { return ret; }
-    Py_INCREF(args);
-
-    /* we also reuse the dsn to hold the result of the make_dsn() call */
-    Py_INCREF(dsn);
-
-    if (!(ext = PyImport_ImportModule("psycopg2.extensions"))) { goto exit; }
-    if (!(make_dsn = PyObject_GetAttrString(ext, "make_dsn"))) { goto exit; }
+    if (!(dsnopts = PyDict_New())) { return ret; }
 
     /* all the nice stuff is located in python-level ReplicationCursor class */
     if (!(extras = PyImport_ImportModule("psycopg2.extras"))) { goto exit; }
@@ -101,7 +94,7 @@ replicationConnection_init(PyObject *obj, PyObject *args, PyObject *kwargs)
 
 #define SET_ITEM(k, v) \
         if (!(item = Text_FromUTF8(#v))) { goto exit; } \
-        if (PyDict_SetItemString(kwargs, #k, item) != 0) { goto exit; } \
+        if (PyDict_SetItemString(dsnopts, #k, item) != 0) { goto exit; } \
         Py_DECREF(item); \
         item = NULL;
 
@@ -118,30 +111,25 @@ replicationConnection_init(PyObject *obj, PyObject *args, PyObject *kwargs)
         goto exit;
     }
 
-    Py_DECREF(args);
-    if (!(args = PyTuple_Pack(1, dsn))) { goto exit; }
-
-    Py_DECREF(dsn);
-    if (!(dsn = PyObject_Call(make_dsn, args, kwargs))) { goto exit; }
-
-    Py_DECREF(args);
-    if (!(args = Py_BuildValue("(Oi)", dsn, async))) { goto exit; }
+    if (!(newdsn = psycopg_make_dsn(dsn, dsnopts))) { goto exit; }
+    if (!(newargs = PyTuple_Pack(2, newdsn, async))) { goto exit; }
 
     /* only attempt the connection once we've handled all possible errors */
-    if ((ret = connectionType.tp_init(obj, args, NULL)) < 0) { goto exit; }
+    if ((ret = connectionType.tp_init((PyObject *)self, newargs, NULL)) < 0) {
+        goto exit;
+    }
 
     self->conn.autocommit = 1;
-    Py_INCREF(self->conn.cursor_factory = cursor);
+    Py_INCREF(cursor);
+    self->conn.cursor_factory = cursor;
 
 exit:
     Py_XDECREF(item);
-    Py_XDECREF(ext);
-    Py_XDECREF(make_dsn);
     Py_XDECREF(extras);
     Py_XDECREF(cursor);
-    Py_XDECREF(dsn);
-    Py_XDECREF(args);
-    Py_XDECREF(kwargs);
+    Py_XDECREF(newdsn);
+    Py_XDECREF(newargs);
+    Py_XDECREF(dsnopts);
 
     return ret;
 }
@@ -212,7 +200,7 @@ PyTypeObject replicationConnectionType = {
     0,          /*tp_descr_get*/
     0,          /*tp_descr_set*/
     0,          /*tp_dictoffset*/
-    replicationConnection_init, /*tp_init*/
+    (initproc)replicationConnection_init, /*tp_init*/
     0,          /*tp_alloc*/
     0,          /*tp_new*/
 };
diff --git a/psycopg/utils.c b/psycopg/utils.c
index 7f6b6e6..7073504 100644
--- a/psycopg/utils.c
+++ b/psycopg/utils.c
@@ -280,6 +280,30 @@ exit:
 }
 
 
+/* Make a connection string out of a string and a dictionary of arguments.
+ *
+ * Helper to call psycopg2.extensions.make_dns()
+ */
+PyObject *
+psycopg_make_dsn(PyObject *dsn, PyObject *kwargs)
+{
+    PyObject *ext = NULL, *make_dsn = NULL;
+    PyObject *args = NULL, *rv = NULL;
+
+    if (!(ext = PyImport_ImportModule("psycopg2.extensions"))) { goto exit; }
+    if (!(make_dsn = PyObject_GetAttrString(ext, "make_dsn"))) { goto exit; }
+
+    if (!(args = PyTuple_Pack(1, dsn))) { goto exit; }
+    rv = PyObject_Call(make_dsn, args, kwargs);
+
+exit:
+    Py_XDECREF(args);
+    Py_XDECREF(make_dsn);
+    Py_XDECREF(ext);
+
+    return rv;
+}
+
 /* Convert a C string into Python Text using a specified codec.
  *
  * The codec is the python function codec.getdecoder(enc). It is only used on
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 385d979..f61b099 100755
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -1479,6 +1479,43 @@ class AutocommitTests(ConnectingTestCase):
         self.assertEqual(cur.fetchone()[0], 'on')
 
 
+class PasswordLeakTestCase(ConnectingTestCase):
+    def setUp(self):
+        super(PasswordLeakTestCase, self).setUp()
+        PasswordLeakTestCase.dsn = None
+
+    class GrassingConnection(ext.connection):
+        """A connection snitching the dsn away.
+
+        This connection passes the dsn to the test case class even if init
+        fails (e.g. connection error). Test that we mangle the dsn ok anyway.
+        """
+
+        def __init__(self, *args, **kwargs):
+            try:
+                super(PasswordLeakTestCase.GrassingConnection, self).__init__(
+                    *args, **kwargs)
+            finally:
+                # The connection is not initialized entirely, however the C
+                # code should have set the dsn, and it should have scrubbed
+                # the password away
+                PasswordLeakTestCase.dsn = self.dsn
+
+    def test_leak(self):
+        self.assertRaises(psycopg2.DatabaseError,
+            self.GrassingConnection, "dbname=nosuch password=whateva")
+        self.assertDsnEqual(self.dsn, "dbname=nosuch password=xxx")
+
+    @skip_before_libpq(9, 2)
+    def test_url_leak(self):
+        self.assertRaises(psycopg2.DatabaseError,
+            self.GrassingConnection,
+            "postgres://someone:whateva@localhost/nosuch")
+
+        self.assertDsnEqual(self.dsn,
+            "user=someone password=xxx host=localhost dbname=nosuch")
+
+
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)