lib/rainbows/max_body.rb

   1 # -*- encoding: binary -*-
   2
   3 # Middleware used to enforce client_max_body_size for TeeInput users.
   4 #
   5 # There is no need to configure this middleware manually, it will
   6 # automatically be configured for you based on the client_max_body_size
   7 # setting.
   8 #
   9 # For more fine-grained conrol, you may also define it per-endpoint in
  10 # your Rack config.ru like this:
  11 #
  12 #        map "/limit_1M" do
  13 #          use Rainbows::MaxBody, 1024*1024
  14 #          run MyApp
  15 #        end
  16 #        map "/limit_10M" do
  17 #          use Rainbows::MaxBody, 1024*1024*10
  18 #          run MyApp
  19 #        end
  20
  21 class Rainbows::MaxBody
  22
  23
  24   # :call-seq:
  25   #   # in config.ru:
  26   #   use Rainbows::MaxBody, 4096
  27   #   run YourApplication.new
  28   def initialize(app, limit = Rainbows.max_bytes)
  29     Integer === limit or raise ArgumentError, "limit not an Integer"
  30     @app, @limit = app, limit
  31   end
  32
  33   # :stopdoc:
  34   RACK_INPUT = "rack.input".freeze
  35   CONTENT_LENGTH = "CONTENT_LENGTH"
  36   HTTP_TRANSFER_ENCODING = "HTTP_TRANSFER_ENCODING"
  37
  38   # our main Rack middleware endpoint
  39   def call(env)
  40     catch(:rainbows_EFBIG) do
  41       len = env[CONTENT_LENGTH]
  42       if len && len.to_i > @limit
  43         return err
  44       elsif /\Achunked\z/i =~ env[HTTP_TRANSFER_ENCODING]
  45         limit_input!(env)
  46       end
  47       @app.call(env)
  48     end || err
  49   end
  50
  51   # this is called after forking, so it won't ever affect the master
  52   # if it's reconfigured
  53   def self.setup # :nodoc:
  54     Rainbows.max_bytes or return
  55     case Rainbows.server.use
  56     when :Rev, :Coolio, :EventMachine, :NeverBlock,
  57          :RevThreadSpawn, :RevThreadPool,
  58          :CoolioThreadSpawn, :CoolioThreadPool,
  59          :Epoll
  60       return
  61     end
  62
  63     # force ourselves to the outermost middleware layer
  64     Rainbows.server.app = self.new(Rainbows.server.app)
  65   end
  66
  67   # Rack response returned when there's an error
  68   def err # :nodoc:
  69     [ 413, { 'Content-Length' => '0', 'Content-Type' => 'text/plain' }, [] ]
  70   end
  71
  72   def limit_input!(env)
  73     input = env[RACK_INPUT]
  74     klass = input.respond_to?(:rewind) ? RewindableWrapper : Wrapper
  75     env[RACK_INPUT] = klass.new(input, @limit)
  76   end
  77
  78   # :startdoc:
  79 end
  80 require 'rainbows/max_body/wrapper'
  81 require 'rainbows/max_body/rewindable_wrapper'