Bound STDIO client reads
Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>
This commit is contained in:
parent
3b72eec0b8
commit
3477fcd0cd
3 changed files with 144 additions and 1 deletions
|
|
@ -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);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
Loading…
Reference in a new issue