Bound HTTP client reads
Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>
This commit is contained in:
Daniel Garnier-Moiroux
parent
8c0fdd1b1f
commit
2e5dbd36e0
9 changed files with 1302 additions and 138 deletions
|
|
@ -82,6 +82,11 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
/** Default SSE endpoint path */
|
||||
private static final String DEFAULT_SSE_ENDPOINT = "/sse";
|
||||
|
||||
/**
|
||||
* Default maximum number of bytes read for a single inbound message.
|
||||
*/
|
||||
private static final int DEFAULT_MAX_RESPONSE_SIZE = 16 * 1024 * 1024; // 16MiB
|
||||
|
||||
/** Base URI for the MCP server */
|
||||
private final URI baseUri;
|
||||
|
||||
|
|
@ -122,6 +127,12 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
*/
|
||||
private final SseMessageEndpointValidator messageEndpointValidator;
|
||||
|
||||
/**
|
||||
* Maximum number of bytes read for a single inbound message, whether it arrives on
|
||||
* the SSE stream or as the response to a posted message.
|
||||
*/
|
||||
private final int maxResponseSize;
|
||||
|
||||
/**
|
||||
* Creates a new transport instance with custom HTTP client builder, object mapper,
|
||||
* and headers.
|
||||
|
|
@ -133,11 +144,13 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
* @param httpRequestCustomizer customizer for the requestBuilder before executing
|
||||
* requests
|
||||
* @param messageEndpointValidator validator for the message endpoint
|
||||
* @param maxResponseSize the maximum number of bytes read for a single inbound
|
||||
* message
|
||||
* @throws IllegalArgumentException if objectMapper, clientBuilder, or headers is null
|
||||
*/
|
||||
HttpClientSseClientTransport(HttpClient httpClient, HttpRequest.Builder requestBuilder, String baseUri,
|
||||
String sseEndpoint, McpJsonMapper jsonMapper, McpAsyncHttpClientRequestCustomizer httpRequestCustomizer,
|
||||
SseMessageEndpointValidator messageEndpointValidator) {
|
||||
SseMessageEndpointValidator messageEndpointValidator, int maxResponseSize) {
|
||||
Assert.notNull(jsonMapper, "jsonMapper must not be null");
|
||||
Assert.hasText(baseUri, "baseUri must not be empty");
|
||||
Assert.hasText(sseEndpoint, "sseEndpoint must not be empty");
|
||||
|
|
@ -145,6 +158,7 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
Assert.notNull(requestBuilder, "requestBuilder must not be null");
|
||||
Assert.notNull(httpRequestCustomizer, "httpRequestCustomizer must not be null");
|
||||
Assert.notNull(messageEndpointValidator, "messageEndpointValidator must not be null");
|
||||
Assert.isTrue(maxResponseSize > 0, "maxResponseSize must be positive");
|
||||
this.baseUri = URI.create(baseUri);
|
||||
this.sseEndpoint = sseEndpoint;
|
||||
this.jsonMapper = jsonMapper;
|
||||
|
|
@ -152,6 +166,7 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
this.requestBuilder = requestBuilder;
|
||||
this.httpRequestCustomizer = httpRequestCustomizer;
|
||||
this.messageEndpointValidator = messageEndpointValidator;
|
||||
this.maxResponseSize = maxResponseSize;
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
@ -189,6 +204,8 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
|
||||
private SseMessageEndpointValidator messageEndpointValidator = new DefaultSseMessageEndpointValidator();
|
||||
|
||||
private int maxResponseSize = DEFAULT_MAX_RESPONSE_SIZE;
|
||||
|
||||
/**
|
||||
* Creates a new builder instance.
|
||||
*/
|
||||
|
|
@ -344,6 +361,26 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets the maximum number of bytes read for a single inbound message, whether it
|
||||
* arrives on the SSE stream or as the response to a posted message. A peer that
|
||||
* sends a larger message (or never terminates one) has its stream aborted instead
|
||||
* of forcing the transport to buffer it in memory. Defaults to 16MiB.
|
||||
*
|
||||
* <p>
|
||||
* The bound applies per message, not to the stream as a whole: a long-lived SSE
|
||||
* stream may deliver any number of messages, each up to this size. SSE field
|
||||
* framing is allowed a small amount of headroom on top of this size, so a message
|
||||
* of exactly this many bytes is still accepted.
|
||||
* @param maxResponseSize the maximum inbound message size, in bytes
|
||||
* @return this builder
|
||||
*/
|
||||
public Builder maxResponseSize(int maxResponseSize) {
|
||||
Assert.isTrue(maxResponseSize > 0, "maxResponseSize must be positive");
|
||||
this.maxResponseSize = maxResponseSize;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Builds a new {@link HttpClientSseClientTransport} instance.
|
||||
* @return a new transport instance
|
||||
|
|
@ -352,7 +389,7 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
HttpClient httpClient = this.clientBuilder.connectTimeout(this.connectTimeout).build();
|
||||
return new HttpClientSseClientTransport(httpClient, requestBuilder, baseUri, sseEndpoint,
|
||||
jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper, httpRequestCustomizer,
|
||||
messageEndpointValidator);
|
||||
messageEndpointValidator, maxResponseSize);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -371,13 +408,15 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
var transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY);
|
||||
return Mono.from(this.httpRequestCustomizer.customize(builder, "GET", uri, null, transportContext));
|
||||
}).flatMap(requestBuilder -> Mono.create(sink -> {
|
||||
Disposable connection = Flux.<ResponseEvent>create(sseSink -> this.httpClient
|
||||
.sendAsync(requestBuilder.build(),
|
||||
responseInfo -> ResponseSubscribers.sseToBodySubscriber(responseInfo, sseSink))
|
||||
.exceptionallyCompose(e -> {
|
||||
sseSink.error(e);
|
||||
return CompletableFuture.failedFuture(e);
|
||||
}))
|
||||
Disposable connection = Flux.<ResponseEvent>create(
|
||||
sseSink -> this.httpClient
|
||||
.sendAsync(requestBuilder.build(),
|
||||
responseInfo -> ResponseSubscribers.sseToBodySubscriber(responseInfo, sseSink,
|
||||
this.maxResponseSize))
|
||||
.exceptionallyCompose(e -> {
|
||||
sseSink.error(e);
|
||||
return CompletableFuture.failedFuture(e);
|
||||
}))
|
||||
.map(responseEvent -> (ResponseSubscribers.SseResponseEvent) responseEvent)
|
||||
.flatMap(responseEvent -> {
|
||||
if (isClosing) {
|
||||
|
|
@ -508,7 +547,8 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
return Mono.from(this.httpRequestCustomizer.customize(builder, "POST", requestUri, body, transportContext));
|
||||
}).flatMap(customizedBuilder -> {
|
||||
var request = customizedBuilder.build();
|
||||
return Mono.fromFuture(httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofString()));
|
||||
return Mono.fromFuture(
|
||||
httpClient.sendAsync(request, ResponseSubscribers.boundedStringBodyHandler(this.maxResponseSize)));
|
||||
});
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@ import java.io.IOException;
|
|||
import java.net.URI;
|
||||
import java.net.http.HttpClient;
|
||||
import java.net.http.HttpRequest;
|
||||
import java.net.http.HttpResponse;
|
||||
import java.net.http.HttpResponse.BodyHandler;
|
||||
import java.time.Duration;
|
||||
import java.util.Collections;
|
||||
|
|
@ -82,6 +81,11 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
|
||||
private static final String DEFAULT_ENDPOINT = "/mcp";
|
||||
|
||||
/**
|
||||
* Default maximum number of bytes read for a single inbound message.
|
||||
*/
|
||||
private static final int DEFAULT_MAX_RESPONSE_SIZE = 16 * 1024 * 1024; // 16MiB
|
||||
|
||||
/**
|
||||
* HTTP client for sending messages to the server. Uses HTTP POST over the message
|
||||
* endpoint
|
||||
|
|
@ -129,10 +133,17 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
|
||||
private final String latestSupportedProtocolVersion;
|
||||
|
||||
/**
|
||||
* Maximum number of bytes read for a single inbound message, whether it arrives on an
|
||||
* SSE stream or as a JSON response body.
|
||||
*/
|
||||
private final int maxResponseSize;
|
||||
|
||||
private HttpClientStreamableHttpTransport(McpJsonMapper jsonMapper, HttpClient httpClient,
|
||||
HttpRequest.Builder requestBuilder, String baseUri, String endpoint, boolean resumableStreams,
|
||||
boolean openConnectionOnStartup, McpAsyncHttpClientRequestCustomizer httpRequestCustomizer,
|
||||
List<String> supportedProtocolVersions) {
|
||||
List<String> supportedProtocolVersions, int maxResponseSize) {
|
||||
Assert.isTrue(maxResponseSize > 0, "maxResponseSize must be positive");
|
||||
this.jsonMapper = jsonMapper;
|
||||
this.httpClient = httpClient;
|
||||
this.requestBuilder = requestBuilder;
|
||||
|
|
@ -147,6 +158,7 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
.sorted(Comparator.reverseOrder())
|
||||
.findFirst()
|
||||
.get();
|
||||
this.maxResponseSize = maxResponseSize;
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
@ -203,7 +215,8 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
return Mono.from(this.httpRequestCustomizer.customize(builder, "DELETE", uri, null, transportContext));
|
||||
}).flatMap(requestBuilder -> {
|
||||
var request = requestBuilder.build();
|
||||
return Mono.fromFuture(() -> this.httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofString()));
|
||||
return Mono.fromFuture(() -> this.httpClient.sendAsync(request,
|
||||
ResponseSubscribers.boundedStringBodyHandler(this.maxResponseSize)));
|
||||
}).then();
|
||||
}
|
||||
|
||||
|
|
@ -275,121 +288,116 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
var transportContext = connectionCtx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY);
|
||||
return Mono.from(this.httpRequestCustomizer.customize(builder, "GET", uri, null, transportContext));
|
||||
})
|
||||
.flatMapMany(
|
||||
requestBuilder -> Flux.<ResponseEvent>create(
|
||||
sseSink -> this.httpClient
|
||||
.sendAsync(requestBuilder.build(),
|
||||
responseInfo -> ResponseSubscribers.sseToBodySubscriber(responseInfo,
|
||||
sseSink))
|
||||
.whenComplete((response, throwable) -> {
|
||||
if (throwable != null) {
|
||||
sseSink.error(throwable);
|
||||
}
|
||||
else {
|
||||
logger.debug("SSE connection established successfully");
|
||||
}
|
||||
}))
|
||||
.map(responseEvent -> (ResponseSubscribers.SseResponseEvent) responseEvent)
|
||||
.flatMap(responseEvent -> {
|
||||
int statusCode = responseEvent.responseInfo().statusCode();
|
||||
|
||||
if (statusCode >= 200 && statusCode < 300) {
|
||||
|
||||
if (MESSAGE_EVENT_TYPE.equals(responseEvent.sseEvent().event())) {
|
||||
String data = responseEvent.sseEvent().data();
|
||||
// Per 2025-11-25 spec (SEP-1699), servers may
|
||||
// send SSE events
|
||||
// with empty data to prime the client for
|
||||
// reconnection.
|
||||
// Skip these events as they contain no JSON-RPC
|
||||
// message.
|
||||
if (data == null || data.isBlank()) {
|
||||
logger.debug("Skipping SSE event with empty data (stream primer)");
|
||||
return Flux.empty();
|
||||
}
|
||||
try {
|
||||
// We don't support batching ATM and probably
|
||||
// won't since the next version considers
|
||||
// removing it.
|
||||
McpSchema.JSONRPCMessage message = McpSchema
|
||||
.deserializeJsonRpcMessage(this.jsonMapper, data);
|
||||
|
||||
Tuple2<Optional<String>, Iterable<McpSchema.JSONRPCMessage>> idWithMessages = Tuples
|
||||
.of(Optional.ofNullable(responseEvent.sseEvent().id()),
|
||||
List.of(message));
|
||||
|
||||
McpTransportStream<Disposable> sessionStream = stream != null ? stream
|
||||
: new DefaultMcpTransportStream<>(this.resumableStreams,
|
||||
this::reconnect);
|
||||
logger.debug("Connected stream {}", sessionStream.streamId());
|
||||
|
||||
return Flux.from(sessionStream.consumeSseStream(Flux.just(idWithMessages)));
|
||||
|
||||
}
|
||||
catch (IOException ioException) {
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
|
||||
"Error parsing JSON-RPC message: " + responseEvent, ioException));
|
||||
}
|
||||
}
|
||||
else {
|
||||
logger.debug("Received SSE event with type: {}", responseEvent.sseEvent());
|
||||
return Flux.empty();
|
||||
}
|
||||
.flatMapMany(requestBuilder -> Flux.<ResponseEvent>create(
|
||||
sseSink -> this.httpClient
|
||||
.sendAsync(requestBuilder.build(),
|
||||
responseInfo -> ResponseSubscribers.sseToBodySubscriber(responseInfo, sseSink,
|
||||
maxResponseSize))
|
||||
.whenComplete((response, throwable) -> {
|
||||
if (throwable != null) {
|
||||
sseSink.error(throwable);
|
||||
}
|
||||
else if (statusCode == METHOD_NOT_ALLOWED) { // NotAllowed
|
||||
logger
|
||||
.debug("The server does not support SSE streams, using request-response mode.");
|
||||
return Flux.empty();
|
||||
}
|
||||
else if (statusCode == NOT_FOUND) {
|
||||
|
||||
if (transportSession != null && transportSession.sessionId().isPresent()) {
|
||||
// only if the request was sent with a session id
|
||||
// and the response is 404, we consider it a
|
||||
// session not found error.
|
||||
logger.debug("Session not found for session ID: {}",
|
||||
transportSession.sessionId().get());
|
||||
String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession);
|
||||
McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException(
|
||||
"Session not found for session ID: " + sessionIdRepresentation);
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(exception);
|
||||
}
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(
|
||||
new McpTransportException("Server Not Found. Status code:" + statusCode
|
||||
+ ", response-event:" + responseEvent));
|
||||
}
|
||||
else if (statusCode == BAD_REQUEST) {
|
||||
if (transportSession != null && transportSession.sessionId().isPresent()) {
|
||||
// only if the request was sent with a session id
|
||||
// and thre response is 404, we consider it a
|
||||
// session not found error.
|
||||
String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession);
|
||||
McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException(
|
||||
"Session not found for session ID: " + sessionIdRepresentation);
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(exception);
|
||||
}
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(
|
||||
new McpTransportException("Bad Request. Status code:" + statusCode
|
||||
+ ", response-event:" + responseEvent));
|
||||
|
||||
}
|
||||
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
|
||||
"Received unrecognized SSE event type: " + responseEvent.sseEvent().event()));
|
||||
}).<McpSchema
|
||||
.JSONRPCMessage>flatMap(
|
||||
jsonrpcMessage -> this.handler.get().apply(Mono.just(jsonrpcMessage)))
|
||||
.onErrorMap(CompletionException.class, t -> t.getCause())
|
||||
.onErrorComplete(t -> {
|
||||
this.handleException(t);
|
||||
return true;
|
||||
})
|
||||
.doFinally(s -> {
|
||||
Disposable ref = disposableRef.getAndSet(null);
|
||||
if (ref != null) {
|
||||
transportSession.removeConnection(ref);
|
||||
else {
|
||||
logger.debug("SSE connection established successfully");
|
||||
}
|
||||
}))
|
||||
.map(responseEvent -> (ResponseSubscribers.SseResponseEvent) responseEvent)
|
||||
.flatMap(responseEvent -> {
|
||||
int statusCode = responseEvent.responseInfo().statusCode();
|
||||
|
||||
if (statusCode >= 200 && statusCode < 300) {
|
||||
|
||||
if (MESSAGE_EVENT_TYPE.equals(responseEvent.sseEvent().event())) {
|
||||
String data = responseEvent.sseEvent().data();
|
||||
// Per 2025-11-25 spec (SEP-1699), servers may
|
||||
// send SSE events
|
||||
// with empty data to prime the client for
|
||||
// reconnection.
|
||||
// Skip these events as they contain no JSON-RPC
|
||||
// message.
|
||||
if (data == null || data.isBlank()) {
|
||||
logger.debug("Skipping SSE event with empty data (stream primer)");
|
||||
return Flux.empty();
|
||||
}
|
||||
try {
|
||||
// We don't support batching ATM and probably
|
||||
// won't since the next version considers
|
||||
// removing it.
|
||||
McpSchema.JSONRPCMessage message = McpSchema
|
||||
.deserializeJsonRpcMessage(this.jsonMapper, data);
|
||||
|
||||
Tuple2<Optional<String>, Iterable<McpSchema.JSONRPCMessage>> idWithMessages = Tuples
|
||||
.of(Optional.ofNullable(responseEvent.sseEvent().id()), List.of(message));
|
||||
|
||||
McpTransportStream<Disposable> sessionStream = stream != null ? stream
|
||||
: new DefaultMcpTransportStream<>(this.resumableStreams, this::reconnect);
|
||||
logger.debug("Connected stream {}", sessionStream.streamId());
|
||||
|
||||
return Flux.from(sessionStream.consumeSseStream(Flux.just(idWithMessages)));
|
||||
|
||||
}
|
||||
catch (IOException ioException) {
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
|
||||
"Error parsing JSON-RPC message: " + responseEvent, ioException));
|
||||
}
|
||||
}
|
||||
else {
|
||||
logger.debug("Received SSE event with type: {}", responseEvent.sseEvent());
|
||||
return Flux.empty();
|
||||
}
|
||||
}
|
||||
else if (statusCode == METHOD_NOT_ALLOWED) { // NotAllowed
|
||||
logger.debug("The server does not support SSE streams, using request-response mode.");
|
||||
return Flux.empty();
|
||||
}
|
||||
else if (statusCode == NOT_FOUND) {
|
||||
|
||||
if (transportSession != null && transportSession.sessionId().isPresent()) {
|
||||
// only if the request was sent with a session id
|
||||
// and the response is 404, we consider it a
|
||||
// session not found error.
|
||||
logger.debug("Session not found for session ID: {}",
|
||||
transportSession.sessionId().get());
|
||||
String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession);
|
||||
McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException(
|
||||
"Session not found for session ID: " + sessionIdRepresentation);
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(exception);
|
||||
}
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(
|
||||
new McpTransportException("Server Not Found. Status code:" + statusCode
|
||||
+ ", response-event:" + responseEvent));
|
||||
}
|
||||
else if (statusCode == BAD_REQUEST) {
|
||||
if (transportSession != null && transportSession.sessionId().isPresent()) {
|
||||
// only if the request was sent with a session id
|
||||
// and thre response is 404, we consider it a
|
||||
// session not found error.
|
||||
String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession);
|
||||
McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException(
|
||||
"Session not found for session ID: " + sessionIdRepresentation);
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(exception);
|
||||
}
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
|
||||
"Bad Request. Status code:" + statusCode + ", response-event:" + responseEvent));
|
||||
|
||||
}
|
||||
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
|
||||
"Received unrecognized SSE event type: " + responseEvent.sseEvent().event()));
|
||||
}).<McpSchema
|
||||
.JSONRPCMessage>flatMap(
|
||||
jsonrpcMessage -> this.handler.get().apply(Mono.just(jsonrpcMessage)))
|
||||
.onErrorMap(CompletionException.class, t -> t.getCause())
|
||||
.onErrorComplete(t -> {
|
||||
this.handleException(t);
|
||||
return true;
|
||||
})
|
||||
.doFinally(s -> {
|
||||
Disposable ref = disposableRef.getAndSet(null);
|
||||
if (ref != null) {
|
||||
transportSession.removeConnection(ref);
|
||||
}
|
||||
}))
|
||||
.contextWrite(ctx)
|
||||
.subscribe();
|
||||
|
||||
|
|
@ -409,16 +417,16 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
if (contentType.contains(TEXT_EVENT_STREAM)) {
|
||||
// For SSE streams, use line subscriber that returns Void
|
||||
logger.debug("Received SSE stream response, using line subscriber");
|
||||
return ResponseSubscribers.sseToBodySubscriber(responseInfo, sink);
|
||||
return ResponseSubscribers.sseToBodySubscriber(responseInfo, sink, this.maxResponseSize);
|
||||
}
|
||||
else if (contentType.contains(APPLICATION_JSON)) {
|
||||
// For JSON responses and others, use string subscriber
|
||||
logger.debug("Received response, using string subscriber");
|
||||
return ResponseSubscribers.aggregateBodySubscriber(responseInfo, sink);
|
||||
return ResponseSubscribers.aggregateBodySubscriber(responseInfo, sink, this.maxResponseSize);
|
||||
}
|
||||
|
||||
logger.debug("Received Bodyless response, using discarding subscriber");
|
||||
return ResponseSubscribers.bodilessBodySubscriber(responseInfo, sink);
|
||||
return ResponseSubscribers.bodilessBodySubscriber(responseInfo, sink, this.maxResponseSize);
|
||||
};
|
||||
|
||||
return responseBodyHandler;
|
||||
|
|
@ -668,6 +676,8 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
private List<String> supportedProtocolVersions = List.of(ProtocolVersions.MCP_2024_11_05,
|
||||
ProtocolVersions.MCP_2025_03_26, ProtocolVersions.MCP_2025_06_18, ProtocolVersions.MCP_2025_11_25);
|
||||
|
||||
private int maxResponseSize = DEFAULT_MAX_RESPONSE_SIZE;
|
||||
|
||||
/**
|
||||
* Creates a new builder with the specified base URI.
|
||||
* @param baseUri the base URI of the MCP server
|
||||
|
|
@ -840,6 +850,26 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets the maximum number of bytes read for a single inbound message, whether it
|
||||
* arrives on an SSE stream or as a JSON response body. A peer that sends a larger
|
||||
* message (or never terminates one) has its stream aborted instead of forcing the
|
||||
* transport to buffer it in memory. Defaults to 16MiB.
|
||||
*
|
||||
* <p>
|
||||
* The bound applies per message, not to the stream as a whole: a long-lived SSE
|
||||
* stream may deliver any number of messages, each up to this size. SSE field
|
||||
* framing is allowed a small amount of headroom on top of this size, so a message
|
||||
* of exactly this many bytes is still accepted.
|
||||
* @param maxResponseSize the maximum inbound message size, in bytes
|
||||
* @return this builder
|
||||
*/
|
||||
public Builder maxResponseSize(int maxResponseSize) {
|
||||
Assert.isTrue(maxResponseSize > 0, "maxResponseSize must be positive");
|
||||
this.maxResponseSize = maxResponseSize;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Construct a fresh instance of {@link HttpClientStreamableHttpTransport} using
|
||||
* the current builder configuration.
|
||||
|
|
@ -849,7 +879,7 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
HttpClient httpClient = this.clientBuilder.connectTimeout(this.connectTimeout).build();
|
||||
return new HttpClientStreamableHttpTransport(jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper,
|
||||
httpClient, requestBuilder, baseUri, endpoint, resumableStreams, openConnectionOnStartup,
|
||||
httpRequestCustomizer, supportedProtocolVersions);
|
||||
httpRequestCustomizer, supportedProtocolVersions, maxResponseSize);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,8 +5,13 @@
|
|||
package io.modelcontextprotocol.client.transport;
|
||||
|
||||
import java.net.http.HttpResponse;
|
||||
import java.net.http.HttpResponse.BodyHandler;
|
||||
import java.net.http.HttpResponse.BodySubscriber;
|
||||
import java.net.http.HttpResponse.ResponseInfo;
|
||||
import java.nio.ByteBuffer;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.CompletionStage;
|
||||
import java.util.concurrent.Flow;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.regex.Pattern;
|
||||
|
||||
|
|
@ -30,11 +35,21 @@ import reactor.core.publisher.FluxSink;
|
|||
*
|
||||
* @author Christian Tzolov
|
||||
* @author Dariusz Jędrzejczyk
|
||||
* @author Daniel Garnier-Moiroux
|
||||
*/
|
||||
class ResponseSubscribers {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(ResponseSubscribers.class);
|
||||
|
||||
/**
|
||||
* Bytes of SSE field framing a single line may carry on top of the message payload:
|
||||
* {@code "event: "} is the longest field prefix this parser recognises. Line
|
||||
* terminators are not counted, as they reset the running line length. Without this
|
||||
* allowance, an event carrying exactly the maximum message size would be rejected
|
||||
* because of the bytes the SSE wire format adds around it.
|
||||
*/
|
||||
private static final int SSE_FRAMING_OVERHEAD = "event: ".length();
|
||||
|
||||
record SseEvent(String id, String event, String data) {
|
||||
}
|
||||
|
||||
|
|
@ -54,19 +69,80 @@ class ResponseSubscribers {
|
|||
record AggregateResponseEvent(ResponseInfo responseInfo, String data) implements ResponseEvent {
|
||||
}
|
||||
|
||||
static BodySubscriber<Void> sseToBodySubscriber(ResponseInfo responseInfo, FluxSink<ResponseEvent> sink) {
|
||||
return HttpResponse.BodySubscribers
|
||||
.fromLineSubscriber(FlowAdapters.toFlowSubscriber(new SseLineSubscriber(responseInfo, sink)));
|
||||
/**
|
||||
* Creates a {@link BodySubscriber} that parses a Server-Sent Events stream, bounding
|
||||
* how much memory a single inbound message may occupy. Both the size of an individual
|
||||
* line (as read off the wire before a terminator is seen) and the accumulated size of
|
||||
* a multi-line SSE event are capped at {@code maxSize}; a peer exceeding either limit
|
||||
* has its stream aborted instead of forcing the transport to buffer it in memory. The
|
||||
* line bound is allowed {@link #SSE_FRAMING_OVERHEAD} extra bytes so that the SSE
|
||||
* framing around a payload does not count against the payload's own budget.
|
||||
* @param responseInfo the HTTP response information
|
||||
* @param sink the sink to emit parsed events to
|
||||
* @param maxSize the maximum number of bytes read for a single inbound message
|
||||
*/
|
||||
static BodySubscriber<Void> sseToBodySubscriber(ResponseInfo responseInfo, FluxSink<ResponseEvent> sink,
|
||||
int maxSize) {
|
||||
BodySubscriber<Void> lineSubscriber = HttpResponse.BodySubscribers
|
||||
.fromLineSubscriber(FlowAdapters.toFlowSubscriber(new SseLineSubscriber(responseInfo, sink, maxSize)));
|
||||
return new BoundedLineBodySubscriber(lineSubscriber, plusFramingOverhead(maxSize));
|
||||
}
|
||||
|
||||
static BodySubscriber<Void> aggregateBodySubscriber(ResponseInfo responseInfo, FluxSink<ResponseEvent> sink) {
|
||||
return HttpResponse.BodySubscribers
|
||||
.fromLineSubscriber(FlowAdapters.toFlowSubscriber(new AggregateSubscriber(responseInfo, sink)));
|
||||
/**
|
||||
* Adds {@link #SSE_FRAMING_OVERHEAD} to {@code maxSize}, saturating at
|
||||
* {@link Integer#MAX_VALUE} rather than overflowing into a negative bound that would
|
||||
* reject everything.
|
||||
*/
|
||||
private static int plusFramingOverhead(int maxSize) {
|
||||
return maxSize > Integer.MAX_VALUE - SSE_FRAMING_OVERHEAD ? Integer.MAX_VALUE : maxSize + SSE_FRAMING_OVERHEAD;
|
||||
}
|
||||
|
||||
static BodySubscriber<Void> bodilessBodySubscriber(ResponseInfo responseInfo, FluxSink<ResponseEvent> sink) {
|
||||
return HttpResponse.BodySubscribers
|
||||
/**
|
||||
* Creates a {@link BodySubscriber} that aggregates the whole response body into a
|
||||
* single event, bounding how much memory it may occupy. Both the size of an
|
||||
* individual line (as read off the wire before a terminator is seen) and the total
|
||||
* accumulated body are capped at {@code maxSize}; a peer exceeding either limit has
|
||||
* its response aborted instead of forcing the transport to buffer it in memory.
|
||||
* @param responseInfo the HTTP response information
|
||||
* @param sink the sink to emit the aggregated event to
|
||||
* @param maxSize the maximum number of bytes read for the response body
|
||||
*/
|
||||
static BodySubscriber<Void> aggregateBodySubscriber(ResponseInfo responseInfo, FluxSink<ResponseEvent> sink,
|
||||
int maxSize) {
|
||||
BodySubscriber<Void> lineSubscriber = HttpResponse.BodySubscribers
|
||||
.fromLineSubscriber(FlowAdapters.toFlowSubscriber(new AggregateSubscriber(responseInfo, sink, maxSize)));
|
||||
return new BoundedLineBodySubscriber(lineSubscriber, maxSize);
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a {@link BodySubscriber} that discards the response body, bounding how much
|
||||
* memory reading it may occupy. The body is discarded as it arrives, but the
|
||||
* underlying line subscriber still buffers each line before handing it over, so a
|
||||
* peer sending a line longer than {@code maxSize} has its response aborted.
|
||||
* @param responseInfo the HTTP response information
|
||||
* @param sink the sink to emit the completion event to
|
||||
* @param maxSize the maximum number of bytes read for a single line
|
||||
*/
|
||||
static BodySubscriber<Void> bodilessBodySubscriber(ResponseInfo responseInfo, FluxSink<ResponseEvent> sink,
|
||||
int maxSize) {
|
||||
BodySubscriber<Void> lineSubscriber = HttpResponse.BodySubscribers
|
||||
.fromLineSubscriber(FlowAdapters.toFlowSubscriber(new BodilessResponseLineSubscriber(responseInfo, sink)));
|
||||
return new BoundedLineBodySubscriber(lineSubscriber, maxSize);
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a {@link BodyHandler} that reads the response body into a string, bounding
|
||||
* how much memory it may occupy. A peer sending more than {@code maxSize} bytes has
|
||||
* its response aborted instead of forcing the transport to buffer it in memory.
|
||||
*
|
||||
* <p>
|
||||
* Decoding matches {@link HttpResponse.BodyHandlers#ofString()}, including its
|
||||
* handling of the charset declared in the {@code Content-Type} header.
|
||||
* @param maxSize the maximum number of bytes read for the response body
|
||||
*/
|
||||
static BodyHandler<String> boundedStringBodyHandler(int maxSize) {
|
||||
BodyHandler<String> delegate = HttpResponse.BodyHandlers.ofString();
|
||||
return responseInfo -> new BoundedTotalBodySubscriber<>(delegate.apply(responseInfo), maxSize);
|
||||
}
|
||||
|
||||
static class SseLineSubscriber extends BaseSubscriber<String> {
|
||||
|
|
@ -112,18 +188,30 @@ class ResponseSubscribers {
|
|||
*/
|
||||
private ResponseInfo responseInfo;
|
||||
|
||||
/**
|
||||
* The maximum number of bytes that may accumulate for a single SSE event. A peer
|
||||
* that never terminates an event (e.g. an endless stream of {@code data:} lines)
|
||||
* has its stream aborted instead of exhausting memory. The accumulated data is
|
||||
* measured in characters, which for UTF-8 is never more than the number of bytes
|
||||
* it was decoded from.
|
||||
*/
|
||||
private final int maxSize;
|
||||
|
||||
/**
|
||||
* Creates a new LineSubscriber that will emit parsed SSE events to the provided
|
||||
* sink.
|
||||
* @param sink the {@link FluxSink} to emit parsed {@link ResponseEvent} objects
|
||||
* to
|
||||
* @param maxSize the maximum number of bytes that may accumulate for a single SSE
|
||||
* event
|
||||
*/
|
||||
public SseLineSubscriber(ResponseInfo responseInfo, FluxSink<ResponseEvent> sink) {
|
||||
public SseLineSubscriber(ResponseInfo responseInfo, FluxSink<ResponseEvent> sink, int maxSize) {
|
||||
this.sink = sink;
|
||||
this.eventBuilder = new StringBuilder();
|
||||
this.currentEventId = new AtomicReference<>();
|
||||
this.currentEventType = new AtomicReference<>();
|
||||
this.responseInfo = responseInfo;
|
||||
this.maxSize = maxSize;
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
@ -155,7 +243,18 @@ class ResponseSubscribers {
|
|||
if (line.startsWith("data:")) {
|
||||
var matcher = EVENT_DATA_PATTERN.matcher(line);
|
||||
if (matcher.find()) {
|
||||
this.eventBuilder.append(matcher.group(1).trim()).append("\n");
|
||||
String data = matcher.group(1).trim();
|
||||
// Measured before appending, so that an event carrying exactly
|
||||
// maxSize of data is accepted: the trailing separator below is
|
||||
// stripped again before the event is emitted.
|
||||
if (this.eventBuilder.length() + data.length() > this.maxSize) {
|
||||
upstream().cancel();
|
||||
this.sink.error(
|
||||
new McpTransportException("Inbound SSE event exceeds the maximum allowed size of "
|
||||
+ this.maxSize + " bytes"));
|
||||
return;
|
||||
}
|
||||
this.eventBuilder.append(data).append("\n");
|
||||
}
|
||||
upstream().request(1);
|
||||
}
|
||||
|
|
@ -225,15 +324,26 @@ class ResponseSubscribers {
|
|||
|
||||
volatile boolean hasRequestedDemand = false;
|
||||
|
||||
/**
|
||||
* The maximum number of bytes that may accumulate for the aggregated response
|
||||
* body. A peer that sends a larger body has its response aborted instead of
|
||||
* exhausting memory. The accumulated body is measured in characters, which for
|
||||
* UTF-8 is never more than the number of bytes it was decoded from.
|
||||
*/
|
||||
private final int maxSize;
|
||||
|
||||
/**
|
||||
* Creates a new JsonLineSubscriber that will emit parsed JSON-RPC messages.
|
||||
* @param sink the {@link FluxSink} to emit parsed {@link ResponseEvent} objects
|
||||
* to
|
||||
* @param maxSize the maximum number of bytes that may accumulate for the
|
||||
* aggregated response body
|
||||
*/
|
||||
public AggregateSubscriber(ResponseInfo responseInfo, FluxSink<ResponseEvent> sink) {
|
||||
public AggregateSubscriber(ResponseInfo responseInfo, FluxSink<ResponseEvent> sink, int maxSize) {
|
||||
this.sink = sink;
|
||||
this.eventBuilder = new StringBuilder();
|
||||
this.responseInfo = responseInfo;
|
||||
this.maxSize = maxSize;
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
@ -252,6 +362,15 @@ class ResponseSubscribers {
|
|||
|
||||
@Override
|
||||
protected void hookOnNext(String line) {
|
||||
// Measured before appending, so that a body of exactly maxSize is accepted.
|
||||
// The separator this adds back for each line stands in for the terminator the
|
||||
// peer sent, which the line subscriber has already stripped.
|
||||
if (this.eventBuilder.length() + line.length() > this.maxSize) {
|
||||
upstream().cancel();
|
||||
this.sink.error(new McpTransportException(
|
||||
"Inbound response body exceeds the maximum allowed size of " + this.maxSize + " bytes"));
|
||||
return;
|
||||
}
|
||||
this.eventBuilder.append(line).append("\n");
|
||||
}
|
||||
|
||||
|
|
@ -324,4 +443,177 @@ class ResponseSubscribers {
|
|||
|
||||
}
|
||||
|
||||
/**
|
||||
* Base for {@link BodySubscriber} wrappers that transparently forward the response
|
||||
* body to a delegate, but abort it once the peer exceeds a size bound.
|
||||
*
|
||||
* <p>
|
||||
* Aborting cancels the upstream subscription, which closes the connection, and
|
||||
* signals a {@link McpTransportException} to the delegate so the failure surfaces
|
||||
* both through the body's {@link CompletionStage} and through any sink the delegate
|
||||
* feeds.
|
||||
*/
|
||||
abstract static class BoundedBodySubscriber<T> implements BodySubscriber<T> {
|
||||
|
||||
private final BodySubscriber<T> delegate;
|
||||
|
||||
protected final int maxSize;
|
||||
|
||||
/**
|
||||
* What the bound applies to, e.g. {@code "Inbound line"}, used to build the
|
||||
* failure message.
|
||||
*/
|
||||
private final String boundedEntity;
|
||||
|
||||
private Flow.Subscription subscription;
|
||||
|
||||
private volatile boolean done = false;
|
||||
|
||||
BoundedBodySubscriber(BodySubscriber<T> delegate, int maxSize, String boundedEntity) {
|
||||
this.delegate = delegate;
|
||||
this.maxSize = maxSize;
|
||||
this.boundedEntity = boundedEntity;
|
||||
}
|
||||
|
||||
@Override
|
||||
public CompletionStage<T> getBody() {
|
||||
return this.delegate.getBody();
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onSubscribe(Flow.Subscription subscription) {
|
||||
this.subscription = subscription;
|
||||
this.delegate.onSubscribe(subscription);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onNext(List<ByteBuffer> buffers) {
|
||||
if (this.done) {
|
||||
return;
|
||||
}
|
||||
for (ByteBuffer buffer : buffers) {
|
||||
if (!checkSize(buffer)) {
|
||||
this.done = true;
|
||||
this.subscription.cancel();
|
||||
this.delegate.onError(new McpTransportException(
|
||||
this.boundedEntity + " exceeds the maximum allowed size of " + this.maxSize + " bytes"));
|
||||
return;
|
||||
}
|
||||
}
|
||||
this.delegate.onNext(buffers);
|
||||
}
|
||||
|
||||
/**
|
||||
* Accounts for the bytes in {@code buffer}, which must be inspected with absolute
|
||||
* reads only so the delegate still sees the original position.
|
||||
* @param buffer the buffer about to be handed to the delegate
|
||||
* @return {@code true} to accept the buffer, or {@code false} to abort the
|
||||
* response because the bound has been exceeded
|
||||
*/
|
||||
protected abstract boolean checkSize(ByteBuffer buffer);
|
||||
|
||||
@Override
|
||||
public void onError(Throwable throwable) {
|
||||
if (this.done) {
|
||||
return;
|
||||
}
|
||||
this.done = true;
|
||||
this.delegate.onError(throwable);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onComplete() {
|
||||
if (this.done) {
|
||||
return;
|
||||
}
|
||||
this.done = true;
|
||||
this.delegate.onComplete();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* A {@link BoundedBodySubscriber} that aborts the response once a single line (a run
|
||||
* of bytes with no CR/LF terminator) exceeds {@code maxSize} bytes.
|
||||
*
|
||||
* <p>
|
||||
* {@link HttpResponse.BodySubscribers#fromLineSubscriber} buffers characters until it
|
||||
* encounters a line terminator, so a peer that never terminates a line (or sends an
|
||||
* enormous one) would force the transport to buffer it in memory. This wrapper counts
|
||||
* bytes as they arrive off the wire and cancels the subscription before that buffer
|
||||
* can grow without bound.
|
||||
*/
|
||||
static final class BoundedLineBodySubscriber extends BoundedBodySubscriber<Void> {
|
||||
|
||||
private long bytesSinceLineTerminator = 0;
|
||||
|
||||
BoundedLineBodySubscriber(BodySubscriber<Void> delegate, int maxSize) {
|
||||
super(delegate, maxSize, "Inbound line");
|
||||
}
|
||||
|
||||
@Override
|
||||
protected boolean checkSize(ByteBuffer buffer) {
|
||||
int position = buffer.position();
|
||||
int limit = buffer.limit();
|
||||
if (position == limit) {
|
||||
return true;
|
||||
}
|
||||
if (this.bytesSinceLineTerminator + (limit - position) <= this.maxSize) {
|
||||
// No line ending in this buffer can exceed the limit, because there are
|
||||
// not enough bytes since the last terminator for one to. Only the
|
||||
// trailing (still unterminated) run matters, so scan back to the last
|
||||
// terminator instead of walking every byte.
|
||||
this.bytesSinceLineTerminator = lengthOfTrailingRun(buffer, position, limit);
|
||||
return true;
|
||||
}
|
||||
// The limit is within reach, so account for every line exactly.
|
||||
for (int i = position; i < limit; i++) {
|
||||
byte b = buffer.get(i);
|
||||
if (b == '\n' || b == '\r') {
|
||||
this.bytesSinceLineTerminator = 0;
|
||||
}
|
||||
else if (++this.bytesSinceLineTerminator > this.maxSize) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the number of bytes after the last line terminator in the buffer, or
|
||||
* the whole span added to the running count when the buffer holds no terminator.
|
||||
*/
|
||||
private long lengthOfTrailingRun(ByteBuffer buffer, int position, int limit) {
|
||||
for (int i = limit - 1; i >= position; i--) {
|
||||
byte b = buffer.get(i);
|
||||
if (b == '\n' || b == '\r') {
|
||||
return limit - 1 - i;
|
||||
}
|
||||
}
|
||||
return this.bytesSinceLineTerminator + (limit - position);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* A {@link BoundedBodySubscriber} that aborts the response once the body as a whole
|
||||
* exceeds {@code maxSize} bytes. Suitable for delegates that aggregate the entire
|
||||
* body in memory, such as {@link HttpResponse.BodyHandlers#ofString()}.
|
||||
*/
|
||||
static final class BoundedTotalBodySubscriber<T> extends BoundedBodySubscriber<T> {
|
||||
|
||||
private long totalBytes = 0;
|
||||
|
||||
BoundedTotalBodySubscriber(BodySubscriber<T> delegate, int maxSize) {
|
||||
super(delegate, maxSize, "Inbound response body");
|
||||
}
|
||||
|
||||
@Override
|
||||
protected boolean checkSize(ByteBuffer buffer) {
|
||||
this.totalBytes += buffer.remaining();
|
||||
return this.totalBytes <= this.maxSize;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,334 @@
|
|||
/*
|
||||
* Copyright 2024-2026 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client.transport;
|
||||
|
||||
import java.net.http.HttpResponse.BodySubscriber;
|
||||
import java.nio.ByteBuffer;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.CompletableFuture;
|
||||
import java.util.concurrent.CompletionStage;
|
||||
import java.util.concurrent.Flow;
|
||||
|
||||
import io.modelcontextprotocol.client.transport.ResponseSubscribers.BoundedLineBodySubscriber;
|
||||
import io.modelcontextprotocol.client.transport.ResponseSubscribers.BoundedTotalBodySubscriber;
|
||||
import io.modelcontextprotocol.spec.McpTransportException;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* Tests the size accounting in {@link ResponseSubscribers.BoundedBodySubscriber} and its
|
||||
* two implementations. These bound how much of a response the transport will buffer, so
|
||||
* the accounting is exercised directly rather than only through a live HTTP exchange:
|
||||
* buffer boundaries, line terminators split across buffers, and the exact limit are all
|
||||
* places where an off-by-one either lets a peer past the bound or rejects a legitimate
|
||||
* message.
|
||||
*
|
||||
* @author Daniel Garnier-Moiroux
|
||||
*/
|
||||
class BoundedBodySubscriberTests {
|
||||
|
||||
private static final int MAX_SIZE = 16;
|
||||
|
||||
private final RecordingBodySubscriber delegate = new RecordingBodySubscriber();
|
||||
|
||||
private final RecordingSubscription subscription = new RecordingSubscription();
|
||||
|
||||
// --- BoundedLineBodySubscriber: per-line accounting -----------------------
|
||||
|
||||
@Test
|
||||
void lineSubscriberAcceptsEmptyBuffer() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
|
||||
assertThat(subscriber.checkSize(buffer(""))).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void lineSubscriberAcceptsLineOfExactlyMaxSize() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
|
||||
assertThat(subscriber.checkSize(buffer("a".repeat(MAX_SIZE)))).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void lineSubscriberRejectsLineOneByteOverMaxSize() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
|
||||
assertThat(subscriber.checkSize(buffer("a".repeat(MAX_SIZE + 1)))).isFalse();
|
||||
}
|
||||
|
||||
@Test
|
||||
void lineSubscriberAccumulatesAcrossBuffers() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
|
||||
assertThat(subscriber.checkSize(buffer("a".repeat(10)))).isTrue();
|
||||
assertThat(subscriber.checkSize(buffer("a".repeat(6)))).isTrue();
|
||||
// 17th byte of the same unterminated line.
|
||||
assertThat(subscriber.checkSize(buffer("a"))).isFalse();
|
||||
}
|
||||
|
||||
@Test
|
||||
void lineSubscriberAcceptsUnboundedTotalOfTerminatedLines() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
|
||||
// Far more than MAX_SIZE in total, but no single line comes close to it.
|
||||
for (int i = 0; i < 100; i++) {
|
||||
assertThat(subscriber.checkSize(buffer("aaaa\n"))).isTrue();
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void lineSubscriberResetsOnTerminatorAtEndOfBuffer() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
|
||||
assertThat(subscriber.checkSize(buffer("a".repeat(10) + "\n"))).isTrue();
|
||||
// A fresh line, so the previous 10 bytes must not count towards it.
|
||||
assertThat(subscriber.checkSize(buffer("a".repeat(MAX_SIZE)))).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void lineSubscriberResetsOnTerminatorAtStartOfBuffer() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
|
||||
assertThat(subscriber.checkSize(buffer("a".repeat(MAX_SIZE)))).isTrue();
|
||||
assertThat(subscriber.checkSize(buffer("\n" + "a".repeat(MAX_SIZE)))).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void lineSubscriberHandlesCrLfSplitAcrossBuffers() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
|
||||
assertThat(subscriber.checkSize(buffer("a".repeat(12) + "\r"))).isTrue();
|
||||
assertThat(subscriber.checkSize(buffer("\n" + "a".repeat(MAX_SIZE)))).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void lineSubscriberAcceptsBufferLargerThanMaxSizeHoldingOnlyShortLines() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
|
||||
// Forces the exact per-byte accounting path: the buffer alone is well over the
|
||||
// limit, yet every line in it is legitimate.
|
||||
assertThat(subscriber.checkSize(buffer("aaaa\n".repeat(20)))).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void lineSubscriberRejectsRunSpanningManyBuffers() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
|
||||
boolean accepted = true;
|
||||
for (int i = 0; i < 10 && accepted; i++) {
|
||||
accepted = subscriber.checkSize(buffer("aa"));
|
||||
}
|
||||
|
||||
assertThat(accepted).isFalse();
|
||||
}
|
||||
|
||||
@Test
|
||||
void lineSubscriberOnlyCountsFromTheBufferPosition() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
ByteBuffer partiallyConsumed = buffer("a".repeat(MAX_SIZE * 2));
|
||||
partiallyConsumed.position(MAX_SIZE * 2 - 4);
|
||||
|
||||
assertThat(subscriber.checkSize(partiallyConsumed)).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void lineSubscriberDoesNotConsumeTheBuffer() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
ByteBuffer buffer = buffer("aaaa\nbbbb");
|
||||
buffer.position(2);
|
||||
|
||||
subscriber.checkSize(buffer);
|
||||
|
||||
assertThat(buffer.position()).isEqualTo(2);
|
||||
assertThat(buffer.limit()).isEqualTo(9);
|
||||
}
|
||||
|
||||
// --- BoundedTotalBodySubscriber: whole-body accounting --------------------
|
||||
|
||||
@Test
|
||||
void totalSubscriberAcceptsBodyOfExactlyMaxSize() {
|
||||
BoundedTotalBodySubscriber<Void> subscriber = totalSubscriber();
|
||||
|
||||
assertThat(subscriber.checkSize(buffer("a".repeat(8)))).isTrue();
|
||||
assertThat(subscriber.checkSize(buffer("a".repeat(8)))).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void totalSubscriberRejectsBodyOneByteOverMaxSize() {
|
||||
BoundedTotalBodySubscriber<Void> subscriber = totalSubscriber();
|
||||
|
||||
assertThat(subscriber.checkSize(buffer("a".repeat(8)))).isTrue();
|
||||
assertThat(subscriber.checkSize(buffer("a".repeat(9)))).isFalse();
|
||||
}
|
||||
|
||||
@Test
|
||||
void totalSubscriberIsNotResetByLineTerminators() {
|
||||
BoundedTotalBodySubscriber<Void> subscriber = totalSubscriber();
|
||||
|
||||
// Unlike the per-line bound, terminated lines still count towards the total.
|
||||
assertThat(subscriber.checkSize(buffer("aaaa\n".repeat(3)))).isTrue();
|
||||
assertThat(subscriber.checkSize(buffer("aaaa\n"))).isFalse();
|
||||
}
|
||||
|
||||
@Test
|
||||
void totalSubscriberOnlyCountsFromTheBufferPosition() {
|
||||
BoundedTotalBodySubscriber<Void> subscriber = totalSubscriber();
|
||||
ByteBuffer partiallyConsumed = buffer("a".repeat(MAX_SIZE * 2));
|
||||
partiallyConsumed.position(MAX_SIZE);
|
||||
|
||||
assertThat(subscriber.checkSize(partiallyConsumed)).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void totalSubscriberDoesNotConsumeTheBuffer() {
|
||||
BoundedTotalBodySubscriber<Void> subscriber = totalSubscriber();
|
||||
ByteBuffer buffer = buffer("aaaa");
|
||||
|
||||
subscriber.checkSize(buffer);
|
||||
|
||||
assertThat(buffer.position()).isZero();
|
||||
assertThat(buffer.remaining()).isEqualTo(4);
|
||||
}
|
||||
|
||||
// --- onNext: what a failed check does ------------------------------------
|
||||
|
||||
@Test
|
||||
void forwardsBuffersWhileWithinBounds() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
List<ByteBuffer> buffers = List.of(buffer("aaaa\n"), buffer("bbbb\n"));
|
||||
|
||||
subscriber.onNext(buffers);
|
||||
|
||||
assertThat(this.delegate.received).containsExactly(buffers);
|
||||
assertThat(this.delegate.error).isNull();
|
||||
assertThat(this.subscription.cancellations).isZero();
|
||||
}
|
||||
|
||||
@Test
|
||||
void abortsTheResponseWhenTheLineBoundIsExceeded() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
|
||||
subscriber.onNext(List.of(buffer("a".repeat(MAX_SIZE + 1))));
|
||||
|
||||
assertThat(this.subscription.cancellations).isEqualTo(1);
|
||||
assertThat(this.delegate.received).isEmpty();
|
||||
assertThat(this.delegate.error).isInstanceOf(McpTransportException.class)
|
||||
.hasMessage("Inbound line exceeds the maximum allowed size of " + MAX_SIZE + " bytes");
|
||||
}
|
||||
|
||||
@Test
|
||||
void abortsTheResponseWhenTheTotalBoundIsExceeded() {
|
||||
BoundedTotalBodySubscriber<Void> subscriber = totalSubscriber();
|
||||
|
||||
subscriber.onNext(List.of(buffer("a".repeat(MAX_SIZE + 1))));
|
||||
|
||||
assertThat(this.subscription.cancellations).isEqualTo(1);
|
||||
assertThat(this.delegate.received).isEmpty();
|
||||
assertThat(this.delegate.error).isInstanceOf(McpTransportException.class)
|
||||
.hasMessage("Inbound response body exceeds the maximum allowed size of " + MAX_SIZE + " bytes");
|
||||
}
|
||||
|
||||
@Test
|
||||
void withholdsTheWholeListWhenALaterBufferExceedsTheBound() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
|
||||
subscriber.onNext(List.of(buffer("a".repeat(8)), buffer("a".repeat(9))));
|
||||
|
||||
assertThat(this.delegate.received).isEmpty();
|
||||
assertThat(this.delegate.error).isInstanceOf(McpTransportException.class);
|
||||
}
|
||||
|
||||
@Test
|
||||
void ignoresFurtherSignalsOnceAborted() {
|
||||
BoundedLineBodySubscriber subscriber = lineSubscriber();
|
||||
subscriber.onNext(List.of(buffer("a".repeat(MAX_SIZE + 1))));
|
||||
Throwable firstError = this.delegate.error;
|
||||
|
||||
// The HTTP client may still signal after the subscription is cancelled.
|
||||
subscriber.onNext(List.of(buffer("aaaa")));
|
||||
subscriber.onError(new RuntimeException("late failure"));
|
||||
subscriber.onComplete();
|
||||
|
||||
assertThat(this.subscription.cancellations).isEqualTo(1);
|
||||
assertThat(this.delegate.received).isEmpty();
|
||||
assertThat(this.delegate.error).isSameAs(firstError);
|
||||
assertThat(this.delegate.completed).isFalse();
|
||||
}
|
||||
|
||||
// --- fixtures ------------------------------------------------------------
|
||||
|
||||
private BoundedLineBodySubscriber lineSubscriber() {
|
||||
BoundedLineBodySubscriber subscriber = new BoundedLineBodySubscriber(this.delegate, MAX_SIZE);
|
||||
subscriber.onSubscribe(this.subscription);
|
||||
return subscriber;
|
||||
}
|
||||
|
||||
private BoundedTotalBodySubscriber<Void> totalSubscriber() {
|
||||
BoundedTotalBodySubscriber<Void> subscriber = new BoundedTotalBodySubscriber<>(this.delegate, MAX_SIZE);
|
||||
subscriber.onSubscribe(this.subscription);
|
||||
return subscriber;
|
||||
}
|
||||
|
||||
private static ByteBuffer buffer(String content) {
|
||||
return ByteBuffer.wrap(content.getBytes(StandardCharsets.US_ASCII));
|
||||
}
|
||||
|
||||
private static final class RecordingBodySubscriber implements BodySubscriber<Void> {
|
||||
|
||||
private final List<List<ByteBuffer>> received = new ArrayList<>();
|
||||
|
||||
private final CompletableFuture<Void> body = new CompletableFuture<>();
|
||||
|
||||
private Throwable error;
|
||||
|
||||
private boolean completed;
|
||||
|
||||
@Override
|
||||
public CompletionStage<Void> getBody() {
|
||||
return this.body;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onSubscribe(Flow.Subscription subscription) {
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onNext(List<ByteBuffer> item) {
|
||||
this.received.add(item);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onError(Throwable throwable) {
|
||||
this.error = throwable;
|
||||
this.body.completeExceptionally(throwable);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onComplete() {
|
||||
this.completed = true;
|
||||
this.body.complete(null);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
private static final class RecordingSubscription implements Flow.Subscription {
|
||||
|
||||
private int cancellations;
|
||||
|
||||
@Override
|
||||
public void request(long n) {
|
||||
}
|
||||
|
||||
@Override
|
||||
public void cancel() {
|
||||
this.cancellations++;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,123 @@
|
|||
/*
|
||||
* Copyright 2024-2026 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client.transport;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.OutputStream;
|
||||
import java.net.InetSocketAddress;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.concurrent.Executors;
|
||||
|
||||
import com.sun.net.httpserver.HttpExchange;
|
||||
import com.sun.net.httpserver.HttpServer;
|
||||
import io.modelcontextprotocol.server.transport.TomcatTestUtil;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
|
||||
/**
|
||||
* Shared fixture for the transport bounded-read tests: a bare {@link HttpServer} whose
|
||||
* response body is written by a per-test {@link Responder}.
|
||||
*
|
||||
* @author Daniel Garnier-Moiroux
|
||||
*/
|
||||
abstract class HttpClientBoundedReadTestSupport {
|
||||
|
||||
protected static final int MAX_SIZE = 1024;
|
||||
|
||||
private HttpServer server;
|
||||
|
||||
protected String host;
|
||||
|
||||
/**
|
||||
* Writes a response body to the exchange.
|
||||
*/
|
||||
@FunctionalInterface
|
||||
protected interface Responder {
|
||||
|
||||
void respond(OutputStream body) throws IOException;
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* The path the transport under test talks to.
|
||||
*/
|
||||
protected abstract String endpoint();
|
||||
|
||||
@BeforeEach
|
||||
void startServer() throws IOException {
|
||||
int port = TomcatTestUtil.findAvailablePort();
|
||||
this.host = "http://localhost:" + port;
|
||||
this.server = HttpServer.create(new InetSocketAddress(port), 0);
|
||||
this.server.setExecutor(Executors.newCachedThreadPool());
|
||||
this.server.start();
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
void stopServer() {
|
||||
if (this.server != null) {
|
||||
this.server.stop(0);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers a handler that responds with the given content type and body.
|
||||
*/
|
||||
protected void respondWith(String path, String contentType, Responder responder) {
|
||||
this.server.createContext(path, exchange -> {
|
||||
exchange.getResponseHeaders().set("Content-Type", contentType);
|
||||
exchange.sendResponseHeaders(200, 0);
|
||||
try (OutputStream body = exchange.getResponseBody()) {
|
||||
responder.respond(body);
|
||||
}
|
||||
catch (IOException ignored) {
|
||||
// The client aborts the response once the limit is exceeded, which closes
|
||||
// the connection and makes further writes fail. That is the behaviour
|
||||
// under test.
|
||||
}
|
||||
finally {
|
||||
exchange.close();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* A responder that writes {@code chunks} blocks of {@code 'a'} with no line
|
||||
* terminator anywhere, so nothing downstream can ever flush a line.
|
||||
*/
|
||||
protected static Responder unterminatedLine(int chunks) {
|
||||
return body -> {
|
||||
byte[] chunk = new byte[MAX_SIZE];
|
||||
java.util.Arrays.fill(chunk, (byte) 'a');
|
||||
for (int i = 0; i < chunks; i++) {
|
||||
body.write(chunk);
|
||||
body.flush();
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* A responder that writes enough short, properly terminated lines to exceed the limit
|
||||
* in aggregate.
|
||||
*/
|
||||
protected static Responder manyShortLines(String prefix) {
|
||||
return body -> {
|
||||
byte[] line = (prefix + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n").getBytes(StandardCharsets.UTF_8);
|
||||
for (int i = 0; i < (MAX_SIZE / line.length) + 64; i++) {
|
||||
body.write(line);
|
||||
body.flush();
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
protected static boolean messageContains(Throwable t, String expected) {
|
||||
for (Throwable current = t; current != null; current = current.getCause()) {
|
||||
if (current.getMessage() != null && current.getMessage().contains(expected)) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,108 @@
|
|||
/*
|
||||
* Copyright 2024-2026 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client.transport;
|
||||
|
||||
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.function.Function;
|
||||
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.JSONRPCRequest;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
/**
|
||||
* Verifies that {@link HttpClientSseClientTransport} bounds the amount of memory a single
|
||||
* inbound message can occupy, so a malicious or buggy server cannot exhaust the client's
|
||||
* memory by streaming an unterminated line, an endless event, or an oversized response to
|
||||
* a posted message.
|
||||
*
|
||||
* @author Daniel Garnier-Moiroux
|
||||
*/
|
||||
@Timeout(15)
|
||||
class HttpClientSseClientTransportBoundedReadTests extends HttpClientBoundedReadTestSupport {
|
||||
|
||||
private static final String MESSAGE_ENDPOINT = "/message";
|
||||
|
||||
private final CountDownLatch keepStreamOpen = new CountDownLatch(1);
|
||||
|
||||
@Override
|
||||
protected String endpoint() {
|
||||
return "/sse";
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
void releaseStream() {
|
||||
this.keepStreamOpen.countDown();
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldRejectSingleLineExceedingMaxSize() {
|
||||
// A line that never terminates, so the line buffer underneath the SSE parser
|
||||
// would grow without limit before any event could be flushed.
|
||||
respondWith(endpoint(), "text/event-stream", unterminatedLine(8));
|
||||
|
||||
StepVerifier.create(connect())
|
||||
.verifyErrorMatches(t -> messageContains(t, "Inbound line exceeds the maximum allowed size"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldRejectEventExceedingMaxSize() {
|
||||
// Many short, terminated "data:" lines with no blank line to end the event. Each
|
||||
// line is small, but the accumulated event data would grow without limit.
|
||||
respondWith(endpoint(), "text/event-stream", manyShortLines("data:"));
|
||||
|
||||
StepVerifier.create(connect())
|
||||
.verifyErrorMatches(t -> messageContains(t, "Inbound SSE event exceeds the maximum allowed size"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldRejectPostResponseExceedingMaxSize() throws Exception {
|
||||
// The response to a posted message is read into a string in full, so an oversized
|
||||
// one must abort rather than accumulate.
|
||||
respondWith(endpoint(), "text/event-stream", body -> {
|
||||
body.write(("event:endpoint\ndata:" + MESSAGE_ENDPOINT + "\n\n").getBytes(StandardCharsets.UTF_8));
|
||||
body.flush();
|
||||
awaitTeardown();
|
||||
});
|
||||
respondWith(MESSAGE_ENDPOINT, "application/json", unterminatedLine(64));
|
||||
|
||||
HttpClientSseClientTransport transport = transport();
|
||||
transport.connect(Function.identity()).block(Duration.ofSeconds(5));
|
||||
|
||||
StepVerifier.create(sendMessage(transport))
|
||||
.verifyErrorMatches(t -> messageContains(t, "Inbound response body exceeds the maximum allowed size"));
|
||||
}
|
||||
|
||||
private void awaitTeardown() {
|
||||
try {
|
||||
this.keepStreamOpen.await(10, TimeUnit.SECONDS);
|
||||
}
|
||||
catch (InterruptedException ex) {
|
||||
Thread.currentThread().interrupt();
|
||||
}
|
||||
}
|
||||
|
||||
private HttpClientSseClientTransport transport() {
|
||||
return HttpClientSseClientTransport.builder(this.host).maxResponseSize(MAX_SIZE).build();
|
||||
}
|
||||
|
||||
private Mono<Void> connect() {
|
||||
return transport().connect(Function.identity());
|
||||
}
|
||||
|
||||
private Mono<Void> sendMessage(HttpClientSseClientTransport transport) {
|
||||
JSONRPCRequest request = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
return transport.sendMessage(request);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -80,7 +80,7 @@ class HttpClientSseClientTransportTests {
|
|||
SseMessageEndpointValidator sseMessageEndpointValidator) {
|
||||
super(HttpClient.newBuilder().version(HttpClient.Version.HTTP_1_1).build(),
|
||||
HttpRequest.newBuilder().header("Content-Type", "application/json"), baseUri, "/sse", JSON_MAPPER,
|
||||
McpAsyncHttpClientRequestCustomizer.NOOP, sseMessageEndpointValidator);
|
||||
McpAsyncHttpClientRequestCustomizer.NOOP, sseMessageEndpointValidator, 16 * 1024 * 1024);
|
||||
}
|
||||
|
||||
public int getInboundMessageCount() {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,111 @@
|
|||
/*
|
||||
* Copyright 2024-2026 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client.transport;
|
||||
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.Map;
|
||||
import java.util.function.Function;
|
||||
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.JSONRPCRequest;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
/**
|
||||
* Verifies that {@link HttpClientStreamableHttpTransport} bounds the amount of memory a
|
||||
* single inbound message can occupy, so a malicious or buggy server cannot exhaust the
|
||||
* client's memory by streaming an unterminated line or an endless event.
|
||||
*
|
||||
* @author Daniel Garnier-Moiroux
|
||||
*/
|
||||
@Timeout(15)
|
||||
class HttpClientStreamableHttpTransportBoundedReadTests extends HttpClientBoundedReadTestSupport {
|
||||
|
||||
@Override
|
||||
protected String endpoint() {
|
||||
return "/mcp";
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldRejectSingleLineExceedingMaxSize() {
|
||||
// A line that never terminates, so the line buffer underneath the SSE parser
|
||||
// would grow without limit before any event could be flushed.
|
||||
respondWith(endpoint(), "text/event-stream", unterminatedLine(8));
|
||||
|
||||
StepVerifier.create(sendMessage())
|
||||
.verifyErrorMatches(t -> messageContains(t, "Inbound line exceeds the maximum allowed size"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldRejectEventExceedingMaxSize() {
|
||||
// Many short, terminated "data:" lines with no blank line to end the event. Each
|
||||
// line is small, but the accumulated event data would grow without limit.
|
||||
respondWith(endpoint(), "text/event-stream", manyShortLines("data:"));
|
||||
|
||||
StepVerifier.create(sendMessage())
|
||||
.verifyErrorMatches(t -> messageContains(t, "Inbound SSE event exceeds the maximum allowed size"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldRejectJsonResponseExceedingMaxSize() {
|
||||
// A multi-line application/json response whose total size exceeds the limit. Each
|
||||
// line is small, but the aggregated body would grow without limit.
|
||||
respondWith(endpoint(), "application/json", manyShortLines(""));
|
||||
|
||||
StepVerifier.create(sendMessage())
|
||||
.verifyErrorMatches(t -> messageContains(t, "Inbound response body exceeds the maximum allowed size"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldRejectDiscardedResponseExceedingMaxSize() {
|
||||
// A content type the transport neither parses as SSE nor as JSON, so the body is
|
||||
// discarded. The line subscriber underneath still buffers each line, so an
|
||||
// unterminated one must abort the response rather than accumulate.
|
||||
respondWith(endpoint(), "text/plain", unterminatedLine(8));
|
||||
|
||||
StepVerifier.create(sendMessage())
|
||||
.verifyErrorMatches(t -> messageContains(t, "Inbound line exceeds the maximum allowed size"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldAcceptEventOfExactlyMaxSize() {
|
||||
// The bound is inclusive and the SSE framing around the payload is given its own
|
||||
// headroom, so a message of exactly maxResponseSize must still be delivered.
|
||||
respondWith(endpoint(), "text/event-stream", body -> body
|
||||
.write(("data:" + jsonRpcResponseOfExactly(MAX_SIZE) + "\n\n").getBytes(StandardCharsets.UTF_8)));
|
||||
|
||||
StepVerifier.create(sendMessage()).verifyComplete();
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldAcceptJsonResponseOfExactlyMaxSize() {
|
||||
// Same inclusive bound on the aggregated body.
|
||||
respondWith(endpoint(), "application/json",
|
||||
body -> body.write(jsonRpcResponseOfExactly(MAX_SIZE).getBytes(StandardCharsets.UTF_8)));
|
||||
|
||||
StepVerifier.create(sendMessage()).verifyComplete();
|
||||
}
|
||||
|
||||
/**
|
||||
* A single-line JSON-RPC response padded to exactly {@code size} bytes.
|
||||
*/
|
||||
private static String jsonRpcResponseOfExactly(int size) {
|
||||
String prefix = "{\"jsonrpc\":\"2.0\",\"id\":\"test-id\",\"result\":{\"pad\":\"";
|
||||
String suffix = "\"}}";
|
||||
return prefix + "a".repeat(size - prefix.length() - suffix.length()) + suffix;
|
||||
}
|
||||
|
||||
private Mono<Void> sendMessage() {
|
||||
HttpClientStreamableHttpTransport transport = HttpClientStreamableHttpTransport.builder(this.host)
|
||||
.maxResponseSize(MAX_SIZE)
|
||||
.build();
|
||||
JSONRPCRequest request = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
return transport.connect(Function.identity()).then(transport.sendMessage(request));
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,126 @@
|
|||
/*
|
||||
* Copyright 2024-2026 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client.transport;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.OutputStream;
|
||||
import java.net.InetSocketAddress;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.time.Duration;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.Executors;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.atomic.AtomicLong;
|
||||
|
||||
import com.sun.net.httpserver.HttpServer;
|
||||
import io.modelcontextprotocol.server.transport.TomcatTestUtil;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.JSONRPCRequest;
|
||||
import io.modelcontextprotocol.spec.McpTransportException;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
/**
|
||||
* Verifies that {@link HttpClientStreamableHttpTransport} bounds the response to the
|
||||
* session-terminating DELETE, so a server cannot exhaust the client's memory on the way
|
||||
* out. Checks both that the failure surfaces to the caller and that the client actually
|
||||
* stops reading rather than draining whatever the server sends.
|
||||
*
|
||||
* @author Daniel Garnier-Moiroux
|
||||
*/
|
||||
@Timeout(30)
|
||||
class HttpClientStreamableHttpTransportDeleteBoundedReadTests {
|
||||
|
||||
private static final int MAX_SIZE = 1024;
|
||||
|
||||
private static final int FLOOD_BYTES = 32 * 1024 * 1024;
|
||||
|
||||
private HttpServer server;
|
||||
|
||||
private String host;
|
||||
|
||||
private final AtomicLong deleteBytesWritten = new AtomicLong();
|
||||
|
||||
private final CountDownLatch deleteFinished = new CountDownLatch(1);
|
||||
|
||||
@BeforeEach
|
||||
void setUp() throws IOException {
|
||||
int port = TomcatTestUtil.findAvailablePort();
|
||||
this.host = "http://localhost:" + port;
|
||||
this.server = HttpServer.create(new InetSocketAddress(port), 0);
|
||||
this.server.createContext("/mcp", exchange -> {
|
||||
if ("DELETE".equals(exchange.getRequestMethod())) {
|
||||
floodDeleteResponse(exchange);
|
||||
return;
|
||||
}
|
||||
// Hand out a session id so the transport issues a DELETE on close.
|
||||
exchange.getResponseHeaders().set("Content-Type", "application/json");
|
||||
exchange.getResponseHeaders().set("Mcp-Session-Id", "test-session");
|
||||
byte[] body = "{\"jsonrpc\":\"2.0\",\"id\":\"test-id\",\"result\":{}}".getBytes(StandardCharsets.UTF_8);
|
||||
exchange.sendResponseHeaders(200, body.length);
|
||||
try (OutputStream out = exchange.getResponseBody()) {
|
||||
out.write(body);
|
||||
}
|
||||
});
|
||||
this.server.setExecutor(Executors.newCachedThreadPool());
|
||||
this.server.start();
|
||||
}
|
||||
|
||||
private void floodDeleteResponse(com.sun.net.httpserver.HttpExchange exchange) throws IOException {
|
||||
exchange.getResponseHeaders().set("Content-Type", "application/json");
|
||||
exchange.sendResponseHeaders(200, 0);
|
||||
try (OutputStream body = exchange.getResponseBody()) {
|
||||
byte[] chunk = new byte[64 * 1024];
|
||||
java.util.Arrays.fill(chunk, (byte) 'a'); // no line terminator, ever
|
||||
while (this.deleteBytesWritten.get() < FLOOD_BYTES) {
|
||||
body.write(chunk);
|
||||
body.flush();
|
||||
this.deleteBytesWritten.addAndGet(chunk.length);
|
||||
}
|
||||
}
|
||||
catch (IOException ignored) {
|
||||
// Expected: the client aborts the response once the limit is exceeded.
|
||||
}
|
||||
finally {
|
||||
exchange.close();
|
||||
this.deleteFinished.countDown();
|
||||
}
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
void tearDown() {
|
||||
if (this.server != null) {
|
||||
this.server.stop(0);
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldStopReadingDeleteResponseOnceLimitExceeded() throws Exception {
|
||||
HttpClientStreamableHttpTransport transport = HttpClientStreamableHttpTransport.builder(this.host)
|
||||
.maxResponseSize(MAX_SIZE)
|
||||
.build();
|
||||
transport.connect(m -> m).block(Duration.ofSeconds(5));
|
||||
JSONRPCRequest request = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
transport.sendMessage(request).block(Duration.ofSeconds(5));
|
||||
|
||||
assertThatThrownBy(() -> transport.closeGracefully().block(Duration.ofSeconds(10)))
|
||||
.isInstanceOf(McpTransportException.class)
|
||||
.hasMessageContaining("Inbound response body exceeds the maximum allowed size");
|
||||
|
||||
assertThat(this.deleteFinished.await(15, TimeUnit.SECONDS))
|
||||
.as("the server's DELETE response should be cut short by the client")
|
||||
.isTrue();
|
||||
assertThat(this.deleteBytesWritten.get()).as("bytes the client let the server stream on DELETE")
|
||||
.isLessThan(FLOOD_BYTES);
|
||||
}
|
||||
|
||||
}
|
||||
Loading…
Reference in a new issue