From e3587fcbb17d17f032c94e64ddc4c8baa8970aeb Mon Sep 17 00:00:00 2001
From: James Coglan <jcoglan@gmail.com>
Date: Mon, 2 Dec 2013 21:55:43 +0000
Subject: [PATCH] Add a max_length option to the Hybi/Client class.

---
 README.md                          |  2 ++
 lib/websocket/driver/hybi.rb       | 17 ++++++++++++-----
 spec/websocket/driver/hybi_spec.rb |  7 +++++++
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index d8ba03a..150c78f 100644
--- a/README.md
+++ b/README.md
@@ -212,6 +212,8 @@ enabled on outgoing frames.
 The `options` argument is optional, and is a hash. It may contain the following
 keys:
 
+* `:max_length` - the maximum allowed size of incoming message frames, in bytes.
+  The default value is `2^30 - 1`, or 1 byte short of 1GiB.
 * `:protocols` - an array of strings representing acceptable subprotocols for
   use over the socket. The driver will negotiate one of these to use via the
   `Sec-WebSocket-Protocol` header if supported by the other peer.
diff --git a/lib/websocket/driver/hybi.rb b/lib/websocket/driver/hybi.rb
index 9ba549f..7089f0f 100644
--- a/lib/websocket/driver/hybi.rb
+++ b/lib/websocket/driver/hybi.rb
@@ -32,6 +32,8 @@ module WebSocket
       FRAGMENTED_OPCODES = OPCODES.values_at(:continuation, :text, :binary)
       OPENING_OPCODES    = OPCODES.values_at(:text, :binary)
 
+      MAX_LENGTH = 0x3fffffff
+
       ERRORS = {
         :normal_closure       => 1000,
         :going_away           => 1001,
@@ -52,11 +54,12 @@ module WebSocket
         super
         reset
 
-        @reader    = StreamReader.new
-        @stage     = 0
-        @masking   = options[:masking]
-        @protocols = options[:protocols] || []
-        @protocols = @protocols.strip.split(/\s*,\s*/) if String === @protocols
+        @reader     = StreamReader.new
+        @stage      = 0
+        @masking    = options[:masking]
+        @protocols  = options[:protocols] || []
+        @protocols  = @protocols.strip.split(/\s*,\s*/) if String === @protocols
+        @max_length = options[:max_length] || MAX_LENGTH
 
         @require_masking = options[:require_masking]
         @ping_callbacks  = {}
@@ -280,6 +283,10 @@ module WebSocket
           return fail(:protocol_error, "Received control frame having too long payload: #{@length}")
         end
 
+        if @length > @max_length
+          return fail(:too_large, 'WebSocket frame length too large')
+        end
+
         @stage  = @masked ? 3 : 4
       end
 
diff --git a/spec/websocket/driver/hybi_spec.rb b/spec/websocket/driver/hybi_spec.rb
index 9f1d4ab..e9c0e53 100644
--- a/spec/websocket/driver/hybi_spec.rb
+++ b/spec/websocket/driver/hybi_spec.rb
@@ -287,6 +287,13 @@ describe WebSocket::Driver::Hybi do
         @message.should == "Hello" * 40
       end
 
+      it "returns an error for too-large frames" do
+        driver.parse [0x81, 0x7f, 0x00, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00]
+        @error.message.should == "WebSocket frame length too large"
+        @close.should == [1009, "WebSocket frame length too large"]
+        driver.state.should == :closed
+      end
+
       it "parses masked medium-length text frames" do
         driver.parse [0x81, 0xfe, 0x00, 0xc8] + mask + mask_message(*([0x48, 0x65, 0x6c, 0x6c, 0x6f] * 40))
         @message.should == "Hello" * 40