5 files changed, 146 insertions, 1 deletions
diff --git a/lib/rainbows/base.rb b/lib/rainbows/base.rb
index 0cbc711..864b847 100644
--- a/lib/rainbows/base.rb
+++ b/lib/rainbows/base.rb
@@ -12,6 +12,7 @@ module Rainbows
  
      def init_worker_process(worker)
        super(worker)
+      MaxBody.setup
        G.tmp = worker.tmp
  
        # avoid spurious wakeups and blocking-accept() with 1.8 green threads
diff --git a/lib/rainbows/const.rb b/lib/rainbows/const.rb
index 08c4821..42906d3 100644
--- a/lib/rainbows/const.rb
+++ b/lib/rainbows/const.rb
@@ -24,5 +24,7 @@ module Rainbows
      # of the official spec, but for now it is "hack.io"
      CLIENT_IO = "hack.io".freeze
  
+    ERROR_413_RESPONSE = "HTTP/1.1 413 Request Entity Too Large\r\n\r\n"
+
    end
  end
diff --git a/lib/rainbows/ev_core.rb b/lib/rainbows/ev_core.rb
index 682bdd6..d4ad040 100644
--- a/lib/rainbows/ev_core.rb
+++ b/lib/rainbows/ev_core.rb
@@ -49,7 +49,7 @@ module Rainbows
              write(EXPECT_100_RESPONSE)
              @env.delete(HTTP_EXPECT)
            end
-          @input = len && len <= MAX_BODY ? StringIO.new("") : Util.tmpio
+          @input = CapInput.new(len, self)
            @hp.filter_body(@buf2 = "", @buf)
            @input << @buf2
            on_read("")
@@ -73,5 +73,45 @@ module Rainbows
          handle_error(e)
      end
  
+    class CapInput < Struct.new(:io, :client, :bytes_left)
+      MAX_BODY = Unicorn::Const::MAX_BODY
+      Util = Unicorn::Util
+
+      def self.err(client, msg)
+        client.write(Const::ERROR_413_RESPONSE)
+        client.quit
+
+        # zip back up the stack
+        raise IOError, msg, []
+      end
+
+      def self.new(len, client)
+        max = Rainbows.max_bytes
+        if len
+          if max && (len > max)
+            err(client, "Content-Length too big: #{len} > #{max}")
+          end
+          len <= MAX_BODY ? StringIO.new("") : Util.tmpio
+        else
+          max ? super(Util.tmpio, client, max) : Util.tmpio
+        end
+      end
+
+      def <<(buf)
+        if (self.bytes_left -= buf.size) < 0
+          io.close
+          CapInput.err(client, "chunked request body too big")
+        end
+        io << buf
+      end
+
+      def gets; io.gets; end
+      def each(&block); io.each(&block); end
+      def size; io.size; end
+      def rewind; io.rewind; end
+      def read(*args); io.read(*args); end
+
+    end
+
    end
  end
diff --git a/lib/rainbows/http_server.rb b/lib/rainbows/http_server.rb
index ea2e23f..50231ff 100644
--- a/lib/rainbows/http_server.rb
+++ b/lib/rainbows/http_server.rb
@@ -85,6 +85,18 @@ module Rainbows
          raise ArgumentError, "keepalive must be a non-negative Integer"
        G.kato = nr
      end
+
+    def client_max_body_size(nr)
+      err = "client_max_body_size must be nil or a non-negative Integer"
+      case nr
+      when nil
+      when Integer
+        nr >= 0 or raise ArgumentError, err
+      else
+        raise ArgumentError, err
+      end
+      Rainbows.max_bytes = nr
+    end
    end
  
  end
diff --git a/lib/rainbows/max_body.rb b/lib/rainbows/max_body.rb
new file mode 100644
index 0000000..7450b2a
--- /dev/null
+++ b/lib/rainbows/max_body.rb
@@ -0,0 +1,90 @@
+# -*- encoding: binary -*-
+module Rainbows
+
+# middleware used to enforce client_max_body_size for TeeInput users,
+# there is no need to configure this middleware manually, it will
+# automatically be configured for you based on the client_max_body_size
+# setting
+class MaxBody < Struct.new(:app)
+
+  # this is meant to be included in Unicorn::TeeInput (and derived
+  # classes) to limit body sizes
+  module Limit
+    Util = Unicorn::Util
+
+    def initialize(socket, req, parser, buf)
+      self.len = parser.content_length
+
+      max = Rainbows.max_bytes # never nil, see MaxBody.setup
+      if len && len > max
+        socket.write(Const::ERROR_413_RESPONSE)
+        socket.close
+        raise IOError, "Content-Length too big: #{len} > #{max}", []
+      end
+
+      self.req = req
+      self.parser = parser
+      self.buf = buf
+      self.socket = socket
+      self.buf2 = ""
+      if buf.size > 0
+        parser.filter_body(buf2, buf) and finalize_input
+        buf2.size > max and raise IOError, "chunked request body too big", []
+      end
+      self.tmp = len && len < Const::MAX_BODY ? StringIO.new("") : Util.tmpio
+      if buf2.size > 0
+        tmp.write(buf2)
+        tmp.seek(0)
+        max -= buf2.size
+      end
+      @max_body = max
+    end
+
+    def tee(length, dst)
+      rv = _tee(length, dst)
+      if rv && ((@max_body -= rv.size) < 0)
+        $stderr.puts "#@max_body  TOO SMALL"
+        # make HttpParser#keepalive? => false to force an immediate disconnect
+        # after we write
+        parser.reset
+        throw :rainbows_EFBIG
+      end
+      rv
+    end
+
+  end
+
+  # this is called after forking, so it won't ever affect the master
+  # if it's reconfigured
+  def self.setup
+    Rainbows.max_bytes or return
+    case G.server.use
+    when :Rev, :EventMachine, :NeverBlock
+      return
+    when :Revactor
+      Rainbows::Revactor::TeeInput
+    else
+      Unicorn::TeeInput
+    end.class_eval do
+      alias _tee tee # can't use super here :<
+      remove_method :tee
+      remove_method :initialize if G.server.use != :Revactor # FIXME CODE SMELL
+      include Limit
+    end
+
+    # force ourselves to the outermost middleware layer
+    G.server.app = MaxBody.new(G.server.app)
+  end
+
+  # Rack response returned when there's an error
+  def err(env)
+    [ 413, [ %w(Content-Length 0), %w(Content-Type text/plain) ], [] ]
+  end
+
+  # our main Rack middleware endpoint
+  def call(env)
+    catch(:rainbows_EFBIG) { app.call(env) } || err(env)
+  end
+
+end # class
+end # module