diff --git a/lib/core/database/mysql_connection.dart b/lib/core/database/mysql_connection.dart index fd467dd7..6cf9f705 100644 --- a/lib/core/database/mysql_connection.dart +++ b/lib/core/database/mysql_connection.dart @@ -270,7 +270,12 @@ class MysqlConnection { if (!isConnected || _conn == null) { throw StateError('Not connected to MySQL'); } - return _conn!.execute(sql, params, iterable); + try { + return await _conn!.execute(sql, params, iterable); + } on TimeoutException { + unawaited(forceClose()); + rethrow; + } } /// Runs [execute] with an application-level [timeout] (driver limits still apply). @@ -282,7 +287,12 @@ class MysqlConnection { }) async { final f = execute(sql, params, iterable); if (timeout == null) return f; - return f.timeout(timeout); + try { + return await f.timeout(timeout); + } on TimeoutException { + unawaited(forceClose()); + rethrow; + } } /// Lists user-visible databases (excludes typical system schemas). diff --git a/lib/core/database/postgres_connection.dart b/lib/core/database/postgres_connection.dart index 2ba4162f..c6929e16 100644 --- a/lib/core/database/postgres_connection.dart +++ b/lib/core/database/postgres_connection.dart @@ -1,3 +1,4 @@ +import 'dart:async'; import 'dart:io' show SecurityContext; import 'package:flutter/foundation.dart'; @@ -252,7 +253,27 @@ class PostgresConnection { if (!isConnected || _conn == null) { throw StateError('Not connected to PostgreSQL'); } - return _conn!.execute(sql, timeout: timeout); + try { + return await _conn!.execute(sql, timeout: timeout); + } on TimeoutException { + unawaited(forceClose()); + rethrow; + } + } + + /// Runs [execute] with an application-level [timeout] (in addition to driver timeout). + Future executeWithTimeout( + String sql, { + Duration? timeout, + }) async { + final f = execute(sql, timeout: timeout); + if (timeout == null) return f; + try { + return await f.timeout(timeout); + } on TimeoutException { + unawaited(forceClose()); + rethrow; + } } /// Whether the session has an open transaction (PostgreSQL 13+). diff --git a/lib/core/database/sqlite_connection.dart b/lib/core/database/sqlite_connection.dart index 3663fce8..d1e9cbae 100644 --- a/lib/core/database/sqlite_connection.dart +++ b/lib/core/database/sqlite_connection.dart @@ -114,11 +114,32 @@ class SqliteConnection { throw StateError('Database connection is read-only'); } - if (isReadOnlyQuery || hasReturning) { - return await _db!.rawQuery(sql, arguments); - } else { - await _db!.execute(sql, arguments); - return []; + try { + if (isReadOnlyQuery || hasReturning) { + return await _db!.rawQuery(sql, arguments); + } else { + await _db!.execute(sql, arguments); + return []; + } + } on TimeoutException { + unawaited(forceClose()); + rethrow; + } + } + + /// Runs [execute] with an application-level [timeout]. + Future>> executeWithTimeout( + String sql, { + Duration? timeout, + List? arguments, + }) async { + final f = execute(sql, arguments); + if (timeout == null) return f; + try { + return await f.timeout(timeout); + } on TimeoutException { + unawaited(forceClose()); + rethrow; } } diff --git a/lib/features/mysql/mysql_sql_workspace.dart b/lib/features/mysql/mysql_sql_workspace.dart index 8f518de0..a63c7d71 100644 --- a/lib/features/mysql/mysql_sql_workspace.dart +++ b/lib/features/mysql/mysql_sql_workspace.dart @@ -244,6 +244,7 @@ class _MysqlSqlWorkspaceState extends material.State { ); } } on TimeoutException catch (e) { + unawaited(_lease?.connection.forceClose()); if (mounted) { setState(() { _error = e.toString(); diff --git a/lib/features/postgresql/postgres_sql_workspace.dart b/lib/features/postgresql/postgres_sql_workspace.dart index 26e579f1..34204167 100644 --- a/lib/features/postgresql/postgres_sql_workspace.dart +++ b/lib/features/postgresql/postgres_sql_workspace.dart @@ -248,6 +248,14 @@ class _PostgresSqlWorkspaceState extends material.State { _statusLine = 'OK: $cmd'; _running = false; }); + } on TimeoutException catch (e) { + unawaited(_lease?.connection.forceClose()); + if (mounted) { + setState(() { + _error = 'Query timed out: ${e.message ?? e}'; + _running = false; + }); + } } on pg.ServerException catch (e) { if (mounted) { setState(() { @@ -375,6 +383,14 @@ class _PostgresSqlWorkspaceState extends material.State { ), ); } + } on TimeoutException catch (e) { + unawaited(_lease?.connection.forceClose()); + if (mounted) { + setState(() { + _error = 'Query timed out: ${e.message ?? e}'; + _running = false; + }); + } } on pg.ServerException catch (e) { if (mounted) { setState(() { diff --git a/lib/features/sqlite/sqlite_sql_workspace.dart b/lib/features/sqlite/sqlite_sql_workspace.dart index c34c62fa..0956e858 100644 --- a/lib/features/sqlite/sqlite_sql_workspace.dart +++ b/lib/features/sqlite/sqlite_sql_workspace.dart @@ -200,6 +200,14 @@ class _SqliteSqlWorkspaceState extends material.State { ), ); } + } on TimeoutException catch (e) { + unawaited(_lease?.connection.forceClose()); + if (mounted) { + setState(() { + _error = 'Query timed out: ${e.message ?? e}'; + _running = false; + }); + } } catch (e) { if (mounted) { setState(() { diff --git a/test/core/database/connection_timeout_protocol_test.dart b/test/core/database/connection_timeout_protocol_test.dart new file mode 100644 index 00000000..6a0a5e79 --- /dev/null +++ b/test/core/database/connection_timeout_protocol_test.dart @@ -0,0 +1,157 @@ +import 'dart:async'; +import 'package:flutter_test/flutter_test.dart'; +import 'package:mysql_client/mysql_client.dart'; +import 'package:querya_desktop/core/database/mysql_connection.dart'; +import 'package:querya_desktop/core/database/postgres_connection.dart'; +import 'package:querya_desktop/core/database/sqlite_connection.dart'; +import 'package:postgres/postgres.dart' as pg; + +class FakeSlowMysqlConnection extends MysqlConnection { + FakeSlowMysqlConnection({super.id = 1}) + : super( + name: 'fake_slow_mysql', + host: 'localhost', + port: 3306, + database: 'testdb', + ); + + bool _connected = true; + int forceCloseCount = 0; + + @override + bool get isConnected => _connected; + + @override + Future execute( + String sql, [ + Map? params, + bool iterable = false, + ]) async { + await Future.delayed(const Duration(seconds: 10)); + throw Exception('should not reach here'); + } + + @override + Future forceClose() async { + forceCloseCount++; + _connected = false; + } +} + +class FakeSlowPostgresConnection extends PostgresConnection { + FakeSlowPostgresConnection({super.id = 1}) + : super( + name: 'fake_slow_pg', + host: 'localhost', + port: 5432, + database: 'postgres', + ); + + bool _connected = true; + int forceCloseCount = 0; + + @override + bool get isConnected => _connected; + + @override + Future execute(String sql, {Duration? timeout}) async { + await Future.delayed(const Duration(seconds: 10)); + throw Exception('should not reach here'); + } + + @override + Future forceClose() async { + forceCloseCount++; + _connected = false; + } +} + +class FakeSlowSqliteConnection extends SqliteConnection { + FakeSlowSqliteConnection({super.id = 1}) + : super( + name: 'fake_slow_sqlite', + path: ':memory:', + ); + + bool _connected = true; + int forceCloseCount = 0; + + @override + bool get isConnected => _connected; + + @override + Future>> execute( + String sql, [ + List? arguments, + ]) async { + await Future.delayed(const Duration(seconds: 10)); + throw Exception('should not reach here'); + } + + @override + Future forceClose() async { + forceCloseCount++; + _connected = false; + } +} + +void main() { + group('Connection Timeout Protocol Protection', () { + test('MysqlConnection.executeWithTimeout force-closes on timeout', () async { + final conn = FakeSlowMysqlConnection(); + expect(conn.isConnected, isTrue); + + try { + await conn.executeWithTimeout( + 'SELECT sleep(100)', + timeout: const Duration(milliseconds: 20), + ); + fail('Should have thrown TimeoutException'); + } on TimeoutException { + // Expected + } + + await Future.delayed(const Duration(milliseconds: 10)); + expect(conn.forceCloseCount, 1); + expect(conn.isConnected, isFalse); + }); + + test('PostgresConnection.executeWithTimeout force-closes on TimeoutException', () async { + final conn = FakeSlowPostgresConnection(); + expect(conn.isConnected, isTrue); + + try { + await conn.executeWithTimeout( + 'SELECT pg_sleep(100)', + timeout: const Duration(milliseconds: 20), + ); + fail('Should have thrown TimeoutException'); + } on TimeoutException { + // Expected + } + + await Future.delayed(const Duration(milliseconds: 10)); + expect(conn.forceCloseCount, 1); + expect(conn.isConnected, isFalse); + }); + + test('SqliteConnection.executeWithTimeout force-closes on TimeoutException', () async { + final conn = FakeSlowSqliteConnection(); + expect(conn.isConnected, isTrue); + + try { + await conn.executeWithTimeout( + 'SELECT 1', + timeout: const Duration(milliseconds: 20), + ); + fail('Should have thrown TimeoutException'); + } on TimeoutException { + // Expected + } + + await Future.delayed(const Duration(milliseconds: 10)); + expect(conn.forceCloseCount, 1); + expect(conn.isConnected, isFalse); + }); + }); +}