From 237f258d4fe8625e831b773816d7ba665dea28fc Mon Sep 17 00:00:00 2001 From: Daniel Garnier-Moiroux Date: Thu, 2 Jul 2026 18:48:58 +0200 Subject: [PATCH 1/2] Bound STDIO server reads Signed-off-by: Daniel Garnier-Moiroux --- .../StdioServerTransportProvider.java | 56 ++++++++++++++++++- .../StdioServerTransportProviderTests.java | 48 +++++++++++++++- 2 files changed, 101 insertions(+), 3 deletions(-) diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java index 045d7e3a9..cc5123336 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java @@ -41,6 +41,8 @@ */ public class StdioServerTransportProvider implements McpServerTransportProvider { + private static final int DEFAULT_INPUT_MAX_SIZE = 16 * 1024 * 1024; // 16MB + private static final Logger logger = LoggerFactory.getLogger(StdioServerTransportProvider.class); private final McpJsonMapper jsonMapper; @@ -49,6 +51,8 @@ public class StdioServerTransportProvider implements McpServerTransportProvider private final OutputStream outputStream; + private final int inputMaxSize; + private McpServerSession session; private final AtomicBoolean isClosing = new AtomicBoolean(false); @@ -72,13 +76,29 @@ public StdioServerTransportProvider(McpJsonMapper jsonMapper) { * @param outputStream The output stream to write to */ public StdioServerTransportProvider(McpJsonMapper jsonMapper, InputStream inputStream, OutputStream outputStream) { + this(jsonMapper, inputStream, outputStream, DEFAULT_INPUT_MAX_SIZE); + } + + /** + * Creates a new StdioServerTransportProvider. + * @param jsonMapper The JsonMapper to use for JSON serialization/deserialization + * @param inputStream The input stream to read from + * @param outputStream The output stream to write to + * @param inputMaxSize The maximum number of characters read for a single inbound + * message. A peer that sends a longer message (or never terminates a line) has its + * message rejected instead of forcing the transport to buffer it in memory. + */ + public StdioServerTransportProvider(McpJsonMapper jsonMapper, InputStream inputStream, OutputStream outputStream, + int inputMaxSize) { Assert.notNull(jsonMapper, "The JsonMapper can not be null"); Assert.notNull(inputStream, "The InputStream can not be null"); Assert.notNull(outputStream, "The OutputStream can not be null"); + Assert.isTrue(inputMaxSize > 0, "inputMaxSize must be positive"); this.jsonMapper = jsonMapper; this.inputStream = inputStream; this.outputStream = outputStream; + this.inputMaxSize = inputMaxSize; } @Override @@ -211,7 +231,7 @@ private void startInboundProcessing() { reader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8)); while (!isClosing.get()) { try { - String line = reader.readLine(); + String line = readLine(reader, inputMaxSize); if (line == null || isClosing.get()) { break; } @@ -232,6 +252,10 @@ private void startInboundProcessing() { break; } } + catch (MaxSizeExceededException e) { + logIfNotClosing("Inbound message exceeds the maximum allowed size", e); + break; + } catch (IOException e) { logIfNotClosing("Error reading from stdin", e); break; @@ -304,6 +328,36 @@ else if (isClosing.get()) { outboundConsumer.apply(outboundSink.asFlux()).subscribe(); } // @formatter:on + /** + * Read line with a max size. + */ + private static String readLine(BufferedReader reader, int maxSize) + throws IOException, MaxSizeExceededException { + StringBuilder sb = new StringBuilder(); + int c; + while ((c = reader.read()) != -1) { + if (c == '\n') { + return sb.toString(); + } + if (c == '\r') { + // Consume an optional trailing '\n' so that "\r\n" is treated as a + // single terminator, mirroring BufferedReader#readLine(). + reader.mark(1); + int next = reader.read(); + if (next != '\n' && next != -1) { + reader.reset(); + } + return sb.toString(); + } + if (sb.length() >= maxSize) { + throw new MaxSizeExceededException( + "Inbound message exceeds the maximum allowed size of " + maxSize + " characters"); + } + sb.append((char) c); + } + return sb.isEmpty() ? null : sb.toString(); + } + private void logIfNotClosing(String message, Exception e) { if (!isClosing.get()) { logger.error(message, e); diff --git a/mcp-test/src/test/java/io/modelcontextprotocol/server/transport/StdioServerTransportProviderTests.java b/mcp-test/src/test/java/io/modelcontextprotocol/server/transport/StdioServerTransportProviderTests.java index 6c2cc2bf4..ab1068954 100644 --- a/mcp-test/src/test/java/io/modelcontextprotocol/server/transport/StdioServerTransportProviderTests.java +++ b/mcp-test/src/test/java/io/modelcontextprotocol/server/transport/StdioServerTransportProviderTests.java @@ -11,24 +11,25 @@ import java.io.InputStreamReader; import java.io.PrintStream; import java.nio.charset.StandardCharsets; +import java.time.Duration; import java.util.Map; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; import io.modelcontextprotocol.json.McpJsonDefaults; -import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpServerSession; import io.modelcontextprotocol.spec.McpServerTransport; +import org.awaitility.Awaitility; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import reactor.core.publisher.Mono; import reactor.test.StepVerifier; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; @@ -246,6 +247,49 @@ void shouldHandleInvalidJsonMessage() throws Exception { .verifyComplete(); } + @Test + void shouldRejectInboundMessageExceedingMaxSize() throws Exception { + // A line larger than the configured limit that never terminates with a newline. + // BufferedReader#readLine would buffer the whole thing; the bounded reader must + // abort instead. + int maxSize = 1024; + String oversized = "a".repeat(maxSize + 10); + InputStream stream = new ByteArrayInputStream(oversized.getBytes(StandardCharsets.UTF_8)); + + transportProvider = new StdioServerTransportProvider(McpJsonDefaults.getMapper(), stream, testOutPrintStream, + maxSize); + + AtomicReference capturedMessage = new AtomicReference<>(); + McpServerSession.Factory realSessionFactory = transport -> { + McpServerSession session = mock(McpServerSession.class); + when(session.handle(any())).thenAnswer(invocation -> { + capturedMessage.set(invocation.getArgument(0)); + return Mono.empty(); + }); + when(session.closeGracefully()).thenReturn(Mono.empty()); + return session; + }; + + transportProvider.setSessionFactory(realSessionFactory); + + Awaitility.await() + .atMost(Duration.ofSeconds(5)) + .pollInterval(Duration.ofMillis(100)) + .untilAsserted( + () -> assertThat(testErr.toString()).contains("Inbound message exceeds the maximum allowed size")); + + // message is never processed + assertThat(capturedMessage.get()).isNull(); + } + + @Test + void shouldRejectNonPositiveMaxSize() { + assertThatThrownBy( + () -> new StdioServerTransportProvider(McpJsonDefaults.getMapper(), System.in, System.out, 0)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("inputMaxSize must be positive"); + } + @Test void shouldHandleSessionClose() throws Exception { // Set session factory From 79b6706dfb88fd8424b7bcd408189dcbbe365d38 Mon Sep 17 00:00:00 2001 From: Daniel Garnier-Moiroux Date: Wed, 8 Jul 2026 10:47:09 +0200 Subject: [PATCH 2/2] Bound STDIO client reads Signed-off-by: Daniel Garnier-Moiroux --- .../transport/MaxSizeExceededException.java | 18 +++++ .../transport/StdioClientTransport.java | 57 ++++++++++++++- .../transport/StdioClientTransportTests.java | 70 +++++++++++++++++++ 3 files changed, 144 insertions(+), 1 deletion(-) create mode 100644 mcp-core/src/main/java/io/modelcontextprotocol/client/transport/MaxSizeExceededException.java create mode 100644 mcp-test/src/test/java/io/modelcontextprotocol/client/transport/StdioClientTransportTests.java diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/MaxSizeExceededException.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/MaxSizeExceededException.java new file mode 100644 index 000000000..807034057 --- /dev/null +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/MaxSizeExceededException.java @@ -0,0 +1,18 @@ +/* + * Copyright 2026-2026 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +/** + * Thrown when reading an inbound message would exceed the configured maximum size. + * + * @author Daniel Garnier-Moiroux + */ +class MaxSizeExceededException extends Exception { + + MaxSizeExceededException(String message) { + super(message); + } + +} diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/StdioClientTransport.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/StdioClientTransport.java index e73e43ef5..73522ae6d 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/StdioClientTransport.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/StdioClientTransport.java @@ -44,6 +44,8 @@ public class StdioClientTransport implements McpClientTransport { private static final Logger logger = LoggerFactory.getLogger(StdioClientTransport.class); + private static final int DEFAULT_INPUT_MAX_SIZE = 16 * 1024 * 1024; // 16MB + // @formatter:off private static final Set EXIT_SUCCESS_CODES = Set.of( 0, // success @@ -76,6 +78,8 @@ public class StdioClientTransport implements McpClientTransport { private final Sinks.Many errorSink; + private final int inputMaxSize; + private volatile boolean isClosing = false; // visible for tests @@ -87,8 +91,21 @@ public class StdioClientTransport implements McpClientTransport { * @param jsonMapper The JsonMapper to use for JSON serialization/deserialization */ public StdioClientTransport(ServerParameters params, McpJsonMapper jsonMapper) { + this(params, jsonMapper, DEFAULT_INPUT_MAX_SIZE); + } + + /** + * Creates a new StdioClientTransport with the specified parameters and JsonMapper. + * @param params The parameters for configuring the server process + * @param jsonMapper The JsonMapper to use for JSON serialization/deserialization + * @param inputMaxSize The maximum number of characters read for a single inbound + * message. A peer that sends a longer message (or never terminates a line) has its + * message rejected instead of forcing the transport to buffer it in memory. + */ + public StdioClientTransport(ServerParameters params, McpJsonMapper jsonMapper, int inputMaxSize) { Assert.notNull(params, "The params can not be null"); Assert.notNull(jsonMapper, "The JsonMapper can not be null"); + Assert.isTrue(inputMaxSize > 0, "inputMaxSize must be positive"); this.inboundSink = Sinks.many().unicast().onBackpressureBuffer(); this.outboundSink = Sinks.many().unicast().onBackpressureBuffer(); @@ -97,6 +114,8 @@ public StdioClientTransport(ServerParameters params, McpJsonMapper jsonMapper) { this.jsonMapper = jsonMapper; + this.inputMaxSize = inputMaxSize; + this.errorSink = Sinks.many().unicast().onBackpressureBuffer(); // Start threads @@ -260,7 +279,7 @@ private void startInboundProcessing() { this.inboundScheduler.schedule(() -> { try (BufferedReader processReader = new BufferedReader(new InputStreamReader(process.getInputStream()))) { String line; - while (!isClosing && (line = processReader.readLine()) != null) { + while (!isClosing && (line = readLine(processReader, inputMaxSize)) != null) { try { JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.jsonMapper, line); if (!this.inboundSink.tryEmitNext(message).isSuccess()) { @@ -278,6 +297,11 @@ private void startInboundProcessing() { } } } + catch (MaxSizeExceededException e) { + if (!isClosing) { + logger.error("Inbound message exceeds the maximum allowed size", e); + } + } catch (IOException e) { if (!isClosing) { logger.error("Error reading from input stream", e); @@ -290,6 +314,37 @@ private void startInboundProcessing() { }); } + /** + * Reads a single line, mirroring {@link BufferedReader#readLine()}, but aborting once + * more than {@code maxSize} characters have been read without encountering a line + * terminator. This bounds how much memory a single inbound message can occupy. + */ + static String readLine(BufferedReader reader, int maxSize) throws IOException, MaxSizeExceededException { + StringBuilder sb = new StringBuilder(); + int c; + while ((c = reader.read()) != -1) { + if (c == '\n') { + return sb.toString(); + } + if (c == '\r') { + // Consume an optional trailing '\n' so that "\r\n" is treated as a + // single terminator, mirroring BufferedReader#readLine(). + reader.mark(1); + int next = reader.read(); + if (next != '\n' && next != -1) { + reader.reset(); + } + return sb.toString(); + } + if (sb.length() >= maxSize) { + throw new MaxSizeExceededException( + "Inbound message exceeds the maximum allowed size of " + maxSize + " characters"); + } + sb.append((char) c); + } + return sb.isEmpty() ? null : sb.toString(); + } + /** * Starts the outbound processing thread that writes JSON-RPC messages to the * process's output stream. Messages are serialized to JSON and written with a newline diff --git a/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/StdioClientTransportTests.java b/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/StdioClientTransportTests.java new file mode 100644 index 000000000..0aad3934a --- /dev/null +++ b/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/StdioClientTransportTests.java @@ -0,0 +1,70 @@ +/* + * Copyright 2026-2026 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +import java.io.ByteArrayOutputStream; +import java.io.PrintStream; +import java.time.Duration; + +import org.awaitility.Awaitility; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import reactor.test.StepVerifier; + +import static io.modelcontextprotocol.util.McpJsonMapperUtils.JSON_MAPPER; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration tests for {@link StdioClientTransport}. + * + * @author Daniel Garnier-Moiroux + */ +class StdioClientTransportTests { + + private final PrintStream originalOut = System.out; + + private final PrintStream originalErr = System.err; + + private ByteArrayOutputStream testErr; + + @BeforeEach + void setUp() { + testErr = new ByteArrayOutputStream(); + PrintStream testOutPrintStream = new PrintStream(testErr, true); + System.setOut(testOutPrintStream); + System.setErr(testOutPrintStream); + } + + @AfterEach + void tearDown() { + System.setOut(originalOut); + System.setErr(originalErr); + } + + @Test + void shouldRejectInboundMessageExceedingMaxSize() throws Exception { + // A server process that emits an endless line with no newline terminator. A + // plain BufferedReader#readLine would buffer it all; the bounded reader must + // abort instead of exhausting memory. + int maxSize = 1024; + ServerParameters params = ServerParameters.builder("sh").args("-c", "while :; do printf a; done").build(); + + StdioClientTransport transport = new StdioClientTransport(params, JSON_MAPPER, maxSize); + try { + StepVerifier.create(transport.connect(msg -> msg)).verifyComplete(); + + Awaitility.await() + .atMost(Duration.ofSeconds(5)) + .pollInterval(Duration.ofMillis(100)) + .untilAsserted(() -> assertThat(testErr.toString()) + .contains("Inbound message exceeds the maximum allowed size")); + } + finally { + StepVerifier.create(transport.closeGracefully()).verifyComplete(); + } + } + +}