diff --git a/Lib/sqlite3/test/userfunctions.py b/Lib/sqlite3/test/userfunctions.py index 1ed090e3d92..b3da3c425b8 100644 --- a/Lib/sqlite3/test/userfunctions.py +++ b/Lib/sqlite3/test/userfunctions.py @@ -53,6 +53,8 @@ def with_tracebacks(strings): def func_returntext(): return "foo" +def func_returntextwithnull(): + return "1\x002" def func_returnunicode(): return "bar" def func_returnint(): @@ -163,11 +165,21 @@ class AggrSum: def finalize(self): return self.val +class AggrText: + def __init__(self): + self.txt = "" + def step(self, txt): + self.txt = self.txt + txt + def finalize(self): + return self.txt + + class FunctionTests(unittest.TestCase): def setUp(self): self.con = sqlite.connect(":memory:") self.con.create_function("returntext", 0, func_returntext) + self.con.create_function("returntextwithnull", 0, func_returntextwithnull) self.con.create_function("returnunicode", 0, func_returnunicode) self.con.create_function("returnint", 0, func_returnint) self.con.create_function("returnfloat", 0, func_returnfloat) @@ -211,6 +223,12 @@ class FunctionTests(unittest.TestCase): self.assertEqual(type(val), str) self.assertEqual(val, "foo") + def test_func_return_text_with_null_char(self): + cur = self.con.cursor() + res = cur.execute("select returntextwithnull()").fetchone()[0] + self.assertEqual(type(res), str) + self.assertEqual(res, "1\x002") + def test_func_return_unicode(self): cur = self.con.cursor() cur.execute("select returnunicode()") @@ -390,6 +408,7 @@ class AggregateTests(unittest.TestCase): self.con.create_aggregate("checkType", 2, AggrCheckType) self.con.create_aggregate("checkTypes", -1, AggrCheckTypes) self.con.create_aggregate("mysum", 1, AggrSum) + self.con.create_aggregate("aggtxt", 1, AggrText) def tearDown(self): #self.cur.close() @@ -486,6 +505,15 @@ class AggregateTests(unittest.TestCase): val = cur.fetchone()[0] self.assertIsNone(val) + def test_aggr_text(self): + cur = self.con.cursor() + for txt in ["foo", "1\x002"]: + with self.subTest(txt=txt): + cur.execute("select aggtxt(?) from test", (txt,)) + val = cur.fetchone()[0] + self.assertEqual(val, txt) + + class AuthorizerTests(unittest.TestCase): @staticmethod def authorizer_cb(action, arg1, arg2, dbname, source): diff --git a/Misc/NEWS.d/next/Library/2021-08-04-12-29-00.bpo-44822.zePNXA.rst b/Misc/NEWS.d/next/Library/2021-08-04-12-29-00.bpo-44822.zePNXA.rst new file mode 100644 index 00000000000..d078142886d --- /dev/null +++ b/Misc/NEWS.d/next/Library/2021-08-04-12-29-00.bpo-44822.zePNXA.rst @@ -0,0 +1,3 @@ +:mod:`sqlite3` user-defined functions and aggregators returning +:class:`strings ` with embedded NUL characters are no longer +truncated. Patch by Erlend E. Aasland. diff --git a/Modules/_sqlite/connection.c b/Modules/_sqlite/connection.c index f75cf833807..aae6c66d63f 100644 --- a/Modules/_sqlite/connection.c +++ b/Modules/_sqlite/connection.c @@ -519,10 +519,17 @@ _pysqlite_set_result(sqlite3_context* context, PyObject* py_val) } else if (PyFloat_Check(py_val)) { sqlite3_result_double(context, PyFloat_AsDouble(py_val)); } else if (PyUnicode_Check(py_val)) { - const char *str = PyUnicode_AsUTF8(py_val); - if (str == NULL) + Py_ssize_t sz; + const char *str = PyUnicode_AsUTF8AndSize(py_val, &sz); + if (str == NULL) { return -1; - sqlite3_result_text(context, str, -1, SQLITE_TRANSIENT); + } + if (sz > INT_MAX) { + PyErr_SetString(PyExc_OverflowError, + "string is longer than INT_MAX bytes"); + return -1; + } + sqlite3_result_text(context, str, (int)sz, SQLITE_TRANSIENT); } else if (PyObject_CheckBuffer(py_val)) { Py_buffer view; if (PyObject_GetBuffer(py_val, &view, PyBUF_SIMPLE) != 0) {