Add configurability for underlying WebSocket engine

Issue: SPR-10844

Author: rstoyanchev

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

@@ -30,14 +30,16 @@
 /**
  * An abstract base class for implementations of {@link WebSocketSession}.
  *
+ * @param T the type of the native (or delegate) WebSocket session
+ *
  * @author Rossen Stoyanchev
  * @since 4.0
  */
-public abstract class AbstractWebSocketSesssion<T> implements DelegatingWebSocketSession<T> {
+public abstract class AbstractWebSocketSesssion<T> implements WebSocketSession, NativeWebSocketSession {
 
 	protected final Log logger = LogFactory.getLog(getClass());
 
-	private T delegateSession;
+	private T nativeSession;
 
 	private final Map<String, Object> handshakeAttributes;
 
@@ -58,28 +60,35 @@ public Map<String, Object> getHandshakeAttributes() {
 		return this.handshakeAttributes;
 	}
 
-	/**
-	 * @return the WebSocket session to delegate to
-	 */
-	public T getDelegateSession() {
-		return this.delegateSession;
+	@Override
+	public T getNativeSession() {
+		return this.nativeSession;
 	}
 
-
+	@SuppressWarnings("unchecked")
 	@Override
-	public void afterSessionInitialized(T session) {
+	public <R> R getNativeSession(Class<R> requiredType) {
+		if (requiredType != null) {
+			if (requiredType.isInstance(this.nativeSession)) {
+				return (R) this.nativeSession;
+			}
+		}
+		return null;
+	}
+
+	public void initializeNativeSession(T session) {
 		Assert.notNull(session, "session must not be null");
-		this.delegateSession = session;
+		this.nativeSession = session;
 	}
 
-	protected final void checkDelegateSessionInitialized() {
-		Assert.state(this.delegateSession != null, "WebSocket session is not yet initialized");
+	protected final void checkNativeSessionInitialized() {
+		Assert.state(this.nativeSession != null, "WebSocket session is not yet initialized");
 	}
 
 	@Override
 	public final void sendMessage(WebSocketMessage message) throws IOException {
 
-		checkDelegateSessionInitialized();
+		checkNativeSessionInitialized();
 		Assert.isTrue(isOpen(), "Cannot send message after connection closed.");
 
 		if (logger.isTraceEnabled()) {
@@ -109,7 +118,7 @@ public final void close() throws IOException {
 
 	@Override
 	public final void close(CloseStatus status) throws IOException {
-		checkDelegateSessionInitialized();
+		checkNativeSessionInitialized();
 		if (logger.isDebugEnabled()) {
 			logger.debug("Closing " + this);
 		}

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

@@ -54,7 +54,7 @@ public JettyWebSocketHandlerAdapter(WebSocketHandler webSocketHandler, JettyWebS
 	@Override
 	public void onWebSocketConnect(Session session) {
 		try {
-			this.wsSession.afterSessionInitialized(session);
+			this.wsSession.initializeNativeSession(session);
 			this.webSocketHandler.afterConnectionEstablished(this.wsSession);
 		}
 		catch (Throwable t) {

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

@@ -58,22 +58,22 @@ public JettyWebSocketSession(Principal principal, Map<String, Object> handshakeA
 
 	@Override
 	public String getId() {
-		checkDelegateSessionInitialized();
-		return ObjectUtils.getIdentityHexString(getDelegateSession());
+		checkNativeSessionInitialized();
+		return ObjectUtils.getIdentityHexString(getNativeSession());
 	}
 
 	@Override
 	public URI getUri() {
-		checkDelegateSessionInitialized();
-		return getDelegateSession().getUpgradeRequest().getRequestURI();
+		checkNativeSessionInitialized();
+		return getNativeSession().getUpgradeRequest().getRequestURI();
 	}
 
 	@Override
 	public HttpHeaders getHandshakeHeaders() {
-		checkDelegateSessionInitialized();
+		checkNativeSessionInitialized();
 		if (this.headers == null) {
 			this.headers = new HttpHeaders();
-			this.headers.putAll(getDelegateSession().getUpgradeRequest().getHeaders());
+			this.headers.putAll(getNativeSession().getUpgradeRequest().getHeaders());
 			this.headers = HttpHeaders.readOnlyHttpHeaders(headers);
 		}
 		return this.headers;
@@ -86,40 +86,40 @@ public Principal getPrincipal() {
 
 	@Override
 	public InetSocketAddress getLocalAddress() {
-		checkDelegateSessionInitialized();
-		return getDelegateSession().getLocalAddress();
+		checkNativeSessionInitialized();
+		return getNativeSession().getLocalAddress();
 	}
 
 	@Override
 	public InetSocketAddress getRemoteAddress() {
-		checkDelegateSessionInitialized();
-		return getDelegateSession().getRemoteAddress();
+		checkNativeSessionInitialized();
+		return getNativeSession().getRemoteAddress();
 	}
 
 	@Override
 	public String getAcceptedProtocol() {
-		checkDelegateSessionInitialized();
-		return getDelegateSession().getUpgradeResponse().getAcceptedSubProtocol();
+		checkNativeSessionInitialized();
+		return getNativeSession().getUpgradeResponse().getAcceptedSubProtocol();
 	}
 
 	@Override
 	public boolean isOpen() {
-		return ((getDelegateSession() != null) && getDelegateSession().isOpen());
+		return ((getNativeSession() != null) && getNativeSession().isOpen());
 	}
 
 	@Override
 	protected void sendTextMessage(TextMessage message) throws IOException {
-		getDelegateSession().getRemote().sendString(message.getPayload());
+		getNativeSession().getRemote().sendString(message.getPayload());
 	}
 
 	@Override
 	protected void sendBinaryMessage(BinaryMessage message) throws IOException {
-		getDelegateSession().getRemote().sendBytes(message.getPayload());
+		getNativeSession().getRemote().sendBytes(message.getPayload());
 	}
 
 	@Override
 	protected void closeInternal(CloseStatus status) throws IOException {
-		getDelegateSession().close(status.getCode(), status.getReason());
+		getNativeSession().close(status.getCode(), status.getReason());
 	}
 
 }

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

@@ -20,21 +20,26 @@
 
 
 /**
- * A contract for a {@link WebSocketSession} that delegates to another WebSocket session
- * (e.g. a native session).
- *
- * @param T the type of the delegate WebSocket session
+ * A {@link WebSocketSession} that exposes the underlying, native WebSocketSession
+ * through a getter.
  *
  * @author Rossen Stoyanchev
  * @since 4.0
  */
-public interface DelegatingWebSocketSession<T> extends WebSocketSession {
+public interface NativeWebSocketSession extends WebSocketSession {
 
 
 	/**
-	 * Invoked when the delegate WebSocket session has been initialized.
+	 * Return the underlying native WebSocketSession, if available.
+	 * @return the native session or {@code null}
 	 */
-	void afterSessionInitialized(T session);
+	Object getNativeSession();
 
+	/**
+	 * Return the underlying native WebSocketSession, if available.
+	 * @param requiredType the required type of the session
+	 * @return the native session of the required type or {@code null}
+	 */
+	<T> T getNativeSession(Class<T> requiredType);
 
 }

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

@@ -58,7 +58,7 @@ public StandardWebSocketHandlerAdapter(WebSocketHandler handler, StandardWebSock
 	@Override
 	public void onOpen(final javax.websocket.Session session, EndpointConfig config) {
 
-		this.wsSession.afterSessionInitialized(session);
+		this.wsSession.initializeNativeSession(session);
 
 		if (this.handler.supportsPartialMessages()) {
 			session.addMessageHandler(new MessageHandler.Partial<String>() {

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

@@ -68,14 +68,14 @@ public StandardWebSocketSession(HttpHeaders handshakeHeaders, Map<String, Object
 
 	@Override
 	public String getId() {
-		checkDelegateSessionInitialized();
-		return getDelegateSession().getId();
+		checkNativeSessionInitialized();
+		return getNativeSession().getId();
 	}
 
 	@Override
 	public URI getUri() {
-		checkDelegateSessionInitialized();
-		return getDelegateSession().getRequestURI();
+		checkNativeSessionInitialized();
+		return getNativeSession().getRequestURI();
 	}
 
 	@Override
@@ -85,8 +85,8 @@ public HttpHeaders getHandshakeHeaders() {
 
 	@Override
 	public Principal getPrincipal() {
-		checkDelegateSessionInitialized();
-		return getDelegateSession().getUserPrincipal();
+		checkNativeSessionInitialized();
+		return getNativeSession().getUserPrincipal();
 	}
 
 	@Override
@@ -101,29 +101,29 @@ public InetSocketAddress getRemoteAddress() {
 
 	@Override
 	public String getAcceptedProtocol() {
-		checkDelegateSessionInitialized();
-		String protocol = getDelegateSession().getNegotiatedSubprotocol();
+		checkNativeSessionInitialized();
+		String protocol = getNativeSession().getNegotiatedSubprotocol();
 		return StringUtils.isEmpty(protocol)? null : protocol;
 	}
 
 	@Override
 	public boolean isOpen() {
-		return ((getDelegateSession() != null) && getDelegateSession().isOpen());
+		return ((getNativeSession() != null) && getNativeSession().isOpen());
 	}
 
 	@Override
 	protected void sendTextMessage(TextMessage message) throws IOException {
-		getDelegateSession().getBasicRemote().sendText(message.getPayload(), message.isLast());
+		getNativeSession().getBasicRemote().sendText(message.getPayload(), message.isLast());
 	}
 
 	@Override
 	protected void sendBinaryMessage(BinaryMessage message) throws IOException {
-		getDelegateSession().getBasicRemote().sendBinary(message.getPayload(), message.isLast());
+		getNativeSession().getBasicRemote().sendBinary(message.getPayload(), message.isLast());
 	}
 
 	@Override
 	protected void closeInternal(CloseStatus status) throws IOException {
-		getDelegateSession().close(new CloseReason(CloseCodes.getCloseCode(status.getCode()), status.getReason()));
+		getNativeSession().close(new CloseReason(CloseCodes.getCloseCode(status.getCode()), status.getReason()));
 	}
 
 }

spring-websocket/src/main/java/org/springframework/web/socket/client/endpoint/StandardWebSocketClient.java

@@ -52,20 +52,30 @@ public class StandardWebSocketClient extends AbstractWebSocketClient {
 	private final WebSocketContainer webSocketContainer;
 
 
+	/**
+	 * Default constructor that calls {@code ContainerProvider.getWebSocketContainer()} to
+	 * obtain a {@link WebSocketContainer} instance.
+	 */
 	public StandardWebSocketClient() {
 		this.webSocketContainer = ContainerProvider.getWebSocketContainer();
 	}
 
+	/**
+	 * Constructor that accepts a pre-configured {@link WebSocketContainer} instance. If
+	 * using XML configuration see {@link WebSocketContainerFactoryBean}. In Java
+	 * configuration use {@code ContainerProvider.getWebSocketContainer()} to obtain
+	 * a container instance.
+	 */
 	public StandardWebSocketClient(WebSocketContainer webSocketContainer) {
 		Assert.notNull(webSocketContainer, "webSocketContainer must not be null");
 		this.webSocketContainer = webSocketContainer;
 	}
 
 
 	@Override
-	protected WebSocketSession doHandshakeInternal(WebSocketHandler webSocketHandler, HttpHeaders headers,
-			URI uri, List<String> protocols, Map<String, Object> handshakeAttributes)
-					throws WebSocketConnectFailureException {
+	protected WebSocketSession doHandshakeInternal(WebSocketHandler webSocketHandler,
+			HttpHeaders headers, URI uri, List<String> protocols,
+			Map<String, Object> handshakeAttributes) throws WebSocketConnectFailureException {
 
 		int port = getPort(uri);
 		InetSocketAddress localAddress = new InetSocketAddress(getLocalHost(), port);

spring-websocket/src/main/java/org/springframework/web/socket/client/jetty/JettyWebSocketClient.java

@@ -22,6 +22,7 @@
 import java.util.Map;
 
 import org.eclipse.jetty.websocket.client.ClientUpgradeRequest;
+import org.eclipse.jetty.websocket.client.WebSocketClient;
 import org.springframework.context.SmartLifecycle;
 import org.springframework.http.HttpHeaders;
 import org.springframework.web.socket.WebSocketHandler;
@@ -51,12 +52,22 @@ public class JettyWebSocketClient extends AbstractWebSocketClient implements Sma
 	private final Object lifecycleMonitor = new Object();
 
 
+	/**
+	 * Default constructor that creates an instance of
+	 * {@link org.eclipse.jetty.websocket.client.WebSocketClient} with default settings.
+	 */
 	public JettyWebSocketClient() {
 		this.client = new org.eclipse.jetty.websocket.client.WebSocketClient();
 	}
 
+	/**
+	 * Constructor that accepts a pre-configured {@link WebSocketClient}.
+	 */
+	public JettyWebSocketClient(WebSocketClient client) {
+		super();
+		this.client = client;
+	}
 
-	// TODO: configure Jetty WebSocketClient properties
 
 	public void setAutoStartup(boolean autoStartup) {
 		this.autoStartup = autoStartup;

spring-websocket/src/main/java/org/springframework/web/socket/server/support/GlassFishRequestUpgradeStrategy.java

@@ -50,10 +50,15 @@
 import org.springframework.util.StringUtils;
 import org.springframework.web.socket.server.HandshakeFailureException;
 import org.springframework.web.socket.server.endpoint.ServerEndpointRegistration;
+import org.springframework.web.socket.server.endpoint.ServletServerContainerFactoryBean;
 
 
 /**
- * GlassFish support for upgrading a request during a WebSocket handshake.
+ * GlassFish support for upgrading a request during a WebSocket handshake. To modify
+ * properties of the underlying {@link javax.websocket.server.ServerContainer} you can use
+ * {@link ServletServerContainerFactoryBean} in XML configuration or if using Java
+ * configuration, access the container instance through the
+ * "javax.websocket.server.ServerContainer" ServletContext attribute.
  *
  * @author Rossen Stoyanchev
  * @since 4.0

spring-websocket/src/main/java/org/springframework/web/socket/server/support/JettyRequestUpgradeStrategy.java

@@ -24,6 +24,7 @@
 
 import org.eclipse.jetty.websocket.api.UpgradeRequest;
 import org.eclipse.jetty.websocket.api.UpgradeResponse;
+import org.eclipse.jetty.websocket.api.WebSocketPolicy;
 import org.eclipse.jetty.websocket.server.HandshakeRFC6455;
 import org.eclipse.jetty.websocket.server.WebSocketServerFactory;
 import org.eclipse.jetty.websocket.servlet.ServletUpgradeRequest;
@@ -55,10 +56,19 @@ public class JettyRequestUpgradeStrategy implements RequestUpgradeStrategy {
 
 
 	/**
-	 * Default constructor.
+	 * Default constructor that creates {@link WebSocketServerFactory} through its default
+	 * constructor thus using a default {@link WebSocketPolicy}.
 	 */
 	public JettyRequestUpgradeStrategy() {
-		this.factory = new WebSocketServerFactory();
+		this(new WebSocketServerFactory());
+	}
+
+	/**
+	 * A constructor accepting a {@link WebSocketServerFactory}. This may be useful for
+	 * modifying the factory's {@link WebSocketPolicy} via
+	 * {@link WebSocketServerFactory#getPolicy()}.
+	 */
+	public JettyRequestUpgradeStrategy(WebSocketServerFactory factory) {
 		this.factory.setCreator(new WebSocketCreator() {
 			@Override
 			public Object createWebSocket(UpgradeRequest request, UpgradeResponse response) {