about summary refs log tree commit homepage
path: root/lib/rainbows
diff options
context:
space:
mode:
authorEric Wong <normalperson@yhbt.net>2010-11-16 16:16:42 -0800
committerEric Wong <normalperson@yhbt.net>2010-11-16 16:21:33 -0800
commit31cf77e7aa2f2e6065e7ace44e55c3f042b51f1b (patch)
tree7dfd53c499cbf7a52165d3480d56d8576a743781 /lib/rainbows
parent42747db815ad668b20849afb2a9dcdd1319713ae (diff)
downloadrainbows-31cf77e7aa2f2e6065e7ace44e55c3f042b51f1b.tar.gz
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).
Diffstat (limited to 'lib/rainbows')
-rw-r--r--lib/rainbows/max_body.rb103
-rw-r--r--lib/rainbows/max_body/rewindable_wrapper.rb17
-rw-r--r--lib/rainbows/max_body/wrapper.rb26
-rw-r--r--lib/rainbows/process_client.rb2
-rw-r--r--lib/rainbows/revactor.rb2
-rw-r--r--lib/rainbows/tee_input.rb18
6 files changed, 95 insertions, 73 deletions
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