diff --git a/backend/agent_vm/main.py b/backend/agent_vm/main.py index 5e0a8062068..4550b8dbf61 100644 --- a/backend/agent_vm/main.py +++ b/backend/agent_vm/main.py @@ -480,6 +480,16 @@ async def execute_backend_tool(name: str, params: dict[str, Any]) -> str: return result.get("result") or json.dumps(result, default=str) +def read_only_sql_authorizer( + action: int, arg1: str | None, arg2: str | None, dbname: str | None, source: str | None +) -> int: + if action in (sqlite3.SQLITE_SELECT, sqlite3.SQLITE_READ, sqlite3.SQLITE_FUNCTION): + return sqlite3.SQLITE_OK + if action == sqlite3.SQLITE_PRAGMA and arg1 is not None and arg1.casefold() == "data_version" and arg2 is None: + return sqlite3.SQLITE_OK + return sqlite3.SQLITE_DENY + + def execute_sql(query: str) -> str: if runtime.db is None: return json.dumps({"error": "Database not loaded. Upload omi.db first."}) @@ -496,8 +506,12 @@ def execute_sql(query: str) -> str: query = query.rstrip().rstrip(";") + " LIMIT 200" try: with runtime.lock: - cursor = runtime.db.execute(query) - rows = [dict(row) for row in cursor.fetchall()] + try: + runtime.db.set_authorizer(read_only_sql_authorizer) + cursor = runtime.db.execute(query) + rows = [dict(row) for row in cursor.fetchall()] + finally: + runtime.db.set_authorizer(None) return json.dumps({"rows": rows, "count": len(rows)}, default=str) except sqlite3.Error as exc: return json.dumps({"error": str(exc)}) diff --git a/backend/tests/unit/test_agent_vm_protocol.py b/backend/tests/unit/test_agent_vm_protocol.py index 491b7a9b0d8..71178ad3ab6 100644 --- a/backend/tests/unit/test_agent_vm_protocol.py +++ b/backend/tests/unit/test_agent_vm_protocol.py @@ -687,6 +687,61 @@ def test_execute_sql_serializes_sqlite_rows(tmp_path: Path) -> None: } +def test_execute_sql_allows_fts5_reads(tmp_path: Path) -> None: + _, module = load_app(tmp_path) + connection = sqlite3.connect(module.runtime.db_path) + connection.execute("CREATE VIRTUAL TABLE documents USING fts5(title, body)") + connection.executemany( + "INSERT INTO documents (title, body) VALUES (?, ?)", + [("one", "hello world"), ("two", "other text")], + ) + connection.commit() + connection.close() + assert module.runtime.open_database() + + assert json.loads(module.execute_sql("SELECT rowid, title FROM documents WHERE documents MATCH 'hello'")) == { + "rows": [{"rowid": 1, "title": "one"}], + "count": 1, + } + + +def test_execute_sql_clears_authorizer_after_error(tmp_path: Path) -> None: + _, module = load_app(tmp_path) + connection = sqlite3.connect(module.runtime.db_path) + connection.execute("CREATE TABLE screenshots (id TEXT)") + connection.commit() + connection.close() + assert module.runtime.open_database() + + result = json.loads(module.execute_sql("SELECT missing FROM screenshots")) + + assert result["error"] + module.runtime.db.execute("CREATE TABLE after_authorizer_cleanup (value TEXT)") + + +@pytest.mark.parametrize( + "query", + [ + "DELETE FROM screenshots", + "UPDATE screenshots SET id = 'changed'", + "SELECT 1; DROP TABLE screenshots", + ], +) +def test_execute_sql_denies_destructive_queries(tmp_path: Path, query: str) -> None: + _, module = load_app(tmp_path) + connection = sqlite3.connect(module.runtime.db_path) + connection.execute("CREATE TABLE screenshots (id TEXT)") + connection.execute("INSERT INTO screenshots VALUES ('one')") + connection.commit() + connection.close() + assert module.runtime.open_database() + + result = json.loads(module.execute_sql(query)) + + assert result["error"] + assert [tuple(row) for row in module.runtime.db.execute("SELECT id FROM screenshots").fetchall()] == [("one",)] + + def test_sync_groups_rows_by_present_columns(tmp_path: Path) -> None: app, module = load_app(tmp_path) connection = sqlite3.connect(module.runtime.db_path)