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 */
017package org.apache.activemq.transport.stomp;
018
019import java.io.DataInput;
020import java.io.DataInputStream;
021import java.io.DataOutput;
022import java.io.DataOutputStream;
023import java.io.IOException;
024import java.io.InputStream;
025import java.io.PushbackInputStream;
026import java.util.HashMap;
027import java.util.Map;
028
029import org.apache.activemq.util.ByteArrayInputStream;
030import org.apache.activemq.util.ByteArrayOutputStream;
031import org.apache.activemq.util.ByteSequence;
032import org.apache.activemq.wireformat.WireFormat;
033
034/**
035 * Implements marshalling and unmarsalling the <a
036 * href="http://stomp.codehaus.org/">Stomp</a> protocol.
037 */
038public 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 boolean encodingEnabled = false;
049    private int version = 1;
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").trim();
210                    String value = decodeHeader(headerLine);
211                    headers.put(name, value);
212                } catch (Exception e) {
213                    throw new ProtocolException("Unable to parser header line [" + line + "]", true);
214                }
215            } else {
216                break;
217            }
218        }
219        return headers;
220    }
221
222    protected int parseContentLength(String contentLength) throws ProtocolException {
223        int length;
224        try {
225            length = Integer.parseInt(contentLength.trim());
226        } catch (NumberFormatException e) {
227            throw new ProtocolException("Specified content-length is not a valid integer", true);
228        }
229
230        if (length > MAX_DATA_LENGTH) {
231            throw new ProtocolException("The maximum data length was exceeded", true);
232        }
233
234        return length;
235    }
236
237    private String encodeHeader(String header) throws IOException {
238        String result = header;
239        if (this.encodingEnabled) {
240            byte[] utf8buf = header.getBytes("UTF-8");
241            ByteArrayOutputStream stream = new ByteArrayOutputStream(utf8buf.length);
242            for(byte val : utf8buf) {
243                switch(val) {
244                case Stomp.ESCAPE:
245                    stream.write(Stomp.ESCAPE_ESCAPE_SEQ);
246                    break;
247                case Stomp.BREAK:
248                    stream.write(Stomp.NEWLINE_ESCAPE_SEQ);
249                    break;
250                case Stomp.COLON:
251                    stream.write(Stomp.COLON_ESCAPE_SEQ);
252                    break;
253                default:
254                    stream.write(val);
255                }
256            }
257            result =  new String(stream.toByteArray(), "UTF-8");
258        }
259
260        return result;
261    }
262
263    private String decodeHeader(InputStream header) throws IOException {
264        ByteArrayOutputStream decoded = new ByteArrayOutputStream();
265        PushbackInputStream stream = new PushbackInputStream(header);
266
267        int value = -1;
268        while( (value = stream.read()) != -1) {
269            if (value == 92) {
270
271                int next = stream.read();
272                if (next != -1) {
273                    switch(next) {
274                    case 110:
275                        decoded.write(Stomp.BREAK);
276                        break;
277                    case 99:
278                        decoded.write(Stomp.COLON);
279                        break;
280                    case 92:
281                        decoded.write(Stomp.ESCAPE);
282                        break;
283                    default:
284                        stream.unread(next);
285                        decoded.write(value);
286                    }
287                } else {
288                    decoded.write(value);
289                }
290
291            } else {
292                decoded.write(value);
293            }
294        }
295
296        return new String(decoded.toByteArray(), "UTF-8");
297    }
298
299    public int getVersion() {
300        return version;
301    }
302
303    public void setVersion(int version) {
304        this.version = version;
305    }
306
307    public boolean isEncodingEnabled() {
308        return this.encodingEnabled;
309    }
310
311    public void setEncodingEnabled(boolean value) {
312        this.encodingEnabled = value;
313    }
314
315}