6 files changed, 158 insertions, 1 deletions
diff --git a/lib/rainbows.rb b/lib/rainbows.rb
index ccf211e..ad4e564 100644
--- a/lib/rainbows.rb
+++ b/lib/rainbows.rb
@@ -31,6 +31,7 @@ module Rainbows
    require 'rainbows/base'
    autoload :AppPool, 'rainbows/app_pool'
    autoload :DevFdResponse, 'rainbows/dev_fd_response'
+  autoload :MaxBody, 'rainbows/max_body'
  
    class << self
  
@@ -81,6 +82,12 @@ module Rainbows
        io.respond_to?(:peeraddr) ?
                          io.peeraddr.last : Unicorn::HttpRequest::LOCALHOST
      end
+
+    # the default max body size is 1 megabyte (1024 * 1024 bytes)
+    @@max_bytes = 1024 * 1024
+
+    def max_bytes; @@max_bytes; end
+    def max_bytes=(nr); @@max_bytes = nr; end
    end
  
    # configures \Rainbows! with a given concurrency model to +use+ and
@@ -91,6 +98,7 @@ module Rainbows
    #     use :Revactor # this may also be :ThreadSpawn or :ThreadPool
    #     worker_connections 400
    #     keepalive_timeout 0 # zero disables keepalives entirely
+  #     client_max_body_size 5*1024*1024 # 5 megabytes
    #   end
    #
    #   # the rest of the Unicorn configuration
@@ -107,6 +115,10 @@ module Rainbows
    # start retrieving extra elements for.  Increasing this beyond 5
    # seconds is not recommended.  Zero disables keepalive entirely
    # (but pipelining fully-formed requests is still works).
+  #
+  # The default +client_max_body_size+ is 1 megabyte (1024 * 1024 bytes),
+  # setting this to +nil+ will disable body size checks and allow any
+  # size to be specified.
    def Rainbows!(&block)
      block_given? or raise ArgumentError, "Rainbows! requires a block"
      HttpServer.setup(block)
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