More precise mapping for WebSocket handshake requests

Closes gh-26565

spring-webflux/src/main/java/org/springframework/web/reactive/handler/AbstractUrlHandlerMapping.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@ import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.BiPredicate;
 
 import reactor.core.publisher.Mono;
 
@@ -57,6 +58,9 @@ public abstract class AbstractUrlHandlerMapping extends AbstractHandlerMapping {
 
 	private final Map<PathPattern, Object> handlerMap = new LinkedHashMap<>();
 
+	@Nullable
+	private BiPredicate<Object, ServerWebExchange> handlerPredicate;
+
 
 	/**
 	 * Set whether to lazily initialize handlers. Only applicable to
@@ -81,6 +85,23 @@ public abstract class AbstractUrlHandlerMapping extends AbstractHandlerMapping {
 		return Collections.unmodifiableMap(this.handlerMap);
 	}
 
+	/**
+	 * Configure a predicate for extended matching of the handler that was
+	 * matched by URL path. This allows for further narrowing of the mapping by
+	 * checking additional properties of the request. If the predicate returns
+	 * "false", it result in a no-match, which allows another
+	 * {@link org.springframework.web.reactive.HandlerMapping} to match or
+	 * result in a 404 (NOT_FOUND) response.
+	 * @param handlerPredicate a bi-predicate to match the candidate handler
+	 * against the current exchange.
+	 * @since 5.3.5
+	 * @see org.springframework.web.reactive.socket.server.support.WebSocketUpgradeHandlerPredicate
+	 */
+	public void setHandlerPredicate(BiPredicate<Object, ServerWebExchange> handlerPredicate) {
+		this.handlerPredicate = (this.handlerPredicate != null ?
+				this.handlerPredicate.and(handlerPredicate) : handlerPredicate);
+	}
+
 
 	@Override
 	public Mono<Object> getHandlerInternal(ServerWebExchange exchange) {
@@ -129,11 +150,7 @@ public abstract class AbstractUrlHandlerMapping extends AbstractHandlerMapping {
 		PathPattern.PathMatchInfo matchInfo = pattern.matchAndExtract(lookupPath);
 		Assert.notNull(matchInfo, "Expected a match");
 
-		return handleMatch(this.handlerMap.get(pattern), pattern, pathWithinMapping, matchInfo, exchange);
-	}
-
-	private Object handleMatch(Object handler, PathPattern bestMatch, PathContainer pathWithinMapping,
-			PathPattern.PathMatchInfo matchInfo, ServerWebExchange exchange) {
+		Object handler = this.handlerMap.get(pattern);
 
 		// Bean name or resolved handler?
 		if (handler instanceof String) {
@@ -141,10 +158,14 @@ public abstract class AbstractUrlHandlerMapping extends AbstractHandlerMapping {
 			handler = obtainApplicationContext().getBean(handlerName);
 		}
 
+		if (this.handlerPredicate != null && !this.handlerPredicate.test(handler, exchange)) {
+			return null;
+		}
+
 		validateHandler(handler, exchange);
 
 		exchange.getAttributes().put(BEST_MATCHING_HANDLER_ATTRIBUTE, handler);
-		exchange.getAttributes().put(BEST_MATCHING_PATTERN_ATTRIBUTE, bestMatch);
+		exchange.getAttributes().put(BEST_MATCHING_PATTERN_ATTRIBUTE, pattern);
 		exchange.getAttributes().put(PATH_WITHIN_HANDLER_MAPPING_ATTRIBUTE, pathWithinMapping);
 		exchange.getAttributes().put(URI_TEMPLATE_VARIABLES_ATTRIBUTE, matchInfo.getUriVariables());
 

spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/support/WebSocketUpgradeHandlerPredicate.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2002-2021 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.web.reactive.socket.server.support;
+
+import java.util.function.BiPredicate;
+
+import org.springframework.web.reactive.socket.WebSocketHandler;
+import org.springframework.web.server.ServerWebExchange;
+
+/**
+ * A predicate for use with
+ * {@link org.springframework.web.reactive.handler.AbstractUrlHandlerMapping#setHandlerPredicate(BiPredicate)}
+ * to ensure only WebSocket handshake requests are matched to handlers of
+ * type {@link WebSocketHandler}.
+ *
+ * @author Rossen Stoyanchev
+ * @since 5.3.5
+ */
+public class WebSocketUpgradeHandlerPredicate implements BiPredicate<Object, ServerWebExchange> {
+
+
+	@Override
+	public boolean test(Object handler, ServerWebExchange exchange) {
+		if (handler instanceof WebSocketHandler) {
+			String method = exchange.getRequest().getMethodValue();
+			String header = exchange.getRequest().getHeaders().getUpgrade();
+			return (method.equals("GET") && header != null && header.equalsIgnoreCase("websocket"));
+		}
+		return true;
+	}
+
+}

spring-webflux/src/test/java/org/springframework/web/reactive/socket/server/support/WebSocketUpgradeHandlerPredicateTests.java

@@ -0,0 +1,97 @@
+/*
+ * Copyright 2002-2021 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.web.reactive.socket.server.support;
+
+import java.util.Collections;
+
+import org.junit.jupiter.api.Test;
+
+import org.springframework.http.HttpHeaders;
+import org.springframework.web.context.support.StaticWebApplicationContext;
+import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping;
+import org.springframework.web.reactive.socket.WebSocketHandler;
+import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.testfixture.http.server.reactive.MockServerHttpRequest;
+import org.springframework.web.testfixture.server.MockServerWebExchange;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+
+/**
+ * Unit tests for and related to the use of {@link WebSocketUpgradeHandlerPredicate}.
+ *
+ * @author Rossen Stoyanchev
+ */
+public class WebSocketUpgradeHandlerPredicateTests {
+
+	private final WebSocketUpgradeHandlerPredicate predicate = new WebSocketUpgradeHandlerPredicate();
+
+	private final WebSocketHandler webSocketHandler = mock(WebSocketHandler.class);
+
+	ServerWebExchange httpGetExchange =
+			MockServerWebExchange.from(MockServerHttpRequest.get("/path"));
+
+	ServerWebExchange httpPostExchange =
+			MockServerWebExchange.from(MockServerHttpRequest.post("/path"));
+
+	ServerWebExchange webSocketExchange =
+			MockServerWebExchange.from(MockServerHttpRequest.get("/path").header(HttpHeaders.UPGRADE, "websocket"));
+
+
+	@Test
+	void match() {
+		assertThat(this.predicate.test(this.webSocketHandler, this.webSocketExchange))
+				.as("Should match WebSocketHandler to WebSocket upgrade")
+				.isTrue();
+
+		assertThat(this.predicate.test(new Object(), this.httpGetExchange))
+				.as("Should match non-WebSocketHandler to any request")
+				.isTrue();
+	}
+
+	@Test
+	void noMatch() {
+		assertThat(this.predicate.test(this.webSocketHandler, this.httpGetExchange))
+				.as("Should not match WebSocket handler to HTTP GET")
+				.isFalse();
+
+		assertThat(this.predicate.test(this.webSocketHandler, this.httpPostExchange))
+				.as("Should not match WebSocket handler to HTTP POST")
+				.isFalse();
+	}
+
+	@Test
+	void simpleUrlHandlerMapping() {
+		SimpleUrlHandlerMapping mapping = new SimpleUrlHandlerMapping();
+		mapping.setUrlMap(Collections.singletonMap("/path", this.webSocketHandler));
+		mapping.setApplicationContext(new StaticWebApplicationContext());
+
+		Object actual = mapping.getHandler(httpGetExchange).block();
+		assertThat(actual).as("Should match HTTP GET by URL path").isSameAs(this.webSocketHandler);
+
+		mapping.setHandlerPredicate(new WebSocketUpgradeHandlerPredicate());
+
+		actual = mapping.getHandler(this.httpGetExchange).block();
+		assertThat(actual).as("Should not match if not a WebSocket upgrade").isNull();
+
+		actual = mapping.getHandler(this.httpPostExchange).block();
+		assertThat(actual).as("Should not match if not a WebSocket upgrade").isNull();
+
+		actual = mapping.getHandler(this.webSocketExchange).block();
+		assertThat(actual).as("Should match WebSocket upgrade").isSameAs(this.webSocketHandler);
+	}
+
+}