View Javadoc
1   /*
2    * Copyright 2014 The Netty Project
3    *
4    * The Netty Project licenses this file to you under the Apache License,
5    * version 2.0 (the "License"); you may not use this file except in compliance
6    * with the License. You may obtain a copy of the License at:
7    *
8    *   https://www.apache.org/licenses/LICENSE-2.0
9    *
10   * Unless required by applicable law or agreed to in writing, software
11   * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
12   * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
13   * License for the specific language governing permissions and limitations
14   * under the License.
15   */
16  package io.netty.handler.codec.stomp;
17  
18  import io.netty.buffer.ByteBuf;
19  import io.netty.buffer.ByteBufUtil;
20  import io.netty.channel.ChannelHandlerContext;
21  import io.netty.handler.codec.MessageToMessageEncoder;
22  import io.netty.util.concurrent.FastThreadLocal;
23  import io.netty.util.internal.AppendableCharSequence;
24  import io.netty.util.internal.PlatformDependent;
25  
26  import java.util.LinkedHashMap;
27  import java.util.List;
28  import java.util.Map.Entry;
29  
30  import static io.netty.handler.codec.stomp.StompConstants.NUL;
31  import static io.netty.handler.codec.stomp.StompHeaders.ACCEPT_VERSION;
32  import static io.netty.handler.codec.stomp.StompHeaders.ACK;
33  import static io.netty.handler.codec.stomp.StompHeaders.CONTENT_LENGTH;
34  import static io.netty.handler.codec.stomp.StompHeaders.CONTENT_TYPE;
35  import static io.netty.handler.codec.stomp.StompHeaders.DESTINATION;
36  import static io.netty.handler.codec.stomp.StompHeaders.HEART_BEAT;
37  import static io.netty.handler.codec.stomp.StompHeaders.HOST;
38  import static io.netty.handler.codec.stomp.StompHeaders.ID;
39  import static io.netty.handler.codec.stomp.StompHeaders.LOGIN;
40  import static io.netty.handler.codec.stomp.StompHeaders.MESSAGE;
41  import static io.netty.handler.codec.stomp.StompHeaders.MESSAGE_ID;
42  import static io.netty.handler.codec.stomp.StompHeaders.PASSCODE;
43  import static io.netty.handler.codec.stomp.StompHeaders.RECEIPT;
44  import static io.netty.handler.codec.stomp.StompHeaders.RECEIPT_ID;
45  import static io.netty.handler.codec.stomp.StompHeaders.SERVER;
46  import static io.netty.handler.codec.stomp.StompHeaders.SESSION;
47  import static io.netty.handler.codec.stomp.StompHeaders.SUBSCRIPTION;
48  import static io.netty.handler.codec.stomp.StompHeaders.TRANSACTION;
49  import static io.netty.handler.codec.stomp.StompHeaders.VERSION;
50  
51  /**
52   * Encodes a {@link StompFrame} or a {@link StompSubframe} into a {@link ByteBuf}.
53   */
54  public class StompSubframeEncoder extends MessageToMessageEncoder<StompSubframe> {
55  
56      private static final int ESCAPE_HEADER_KEY_CACHE_LIMIT = 32;
57      private static final float DEFAULT_LOAD_FACTOR = 0.75f;
58      private static final FastThreadLocal<LinkedHashMap<CharSequence, CharSequence>> ESCAPE_HEADER_KEY_CACHE =
59              new FastThreadLocal<LinkedHashMap<CharSequence, CharSequence>>() {
60                  @Override
61                  protected LinkedHashMap<CharSequence, CharSequence> initialValue() throws Exception {
62                      LinkedHashMap<CharSequence, CharSequence> cache = new LinkedHashMap<CharSequence, CharSequence>(
63                              ESCAPE_HEADER_KEY_CACHE_LIMIT, DEFAULT_LOAD_FACTOR, true) {
64  
65                          @Override
66                          protected boolean removeEldestEntry(Entry eldest) {
67                              return size() > ESCAPE_HEADER_KEY_CACHE_LIMIT;
68                          }
69                      };
70  
71                      cache.put(ACCEPT_VERSION, ACCEPT_VERSION);
72                      cache.put(HOST, HOST);
73                      cache.put(LOGIN, LOGIN);
74                      cache.put(PASSCODE, PASSCODE);
75                      cache.put(HEART_BEAT, HEART_BEAT);
76                      cache.put(VERSION, VERSION);
77                      cache.put(SESSION, SESSION);
78                      cache.put(SERVER, SERVER);
79                      cache.put(DESTINATION, DESTINATION);
80                      cache.put(ID, ID);
81                      cache.put(ACK, ACK);
82                      cache.put(TRANSACTION, TRANSACTION);
83                      cache.put(RECEIPT, RECEIPT);
84                      cache.put(MESSAGE_ID, MESSAGE_ID);
85                      cache.put(SUBSCRIPTION, SUBSCRIPTION);
86                      cache.put(RECEIPT_ID, RECEIPT_ID);
87                      cache.put(MESSAGE, MESSAGE);
88                      cache.put(CONTENT_LENGTH, CONTENT_LENGTH);
89                      cache.put(CONTENT_TYPE, CONTENT_TYPE);
90  
91                      return cache;
92                  }
93              };
94  
95      public StompSubframeEncoder() {
96          super(StompSubframe.class);
97      }
98  
99      @Override
100     protected void encode(ChannelHandlerContext ctx, StompSubframe msg, List<Object> out) throws Exception {
101         if (msg instanceof StompFrame) {
102             StompFrame stompFrame = (StompFrame) msg;
103             ByteBuf buf = encodeFullFrame(stompFrame, ctx);
104 
105             out.add(convertFullFrame(stompFrame, buf));
106         } else if (msg instanceof StompHeadersSubframe) {
107             StompHeadersSubframe stompHeadersSubframe = (StompHeadersSubframe) msg;
108             ByteBuf buf = ctx.alloc().buffer(headersSubFrameSize(stompHeadersSubframe));
109             try {
110                 encodeHeaders(stompHeadersSubframe, buf);
111             } catch (Exception e) {
112                 buf.release();
113                 PlatformDependent.throwException(e);
114             }
115 
116             out.add(convertHeadersSubFrame(stompHeadersSubframe, buf));
117         } else if (msg instanceof StompContentSubframe) {
118             StompContentSubframe stompContentSubframe = (StompContentSubframe) msg;
119             ByteBuf buf = encodeContent(stompContentSubframe, ctx);
120 
121             out.add(convertContentSubFrame(stompContentSubframe, buf));
122         }
123     }
124 
125     /**
126      * An extension method to convert a STOMP encoded buffer to a different message type
127      * based on an original {@link StompFrame} full frame.
128      *
129      * <p>By default an encoded buffer is returned as is.
130      */
131     protected Object convertFullFrame(StompFrame original, ByteBuf encoded) {
132         return encoded;
133     }
134 
135     /**
136      * An extension method to convert a STOMP encoded buffer to a different message type
137      * based on an original {@link StompHeadersSubframe} headers sub frame.
138      *
139      * <p>By default an encoded buffer is returned as is.
140      */
141     protected Object convertHeadersSubFrame(StompHeadersSubframe original, ByteBuf encoded) {
142         return encoded;
143     }
144 
145     /**
146      * An extension method to convert a STOMP encoded buffer to a different message type
147      * based on an original {@link StompHeadersSubframe} content sub frame.
148      *
149      * <p>By default an encoded buffer is returned as is.
150      */
151     protected Object convertContentSubFrame(StompContentSubframe original, ByteBuf encoded) {
152         return encoded;
153     }
154 
155     /**
156      * Returns a heuristic size for headers (32 bytes per header line) + (2 bytes for colon and eol) + (additional
157      * command buffer).
158      */
159     protected int headersSubFrameSize(StompHeadersSubframe headersSubframe) {
160         int estimatedSize = headersSubframe.headers().size() * 34 + 48;
161         if (estimatedSize < 128) {
162             return 128;
163         }
164 
165         return Math.max(estimatedSize, 256);
166     }
167 
168     private ByteBuf encodeFullFrame(StompFrame frame, ChannelHandlerContext ctx) {
169         int contentReadableBytes = frame.content().readableBytes();
170         ByteBuf buf = ctx.alloc().buffer(headersSubFrameSize(frame) + contentReadableBytes);
171         try {
172             encodeHeaders(frame, buf);
173         } catch (Exception e) {
174             buf.release();
175             PlatformDependent.throwException(e);
176         }
177 
178         if (contentReadableBytes > 0) {
179             buf.writeBytes(frame.content());
180         }
181 
182         return buf.writeByte(NUL);
183     }
184 
185     private static void encodeHeaders(StompHeadersSubframe frame, ByteBuf buf) {
186         StompCommand command = frame.command();
187         ByteBufUtil.writeUtf8(buf, command.toString());
188         buf.writeByte(StompConstants.LF);
189 
190         boolean shouldEscape = shouldEscape(command);
191         LinkedHashMap<CharSequence, CharSequence> cache = ESCAPE_HEADER_KEY_CACHE.get();
192         for (Entry<CharSequence, CharSequence> entry : frame.headers()) {
193             CharSequence headerKey = entry.getKey();
194             CharSequence headerValue = entry.getValue();
195             if (headerKey.length() == 0) {
196                 throw new IllegalArgumentException("STOMP " + command + " contains empty header name");
197             }
198             if (shouldEscape) {
199                 CharSequence cachedHeaderKey = cache.get(headerKey);
200                 if (cachedHeaderKey == null) {
201                     cachedHeaderKey = escape(command, "header name", headerKey);
202                     cache.put(headerKey, cachedHeaderKey);
203                 }
204                 headerKey = cachedHeaderKey;
205                 headerValue = escape(command, "header value", headerValue);
206             } else {
207                 // For CONNECT/CONNECTED: don't escape but REJECT illegal characters
208                 validateNoIllegalCharacters(command, headerKey, "header name");
209                 validateNoIllegalCharacters(command, headerValue, "header value");
210             }
211 
212             ByteBufUtil.writeUtf8(buf, headerKey);
213             buf.writeByte(StompConstants.COLON);
214 
215             ByteBufUtil.writeUtf8(buf, headerValue);
216             buf.writeByte(StompConstants.LF);
217         }
218 
219         buf.writeByte(StompConstants.LF);
220     }
221 
222     private static ByteBuf encodeContent(StompContentSubframe content, ChannelHandlerContext ctx) {
223         if (content instanceof LastStompContentSubframe) {
224             ByteBuf buf = ctx.alloc().buffer(content.content().readableBytes() + 1);
225             buf.writeBytes(content.content());
226             buf.writeByte(StompConstants.NUL);
227             return buf;
228         }
229 
230         return content.content().retain();
231     }
232 
233     private static boolean shouldEscape(StompCommand command) {
234         return command != StompCommand.CONNECT && command != StompCommand.CONNECTED;
235     }
236 
237     private static void validateNoIllegalCharacters(StompCommand command, CharSequence value, String type) {
238         for (int i = 0; i < value.length(); i++) {
239             char c = value.charAt(i);
240             if (c == '\n' || c == '\r' || c == ':' || c == '\0') {
241                 throw newIllegalCharacterException(command, type, i);
242             }
243         }
244     }
245 
246     private static IllegalArgumentException newIllegalCharacterException(StompCommand command, String type, int index) {
247         return new IllegalArgumentException(
248                 "STOMP " + command + " " + type + " contains illegal character at index " + index);
249     }
250 
251     private static CharSequence escape(StompCommand command, String type, CharSequence input) {
252         AppendableCharSequence builder = null;
253         for (int i = 0; i < input.length(); i++) {
254             char chr = input.charAt(i);
255             if (chr == '\\') {
256                 builder = escapeBuilder(builder, input, i);
257                 builder.append("\\\\");
258             } else if (chr == ':') {
259                 builder = escapeBuilder(builder, input, i);
260                 builder.append("\\c");
261             } else if (chr == '\n') {
262                 builder = escapeBuilder(builder, input, i);
263                 builder.append("\\n");
264             } else if (chr == '\r') {
265                 builder = escapeBuilder(builder, input, i);
266                 builder.append("\\r");
267             } else if (chr == '\0') {
268                 // The NUL character has no escape and is always illegal.
269                 throw newIllegalCharacterException(command, type, i);
270             } else if (builder != null) {
271                 builder.append(chr);
272             }
273         }
274 
275         return builder != null? builder : input;
276     }
277 
278     private static AppendableCharSequence escapeBuilder(AppendableCharSequence builder, CharSequence input,
279                                                         int offset) {
280         if (builder != null) {
281             return builder;
282         }
283 
284         // Add extra overhead to the input char sequence to avoid resizing during escaping.
285         return new AppendableCharSequence(input.length() + 8).append(input, 0, offset);
286     }
287 }