2 files changed, 56 insertions, 11 deletions
diff --git a/lib/rainbows/max_body/rewindable_wrapper.rb b/lib/rainbows/max_body/rewindable_wrapper.rb
index b52726e..5693ead 100644
--- a/lib/rainbows/max_body/rewindable_wrapper.rb
+++ b/lib/rainbows/max_body/rewindable_wrapper.rb
@@ -8,6 +8,7 @@ class Rainbows::MaxBody::RewindableWrapper < Rainbows::MaxBody::Wrapper
  
    def rewind
      @limit = @orig_limit
+    @rbuf = ''
      @input.rewind
    end
  
diff --git a/lib/rainbows/max_body/wrapper.rb b/lib/rainbows/max_body/wrapper.rb
index 3c38ca6..15faeeb 100644
--- a/lib/rainbows/max_body/wrapper.rb
+++ b/lib/rainbows/max_body/wrapper.rb
@@ -1,26 +1,70 @@
  # -*- encoding: binary -*-
  # :enddoc:
+#
+# This is only used for chunked request bodies, which are rare
  class Rainbows::MaxBody::Wrapper
    def initialize(rack_input, limit)
-    @input, @limit = rack_input, limit
-  end
-
-  def check(rv)
-    throw :rainbows_EFBIG if rv && ((@limit -= rv.size) < 0)
-    rv
+    @input, @limit, @rbuf = rack_input, limit, ''
    end
  
    def each(&block)
-    while line = @input.gets
-      yield check(line)
+    while line = gets
+      yield line
      end
    end
  
-  def read(*args)
-    check(@input.read(*args))
+  # chunked encoding means this method behaves more like readpartial,
+  # since Rack does not support a method named "readpartial"
+  def read(length = nil, rv = '')
+    if length
+      if length <= @rbuf.size
+        length < 0 and raise ArgumentError, "negative length #{length} given"
+        rv.replace(@rbuf.slice!(0, length))
+      elsif @rbuf.empty?
+        checked_read(length, rv) or return
+      else
+        rv.replace(@rbuf.slice!(0, @rbuf.size))
+      end
+      rv.empty? && length != 0 ? nil : rv
+    else
+      rv.replace(read_all)
+    end
    end
  
    def gets
-    check(@input.gets)
+    sep = $/
+    if sep.nil?
+      rv = read_all
+      return rv.empty? ? nil : rv
+    end
+    re = /\A(.*?#{Regexp.escape(sep)})/
+
+    begin
+      @rbuf.sub!(re, '') and return $1
+
+      if tmp = checked_read(16384)
+        @rbuf << tmp
+      elsif @rbuf.empty? # EOF
+        return nil
+      else # EOF, return whatever is left
+        return @rbuf.slice!(0, @rbuf.size)
+      end
+    end while true
+  end
+
+  def checked_read(length = 16384, buf = '')
+    if @input.read(length, buf)
+      throw :rainbows_EFBIG if ((@limit -= buf.size) < 0)
+      return buf
+    end
+  end
+
+  def read_all
+    rv = @rbuf.slice!(0, @rbuf.size)
+    tmp = ''
+    while checked_read(16384, tmp)
+      rv << tmp
+    end
+    rv
    end
  end