Add WebSocketHandlerDecorator

spring-websocket/src/main/java/org/springframework/sockjs/AbstractSockJsSession.java

@@ -26,7 +26,6 @@ import org.springframework.websocket.CloseStatus;
 import org.springframework.websocket.TextMessage;
 import org.springframework.websocket.WebSocketHandler;
 import org.springframework.websocket.WebSocketSession;
-import org.springframework.websocket.adapter.WebSocketHandlerInvoker;
 
 
 /**
@@ -42,7 +41,7 @@ public abstract class AbstractSockJsSession implements WebSocketSession {
 
 	private final String sessionId;
 
-	private WebSocketHandlerInvoker handler;
+	private WebSocketHandler handler;
 
 	private State state = State.NEW;
 
@@ -59,7 +58,7 @@ public abstract class AbstractSockJsSession implements WebSocketSession {
 		Assert.notNull(sessionId, "sessionId is required");
 		Assert.notNull(webSocketHandler, "webSocketHandler is required");
 		this.sessionId = sessionId;
-		this.handler = new WebSocketHandlerInvoker(webSocketHandler).setLogger(logger);
+		this.handler = webSocketHandler;
 	}
 
 	public String getId() {
@@ -129,7 +128,13 @@ public abstract class AbstractSockJsSession implements WebSocketSession {
 	 */
 	protected void tryCloseWithSockJsTransportError(Throwable ex, CloseStatus closeStatus) {
 		delegateError(ex);
-		this.handler.tryCloseWithError(this, ex, closeStatus);
+		try {
+			logger.error("Closing due to transport error for " + this, ex);
+			close(closeStatus);
+		}
+		catch (Throwable t) {
+			// ignore
+		}
 	}
 
 	public void delegateMessages(String[] messages) {

spring-websocket/src/main/java/org/springframework/sockjs/server/support/SockJsHttpRequestHandler.java

@@ -32,6 +32,8 @@ import org.springframework.web.HttpRequestHandler;
 import org.springframework.web.util.NestedServletException;
 import org.springframework.web.util.UrlPathHelper;
 import org.springframework.websocket.WebSocketHandler;
+import org.springframework.websocket.support.ExceptionWebSocketHandlerDecorator;
+import org.springframework.websocket.support.LoggingWebSocketHandlerDecorator;
 
 /**
  * @author Rossen Stoyanchev
@@ -63,7 +65,18 @@ public class SockJsHttpRequestHandler implements HttpRequestHandler {
 
 		this.prefix = prefix;
 		this.sockJsService = sockJsService;
-		this.webSocketHandler = webSocketHandler;
+		this.webSocketHandler = decorateWebSocketHandler(webSocketHandler);
+	}
+
+	/**
+	 * Decorate the WebSocketHandler provided to the class constructor.
+	 * <p>
+	 * By default {@link ExceptionWebSocketHandlerDecorator} and
+	 * {@link LoggingWebSocketHandlerDecorator} are applied are added.
+	 */
+	protected WebSocketHandler decorateWebSocketHandler(WebSocketHandler handler) {
+		handler = new ExceptionWebSocketHandlerDecorator(handler);
+		return new LoggingWebSocketHandlerDecorator(handler);
 	}
 
 	public String getPrefix() {

spring-websocket/src/main/java/org/springframework/websocket/WebSocketHandler.java

@@ -48,4 +48,9 @@ public interface WebSocketHandler {
 	 */
 	void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus);
 
+	/**
+	 * Whether this WebSocketHandler wishes to receive messages broken up in parts.
+	 */
+	boolean isStreaming();
+
 }

spring-websocket/src/main/java/org/springframework/websocket/adapter/JettyWebSocketListenerAdapter.java

@@ -16,8 +16,6 @@
 
 package org.springframework.websocket.adapter;
 
-import org.apache.commons.logging.Log;
-import org.apache.commons.logging.LogFactory;
 import org.eclipse.jetty.websocket.api.Session;
 import org.eclipse.jetty.websocket.api.WebSocketListener;
 import org.springframework.util.Assert;
@@ -35,8 +33,6 @@ import org.springframework.websocket.WebSocketSession;
  */
 public class JettyWebSocketListenerAdapter implements WebSocketListener {
 
-	private static Log logger = LogFactory.getLog(JettyWebSocketListenerAdapter.class);
-
 	private final WebSocketHandler webSocketHandler;
 
 	private WebSocketSession wsSession;
@@ -44,7 +40,7 @@ public class JettyWebSocketListenerAdapter implements WebSocketListener {
 
 	public JettyWebSocketListenerAdapter(WebSocketHandler webSocketHandler) {
 		Assert.notNull(webSocketHandler, "webSocketHandler is required");
-		this.webSocketHandler = new WebSocketHandlerInvoker(webSocketHandler).setLogger(logger);
+		this.webSocketHandler = webSocketHandler;
 	}
 
 
@@ -76,4 +72,5 @@ public class JettyWebSocketListenerAdapter implements WebSocketListener {
 	public void onWebSocketError(Throwable cause) {
 		this.webSocketHandler.handleTransportError(this.wsSession, cause);
 	}
+
 }

spring-websocket/src/main/java/org/springframework/websocket/adapter/StandardEndpointAdapter.java

@@ -23,12 +23,9 @@ import javax.websocket.Endpoint;
 import javax.websocket.EndpointConfig;
 import javax.websocket.MessageHandler;
 
-import org.apache.commons.logging.Log;
-import org.apache.commons.logging.LogFactory;
 import org.springframework.util.Assert;
 import org.springframework.websocket.BinaryMessage;
 import org.springframework.websocket.CloseStatus;
-import org.springframework.websocket.PartialMessageHandler;
 import org.springframework.websocket.TextMessage;
 import org.springframework.websocket.WebSocketHandler;
 import org.springframework.websocket.WebSocketSession;
@@ -42,26 +39,23 @@ import org.springframework.websocket.WebSocketSession;
  */
 public class StandardEndpointAdapter extends Endpoint {
 
-	private static Log logger = LogFactory.getLog(StandardEndpointAdapter.class);
+	private final WebSocketHandler handler;
-
-	private final WebSocketHandlerInvoker handler;
-
-	private final Class<?> handlerClass;
 
 	private WebSocketSession wsSession;
 
 
-
 	public StandardEndpointAdapter(WebSocketHandler webSocketHandler) {
 		Assert.notNull(webSocketHandler, "webSocketHandler is required");
-		this.handler = new WebSocketHandlerInvoker(webSocketHandler).setLogger(logger);
+		this.handler = webSocketHandler;
-		this.handlerClass= webSocketHandler.getClass();
 	}
 
 
 	@Override
 	public void onOpen(final javax.websocket.Session session, EndpointConfig config) {
 
+		this.wsSession = new StandardWebSocketSessionAdapter(session);
+		this.handler.afterConnectionEstablished(this.wsSession);
+
 		session.addMessageHandler(new MessageHandler.Whole<String>() {
 			@Override
 			public void onMessage(String message) {
@@ -69,9 +63,7 @@ public class StandardEndpointAdapter extends Endpoint {
 			}
 		});
 
-		// TODO: per-connection proxy
+		if (!this.handler.isStreaming()) {
-
-		if (!PartialMessageHandler.class.isAssignableFrom(this.handlerClass)) {
 			session.addMessageHandler(new MessageHandler.Whole<ByteBuffer>() {
 				@Override
 				public void onMessage(ByteBuffer message) {
@@ -88,8 +80,6 @@ public class StandardEndpointAdapter extends Endpoint {
 			});
 		}
 
-		this.wsSession = new StandardWebSocketSessionAdapter(session);
-		this.handler.afterConnectionEstablished(this.wsSession);
 	}
 
 	private void handleTextMessage(javax.websocket.Session session, String payload) {

spring-websocket/src/main/java/org/springframework/websocket/adapter/WebSocketHandlerAdapter.java

@@ -68,4 +68,9 @@ public class WebSocketHandlerAdapter implements WebSocketHandler {
 	public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
 	}
 
+	@Override
+	public boolean isStreaming() {
+		return false;
+	}
+
 }

spring-websocket/src/main/java/org/springframework/websocket/adapter/WebSocketHandlerInvoker.java

@@ -1,126 +0,0 @@
-/*
- * Copyright 2002-2013 the original author or authors.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.springframework.websocket.adapter;
-
-import java.util.concurrent.atomic.AtomicInteger;
-
-import org.apache.commons.logging.Log;
-import org.apache.commons.logging.LogFactory;
-import org.springframework.util.Assert;
-import org.springframework.websocket.CloseStatus;
-import org.springframework.websocket.WebSocketHandler;
-import org.springframework.websocket.WebSocketMessage;
-import org.springframework.websocket.WebSocketSession;
-
-/**
- * A class for managing and delegating to a {@link WebSocketHandler} instance, ensuring
- * the handler is initialized and destroyed, that any unhandled exceptions from handler
- * are caught (and handled by closing the session), as well as adding logging.
- *
- * @author Rossen Stoyanchev
- * @author Phillip Webb
- * @since 4.0
- */
-public class WebSocketHandlerInvoker implements WebSocketHandler {
-
-	private Log logger = LogFactory.getLog(WebSocketHandlerInvoker.class);
-
-	private final WebSocketHandler handler;
-
-	private final AtomicInteger sessionCount = new AtomicInteger(0);
-
-
-	public WebSocketHandlerInvoker(WebSocketHandler webSocketHandler) {
-		Assert.notNull(webSocketHandler, "webSocketHandler is required");
-		this.handler = webSocketHandler;
-	}
-
-	public WebSocketHandlerInvoker setLogger(Log logger) {
-		this.logger = logger;
-		return this;
-	}
-
-	@Override
-	public void afterConnectionEstablished(WebSocketSession session) {
-		if (logger.isDebugEnabled()) {
-			logger.debug("Connection established, "	+ session + ", uri=" + session.getURI());
-		}
-		try {
-			Assert.isTrue(this.sessionCount.compareAndSet(0, 1), "Unexpected new session");
-			this.handler.afterConnectionEstablished(session);
-		}
-		catch (Throwable ex) {
-			tryCloseWithError(session, ex);
-		}
-	}
-
-	public void tryCloseWithError(WebSocketSession session, Throwable ex) {
-		tryCloseWithError(session, ex, null);
-	}
-
-	public void tryCloseWithError(WebSocketSession session, Throwable exeption, CloseStatus status) {
-		logger.error("Closing due to exception for " + session, exeption);
-		if (session.isOpen()) {
-			try {
-				session.close((status != null) ? status : CloseStatus.SERVER_ERROR);
-			}
-			catch (Throwable t) {
-				// ignore
-			}
-		}
-	}
-
-	@Override
-	public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) {
-		if (logger.isTraceEnabled()) {
-			logger.trace("Received " + message + ", " + session);
-		}
-		try {
-			this.handler.handleMessage(session, message);
-		}
-		catch (Throwable ex) {
-			tryCloseWithError(session,ex);
-		}
-	}
-
-	@Override
-	public void handleTransportError(WebSocketSession session, Throwable exception) {
-		if (logger.isDebugEnabled()) {
-			logger.debug("Transport error for " + session, exception);
-		}
-		try {
-			this.handler.handleTransportError(session, exception);
-		}
-		catch (Throwable ex) {
-			tryCloseWithError(session, ex);
-		}
-	}
-
-	@Override
-	public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) {
-		if (logger.isDebugEnabled()) {
-			logger.debug("Connection closed for " + session + ", " + closeStatus);
-		}
-		try {
-			this.handler.afterConnectionClosed(session, closeStatus);
-		}
-		catch (Throwable ex) {
-			logger.error("Unhandled error for " + this, ex);
-		}
-	}
-
-}

spring-websocket/src/main/java/org/springframework/websocket/client/WebSocketConnectionManager.java

@@ -23,6 +23,8 @@ import org.springframework.http.HttpHeaders;
 import org.springframework.util.CollectionUtils;
 import org.springframework.websocket.WebSocketHandler;
 import org.springframework.websocket.WebSocketSession;
+import org.springframework.websocket.support.ExceptionWebSocketHandlerDecorator;
+import org.springframework.websocket.support.LoggingWebSocketHandlerDecorator;
 
 /**
  * @author Rossen Stoyanchev
@@ -44,7 +46,18 @@ public class WebSocketConnectionManager extends AbstractWebSocketConnectionManag
 
 		super(uriTemplate, uriVariables);
 		this.client = webSocketClient;
-		this.webSocketHandler = webSocketHandler;
+		this.webSocketHandler = decorateWebSocketHandler(webSocketHandler);
+	}
+
+	/**
+	 * Decorate the WebSocketHandler provided to the class constructor.
+	 * <p>
+	 * By default {@link ExceptionWebSocketHandlerDecorator} and
+	 * {@link LoggingWebSocketHandlerDecorator} are applied are added.
+	 */
+	protected WebSocketHandler decorateWebSocketHandler(WebSocketHandler handler) {
+		handler = new ExceptionWebSocketHandlerDecorator(handler);
+		return new LoggingWebSocketHandlerDecorator(handler);
 	}
 
 	public void setSubProtocols(List<String> subProtocols) {

spring-websocket/src/main/java/org/springframework/websocket/server/support/WebSocketHttpRequestHandler.java

@@ -32,6 +32,8 @@ import org.springframework.web.util.NestedServletException;
 import org.springframework.websocket.WebSocketHandler;
 import org.springframework.websocket.server.DefaultHandshakeHandler;
 import org.springframework.websocket.server.HandshakeHandler;
+import org.springframework.websocket.support.ExceptionWebSocketHandlerDecorator;
+import org.springframework.websocket.support.LoggingWebSocketHandlerDecorator;
 
 /**
  * An {@link HttpRequestHandler} that wraps the invocation of a {@link HandshakeHandler}.
@@ -53,10 +55,21 @@ public class WebSocketHttpRequestHandler implements HttpRequestHandler {
 	public WebSocketHttpRequestHandler(	WebSocketHandler webSocketHandler, HandshakeHandler handshakeHandler) {
 		Assert.notNull(webSocketHandler, "webSocketHandler is required");
 		Assert.notNull(handshakeHandler, "handshakeHandler is required");
-		this.webSocketHandler = webSocketHandler;
+		this.webSocketHandler = decorateWebSocketHandler(webSocketHandler);
 		this.handshakeHandler = new DefaultHandshakeHandler();
 	}
 
+	/**
+	 * Decorate the WebSocketHandler provided to the class constructor.
+	 * <p>
+	 * By default {@link ExceptionWebSocketHandlerDecorator} and
+	 * {@link LoggingWebSocketHandlerDecorator} are applied are added.
+	 */
+	protected WebSocketHandler decorateWebSocketHandler(WebSocketHandler handler) {
+		handler = new ExceptionWebSocketHandlerDecorator(handler);
+		return new LoggingWebSocketHandlerDecorator(handler);
+	}
+
 	@Override
 	public void handleRequest(HttpServletRequest request, HttpServletResponse response)
 			throws ServletException, IOException {

spring-websocket/src/main/java/org/springframework/websocket/support/ExceptionWebSocketHandlerDecorator.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2002-2013 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.websocket.support;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.springframework.websocket.CloseStatus;
+import org.springframework.websocket.WebSocketHandler;
+import org.springframework.websocket.WebSocketMessage;
+import org.springframework.websocket.WebSocketSession;
+
+
+/**
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public class ExceptionWebSocketHandlerDecorator extends WebSocketHandlerDecorator {
+
+	private Log logger = LogFactory.getLog(ExceptionWebSocketHandlerDecorator.class);
+
+
+	public ExceptionWebSocketHandlerDecorator(WebSocketHandler delegate) {
+		super(delegate);
+	}
+
+	@Override
+	public void afterConnectionEstablished(WebSocketSession session) {
+		try {
+			getDelegate().afterConnectionEstablished(session);
+		}
+		catch (Throwable ex) {
+			tryCloseWithError(session, ex);
+		}
+	}
+
+	private void tryCloseWithError(WebSocketSession session, Throwable exception) {
+		logger.error("Closing due to exception for " + session, exception);
+		if (session.isOpen()) {
+			try {
+				session.close(CloseStatus.SERVER_ERROR);
+			}
+			catch (Throwable t) {
+				// ignore
+			}
+		}
+	}
+
+	@Override
+	public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) {
+		try {
+			getDelegate().handleMessage(session, message);
+		}
+		catch (Throwable ex) {
+			tryCloseWithError(session,ex);
+		}
+	}
+
+	@Override
+	public void handleTransportError(WebSocketSession session, Throwable exception) {
+		try {
+			getDelegate().handleTransportError(session, exception);
+		}
+		catch (Throwable ex) {
+			tryCloseWithError(session, ex);
+		}
+	}
+
+	@Override
+	public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) {
+		try {
+			getDelegate().afterConnectionClosed(session, closeStatus);
+		}
+		catch (Throwable ex) {
+			logger.error("Unhandled error for " + this, ex);
+		}
+	}
+
+}

spring-websocket/src/main/java/org/springframework/websocket/support/LoggingWebSocketHandlerDecorator.java

@@ -0,0 +1,73 @@
+/*
+ * Copyright 2002-2013 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.websocket.support;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.springframework.websocket.CloseStatus;
+import org.springframework.websocket.WebSocketHandler;
+import org.springframework.websocket.WebSocketMessage;
+import org.springframework.websocket.WebSocketSession;
+
+
+/**
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public class LoggingWebSocketHandlerDecorator extends WebSocketHandlerDecorator {
+
+	private Log logger = LogFactory.getLog(LoggingWebSocketHandlerDecorator.class);
+
+
+	public LoggingWebSocketHandlerDecorator(WebSocketHandler delegate) {
+		super(delegate);
+	}
+
+
+	@Override
+	public void afterConnectionEstablished(WebSocketSession session) {
+		if (logger.isDebugEnabled()) {
+			logger.debug("Connection established, "	+ session + ", uri=" + session.getURI());
+		}
+		super.afterConnectionEstablished(session);
+	}
+
+	@Override
+	public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) {
+		if (logger.isTraceEnabled()) {
+			logger.trace("Received " + message + ", " + session);
+		}
+		super.handleMessage(session, message);
+	}
+
+	@Override
+	public void handleTransportError(WebSocketSession session, Throwable exception) {
+		if (logger.isDebugEnabled()) {
+			logger.debug("Transport error for " + session, exception);
+		}
+		super.handleTransportError(session, exception);
+	}
+
+	@Override
+	public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) {
+		if (logger.isDebugEnabled()) {
+			logger.debug("Connection closed for " + session + ", " + closeStatus);
+		}
+		super.afterConnectionClosed(session, closeStatus);
+	}
+
+}

spring-websocket/src/main/java/org/springframework/websocket/support/PerConnectionWebSocketHandlerProxy.java

@@ -57,11 +57,23 @@ public class PerConnectionWebSocketHandlerProxy implements WebSocketHandler, Bea
 	private Map<WebSocketSession, WebSocketHandler> handlers =
 			new ConcurrentHashMap<WebSocketSession, WebSocketHandler>();
 
+	private boolean streaming;
+
 
 	public PerConnectionWebSocketHandlerProxy(Class<? extends WebSocketHandler> handlerType) {
-		this.provider = new BeanCreatingHandlerProvider<WebSocketHandler>(handlerType);
+		this(handlerType, false);
 	}
 
+	public PerConnectionWebSocketHandlerProxy(Class<? extends WebSocketHandler> handlerType, boolean isStreaming) {
+		this.provider = new BeanCreatingHandlerProvider<WebSocketHandler>(handlerType);
+		this.streaming = isStreaming;
+	}
+
+
+	@Override
+	public boolean isStreaming() {
+		return this.streaming;
+	}
 
 	@Override
 	public void setBeanFactory(BeanFactory beanFactory) throws BeansException {

spring-websocket/src/main/java/org/springframework/websocket/support/WebSocketHandlerDecorator.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2002-2013 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.websocket.support;
+
+import org.springframework.util.Assert;
+import org.springframework.websocket.CloseStatus;
+import org.springframework.websocket.WebSocketHandler;
+import org.springframework.websocket.WebSocketMessage;
+import org.springframework.websocket.WebSocketSession;
+
+
+/**
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public class WebSocketHandlerDecorator implements WebSocketHandler {
+
+	private final WebSocketHandler delegate;
+
+
+	public WebSocketHandlerDecorator(WebSocketHandler delegate) {
+		Assert.notNull(delegate, "delegate is required");
+		this.delegate = delegate;
+	}
+
+
+	protected WebSocketHandler getDelegate() {
+		return this.delegate;
+	}
+
+	@Override
+	public void afterConnectionEstablished(WebSocketSession session) {
+		this.delegate.afterConnectionEstablished(session);
+	}
+
+	@Override
+	public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) {
+		this.delegate.handleMessage(session, message);
+	}
+
+	@Override
+	public void handleTransportError(WebSocketSession session, Throwable exception) {
+		this.delegate.handleTransportError(session, exception);
+	}
+
+	@Override
+	public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) {
+		this.delegate.afterConnectionClosed(session, closeStatus);
+	}
+
+	@Override
+	public boolean isStreaming() {
+		return this.delegate.isStreaming();
+	}
+
+}