summary refs log tree commit homepage
path: root/lib/rainbows/max_body
diff options
context:
space:
mode:
authorEric Wong <normalperson@yhbt.net>2010-11-19 10:19:45 +0000
committerEric Wong <normalperson@yhbt.net>2010-11-19 16:59:01 -0800
commitc6ffae22748bc22d5ef88fea2a3ca67f480ee74b (patch)
treed64947098657f2bbdbca04a6db2e43645060a223 /lib/rainbows/max_body
parent3cee07d750f678af92318c14110c803be3f9b97f (diff)
To avoid denial-of-service attacks, the wrappers need to
intercept requests *before* they hit the memory allocator, so we
need to reimplement the read(all) and gets cases to use
smaller buffers whenever the application does not specify one.
Diffstat (limited to 'lib/rainbows/max_body')
-rw-r--r--lib/rainbows/max_body/rewindable_wrapper.rb1
-rw-r--r--lib/rainbows/max_body/wrapper.rb66
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