Correctly count the key space size for nested param queries.

4 files changed, 59 insertions, 41 deletions

diff --git a/lib/rack/multipart/parser.rb b/lib/rack/multipart/parser.rb
index 2b55cf94..c2973205 100644
--- a/lib/rack/multipart/parser.rb
+++ b/lib/rack/multipart/parser.rb
@@ -14,9 +14,6 @@ module Rack
 
         fast_forward_to_first_boundary
 
-        max_key_space = Utils.key_space_limit
-        bytes = 0
-
         loop do
           head, filename, content_type, name, body =
             get_current_head_and_filename_and_content_type_and_name_and_body
@@ -31,13 +28,6 @@ module Rack
 
           filename, data = get_data(filename, body, content_type, name, head)
 
-          if name
-            bytes += name.size
-            if bytes > max_key_space
-              raise RangeError, "exceeded available parameter key space"
-            end
-          end
-
           Utils.normalize_params(@params, name, data) unless data.nil?
 
           # break if we're at the end of a buffer, but not if it is the end of a field
@@ -46,7 +36,7 @@ module Rack
 
         @io.rewind
 
-        @params
+        @params.to_params_hash
       end
 
       private
@@ -56,7 +46,7 @@ module Rack
         @boundary = "--#{$1}"
 
         @buf = ""
-        @params = {}
+        @params = Utils::KeySpaceConstrainedParams.new
 
         @content_length = @env['CONTENT_LENGTH'].to_i
         @io = @env['rack.input']
diff --git a/lib/rack/utils.rb b/lib/rack/utils.rb
index 7bceb458..c1f71c0e 100644
--- a/lib/rack/utils.rb
+++ b/lib/rack/utils.rb
@@ -61,21 +61,11 @@ module Rack
     # cookies by changing the characters used in the second
     # parameter (which defaults to '&;').
     def parse_query(qs, d = nil)
-      params = {}
-
-      max_key_space = Utils.key_space_limit
-      bytes = 0
+      params = KeySpaceConstrainedParams.new
 
       (qs || '').split(d ? /[#{d}] */n : DEFAULT_SEP).each do |p|
         k, v = p.split('=', 2).map { |x| unescape(x) }
 
-        if k
-          bytes += k.size
-          if bytes > max_key_space
-            raise RangeError, "exceeded available parameter key space"
-          end
-        end
-
         if cur = params[k]
           if cur.class == Array
             params[k] << v
@@ -87,30 +77,20 @@ module Rack
         end
       end
 
-      return params
+      return params.to_params_hash
     end
     module_function :parse_query
 
     def parse_nested_query(qs, d = nil)
-      params = {}
-
-      max_key_space = Utils.key_space_limit
-      bytes = 0
+      params = KeySpaceConstrainedParams.new
 
       (qs || '').split(d ? /[#{d}] */n : DEFAULT_SEP).each do |p|
         k, v = p.split('=', 2).map { |s| unescape(s) }
 
-        if k
-          bytes += k.size
-          if bytes > max_key_space
-            raise RangeError, "exceeded available parameter key space"
-          end
-        end
-
         normalize_params(params, k, v)
       end
 
-      return params
+      return params.to_params_hash
     end
     module_function :parse_nested_query
 
@@ -131,14 +111,14 @@ module Rack
         child_key = $1
         params[k] ||= []
         raise TypeError, "expected Array (got #{params[k].class.name}) for param `#{k}'" unless params[k].is_a?(Array)
-        if params[k].last.is_a?(Hash) && !params[k].last.key?(child_key)
+        if params[k].last.is_a?(KeySpaceConstrainedParams) && !params[k].last.key?(child_key)
           normalize_params(params[k].last, child_key, v)
         else
-          params[k] << normalize_params({}, child_key, v)
+          params[k] << normalize_params(KeySpaceConstrainedParams.new, child_key, v)
         end
       else
-        params[k] ||= {}
-        raise TypeError, "expected Hash (got #{params[k].class.name}) for param `#{k}'" unless params[k].is_a?(Hash)
+        params[k] ||= KeySpaceConstrainedParams.new
+        raise TypeError, "expected Hash (got #{params[k].class.name}) for param `#{k}'" unless params[k].is_a?(KeySpaceConstrainedParams)
         params[k] = normalize_params(params[k], after, v)
       end
 
@@ -445,6 +425,41 @@ module Rack
       end
     end
 
+    class KeySpaceConstrainedParams
+      def initialize(limit = Utils.key_space_limit)
+        @limit  = limit
+        @size   = 0
+        @params = {}
+      end
+
+      def [](key)
+        @params[key]
+      end
+
+      def []=(key, value)
+        @size += key.size unless @params.key?(key)
+        raise RangeError, 'exceeded available parameter key space' if @size > @limit
+        @params[key] = value
+      end
+
+      def key?(key)
+        @params.key?(key)
+      end
+
+      def to_params_hash
+        hash = @params
+        hash.keys.each do |key|
+          value = hash[key]
+          if value.kind_of?(self.class)
+            hash[key] = value.to_params_hash
+          elsif value.kind_of?(Array)
+            value.map! {|x| x.kind_of?(self.class) ? x.to_params_hash : x}
+          end
+        end
+        hash
+      end
+    end
+
     # Every standard HTTP code mapped to the appropriate message.
     # Generated with:
     #   curl -s http://www.iana.org/assignments/http-status-codes | \
diff --git a/test/spec_request.rb b/test/spec_request.rb
index d20585c9..57d34abc 100644
--- a/test/spec_request.rb
+++ b/test/spec_request.rb
@@ -137,6 +137,19 @@ describe Rack::Request do
     end
   end
 
+  should "limit the key size per nested params hash" do
+    nested_query = Rack::MockRequest.env_for("/?foo[bar][baz][qux]=1")
+    plain_query  = Rack::MockRequest.env_for("/?foo_bar__baz__qux_=1")
+
+    old, Rack::Utils.key_space_limit = Rack::Utils.key_space_limit, 3
+    begin
+      lambda { Rack::Request.new(nested_query).GET }.should.not.raise(RangeError)
+      lambda { Rack::Request.new(plain_query).GET  }.should.raise(RangeError)
+    ensure
+      Rack::Utils.key_space_limit = old
+    end
+  end
+
   should "not unify GET and POST when calling params" do
     mr = Rack::MockRequest.env_for("/?foo=quux",
       "REQUEST_METHOD" => 'POST',
diff --git a/test/spec_utils.rb b/test/spec_utils.rb
index 069e2299..32d0a6f5 100644
--- a/test/spec_utils.rb
+++ b/test/spec_utils.rb
@@ -188,7 +188,7 @@ describe Rack::Utils do
 
     lambda { Rack::Utils.parse_nested_query("x[y]=1&x[]=1") }.
       should.raise(TypeError).
-      message.should.equal "expected Array (got Hash) for param `x'"
+      message.should.match /expected Array \(got [^)]*\) for param `x'/
 
     lambda { Rack::Utils.parse_nested_query("x[y]=1&x[y][][w]=2") }.
       should.raise(TypeError).