Polish WebSocketSession

Update methods available on WebSocketSession interface.
Introduce DelegatingWebSocketSession interface.

Author: rstoyanchev

spring-messaging/src/main/java/org/springframework/messaging/handler/websocket/SubProtocolWebSocketHandler.java

@@ -30,6 +30,7 @@
 import org.springframework.messaging.MessagingException;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
+import org.springframework.util.StringUtils;
 import org.springframework.web.socket.CloseStatus;
 import org.springframework.web.socket.WebSocketHandler;
 import org.springframework.web.socket.WebSocketMessage;
@@ -133,7 +134,7 @@ public void afterConnectionEstablished(WebSocketSession session) throws Exceptio
 	protected final SubProtocolHandler getProtocolHandler(WebSocketSession session) {
 		SubProtocolHandler handler;
 		String protocol = session.getAcceptedProtocol();
-		if (protocol != null) {
+		if (!StringUtils.isEmpty(protocol)) {
 			handler = this.protocolHandlers.get(protocol);
 			Assert.state(handler != null,
 					"No handler for sub-protocol '" + protocol + "', handlers=" + this.protocolHandlers);

spring-messaging/src/test/java/org/springframework/messaging/handler/websocket/SubProtocolWebSocketHandlerTests.java

@@ -91,7 +91,18 @@ public void subProtocolNoMatch() throws Exception {
 	}
 
 	@Test
-	public void noSubProtocol() throws Exception {
+	public void nullSubProtocol() throws Exception {
+		this.webSocketHandler.setDefaultProtocolHandler(defaultHandler);
+		this.webSocketHandler.afterConnectionEstablished(session);
+
+		verify(this.defaultHandler).afterSessionStarted(session, this.channel);
+		verify(this.stompHandler, times(0)).afterSessionStarted(session, this.channel);
+		verify(this.mqttHandler, times(0)).afterSessionStarted(session, this.channel);
+	}
+
+	@Test
+	public void emptySubProtocol() throws Exception {
+		this.session.setAcceptedProtocol("");
 		this.webSocketHandler.setDefaultProtocolHandler(defaultHandler);
 		this.webSocketHandler.afterConnectionEstablished(session);
 

spring-web/src/main/java/org/springframework/http/server/ServerHttpRequest.java

@@ -16,6 +16,7 @@
 
 package org.springframework.http.server;
 
+import java.net.InetSocketAddress;
 import java.security.Principal;
 import java.util.Map;
 
@@ -51,14 +52,14 @@ public interface ServerHttpRequest extends HttpRequest, HttpInputMessage {
 	Principal getPrincipal();
 
 	/**
-	 * Return the host name of the endpoint on the other end.
+	 * Return the address on which the request was received.
 	 */
-	String getRemoteHostName();
+	InetSocketAddress getLocalAddress();
 
 	/**
-	 * Return the IP address of the endpoint on the other end.
+	 * Return the address of the remote client.
 	 */
-	String getRemoteAddress();
+	InetSocketAddress getRemoteAddress();
 
 	/**
 	 * Return a control that allows putting the request in asynchronous mode so the

spring-web/src/main/java/org/springframework/http/server/ServletServerHttpRequest.java

@@ -22,6 +22,7 @@
 import java.io.InputStream;
 import java.io.OutputStreamWriter;
 import java.io.Writer;
+import java.net.InetSocketAddress;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.URLEncoder;
@@ -147,13 +148,13 @@ public Principal getPrincipal() {
 	}
 
 	@Override
-	public String getRemoteHostName() {
-		return this.servletRequest.getRemoteHost();
+	public InetSocketAddress getLocalAddress() {
+		return new InetSocketAddress(this.servletRequest.getLocalName(), this.servletRequest.getLocalPort());
 	}
 
 	@Override
-	public String getRemoteAddress() {
-		return this.servletRequest.getRemoteAddr();
+	public InetSocketAddress getRemoteAddress() {
+		return new InetSocketAddress(this.servletRequest.getRemoteHost(), this.servletRequest.getRemotePort());
 	}
 
 	@Override

spring-websocket/src/main/java/org/springframework/web/socket/WebSocketSession.java

@@ -17,9 +17,12 @@
 package org.springframework.web.socket;
 
 import java.io.IOException;
+import java.net.InetSocketAddress;
 import java.net.URI;
 import java.security.Principal;
 
+import org.springframework.http.HttpHeaders;
+
 /**
  * A WebSocket session abstraction. Allows sending messages over a WebSocket connection
  * and closing it.
@@ -29,6 +32,7 @@
  */
 public interface WebSocketSession {
 
+
 	/**
 	 * Return a unique session identifier.
 	 */
@@ -40,9 +44,9 @@ public interface WebSocketSession {
 	URI getUri();
 
 	/**
-     * Return whether the underlying socket is using a secure transport.
+	 * Return the headers used in the handshake request.
 	 */
-	boolean isSecure();
+	HttpHeaders getHandshakeHeaders();
 
 	/**
 	 * Return a {@link java.security.Principal} instance containing the name of the
@@ -52,17 +56,18 @@ public interface WebSocketSession {
 	Principal getPrincipal();
 
 	/**
-	 * Return the host name of the endpoint on the other end.
+	 * Return the address on which the request was received.
 	 */
-	String getRemoteHostName();
+	InetSocketAddress getLocalAddress();
 
 	/**
-	 * Return the IP address of the endpoint on the other end.
+	 * Return the address of the remote client.
 	 */
-	String getRemoteAddress();
+	InetSocketAddress getRemoteAddress();
 
 	/**
-	 * Return the negotiated sub-protocol or {@code null} if none was specified.
+	 * Return the negotiated sub-protocol or {@code null} if none was specified or
+	 * negotiated successfully.
 	 */
 	String getAcceptedProtocol();
 

spring-websocket/src/main/java/org/springframework/web/socket/adapter/AbstractWebSocketSesssionAdapter.java renamed to spring-websocket/src/main/java/org/springframework/web/socket/adapter/AbstractWebSocketSesssion.java

@@ -32,19 +32,41 @@
  * @author Rossen Stoyanchev
  * @since 4.0
  */
-public abstract class AbstractWebSocketSesssionAdapter<T> implements ConfigurableWebSocketSession {
+public abstract class AbstractWebSocketSesssion<T> implements DelegatingWebSocketSession<T> {
 
 	protected final Log logger = LogFactory.getLog(getClass());
 
+	private T delegateSession;
 
-	public abstract void initSession(T session);
+
+	/**
+	 * @return the WebSocket session to delegate to
+	 */
+	public T getDelegateSession() {
+		return this.delegateSession;
+	}
+
+
+	@Override
+	public void afterSessionInitialized(T session) {
+		Assert.notNull(session, "session must not be null");
+		this.delegateSession = session;
+	}
+
+	protected final void checkDelegateSessionInitialized() {
+		Assert.state(this.delegateSession != null, "WebSocket session is not yet initialized");
+	}
 
 	@Override
 	public final void sendMessage(WebSocketMessage message) throws IOException {
+
+		checkDelegateSessionInitialized();
+		Assert.isTrue(isOpen(), "Cannot send message after connection closed.");
+
 		if (logger.isTraceEnabled()) {
 			logger.trace("Sending " + message + ", " + this);
 		}
-		Assert.isTrue(isOpen(), "Cannot send message after connection closed.");
+
 		if (message instanceof TextMessage) {
 			sendTextMessage((TextMessage) message);
 		}
@@ -60,13 +82,15 @@ else if (message instanceof BinaryMessage) {
 
 	protected abstract void sendBinaryMessage(BinaryMessage message) throws IOException ;
 
+
 	@Override
-	public void close() throws IOException {
+	public final void close() throws IOException {
 		close(CloseStatus.NORMAL);
 	}
 
 	@Override
 	public final void close(CloseStatus status) throws IOException {
+		checkDelegateSessionInitialized();
 		if (logger.isDebugEnabled()) {
 			logger.debug("Closing " + this);
 		}
@@ -75,6 +99,7 @@ public final void close(CloseStatus status) throws IOException {
 
 	protected abstract void closeInternal(CloseStatus status) throws IOException;
 
+
 	@Override
 	public String toString() {
 		return "WebSocket session id=" + getId();

spring-websocket/src/main/java/org/springframework/web/socket/adapter/ConfigurableWebSocketSession.java renamed to spring-websocket/src/main/java/org/springframework/web/socket/adapter/DelegatingWebSocketSession.java

@@ -16,34 +16,25 @@
 
 package org.springframework.web.socket.adapter;
 
-import java.net.URI;
-import java.security.Principal;
-
 import org.springframework.web.socket.WebSocketSession;
-import org.springframework.web.socket.server.DefaultHandshakeHandler;
+
 
 /**
- * A WebSocketSession with configurable properties.
+ * A contract for {@link WebSocketSession} implementations that delegate to another
+ * WebSocket session (e.g. a native session).
+ *
+ * @param T the type of the delegate WebSocket session
  *
  * @author Rossen Stoyanchev
  * @since 4.0
  */
-public interface ConfigurableWebSocketSession extends WebSocketSession {
-
-	void setUri(URI uri);
+public interface DelegatingWebSocketSession<T> extends WebSocketSession {
 
-	void setRemoteHostName(String name);
-
-	void setRemoteAddress(String address);
-
-	void setPrincipal(Principal principal);
 
 	/**
-	 * Set the protocol accepted as part of the WebSocket handshake. This property can be
-	 * used when the WebSocket handshake is performed through
-	 * {@link DefaultHandshakeHandler} rather than the underlying WebSocket runtime, or
-	 * when there is no WebSocket handshake (e.g. SockJS HTTP fallback options)
+	 * Invoked when the delegate WebSocket session has been initialized.
 	 */
-	void setAcceptedProtocol(String protocol);
+	void afterSessionInitialized(T session);
+
 
 }

spring-websocket/src/main/java/org/springframework/web/socket/adapter/JettyWebSocketListenerAdapter.java renamed to spring-websocket/src/main/java/org/springframework/web/socket/adapter/JettyWebSocketHandlerAdapter.java

@@ -28,21 +28,22 @@
 import org.springframework.web.socket.support.ExceptionWebSocketHandlerDecorator;
 
 /**
- * Adapts {@link WebSocketHandler} to the Jetty 9 {@link WebSocketListener}.
+ * Adapts {@link WebSocketHandler} to the Jetty 9 WebSocket API.
  *
  * @author Phillip Webb
+ * @author Rossen Stoyanchev
  * @since 4.0
  */
-public class JettyWebSocketListenerAdapter implements WebSocketListener {
+public class JettyWebSocketHandlerAdapter implements WebSocketListener {
 
-	private static final Log logger = LogFactory.getLog(JettyWebSocketListenerAdapter.class);
+	private static final Log logger = LogFactory.getLog(JettyWebSocketHandlerAdapter.class);
 
 	private final WebSocketHandler webSocketHandler;
 
-	private final JettyWebSocketSessionAdapter wsSession;
+	private final JettyWebSocketSession wsSession;
 
 
-	public JettyWebSocketListenerAdapter(WebSocketHandler webSocketHandler, JettyWebSocketSessionAdapter wsSession) {
+	public JettyWebSocketHandlerAdapter(WebSocketHandler webSocketHandler, JettyWebSocketSession wsSession) {
 		Assert.notNull(webSocketHandler, "webSocketHandler must not be null");
 		Assert.notNull(wsSession, "wsSession must not be null");
 		this.webSocketHandler = webSocketHandler;
@@ -52,8 +53,8 @@ public JettyWebSocketListenerAdapter(WebSocketHandler webSocketHandler, JettyWeb
 
 	@Override
 	public void onWebSocketConnect(Session session) {
-		this.wsSession.initSession(session);
 		try {
+			this.wsSession.afterSessionInitialized(session);
 			this.webSocketHandler.afterConnectionEstablished(this.wsSession);
 		}
 		catch (Throwable t) {