Add partial WebSocketMessage support

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

@@ -15,8 +15,6 @@
  */
 package org.springframework.web.socket;
 
-import java.io.ByteArrayInputStream;
-import java.io.InputStream;
 import java.nio.ByteBuffer;
 
 
@@ -34,11 +32,9 @@ public final class BinaryMessage extends WebSocketMessage<ByteBuffer> {
 	/**
 	 * Create a new {@link BinaryMessage} instance.
 	 * @param payload a non-null payload
-	 * @param isLast if the message is the last of a series of partial messages
 	 */
 	public BinaryMessage(ByteBuffer payload) {
-		super(payload);
-		this.bytes = null;
+		this(payload, true);
 	}
 
 	/**
@@ -46,8 +42,26 @@ public final class BinaryMessage extends WebSocketMessage<ByteBuffer> {
 	 * @param payload a non-null payload
 	 * @param isLast if the message is the last of a series of partial messages
 	 */
+	public BinaryMessage(ByteBuffer payload, boolean isLast) {
+		super(payload, isLast);
+		this.bytes = null;
+	}
+
+	/**
+	 * Create a new {@link BinaryMessage} instance.
+	 * @param payload a non-null payload
+	 */
 	public BinaryMessage(byte[] payload) {
-		this(payload, 0, (payload == null ? 0 : payload.length));
+		this(payload, 0, (payload == null ? 0 : payload.length), true);
+	}
+
+	/**
+	 * Create a new {@link BinaryMessage} instance.
+	 * @param payload a non-null payload
+	 * @param isLast if the message is the last of a series of partial messages
+	 */
+	public BinaryMessage(byte[] payload, boolean isLast) {
+		this(payload, 0, (payload == null ? 0 : payload.length), isLast);
 	}
 
 	/**
@@ -58,8 +72,8 @@ public final class BinaryMessage extends WebSocketMessage<ByteBuffer> {
 	 * @param len the length of the array considered for the payload
 	 * @param isLast if the message is the last of a series of partial messages
 	 */
-	public BinaryMessage(byte[] payload, int offset, int len) {
-		super(payload != null ? ByteBuffer.wrap(payload, offset, len) : null);
+	public BinaryMessage(byte[] payload, int offset, int len, boolean isLast) {
+		super(payload != null ? ByteBuffer.wrap(payload, offset, len) : null, isLast);
 		if(offset == 0 && len == payload.length) {
 			this.bytes = payload;
 		}
@@ -82,18 +96,9 @@ public final class BinaryMessage extends WebSocketMessage<ByteBuffer> {
 		return result;
 	}
 
-	/**
-	 * Returns access to the message payload as an {@link InputStream}.
-	 */
-	public InputStream getInputStream() {
-		byte[] array = getByteArray();
-		return (array != null) ? new ByteArrayInputStream(array) : null;
-	}
-
 	@Override
-	public String toString() {
-		int size = (getPayload() != null) ? getPayload().remaining() : 0;
-		return "WebSocket binary message size=" + size;
+	protected int getPayloadSize() {
+		return (getPayload() != null) ? getPayload().remaining() : 0;
 	}
 
 }

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

@@ -27,12 +27,23 @@ import java.io.StringReader;
  */
 public final class TextMessage extends WebSocketMessage<String> {
 
+
 	/**
 	 * Create a new {@link TextMessage} instance.
 	 * @param payload the payload
+	 * @param isLast whether this the last part of a message received or transmitted in parts
 	 */
 	public TextMessage(CharSequence payload) {
-		super(payload.toString());
+		super(payload.toString(), true);
+	}
+
+	/**
+	 * Create a new {@link TextMessage} instance.
+	 * @param payload the payload
+	 * @param isLast whether this the last part of a message received or transmitted in parts
+	 */
+	public TextMessage(CharSequence payload, boolean isLast) {
+		super(payload.toString(), isLast);
 	}
 
 	/**
@@ -42,4 +53,9 @@ public final class TextMessage extends WebSocketMessage<String> {
 		return new StringReader(getPayload());
 	}
 
+	@Override
+	protected int getPayloadSize() {
+		return getPayload().length();
+	}
+
 }

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

@@ -73,4 +73,9 @@ public interface WebSocketHandler {
 	 */
 	void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception;
 
+	/**
+	 * Whether the WebSocketHandler handles messages in parts.
+	 */
+	boolean supportsPartialMessages();
+
 }

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

@@ -31,14 +31,17 @@ public abstract class WebSocketMessage<T> {
 
 	private final T payload;
 
+	private final boolean last;
+
 
 	/**
 	 * Create a new {@link WebSocketMessage} instance with the given payload.
 	 * @param payload a non-null payload
 	 */
-	WebSocketMessage(T payload) {
+	WebSocketMessage(T payload, boolean isLast) {
 		Assert.notNull(payload, "Payload must not be null");
 		this.payload = payload;
+		this.last = isLast;
 	}
 
 	/**
@@ -48,9 +51,13 @@ public abstract class WebSocketMessage<T> {
 		return this.payload;
 	}
 
-	@Override
-	public String toString() {
-		return getClass().getSimpleName() + " [payload=" + this.payload + "]";
+	/**
+	 * Whether this is the last part of a message, when partial message support on a
+	 * {@link WebSocketHandler} is enabled. If partial message support is not enabled the
+	 * returned value is always {@code true}.
+	 */
+	public boolean isLast() {
+		return this.last;
 	}
 
 	@Override
@@ -70,4 +77,11 @@ public abstract class WebSocketMessage<T> {
 		return ObjectUtils.nullSafeEquals(this.payload, otherMessage.payload);
 	}
 
+	@Override
+	public String toString() {
+		return getClass().getSimpleName() + " [payload length=" + getPayloadSize() + ", last=" + isLast() + "]";
+	}
+
+	protected abstract int getPayloadSize();
+
 }

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

@@ -74,7 +74,7 @@ public class JettyWebSocketListenerAdapter implements WebSocketListener {
 
 	@Override
 	public void onWebSocketBinary(byte[] payload, int offset, int len) {
-		BinaryMessage message = new BinaryMessage(payload, offset, len);
+		BinaryMessage message = new BinaryMessage(payload, offset, len, true);
 		try {
 			this.webSocketHandler.handleMessage(this.wsSession, message);
 		}

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

@@ -61,6 +61,35 @@ public class StandardEndpointAdapter extends Endpoint {
 
 		this.wsSession.initSession(session);
 
+		if (this.handler.supportsPartialMessages()) {
+			session.addMessageHandler(new MessageHandler.Partial<String>() {
+				@Override
+				public void onMessage(String message, boolean isLast) {
+					handleTextMessage(session, message, isLast);
+				}
+			});
+			session.addMessageHandler(new MessageHandler.Partial<ByteBuffer>() {
+				@Override
+				public void onMessage(ByteBuffer message, boolean isLast) {
+					handleBinaryMessage(session, message, isLast);
+				}
+			});
+		}
+		else {
+			session.addMessageHandler(new MessageHandler.Whole<String>() {
+				@Override
+				public void onMessage(String message) {
+					handleTextMessage(session, message, true);
+				}
+			});
+			session.addMessageHandler(new MessageHandler.Whole<ByteBuffer>() {
+				@Override
+				public void onMessage(ByteBuffer message) {
+					handleBinaryMessage(session, message, true);
+				}
+			});
+		}
+
 		try {
 			this.handler.afterConnectionEstablished(this.wsSession);
 		}
@@ -68,23 +97,10 @@ public class StandardEndpointAdapter extends Endpoint {
 			ExceptionWebSocketHandlerDecorator.tryCloseWithError(this.wsSession, t, logger);
 			return;
 		}
-
-		session.addMessageHandler(new MessageHandler.Whole<String>() {
-			@Override
-			public void onMessage(String message) {
-				handleTextMessage(session, message);
-			}
-		});
-		session.addMessageHandler(new MessageHandler.Whole<ByteBuffer>() {
-			@Override
-			public void onMessage(ByteBuffer message) {
-				handleBinaryMessage(session, message);
-			}
-		});
 	}
 
-	private void handleTextMessage(javax.websocket.Session session, String payload) {
-		TextMessage textMessage = new TextMessage(payload);
+	private void handleTextMessage(javax.websocket.Session session, String payload, boolean isLast) {
+		TextMessage textMessage = new TextMessage(payload, isLast);
 		try {
 			this.handler.handleMessage(this.wsSession, textMessage);
 		}
@@ -93,8 +109,8 @@ public class StandardEndpointAdapter extends Endpoint {
 		}
 	}
 
-	private void handleBinaryMessage(javax.websocket.Session session, ByteBuffer payload) {
-		BinaryMessage binaryMessage = new BinaryMessage(payload);
+	private void handleBinaryMessage(javax.websocket.Session session, ByteBuffer payload, boolean isLast) {
+		BinaryMessage binaryMessage = new BinaryMessage(payload, isLast);
 		try {
 			this.handler.handleMessage(this.wsSession, binaryMessage);
 		}

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

@@ -110,12 +110,12 @@ public class StandardWebSocketSessionAdapter extends AbstractWebSocketSesssionAd
 
 	@Override
 	protected void sendTextMessage(TextMessage message) throws IOException {
-		this.session.getBasicRemote().sendText(message.getPayload());
+		this.session.getBasicRemote().sendText(message.getPayload(), message.isLast());
 	}
 
 	@Override
 	protected void sendBinaryMessage(BinaryMessage message) throws IOException {
-		this.session.getBasicRemote().sendBinary(message.getPayload());
+		this.session.getBasicRemote().sendBinary(message.getPayload(), message.isLast());
 	}
 
 	@Override

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

@@ -38,7 +38,7 @@ public class WebSocketHandlerAdapter implements WebSocketHandler {
 	}
 
 	@Override
-	public final void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
+	public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
 		if (message instanceof TextMessage) {
 			handleTextMessage(session, (TextMessage) message);
 		}
@@ -46,7 +46,6 @@ public class WebSocketHandlerAdapter implements WebSocketHandler {
 			handleBinaryMessage(session, (BinaryMessage) message);
 		}
 		else {
-			// should not happen
 			throw new IllegalStateException("Unexpected WebSocket message type: " + message);
 		}
 	}
@@ -65,4 +64,9 @@ public class WebSocketHandlerAdapter implements WebSocketHandler {
 	public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
 	}
 
+	@Override
+	public boolean supportsPartialMessages() {
+		return false;
+	}
+
 }

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

@@ -50,8 +50,8 @@ public class LoggingWebSocketHandlerDecorator extends WebSocketHandlerDecorator
 
 	@Override
 	public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
-		if (logger.isTraceEnabled()) {
-			logger.trace("Received " + message + ", " + session);
+		if (logger.isDebugEnabled()) {
+			logger.debug("Received " + message + ", " + session);
 		}
 		super.handleMessage(session, message);
 	}

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

@@ -57,9 +57,16 @@ public class PerConnectionWebSocketHandler implements WebSocketHandler, BeanFact
 	private final Map<WebSocketSession, WebSocketHandler> handlers =
 			new ConcurrentHashMap<WebSocketSession, WebSocketHandler>();
 
+	private final boolean supportsPartialMessages;
+
 
 	public PerConnectionWebSocketHandler(Class<? extends WebSocketHandler> handlerType) {
+		this(handlerType, false);
+	}
+
+	public PerConnectionWebSocketHandler(Class<? extends WebSocketHandler> handlerType, boolean supportsPartialMessages) {
 		this.provider = new BeanCreatingHandlerProvider<WebSocketHandler>(handlerType);
+		this.supportsPartialMessages = supportsPartialMessages;
 	}
 
 	@Override
@@ -112,6 +119,11 @@ public class PerConnectionWebSocketHandler implements WebSocketHandler, BeanFact
 		}
 	}
 
+	@Override
+	public boolean supportsPartialMessages() {
+		return this.supportsPartialMessages;
+	}
+
 	@Override
 	public String toString() {
 		return "PerConnectionWebSocketHandlerProxy [handlerType=" + this.provider.getHandlerType() + "]";

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

@@ -62,6 +62,12 @@ public class WebSocketHandlerDecorator implements WebSocketHandler {
 		this.delegate.afterConnectionClosed(session, closeStatus);
 	}
 
+	@Override
+	public boolean supportsPartialMessages() {
+		return this.delegate.supportsPartialMessages();
+	}
+
+
 	@Override
 	public String toString() {
 		return getClass().getSimpleName() + " [delegate=" + this.delegate + "]";