Validate message endpoint in SSE WebFlux client transport

Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>

Author: Kehrlann

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<URI> 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 WebFluxSseClientTransport(WebClient.Builder webClientBuilder, McpJsonMapp
 	 * @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 Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> h
 		this.inboundSubscription = events.concatMap(event -> Mono.just(event).<JSONRPCMessage>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 Mono<Void> sendMessage(JSONRPCMessage message) {
 	 * Includes automatic retry logic for handling transient connection failures.
 	 */
 	// visible for tests
-	protected Flux<ServerSentEvent<String>> eventStream() {// @formatter:off
-		return this.webClient
-			.get()
+	protected Flux<ServerSentEvent<String>> 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 static class Builder {
 
 		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 Builder jsonMapper(McpJsonMapper jsonMapper) {
 			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);
 		}
 
 	}

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.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 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,15 +62,18 @@ class WebFluxSseClientTransportTests {
 
 	private WebClient.Builder webClientBuilder;
 
+	private SseMessageEndpointValidator sseMessageEndpointValidator = mock(SseMessageEndpointValidator.class);
+
 	// Test class to access protected methods
 	static class TestSseClientTransport extends WebFluxSseClientTransport {
 
 		private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
 
 		private Sinks.Many<ServerSentEvent<String>> 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 @@ static void cleanup() {
 	@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 @@ void testMessageOrderPreservation() {
 		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;
+	}
+
 }