From 7481d73456b0e5a845912efe323ce8fddd8eb712 Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@pivotal.io>
Date: Wed, 19 Sep 2018 16:53:57 -0400
Subject: [PATCH] WebFlux HandshakeInfo exposes the remoteAddress

Issue: SPR-17192
---
 .../web/reactive/socket/HandshakeInfo.java    | 27 ++++++++++++++++---
 .../support/HandshakeWebSocketService.java    |  4 ++-
 2 files changed, 26 insertions(+), 5 deletions(-)

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 1aa678dda7..5281474aa1 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
@@ -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,6 +47,9 @@ public class HandshakeInfo {
 	@Nullable
 	private final String protocol;
 
+	@Nullable
+	private final InetSocketAddress remoteAddress;
+
 	private final Map<String, Object> attributes;
 
 	@Nullable
@@ -53,29 +57,33 @@ public class HandshakeInfo {
 
 
 	/**
-	 * 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 class HandshakeInfo {
 		this.headers = headers;
 		this.principalMono = principal;
 		this.protocol = protocol;
+		this.remoteAddress = remoteAddress;
 		this.attributes = attributes;
 		this.logPrefix = logPrefix;
 	}
@@ -123,6 +132,16 @@ public class HandshakeInfo {
 		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.
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 e9d7e3bee2..f5061f960e 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,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 @@ public class HandshakeWebSocketService implements WebSocketService, Lifecycle {
 		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);
 	}
 
 }