Add BufferingStompDecoder

Before this change the StompDecoder decoded and returned only the first
Message in the ByteBuffer passed to it. So to obtain all messages from
the buffer, one had to loop passing the same buffer in until no more
complete STOMP frames could be decoded.

This chage modifies StompDecoder to return List<Message> after
exhaustively decoding all available STOMP frames from the input buffer.
Also an overloaded decode method allows passing in Map that will be
populated with any headers successfully parsed, which is useful for
"peeking" at the "content-length" header.

This change also adds a BufferingStompDecoder sub-class which buffers
any content left in the input buffer after parsing one or more STOMP
frames. This sub-class can also deal with fragmented messages,
re-assembling them and parsing as a whole message.

Issue: SPR-11527

Author: rstoyanchev

spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/BufferingStompDecoder.java

@@ -0,0 +1,134 @@
+/*
+ * Copyright 2002-2014 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.messaging.simp.stomp;
+
+
+import org.springframework.messaging.Message;
+import org.springframework.util.Assert;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Queue;
+import java.util.concurrent.LinkedBlockingQueue;
+
+
+/**
+ * A an extension of {@link org.springframework.messaging.simp.stomp.StompDecoder}
+ * that chunks any bytes remaining after a single full STOMP frame has been read.
+ * The remaining bytes may contain more STOMP frames or an incomplete STOMP frame.
+ *
+ * <p>Similarly if there is not enough content for a full STOMP frame, the content
+ * is buffered until more input is received. That means the
+ * {@link #decode(java.nio.ByteBuffer)} effectively never returns {@code null} as
+ * the parent class does.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0.3
+ */
+public class BufferingStompDecoder extends StompDecoder {
+
+	private final int bufferSizeLimit;
+
+	private final Queue<ByteBuffer> chunks = new LinkedBlockingQueue<ByteBuffer>();
+
+	private volatile Integer expectedContentLength;
+
+
+	public BufferingStompDecoder(int bufferSizeLimit) {
+		Assert.isTrue(bufferSizeLimit > 0, "Buffer size must be greater than 0");
+		this.bufferSizeLimit = bufferSizeLimit;
+	}
+
+
+	public int getBufferSizeLimit() {
+		return this.bufferSizeLimit;
+	}
+
+	public int getBufferSize() {
+		int size = 0;
+		for (ByteBuffer buffer : this.chunks) {
+			size = size + buffer.remaining();
+		}
+		return size;
+	}
+
+	public Integer getExpectedContentLength() {
+		return this.expectedContentLength;
+	}
+
+
+	@Override
+	public List<Message<byte[]>> decode(ByteBuffer newData) {
+
+		this.chunks.add(newData);
+
+		checkBufferLimits();
+
+		if (getExpectedContentLength() != null && getBufferSize() < this.expectedContentLength) {
+			return Collections.<Message<byte[]>>emptyList();
+		}
+
+		ByteBuffer buffer = assembleChunksAndReset();
+
+		MultiValueMap<String, String> headers = new LinkedMultiValueMap<String, String>();
+		List<Message<byte[]>> messages = decode(buffer, headers);
+
+		if (buffer.hasRemaining()) {
+			this.chunks.add(buffer);
+			this.expectedContentLength = getContentLength(headers);
+		}
+
+		return messages;
+	}
+
+	private void checkBufferLimits() {
+		if (getExpectedContentLength() != null) {
+			if (getExpectedContentLength() > getBufferSizeLimit()) {
+				throw new StompConversionException(
+						"The 'content-length' header " + getExpectedContentLength() +
+								"  exceeds the configured message buffer size limit " + getBufferSizeLimit());
+			}
+		}
+		if (getBufferSize() > getBufferSizeLimit()) {
+			throw new StompConversionException("The configured stomp frame buffer size limit of " +
+					getBufferSizeLimit() + " bytes has been exceeded");
+
+		}
+	}
+
+	private ByteBuffer assembleChunksAndReset() {
+		ByteBuffer result;
+		if (this.chunks.size() == 1) {
+			result = this.chunks.remove();
+		}
+		else {
+			result = ByteBuffer.allocate(getBufferSize());
+			for (ByteBuffer partial : this.chunks) {
+				result.put(partial);
+			}
+			result.flip();
+		}
+		this.chunks.clear();
+		this.expectedContentLength = null;
+		return result;
+	}
+
+}

spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompCodec.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2013 the original author or authors.
+ * Copyright 2002-2014 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.
@@ -23,10 +23,13 @@
 import reactor.io.Buffer;
 import reactor.tcp.encoding.Codec;
 
+import java.util.List;
+
 /**
- * A Reactor TCP {@link Codec} for sending and receiving STOMP messages
+ * A Reactor TCP {@link Codec} for sending and receiving STOMP messages.
  *
  * @author Andy Wilkinson
+ * @author Rossen Stoyanchev
  * @since 4.0
  */
 public class StompCodec implements Codec<Buffer, Message<byte[]>, Message<byte[]>> {
@@ -49,14 +52,8 @@ public Function<Buffer, Message<byte[]>> decoder(final Consumer<Message<byte[]>>
 
 			@Override
 			public Message<byte[]> apply(Buffer buffer) {
-				while (buffer.remaining() > 0) {
-					Message<byte[]> message = DECODER.decode(buffer.byteBuffer());
-					if (message != null) {
-						next.accept(message);
-					}
-					else {
-						break;
-					}
+				for (Message<byte[]> message : DECODER.decode(buffer.byteBuffer())) {
+					next.accept(message);
 				}
 				return null;
 			}

spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompDecoder.java

@@ -19,6 +19,8 @@
 import java.io.ByteArrayOutputStream;
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
+import java.util.ArrayList;
+import java.util.List;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -30,9 +32,10 @@
 import org.springframework.util.MultiValueMap;
 
 /**
- * Decodes STOMP frames from a {@link ByteBuffer}. If the buffer does not contain
- * enough data to form a complete STOMP frame, the buffer is reset and the value
- * returned is {@code null} indicating that no message could be read.
+ * Decodes one or more STOMP frames from a {@link ByteBuffer}. If the buffer
+ * contains any additional (incomplete) data, or perhaps not enough data to
+ * form even one Message, the the buffer is reset and the value returned is
+ * an empty list indicating that no more message can be read.
  *
  * @author Andy Wilkinson
  * @author Rossen Stoyanchev
@@ -47,21 +50,66 @@ public class StompDecoder {
 	private final Log logger = LogFactory.getLog(StompDecoder.class);
 
 
+
+	/**
+	 * Decodes one or more STOMP frames from the given {@code buffer} into a
+	 * list of {@link Message}s.
+	 *
+	 * <p>If the given ByteBuffer contains partial STOMP frame content, or additional
+	 * content with a partial STOMP frame, the buffer is reset and {@code null} is
+	 * returned.
+	 *
+	 * @param buffer The buffer to decode the STOMP frame from
+	 *
+	 * @return the decoded messages or an empty list
+	 */
+	public List<Message<byte[]>> decode(ByteBuffer buffer) {
+		return decode(buffer, new LinkedMultiValueMap<String, String>());
+	}
+
+	/**
+	 * Decodes one or more STOMP frames from the given {@code buffer} into a
+	 * list of {@link Message}s.
+	 *
+	 * <p>If the given ByteBuffer contains partial STOMP frame content, or additional
+	 * content with a partial STOMP frame, the buffer is reset and {@code null} is
+	 * returned.
+	 *
+	 * @param buffer The buffer to decode the STOMP frame from
+	 * @param headers an empty map that will be filled with the successfully parsed
+	 * 	headers of the last decoded message, or the last attempt at decoding an
+	 *  (incomplete) STOMP frame. This can be useful for detecting 'content-length'.
+	 *
+	 * @return the decoded messages or an empty list
+	 */
+	public List<Message<byte[]>> decode(ByteBuffer buffer, MultiValueMap<String, String> headers) {
+		List<Message<byte[]>> messages = new ArrayList<Message<byte[]>>();
+		while (buffer.hasRemaining()) {
+			headers.clear();
+			Message<byte[]> m = decodeMessage(buffer, headers);
+			if (m != null) {
+				messages.add(m);
+			}
+			else {
+				break;
+			}
+		}
+		return messages;
+	}
+
 	/**
-	 * Decodes a STOMP frame in the given {@code buffer} into a {@link Message}.
-	 * If the given ByteBuffer contains partial STOMP frame content, the method
-	 * resets the buffer and returns {@code null}.
-	 * @param buffer the buffer to decode the frame from
-	 * @return the decoded message or {@code null}
+	 * Decode a single STOMP frame from the given {@code buffer} into a {@link Message}.
 	 */
-	public Message<byte[]> decode(ByteBuffer buffer) {
+	private Message<byte[]> decodeMessage(ByteBuffer buffer, MultiValueMap<String, String> headers) {
+
 		Message<byte[]> decodedMessage = null;
 		skipLeadingEol(buffer);
 		buffer.mark();
 
 		String command = readCommand(buffer);
 		if (command.length() > 0) {
-			MultiValueMap<String, String> headers = readHeaders(buffer);
+
+			readHeaders(buffer, headers);
 			byte[] payload = readPayload(buffer, headers);
 
 			if (payload != null) {
@@ -78,7 +126,7 @@ public Message<byte[]> decode(ByteBuffer buffer) {
 			}
 			else {
 				if (logger.isTraceEnabled()) {
-					logger.trace("Received incomplete frame. Resetting buffer");
+					logger.trace("Received incomplete frame. Resetting buffer.");
 				}
 				buffer.reset();
 			}
@@ -93,27 +141,31 @@ public Message<byte[]> decode(ByteBuffer buffer) {
 		return decodedMessage;
 	}
 
-	private void skipLeadingEol(ByteBuffer buffer) {
+
+	/**
+	 * Skip one ore more EOL characters at the start of the given ByteBuffer.
+	 * Those are STOMP heartbeat frames.
+	 */
+	protected void skipLeadingEol(ByteBuffer buffer) {
 		while (true) {
-			if (!isEol(buffer)) {
+			if (!tryConsumeEndOfLine(buffer)) {
 				break;
 			}
 		}
 	}
 
 	private String readCommand(ByteBuffer buffer) {
 		ByteArrayOutputStream command = new ByteArrayOutputStream(256);
-		while (buffer.remaining() > 0 && !isEol(buffer)) {
+		while (buffer.remaining() > 0 && !tryConsumeEndOfLine(buffer)) {
 			command.write(buffer.get());
 		}
 		return new String(command.toByteArray(), UTF8_CHARSET);
 	}
 
-	private MultiValueMap<String, String> readHeaders(ByteBuffer buffer) {
-		MultiValueMap<String, String> headers = new LinkedMultiValueMap<String, String>();
+	private void readHeaders(ByteBuffer buffer, MultiValueMap<String, String> headers) {
 		while (true) {
 			ByteArrayOutputStream headerStream = new ByteArrayOutputStream(256);
-			while (buffer.remaining() > 0 && !isEol(buffer)) {
+			while (buffer.remaining() > 0 && !tryConsumeEndOfLine(buffer)) {
 				headerStream.write(buffer.get());
 			}
 			if (headerStream.size() > 0) {
@@ -135,7 +187,6 @@ private MultiValueMap<String, String> readHeaders(ByteBuffer buffer) {
 				break;
 			}
 		}
-		return headers;
 	}
 
 	private String unescape(String input) {
@@ -146,16 +197,7 @@ private String unescape(String input) {
 	}
 
 	private byte[] readPayload(ByteBuffer buffer, MultiValueMap<String, String> headers) {
-		Integer contentLength = null;
-		if (headers.containsKey("content-length")) {
-			String rawContentLength = headers.getFirst("content-length");
-			try {
-				contentLength = Integer.valueOf(rawContentLength);
-			}
-			catch (NumberFormatException ex) {
-				logger.warn("Ignoring invalid content-length header value: '" + rawContentLength + "'");
-			}
-		}
+		Integer contentLength = getContentLength(headers);
 		if (contentLength != null && contentLength >= 0) {
 			if (buffer.remaining() > contentLength) {
 				byte[] payload = new byte[contentLength];
@@ -184,7 +226,25 @@ private byte[] readPayload(ByteBuffer buffer, MultiValueMap<String, String> head
 		return null;
 	}
 
-	private boolean isEol(ByteBuffer buffer) {
+	protected Integer getContentLength(MultiValueMap<String, String> headers) {
+		if (headers.containsKey(StompHeaderAccessor.STOMP_CONTENT_LENGTH_HEADER)) {
+			String rawContentLength = headers.getFirst(StompHeaderAccessor.STOMP_CONTENT_LENGTH_HEADER);
+			try {
+				return Integer.valueOf(rawContentLength);
+			}
+			catch (NumberFormatException ex) {
+				logger.warn("Ignoring invalid content-length header value: '" + rawContentLength + "'");
+			}
+		}
+		return null;
+	}
+
+	/**
+	 * Try to read an EOL incrementing the buffer position if successful.
+	 *
+	 * @return whether an EOL was consumed
+	 */
+	private boolean tryConsumeEndOfLine(ByteBuffer buffer) {
 		if (buffer.remaining() > 0) {
 			byte b = buffer.get();
 			if (b == '\n') {