summary refs log tree commit homepage
path: root/lib/rainbows/max_body.rb
diff options
context:
space:
mode:
authorEric Wong <normalperson@yhbt.net>2010-11-16 16:16:42 -0800
committerEric Wong <normalperson@yhbt.net>2010-11-16 16:21:33 -0800
commit31cf77e7aa2f2e6065e7ace44e55c3f042b51f1b (patch)
tree7dfd53c499cbf7a52165d3480d56d8576a743781 /lib/rainbows/max_body.rb
parent42747db815ad668b20849afb2a9dcdd1319713ae (diff)
This allows the client_max_body_size implementation to not rely
on Unicorn::TeeInput internals, allowing it to be used with
Unicorn::StreamInput (or any other (nearly)
Rack::Lint-compatible input object).
Diffstat (limited to 'lib/rainbows/max_body.rb')
-rw-r--r--lib/rainbows/max_body.rb103
1 files changed, 50 insertions, 53 deletions
diff --git a/lib/rainbows/max_body.rb b/lib/rainbows/max_body.rb
index 223c57e..5e2cb64 100644
--- a/lib/rainbows/max_body.rb
+++ b/lib/rainbows/max_body.rb
@@ -1,81 +1,78 @@
 # -*- encoding: binary -*-
-# :enddoc:
 
-# middleware used to enforce client_max_body_size for TeeInput users,
-# there is no need to configure this middleware manually, it will
+# 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 Rainbows::MaxBody < Struct.new(:app)
+# setting.
+#
+# For more fine-grained conrol, you may also define it per-endpoint in
+# your Rack config.ru like this:
+#
+#        map "/limit_1M" do
+#          use Rainbows::MaxBody, 1024*1024
+#          run MyApp
+#        end
+#        map "/limit_10M" do
+#          use Rainbows::MaxBody, 1024*1024*10
+#          run MyApp
+#        end
 
-  # this is meant to be included in Rainbows::TeeInput (and derived
-  # classes) to limit body sizes
-  module Limit
-    TmpIO = Unicorn::TmpIO
-    MAX_BODY = Rainbows::Const::MAX_BODY
+class Rainbows::MaxBody
 
-    def initialize(socket, request)
-      @parser = request
-      @buf = request.buf
-      @env = request.env
-      @len = request.content_length
 
-      max = Rainbows.max_bytes # never nil, see MaxBody.setup
-      if @len && @len > max
-        socket.write(Rainbows::Const::ERROR_413_RESPONSE)
-        socket.close
-        raise IOError, "Content-Length too big: #@len > #{max}", []
-      end
+  # :call-seq:
+  #   # in config.ru:
+  #   use Rainbows::MaxBody, 4096
+  #   run YourApplication.new
+  def initialize(app, limit = Rainbows.max_bytes)
+    Integer === limit or raise ArgumentError, "limit not an Integer"
+    @app, @limit = app, limit
+  end
 
-      @socket = socket
-      @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
-      @tmp = @len && @len < MAX_BODY ? StringIO.new("") : TmpIO.new
-      if @buf2.size > 0
-        @tmp.write(@buf2)
-        @tmp.rewind
-        max -= @buf2.size
-      end
-      @max_body = max
-    end
+  # :stopdoc:
+  RACK_INPUT = "rack.input".freeze
+  CONTENT_LENGTH = "CONTENT_LENGTH"
+  HTTP_TRANSFER_ENCODING = "HTTP_TRANSFER_ENCODING"
 
-    def tee(length, dst)
-      rv = super
-      if rv && ((@max_body -= rv.size) < 0)
-        # make HttpParser#keepalive? => false to force an immediate disconnect
-        # after we write
-        @parser.reset
-        throw :rainbows_EFBIG
+  # our main Rack middleware endpoint
+  def call(env)
+    catch(:rainbows_EFBIG) do
+      len = env[CONTENT_LENGTH]
+      if len && len.to_i > @limit
+        return err
+      elsif /\Achunked\z/i =~ env[HTTP_TRANSFER_ENCODING]
+        limit_input!(env)
       end
-      rv
-    end
-
+      @app.call(env)
+    end || err
   end
 
   # this is called after forking, so it won't ever affect the master
   # if it's reconfigured
-  def self.setup
+  def self.setup # :nodoc:
     Rainbows.max_bytes or return
     case Rainbows::G.server.use
     when :Rev, :EventMachine, :NeverBlock
       return
     end
 
-    Rainbows::TeeInput.__send__(:include, Limit)
-
     # force ourselves to the outermost middleware layer
     Rainbows::G.server.app = self.new(Rainbows::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) ], [] ]
+  def err # :nodoc:
+    [ 413, { 'Content-Length' => '0', 'Content-Type' => 'text/plain' }, [] ]
   end
 
-  # our main Rack middleware endpoint
-  def call(env)
-    catch(:rainbows_EFBIG) { app.call(env) } || err(env)
+  def limit_input!(env)
+    input = env[RACK_INPUT]
+    klass = input.respond_to?(:rewind) ? RewindableWrapper : Wrapper
+    env[RACK_INPUT] = klass.new(input, @limit)
   end
+
+  # :startdoc:
 end
+require 'rainbows/max_body/wrapper'
+require 'rainbows/max_body/rewindable_wrapper'