max_body.rb in rainbows-2.0.0

- old
+ new
@@ -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
+    when :Rev, :EventMachine, :NeverBlock, :RevThreadSpawn, :RevThreadPool
       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'