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 control, 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 # This is only compatible with concurrency models that expose a streaming
  22 # "rack.input" to the Rack application.  Thus it is NOT compatible with
  23 # any of the following as they fully buffer the request body before
  24 # the application dispatch:
  25 #
  26 # * :Coolio
  27 # * :CoolioThreadPool
  28 # * :CoolioThreadSpawn
  29 # * :Epoll
  30 # * :EventMachine
  31 # * :NeverBlock
  32 # * :Rev
  33 # * :RevThreadPool
  34 # * :RevThreadSpawn
  35 # * :XEpoll
  36 #
  37 # However, the global Rainbows::Configurator#client_max_body_size is compatible
  38 # with all concurrency models \Rainbows! supports.
  39 class Rainbows::MaxBody
  40
  41   # This is automatically called when used with Rack::Builder#use
  42   def initialize(app, limit = nil)
  43     case limit
  44     when Integer, nil
  45     else
  46       raise ArgumentError, "limit not an Integer"
  47     end
  48     @app, @limit = app, limit
  49   end
  50
  51   # :stopdoc:
  52   RACK_INPUT = "rack.input".freeze
  53   CONTENT_LENGTH = "CONTENT_LENGTH"
  54   HTTP_TRANSFER_ENCODING = "HTTP_TRANSFER_ENCODING"
  55
  56   # our main Rack middleware endpoint
  57   def call(env)
  58     @limit = Rainbows.server.client_max_body_size if nil == @limit
  59     catch(:rainbows_EFBIG) do
  60       len = env[CONTENT_LENGTH]
  61       if len && len.to_i > @limit
  62         return err
  63       elsif /\Achunked\z/i =~ env[HTTP_TRANSFER_ENCODING]
  64         limit_input!(env)
  65       end
  66       @app.call(env)
  67     end || err
  68   end
  69
  70   # this is called after forking, so it won't ever affect the master
  71   # if it's reconfigured
  72   def self.setup # :nodoc:
  73     Rainbows.server.client_max_body_size or return
  74     case Rainbows.server.use
  75     when :Rev, :Coolio, :EventMachine, :NeverBlock,
  76          :RevThreadSpawn, :RevThreadPool,
  77          :CoolioThreadSpawn, :CoolioThreadPool,
  78          :Epoll, :XEpoll
  79       return
  80     end
  81
  82     # force ourselves to the outermost middleware layer
  83     Rainbows.server.app = self.new(Rainbows.server.app)
  84   end
  85
  86   # Rack response returned when there's an error
  87   def err # :nodoc:
  88     [ 413, { 'Content-Length' => '0', 'Content-Type' => 'text/plain' }, [] ]
  89   end
  90
  91   def limit_input!(env)
  92     input = env[RACK_INPUT]
  93     klass = input.respond_to?(:rewind) ? RewindableWrapper : Wrapper
  94     env[RACK_INPUT] = klass.new(input, @limit)
  95   end
  96
  97   # :startdoc:
  98 end
  99 require 'rainbows/max_body/wrapper'
 100 require 'rainbows/max_body/rewindable_wrapper'