Polish

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

@@ -52,6 +52,7 @@ public enum StompCommand {
 
 	private static Collection<StompCommand> destinationRequired = Arrays.asList(SEND, SUBSCRIBE, MESSAGE);
 	private static Collection<StompCommand> subscriptionIdRequired = Arrays.asList(SUBSCRIBE, UNSUBSCRIBE, MESSAGE);
+	private static Collection<StompCommand> contentLengthRequired = Arrays.asList(SEND, MESSAGE, ERROR);
 	private static Collection<StompCommand> 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);
 	}

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

@@ -55,19 +55,21 @@ public final class StompEncoder  {
 	 */
 	public byte[] encode(Message<byte[]> 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<byte[]> message, DataOutputStream output)
 			throws IOException {
 
+		StompCommand command = headers.getCommand();
 		Map<String,List<String>> stompHeaders = headers.toStompHeaderMap();
-		if (SimpMessageType.HEARTBEAT.equals(headers.getMessageType())) {
-			logger.trace("Encoded heartbeat");
-		}
-		else if (logger.isDebugEnabled()) {
-			logger.debug("Encoded STOMP command=" + headers.getCommand() + " headers=" + stompHeaders);
+		boolean shouldEscape = (command != StompCommand.CONNECT && command != StompCommand.CONNECTED);
+
+		if (logger.isDebugEnabled()) {
+			logger.debug("Encoded STOMP " + command + ", headers=" + stompHeaders);
 		}
-		boolean escapeHeaders = shouldEscapeHeaders(headers);
+
 		for (Entry<String, List<String>> 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);
 	}