add client_max_body_size config directive

Since Rainbows! is supported when exposed directly to the
Internet, administrators may want to limit the amount of data a
user may upload in a single request body to prevent a
denial-of-service via disk space exhaustion.

This amount may be specified in bytes, the default limit being
1024*1024 bytes (1 megabyte).  To override this default, a user
may specify `client_max_body_size' in the Rainbows! block
of their server config file:

  Rainbows! do
    client_max_body_size 10 * 1024 * 1024
  end

Clients that exceed the limit will get a "413 Request Entity Too
Large" response if the request body is too large and the
connection will close.

For chunked requests, we have no choice but to interrupt during
the client upload since we have no prior knowledge of the
request body size.

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