epoll/client: factor out on_close method
[rainbows.git] / lib / rainbows / max_body.rb
blob105bbb0b4e7df679e4482a7df08dbcd1e0b219d3
1 # -*- encoding: binary -*-
3 # 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.
9 # For more fine-grained conrol, you may also define it per-endpoint in
10 # your Rack config.ru like this:
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
21 class Rainbows::MaxBody
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
33   # :stopdoc:
34   RACK_INPUT = "rack.input".freeze
35   CONTENT_LENGTH = "CONTENT_LENGTH"
36   HTTP_TRANSFER_ENCODING = "HTTP_TRANSFER_ENCODING"
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
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
63     # force ourselves to the outermost middleware layer
64     Rainbows.server.app = self.new(Rainbows.server.app)
65   end
67   # Rack response returned when there's an error
68   def err # :nodoc:
69     [ 413, { 'Content-Length' => '0', 'Content-Type' => 'text/plain' }, [] ]
70   end
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
78   # :startdoc:
79 end
80 require 'rainbows/max_body/wrapper'
81 require 'rainbows/max_body/rewindable_wrapper'