From f016a363c15dd7f659531438d7721a3663aa6e44 Mon Sep 17 00:00:00 2001 From: Widthdom Date: Wed, 3 Jun 2026 04:49:20 +0900 Subject: [PATCH] Limit MCP response serialization (#2860) --- changelog.d/unreleased/2860.security.md | 17 ++++ src/CodeIndex/Mcp/McpServer.cs | 121 ++++++++++++++++++++++-- src/CodeIndex/Mcp/McpToolHandlers.cs | 11 ++- tests/CodeIndex.Tests/McpServerTests.cs | 34 +++++++ 4 files changed, 173 insertions(+), 10 deletions(-) create mode 100644 changelog.d/unreleased/2860.security.md diff --git a/changelog.d/unreleased/2860.security.md b/changelog.d/unreleased/2860.security.md new file mode 100644 index 0000000000..e2bef6ae6c --- /dev/null +++ b/changelog.d/unreleased/2860.security.md @@ -0,0 +1,17 @@ +--- +category: security +issues: + - 2860 +affected: + - src/CodeIndex/Mcp/McpServer.cs + - src/CodeIndex/Mcp/McpToolHandlers.cs + - tests/CodeIndex.Tests/McpServerTests.cs +--- + +## English + +- **MCP response byte limits now stop JSON serialization before full string materialization (#2860)** — oversized MCP responses are measured with a bounded UTF-8 JSON writer so the server can reject them as soon as the configured byte cap is crossed, with response-too-large errors marking whether the reported byte count is exact. + +## 日本語 + +- **MCP response byte limit が JSON 全体の文字列化前に serialization を停止するようになりました (#2860)** — oversized な MCP response は上限付き UTF-8 JSON writer で測定し、設定された byte cap を超えた時点で拒否し、response-too-large error では報告 byte count が exact かどうかも示すようにしました。 diff --git a/src/CodeIndex/Mcp/McpServer.cs b/src/CodeIndex/Mcp/McpServer.cs index fdadbee598..53efee1e2c 100644 --- a/src/CodeIndex/Mcp/McpServer.cs +++ b/src/CodeIndex/Mcp/McpServer.cs @@ -32,6 +32,7 @@ public partial class McpServer : IDisposable private readonly string _version; private readonly JsonSerializerOptions _jsonOptions; private readonly Func _serializeResponse; + private readonly bool _usesDefaultResponseSerializer; private readonly IMcpAuthenticator _authenticator; private readonly McpToolFilter _toolFilter; private readonly TimeProvider _timeProvider; @@ -262,6 +263,7 @@ internal McpServer(string dbPath, string version, bool dbPathExplicit, Func node.ToJsonString(_jsonOptions)); _authenticator = authenticator ?? LocalStdioAuthenticator.Instance; _toolFilter = toolFilter ?? McpToolFilter.FromEnvironment(); @@ -1042,9 +1044,17 @@ private string SerializeResponseOrFallback(JsonNode response, bool hasId, JsonNo { try { + var responseLimit = GetMaxResponseBytes(); + if (_usesDefaultResponseSerializer) + { + if (!TrySerializeJsonNodeWithinByteLimit(response, _jsonOptions, responseLimit, captureSerialized: true, out var boundedSerialized, out var boundedResponseBytes)) + return CreateResponseTooLargeError(hasId, id, boundedResponseBytes, responseLimit, actualBytesExact: false).ToJsonString(_jsonOptions); + + return boundedSerialized!; + } + var serialized = _serializeResponse(response); var responseBytes = Encoding.UTF8.GetByteCount(serialized); - var responseLimit = GetMaxResponseBytes(); if (responseBytes <= responseLimit) return serialized; @@ -2991,7 +3001,7 @@ private static JsonObject CreateCancelledResponse(JsonNode? id) /// Create a tool result response (MCP format). /// ツール結果レスポンスを作成(MCP形式)。 /// - private static JsonObject CreateToolResult(JsonNode? id, string text, JsonNode? structuredContent = null, string? mimeType = null) + private JsonObject CreateToolResult(JsonNode? id, string text, JsonNode? structuredContent = null, string? mimeType = null) { mimeType ??= structuredContent is null ? "text/plain" : "application/json"; var result = new JsonObject @@ -3009,15 +3019,113 @@ private static JsonObject CreateToolResult(JsonNode? id, string text, JsonNode? if (structuredContent != null) result["structuredContent"] = structuredContent; var response = CreateSuccessResponse(true, id, result); - var responseBytes = Encoding.UTF8.GetByteCount(response.ToJsonString()); var responseLimit = GetMaxResponseBytes(); - if (responseBytes <= responseLimit) + if (TryMeasureJsonUtf8BytesWithinLimit(response, _jsonOptions, responseLimit, out var responseBytes)) return response; - return CreateResponseTooLargeError(true, id, responseBytes, responseLimit); + return CreateResponseTooLargeError(true, id, responseBytes, responseLimit, actualBytesExact: false); + } + + internal bool TrySerializeJsonNodeWithinByteLimitForTests(JsonNode node, int maxBytes, out string? serialized, out int bytesWritten) + => TrySerializeJsonNodeWithinByteLimit(node, _jsonOptions, maxBytes, captureSerialized: true, out serialized, out bytesWritten); + + private static bool TryMeasureJsonUtf8BytesWithinLimit(JsonNode node, JsonSerializerOptions options, int maxBytes, out int bytesWritten) + => TrySerializeJsonNodeWithinByteLimit(node, options, maxBytes, captureSerialized: false, out _, out bytesWritten); + + private static bool TrySerializeJsonNodeWithinByteLimit(JsonNode node, JsonSerializerOptions options, int maxBytes, bool captureSerialized, out string? serialized, out int bytesWritten) + { + if (maxBytes < 0) + throw new ArgumentOutOfRangeException(nameof(maxBytes), maxBytes, "JSON byte limit must be non-negative."); + + serialized = null; + using var stream = new BoundedJsonUtf8Stream(maxBytes, captureSerialized); + var writerOptions = new JsonWriterOptions + { + Encoder = options.Encoder, + Indented = options.WriteIndented, + }; + + try + { + using var writer = new Utf8JsonWriter(stream, writerOptions); + node.WriteTo(writer, options); + writer.Flush(); + bytesWritten = stream.BytesWritten; + serialized = stream.GetCapturedString(); + return true; + } + catch (JsonResponseByteLimitExceededException ex) + { + bytesWritten = ex.BytesWritten; + return false; + } + } + + private sealed class JsonResponseByteLimitExceededException(int bytesWritten) : Exception + { + public int BytesWritten { get; } = bytesWritten; + } + + private sealed class BoundedJsonUtf8Stream(int maxBytes, bool captureSerialized) : Stream + { + private readonly MemoryStream? _buffer = captureSerialized ? new MemoryStream(Math.Min(Math.Max(maxBytes, 0), 16 * 1024)) : null; + + public int BytesWritten { get; private set; } + + public override bool CanRead => false; + public override bool CanSeek => false; + public override bool CanWrite => true; + public override long Length => throw new NotSupportedException(); + + public override long Position + { + get => throw new NotSupportedException(); + set => throw new NotSupportedException(); + } + + public string? GetCapturedString() + { + if (_buffer is null) + return null; + return Encoding.UTF8.GetString(_buffer.GetBuffer(), 0, (int)_buffer.Length); + } + + public override void Flush() + { + } + + public override int Read(byte[] buffer, int offset, int count) + => throw new NotSupportedException(); + + public override long Seek(long offset, SeekOrigin origin) + => throw new NotSupportedException(); + + public override void SetLength(long value) + => throw new NotSupportedException(); + + public override void Write(byte[] buffer, int offset, int count) + => Write(buffer.AsSpan(offset, count)); + + public override void Write(ReadOnlySpan buffer) + { + if (buffer.Length == 0) + return; + + var remaining = maxBytes - BytesWritten; + if (remaining < buffer.Length) + { + if (remaining > 0) + _buffer?.Write(buffer[..remaining]); + BytesWritten = maxBytes == int.MaxValue ? int.MaxValue : maxBytes + 1; + throw new JsonResponseByteLimitExceededException(BytesWritten); + } + + _buffer?.Write(buffer); + BytesWritten += buffer.Length; + } } - private static JsonObject CreateResponseTooLargeError(bool hasId, JsonNode? id, int responseBytes, int responseLimit) + private static JsonObject CreateResponseTooLargeError(bool hasId, JsonNode? id, int responseBytes, int responseLimit, bool actualBytesExact = true) { return CreateErrorResponse( hasId: hasId, @@ -3032,6 +3140,7 @@ private static JsonObject CreateResponseTooLargeError(bool hasId, JsonNode? id, ["reason"] = "response_too_large", ["limit_bytes"] = responseLimit, ["actual_bytes"] = responseBytes, + ["actual_bytes_exact"] = actualBytesExact, }); } diff --git a/src/CodeIndex/Mcp/McpToolHandlers.cs b/src/CodeIndex/Mcp/McpToolHandlers.cs index 57236f7b3d..dabb0bfbb4 100644 --- a/src/CodeIndex/Mcp/McpToolHandlers.cs +++ b/src/CodeIndex/Mcp/McpToolHandlers.cs @@ -2792,7 +2792,7 @@ string BuildSummary() } summary = BuildSummary(); - estimatedResponseBytes = EstimateJsonUtf8Bytes(CreateToolResult(id, summary, payload.DeepClone())); + estimatedResponseBytes = EstimateJsonUtf8Bytes(CreateToolResult(id, summary, payload.DeepClone()), responseByteLimit); if (estimatedResponseBytes <= responseByteLimit) break; if (resultsArray.Count > 0) @@ -2827,8 +2827,11 @@ private static int GetBatchQueryResponseByteLimit() MaxBatchQueryResponseByteLimit, "MCP batch_query response byte limit"); - private int EstimateJsonUtf8Bytes(JsonNode node) => - Encoding.UTF8.GetByteCount(node.ToJsonString(_jsonOptions)); + private int EstimateJsonUtf8Bytes(JsonNode node, int maxBytes = int.MaxValue) + { + _ = TryMeasureJsonUtf8BytesWithinLimit(node, _jsonOptions, maxBytes, out var bytesWritten); + return bytesWritten; + } private int EstimateBatchResponseBytes(JsonNode? id, string summary, int submittedCount, int successCount, int failureCount, string failureScope, int? cascadeStartedAtIndex, int responseByteLimit, JsonArray resultsArray, bool truncated, JsonArray truncatedQueries) @@ -2861,7 +2864,7 @@ private int EstimateBatchResponseBytes(JsonNode? id, string summary, int submitt payload["truncated_queries"] = truncatedQueries.DeepClone(); } - return EstimateJsonUtf8Bytes(CreateToolResult(id, summary, payload)); + return EstimateJsonUtf8Bytes(CreateToolResult(id, summary, payload), responseByteLimit); } private static string GetBatchFailureScope(int submittedCount, int successCount, int failureCount, int? cascadeStartedAtIndex) diff --git a/tests/CodeIndex.Tests/McpServerTests.cs b/tests/CodeIndex.Tests/McpServerTests.cs index 2448803995..a967b1a9c0 100644 --- a/tests/CodeIndex.Tests/McpServerTests.cs +++ b/tests/CodeIndex.Tests/McpServerTests.cs @@ -5318,6 +5318,40 @@ public void ToolsCall_ResponseOverByteLimit_ReturnsStructuredError() Assert.Equal("response_too_large", response["error"]!["data"]!["reason"]!.GetValue()); Assert.Equal(256, response["error"]!["data"]!["limit_bytes"]!.GetValue()); Assert.True(response["error"]!["data"]!["actual_bytes"]!.GetValue() > 256); + Assert.False(response["error"]!["data"]!["actual_bytes_exact"]!.GetValue()); + } + + [Fact] + public void ResponseLimitSerializer_StopsBeforeFullStringMaterialization_Issue2860() + { + var payload = new JsonObject + { + ["value"] = new string('x', 10_000), + }; + + var withinLimit = _server.TrySerializeJsonNodeWithinByteLimitForTests(payload, 256, out var serialized, out var bytesWritten); + + Assert.False(withinLimit); + Assert.Null(serialized); + Assert.True(bytesWritten > 256); + Assert.True(bytesWritten < 10_000); + } + + [Fact] + public void ResponseLimitSerializer_ReturnsCapturedJsonWhenWithinLimit_Issue2860() + { + var payload = new JsonObject + { + ["value"] = "ok", + }; + + var withinLimit = _server.TrySerializeJsonNodeWithinByteLimitForTests(payload, 256, out var serialized, out var bytesWritten); + + Assert.True(withinLimit); + Assert.NotNull(serialized); + Assert.Equal(Encoding.UTF8.GetByteCount(serialized), bytesWritten); + using var parsed = JsonDocument.Parse(serialized); + Assert.Equal("ok", parsed.RootElement.GetProperty("value").GetString()); } [Fact]