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<String> on HandshakeWebSocketService.

Issue: SPR-16212

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<String, Object> 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> 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> principal,
+			@Nullable String protocol, Map<String, Object> 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<String, Object> getAttributes() {
+		return this.attributes;
+	}
+
 
 	@Override
 	public String toString() {

spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/AbstractWebSocketSession.java

@@ -54,7 +54,7 @@ public abstract class AbstractWebSocketSession<T> 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<T> implements WebSocketSession {
 		this.id = id;
 		this.handshakeInfo = handshakeInfo;
 		this.bufferFactory = bufferFactory;
+		this.attributes.putAll(handshakeInfo.getAttributes());
 	}
 
 

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<Void>} to indicate the outcome of the
 	 * WebSocket session handling.
+	 * @deprecated as of 5.1 in favor of
+	 * {@link #upgrade(ServerWebExchange, WebSocketHandler, String, Supplier)}
 	 */
-	Mono<Void> upgrade(ServerWebExchange exchange, WebSocketHandler webSocketHandler,
-			@Nullable String subProtocol);
+	@Deprecated
+	default Mono<Void> 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<Void>} to indicate the outcome of the
+	 * WebSocket session handling.
+	 * @since 5.1
+	 */
+	@SuppressWarnings("deprecation")
+	default Mono<Void> upgrade(ServerWebExchange exchange, WebSocketHandler webSocketHandler,
+			@Nullable String subProtocol, Supplier<HandshakeInfo> handshakeInfoFactory) {
+
+		return upgrade(exchange, webSocketHandler, subProtocol);
+	}
 
 }

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<Void>} that completes when application handling of
+	 * the WebSocket session completes.
 	 */
 	Mono<Void> handleRequest(ServerWebExchange exchange, WebSocketHandler webSocketHandler);
 

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<Map<String, Object>> 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<String> 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.
+	 * <p>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<String> predicate) {
+		this.sessionAttributePredicate = predicate;
+	}
+
+	/**
+	 * Return the configured predicate for initialization WebSocket session
+	 * attributes from {@code WebSession} attributes.
+	 * @since 5.1
+	 */
+	@Nullable
+	public Predicate<String> 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<Void> handleBadRequest(String reason) {
@@ -224,4 +260,21 @@ public class HandshakeWebSocketService implements WebSocketService, Lifecycle {
 		return null;
 	}
 
+	private Mono<Map<String, Object>> 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<String, Object> attributes) {
+
+		Mono<Principal> principal = exchange.getPrincipal();
+		return new HandshakeInfo(request.getURI(), request.getHeaders(), principal, protocol, attributes);
+	}
+
 }

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;
 	}

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<Void> upgrade(ServerWebExchange exchange, WebSocketHandler handler,
-			@Nullable String subProtocol) {
+			@Nullable String subProtocol, Supplier<HandshakeInfo> 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> principal = exchange.getPrincipal();
-		return new HandshakeInfo(request.getURI(), request.getHeaders(), principal, protocol);
-	}
-
 	private void startLazily(HttpServletRequest request) {
 		if (this.servletContext != null) {
 			return;

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<Void> upgrade(ServerWebExchange exchange, WebSocketHandler handler, @Nullable String subProtocol) {
+	public Mono<Void> upgrade(ServerWebExchange exchange, WebSocketHandler handler,
+			@Nullable String subProtocol, Supplier<HandshakeInfo> 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> principal = exchange.getPrincipal();
-		return new HandshakeInfo(request.getURI(), request.getHeaders(), principal, protocol);
+				(in, out) -> handler.handle(new ReactorNettyWebSocketSession(in, out, handshakeInfo, bufferFactory)));
 	}
 
 }

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<Void> upgrade(ServerWebExchange exchange, WebSocketHandler handler,
-			@Nullable String subProtocol){
+			@Nullable String subProtocol, Supplier<HandshakeInfo> 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> 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);

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<Void> upgrade(ServerWebExchange exchange, WebSocketHandler handler, @Nullable String subProtocol) {
+	public Mono<Void> upgrade(ServerWebExchange exchange, WebSocketHandler handler,
+			@Nullable String subProtocol, Supplier<HandshakeInfo> 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<Handshake> handshakes = Collections.singletonList(handshake);
 
-		URI url = request.getURI();
-		HttpHeaders headers = request.getHeaders();
-		Mono<Principal> 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) {

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<String, Object> 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<Void> upgrade(ServerWebExchange exchange, WebSocketHandler webSocketHandler,
+				@Nullable  String subProtocol, Supplier<HandshakeInfo> handshakeInfoFactory) {
+
+			this.handshakeInfo = handshakeInfoFactory.get();
+			return Mono.empty();
+		}
+	}
+
+}

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<String>` 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]#<<web.adoc#websocket-server-runtime-configuration,Same in Servlet stack>>#
 
-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: