psycopg/psycopg2
1
1
Fork 0
mirror of https://github.com/psycopg/psycopg2.git synced 2024-11-11 19:46:36 +03:00
Code Issues Packages Projects Releases Wiki Activity

callproc using a dict now uses connection encoding and sanitizes parameter names

Browse Source
This commit is contained in:
mrmilosz 2014-05-28 01:05:17 -04:00
parent 1205bf9c2b
commit 31c95c0922
2 changed files with 46 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

53
psycopg/cursor_type.c
View File

@ -1025,6 +1025,10 @@ psyco_curs_callproc(cursorObject *self, PyObject *args)
PyObject *operation = NULL; PyObject *operation = NULL;
PyObject *res = NULL; PyObject *res = NULL;
PyObject *parameter_name = NULL; PyObject *parameter_name = NULL;
PyObject *parameter_name_bytes = NULL;
char *parameter_name_cstr = NULL;
char *parameter_name_cstr_sanitized = NULL;
char **parameter_name_cstr_sanitized_CACHE = NULL;
PyObject *parameter_names = NULL; PyObject *parameter_names = NULL;
if (!PyArg_ParseTuple(args, "s#|O", if (!PyArg_ParseTuple(args, "s#|O",
@ -1048,40 +1052,71 @@ psyco_curs_callproc(cursorObject *self, PyObject *args)
/* allocate some memory, build the SQL and create a PyString from it */ /* allocate some memory, build the SQL and create a PyString from it */
/* a dict requires special handling: we put the parameter names into the SQL */
if (nparameters > 0 && PyDict_Check(parameters)) { if (nparameters > 0 && PyDict_Check(parameters)) {
/* for a dict, we put the parameter names into the SQL */
parameter_names = PyDict_Keys(parameters); parameter_names = PyDict_Keys(parameters);
/* first we need to figure out how much space we need for the SQL */ /* first we need to ensure the dict's keys are text */
sl = procname_len + 17 + nparameters*5 - (nparameters ? 1 : 0);
for(i=0; i<nparameters; i++) { for(i=0; i<nparameters; i++) {
parameter_name = PyList_GetItem(parameter_names, i); if (!Text_Check(PyList_GetItem(parameter_names, i))) {
if (!Text_Check(parameter_name)) { PyErr_SetString(PyExc_TypeError, "argument 2 must have only string keys if Dict");
PyErr_SetString(PyExc_TypeError, "argument 2 must have string keys if Dict"); Py_DECREF(parameter_names);
goto exit; goto exit;
} }
sl += strlen(Text_AsUTF8(parameter_name));
} }
/* we need one extra pass to determine the amount of memory to allocate for the SQL */
sl = procname_len + 17 + nparameters*5 - (nparameters ? 1 : 0);
/* we will throw the sanitized C strings into a cache to not redo the work later */
parameter_name_cstr_sanitized_CACHE = PyMem_New(char *, nparameters);
for(i=0; i<nparameters; i++) {
parameter_name = PyList_GetItem(parameter_names, i);
Py_INCREF(parameter_name);
parameter_name_bytes = psycopg_ensure_bytes(parameter_name);
parameter_name_cstr = Bytes_AsString(parameter_name_bytes);
parameter_name_cstr_sanitized = PQescapeIdentifier(self->conn->pgconn, parameter_name_cstr, strlen(parameter_name_cstr));
Py_DECREF(parameter_name);
/* must add the length of the sanitized string to the length of the SQL string */
sl += strlen(parameter_name_cstr_sanitized);
parameter_name_cstr_sanitized_CACHE[i] = parameter_name_cstr_sanitized;
}
Py_DECREF(parameter_names);
sql = (char*)PyMem_Malloc(sl); sql = (char*)PyMem_Malloc(sl);
if (sql == NULL) { if (sql == NULL) {
PyErr_NoMemory(); PyErr_NoMemory();
for(i=0; i<nparameters; i++) {
PQfreemem(parameter_name_cstr_sanitized_CACHE[i]);
}
PyMem_Del(parameter_name_cstr_sanitized_CACHE);
goto exit; goto exit;
} }
sprintf(sql, "SELECT * FROM %s(", procname); sprintf(sql, "SELECT * FROM %s(", procname);
for(i=0; i<nparameters; i++) { for(i=0; i<nparameters; i++) {
parameter_name_cstr_sanitized = parameter_name_cstr_sanitized_CACHE[i];
/* like procname, param names are not sanitized. don't SQL inject yourself. */ /* like procname, param names are not sanitized. don't SQL inject yourself. */
strcat(sql, Text_AsUTF8(PyList_GetItem(parameter_names, i))); strcat(sql, parameter_name_cstr_sanitized);
strcat(sql, ":=%s,"); strcat(sql, ":=%s,");
PQfreemem(parameter_name_cstr_sanitized);
} }
PyMem_Del(parameter_name_cstr_sanitized_CACHE);
sql[sl-2] = ')'; sql[sl-2] = ')';
sql[sl-1] = '\0'; sql[sl-1] = '\0';
Py_DECREF(parameter_names);
/* now that we have the query string we can discard the keys and proceed normally */ /* now that we have the query string we can discard the keys and proceed normally */
parameters = PyDict_Values(parameters); parameters = PyDict_Values(parameters);
} }
/* a list (or None) is a little bit simpler */
else { else {
sl = procname_len + 17 + nparameters*3 - (nparameters ? 1 : 0); sl = procname_len + 17 + nparameters*3 - (nparameters ? 1 : 0);

2
psycopg/python.h
View File

@ -71,14 +71,12 @@ typedef unsigned long Py_uhash_t;
#define Text_Format(f,a) PyString_Format(f,a) #define Text_Format(f,a) PyString_Format(f,a)
#define Text_FromUTF8(s) PyString_FromString(s) #define Text_FromUTF8(s) PyString_FromString(s)
#define Text_FromUTF8AndSize(s,n) PyString_FromStringAndSize(s,n) #define Text_FromUTF8AndSize(s,n) PyString_FromStringAndSize(s,n)
#define Text_AsUTF8(s) PyString_AsString(s)
#else #else
#define Text_Type PyUnicode_Type #define Text_Type PyUnicode_Type
#define Text_Check(s) PyUnicode_Check(s) #define Text_Check(s) PyUnicode_Check(s)
#define Text_Format(f,a) PyUnicode_Format(f,a) #define Text_Format(f,a) PyUnicode_Format(f,a)
#define Text_FromUTF8(s) PyUnicode_FromString(s) #define Text_FromUTF8(s) PyUnicode_FromString(s)
#define Text_FromUTF8AndSize(s,n) PyUnicode_FromStringAndSize(s,n) #define Text_FromUTF8AndSize(s,n) PyUnicode_FromStringAndSize(s,n)
#define Text_AsUTF8(s) PyUnicode_AsUTF8(s)
#endif #endif
#if PY_MAJOR_VERSION > 2 #if PY_MAJOR_VERSION > 2