Bound STDIO client reads

Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>
This commit is contained in:
Daniel Garnier-Moiroux 2026-07-08 10:47:09 +02:00 • committed by Daniel Garnier-Moiroux
commit 3477fcd0cd
3 changed files with 144 additions and 1 deletions

View file

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

View file

@ -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
@ -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
@ -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 class StdioClientTransport implements McpClientTransport {
this.jsonMapper = jsonMapper;
this.inputMaxSize = inputMaxSize;
this.errorSink = Sinks.many().unicast().onBackpressureBuffer();
// Start threads
@ -260,7 +279,7 @@ public class StdioClientTransport implements McpClientTransport {
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 @@ public class StdioClientTransport implements McpClientTransport {
}
}
}
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 @@ public class StdioClientTransport implements McpClientTransport {
});
}
/**
* 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

View file

@ -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();
}
}
}