diff --git a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransport.java b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransport.java index 304a343..350a204 100644 --- a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransport.java +++ b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransport.java @@ -5,7 +5,9 @@ package io.modelcontextprotocol.client.transport; import java.io.IOException; +import java.net.URI; import java.util.List; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiConsumer; import java.util.function.Function; @@ -128,7 +130,17 @@ public class WebFluxSseClientTransport implements McpClientTransport { * The SSE endpoint URI provided by the server. Used for sending outbound messages via * HTTP POST requests. */ - private String sseEndpoint; + private final String sseEndpoint; + + /** + * Used to capture the full SSE URI from the web client when connecting. + */ + private final AtomicReference sseUri = new AtomicReference<>(); + + /** + * Validator for the message endpoint. + */ + private final SseMessageEndpointValidator messageEndpointValidator; /** * Constructs a new SseClientTransport with the specified WebClient builder and @@ -152,13 +164,30 @@ public class WebFluxSseClientTransport implements McpClientTransport { * @throws IllegalArgumentException if either parameter is null */ public WebFluxSseClientTransport(WebClient.Builder webClientBuilder, McpJsonMapper jsonMapper, String sseEndpoint) { + this(webClientBuilder, jsonMapper, sseEndpoint, new DefaultSseMessageEndpointValidator()); + } + + /** + * Constructs a new SseClientTransport with the specified WebClient builder and + * ObjectMapper. Initializes both inbound and outbound message processing pipelines. + * @param webClientBuilder the WebClient.Builder to use for creating the WebClient + * instance + * @param jsonMapper the ObjectMapper to use for JSON processing + * @param sseEndpoint the SSE endpoint URI to use for establishing the connection + * @param messageEndpointValidator validator for the message endpoint + * @throws IllegalArgumentException if either parameter is null + */ + public WebFluxSseClientTransport(WebClient.Builder webClientBuilder, McpJsonMapper jsonMapper, String sseEndpoint, + SseMessageEndpointValidator messageEndpointValidator) { Assert.notNull(jsonMapper, "jsonMapper must not be null"); Assert.notNull(webClientBuilder, "WebClient.Builder must not be null"); + Assert.notNull(messageEndpointValidator, "messageEndpointValidator must not be null"); Assert.hasText(sseEndpoint, "SSE endpoint must not be null or empty"); this.jsonMapper = jsonMapper; this.webClient = webClientBuilder.build(); this.sseEndpoint = sseEndpoint; + this.messageEndpointValidator = messageEndpointValidator; } @Override @@ -195,6 +224,14 @@ public class WebFluxSseClientTransport implements McpClientTransport { this.inboundSubscription = events.concatMap(event -> Mono.just(event).handle((e, s) -> { if (ENDPOINT_EVENT_TYPE.equals(event.event())) { String messageEndpointUri = event.data(); + try { + this.messageEndpointValidator.validate(this.sseUri.get(), messageEndpointUri); + } + catch (InvalidSseMessageEndpointException ex) { + messageEndpointSink.tryEmitError(ex); + s.error(ex); + return; + } if (messageEndpointSink.tryEmitValue(messageEndpointUri).isSuccess()) { s.complete(); } @@ -276,16 +313,17 @@ public class WebFluxSseClientTransport implements McpClientTransport { * Includes automatic retry logic for handling transient connection failures. */ // visible for tests - protected Flux> eventStream() {// @formatter:off - return this.webClient - .get() + protected Flux> eventStream() { + return this.webClient.get() .uri(this.sseEndpoint) .accept(MediaType.TEXT_EVENT_STREAM) .header(HttpHeaders.PROTOCOL_VERSION, MCP_PROTOCOL_VERSION) - .retrieve() - .bodyToFlux(SSE_TYPE) + .exchangeToFlux(exchange -> { + this.sseUri.set(exchange.request().getURI()); + return exchange.bodyToFlux(SSE_TYPE); + }) .retryWhen(Retry.from(retrySignal -> retrySignal.handle(inboundRetryHandler))); - } // @formatter:on + } /** * Retry handler for the inbound SSE stream. Implements the retry logic for handling @@ -368,6 +406,8 @@ public class WebFluxSseClientTransport implements McpClientTransport { private McpJsonMapper jsonMapper; + private SseMessageEndpointValidator messageEndpointValidator = new DefaultSseMessageEndpointValidator(); + /** * Creates a new builder with the specified WebClient.Builder. * @param webClientBuilder the WebClient.Builder to use @@ -399,13 +439,26 @@ public class WebFluxSseClientTransport implements McpClientTransport { return this; } + /** + * Sets the validator that ensure the message endpoint returned over the SSE + * connection is valid. + * @param messageEndpointValidator the validator + * @return this builder + */ + public Builder messageEndpointValidator(SseMessageEndpointValidator messageEndpointValidator) { + Assert.notNull(messageEndpointValidator, "messageEndpointValidator must not be null"); + this.messageEndpointValidator = messageEndpointValidator; + return this; + } + /** * Builds a new {@link WebFluxSseClientTransport} instance. * @return a new transport instance */ public WebFluxSseClientTransport build() { return new WebFluxSseClientTransport(webClientBuilder, - jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper, sseEndpoint); + jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper, sseEndpoint, + messageEndpointValidator); } } diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransportTests.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransportTests.java index 4b0d4e5..d109a32 100644 --- a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransportTests.java +++ b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransportTests.java @@ -4,6 +4,7 @@ package io.modelcontextprotocol.client.transport; +import java.net.URI; import java.time.Duration; import java.util.Map; import java.util.concurrent.CopyOnWriteArrayList; @@ -21,6 +22,7 @@ import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; +import org.mockito.ArgumentCaptor; import org.testcontainers.containers.GenericContainer; import org.testcontainers.containers.wait.strategy.Wait; import reactor.core.publisher.Flux; @@ -35,6 +37,9 @@ import static io.modelcontextprotocol.util.McpJsonMapperUtils.JSON_MAPPER; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatCode; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.matches; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; /** * Tests for the {@link WebFluxSseClientTransport} class. @@ -57,6 +62,8 @@ class WebFluxSseClientTransportTests { private WebClient.Builder webClientBuilder; + private SseMessageEndpointValidator sseMessageEndpointValidator = mock(SseMessageEndpointValidator.class); + // Test class to access protected methods static class TestSseClientTransport extends WebFluxSseClientTransport { @@ -64,8 +71,9 @@ class WebFluxSseClientTransportTests { private Sinks.Many> events = Sinks.many().unicast().onBackpressureBuffer(); - public TestSseClientTransport(WebClient.Builder webClientBuilder, McpJsonMapper jsonMapper) { - super(webClientBuilder, jsonMapper); + public TestSseClientTransport(WebClient.Builder webClientBuilder, McpJsonMapper jsonMapper, + SseMessageEndpointValidator sseMessageEndpointValidator) { + super(webClientBuilder, jsonMapper, "/sse", sseMessageEndpointValidator); } @Override @@ -113,7 +121,7 @@ class WebFluxSseClientTransportTests { @BeforeEach void setUp() { webClientBuilder = WebClient.builder().baseUrl(host); - transport = new TestSseClientTransport(webClientBuilder, JSON_MAPPER); + transport = new TestSseClientTransport(webClientBuilder, JSON_MAPPER, sseMessageEndpointValidator); transport.connect(Function.identity()).block(); } @@ -368,4 +376,44 @@ class WebFluxSseClientTransportTests { assertThat(transport.getInboundMessageCount()).isEqualTo(3); } + @Test + void testMessageEndpointValidation() throws InvalidSseMessageEndpointException { + var uriCaptor = ArgumentCaptor.forClass(URI.class); + verify(sseMessageEndpointValidator).validate(uriCaptor.capture(), matches("/message\\?sessionId=[a-z0-9-]+")); + assertThat(uriCaptor.getValue().toString()).matches("http://localhost:\\d+/sse"); + } + + @Test + void testMessageEndpointValidationRejects() { + TestSseClientTransport transport = new TestSseClientTransport(webClientBuilder, JSON_MAPPER, + (sseUri, messageEndpoint) -> { + throw new InvalidSseMessageEndpointException("boom", messageEndpoint); + }); + + try { + // fails to connect + StepVerifier.create(transport.connect(Function.identity())) + .verifyErrorMatches(WebFluxSseClientTransportTests::isInvalidEndpointError); + + // Since connection failed, there is no message endpoint, and no message can + // be sent + JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id", + Map.of("key", "value")); + + StepVerifier.create(transport.sendMessage(testMessage)) + .verifyErrorMatches(WebFluxSseClientTransportTests::isInvalidEndpointError); + } + finally { + transport.closeGracefully(); + } + } + + private static boolean isInvalidEndpointError(Throwable e) { + if (e instanceof InvalidSseMessageEndpointException ismee) { + return ismee.getMessageEndpoint().matches("/message\\?sessionId=[a-z0-9-]+") + && ismee.getMessage().equals("boom"); + } + return false; + } + }