Bound STDIO server reads

Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>

Author: Kehrlann

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);

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<McpSchema.JSONRPCMessage> 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