Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
169 changes: 169 additions & 0 deletions Lib/test/test_sqlite3/test_userfunctions.py
Original file line number Diff line number Diff line change
Expand Up @@ -587,6 +587,33 @@ def value(self): return 1 << 65
self.assertRaisesRegex(sqlite.DataError, "string or blob too big",
self.cur.execute, self.query % "err_val_ret")

def test_close_conn_in_window_func_value(self):
# gh-145040: closing connection in window function value() callback.
con = sqlite.connect(":memory:", autocommit=True)
con.execute("CREATE TABLE t(x INTEGER)")
con.executemany("INSERT INTO t VALUES(?)",
[(i,) for i in range(20)])

class CloseConnWindow:
def step(self, value):
pass
def finalize(self):
return 0
def value(self):
con.close()
return 0
def inverse(self, value):
pass

con.create_window_function("evil_win", 1, CloseConnWindow)
msg = "from within a callback"
with self.assertRaisesRegex(sqlite.ProgrammingError, msg):
cursor = con.execute(
"SELECT evil_win(x) OVER "
"(ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) FROM t"
)
list(cursor)


class AggregateTests(unittest.TestCase):
def setUp(self):
Expand Down Expand Up @@ -723,6 +750,148 @@ def test_agg_keyword_args(self):
'takes exactly 3 positional arguments'):
self.con.create_aggregate("test", 1, aggregate_class=AggrText)

def test_aggr_close_conn_in_step(self):
# Connection.close() in an aggregate step callback must not crash.
con = sqlite.connect(":memory:", autocommit=True)
cur = con.cursor()
cur.execute("CREATE TABLE t(x INTEGER)")
for i in range(50):
cur.execute("INSERT INTO t VALUES (?)", (i,))

class CloseConnAgg:
def __init__(self):
self.total = 0

def step(self, value):
self.total += value
con.close()

def finalize(self):
return self.total

con.create_aggregate("agg_close", 1, CloseConnAgg)
msg = "from within a callback"
with self.assertRaisesRegex(sqlite.ProgrammingError, msg):
con.execute("SELECT agg_close(x) FROM t")

def test_close_conn_in_nested_callback(self):
# gh-145040: close() must be prevented even in nested callbacks.
con = sqlite.connect(":memory:", autocommit=True)
con.execute("CREATE TABLE t(x INTEGER)")
for i in range(5):
con.execute("INSERT INTO t VALUES(?)", (i,))

def outer_func(x):
con.close()
return x

def inner_func(x):
return x * 10

con.create_function("outer_func", 1, outer_func)
con.create_function("inner_func", 1, inner_func)
msg = "from within a callback"
with self.assertRaisesRegex(sqlite.ProgrammingError, msg):
con.execute("SELECT outer_func(inner_func(x)) FROM t")
# Connection must still be usable after the failed close attempt.
self.assertEqual(con.execute("SELECT 1").fetchone(), (1,))
con.close()

def test_close_conn_in_nested_callback_caught(self):
# gh-145040: close attempt must propagate even if the exception
# is caught inside the callback and a nested execute consumes
# the flag.
con = sqlite.connect(":memory:", autocommit=True)
con.execute("CREATE TABLE t(x INTEGER)")
con.execute("INSERT INTO t VALUES(1)")

def swallow_close(x):
try:
con.close()
except sqlite.ProgrammingError:
pass
try:
con.execute("SELECT 1")
except sqlite.ProgrammingError:
pass
return x

con.create_function("swallow_close", 1, swallow_close)
msg = "from within a callback"
with self.assertRaisesRegex(sqlite.ProgrammingError, msg):
con.execute("SELECT swallow_close(x) FROM t")
# Connection must still be usable.
self.assertEqual(con.execute("SELECT 1").fetchone(), (1,))
con.close()

def test_close_conn_in_udf_during_executemany(self):
# gh-145040: closing connection in UDF during executemany.
con = sqlite.connect(":memory:", autocommit=True)
con.execute("CREATE TABLE t(x)")

def close_conn(x):
con.close()
return x

con.create_function("close_conn", 1, close_conn)
msg = "from within a callback"
with self.assertRaisesRegex(sqlite.ProgrammingError, msg):
con.executemany("INSERT INTO t VALUES(close_conn(?))",
[(i,) for i in range(10)])

def test_close_conn_in_progress_handler_during_iternext(self):
# gh-145040: closing connection in progress handler during iteration.
con = sqlite.connect(":memory:", autocommit=True)
con.execute("CREATE TABLE t(x)")
con.executemany("INSERT INTO t VALUES(?)",
[(i,) for i in range(100)])

count = 0
def close_progress():
nonlocal count
count += 1
if count >= 5:
con.close()
return 1
return 0

cursor = con.execute("SELECT * FROM t")
con.set_progress_handler(close_progress, 1)
msg = "from within a callback"
import test.support
with test.support.catch_unraisable_exception():
with self.assertRaisesRegex(sqlite.ProgrammingError, msg):
for row in cursor:
pass
del cursor
gc_collect()

def test_close_conn_in_collation_callback(self):
# gh-145040: closing connection in collation callback.
con = sqlite.connect(":memory:", autocommit=True)
con.execute("CREATE TABLE t(x TEXT)")
con.executemany("INSERT INTO t VALUES(?)",
[(f"item_{i}",) for i in range(50)])

count = 0
def evil_collation(a, b):
nonlocal count
count += 1
if count == 10:
con.close()
if a < b:
return -1
elif a > b:
return 1
return 0

con.create_collation("evil_coll", evil_collation)
msg = "from within a callback"
with self.assertRaisesRegex(sqlite.ProgrammingError, msg):
con.execute(
"SELECT * FROM t ORDER BY x COLLATE evil_coll"
)


class AuthorizerTests(unittest.TestCase):
@staticmethod
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
Fixed a crash in the :mod:`sqlite3` module caused by closing the database
connection from within a callback function invoked during
``sqlite3_step()`` (e.g., an aggregate ``step``, a user-defined function
via :meth:`~sqlite3.Connection.create_function`, a progress handler, or a
collation callback). Raise :exc:`~sqlite3.ProgrammingError` instead of
crashing.
12 changes: 12 additions & 0 deletions Modules/_sqlite/connection.c
Original file line number Diff line number Diff line change
Expand Up @@ -302,6 +302,8 @@ pysqlite_connection_init_impl(pysqlite_Connection *self, PyObject *database,
self->blobs = blobs;
self->row_factory = Py_NewRef(Py_None);
self->text_factory = Py_NewRef(&PyUnicode_Type);
self->in_callback = 0;
self->close_attempted_in_callback = 0;
self->trace_ctx = NULL;
self->progress_ctx = NULL;
self->authorizer_ctx = NULL;
Expand Down Expand Up @@ -655,6 +657,16 @@ pysqlite_connection_close_impl(pysqlite_Connection *self)
return NULL;
}

if (self->in_callback > 0) {
self->close_attempted_in_callback = 1;
PyTypeObject *tp = Py_TYPE(self);
pysqlite_state *state = pysqlite_get_state_by_type(tp);
PyErr_SetString(state->ProgrammingError,
"Cannot close the database connection "
"from within a callback function.");
return NULL;
}

pysqlite_close_all_blobs(self);
Py_CLEAR(self->statement_cache);
if (connection_close(self) < 0) {
Expand Down
9 changes: 9 additions & 0 deletions Modules/_sqlite/connection.h
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,15 @@ typedef struct

int initialized;

/* set to 1 while a SQLite callback (UDF, aggregate, progress handler,
* etc.) is executing; used to prevent closing the connection from
* within a callback, which is illegal per the SQLite C API docs */
int in_callback;

/* set to 1 when close() is attempted during a callback; checked after
* stmt_step() returns to raise the appropriate ProgrammingError */
int close_attempted_in_callback;

/* thread identification of the thread the connection was created in */
unsigned long thread_ident;

Expand Down
26 changes: 26 additions & 0 deletions Modules/_sqlite/cursor.c
Original file line number Diff line number Diff line change
Expand Up @@ -905,7 +905,19 @@ _pysqlite_query_execute(pysqlite_Cursor* self, int multiple, PyObject* operation
goto error;
}

self->connection->in_callback++;
rc = stmt_step(self->statement->st);
self->connection->in_callback--;
if (self->connection->close_attempted_in_callback
&& self->connection->in_callback == 0)
{
self->connection->close_attempted_in_callback = 0;
PyErr_Clear();
PyErr_SetString(state->ProgrammingError,
"Cannot close the database connection "
"from within a callback function.");
goto error;
}
if (rc != SQLITE_DONE && rc != SQLITE_ROW) {
if (PyErr_Occurred()) {
/* there was an error that occurred in a user-defined callback */
Expand Down Expand Up @@ -1155,7 +1167,21 @@ pysqlite_cursor_iternext(PyObject *op)
if (row == NULL) {
return NULL;
}
self->connection->in_callback++;
int rc = stmt_step(stmt);
self->connection->in_callback--;
if (self->connection->close_attempted_in_callback
&& self->connection->in_callback == 0)
{
self->connection->close_attempted_in_callback = 0;
Py_DECREF(row);
Py_CLEAR(self->statement);
PyErr_Clear();
PyErr_SetString(self->connection->state->ProgrammingError,
"Cannot close the database connection "
"from within a callback function.");
return NULL;
}
if (rc == SQLITE_DONE) {
if (self->statement->is_dml) {
self->rowcount = (long)sqlite3_changes(self->connection->db);
Expand Down
Loading