Add cookies to the WebSocket HandshakeInfo

See gh-26674

Author: sokomishalov

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

@@ -22,6 +22,9 @@
 import java.util.Collections;
 import java.util.Map;
 
+import org.springframework.http.HttpCookie;
+import org.springframework.util.CollectionUtils;
+import org.springframework.util.MultiValueMap;
 import reactor.core.publisher.Mono;
 
 import org.springframework.http.HttpHeaders;
@@ -44,6 +47,8 @@ public class HandshakeInfo {
 
 	private final HttpHeaders headers;
 
+	private final MultiValueMap<String, HttpCookie> cookies;
+
 	@Nullable
 	private final String protocol;
 
@@ -67,6 +72,7 @@ public HandshakeInfo(URI uri, HttpHeaders headers, Mono<Principal> principal, @N
 		this(uri, headers, principal, protocol, null, Collections.emptyMap(), null);
 	}
 
+
 	/**
 	 * Constructor targetting server-side use with extra information about the
 	 * handshake, the remote address, and a pre-existing log prefix for
@@ -81,17 +87,39 @@ public HandshakeInfo(URI uri, HttpHeaders headers, Mono<Principal> principal, @N
 	 * messages, if any.
 	 * @since 5.1
 	 */
+	@Deprecated
 	public HandshakeInfo(URI uri, HttpHeaders headers, Mono<Principal> principal,
-			@Nullable String protocol, @Nullable InetSocketAddress remoteAddress,
-			Map<String, Object> attributes, @Nullable String logPrefix) {
+				@Nullable String protocol, @Nullable InetSocketAddress remoteAddress,
+				Map<String, Object> attributes, @Nullable String logPrefix) {
+		this(uri, headers, CollectionUtils.toMultiValueMap(Collections.emptyMap()), principal, protocol, remoteAddress, attributes, logPrefix);
+	}
 
+	/**
+	 * Constructor targetting server-side use with extra information about the
+	 * handshake, the remote address, and a pre-existing log prefix for
+	 * correlation.
+	 * @param uri the endpoint URL
+	 * @param headers request headers for server or response headers or client
+	 * @param cookies request cookies for server
+	 * @param principal the principal for the session
+	 * @param protocol the negotiated sub-protocol (may be {@code null})
+	 * @param remoteAddress the remote address where the handshake came from
+	 * @param attributes initial attributes to use for the WebSocket session
+	 * @param logPrefix log prefix used during the handshake for correlating log
+	 * messages, if any.
+	 * @since 5.4
+	 */
+	public HandshakeInfo(URI uri, HttpHeaders headers, MultiValueMap<String, HttpCookie> cookies,
+				Mono<Principal> principal, @Nullable String protocol, @Nullable InetSocketAddress remoteAddress,
+				Map<String, Object> attributes, @Nullable String logPrefix) {
 		Assert.notNull(uri, "URI is required");
 		Assert.notNull(headers, "HttpHeaders are required");
 		Assert.notNull(principal, "Principal is required");
 		Assert.notNull(attributes, "'attributes' is required");
 
 		this.uri = uri;
 		this.headers = headers;
+		this.cookies = cookies;
 		this.principalMono = principal;
 		this.protocol = protocol;
 		this.remoteAddress = remoteAddress;
@@ -115,6 +143,13 @@ public HttpHeaders getHeaders() {
 		return this.headers;
 	}
 
+	/**
+	 * Return the handshake HTTP cookies.
+	 */
+	public MultiValueMap<String, HttpCookie> getCookies() {
+		return this.cookies;
+	}
+
 	/**
 	 * Return the principal associated with the handshake HTTP request.
 	 */

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

@@ -32,12 +32,15 @@
 import org.springframework.context.Lifecycle;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
+import org.springframework.http.HttpCookie;
 import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
 import org.springframework.util.ReflectionUtils;
 import org.springframework.util.StringUtils;
+import org.springframework.util.MultiValueMap;
+import org.springframework.util.LinkedMultiValueMap;
 import org.springframework.web.reactive.socket.HandshakeInfo;
 import org.springframework.web.reactive.socket.WebSocketHandler;
 import org.springframework.web.reactive.socket.server.RequestUpgradeStrategy;
@@ -282,10 +285,12 @@ private HandshakeInfo createHandshakeInfo(ServerWebExchange exchange, ServerHttp
 		// the server implementation once the handshake HTTP exchange is done.
 		HttpHeaders headers = new HttpHeaders();
 		headers.addAll(request.getHeaders());
+		MultiValueMap<String, HttpCookie> cookies = new LinkedMultiValueMap<>();
+		cookies.addAll(request.getCookies());
 		Mono<Principal> principal = exchange.getPrincipal();
 		String logPrefix = exchange.getLogPrefix();
 		InetSocketAddress remoteAddress = request.getRemoteAddress();
-		return new HandshakeInfo(uri, headers, principal, protocol, remoteAddress, attributes, logPrefix);
+		return new HandshakeInfo(uri, headers, cookies, principal, protocol, remoteAddress, attributes, logPrefix);
 	}
 
 }