Disable extension loading once --load-extension extensions are loaded

Refs GHSA-2mvv-ffvc-q5p6

Co-authored-by: Alex Garcia <15178711+asg017@users.noreply.github.com>

datasette/app.py

@@ -1532,15 +1532,28 @@ class Datasette:
         conn.row_factory = sqlite3.Row
         conn.text_factory = lambda x: str(x, "utf-8", "replace")
         if self.sqlite_extensions and database != INTERNAL_DB_NAME:
+            # Extension loading is only enabled for as long as it takes to
+            # load the configured extensions. Leaving it enabled would let
+            # anyone who can execute SQL call load_extension() themselves.
             conn.enable_load_extension(True)
-            for extension in self.sqlite_extensions:
-                # "extension" is either a string path to the extension
-                # or a 2-item tuple that specifies which entrypoint to load.
-                if isinstance(extension, tuple):
-                    path, entrypoint = extension
-                    conn.execute("SELECT load_extension(?, ?)", [path, entrypoint])
-                else:
-                    conn.execute("SELECT load_extension(?)", [extension])
+            try:
+                for extension in self.sqlite_extensions:
+                    # "extension" is either a string path to the extension
+                    # or a 2-item tuple that specifies which entrypoint to load.
+                    if isinstance(extension, tuple):
+                        path, entrypoint = extension
+                        if sys.version_info >= (3, 12):
+                            conn.load_extension(path, entrypoint=entrypoint)
+                        else:
+                            # Connection.load_extension() only gained the
+                            # entrypoint argument in Python 3.12
+                            conn.execute(
+                                "SELECT load_extension(?, ?)", [path, entrypoint]
+                            )
+                    else:
+                        conn.load_extension(extension)
+            finally:
+                conn.enable_load_extension(False)
         if self.setting("cache_size_kb"):
             conn.execute(f"PRAGMA cache_size=-{self.setting('cache_size_kb')}")
         # pylint: disable=no-member

tests/test_load_extensions.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+from unittest import mock
 
 import pytest
 
@@ -20,6 +21,29 @@ def has_compiled_ext():
     return False
 
 
+@pytest.mark.parametrize("load_fails", (False, True))
+def test_load_extension_is_disabled(load_fails):
+    ds = Datasette(sqlite_extensions=[COMPILED_EXTENSION_PATH])
+    connection = mock.Mock()
+    if load_fails:
+        connection.load_extension.side_effect = RuntimeError
+
+    if load_fails:
+        with pytest.raises(RuntimeError):
+            ds._prepare_connection(connection, "data")
+    else:
+        ds._prepare_connection(connection, "data")
+
+    # Extensions are loaded using the Python API, never via SQL
+    assert connection.load_extension.mock_calls == [
+        mock.call(COMPILED_EXTENSION_PATH),
+    ]
+    assert connection.enable_load_extension.mock_calls == [
+        mock.call(True),
+        mock.call(False),
+    ]
+
+
 @pytest.mark.asyncio
 @pytest.mark.skipif(not has_compiled_ext(), reason="Requires compiled ext.c")
 async def test_load_extension_default_entrypoint():
@@ -64,3 +88,20 @@ async def test_load_extension_multiple_entrypoints():
     response = await ds.client.get("/_memory/-/query.json?_shape=arrays&sql=select+c()")
     assert response.status_code == 200
     assert response.json()["rows"][0][0] == "c"
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not has_compiled_ext(), reason="Requires compiled ext.c")
+async def test_sql_cannot_load_additional_extension():
+    ds = Datasette(sqlite_extensions=[COMPILED_EXTENSION_PATH])
+
+    response = await ds.client.get(
+        "/_memory/-/query.json",
+        params={
+            "sql": "select load_extension(:path, :entrypoint)",
+            "path": COMPILED_EXTENSION_PATH,
+            "entrypoint": "sqlite3_ext_b_init",
+        },
+    )
+    assert response.status_code == 400
+    assert response.json()["error"] == "not authorized"