Skip to content

Commit 1dd7d53

Browse files
rstoyanchevrstoyanchev
committed
More precise mapping for WebSocket handshake requests
Closes gh-26565
1 parent 8535193 commit 1dd7d53

File tree

5 files changed

+284
-12
lines changed

5 files changed

+284
-12
lines changed

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

+28-7
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2020 the original author or authors.
2+
* Copyright 2002-2021 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@
2121
import java.util.LinkedHashMap;
2222
import java.util.List;
2323
import java.util.Map;
24+
import java.util.function.BiPredicate;
2425

2526
import reactor.core.publisher.Mono;
2627

@@ -57,6 +58,9 @@ public abstract class AbstractUrlHandlerMapping extends AbstractHandlerMapping {
5758

5859
private final Map<PathPattern, Object> handlerMap = new LinkedHashMap<>();
5960

61+
@Nullable
62+
private BiPredicate<Object, ServerWebExchange> handlerPredicate;
63+
6064

6165
/**
6266
* Set whether to lazily initialize handlers. Only applicable to
@@ -81,6 +85,23 @@ public final Map<PathPattern, Object> getHandlerMap() {
8185
return Collections.unmodifiableMap(this.handlerMap);
8286
}
8387

88+
/**
89+
* Configure a predicate for extended matching of the handler that was
90+
* matched by URL path. This allows for further narrowing of the mapping by
91+
* checking additional properties of the request. If the predicate returns
92+
* "false", it result in a no-match, which allows another
93+
* {@link org.springframework.web.reactive.HandlerMapping} to match or
94+
* result in a 404 (NOT_FOUND) response.
95+
* @param handlerPredicate a bi-predicate to match the candidate handler
96+
* against the current exchange.
97+
* @since 5.3.5
98+
* @see org.springframework.web.reactive.socket.server.support.WebSocketUpgradeHandlerPredicate
99+
*/
100+
public void setHandlerPredicate(BiPredicate<Object, ServerWebExchange> handlerPredicate) {
101+
this.handlerPredicate = (this.handlerPredicate != null ?
102+
this.handlerPredicate.and(handlerPredicate) : handlerPredicate);
103+
}
104+
84105

85106
@Override
86107
public Mono<Object> getHandlerInternal(ServerWebExchange exchange) {
@@ -129,22 +150,22 @@ protected Object lookupHandler(PathContainer lookupPath, ServerWebExchange excha
129150
PathPattern.PathMatchInfo matchInfo = pattern.matchAndExtract(lookupPath);
130151
Assert.notNull(matchInfo, "Expected a match");
131152

132-
return handleMatch(this.handlerMap.get(pattern), pattern, pathWithinMapping, matchInfo, exchange);
133-
}
134-
135-
private Object handleMatch(Object handler, PathPattern bestMatch, PathContainer pathWithinMapping,
136-
PathPattern.PathMatchInfo matchInfo, ServerWebExchange exchange) {
153+
Object handler = this.handlerMap.get(pattern);
137154

138155
// Bean name or resolved handler?
139156
if (handler instanceof String) {
140157
String handlerName = (String) handler;
141158
handler = obtainApplicationContext().getBean(handlerName);
142159
}
143160

161+
if (this.handlerPredicate != null && !this.handlerPredicate.test(handler, exchange)) {
162+
return null;
163+
}
164+
144165
validateHandler(handler, exchange);
145166

146167
exchange.getAttributes().put(BEST_MATCHING_HANDLER_ATTRIBUTE, handler);
147-
exchange.getAttributes().put(BEST_MATCHING_PATTERN_ATTRIBUTE, bestMatch);
168+
exchange.getAttributes().put(BEST_MATCHING_PATTERN_ATTRIBUTE, pattern);
148169
exchange.getAttributes().put(PATH_WITHIN_HANDLER_MAPPING_ATTRIBUTE, pathWithinMapping);
149170
exchange.getAttributes().put(URI_TEMPLATE_VARIABLES_ATTRIBUTE, matchInfo.getUriVariables());
150171

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

+45
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/*
2+
* Copyright 2002-2021 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.web.reactive.socket.server.support;
17+
18+
import java.util.function.BiPredicate;
19+
20+
import org.springframework.web.reactive.socket.WebSocketHandler;
21+
import org.springframework.web.server.ServerWebExchange;
22+
23+
/**
24+
* A predicate for use with
25+
* {@link org.springframework.web.reactive.handler.AbstractUrlHandlerMapping#setHandlerPredicate(BiPredicate)}
26+
* to ensure only WebSocket handshake requests are matched to handlers of
27+
* type {@link WebSocketHandler}.
28+
*
29+
* @author Rossen Stoyanchev
30+
* @since 5.3.5
31+
*/
32+
public class WebSocketUpgradeHandlerPredicate implements BiPredicate<Object, ServerWebExchange> {
33+
34+
35+
@Override
36+
public boolean test(Object handler, ServerWebExchange exchange) {
37+
if (handler instanceof WebSocketHandler) {
38+
String method = exchange.getRequest().getMethodValue();
39+
String header = exchange.getRequest().getHeaders().getUpgrade();
40+
return (method.equals("GET") && header != null && header.equalsIgnoreCase("websocket"));
41+
}
42+
return true;
43+
}
44+
45+
}

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

+97
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
/*
2+
* Copyright 2002-2021 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.web.reactive.socket.server.support;
17+
18+
import java.util.Collections;
19+
20+
import org.junit.jupiter.api.Test;
21+
22+
import org.springframework.http.HttpHeaders;
23+
import org.springframework.web.context.support.StaticWebApplicationContext;
24+
import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping;
25+
import org.springframework.web.reactive.socket.WebSocketHandler;
26+
import org.springframework.web.server.ServerWebExchange;
27+
import org.springframework.web.testfixture.http.server.reactive.MockServerHttpRequest;
28+
import org.springframework.web.testfixture.server.MockServerWebExchange;
29+
30+
import static org.assertj.core.api.Assertions.assertThat;
31+
import static org.mockito.Mockito.mock;
32+
33+
/**
34+
* Unit tests for and related to the use of {@link WebSocketUpgradeHandlerPredicate}.
35+
*
36+
* @author Rossen Stoyanchev
37+
*/
38+
public class WebSocketUpgradeHandlerPredicateTests {
39+
40+
private final WebSocketUpgradeHandlerPredicate predicate = new WebSocketUpgradeHandlerPredicate();
41+
42+
private final WebSocketHandler webSocketHandler = mock(WebSocketHandler.class);
43+
44+
ServerWebExchange httpGetExchange =
45+
MockServerWebExchange.from(MockServerHttpRequest.get("/path"));
46+
47+
ServerWebExchange httpPostExchange =
48+
MockServerWebExchange.from(MockServerHttpRequest.post("/path"));
49+
50+
ServerWebExchange webSocketExchange =
51+
MockServerWebExchange.from(MockServerHttpRequest.get("/path").header(HttpHeaders.UPGRADE, "websocket"));
52+
53+
54+
@Test
55+
void match() {
56+
assertThat(this.predicate.test(this.webSocketHandler, this.webSocketExchange))
57+
.as("Should match WebSocketHandler to WebSocket upgrade")
58+
.isTrue();
59+
60+
assertThat(this.predicate.test(new Object(), this.httpGetExchange))
61+
.as("Should match non-WebSocketHandler to any request")
62+
.isTrue();
63+
}
64+
65+
@Test
66+
void noMatch() {
67+
assertThat(this.predicate.test(this.webSocketHandler, this.httpGetExchange))
68+
.as("Should not match WebSocket handler to HTTP GET")
69+
.isFalse();
70+
71+
assertThat(this.predicate.test(this.webSocketHandler, this.httpPostExchange))
72+
.as("Should not match WebSocket handler to HTTP POST")
73+
.isFalse();
74+
}
75+
76+
@Test
77+
void simpleUrlHandlerMapping() {
78+
SimpleUrlHandlerMapping mapping = new SimpleUrlHandlerMapping();
79+
mapping.setUrlMap(Collections.singletonMap("/path", this.webSocketHandler));
80+
mapping.setApplicationContext(new StaticWebApplicationContext());
81+
82+
Object actual = mapping.getHandler(httpGetExchange).block();
83+
assertThat(actual).as("Should match HTTP GET by URL path").isSameAs(this.webSocketHandler);
84+
85+
mapping.setHandlerPredicate(new WebSocketUpgradeHandlerPredicate());
86+
87+
actual = mapping.getHandler(this.httpGetExchange).block();
88+
assertThat(actual).as("Should not match if not a WebSocket upgrade").isNull();
89+
90+
actual = mapping.getHandler(this.httpPostExchange).block();
91+
assertThat(actual).as("Should not match if not a WebSocket upgrade").isNull();
92+
93+
actual = mapping.getHandler(this.webSocketExchange).block();
94+
assertThat(actual).as("Should match WebSocket upgrade").isSameAs(this.webSocketHandler);
95+
}
96+
97+
}

‎spring-websocket/src/main/java/org/springframework/web/socket/server/support/WebSocketHandlerMapping.java

+44-5
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2020 the original author or authors.
2+
* Copyright 2002-2021 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -17,26 +17,47 @@
1717
package org.springframework.web.socket.server.support;
1818

1919
import javax.servlet.ServletContext;
20+
import javax.servlet.http.HttpServletRequest;
2021

2122
import org.springframework.context.Lifecycle;
2223
import org.springframework.context.SmartLifecycle;
24+
import org.springframework.http.HttpHeaders;
25+
import org.springframework.lang.Nullable;
2326
import org.springframework.web.context.ServletContextAware;
27+
import org.springframework.web.servlet.HandlerExecutionChain;
2428
import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
2529

2630
/**
27-
* An extension of {@link SimpleUrlHandlerMapping} that is also a
28-
* {@link SmartLifecycle} container and propagates start and stop calls to any
29-
* handlers that implement {@link Lifecycle}. The handlers are typically expected
30-
* to be {@code WebSocketHttpRequestHandler} or {@code SockJsHttpRequestHandler}.
31+
* Extension of {@link SimpleUrlHandlerMapping} with support for more
32+
* precise mapping of WebSocket handshake requests to handlers of type
33+
* {@link WebSocketHttpRequestHandler}. Also delegates {@link Lifecycle}
34+
* methods to handlers in the {@link #getUrlMap()} that implement it.
3135
*
3236
* @author Rossen Stoyanchev
3337
* @since 4.2
3438
*/
3539
public class WebSocketHandlerMapping extends SimpleUrlHandlerMapping implements SmartLifecycle {
3640

41+
private boolean webSocketUpgradeMatch;
42+
3743
private volatile boolean running;
3844

3945

46+
/**
47+
* When this is set, if the matched handler is
48+
* {@link WebSocketHttpRequestHandler}, ensure the request is a WebSocket
49+
* handshake, i.e. HTTP GET with the header {@code "Upgrade:websocket"},
50+
* or otherwise suppress the match and return {@code null} allowing another
51+
* {@link org.springframework.web.servlet.HandlerMapping} to match for the
52+
* same URL path.
53+
* @param match whether to enable matching on {@code "Upgrade: websocket"}
54+
* @since 5.3.5
55+
*/
56+
public void setWebSocketUpgradeMatch(boolean match) {
57+
this.webSocketUpgradeMatch = match;
58+
}
59+
60+
4061
@Override
4162
protected void initServletContext(ServletContext servletContext) {
4263
for (Object handler : getUrlMap().values()) {
@@ -76,4 +97,22 @@ public boolean isRunning() {
7697
return this.running;
7798
}
7899

100+
@Nullable
101+
@Override
102+
protected Object getHandlerInternal(HttpServletRequest request) throws Exception {
103+
Object handler = super.getHandlerInternal(request);
104+
return matchWebSocketUpgrade(handler, request) ? handler : null;
105+
}
106+
107+
private boolean matchWebSocketUpgrade(@Nullable Object handler, HttpServletRequest request) {
108+
handler = (handler instanceof HandlerExecutionChain ?
109+
((HandlerExecutionChain) handler).getHandler() : handler);
110+
if (this.webSocketUpgradeMatch && handler instanceof WebSocketHttpRequestHandler) {
111+
String header = request.getHeader(HttpHeaders.UPGRADE);
112+
return (request.getMethod().equals("GET") &&
113+
header != null && header.equalsIgnoreCase("websocket"));
114+
}
115+
return true;
116+
}
117+
79118
}

‎spring-websocket/src/test/java/org/springframework/web/socket/server/support/WebSocketHandlerMappingTests.java

+70
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
/*
2+
* Copyright 2002-2021 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.web.socket.server.support;
17+
18+
import java.util.Collections;
19+
20+
import org.junit.jupiter.api.Test;
21+
22+
import org.springframework.web.HttpRequestHandler;
23+
import org.springframework.web.context.support.StaticWebApplicationContext;
24+
import org.springframework.web.servlet.HandlerExecutionChain;
25+
import org.springframework.web.socket.WebSocketHandler;
26+
import org.springframework.web.testfixture.servlet.MockHttpServletRequest;
27+
28+
import static org.assertj.core.api.Assertions.assertThat;
29+
import static org.mockito.Mockito.mock;
30+
31+
/**
32+
* Unit tests for {@link WebSocketHandlerMapping}.
33+
*
34+
* @author Rossen Stoyanchev
35+
*/
36+
public class WebSocketHandlerMappingTests {
37+
38+
39+
@Test
40+
void webSocketHandshakeMatch() throws Exception {
41+
HttpRequestHandler handler = new WebSocketHttpRequestHandler(mock(WebSocketHandler.class));
42+
43+
WebSocketHandlerMapping mapping = new WebSocketHandlerMapping();
44+
mapping.setUrlMap(Collections.singletonMap("/path", handler));
45+
mapping.setApplicationContext(new StaticWebApplicationContext());
46+
47+
MockHttpServletRequest request = new MockHttpServletRequest("GET", "/path");
48+
49+
HandlerExecutionChain chain = mapping.getHandler(request);
50+
assertThat(chain).isNotNull();
51+
assertThat(chain.getHandler()).isSameAs(handler);
52+
53+
mapping.setWebSocketUpgradeMatch(true);
54+
55+
chain = mapping.getHandler(request);
56+
assertThat(chain).isNull();
57+
58+
request.addHeader("Upgrade", "websocket");
59+
60+
chain = mapping.getHandler(request);
61+
assertThat(chain).isNotNull();
62+
assertThat(chain.getHandler()).isSameAs(handler);
63+
64+
request.setMethod("POST");
65+
66+
chain = mapping.getHandler(request);
67+
assertThat(chain).isNull();
68+
}
69+
70+
}

0 commit comments

Comments
 (0)