From 192c7a56277a651839cf9275a59eb5001b56ff99 Mon Sep 17 00:00:00 2001 From: Rossen Stoyanchev Date: Fri, 18 May 2018 21:30:13 -0400 Subject: [PATCH] WebSession to WebSocketSession attribute passing This commit makes it possible to pass attributes from the WebSession of a handshake request to the WebSocketSession, by configuring a Predicate on HandshakeWebSocketService. Issue: SPR-16212 --- .../web/reactive/socket/HandshakeInfo.java | 32 +++++- .../adapter/AbstractWebSocketSession.java | 3 +- .../socket/server/RequestUpgradeStrategy.java | 32 +++++- .../socket/server/WebSocketService.java | 7 +- .../support/HandshakeWebSocketService.java | 55 ++++++++++- .../support/WebSocketHandlerAdapter.java | 5 +- .../upgrade/JettyRequestUpgradeStrategy.java | 23 ++--- .../ReactorNettyRequestUpgradeStrategy.java | 20 ++-- .../upgrade/TomcatRequestUpgradeStrategy.java | 24 ++--- .../UndertowRequestUpgradeStrategy.java | 18 ++-- .../HandshakeWebSocketServiceTests.java | 99 +++++++++++++++++++ src/docs/asciidoc/web/webflux-websocket.adoc | 7 +- 12 files changed, 261 insertions(+), 64 deletions(-) create mode 100644 spring-webflux/src/test/java/org/springframework/web/reactive/socket/server/support/HandshakeWebSocketServiceTests.java diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/HandshakeInfo.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/HandshakeInfo.java index 0a7f6d4291..2298e3c7d9 100644 --- a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/HandshakeInfo.java +++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/HandshakeInfo.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2017 the original author or authors. + * Copyright 2002-2018 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. @@ -18,6 +18,8 @@ package org.springframework.web.reactive.socket; import java.net.URI; import java.security.Principal; +import java.util.Collections; +import java.util.Map; import reactor.core.publisher.Mono; @@ -44,6 +46,8 @@ public class HandshakeInfo { @Nullable private final String protocol; + private final Map attributes; + /** * Constructor with information about the handshake. @@ -53,13 +57,30 @@ public class HandshakeInfo { * @param protocol the negotiated sub-protocol (may be {@code null}) */ public HandshakeInfo(URI uri, HttpHeaders headers, Mono principal, @Nullable String protocol) { + this(uri, headers, principal, protocol, Collections.emptyMap()); + } + + /** + * Constructor with information about the handshake. + * @param uri the endpoint URL + * @param headers request headers for server or response headers or client + * @param principal the principal for the session + * @param protocol the negotiated sub-protocol (may be {@code null}) + * @since 5.1 + */ + public HandshakeInfo(URI uri, HttpHeaders headers, Mono principal, + @Nullable String protocol, Map attributes) { + Assert.notNull(uri, "URI is required"); Assert.notNull(headers, "HttpHeaders are required"); Assert.notNull(principal, "Principal is required"); + Assert.notNull(principal, "'attributes' is required"); + this.uri = uri; this.headers = headers; this.principalMono = principal; this.protocol = protocol; + this.attributes = attributes; } @@ -95,6 +116,15 @@ public class HandshakeInfo { return this.protocol; } + /** + * Attributes extracted from the handshake request to be added to the + * WebSocket session. + * @since 5.1 + */ + public Map getAttributes() { + return this.attributes; + } + @Override public String toString() { diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/AbstractWebSocketSession.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/AbstractWebSocketSession.java index 72dc7616cf..2e342d81d4 100644 --- a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/AbstractWebSocketSession.java +++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/AbstractWebSocketSession.java @@ -54,7 +54,7 @@ public abstract class AbstractWebSocketSession implements WebSocketSession { /** - * Create a new instance and associate the given attributes with it. + * Create a new WebSocket session. */ protected AbstractWebSocketSession(T delegate, String id, HandshakeInfo handshakeInfo, DataBufferFactory bufferFactory) { @@ -68,6 +68,7 @@ public abstract class AbstractWebSocketSession implements WebSocketSession { this.id = id; this.handshakeInfo = handshakeInfo; this.bufferFactory = bufferFactory; + this.attributes.putAll(handshakeInfo.getAttributes()); } diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/RequestUpgradeStrategy.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/RequestUpgradeStrategy.java index 568877faae..5c0c960258 100644 --- a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/RequestUpgradeStrategy.java +++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/RequestUpgradeStrategy.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2017 the original author or authors. + * Copyright 2002-2018 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. @@ -16,11 +16,14 @@ package org.springframework.web.reactive.socket.server; +import java.util.function.Supplier; + import reactor.core.publisher.Mono; import org.springframework.http.server.reactive.ServerHttpRequest; import org.springframework.http.server.reactive.ServerHttpResponse; import org.springframework.lang.Nullable; +import org.springframework.web.reactive.socket.HandshakeInfo; import org.springframework.web.reactive.socket.WebSocketHandler; import org.springframework.web.server.ServerWebExchange; @@ -45,8 +48,31 @@ public interface RequestUpgradeStrategy { * @param subProtocol the selected sub-protocol got the handler * @return completion {@code Mono} to indicate the outcome of the * WebSocket session handling. + * @deprecated as of 5.1 in favor of + * {@link #upgrade(ServerWebExchange, WebSocketHandler, String, Supplier)} */ - Mono upgrade(ServerWebExchange exchange, WebSocketHandler webSocketHandler, - @Nullable String subProtocol); + @Deprecated + default Mono upgrade(ServerWebExchange exchange, WebSocketHandler webSocketHandler, + @Nullable String subProtocol) { + + return Mono.error(new UnsupportedOperationException()); + } + + /** + * Upgrade to a WebSocket session and handle it with the given handler. + * @param exchange the current exchange + * @param webSocketHandler handler for the WebSocket session + * @param subProtocol the selected sub-protocol got the handler + * @param handshakeInfoFactory factory to create HandshakeInfo for the WebSocket session + * @return completion {@code Mono} to indicate the outcome of the + * WebSocket session handling. + * @since 5.1 + */ + @SuppressWarnings("deprecation") + default Mono upgrade(ServerWebExchange exchange, WebSocketHandler webSocketHandler, + @Nullable String subProtocol, Supplier handshakeInfoFactory) { + + return upgrade(exchange, webSocketHandler, subProtocol); + } } diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/WebSocketService.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/WebSocketService.java index 54bd939b11..d7bd875237 100644 --- a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/WebSocketService.java +++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/WebSocketService.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2017 the original author or authors. + * Copyright 2002-2018 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. @@ -35,10 +35,11 @@ import org.springframework.web.server.ServerWebExchange; public interface WebSocketService { /** - * Handle the HTTP request and use the given {@link WebSocketHandler}. + * Handle the request with the given {@link WebSocketHandler}. * @param exchange the current exchange * @param webSocketHandler handler for WebSocket session - * @return a completion Mono for the WebSocket session handling + * @return a {@code Mono} that completes when application handling of + * the WebSocket session completes. */ Mono handleRequest(ServerWebExchange exchange, WebSocketHandler webSocketHandler); diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/support/HandshakeWebSocketService.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/support/HandshakeWebSocketService.java index 457c2440c9..8184d1d546 100644 --- a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/support/HandshakeWebSocketService.java +++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/support/HandshakeWebSocketService.java @@ -16,8 +16,12 @@ package org.springframework.web.reactive.socket.server.support; +import java.security.Principal; import java.util.Collections; import java.util.List; +import java.util.Map; +import java.util.function.Predicate; +import java.util.stream.Collectors; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; @@ -32,6 +36,7 @@ import org.springframework.util.Assert; import org.springframework.util.ClassUtils; import org.springframework.util.ReflectionUtils; import org.springframework.util.StringUtils; +import org.springframework.web.reactive.socket.HandshakeInfo; import org.springframework.web.reactive.socket.WebSocketHandler; import org.springframework.web.reactive.socket.server.RequestUpgradeStrategy; import org.springframework.web.reactive.socket.server.WebSocketService; @@ -54,6 +59,8 @@ public class HandshakeWebSocketService implements WebSocketService, Lifecycle { private static final String SEC_WEBSOCKET_PROTOCOL = "Sec-WebSocket-Protocol"; + private static final Mono> EMPTY_ATTRIBUTES = Mono.just(Collections.emptyMap()); + private static final boolean tomcatPresent = ClassUtils.isPresent( "org.apache.tomcat.websocket.server.WsHttpUpgradeHandler", @@ -77,6 +84,9 @@ public class HandshakeWebSocketService implements WebSocketService, Lifecycle { private final RequestUpgradeStrategy upgradeStrategy; + @Nullable + private Predicate sessionAttributePredicate; + private volatile boolean running = false; @@ -135,6 +145,28 @@ public class HandshakeWebSocketService implements WebSocketService, Lifecycle { return this.upgradeStrategy; } + /** + * Configure a predicate to use to extract + * {@link org.springframework.web.server.WebSession WebSession} attributes + * and use them to initialize the WebSocket session with. + *

By default this is not set in which case no attributes are passed. + * @param predicate the predicate + * @since 5.1 + */ + public void setSessionAttributePredicate(@Nullable Predicate predicate) { + this.sessionAttributePredicate = predicate; + } + + /** + * Return the configured predicate for initialization WebSocket session + * attributes from {@code WebSession} attributes. + * @since 5.1 + */ + @Nullable + public Predicate getSessionAttributePredicate() { + return this.sessionAttributePredicate; + } + @Override public void start() { @@ -200,7 +232,11 @@ public class HandshakeWebSocketService implements WebSocketService, Lifecycle { } String protocol = selectProtocol(headers, handler); - return this.upgradeStrategy.upgrade(exchange, handler, protocol); + + return initAttributes(exchange).flatMap(attributes -> + this.upgradeStrategy.upgrade(exchange, handler, protocol, + () -> createHandshakeInfo(exchange, request, protocol, attributes)) + ); } private Mono handleBadRequest(String reason) { @@ -224,4 +260,21 @@ public class HandshakeWebSocketService implements WebSocketService, Lifecycle { return null; } + private Mono> initAttributes(ServerWebExchange exchange) { + if (this.sessionAttributePredicate == null) { + return EMPTY_ATTRIBUTES; + } + return exchange.getSession().map(session -> + session.getAttributes().entrySet().stream() + .filter(entry -> this.sessionAttributePredicate.test(entry.getKey())) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))); + } + + private HandshakeInfo createHandshakeInfo(ServerWebExchange exchange, ServerHttpRequest request, + @Nullable String protocol, Map attributes) { + + Mono principal = exchange.getPrincipal(); + return new HandshakeInfo(request.getURI(), request.getHeaders(), principal, protocol, attributes); + } + } diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/support/WebSocketHandlerAdapter.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/support/WebSocketHandlerAdapter.java index e492a176b7..819c75e021 100644 --- a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/support/WebSocketHandlerAdapter.java +++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/support/WebSocketHandlerAdapter.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2016 the original author or authors. + * Copyright 2002-2018 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. @@ -56,6 +56,9 @@ public class WebSocketHandlerAdapter implements HandlerAdapter { } + /** + * Return the configured {@code WebSocketService} to handle requests. + */ public WebSocketService getWebSocketService() { return this.webSocketService; } diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/JettyRequestUpgradeStrategy.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/JettyRequestUpgradeStrategy.java index 3c43993ea5..81e3bbfc3c 100644 --- a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/JettyRequestUpgradeStrategy.java +++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/JettyRequestUpgradeStrategy.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2017 the original author or authors. + * Copyright 2002-2018 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. @@ -17,7 +17,7 @@ package org.springframework.web.reactive.socket.server.upgrade; import java.io.IOException; -import java.security.Principal; +import java.util.function.Supplier; import javax.servlet.ServletContext; import javax.servlet.http.HttpServletRequest; import javax.servlet.http.HttpServletResponse; @@ -140,7 +140,7 @@ public class JettyRequestUpgradeStrategy implements RequestUpgradeStrategy, Life @Override public Mono upgrade(ServerWebExchange exchange, WebSocketHandler handler, - @Nullable String subProtocol) { + @Nullable String subProtocol, Supplier handshakeInfoFactory) { ServerHttpRequest request = exchange.getRequest(); ServerHttpResponse response = exchange.getResponse(); @@ -148,12 +148,11 @@ public class JettyRequestUpgradeStrategy implements RequestUpgradeStrategy, Life HttpServletRequest servletRequest = getHttpServletRequest(request); HttpServletResponse servletResponse = getHttpServletResponse(response); - JettyWebSocketHandlerAdapter adapter = new JettyWebSocketHandlerAdapter(handler, - session -> { - HandshakeInfo info = getHandshakeInfo(exchange, subProtocol); - DataBufferFactory factory = response.bufferFactory(); - return new JettyWebSocketSession(session, info, factory); - }); + HandshakeInfo handshakeInfo = handshakeInfoFactory.get(); + DataBufferFactory factory = response.bufferFactory(); + + JettyWebSocketHandlerAdapter adapter = new JettyWebSocketHandlerAdapter( + handler, session -> new JettyWebSocketSession(session, handshakeInfo, factory)); startLazily(servletRequest); @@ -185,12 +184,6 @@ public class JettyRequestUpgradeStrategy implements RequestUpgradeStrategy, Life return ((AbstractServerHttpResponse) response).getNativeResponse(); } - private HandshakeInfo getHandshakeInfo(ServerWebExchange exchange, @Nullable String protocol) { - ServerHttpRequest request = exchange.getRequest(); - Mono principal = exchange.getPrincipal(); - return new HandshakeInfo(request.getURI(), request.getHeaders(), principal, protocol); - } - private void startLazily(HttpServletRequest request) { if (this.servletContext != null) { return; diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/ReactorNettyRequestUpgradeStrategy.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/ReactorNettyRequestUpgradeStrategy.java index 556abfe359..8d8af05ad7 100644 --- a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/ReactorNettyRequestUpgradeStrategy.java +++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/ReactorNettyRequestUpgradeStrategy.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2017 the original author or authors. + * Copyright 2002-2018 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. @@ -16,14 +16,13 @@ package org.springframework.web.reactive.socket.server.upgrade; -import java.security.Principal; +import java.util.function.Supplier; import reactor.core.publisher.Mono; import reactor.ipc.netty.http.server.HttpServerResponse; import org.springframework.core.io.buffer.NettyDataBufferFactory; import org.springframework.http.server.reactive.AbstractServerHttpResponse; -import org.springframework.http.server.reactive.ServerHttpRequest; import org.springframework.http.server.reactive.ServerHttpResponse; import org.springframework.lang.Nullable; import org.springframework.web.reactive.socket.HandshakeInfo; @@ -40,21 +39,18 @@ import org.springframework.web.server.ServerWebExchange; */ public class ReactorNettyRequestUpgradeStrategy implements RequestUpgradeStrategy { + @Override - public Mono upgrade(ServerWebExchange exchange, WebSocketHandler handler, @Nullable String subProtocol) { + public Mono upgrade(ServerWebExchange exchange, WebSocketHandler handler, + @Nullable String subProtocol, Supplier handshakeInfoFactory) { + ServerHttpResponse response = exchange.getResponse(); HttpServerResponse nativeResponse = ((AbstractServerHttpResponse) response).getNativeResponse(); - HandshakeInfo info = getHandshakeInfo(exchange, subProtocol); + HandshakeInfo handshakeInfo = handshakeInfoFactory.get(); NettyDataBufferFactory bufferFactory = (NettyDataBufferFactory) response.bufferFactory(); return nativeResponse.sendWebsocket(subProtocol, - (in, out) -> handler.handle(new ReactorNettyWebSocketSession(in, out, info, bufferFactory))); - } - - private HandshakeInfo getHandshakeInfo(ServerWebExchange exchange, @Nullable String protocol) { - ServerHttpRequest request = exchange.getRequest(); - Mono principal = exchange.getPrincipal(); - return new HandshakeInfo(request.getURI(), request.getHeaders(), principal, protocol); + (in, out) -> handler.handle(new ReactorNettyWebSocketSession(in, out, handshakeInfo, bufferFactory))); } } diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/TomcatRequestUpgradeStrategy.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/TomcatRequestUpgradeStrategy.java index a4d396a722..d487bb0d4a 100644 --- a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/TomcatRequestUpgradeStrategy.java +++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/TomcatRequestUpgradeStrategy.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2017 the original author or authors. + * Copyright 2002-2018 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. @@ -17,8 +17,8 @@ package org.springframework.web.reactive.socket.server.upgrade; import java.io.IOException; -import java.security.Principal; import java.util.Collections; +import java.util.function.Supplier; import javax.servlet.ServletException; import javax.servlet.http.HttpServletRequest; import javax.servlet.http.HttpServletResponse; @@ -46,6 +46,7 @@ import org.springframework.web.server.ServerWebExchange; * A {@link RequestUpgradeStrategy} for use with Tomcat. * * @author Violeta Georgieva + * @author Rossen Stoyanchev * @since 5.0 */ public class TomcatRequestUpgradeStrategy implements RequestUpgradeStrategy { @@ -124,7 +125,7 @@ public class TomcatRequestUpgradeStrategy implements RequestUpgradeStrategy { @Override public Mono upgrade(ServerWebExchange exchange, WebSocketHandler handler, - @Nullable String subProtocol){ + @Nullable String subProtocol, Supplier handshakeInfoFactory){ ServerHttpRequest request = exchange.getRequest(); ServerHttpResponse response = exchange.getResponse(); @@ -132,12 +133,11 @@ public class TomcatRequestUpgradeStrategy implements RequestUpgradeStrategy { HttpServletRequest servletRequest = getHttpServletRequest(request); HttpServletResponse servletResponse = getHttpServletResponse(response); - Endpoint endpoint = new StandardWebSocketHandlerAdapter(handler, - session -> { - HandshakeInfo info = getHandshakeInfo(exchange, subProtocol); - DataBufferFactory factory = response.bufferFactory(); - return new TomcatWebSocketSession(session, info, factory); - }); + HandshakeInfo handshakeInfo = handshakeInfoFactory.get(); + DataBufferFactory bufferFactory = response.bufferFactory(); + + Endpoint endpoint = new StandardWebSocketHandlerAdapter( + handler, session -> new TomcatWebSocketSession(session, handshakeInfo, bufferFactory)); String requestURI = servletRequest.getRequestURI(); DefaultServerEndpointConfig config = new DefaultServerEndpointConfig(requestURI, endpoint); @@ -165,12 +165,6 @@ public class TomcatRequestUpgradeStrategy implements RequestUpgradeStrategy { return ((AbstractServerHttpResponse) response).getNativeResponse(); } - private HandshakeInfo getHandshakeInfo(ServerWebExchange exchange, @Nullable String protocol) { - ServerHttpRequest request = exchange.getRequest(); - Mono principal = exchange.getPrincipal(); - return new HandshakeInfo(request.getURI(), request.getHeaders(), principal, protocol); - } - private WsServerContainer getContainer(HttpServletRequest request) { if (this.serverContainer == null) { Object container = request.getServletContext().getAttribute(SERVER_CONTAINER_ATTR); diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/UndertowRequestUpgradeStrategy.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/UndertowRequestUpgradeStrategy.java index 663cabf038..fc940e1434 100644 --- a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/UndertowRequestUpgradeStrategy.java +++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/UndertowRequestUpgradeStrategy.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2017 the original author or authors. + * Copyright 2002-2018 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. @@ -16,11 +16,10 @@ package org.springframework.web.reactive.socket.server.upgrade; -import java.net.URI; -import java.security.Principal; import java.util.Collections; import java.util.List; import java.util.Set; +import java.util.function.Supplier; import io.undertow.server.HttpServerExchange; import io.undertow.websockets.WebSocketConnectionCallback; @@ -32,7 +31,6 @@ import io.undertow.websockets.spi.WebSocketHttpExchange; import reactor.core.publisher.Mono; import org.springframework.core.io.buffer.DataBufferFactory; -import org.springframework.http.HttpHeaders; import org.springframework.http.server.reactive.AbstractServerHttpRequest; import org.springframework.http.server.reactive.ServerHttpRequest; import org.springframework.lang.Nullable; @@ -48,12 +46,15 @@ import org.springframework.web.server.ServerWebExchange; * A {@link RequestUpgradeStrategy} for use with Undertow. * * @author Violeta Georgieva + * @author Rossen Stoyanchev * @since 5.0 */ public class UndertowRequestUpgradeStrategy implements RequestUpgradeStrategy { @Override - public Mono upgrade(ServerWebExchange exchange, WebSocketHandler handler, @Nullable String subProtocol) { + public Mono upgrade(ServerWebExchange exchange, WebSocketHandler handler, + @Nullable String subProtocol, Supplier handshakeInfoFactory) { + ServerHttpRequest request = exchange.getRequest(); Assert.isInstanceOf(AbstractServerHttpRequest.class, request); HttpServerExchange httpExchange = ((AbstractServerHttpRequest) request).getNativeRequest(); @@ -62,14 +63,11 @@ public class UndertowRequestUpgradeStrategy implements RequestUpgradeStrategy { Hybi13Handshake handshake = new Hybi13Handshake(protocols, false); List handshakes = Collections.singletonList(handshake); - URI url = request.getURI(); - HttpHeaders headers = request.getHeaders(); - Mono principal = exchange.getPrincipal(); - HandshakeInfo info = new HandshakeInfo(url, headers, principal, subProtocol); + HandshakeInfo handshakeInfo = handshakeInfoFactory.get(); DataBufferFactory bufferFactory = exchange.getResponse().bufferFactory(); try { - DefaultCallback callback = new DefaultCallback(info, handler, bufferFactory); + DefaultCallback callback = new DefaultCallback(handshakeInfo, handler, bufferFactory); new WebSocketProtocolHandshakeHandler(handshakes, callback).handleRequest(httpExchange); } catch (Exception ex) { diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/socket/server/support/HandshakeWebSocketServiceTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/socket/server/support/HandshakeWebSocketServiceTests.java new file mode 100644 index 0000000000..9a546d9dee --- /dev/null +++ b/spring-webflux/src/test/java/org/springframework/web/reactive/socket/server/support/HandshakeWebSocketServiceTests.java @@ -0,0 +1,99 @@ +/* + * Copyright 2002-2018 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 + * + * http://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.Arrays; +import java.util.Map; +import java.util.function.Supplier; + +import org.hamcrest.Matchers; +import org.junit.Test; +import reactor.core.publisher.Mono; + +import org.springframework.lang.Nullable; +import org.springframework.mock.http.server.reactive.test.MockServerHttpRequest; +import org.springframework.mock.web.test.server.MockServerWebExchange; +import org.springframework.mock.web.test.server.MockWebSession; +import org.springframework.web.reactive.socket.HandshakeInfo; +import org.springframework.web.reactive.socket.WebSocketHandler; +import org.springframework.web.reactive.socket.server.RequestUpgradeStrategy; +import org.springframework.web.server.ServerWebExchange; + +import static org.junit.Assert.*; +import static org.mockito.Mockito.*; + +/** + * Unit tests for {@link HandshakeWebSocketService}. + * + * @author Rossen Stoyanchev + */ +public class HandshakeWebSocketServiceTests { + + + @Test + public void sessionAttributePredicate() { + + MockWebSession session = new MockWebSession(); + session.getAttributes().put("a1", "v1"); + session.getAttributes().put("a2", "v2"); + session.getAttributes().put("a3", "v3"); + session.getAttributes().put("a4", "v4"); + session.getAttributes().put("a5", "v5"); + + MockServerHttpRequest request = initHandshakeRequest(); + MockServerWebExchange exchange = MockServerWebExchange.builder(request).session(session).build(); + + TestRequestUpgradeStrategy upgradeStrategy = new TestRequestUpgradeStrategy(); + HandshakeWebSocketService service = new HandshakeWebSocketService(upgradeStrategy); + service.setSessionAttributePredicate(name -> Arrays.asList("a1", "a3", "a5").contains(name)); + + service.handleRequest(exchange, mock(WebSocketHandler.class)).block(); + + HandshakeInfo info = upgradeStrategy.handshakeInfo; + assertNotNull(info); + + Map attributes = info.getAttributes(); + assertEquals(3, attributes.size()); + assertThat(attributes, Matchers.hasEntry("a1", "v1")); + assertThat(attributes, Matchers.hasEntry("a3", "v3")); + assertThat(attributes, Matchers.hasEntry("a5", "v5")); + } + + private MockServerHttpRequest initHandshakeRequest() { + return MockServerHttpRequest.get("/") + .header("upgrade", "websocket") + .header("connection", "upgrade") + .header("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==") + .header("Sec-WebSocket-Version", "13") + .build(); + } + + + private static class TestRequestUpgradeStrategy implements RequestUpgradeStrategy { + + HandshakeInfo handshakeInfo; + + + @Override + public Mono upgrade(ServerWebExchange exchange, WebSocketHandler webSocketHandler, + @Nullable String subProtocol, Supplier handshakeInfoFactory) { + + this.handshakeInfo = handshakeInfoFactory.get(); + return Mono.empty(); + } + } + +} diff --git a/src/docs/asciidoc/web/webflux-websocket.adoc b/src/docs/asciidoc/web/webflux-websocket.adoc index 17a36d9b48..6e1e917a51 100644 --- a/src/docs/asciidoc/web/webflux-websocket.adoc +++ b/src/docs/asciidoc/web/webflux-websocket.adoc @@ -202,7 +202,10 @@ of `HandshakeWebSocketService`, which performs basic checks on the WebSocket req then uses `RequestUpgradeStrategy` for the server in use. Currently there is built-in support for Reactor Netty, Tomcat, Jetty, and Undertow. -The above are just 3 examples to serve as a starting point. +`HandshakeWebSocketService` exposes a `sessionAttributePredicate` property that allows +setting a `Predicate` to extract attributes from the `WebSession` and insert them +into the attributes of the `WebSocketSession`. + @@ -210,7 +213,7 @@ The above are just 3 examples to serve as a starting point. === Server config [.small]#<># -The `RequestUpgradeStrategy` for each server exposes the WebSocket-related configuration +The `RequestUpgradeStrategy` for each server exposes WebSocket-related configuration options available for the underlying WebSocket engine. Below is an example of setting WebSocket options when running on Tomcat: