Remove ContinuationWebSocketFrame.aggregatedText()

Motivation:
Before we aggregated the full text in the WebSocket08FrameDecoder just to fill in the ContinuationWebSocketFrame.aggregatedText(). The problem was that there was no upper-limit and so it would be possible to see an OOME if the remote peer sends a TextWebSocketFrame + a never ending stream of ContinuationWebSocketFrames. Furthermore the aggregation does not really belong in the WebSocket08FrameDecoder, as we provide an extra ChannelHandler for this anyway (WebSocketFrameAggregator).

Modification:
Remove the ContinuationWebSocketFrame.aggregatedText() method and corresponding constructor. Also refactored WebSocket08FrameDecoder a bit to me more efficient which is now possible as we not need to aggregate here.

Result:
No more risk of OOME because of frames.

Author: Norman Maurer

src/main/java/org/jboss/netty/handler/codec/http/websocketx/ContinuationWebSocketFrame.java

@@ -25,8 +25,6 @@
  */
 public class ContinuationWebSocketFrame extends WebSocketFrame {
 
-    private String aggregatedText;
-
     /**
      * Creates a new empty continuation frame.
      */
@@ -61,26 +59,6 @@ public ContinuationWebSocketFrame(boolean finalFragment, int rsv, ChannelBuffer
         setBinaryData(binaryData);
     }
 
-    /**
-     * Creates a new continuation frame with the specified binary data
-     *
-     * @param finalFragment
-     *            flag indicating if this frame is the final fragment
-     * @param rsv
-     *            reserved bits used for protocol extensions
-     * @param binaryData
-     *            the content of the frame.
-     * @param aggregatedText
-     *            Aggregated text set by decoder on the final continuation frame of a fragmented text message
-     */
-    public ContinuationWebSocketFrame(
-            boolean finalFragment, int rsv, ChannelBuffer binaryData, String aggregatedText) {
-        setFinalFragment(finalFragment);
-        setRsv(rsv);
-        setBinaryData(binaryData);
-        this.aggregatedText = aggregatedText;
-    }
-
     /**
      * Creates a new continuation frame with the specified text data
      *
@@ -125,16 +103,4 @@ public void setText(String text) {
     public String toString() {
         return getClass().getSimpleName() + "(data: " + getBinaryData() + ')';
     }
-
-    /**
-     * Aggregated text returned by decoder on the final continuation frame of a fragmented text message
-     */
-    public String getAggregatedText() {
-        return aggregatedText;
-    }
-
-    public void setAggregatedText(String aggregatedText) {
-        this.aggregatedText = aggregatedText;
-    }
-
 }

src/main/java/org/jboss/netty/handler/codec/http/websocketx/UTF8Exception.java

@@ -35,10 +35,12 @@
  */
 package org.jboss.netty.handler.codec.http.websocketx;
 
+import org.jboss.netty.handler.codec.frame.CorruptedFrameException;
+
 /**
  * Checks UTF8 bytes for validity before converting it into a string
  */
-final class UTF8Output {
+final class Utf8Validator {
     private static final int UTF8_ACCEPT = 0;
     private static final int UTF8_REJECT = 12;
 
@@ -63,39 +65,37 @@ final class UTF8Output {
     @SuppressWarnings("RedundantFieldInitialization")
     private int state = UTF8_ACCEPT;
     private int codep;
+    private boolean checking;
 
-    private final StringBuilder stringBuilder;
-
-    UTF8Output(byte[] bytes) {
-        stringBuilder = new StringBuilder(bytes.length);
-        write(bytes);
-    }
-
-    public void write(byte[] bytes) {
+    public void check(byte[] bytes) throws CorruptedFrameException {
+        checking = true;
         for (byte b : bytes) {
             write(b);
         }
     }
 
-    public void write(int b) {
+    private void write(int b) throws CorruptedFrameException {
         byte type = TYPES[b & 0xFF];
 
         codep = state != UTF8_ACCEPT ? b & 0x3f | codep << 6 : 0xff >> type & b;
 
         state = STATES[state + type];
 
-        if (state == UTF8_ACCEPT) {
-            stringBuilder.append((char) codep);
-        } else if (state == UTF8_REJECT) {
-            throw new UTF8Exception("bytes are not UTF-8");
+        if (state == UTF8_REJECT) {
+            throw new CorruptedFrameException("bytes are not UTF-8");
         }
     }
 
-    @Override
-    public String toString() {
+    public void finish() throws CorruptedFrameException {
+        checking = false;
+        codep = 0;
         if (state != UTF8_ACCEPT) {
-            throw new UTF8Exception("bytes are not UTF-8");
+            state = UTF8_ACCEPT;
+            throw new CorruptedFrameException("bytes are not UTF-8");
         }
-        return stringBuilder.toString();
+    }
+
+    public boolean isChecking() {
+        return checking;
     }
 }

src/main/java/org/jboss/netty/handler/codec/http/websocketx/UTF8Output.java renamed to src/main/java/org/jboss/netty/handler/codec/http/websocketx/Utf8Validator.java

@@ -79,7 +79,7 @@ public class WebSocket08FrameDecoder extends ReplayingDecoder<WebSocket08FrameDe
     private static final byte OPCODE_PING = 0x9;
     private static final byte OPCODE_PONG = 0xA;
 
-    private UTF8Output fragmentedFramesText;
+    private Utf8Validator utf8Validator;
     private int fragmentedFramesCount;
 
     private final long maxFramePayloadLength;
@@ -315,37 +315,33 @@ protected Object decode(ChannelHandlerContext ctx, Channel channel, ChannelBuffe
 
             // Processing for possible fragmented messages for text and binary
             // frames
-            String aggregatedText = null;
             if (frameFinalFlag) {
                 // Final frame of the sequence. Apparently ping frames are
                 // allowed in the middle of a fragmented message
                 if (frameOpcode != OPCODE_PING) {
                     fragmentedFramesCount = 0;
 
                     // Check text for UTF8 correctness
-                    if (frameOpcode == OPCODE_TEXT || fragmentedFramesText != null) {
+                    if (frameOpcode == OPCODE_TEXT || (utf8Validator != null && utf8Validator.isChecking())) {
                         // Check UTF-8 correctness for this payload
                         checkUTF8String(channel, framePayload.array());
 
                         // This does a second check to make sure UTF-8
                         // correctness for entire text message
-                        aggregatedText = fragmentedFramesText.toString();
-
-                        fragmentedFramesText = null;
+                        utf8Validator.finish();
                     }
                 }
             } else {
                 // Not final frame so we can expect more frames in the
                 // fragmented sequence
                 if (fragmentedFramesCount == 0) {
                     // First text or binary frame for a fragmented set
-                    fragmentedFramesText = null;
                     if (frameOpcode == OPCODE_TEXT) {
                         checkUTF8String(channel, framePayload.array());
                     }
                 } else {
                     // Subsequent frames - only check if init frame is text
-                    if (fragmentedFramesText != null) {
+                    if (utf8Validator != null && utf8Validator.isChecking()) {
                         checkUTF8String(channel, framePayload.array());
                     }
                 }
@@ -360,7 +356,7 @@ protected Object decode(ChannelHandlerContext ctx, Channel channel, ChannelBuffe
             } else if (frameOpcode == OPCODE_BINARY) {
                 return new BinaryWebSocketFrame(frameFinalFlag, frameRsv, framePayload);
             } else if (frameOpcode == OPCODE_CONT) {
-                return new ContinuationWebSocketFrame(frameFinalFlag, frameRsv, framePayload, aggregatedText);
+                return new ContinuationWebSocketFrame(frameFinalFlag, frameRsv, framePayload);
             } else {
                 throw new UnsupportedOperationException("Cannot decode web socket frame with opcode: " + frameOpcode);
             }
@@ -382,11 +378,15 @@ private void unmask(ChannelBuffer frame) {
     }
 
     private void protocolViolation(Channel channel, String reason) throws CorruptedFrameException {
+        protocolViolation(channel, new CorruptedFrameException(reason));
+    }
+
+    private void protocolViolation(Channel channel, CorruptedFrameException ex) throws CorruptedFrameException {
         checkpoint(State.CORRUPT);
         if (channel.isConnected()) {
             channel.write(ChannelBuffers.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
         }
-        throw new CorruptedFrameException(reason);
+        throw ex;
     }
 
     private static int toFrameLength(long l) throws TooLongFrameException {
@@ -399,20 +399,12 @@ private static int toFrameLength(long l) throws TooLongFrameException {
 
     private void checkUTF8String(Channel channel, byte[] bytes) throws CorruptedFrameException {
         try {
-            // StringBuilder sb = new StringBuilder("UTF8 " + bytes.length +
-            // " bytes: ");
-            // for (byte b : bytes) {
-            // sb.append(Integer.toHexString(b)).append(" ");
-            // }
-            // logger.debug(sb.toString());
-
-            if (fragmentedFramesText == null) {
-                fragmentedFramesText = new UTF8Output(bytes);
-            } else {
-                fragmentedFramesText.write(bytes);
+            if (utf8Validator == null) {
+                utf8Validator = new Utf8Validator();
             }
-        } catch (UTF8Exception ex) {
-            protocolViolation(channel, "invalid UTF-8 bytes");
+            utf8Validator.check(bytes);
+        } catch (CorruptedFrameException ex) {
+            protocolViolation(channel, ex);
         }
     }
 
@@ -440,9 +432,10 @@ protected void checkCloseFrameBody(Channel channel, ChannelBuffer buffer) throws
             byte[] b = new byte[buffer.readableBytes()];
             buffer.readBytes(b);
             try {
-                new UTF8Output(b);
-            } catch (UTF8Exception ex) {
-                protocolViolation(channel, "Invalid close frame reason text. Invalid UTF-8 bytes");
+                Utf8Validator validator = new Utf8Validator();
+                validator.check(b);
+            } catch (CorruptedFrameException ex) {
+                protocolViolation(channel, ex);
             }
         }
 

src/main/java/org/jboss/netty/handler/codec/http/websocketx/WebSocket08FrameDecoder.java

@@ -0,0 +1,101 @@
+/*
+ * Copyright 2013 The Netty Project
+ *
+ * The Netty Project licenses this file to you under the Apache License,
+ * version 2.0 (the "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at:
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+package org.jboss.netty.handler.codec.http.websocketx;
+
+
+import org.jboss.netty.buffer.ChannelBuffer;
+import org.jboss.netty.buffer.ChannelBuffers;
+import org.jboss.netty.channel.Channel;
+import org.jboss.netty.channel.ChannelHandlerContext;
+import org.jboss.netty.handler.codec.frame.TooLongFrameException;
+import org.jboss.netty.handler.codec.oneone.OneToOneDecoder;
+
+/**
+ * Handler that aggregate fragmented WebSocketFrame's.
+ *
+ * Be aware if PING/PONG/CLOSE frames are send in the middle of a fragmented {@link WebSocketFrame} they will
+ * just get forwarded to the next handler in the pipeline.
+ */
+public class WebSocketFrameAggregator extends OneToOneDecoder {
+    private final int maxFrameSize;
+    private WebSocketFrame currentFrame;
+    private boolean tooLongFrameFound;
+
+    /**
+     * Construct a new instance
+     *
+     * @param maxFrameSize      If the size of the aggregated frame exceeds this value,
+     *                          a {@link TooLongFrameException} is thrown.
+     */
+    public WebSocketFrameAggregator(int maxFrameSize) {
+        if (maxFrameSize < 1) {
+            throw new IllegalArgumentException("maxFrameSize must be > 0");
+        }
+        this.maxFrameSize = maxFrameSize;
+    }
+
+    @Override
+    protected Object decode(ChannelHandlerContext ctx, Channel channel, Object message) throws Exception {
+        if (!(message instanceof WebSocketFrame)) {
+            return message;
+        }
+        WebSocketFrame msg = (WebSocketFrame) message;
+        if (currentFrame == null) {
+            tooLongFrameFound = false;
+            if (msg.isFinalFragment()) {
+                return msg;
+            }
+            ChannelBuffer buf = msg.getBinaryData();
+
+            if (msg instanceof TextWebSocketFrame) {
+                currentFrame = new TextWebSocketFrame(true, msg.getRsv(), buf);
+            } else if (msg instanceof BinaryWebSocketFrame) {
+                currentFrame = new BinaryWebSocketFrame(true, msg.getRsv(), buf);
+            } else {
+                throw new IllegalStateException(
+                        "WebSocket frame was not of type TextWebSocketFrame or BinaryWebSocketFrame");
+            }
+            return null;
+        }
+        if (msg instanceof ContinuationWebSocketFrame) {
+            if (tooLongFrameFound) {
+                if (msg.isFinalFragment()) {
+                    currentFrame = null;
+                }
+                return null;
+            }
+            ChannelBuffer content = currentFrame.getBinaryData();
+            if (content.readableBytes() > maxFrameSize - msg.getBinaryData().readableBytes()) {
+                tooLongFrameFound = true;
+                throw new TooLongFrameException(
+                        "WebSocketFrame length exceeded " + content +
+                                " bytes.");
+            }
+            currentFrame.setBinaryData(ChannelBuffers.wrappedBuffer(content, msg.getBinaryData()));
+
+            if (msg.isFinalFragment()) {
+                WebSocketFrame currentFrame = this.currentFrame;
+                this.currentFrame = null;
+                return currentFrame;
+            } else {
+                return null;
+            }
+        }
+        // It is possible to receive CLOSE/PING/PONG frames during fragmented frames so just pass them to the next
+        // handler in the chain
+        return msg;
+    }
+}