WebFlux HandshakeInfo exposes the remoteAddress

Issue: SPR-17192

Author: rstoyanchev

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

@@ -16,6 +16,7 @@
 
 package org.springframework.web.reactive.socket;
 
+import java.net.InetSocketAddress;
 import java.net.URI;
 import java.security.Principal;
 import java.util.Collections;
@@ -46,36 +47,43 @@ public class HandshakeInfo {
 	@Nullable
 	private final String protocol;
 
+	@Nullable
+	private final InetSocketAddress remoteAddress;
+
 	private final Map<String, Object> attributes;
 
 	@Nullable
 	private final String logPrefix;
 
 
 	/**
-	 * Constructor with information about the handshake.
+	 * Constructor with basic 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})
 	 */
 	public HandshakeInfo(URI uri, HttpHeaders headers, Mono<Principal> principal, @Nullable String protocol) {
-		this(uri, headers, principal, protocol, Collections.emptyMap(), null);
+		this(uri, headers, principal, protocol, null, Collections.emptyMap(), null);
 	}
 
 	/**
-	 * Constructor with information about the handshake.
+	 * 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 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.1
 	 */
 	public HandshakeInfo(URI uri, HttpHeaders headers, Mono<Principal> principal,
-			@Nullable String protocol, Map<String, Object> attributes, @Nullable String logPrefix) {
+			@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");
@@ -86,6 +94,7 @@ public HandshakeInfo(URI uri, HttpHeaders headers, Mono<Principal> principal,
 		this.headers = headers;
 		this.principalMono = principal;
 		this.protocol = protocol;
+		this.remoteAddress = remoteAddress;
 		this.attributes = attributes;
 		this.logPrefix = logPrefix;
 	}
@@ -123,6 +132,16 @@ public String getSubProtocol() {
 		return this.protocol;
 	}
 
+	/**
+	 * For a server-side session this is the remote address where the handshake
+	 * request came from.
+	 * @since 5.1
+	 */
+	@Nullable
+	public InetSocketAddress getRemoteAddress() {
+		return this.remoteAddress;
+	}
+
 	/**
 	 * Attributes extracted from the handshake request to be added to the
 	 * WebSocket session.

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

@@ -16,6 +16,7 @@
 
 package org.springframework.web.reactive.socket.server.support;
 
+import java.net.InetSocketAddress;
 import java.net.URI;
 import java.security.Principal;
 import java.util.Collections;
@@ -274,7 +275,8 @@ private HandshakeInfo createHandshakeInfo(ServerWebExchange exchange, ServerHttp
 		HttpHeaders headers = request.getHeaders();
 		Mono<Principal> principal = exchange.getPrincipal();
 		String logPrefix = exchange.getLogPrefix();
-		return new HandshakeInfo(uri, headers, principal, protocol, attributes, logPrefix);
+		InetSocketAddress remoteAddress = request.getRemoteAddress();
+		return new HandshakeInfo(uri, headers, principal, protocol, remoteAddress, attributes, logPrefix);
 	}
 
 }