diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/BufferingStompDecoder.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/BufferingStompDecoder.java new file mode 100644 index 0000000000..10d46bf26d --- /dev/null +++ b/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. + * + *

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 chunks = new LinkedBlockingQueue(); + + 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> decode(ByteBuffer newData) { + + this.chunks.add(newData); + + checkBufferLimits(); + + if (getExpectedContentLength() != null && getBufferSize() < this.expectedContentLength) { + return Collections.>emptyList(); + } + + ByteBuffer buffer = assembleChunksAndReset(); + + MultiValueMap headers = new LinkedMultiValueMap(); + List> 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; + } + +} diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompCodec.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompCodec.java index f9ffc36c24..1425860bdc 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompCodec.java +++ b/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.function.Function; 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, Message> { @@ -49,14 +52,8 @@ public class StompCodec implements Codec, Message apply(Buffer buffer) { - while (buffer.remaining() > 0) { - Message message = DECODER.decode(buffer.byteBuffer()); - if (message != null) { - next.accept(message); - } - else { - break; - } + for (Message message : DECODER.decode(buffer.byteBuffer())) { + next.accept(message); } return null; } diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompDecoder.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompDecoder.java index 11c80f845d..10e7df6582 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompDecoder.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompDecoder.java @@ -19,6 +19,8 @@ package org.springframework.messaging.simp.stomp; 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.LinkedMultiValueMap; 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 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} + * Decodes one or more STOMP frames from the given {@code buffer} into a + * list of {@link Message}s. + * + *

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 Message decode(ByteBuffer buffer) { + public List> decode(ByteBuffer buffer) { + return decode(buffer, new LinkedMultiValueMap()); + } + + /** + * Decodes one or more STOMP frames from the given {@code buffer} into a + * list of {@link Message}s. + * + *

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> decode(ByteBuffer buffer, MultiValueMap headers) { + List> messages = new ArrayList>(); + while (buffer.hasRemaining()) { + headers.clear(); + Message m = decodeMessage(buffer, headers); + if (m != null) { + messages.add(m); + } + else { + break; + } + } + return messages; + } + + /** + * Decode a single STOMP frame from the given {@code buffer} into a {@link Message}. + */ + private Message decodeMessage(ByteBuffer buffer, MultiValueMap headers) { + Message decodedMessage = null; skipLeadingEol(buffer); buffer.mark(); String command = readCommand(buffer); if (command.length() > 0) { - MultiValueMap headers = readHeaders(buffer); + + readHeaders(buffer, headers); byte[] payload = readPayload(buffer, headers); if (payload != null) { @@ -78,7 +126,7 @@ public class StompDecoder { } else { if (logger.isTraceEnabled()) { - logger.trace("Received incomplete frame. Resetting buffer"); + logger.trace("Received incomplete frame. Resetting buffer."); } buffer.reset(); } @@ -93,9 +141,14 @@ public class StompDecoder { 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; } } @@ -103,17 +156,16 @@ public class StompDecoder { 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 readHeaders(ByteBuffer buffer) { - MultiValueMap headers = new LinkedMultiValueMap(); + private void readHeaders(ByteBuffer buffer, MultiValueMap 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 @@ public class StompDecoder { break; } } - return headers; } private String unescape(String input) { @@ -146,16 +197,7 @@ public class StompDecoder { } private byte[] readPayload(ByteBuffer buffer, MultiValueMap 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 @@ public class StompDecoder { return null; } - private boolean isEol(ByteBuffer buffer) { + protected Integer getContentLength(MultiValueMap 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') { diff --git a/spring-messaging/src/test/java/org/springframework/messaging/simp/stomp/BufferingStompDecoderTests.java b/spring-messaging/src/test/java/org/springframework/messaging/simp/stomp/BufferingStompDecoderTests.java new file mode 100644 index 0000000000..119953cfd8 --- /dev/null +++ b/spring-messaging/src/test/java/org/springframework/messaging/simp/stomp/BufferingStompDecoderTests.java @@ -0,0 +1,183 @@ +/* + * 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.junit.Test; +import org.springframework.messaging.Message; +import org.springframework.messaging.converter.MessageConversionException; + +import java.nio.ByteBuffer; +import java.nio.charset.Charset; +import java.util.Arrays; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.fail; + + +/** + * Unit tests for {@link BufferingStompDecoder}.. + * + * @author Rossen Stoyanchev + * @since 4.0.3 + */ +public class BufferingStompDecoderTests { + + + @Test + public void basic() throws InterruptedException { + + BufferingStompDecoder stompDecoder = new BufferingStompDecoder(128); + String chunk = "SEND\na:alpha\n\nMessage body\0"; + + List> messages = stompDecoder.decode(toByteBuffer(chunk)); + assertEquals(1, messages.size()); + assertEquals("Message body", new String(messages.get(0).getPayload())); + + assertEquals(0, stompDecoder.getBufferSize()); + assertNull(stompDecoder.getExpectedContentLength()); + } + + @Test + public void oneMessageInTwoChunks() throws InterruptedException { + + BufferingStompDecoder stompDecoder = new BufferingStompDecoder(128); + String chunk1 = "SEND\na:alpha\n\nMessage"; + String chunk2 = " body\0"; + + List> messages = stompDecoder.decode(toByteBuffer(chunk1)); + assertEquals(Arrays.asList(), messages); + + messages = stompDecoder.decode(toByteBuffer(chunk2)); + assertEquals(1, messages.size()); + assertEquals("Message body", new String(messages.get(0).getPayload())); + + assertEquals(0, stompDecoder.getBufferSize()); + assertNull(stompDecoder.getExpectedContentLength()); + } + + @Test + public void twoMessagesInOneChunk() throws InterruptedException { + + BufferingStompDecoder stompDecoder = new BufferingStompDecoder(128); + String chunk = "SEND\na:alpha\n\nPayload1\0" + "SEND\na:alpha\n\nPayload2\0"; + List> messages = stompDecoder.decode(toByteBuffer(chunk)); + + assertEquals(2, messages.size()); + assertEquals("Payload1", new String(messages.get(0).getPayload())); + assertEquals("Payload2", new String(messages.get(1).getPayload())); + + assertEquals(0, stompDecoder.getBufferSize()); + assertNull(stompDecoder.getExpectedContentLength()); + } + + @Test + public void oneFullAndOneSplitMessageContentLength() throws InterruptedException { + + int contentLength = "Payload2a-Payload2b".getBytes().length; + + BufferingStompDecoder stompDecoder = new BufferingStompDecoder(128); + String chunk1 = "SEND\na:alpha\n\nPayload1\0SEND\ncontent-length:" + contentLength + "\n"; + List> messages = stompDecoder.decode(toByteBuffer(chunk1)); + + assertEquals(1, messages.size()); + assertEquals("Payload1", new String(messages.get(0).getPayload())); + + assertEquals(23, stompDecoder.getBufferSize()); + assertEquals(contentLength, (int) stompDecoder.getExpectedContentLength()); + + String chunk2 = "\nPayload2a"; + messages = stompDecoder.decode(toByteBuffer(chunk2)); + + assertEquals(0, messages.size()); + assertEquals(33, stompDecoder.getBufferSize()); + assertEquals(contentLength, (int) stompDecoder.getExpectedContentLength()); + + String chunk3 = "-Payload2b\0"; + messages = stompDecoder.decode(toByteBuffer(chunk3)); + + assertEquals(1, messages.size()); + assertEquals("Payload2a-Payload2b", new String(messages.get(0).getPayload())); + assertEquals(0, stompDecoder.getBufferSize()); + assertNull(stompDecoder.getExpectedContentLength()); + } + + @Test + public void oneFullAndOneSplitMessageNoContentLength() throws InterruptedException { + + BufferingStompDecoder stompDecoder = new BufferingStompDecoder(128); + String chunk1 = "SEND\na:alpha\n\nPayload1\0SEND\na:alpha\n"; + List> messages = stompDecoder.decode(toByteBuffer(chunk1)); + + assertEquals(1, messages.size()); + assertEquals("Payload1", new String(messages.get(0).getPayload())); + + assertEquals(13, stompDecoder.getBufferSize()); + assertNull(stompDecoder.getExpectedContentLength()); + + String chunk2 = "\nPayload2a"; + messages = stompDecoder.decode(toByteBuffer(chunk2)); + + assertEquals(0, messages.size()); + assertEquals(23, stompDecoder.getBufferSize()); + assertNull(stompDecoder.getExpectedContentLength()); + + String chunk3 = "-Payload2b\0"; + messages = stompDecoder.decode(toByteBuffer(chunk3)); + + assertEquals(1, messages.size()); + assertEquals("Payload2a-Payload2b", new String(messages.get(0).getPayload())); + assertEquals(0, stompDecoder.getBufferSize()); + assertNull(stompDecoder.getExpectedContentLength()); + } + + @Test + public void oneFullAndOneSplitWithContentLengthExceedingBufferSize() throws InterruptedException { + + BufferingStompDecoder stompDecoder = new BufferingStompDecoder(128); + String chunk1 = "SEND\na:alpha\n\nPayload1\0SEND\ncontent-length:129\n"; + List> messages = stompDecoder.decode(toByteBuffer(chunk1)); + + assertEquals("We should have gotten the 1st message", 1, messages.size()); + assertEquals("Payload1", new String(messages.get(0).getPayload())); + + assertEquals(24, stompDecoder.getBufferSize()); + assertEquals(129, (int) stompDecoder.getExpectedContentLength()); + + try { + String chunk2 = "\nPayload2a"; + stompDecoder.decode(toByteBuffer(chunk2)); + fail("Expected exception"); + } + catch (StompConversionException ex) { + } + } + + @Test(expected = StompConversionException.class) + public void bufferSizeLimit() throws InterruptedException { + BufferingStompDecoder stompDecoder = new BufferingStompDecoder(10); + String payload = "SEND\na:alpha\n\nMessage body"; + stompDecoder.decode(toByteBuffer(payload)); + } + + + private ByteBuffer toByteBuffer(String chunk) { + return ByteBuffer.wrap(chunk.getBytes(Charset.forName("UTF-8"))); + } + +} diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java b/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java index 67f8963cfb..3827eed443 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java +++ b/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.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. @@ -21,7 +21,9 @@ import java.nio.ByteBuffer; import java.security.Principal; import java.util.Arrays; import java.util.List; +import java.util.Map; import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; @@ -29,6 +31,7 @@ import org.apache.commons.logging.LogFactory; import org.springframework.messaging.Message; import org.springframework.messaging.MessageChannel; import org.springframework.messaging.simp.SimpMessageType; +import org.springframework.messaging.simp.stomp.BufferingStompDecoder; import org.springframework.messaging.simp.stomp.StompCommand; import org.springframework.messaging.simp.stomp.StompConversionException; import org.springframework.messaging.simp.stomp.StompDecoder; @@ -66,13 +69,31 @@ public class StompSubProtocolHandler implements SubProtocolHandler { private static final Log logger = LogFactory.getLog(StompSubProtocolHandler.class); - private final StompDecoder stompDecoder = new StompDecoder(); + private int messageBufferSizeLimit = 64 * 1024; + + private final Map decoders = new ConcurrentHashMap(); private final StompEncoder stompEncoder = new StompEncoder(); private UserSessionRegistry userSessionRegistry; + /** + * TODO + * @param messageBufferSizeLimit + */ + public void setMessageBufferSizeLimit(int messageBufferSizeLimit) { + this.messageBufferSizeLimit = messageBufferSizeLimit; + } + + /** + * TODO + * @return + */ + public int getMessageBufferSizeLimit() { + return this.messageBufferSizeLimit; + } + /** * Provide a registry with which to register active user session ids. * @see org.springframework.messaging.simp.user.UserDestinationMessageHandler @@ -99,49 +120,53 @@ public class StompSubProtocolHandler implements SubProtocolHandler { public void handleMessageFromClient(WebSocketSession session, WebSocketMessage webSocketMessage, MessageChannel outputChannel) { - Message message = null; - Throwable decodeFailure = null; + List> messages = null; try { Assert.isInstanceOf(TextMessage.class, webSocketMessage); TextMessage textMessage = (TextMessage) webSocketMessage; ByteBuffer byteBuffer = ByteBuffer.wrap(textMessage.asBytes()); - message = this.stompDecoder.decode(byteBuffer); - if (message == null) { - decodeFailure = new IllegalStateException("Not a valid STOMP frame: " + textMessage.getPayload()); + BufferingStompDecoder decoder = this.decoders.get(session.getId()); + if (decoder == null) { + throw new IllegalStateException("No decoder for session id '" + session.getId() + "'"); + } + + messages = decoder.decode(byteBuffer); + if (messages.isEmpty()) { + logger.debug("Incomplete STOMP frame content received," + "buffered=" + + decoder.getBufferSize() + ", buffer size limit=" + decoder.getBufferSizeLimit()); + return; } } catch (Throwable ex) { - decodeFailure = ex; - } - - if (decodeFailure != null) { - logger.error("Failed to parse WebSocket message as STOMP frame", decodeFailure); - sendErrorMessage(session, decodeFailure); + logger.error("Failed to parse WebSocket message to STOMP frame(s)", ex); + sendErrorMessage(session, ex); return; } - try { - StompHeaderAccessor headers = StompHeaderAccessor.wrap(message); - if (logger.isTraceEnabled()) { - if (SimpMessageType.HEARTBEAT.equals(headers.getMessageType())) { - logger.trace("Received heartbeat from client session=" + session.getId()); - } - else { - logger.trace("Received message from client session=" + session.getId()); + for (Message message : messages) { + try { + StompHeaderAccessor headers = StompHeaderAccessor.wrap(message); + if (logger.isTraceEnabled()) { + if (SimpMessageType.HEARTBEAT.equals(headers.getMessageType())) { + logger.trace("Received heartbeat from client session=" + session.getId()); + } + else { + logger.trace("Received message from client session=" + session.getId()); + } } + + headers.setSessionId(session.getId()); + headers.setSessionAttributes(session.getAttributes()); + headers.setUser(session.getPrincipal()); + + message = MessageBuilder.withPayload(message.getPayload()).setHeaders(headers).build(); + outputChannel.send(message); + } + catch (Throwable ex) { + logger.error("Terminating STOMP session due to failure to send message", ex); + sendErrorMessage(session, ex); } - - headers.setSessionId(session.getId()); - headers.setSessionAttributes(session.getAttributes()); - headers.setUser(session.getPrincipal()); - - message = MessageBuilder.withPayload(message.getPayload()).setHeaders(headers).build(); - outputChannel.send(message); - } - catch (Throwable ex) { - logger.error("Terminating STOMP session due to failure to send message", ex); - sendErrorMessage(session, ex); } } @@ -281,11 +306,14 @@ public class StompSubProtocolHandler implements SubProtocolHandler { @Override public void afterSessionStarted(WebSocketSession session, MessageChannel outputChannel) { + this.decoders.put(session.getId(), new BufferingStompDecoder(getMessageBufferSizeLimit())); } @Override public void afterSessionEnded(WebSocketSession session, CloseStatus closeStatus, MessageChannel outputChannel) { + this.decoders.remove(session.getId()); + Principal principal = session.getPrincipal(); if ((this.userSessionRegistry != null) && (principal != null)) { String userName = resolveNameForUserSessionRegistry(principal); diff --git a/spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupportTests.java b/spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupportTests.java index 4252001147..6314dcd204 100644 --- a/spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupportTests.java +++ b/spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupportTests.java @@ -41,6 +41,7 @@ import org.springframework.stereotype.Controller; import org.springframework.web.servlet.HandlerMapping; import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping; import org.springframework.web.socket.TextMessage; +import org.springframework.web.socket.WebSocketSession; import org.springframework.web.socket.handler.TestWebSocketSession; import org.springframework.web.socket.messaging.StompTextMessageBuilder; import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler; @@ -81,8 +82,11 @@ public class WebSocketMessageBrokerConfigurationSupportTests { TestChannel channel = this.config.getBean("clientInboundChannel", TestChannel.class); SubProtocolWebSocketHandler webSocketHandler = this.config.getBean(SubProtocolWebSocketHandler.class); + WebSocketSession session = new TestWebSocketSession("s1"); + webSocketHandler.afterConnectionEstablished(session); + TextMessage textMessage = StompTextMessageBuilder.create(StompCommand.SEND).headers("destination:/foo").build(); - webSocketHandler.handleMessage(new TestWebSocketSession(), textMessage); + webSocketHandler.handleMessage(session, textMessage); Message message = channel.messages.get(0); StompHeaderAccessor headers = StompHeaderAccessor.wrap(message); diff --git a/spring-websocket/src/test/java/org/springframework/web/socket/messaging/StompSubProtocolHandlerTests.java b/spring-websocket/src/test/java/org/springframework/web/socket/messaging/StompSubProtocolHandlerTests.java index cfc769d033..28718b1716 100644 --- a/spring-websocket/src/test/java/org/springframework/web/socket/messaging/StompSubProtocolHandlerTests.java +++ b/spring-websocket/src/test/java/org/springframework/web/socket/messaging/StompSubProtocolHandlerTests.java @@ -20,6 +20,7 @@ import java.nio.ByteBuffer; import java.util.Arrays; import java.util.Collections; import java.util.HashSet; +import java.util.List; import org.junit.Before; import org.junit.Test; @@ -145,8 +146,8 @@ public class StompSubProtocolHandlerTests { assertEquals(1, this.session.getSentMessages().size()); TextMessage textMessage = (TextMessage) this.session.getSentMessages().get(0); - Message message = new StompDecoder().decode(ByteBuffer.wrap(textMessage.getPayload().getBytes())); - StompHeaderAccessor replyHeaders = StompHeaderAccessor.wrap(message); + List> message = new StompDecoder().decode(ByteBuffer.wrap(textMessage.getPayload().getBytes())); + StompHeaderAccessor replyHeaders = StompHeaderAccessor.wrap(message.get(0)); assertEquals(StompCommand.CONNECTED, replyHeaders.getCommand()); assertEquals("1.1", replyHeaders.getVersion()); @@ -176,6 +177,7 @@ public class StompSubProtocolHandlerTests { TextMessage textMessage = StompTextMessageBuilder.create(StompCommand.CONNECT).headers( "login:guest", "passcode:guest", "accept-version:1.1,1.0", "heart-beat:10000,10000").build(); + this.protocolHandler.afterSessionStarted(this.session, this.channel); this.protocolHandler.handleMessageFromClient(this.session, textMessage, this.channel); verify(this.channel).send(this.messageCaptor.capture());