lib/rainbows/max_body.rb

   1 # -*- encoding: binary -*-
   2 module Rainbows
   3
   4 # middleware used to enforce client_max_body_size for TeeInput users,
   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 class MaxBody < Struct.new(:app)
   9
  10   # this is meant to be included in Rainbows::TeeInput (and derived
  11   # classes) to limit body sizes
  12   module Limit
  13     Util = Unicorn::Util
  14
  15     def initialize(socket, req, parser, buf)
  16       self.len = parser.content_length
  17
  18       max = Rainbows.max_bytes # never nil, see MaxBody.setup
  19       if len && len > max
  20         socket.write(Const::ERROR_413_RESPONSE)
  21         socket.close
  22         raise IOError, "Content-Length too big: #{len} > #{max}", []
  23       end
  24
  25       self.req = req
  26       self.parser = parser
  27       self.buf = buf
  28       self.socket = socket
  29       self.buf2 = ""
  30       if buf.size > 0
  31         parser.filter_body(buf2, buf) and finalize_input
  32         buf2.size > max and raise IOError, "chunked request body too big", []
  33       end
  34       self.tmp = len && len < Const::MAX_BODY ? StringIO.new("") : Util.tmpio
  35       if buf2.size > 0
  36         tmp.write(buf2)
  37         tmp.seek(0)
  38         max -= buf2.size
  39       end
  40       @max_body = max
  41     end
  42
  43     def tee(length, dst)
  44       rv = super
  45       if rv && ((@max_body -= rv.size) < 0)
  46         # make HttpParser#keepalive? => false to force an immediate disconnect
  47         # after we write
  48         parser.reset
  49         throw :rainbows_EFBIG
  50       end
  51       rv
  52     end
  53
  54   end
  55
  56   # this is called after forking, so it won't ever affect the master
  57   # if it's reconfigured
  58   def self.setup
  59     Rainbows.max_bytes or return
  60     case G.server.use
  61     when :Rev, :EventMachine, :NeverBlock
  62       return
  63     end
  64
  65     TeeInput.class_eval { include Limit }
  66
  67     # force ourselves to the outermost middleware layer
  68     G.server.app = MaxBody.new(G.server.app)
  69   end
  70
  71   # Rack response returned when there's an error
  72   def err(env)
  73     [ 413, [ %w(Content-Length 0), %w(Content-Type text/plain) ], [] ]
  74   end
  75
  76   # our main Rack middleware endpoint
  77   def call(env)
  78     catch(:rainbows_EFBIG) { app.call(env) } || err(env)
  79   end
  80
  81 end # class
  82 end # module