Merge pull request #1236 from dreis2211-stomp-header-encoding

Author: rstoyanchev

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

@@ -21,9 +21,11 @@
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.concurrent.ConcurrentHashMap;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -50,6 +52,27 @@ public class StompEncoder  {
 
 	private static final Log logger = LogFactory.getLog(StompEncoder.class);
 
+	private static final int HEADER_KEY_CACHE_LIMIT = 32;
+
+
+	private final Map<String, byte[]> headerKeyAccessCache =
+			new ConcurrentHashMap<>(HEADER_KEY_CACHE_LIMIT);
+
+	@SuppressWarnings("serial")
+	private final Map<String, byte[]> headerKeyUpdateCache =
+			new LinkedHashMap<String, byte[]>(HEADER_KEY_CACHE_LIMIT, 0.75f, true) {
+				@Override
+				protected boolean removeEldestEntry(Map.Entry<String, byte[]> eldest) {
+					if (size() > HEADER_KEY_CACHE_LIMIT) {
+						headerKeyAccessCache.remove(eldest.getKey());
+						return true;
+					}
+					else {
+						return false;
+					}
+				}
+			};
+
 
 	/**
 	 * Encodes the given STOMP {@code message} into a {@code byte[]}
@@ -130,11 +153,11 @@ private void writeHeaders(StompCommand command, Map<String, Object> headers, byt
 				values = Collections.singletonList(StompHeaderAccessor.getPasscode(headers));
 			}
 
-			byte[] encodedKey = encodeHeaderString(entry.getKey(), shouldEscape);
+			byte[] encodedKey = encodeHeaderKey(entry.getKey(), shouldEscape);
 			for (String value : values) {
 				output.write(encodedKey);
 				output.write(COLON);
-				output.write(encodeHeaderString(value, shouldEscape));
+				output.write(encodeHeaderValue(value, shouldEscape));
 				output.write(LF);
 			}
 		}
@@ -147,7 +170,23 @@ private void writeHeaders(StompCommand command, Map<String, Object> headers, byt
 		}
 	}
 
-	private byte[] encodeHeaderString(String input, boolean escape) {
+	private byte[] encodeHeaderKey(String input, boolean escape) {
+		String inputToUse = (escape ? escape(input) : input);
+		if (this.headerKeyAccessCache.containsKey(inputToUse)) {
+			return this.headerKeyAccessCache.get(inputToUse);
+		}
+		synchronized (this.headerKeyUpdateCache) {
+			byte[] bytes = this.headerKeyUpdateCache.get(inputToUse);
+			if (bytes == null) {
+				bytes = inputToUse.getBytes(StandardCharsets.UTF_8);
+				this.headerKeyAccessCache.put(inputToUse, bytes);
+				this.headerKeyUpdateCache.put(inputToUse, bytes);
+			}
+			return bytes;
+		}
+	}
+
+	private byte[] encodeHeaderValue(String input, boolean escape) {
 		String inputToUse = (escape ? escape(input) : input);
 		return inputToUse.getBytes(StandardCharsets.UTF_8);
 	}
@@ -157,26 +196,38 @@ private byte[] encodeHeaderString(String input, boolean escape) {
 	 * <a href="http://stomp.github.io/stomp-specification-1.2.html#Value_Encoding">"Value Encoding"</a>.
 	 */
 	private String escape(String inString) {
-		StringBuilder sb = new StringBuilder(inString.length());
+		StringBuilder sb = null;
 		for (int i = 0; i < inString.length(); i++) {
 			char c = inString.charAt(i);
 			if (c == '\\') {
+				sb = getStringBuilder(sb, inString, i);
 				sb.append("\\\\");
 			}
 			else if (c == ':') {
+				sb = getStringBuilder(sb, inString, i);
 				sb.append("\\c");
 			}
 			else if (c == '\n') {
-				 sb.append("\\n");
+				sb = getStringBuilder(sb, inString, i);
+				sb.append("\\n");
 			}
 			else if (c == '\r') {
+				sb = getStringBuilder(sb, inString, i);
 				sb.append("\\r");
 			}
-			else {
+			else if (sb != null){
 				sb.append(c);
 			}
 		}
-		return sb.toString();
+		return (sb != null ? sb.toString() : inString);
+	}
+
+	private StringBuilder getStringBuilder(StringBuilder sb, String inString, int i) {
+		if (sb == null) {
+			sb = new StringBuilder(inString.length());
+			sb.append(inString.substring(0, i));
+		}
+		return sb;
 	}
 
 	private void writeBody(byte[] payload, DataOutputStream output) throws IOException {

spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistry.java

@@ -142,7 +142,6 @@ protected void setApplicationContext(ApplicationContext applicationContext) {
 		this.stompHandler.setApplicationEventPublisher(applicationContext);
 	}
 
-
 	/**
 	 * Return a handler mapping with the mapped ViewControllers; or {@code null}
 	 * in case of no registrations.

spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java

@@ -93,9 +93,9 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
 
 	private int messageSizeLimit = 64 * 1024;
 
-	private final StompEncoder stompEncoder = new StompEncoder();
+	private StompEncoder stompEncoder = new StompEncoder();
 
-	private final StompDecoder stompDecoder = new StompDecoder();
+	private StompDecoder stompDecoder = new StompDecoder();
 
 	private final Map<String, BufferingStompDecoder> decoders = new ConcurrentHashMap<>();
 
@@ -146,6 +146,24 @@ public int getMessageSizeLimit() {
 		return this.messageSizeLimit;
 	}
 
+	/**
+	 * Configure a {@link StompEncoder} for encoding STOMP frames
+	 * @param encoder the encoder
+	 * @since 4.3.5
+	 */
+	public void setEncoder(StompEncoder encoder) {
+		this.stompEncoder = encoder;
+	}
+
+	/**
+	 * Configure a {@link StompDecoder} for decoding STOMP frames
+	 * @param decoder the decoder
+	 * @since 4.3.5
+	 */
+	public void setDecoder(StompDecoder decoder) {
+		this.stompDecoder = decoder;
+	}
+
 	/**
 	 * Configure a {@link MessageHeaderInitializer} to apply to the headers of all
 	 * messages created from decoded STOMP frames and other messages sent to the