Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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);
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -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<Integer> EXIT_SUCCESS_CODES = Set.of(
0, // success
Expand Down Expand Up @@ -76,6 +78,8 @@ public class StdioClientTransport implements McpClientTransport {

private final Sinks.Many<String> errorSink;

private final int inputMaxSize;

private volatile boolean isClosing = false;

// visible for tests
Expand All @@ -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();
Expand All @@ -97,6 +114,8 @@ public StdioClientTransport(ServerParameters params, McpJsonMapper jsonMapper) {

this.jsonMapper = jsonMapper;

this.inputMaxSize = inputMaxSize;

this.errorSink = Sinks.many().unicast().onBackpressureBuffer();

// Start threads
Expand Down Expand Up @@ -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()) {
Expand All @@ -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);
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);
Expand All @@ -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
Expand Down Expand Up @@ -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;
}
Expand All @@ -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;
Expand Down Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
@@ -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();
}
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down
Loading