diff --git a/lib/core/mcp/mcp_access_store.dart b/lib/core/mcp/mcp_access_store.dart index b737080d..d2d06afb 100644 --- a/lib/core/mcp/mcp_access_store.dart +++ b/lib/core/mcp/mcp_access_store.dart @@ -36,8 +36,18 @@ class McpAccessStore implements McpAccessSettings { return {}; } + /// Two quick toggles each read the list, change it and write it back; run + /// one at a time so the second sees the first (it used to overwrite it). + Future _writes = Future.value(); + @override - Future setReadable(int connectionId, bool readable) async { + Future setReadable(int connectionId, bool readable) { + final run = _writes.then((_) => _setReadable(connectionId, readable)); + _writes = run.then((_) {}, onError: (Object _) {}); + return run; + } + + Future _setReadable(int connectionId, bool readable) async { final ids = await readableIds(); if (readable) { ids.add(connectionId); diff --git a/lib/core/mcp/mcp_mongo_service.dart b/lib/core/mcp/mcp_mongo_service.dart index f2937729..3d501ad4 100644 --- a/lib/core/mcp/mcp_mongo_service.dart +++ b/lib/core/mcp/mcp_mongo_service.dart @@ -58,6 +58,11 @@ abstract final class McpMongoGuard { } } else if (node is List) { for (final item in node) { + // A huge `$in` list is as heavy as as many keys. + if (++keys > maxKeys) { + return McpMongoRefusal('mongo_too_large', + '$what has more than $maxKeys fields or list items.'); + } final inner = walk(item, depth + 1); if (inner != null) return inner; } diff --git a/lib/core/mcp/mcp_server_controller.dart b/lib/core/mcp/mcp_server_controller.dart index 1f9b0c08..7d9fa51a 100644 --- a/lib/core/mcp/mcp_server_controller.dart +++ b/lib/core/mcp/mcp_server_controller.dart @@ -134,7 +134,18 @@ class McpServerController { }, ); host.changes.listen((_) => _publish(host)); - await host.start(); + try { + await host.start(); + } catch (_) { + // It may have bound the socket before it failed; do not leave it + // listening with a token nobody holds. + try { + await host.stop(); + } catch (e, st) { + logIgnored('McpServerController.start cleanup', e, st); + } + rethrow; + } _host = host; _publish(host); } @@ -154,7 +165,11 @@ class McpServerController { } void _onCall(McpCallRecord r) { - if (_logActivity) unawaited(_persist(r)); + if (_logActivity) { + unawaited(_persist(r).catchError((Object e, StackTrace st) { + logIgnored('McpServerController.persist activity', e, st); + })); + } _calls.add(r); } diff --git a/lib/core/mcp/mcp_socket_host.dart b/lib/core/mcp/mcp_socket_host.dart index e25cc3ef..6fc30a09 100644 --- a/lib/core/mcp/mcp_socket_host.dart +++ b/lib/core/mcp/mcp_socket_host.dart @@ -18,6 +18,8 @@ class McpSocketHost { this.handshakeTimeout = const Duration(seconds: 5), this.maxUnauthenticated = 16, this.maxHandshakeBytes = 4096, + this.maxLineBytes = 4 * 1024 * 1024, + this.maxSessions = 16, }); final File endpointFile; @@ -31,6 +33,13 @@ class McpSocketHost { /// Bytes accepted before the token line is complete (a token is ~44). final int maxHandshakeBytes; + /// Bytes of one line from an authenticated client; a longer line (or one + /// that never ends) closes the connection instead of growing in memory. + final int maxLineBytes; + + /// Authenticated connections that may be open at once. + final int maxSessions; + ServerSocket? _server; String? _token; final _sockets = {}; @@ -99,6 +108,7 @@ class McpSocketHost { var authed = false; var waiting = true; var handshakeBytes = 0; + var lineBytes = 0; void stopWaiting() { if (!waiting) return; waiting = false; @@ -124,6 +134,16 @@ class McpSocketHost { _reject(socket, 'Handshake too long.'); return const []; } + if (authed) { + final newline = chunk.lastIndexOf(0x0A); + lineBytes = + newline >= 0 ? chunk.length - newline - 1 : lineBytes + chunk.length; + if (lineBytes > maxLineBytes) { + sub.cancel(); + _closed(socket, incoming, outgoing, authed); + return const []; + } + } return chunk; }) .transform(utf8.decoder) @@ -142,7 +162,15 @@ class McpSocketHost { _reject(socket, 'Invalid Querya MCP token. Restart the MCP client.'); return; } + if (_sessions >= maxSessions) { + timer.cancel(); + sub.cancel(); + stopWaiting(); + _reject(socket, 'Too many Querya MCP sessions.'); + return; + } authed = true; + lineBytes = 0; stopWaiting(); timer.cancel(); _sessions++; diff --git a/test/core/mcp/mcp_mongo_service_test.dart b/test/core/mcp/mcp_mongo_service_test.dart index 3284439a..b432f7b6 100644 --- a/test/core/mcp/mcp_mongo_service_test.dart +++ b/test/core/mcp/mcp_mongo_service_test.dart @@ -181,6 +181,18 @@ void main() { 'mongo_too_large'); }); + test('a huge list in a filter is refused (#1776)', () { + final filter = { + 'id': {r'$in': List.generate(McpMongoGuard.maxKeys + 1, (i) => i)}, + }; + expect(McpMongoGuard.refusal(filter, what: 'filter')?.rule, + 'mongo_too_large'); + final small = { + 'id': {r'$in': List.generate(20, (i) => i)}, + }; + expect(McpMongoGuard.refusal(small, what: 'filter'), isNull); + }); + test('a value that is not a document is refused', () { expect(McpMongoGuard.refusal('abc', what: 'filter')?.rule, 'mongo_not_a_document'); diff --git a/test/core/mcp/mcp_socket_host_test.dart b/test/core/mcp/mcp_socket_host_test.dart index 94d9a834..2595518d 100644 --- a/test/core/mcp/mcp_socket_host_test.dart +++ b/test/core/mcp/mcp_socket_host_test.dart @@ -3,19 +3,27 @@ import 'dart:io'; import 'package:flutter_test/flutter_test.dart'; import 'package:path/path.dart' as p; +import 'package:querya_mcp_bridge/querya_mcp_bridge.dart'; import 'package:querya_desktop/core/mcp/mcp_socket_host.dart'; void main() { late Directory tempDir; late McpSocketHost host; - Future startHost({int maxUnauthenticated = 16}) async { + Future startHost({ + int maxUnauthenticated = 16, + int maxSessions = 16, + int maxLineBytes = 4 * 1024 * 1024, + int maxHandshakeBytes = 64, + }) async { host = McpSocketHost( endpointFile: File(p.join(tempDir.path, 'mcp.json')), version: 'test', onSession: (_) async {}, maxUnauthenticated: maxUnauthenticated, - maxHandshakeBytes: 64, + maxHandshakeBytes: maxHandshakeBytes, + maxSessions: maxSessions, + maxLineBytes: maxLineBytes, ); await host.start(); } @@ -54,4 +62,43 @@ void main() { first.destroy(); second.destroy(); }); + + Future connectAuthed() async { + final endpoint = await McpEndpoint.read(File(p.join(tempDir.path, 'mcp.json'))); + final socket = + await Socket.connect(InternetAddress.loopbackIPv4, host.port!); + socket.write('${endpoint!.token}\n'); + await socket.flush(); + return socket; + } + + test('sessions beyond the cap are turned away (#1776)', () async { + await startHost(maxSessions: 1, maxHandshakeBytes: 4096); + final first = await connectAuthed(); + await Future.delayed(const Duration(milliseconds: 200)); + expect(host.sessionCount, 1); + + final second = await connectAuthed(); + final reply = await utf8.decodeStream(second).timeout(const Duration(seconds: 3)); + + expect(reply, contains('Too many')); + expect(host.sessionCount, 1); + first.destroy(); + second.destroy(); + }); + + test('an authenticated line without a newline is cut off (#1776)', () async { + await startHost(maxLineBytes: 100, maxHandshakeBytes: 4096); + final socket = await connectAuthed(); + await Future.delayed(const Duration(milliseconds: 200)); + expect(host.sessionCount, 1); + + socket.add(List.filled(300, 0x41)); + await socket.flush(); + + await socket.drain().timeout(const Duration(seconds: 3)); + await Future.delayed(const Duration(milliseconds: 200)); + expect(host.sessionCount, 0); + socket.destroy(); + }); }