From ff2e9aa5bc4b182d74e1e3e654f24e253ad4768b Mon Sep 17 00:00:00 2001
From: Phillip Webb <pwebb@vmware.com>
Date: Tue, 23 Apr 2013 12:08:55 -0700
Subject: [PATCH] Minor tweaks to WebSocketMessage

- payload is now required
- byte data is obtained and stored only once (allowing multiple calls)
- minor formatting polish
---
 .../websocket/BinaryMessage.java              | 38 ++++++++++++-------
 .../websocket/TextMessage.java                |  2 -
 .../websocket/WebSocketMessage.java           |  2 +
 3 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/spring-websocket/src/main/java/org/springframework/websocket/BinaryMessage.java b/spring-websocket/src/main/java/org/springframework/websocket/BinaryMessage.java
index 5f136bc31e..dbb12bcee2 100644
--- a/spring-websocket/src/main/java/org/springframework/websocket/BinaryMessage.java
+++ b/spring-websocket/src/main/java/org/springframework/websocket/BinaryMessage.java
@@ -19,6 +19,8 @@ import java.io.ByteArrayInputStream;
 import java.io.InputStream;
 import java.nio.ByteBuffer;
 
+import org.springframework.util.Assert;
+
 
 /**
  * Represents a binary WebSocket message.
@@ -28,7 +30,7 @@ import java.nio.ByteBuffer;
  */
 public final class BinaryMessage extends WebSocketMessage<ByteBuffer> {
 
-	private final byte[] bytes;
+	private byte[] bytes;
 
 	private final boolean last;
 
@@ -48,8 +50,19 @@ public final class BinaryMessage extends WebSocketMessage<ByteBuffer> {
 	}
 
 	public BinaryMessage(byte[] payload, boolean isLast) {
-		super((payload != null) ? ByteBuffer.wrap(payload) : null);
-		this.bytes = payload;
+		this(payload, 0, (payload == null ? 0 : payload.length), isLast);
+	}
+
+	public BinaryMessage(byte[] payload, int offset, int len) {
+		this(payload, offset, len, true);
+	}
+
+	public BinaryMessage(byte[] payload, int offset, int len, boolean isLast) {
+		super(payload != null ? ByteBuffer.wrap(payload, offset, len) : null);
+		if(payload != null && offset == 0 && len == payload.length) {
+			// FIXME better if a message always needs a payload?
+			this.bytes = payload;
+		}
 		this.last = isLast;
 	}
 
@@ -58,17 +71,16 @@ public final class BinaryMessage extends WebSocketMessage<ByteBuffer> {
 	}
 
 	public byte[] getByteArray() {
-		if (this.bytes != null) {
-			return this.bytes;
-		}
-		else if (getPayload() != null){
-			byte[] result = new byte[getPayload().remaining()];
-			getPayload().get(result);
-			return result;
-		}
-		else {
-			return null;
+		if(this.bytes == null && getPayload() != null) {
+			this.bytes = getRemainingBytes(getPayload());
 		}
+		return this.bytes;
+	}
+
+	private byte[] getRemainingBytes(ByteBuffer payload) {
+		byte[] result = new byte[getPayload().remaining()];
+		getPayload().get(result);
+		return result;
 	}
 
 	public InputStream getInputStream() {
diff --git a/spring-websocket/src/main/java/org/springframework/websocket/TextMessage.java b/spring-websocket/src/main/java/org/springframework/websocket/TextMessage.java
index b254ae7f80..61f3916547 100644
--- a/spring-websocket/src/main/java/org/springframework/websocket/TextMessage.java
+++ b/spring-websocket/src/main/java/org/springframework/websocket/TextMessage.java
@@ -18,7 +18,6 @@ package org.springframework.websocket;
 import java.io.Reader;
 import java.io.StringReader;
 
-
 /**
  * Represents a text WebSocket message.
  *
@@ -27,7 +26,6 @@ import java.io.StringReader;
  */
 public final class TextMessage extends WebSocketMessage<String> {
 
-
 	public TextMessage(String payload) {
 		super(payload);
 	}
diff --git a/spring-websocket/src/main/java/org/springframework/websocket/WebSocketMessage.java b/spring-websocket/src/main/java/org/springframework/websocket/WebSocketMessage.java
index 19b79eda74..6ed55410ad 100644
--- a/spring-websocket/src/main/java/org/springframework/websocket/WebSocketMessage.java
+++ b/spring-websocket/src/main/java/org/springframework/websocket/WebSocketMessage.java
@@ -15,6 +15,7 @@
  */
 package org.springframework.websocket;
 
+import org.springframework.util.Assert;
 import org.springframework.util.ObjectUtils;
 
 
@@ -29,6 +30,7 @@ public abstract class WebSocketMessage<T> {
 
 
 	WebSocketMessage(T payload) {
+		Assert.notNull(payload, "Payload must not be null");
 		this.payload = payload;
 	}