From 31cf77e7aa2f2e6065e7ace44e55c3f042b51f1b Mon Sep 17 00:00:00 2001 From: Eric Wong Date: Tue, 16 Nov 2010 16:16:42 -0800 Subject: reimplement client_max_body_size handlers This allows the client_max_body_size implementation to not rely on Unicorn::TeeInput internals, allowing it to be used with Unicorn::StreamInput (or any other (nearly) Rack::Lint-compatible input object). --- lib/rainbows.rb | 1 - lib/rainbows/max_body.rb | 103 ++++++++++++++-------------- lib/rainbows/max_body/rewindable_wrapper.rb | 17 +++++ lib/rainbows/max_body/wrapper.rb | 26 +++++++ lib/rainbows/process_client.rb | 2 +- lib/rainbows/revactor.rb | 2 +- lib/rainbows/tee_input.rb | 18 ----- 7 files changed, 95 insertions(+), 74 deletions(-) create mode 100644 lib/rainbows/max_body/rewindable_wrapper.rb create mode 100644 lib/rainbows/max_body/wrapper.rb delete mode 100644 lib/rainbows/tee_input.rb (limited to 'lib') diff --git a/lib/rainbows.rb b/lib/rainbows.rb index a1f842c..144b539 100644 --- a/lib/rainbows.rb +++ b/lib/rainbows.rb @@ -39,7 +39,6 @@ module Rainbows require 'rainbows/http_server' require 'rainbows/response' require 'rainbows/client' - require 'rainbows/tee_input' require 'rainbows/process_client' autoload :Base, 'rainbows/base' autoload :Sendfile, 'rainbows/sendfile' diff --git a/lib/rainbows/max_body.rb b/lib/rainbows/max_body.rb index 223c57e..5e2cb64 100644 --- a/lib/rainbows/max_body.rb +++ b/lib/rainbows/max_body.rb @@ -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 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' diff --git a/lib/rainbows/max_body/rewindable_wrapper.rb b/lib/rainbows/max_body/rewindable_wrapper.rb new file mode 100644 index 0000000..b52726e --- /dev/null +++ b/lib/rainbows/max_body/rewindable_wrapper.rb @@ -0,0 +1,17 @@ +# -*- encoding: binary -*- +# :enddoc: +class Rainbows::MaxBody::RewindableWrapper < Rainbows::MaxBody::Wrapper + def initialize(rack_input, limit) + @orig_limit = limit + super + end + + def rewind + @limit = @orig_limit + @input.rewind + end + + def size + @input.size + end +end diff --git a/lib/rainbows/max_body/wrapper.rb b/lib/rainbows/max_body/wrapper.rb new file mode 100644 index 0000000..3c38ca6 --- /dev/null +++ b/lib/rainbows/max_body/wrapper.rb @@ -0,0 +1,26 @@ +# -*- encoding: binary -*- +# :enddoc: +class Rainbows::MaxBody::Wrapper + def initialize(rack_input, limit) + @input, @limit = rack_input, limit + end + + def check(rv) + throw :rainbows_EFBIG if rv && ((@limit -= rv.size) < 0) + rv + end + + def each(&block) + while line = @input.gets + yield check(line) + end + end + + def read(*args) + check(@input.read(*args)) + end + + def gets + check(@input.gets) + end +end diff --git a/lib/rainbows/process_client.rb b/lib/rainbows/process_client.rb index d66c1ae..143a7a9 100644 --- a/lib/rainbows/process_client.rb +++ b/lib/rainbows/process_client.rb @@ -6,7 +6,7 @@ module Rainbows::ProcessClient HttpParser = Unicorn::HttpParser NULL_IO = Unicorn::HttpRequest::NULL_IO RACK_INPUT = Unicorn::HttpRequest::RACK_INPUT - TeeInput = Rainbows::TeeInput + TeeInput = Unicorn::TeeInput include Rainbows::Const # once a client is accepted, it is processed in its entirety here diff --git a/lib/rainbows/revactor.rb b/lib/rainbows/revactor.rb index a0b4bbf..9e69251 100644 --- a/lib/rainbows/revactor.rb +++ b/lib/rainbows/revactor.rb @@ -51,7 +51,7 @@ module Rainbows::Revactor env[CLIENT_IO] = client env[RACK_INPUT] = 0 == hp.content_length ? - NULL_IO : TeeInput.new(TeeSocket.new(client), hp) + NULL_IO : Unicorn::TeeInput.new(TeeSocket.new(client), hp) env[REMOTE_ADDR] = remote_addr status, headers, body = app.call(env.update(RACK_DEFAULTS)) diff --git a/lib/rainbows/tee_input.rb b/lib/rainbows/tee_input.rb deleted file mode 100644 index 956c68f..0000000 --- a/lib/rainbows/tee_input.rb +++ /dev/null @@ -1,18 +0,0 @@ -# -*- encoding: binary -*- -# :enddoc: -module Rainbows - - # acts like tee(1) on an input input to provide a input-like stream - # while providing rewindable semantics through a File/StringIO - # backing store. On the first pass, the input is only read on demand - # so your Rack application can use input notification (upload progress - # and like). This should fully conform to the Rack::InputWrapper - # specification on the public API. This class is intended to be a - # strict interpretation of Rack::InputWrapper functionality and will - # not support any deviations from it. - class TeeInput < Unicorn::TeeInput - - # empty class, this is to avoid unecessarily modifying Unicorn::TeeInput - # when MaxBody::Limit is included - end -end -- cgit v1.2.3-24-ge0c7