wire up client_max_body_size limits

This is mostly code imported from Rainbows! (so GPLv2+).  This should
implement everything necessary to prevent clients from DoS-ing us with
overly large bodies.  The default is 1M (same as Rainbows! and nginx).

Yahns::MaxBody may become part of the public API (as the equivalent is
in Rainbows!), since it makes more sense in the rackup (config.ru) file
(since it's endpoint-specific).  However, that's confusing as
Yahns::MaxBody only works when input_buffering is :lazy or false, and
not when it is true (preread).

9 files changed, 373 insertions, 10 deletions

diff --git a/lib/yahns/cap_input.rb b/lib/yahns/cap_input.rb
new file mode 100644
index 0000000..1aa10b6
--- /dev/null
+++ b/lib/yahns/cap_input.rb
@@ -0,0 +1,22 @@
+# -*- encoding: binary -*-
+# Copyright (C) 2009-2013, Eric Wong <normalperson@yhbt.net> et. al.
+# License: GPLv2 or later (https://www.gnu.org/licenses/gpl-2.0.txt)
+
+# This is used as the @input/env["rack.input"] when
+# input_buffering == true or :lazy
+class Yahns::CapInput < Yahns::TmpIO # :nodoc:
+  attr_writer :bytes_left
+
+  def self.new(limit)
+    rv = super()
+    rv.bytes_left = limit
+    rv
+  end
+
+  def write(buf)
+    if (@bytes_left -= buf.size) < 0
+      raise Unicorn::RequestEntityTooLargeError, "chunked body too big", []
+    end
+    super(buf)
+  end
+end
diff --git a/lib/yahns/config.rb b/lib/yahns/config.rb
index 4861f3e..61de74e 100644
--- a/lib/yahns/config.rb
+++ b/lib/yahns/config.rb
@@ -285,7 +285,6 @@ class Yahns::Config # :nodoc:
   {
     # config name, minimum value
     client_body_buffer_size: 1,
-    client_max_body_size: 0,
     client_header_buffer_size: 1,
     client_max_header_size: 1,
     client_timeout: 0,
@@ -298,6 +297,12 @@ class Yahns::Config # :nodoc:
     )
   end
 
+  def client_max_body_size(val)
+    var = _check_in_block(:app, :client_max_body_size)
+    val = _check_int(var, val, 0) if val != nil
+    @block.ctx.__send__("#{var}=", val)
+  end
+
   def input_buffering(val)
     var = _check_in_block(:app, :input_buffering)
     ok = [ :lazy, true, false ]
diff --git a/lib/yahns/http_client.rb b/lib/yahns/http_client.rb
index 198c130..e95bb47 100644
--- a/lib/yahns/http_client.rb
+++ b/lib/yahns/http_client.rb
@@ -45,9 +45,17 @@ class Yahns::HttpClient < Kgio::Socket # :nodoc:
     end while true
   end
 
+  # used only with "input_buffering true"
   def mkinput_preread
+    k = self.class
+    len = @hs.content_length
+    mbs = k.client_max_body_size
+    if mbs && len && len > mbs
+      raise Unicorn::RequestEntityTooLargeError,
+            "Content-Length:#{len} too large (>#{mbs})", []
+    end
     @state = :body
-    @input = self.class.tmpio_for(@hs.content_length)
+    @input = k.tmpio_for(len)
     rbuf = Thread.current[:yahns_rbuf]
     @hs.filter_body(rbuf, @hs.buf)
     @input.write(rbuf)
diff --git a/lib/yahns/http_context.rb b/lib/yahns/http_context.rb
index 97a0f82..547d41f 100644
--- a/lib/yahns/http_context.rb
+++ b/lib/yahns/http_context.rb
@@ -24,7 +24,7 @@ module Yahns::HttpContext # :nodoc:
     @check_client_connection = false
     @client_body_buffer_size = 112 * 1024
     @client_header_buffer_size = 4000
-    @client_max_body_size = 1024 * 1024
+    @client_max_body_size = 1024 * 1024 # nil => infinity
     @input_buffering = true
     @output_buffering = true
     @persistent_connections = true
@@ -34,7 +34,19 @@ module Yahns::HttpContext # :nodoc:
 
   # call this after forking
   def after_fork_init
-    @app = @yahns_rack.app_after_fork
+    @app = __wrap_app(@yahns_rack.app_after_fork)
+  end
+
+  def __wrap_app(app)
+    # input_buffering == false is handled in http_client
+    return app if @client_max_body_size == nil
+
+    require 'yahns/cap_input'
+    return app if @input_buffering == true
+
+    # @input_buffering == false/:lazy
+    require 'yahns/max_body'
+    Yahns::MaxBody.new(app, @client_max_body_size)
   end
 
   # call this immediately after successful accept()/accept4()
@@ -59,7 +71,11 @@ module Yahns::HttpContext # :nodoc:
   end
 
   def tmpio_for(len)
-    len && len <= @client_body_buffer_size ?
-      StringIO.new("") : Yahns::TmpIO.new
+    if len # Content-Length given
+      len <= @client_body_buffer_size ? StringIO.new("") : Yahns::TmpIO.new
+    else # chunked, unknown length
+      mbs = @client_max_body_size
+      mbs ? Yahns::CapInput.new(mbs) : Yahns::TmpIO.new
+    end
   end
 end
diff --git a/lib/yahns/max_body.rb b/lib/yahns/max_body.rb
new file mode 100644
index 0000000..fadbddc
--- /dev/null
+++ b/lib/yahns/max_body.rb
@@ -0,0 +1,60 @@
+# -*- encoding: binary -*-
+# Copyright (C) 2009-2013, Eric Wong <normalperson@yhbt.net> et. al.
+# License: GPLv2 or later (https://www.gnu.org/licenses/gpl-2.0.txt)
+
+# 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.
+#
+# For more fine-grained control, you may also define it per-endpoint in
+# your Rack config.ru like this:
+#
+#        map "/limit_1M" do
+#          use Yahns::MaxBody, 1024*1024
+#          run MyApp
+#        end
+#        map "/limit_10M" do
+#          use Yahns::MaxBody, 1024*1024*10
+#          run MyApp
+#        end
+class Yahns::MaxBody # :nodoc:
+  # This is automatically called when used with Rack::Builder#use
+  # See Yahns::MaxBody
+  def initialize(app, limit)
+    Integer === limit or raise ArgumentError, "limit not an Integer"
+    @app = app
+    @limit = limit
+  end
+
+  RACK_INPUT = "rack.input".freeze # :nodoc:
+  CONTENT_LENGTH = "CONTENT_LENGTH" # :nodoc:
+  HTTP_TRANSFER_ENCODING = "HTTP_TRANSFER_ENCODING" # :nodoc:
+
+  # our main Rack middleware endpoint
+  def call(env) # :nodoc:
+    catch(:yahns_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
+      @app.call(env)
+    end || err
+  end
+
+  # Rack response returned when there's an error
+  def err # :nodoc:
+    [ 413, { 'Content-Length' => '0', 'Content-Type' => 'text/plain' }, [] ]
+  end
+
+  def limit_input!(env) # :nodoc:
+    input = env[RACK_INPUT]
+    klass = input.respond_to?(:rewind) ? RewindableWrapper : Wrapper
+    env[RACK_INPUT] = klass.new(input, @limit)
+  end
+end
+require_relative 'max_body/wrapper'
+require_relative 'max_body/rewindable_wrapper'
diff --git a/lib/yahns/max_body/rewindable_wrapper.rb b/lib/yahns/max_body/rewindable_wrapper.rb
new file mode 100644
index 0000000..5888def
--- /dev/null
+++ b/lib/yahns/max_body/rewindable_wrapper.rb
@@ -0,0 +1,19 @@
+# -*- encoding: binary -*-
+# Copyright (C) 2009-2013, Eric Wong <normalperson@yhbt.net> et. al.
+# License: GPLv2 or later (https://www.gnu.org/licenses/gpl-2.0.txt)
+class Yahns::MaxBody::RewindableWrapper < Yahns::MaxBody::Wrapper # :nodoc:
+  def initialize(rack_input, limit)
+    @orig_limit = limit
+    super
+  end
+
+  def rewind
+    @limit = @orig_limit
+    @rbuf = ''
+    @input.rewind
+  end
+
+  def size
+    @input.size
+  end
+end
diff --git a/lib/yahns/max_body/wrapper.rb b/lib/yahns/max_body/wrapper.rb
new file mode 100644
index 0000000..68f6b5b
--- /dev/null
+++ b/lib/yahns/max_body/wrapper.rb
@@ -0,0 +1,71 @@
+# -*- encoding: binary -*-
+# Copyright (C) 2009-2013, Eric Wong <normalperson@yhbt.net> et. al.
+# License: GPLv2 or later (https://www.gnu.org/licenses/gpl-2.0.txt)
+#
+# This is only used for chunked request bodies, which are rare
+class Yahns::MaxBody::Wrapper # :nodoc:
+  def initialize(rack_input, limit)
+    @input, @limit, @rbuf = rack_input, limit, ''
+  end
+
+  def each
+    while line = gets
+      yield line
+    end
+  end
+
+  # chunked encoding means this method behaves more like readpartial,
+  # since Rack does not support a method named "readpartial"
+  def read(length = nil, rv = '')
+    if length
+      if length <= @rbuf.size
+        length < 0 and raise ArgumentError, "negative length #{length} given"
+        rv.replace(@rbuf.slice!(0, length))
+      elsif @rbuf.empty?
+        checked_read(length, rv) or return
+      else
+        rv.replace(@rbuf.slice!(0, @rbuf.size))
+      end
+      rv.empty? && length != 0 ? nil : rv
+    else
+      rv.replace(read_all)
+    end
+  end
+
+  def gets
+    sep = $/
+    if sep.nil?
+      rv = read_all
+      return rv.empty? ? nil : rv
+    end
+    re = /\A(.*?#{Regexp.escape(sep)})/
+
+    begin
+      @rbuf.sub!(re, '') and return $1
+
+      if tmp = checked_read(16384)
+        @rbuf << tmp
+      elsif @rbuf.empty? # EOF
+        return nil
+      else # EOF, return whatever is left
+        return @rbuf.slice!(0, @rbuf.size)
+      end
+    end while true
+  end
+
+  def checked_read(length = 16384, buf = '')
+    if @input.read(length, buf)
+      throw :yahns_EFBIG if ((@limit -= buf.size) < 0)
+      return buf
+    end
+  end
+
+  def read_all
+    rv = @rbuf.slice!(0, @rbuf.size)
+    tmp = ''
+    while checked_read(16384, tmp)
+      rv << tmp
+    end
+    rv
+  end
+end
diff --git a/lib/yahns/stream_input.rb b/lib/yahns/stream_input.rb
index 8bef95e..b8a3a8f 100644
--- a/lib/yahns/stream_input.rb
+++ b/lib/yahns/stream_input.rb
@@ -60,7 +60,7 @@ class Yahns::StreamInput # :nodoc:
   end
 
   def __rsize
-    @client.class.client_body_buffer_size
+    @client ? @client.class.client_body_buffer_size : nil
   end
 
   # :call-seq:
@@ -79,7 +79,7 @@ class Yahns::StreamInput # :nodoc:
       return rv.empty? ? nil : rv
     end
     re = /\A(.*?#{Regexp.escape(sep)})/
-    rsize = __rsize
+    rsize = __rsize or return
     begin
       @rbuf.sub!(re, '') and return $1
       return @rbuf.empty? ? nil : @rbuf.slice!(0, @rbuf.size) if eof?
@@ -124,8 +124,7 @@ class Yahns::StreamInput # :nodoc:
 
   def read_all(dst)
     dst.replace(@rbuf)
-    @client or return
-    rsize = @client.class.client_body_buffer_size
+    rsize = __rsize or return
     until eof?
       @client.kgio_read(rsize, @buf) or eof!
       filter_body(@rbuf, @buf)
diff --git a/test/test_client_max_body_size.rb b/test/test_client_max_body_size.rb
new file mode 100644
index 0000000..b50e2de
--- /dev/null
+++ b/test/test_client_max_body_size.rb
@@ -0,0 +1,163 @@
+# Copyright (C) 2013, Eric Wong <normalperson@yhbt.net> and all contributors
+# License: GPLv3 or later (https://www.gnu.org/licenses/gpl-3.0.txt)
+require_relative 'server_helper'
+
+class TestClientMaxBodySize < Testcase
+  parallelize_me!
+  include ServerHelper
+  alias setup server_helper_setup
+  alias teardown server_helper_teardown
+  DEFMBS = 1024 * 1024
+
+  DRAINER = lambda do |e|
+    input = e["rack.input"]
+    buf = ""
+    nr = 0
+    while rv = input.read(16384, buf)
+      nr += rv.size
+    end
+    body = nr.to_s
+    h = { "Content-Length" => body.size.to_s, "Content-Type" => 'text/plain' }
+    [ 200, h, [body] ]
+  end
+
+  def identity_req(bytes, body = true)
+    body_bytes = body ? bytes : 0
+    "PUT / HTTP/1.1\r\nConnection: close\r\nHost: example.com\r\n" \
+    "Content-Length: #{bytes}\r\n\r\n#{'*' * body_bytes}"
+  end
+
+  def mkserver(cfg)
+    fork do
+      srv = Yahns::Server.new(cfg)
+      ENV["YAHNS_FD"] = @srv.fileno.to_s
+      srv.start.join
+    end
+  end
+
+  def test_0_lazy; cmbs_test_0(:lazy); end
+  def test_0_true; cmbs_test_0(true); end
+  def test_0_false; cmbs_test_0(false); end
+
+  def cmbs_test_0(btype)
+    err = @err
+    cfg = Yahns::Config.new
+    host, port = @srv.addr[3], @srv.addr[1]
+    cfg.instance_eval do
+      GTL.synchronize {
+        app(:rack, DRAINER) {
+          listen "#{host}:#{port}"
+          input_buffering btype
+          client_max_body_size 0
+        }
+      }
+      logger(Logger.new(err.path))
+    end
+    pid = mkserver(cfg)
+    default_identity_checks(host, port, 0)
+    default_chunked_checks(host, port, 0)
+  ensure
+    quit_wait(pid)
+  end
+
+  def test_cmbs_lazy; cmbs_test(:lazy); end
+  def test_cmbs_true; cmbs_test(true); end
+  def test_cmbs_false; cmbs_test(false); end
+
+  def cmbs_test(btype)
+    err = @err
+    cfg = Yahns::Config.new
+    host, port = @srv.addr[3], @srv.addr[1]
+    cfg.instance_eval do
+      GTL.synchronize {
+        app(:rack, DRAINER) {
+          listen "#{host}:#{port}"
+          input_buffering btype
+        }
+      }
+      logger(Logger.new(err.path))
+    end
+    pid = mkserver(cfg)
+    default_identity_checks(host, port)
+    default_chunked_checks(host, port)
+  ensure
+    quit_wait(pid)
+  end
+
+  def test_inf_false; big_test(false); end
+  def test_inf_true; big_test(true); end
+  def test_inf_lazy; big_test(:lazy); end
+
+  def big_test(btype)
+    err = @err
+    cfg = Yahns::Config.new
+    host, port = @srv.addr[3], @srv.addr[1]
+    cfg.instance_eval do
+      GTL.synchronize {
+        app(:rack, DRAINER) {
+          listen "#{host}:#{port}"
+          input_buffering btype
+          client_max_body_size nil
+        }
+      }
+      logger(Logger.new(err.path))
+    end
+    pid = mkserver(cfg)
+
+    bytes = 10 * 1024 * 1024
+    r = `dd if=/dev/zero bs=#{bytes} count=1 2>/dev/null | \
+         curl -sSf -HExpect: -T- http://#{host}:#{port}/`
+    assert $?.success?, $?.inspect
+    assert_equal bytes.to_s, r
+
+    r = `dd if=/dev/zero bs=#{bytes} count=1 2>/dev/null | \
+         curl -sSf -HExpect: -HContent-Length:#{bytes} -HTransfer-Encoding: \
+         -T- http://#{host}:#{port}/`
+    assert $?.success?, $?.inspect
+    assert_equal bytes.to_s, r
+  ensure
+    quit_wait(pid)
+  end
+
+  def default_chunked_checks(host, port, defmax = DEFMBS)
+    r = `curl -sSf -HExpect: -T- </dev/null http://#{host}:#{port}/`
+    assert $?.success?, $?.inspect
+    assert_equal "0", r
+
+    r = `dd if=/dev/zero bs=#{defmax} count=1 2>/dev/null | \
+         curl -sSf -HExpect: -T- http://#{host}:#{port}/`
+    assert $?.success?, $?.inspect
+    assert_equal "#{defmax}", r
+
+    r = `dd if=/dev/zero bs=#{defmax + 1} count=1 2>/dev/null | \
+         curl -sf -HExpect: -T- --write-out %{http_code} \
+         http://#{host}:#{port}/ 2>&1`
+    refute $?.success?, $?.inspect
+    assert_equal "413", r
+  end
+
+  def default_identity_checks(host, port, defmax = DEFMBS)
+    if defmax >= 666
+      c = TCPSocket.new(host, port)
+      c.write(identity_req(666))
+      assert_equal "666", c.read.split(/\r\n\r\n/)[1]
+      c.close
+    end
+
+    c = TCPSocket.new(host, port)
+    c.write(identity_req(0))
+    assert_equal "0", c.read.split(/\r\n\r\n/)[1]
+    c.close
+
+    c = TCPSocket.new(host, port)
+    c.write(identity_req(defmax))
+    assert_equal "#{defmax}", c.read.split(/\r\n\r\n/)[1]
+    c.close
+
+    toobig = defmax + 1
+    c = TCPSocket.new(host, port)
+    c.write(identity_req(toobig, false))
+    assert_match(%r{\AHTTP/1\.[01] 413 }, Timeout.timeout(10) { c.read })
+    c.close
+  end
+end