From de3bcfe3ba9402bd510f7414df1763b6b99dae47 Mon Sep 17 00:00:00 2001 From: Eric Wong Date: Mon, 3 May 2010 17:56:00 -0700 Subject: cleanup request size limiting for TeeInput users WAvoid mucking with Unicorn::TeeInput, since other apps may depend on that class, so we subclass it as Rainbows::TeeInput and modify as necessary in worker processes. For Revactor, remove the special-cased Rainbows::Revactor::TeeInput class and instead emulate readpartial for Revactor sockets instead. --- lib/rainbows/base.rb | 3 +-- lib/rainbows/max_body.rb | 15 +++-------- lib/rainbows/revactor.rb | 39 +++++++++++++++++++++++++--- lib/rainbows/revactor/tee_input.rb | 52 -------------------------------------- lib/rainbows/tee_input.rb | 16 ++++++++++++ 5 files changed, 57 insertions(+), 68 deletions(-) delete mode 100644 lib/rainbows/revactor/tee_input.rb create mode 100644 lib/rainbows/tee_input.rb (limited to 'lib/rainbows') diff --git a/lib/rainbows/base.rb b/lib/rainbows/base.rb index 864b847..faec951 100644 --- a/lib/rainbows/base.rb +++ b/lib/rainbows/base.rb @@ -68,8 +68,7 @@ module Rainbows env[CLIENT_IO] = client env[RACK_INPUT] = 0 == hp.content_length ? - HttpRequest::NULL_IO : - Unicorn::TeeInput.new(client, env, hp, buf) + HttpRequest::NULL_IO : TeeInput.new(client, env, hp, buf) env[REMOTE_ADDR] = remote_addr status, headers, body = app.call(env.update(RACK_DEFAULTS)) diff --git a/lib/rainbows/max_body.rb b/lib/rainbows/max_body.rb index 71141c1..ca63ea4 100644 --- a/lib/rainbows/max_body.rb +++ b/lib/rainbows/max_body.rb @@ -7,7 +7,7 @@ module Rainbows # setting class MaxBody < Struct.new(:app) - # this is meant to be included in Unicorn::TeeInput (and derived + # this is meant to be included in Rainbows::TeeInput (and derived # classes) to limit body sizes module Limit Util = Unicorn::Util @@ -41,7 +41,7 @@ class MaxBody < Struct.new(:app) end def tee(length, dst) - rv = _tee(length, dst) + rv = super if rv && ((@max_body -= rv.size) < 0) # make HttpParser#keepalive? => false to force an immediate disconnect # after we write @@ -60,17 +60,10 @@ class MaxBody < Struct.new(:app) case G.server.use when :Rev, :EventMachine, :NeverBlock return - when :Revactor - Rainbows::Revactor::TeeInput - else - Unicorn::TeeInput - end.class_eval do - alias _tee tee # can't use super here :< - remove_method :tee - remove_method :initialize if G.server.use != :Revactor # FIXME CODE SMELL - include Limit end + TeeInput.class_eval { include Limit } + # force ourselves to the outermost middleware layer G.server.app = MaxBody.new(G.server.app) end diff --git a/lib/rainbows/revactor.rb b/lib/rainbows/revactor.rb index ed08f2c..21fa72f 100644 --- a/lib/rainbows/revactor.rb +++ b/lib/rainbows/revactor.rb @@ -21,8 +21,6 @@ module Rainbows # concurrency features this model provides. module Revactor - require 'rainbows/revactor/tee_input' - RD_ARGS = {} include Base @@ -52,7 +50,7 @@ module Rainbows env[Const::CLIENT_IO] = client env[Const::RACK_INPUT] = 0 == hp.content_length ? HttpRequest::NULL_IO : - Rainbows::Revactor::TeeInput.new(client, env, hp, buf) + TeeInput.new(PartialSocket.new(client), env, hp, buf) env[Const::REMOTE_ADDR] = remote_addr response = app.call(env.update(RACK_DEFAULTS)) @@ -137,5 +135,40 @@ module Rainbows end end + # Revactor Sockets do not implement readpartial, so we emulate just + # enough to avoid mucking with TeeInput internals. Fortunately + # this code is not heavily used so we can usually avoid the overhead + # of adding a userspace buffer. + class PartialSocket < Struct.new(:socket, :rbuf) + def initialize(socket) + # IO::Buffer is used internally by Rev which Revactor is based on + # so we'll always have it available + super(socket, IO::Buffer.new) + end + + # Revactor socket reads always return an unspecified amount, + # sometimes too much + def readpartial(length, dst = "") + # always check and return from the userspace buffer first + rbuf.size > 0 and return dst.replace(rbuf.read(length)) + + # read off the socket since there was nothing in rbuf + tmp = socket.read + + # we didn't read too much, good, just return it straight back + # to avoid needlessly wasting memory bandwidth + tmp.size <= length and return dst.replace(tmp) + + # ugh, read returned too much, copy + reread to avoid slicing + rbuf << tmp[length, tmp.size] + dst.replace(tmp[0, length]) + end + + # just proxy any remaining methods TeeInput may use + def close + socket.close + end + end + end end diff --git a/lib/rainbows/revactor/tee_input.rb b/lib/rainbows/revactor/tee_input.rb deleted file mode 100644 index 99d1f7d..0000000 --- a/lib/rainbows/revactor/tee_input.rb +++ /dev/null @@ -1,52 +0,0 @@ -# -*- encoding: binary -*- -require 'rainbows/revactor' - -module Rainbows - module Revactor - - # 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 - - private - - # tees off a +length+ chunk of data from the input into the IO - # backing store as well as returning it. +dst+ must be specified. - # returns nil if reading from the input returns nil - def tee(length, dst) - unless parser.body_eof? - if parser.filter_body(dst, buf << socket.read).nil? - tmp.write(dst) - diff = dst.size - length - if diff > 0 - dst.replace(dst[0,length]) - tmp.seek(-diff, IO::SEEK_CUR) - end - return dst - end - end - finalize_input - rescue => e - client_error(e) - end - - def finalize_input - while parser.trailers(req, buf).nil? - # Don't worry about raising ClientShutdown here on EOFError, tee() - # will catch EOFError when app is processing it, otherwise in - # initialize we never get any chance to enter the app so the - # EOFError will just get trapped by Unicorn and not the Rack app - buf << socket.read - end - self.socket = nil - end - - end - end -end diff --git a/lib/rainbows/tee_input.rb b/lib/rainbows/tee_input.rb new file mode 100644 index 0000000..d405a5c --- /dev/null +++ b/lib/rainbows/tee_input.rb @@ -0,0 +1,16 @@ +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