diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompCommand.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompCommand.java index cdd5cc8df3..448d3dc152 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompCommand.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompCommand.java @@ -52,6 +52,7 @@ public enum StompCommand { private static Collection destinationRequired = Arrays.asList(SEND, SUBSCRIBE, MESSAGE); private static Collection subscriptionIdRequired = Arrays.asList(SUBSCRIBE, UNSUBSCRIBE, MESSAGE); + private static Collection contentLengthRequired = Arrays.asList(SEND, MESSAGE, ERROR); private static Collection bodyAllowed = Arrays.asList(SEND, MESSAGE, ERROR); static { @@ -77,6 +78,10 @@ public enum StompCommand { return subscriptionIdRequired.contains(this); } + public boolean requiresContentLength() { + return contentLengthRequired.contains(this); + } + public boolean isBodyAllowed() { return bodyAllowed.contains(this); } diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompEncoder.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompEncoder.java index 0deb332361..ae7e54315a 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompEncoder.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompEncoder.java @@ -55,19 +55,21 @@ public final class StompEncoder { */ public byte[] encode(Message message) { try { - ByteArrayOutputStream baos = new ByteArrayOutputStream(256); + ByteArrayOutputStream baos = new ByteArrayOutputStream(128 + message.getPayload().length); DataOutputStream output = new DataOutputStream(baos); StompHeaderAccessor headers = StompHeaderAccessor.wrap(message); - if (isHeartbeat(headers)) { + if (SimpMessageType.HEARTBEAT == headers.getMessageType()) { + logger.trace("Encoded heartbeat"); output.write(message.getPayload()); } else { - writeCommand(headers, output); + output.write(headers.getCommand().toString().getBytes(UTF8_CHARSET)); + output.write(LF); writeHeaders(headers, message, output); output.write(LF); writeBody(message, output); - output.write((byte)0); + output.write((byte) 0); } return baos.toByteArray(); @@ -77,38 +79,27 @@ public final class StompEncoder { } } - private boolean isHeartbeat(StompHeaderAccessor headers) { - return (headers.getMessageType() == SimpMessageType.HEARTBEAT); - } - - private void writeCommand(StompHeaderAccessor headers, DataOutputStream output) throws IOException { - output.write(headers.getCommand().toString().getBytes(UTF8_CHARSET)); - output.write(LF); - } - private void writeHeaders(StompHeaderAccessor headers, Message message, DataOutputStream output) throws IOException { + StompCommand command = headers.getCommand(); Map> stompHeaders = headers.toStompHeaderMap(); - if (SimpMessageType.HEARTBEAT.equals(headers.getMessageType())) { - logger.trace("Encoded heartbeat"); + boolean shouldEscape = (command != StompCommand.CONNECT && command != StompCommand.CONNECTED); + + if (logger.isDebugEnabled()) { + logger.debug("Encoded STOMP " + command + ", headers=" + stompHeaders); } - else if (logger.isDebugEnabled()) { - logger.debug("Encoded STOMP command=" + headers.getCommand() + " headers=" + stompHeaders); - } - boolean escapeHeaders = shouldEscapeHeaders(headers); + for (Entry> entry : stompHeaders.entrySet()) { - byte[] key = toUtf8Bytes(entry.getKey(), escapeHeaders); + byte[] key = encodeHeaderString(entry.getKey(), shouldEscape); for (String value : entry.getValue()) { output.write(key); output.write(COLON); - output.write(toUtf8Bytes(value, escapeHeaders)); + output.write(encodeHeaderString(value, shouldEscape)); output.write(LF); } } - if ((headers.getCommand() == StompCommand.SEND) || (headers.getCommand() == StompCommand.MESSAGE) || - (headers.getCommand() == StompCommand.ERROR)) { - + if (command.requiresContentLength()) { int contentLength = message.getPayload().length; output.write("content-length:".getBytes(UTF8_CHARSET)); output.write(Integer.toString(contentLength).getBytes(UTF8_CHARSET)); @@ -116,11 +107,7 @@ public final class StompEncoder { } } - private boolean shouldEscapeHeaders(StompHeaderAccessor headers) { - return (headers.getCommand() != StompCommand.CONNECT && headers.getCommand() != StompCommand.CONNECTED); - } - - private byte[] toUtf8Bytes(String input, boolean escape) { + private byte[] encodeHeaderString(String input, boolean escape) { input = escape ? escape(input) : input; return input.getBytes(UTF8_CHARSET); }