001    /**
002     * Licensed to the Apache Software Foundation (ASF) under one or more
003     * contributor license agreements.  See the NOTICE file distributed with
004     * this work for additional information regarding copyright ownership.
005     * The ASF licenses this file to You under the Apache License, Version 2.0
006     * (the "License"); you may not use this file except in compliance with
007     * the License.  You may obtain a copy of the License at
008     *
009     *      http://www.apache.org/licenses/LICENSE-2.0
010     *
011     * Unless required by applicable law or agreed to in writing, software
012     * distributed under the License is distributed on an "AS IS" BASIS,
013     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014     * See the License for the specific language governing permissions and
015     * limitations under the License.
016     */
017    package org.apache.activemq.transport.stomp;
018    
019    import java.io.DataInput;
020    import java.io.DataInputStream;
021    import java.io.DataOutput;
022    import java.io.DataOutputStream;
023    import java.io.IOException;
024    import java.io.InputStream;
025    import java.io.PushbackInputStream;
026    import java.util.HashMap;
027    import java.util.Map;
028    
029    import org.apache.activemq.util.ByteArrayInputStream;
030    import org.apache.activemq.util.ByteArrayOutputStream;
031    import org.apache.activemq.util.ByteSequence;
032    import org.apache.activemq.wireformat.WireFormat;
033    
034    /**
035     * Implements marshalling and unmarsalling the <a
036     * href="http://stomp.codehaus.org/">Stomp</a> protocol.
037     */
038    public class StompWireFormat implements WireFormat {
039    
040        private static final byte[] NO_DATA = new byte[] {};
041        private static final byte[] END_OF_FRAME = new byte[] {0, '\n'};
042    
043        private static final int MAX_COMMAND_LENGTH = 1024;
044        private static final int MAX_HEADER_LENGTH = 1024 * 10;
045        private static final int MAX_HEADERS = 1000;
046        private static final int MAX_DATA_LENGTH = 1024 * 1024 * 100;
047    
048        private int version = 1;
049        private String stompVersion = Stomp.DEFAULT_VERSION;
050    
051        public ByteSequence marshal(Object command) throws IOException {
052            ByteArrayOutputStream baos = new ByteArrayOutputStream();
053            DataOutputStream dos = new DataOutputStream(baos);
054            marshal(command, dos);
055            dos.close();
056            return baos.toByteSequence();
057        }
058    
059        public Object unmarshal(ByteSequence packet) throws IOException {
060            ByteArrayInputStream stream = new ByteArrayInputStream(packet);
061            DataInputStream dis = new DataInputStream(stream);
062            return unmarshal(dis);
063        }
064    
065        public void marshal(Object command, DataOutput os) throws IOException {
066            StompFrame stomp = (org.apache.activemq.transport.stomp.StompFrame)command;
067    
068            if (stomp.getAction().equals(Stomp.Commands.KEEPALIVE)) {
069                os.write(Stomp.BREAK);
070                return;
071            }
072    
073            StringBuilder buffer = new StringBuilder();
074            buffer.append(stomp.getAction());
075            buffer.append(Stomp.NEWLINE);
076    
077            // Output the headers.
078            for (Map.Entry<String, String> entry : stomp.getHeaders().entrySet()) {
079                buffer.append(entry.getKey());
080                buffer.append(Stomp.Headers.SEPERATOR);
081                buffer.append(encodeHeader(entry.getValue()));
082                buffer.append(Stomp.NEWLINE);
083            }
084    
085            // Add a newline to seperate the headers from the content.
086            buffer.append(Stomp.NEWLINE);
087    
088            os.write(buffer.toString().getBytes("UTF-8"));
089            os.write(stomp.getContent());
090            os.write(END_OF_FRAME);
091        }
092    
093        public Object unmarshal(DataInput in) throws IOException {
094    
095            try {
096    
097                // parse action
098                String action = parseAction(in);
099    
100                // Parse the headers
101                HashMap<String, String> headers = parseHeaders(in);
102    
103                // Read in the data part.
104                byte[] data = NO_DATA;
105                String contentLength = headers.get(Stomp.Headers.CONTENT_LENGTH);
106                if ((action.equals(Stomp.Commands.SEND) || action.equals(Stomp.Responses.MESSAGE)) && contentLength != null) {
107    
108                    // Bless the client, he's telling us how much data to read in.
109                    int length = parseContentLength(contentLength);
110    
111                    data = new byte[length];
112                    in.readFully(data);
113    
114                    if (in.readByte() != 0) {
115                        throw new ProtocolException(Stomp.Headers.CONTENT_LENGTH + " bytes were read and " + "there was no trailing null byte", true);
116                    }
117    
118                } else {
119    
120                    // We don't know how much to read.. data ends when we hit a 0
121                    byte b;
122                    ByteArrayOutputStream baos = null;
123                    while ((b = in.readByte()) != 0) {
124    
125                        if (baos == null) {
126                            baos = new ByteArrayOutputStream();
127                        } else if (baos.size() > MAX_DATA_LENGTH) {
128                            throw new ProtocolException("The maximum data length was exceeded", true);
129                        }
130    
131                        baos.write(b);
132                    }
133    
134                    if (baos != null) {
135                        baos.close();
136                        data = baos.toByteArray();
137                    }
138                }
139    
140                return new StompFrame(action, headers, data);
141    
142            } catch (ProtocolException e) {
143                return new StompFrameError(e);
144            }
145        }
146    
147        private String readLine(DataInput in, int maxLength, String errorMessage) throws IOException {
148            ByteSequence sequence = readHeaderLine(in, maxLength, errorMessage);
149            return new String(sequence.getData(), sequence.getOffset(), sequence.getLength(), "UTF-8").trim();
150        }
151    
152        private ByteSequence readHeaderLine(DataInput in, int maxLength, String errorMessage) throws IOException {
153            byte b;
154            ByteArrayOutputStream baos = new ByteArrayOutputStream(maxLength);
155            while ((b = in.readByte()) != '\n') {
156                if (baos.size() > maxLength) {
157                    throw new ProtocolException(errorMessage, true);
158                }
159                baos.write(b);
160            }
161            baos.close();
162            return baos.toByteSequence();
163        }
164    
165        protected String parseAction(DataInput in) throws IOException {
166            String action = null;
167    
168            // skip white space to next real action line
169            while (true) {
170                action = readLine(in, MAX_COMMAND_LENGTH, "The maximum command length was exceeded");
171                if (action == null) {
172                    throw new IOException("connection was closed");
173                } else {
174                    action = action.trim();
175                    if (action.length() > 0) {
176                        break;
177                    }
178                }
179            }
180            return action;
181        }
182    
183        protected HashMap<String, String> parseHeaders(DataInput in) throws IOException {
184            HashMap<String, String> headers = new HashMap<String, String>(25);
185            while (true) {
186                ByteSequence line = readHeaderLine(in, MAX_HEADER_LENGTH, "The maximum header length was exceeded");
187                if (line != null && line.length > 1) {
188    
189                    if (headers.size() > MAX_HEADERS) {
190                        throw new ProtocolException("The maximum number of headers was exceeded", true);
191                    }
192    
193                    try {
194    
195                        ByteArrayInputStream headerLine = new ByteArrayInputStream(line);
196                        ByteArrayOutputStream stream = new ByteArrayOutputStream(line.length);
197    
198                        // First complete the name
199                        int result = -1;
200                        while ((result = headerLine.read()) != -1) {
201                            if (result != ':') {
202                                stream.write(result);
203                            } else {
204                                break;
205                            }
206                        }
207    
208                        ByteSequence nameSeq = stream.toByteSequence();
209                        String name = new String(nameSeq.getData(), nameSeq.getOffset(), nameSeq.getLength(), "UTF-8");
210                        String value = decodeHeader(headerLine);
211                        if (stompVersion.equals(Stomp.V1_0)) {
212                            value = value.trim();
213                        }
214    
215                        if (!headers.containsKey(name)) {
216                            headers.put(name, value);
217                        }
218                    } catch (Exception e) {
219                        throw new ProtocolException("Unable to parser header line [" + line + "]", true);
220                    }
221                } else {
222                    break;
223                }
224            }
225            return headers;
226        }
227    
228        protected int parseContentLength(String contentLength) throws ProtocolException {
229            int length;
230            try {
231                length = Integer.parseInt(contentLength.trim());
232            } catch (NumberFormatException e) {
233                throw new ProtocolException("Specified content-length is not a valid integer", true);
234            }
235    
236            if (length > MAX_DATA_LENGTH) {
237                throw new ProtocolException("The maximum data length was exceeded", true);
238            }
239    
240            return length;
241        }
242    
243        private String encodeHeader(String header) throws IOException {
244            String result = header;
245            if (!stompVersion.equals(Stomp.V1_0)) {
246                byte[] utf8buf = header.getBytes("UTF-8");
247                ByteArrayOutputStream stream = new ByteArrayOutputStream(utf8buf.length);
248                for(byte val : utf8buf) {
249                    switch(val) {
250                    case Stomp.ESCAPE:
251                        stream.write(Stomp.ESCAPE_ESCAPE_SEQ);
252                        break;
253                    case Stomp.BREAK:
254                        stream.write(Stomp.NEWLINE_ESCAPE_SEQ);
255                        break;
256                    case Stomp.COLON:
257                        stream.write(Stomp.COLON_ESCAPE_SEQ);
258                        break;
259                    default:
260                        stream.write(val);
261                    }
262                }
263                result =  new String(stream.toByteArray(), "UTF-8");
264            }
265    
266            return result;
267        }
268    
269        private String decodeHeader(InputStream header) throws IOException {
270            ByteArrayOutputStream decoded = new ByteArrayOutputStream();
271            PushbackInputStream stream = new PushbackInputStream(header);
272    
273            int value = -1;
274            while( (value = stream.read()) != -1) {
275                if (value == 92) {
276    
277                    int next = stream.read();
278                    if (next != -1) {
279                        switch(next) {
280                        case 110:
281                            decoded.write(Stomp.BREAK);
282                            break;
283                        case 99:
284                            decoded.write(Stomp.COLON);
285                            break;
286                        case 92:
287                            decoded.write(Stomp.ESCAPE);
288                            break;
289                        default:
290                            stream.unread(next);
291                            decoded.write(value);
292                        }
293                    } else {
294                        decoded.write(value);
295                    }
296    
297                } else {
298                    decoded.write(value);
299                }
300            }
301    
302            return new String(decoded.toByteArray(), "UTF-8");
303        }
304    
305        public int getVersion() {
306            return version;
307        }
308    
309        public void setVersion(int version) {
310            this.version = version;
311        }
312    
313        public String getStompVersion() {
314            return stompVersion;
315        }
316    
317        public void setStompVersion(String stompVersion) {
318            this.stompVersion = stompVersion;
319        }
320    }