From 0756709698d46e33db701f8927798235fdb9e124 Mon Sep 17 00:00:00 2001 From: Christoph Dreis Date: Fri, 11 Nov 2016 21:41:25 +0100 Subject: [PATCH] SPR-14901 Allow customization of STOMP message header encoding Fixes SPR-14901 --- .../messaging/simp/stomp/StompEncoder.java | 53 ++++++++++++++++--- .../WebMvcStompEndpointRegistry.java | 1 - .../messaging/StompSubProtocolHandler.java | 27 ++++++++-- 3 files changed, 69 insertions(+), 12 deletions(-) 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 67e20886365..1654a9f3e10 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 @@ -21,6 +21,7 @@ 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; @@ -50,6 +51,16 @@ public class StompEncoder { private static final Log logger = LogFactory.getLog(StompEncoder.class); + private static final int HEADER_KEY_CACHE_LIMIT = 32; + + @SuppressWarnings("serial") + private final Map headerKeyCache = + new LinkedHashMap(HEADER_KEY_CACHE_LIMIT, 0.75f, true) { + @Override + protected boolean removeEldestEntry(Map.Entry eldest) { + return size() > HEADER_KEY_CACHE_LIMIT; + } + }; /** * Encodes the given STOMP {@code message} into a {@code byte[]} @@ -130,11 +141,11 @@ private void writeHeaders(StompCommand command, Map 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,9 +158,23 @@ private void writeHeaders(StompCommand command, Map headers, byt } } - private byte[] encodeHeaderString(String input, boolean escape) { + private byte[] encodeHeaderKey(String input, boolean escape) { + String inputToUse = (escape ? escape(input) : input); + if (headerKeyCache.containsKey(inputToUse)) { + return headerKeyCache.get(inputToUse); + } + byte[] bytes = encodeHeaderString(inputToUse); + headerKeyCache.put(inputToUse, bytes); + return bytes; + } + + private byte[] encodeHeaderValue(String input, boolean escape) { String inputToUse = (escape ? escape(input) : input); - return inputToUse.getBytes(StandardCharsets.UTF_8); + return encodeHeaderString(inputToUse); + } + + private byte[] encodeHeaderString(String input) { + return input.getBytes(StandardCharsets.UTF_8); } /** @@ -157,26 +182,38 @@ private byte[] encodeHeaderString(String input, boolean escape) { * "Value Encoding". */ 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 { diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistry.java b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistry.java index e30ddfb9e0b..d3cb5b1a8ef 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistry.java +++ b/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. 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 3c0408bb96b..06cc1c48494 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 @@ -88,14 +88,13 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE private static final byte[] EMPTY_PAYLOAD = new byte[0]; - private StompSubProtocolErrorHandler errorHandler; private int messageSizeLimit = 64 * 1024; - private final StompEncoder stompEncoder = new StompEncoder(); + private StompEncoder stompEncoder; - private final StompDecoder stompDecoder = new StompDecoder(); + private StompDecoder stompDecoder; private final Map decoders = new ConcurrentHashMap<>(); @@ -107,6 +106,10 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE private final Stats stats = new Stats(); + public StompSubProtocolHandler() { + setEncoder(new StompEncoder()); + setDecoder(new StompDecoder()); + } /** * Configure a handler for error messages sent to clients which allows @@ -126,6 +129,24 @@ public StompSubProtocolErrorHandler getErrorHandler() { return this.errorHandler; } + /** + * 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 the maximum size allowed for an incoming STOMP message. * Since a STOMP message can be received in multiple WebSocket messages,