tee_input: make interface more usable outside of Unicorn

TeeInput being needed is now (once again) an uncommon code path
so there's no point in relying on global constants.  While we're
at it, allow StringIO to be used in the presence of small
inputs; too.

3 files changed, 150 insertions, 54 deletions

diff --git a/lib/unicorn/http_request.rb b/lib/unicorn/http_request.rb
index a0d811c..1358ccc 100644
--- a/lib/unicorn/http_request.rb
+++ b/lib/unicorn/http_request.rb
@@ -74,7 +74,7 @@ module Unicorn
     # returns a # Rack environment if successful
     def handle_body(socket)
       REQ[Const::RACK_INPUT] = 0 == PARSER.content_length ?
-                               NULL_IO : Unicorn::TeeInput.new(socket)
+                   NULL_IO : Unicorn::TeeInput.new(socket, REQ, PARSER, BUF)
       REQ.update(DEFAULTS)
     end
 
diff --git a/lib/unicorn/tee_input.rb b/lib/unicorn/tee_input.rb
index 9f56860..fcf4a95 100644
--- a/lib/unicorn/tee_input.rb
+++ b/lib/unicorn/tee_input.rb
@@ -11,23 +11,16 @@
 # not support any deviations from it.
 
 module Unicorn
-  class TeeInput
-
-    # it's so awesome to not have to care for thread safety...
-
-    RAW = HttpRequest::BUF # :nodoc:
-    DST = RAW.dup # :nodoc:
-    PARSER = HttpRequest::PARSER # :nodoc:
-    REQ = HttpRequest::REQ # :nodoc:
-
-    def initialize(socket)
-      @tmp = Util.tmpio
-      @size = PARSER.content_length
-      return(@input = nil) if 0 == @size
-      @input = socket
-      if RAW.size > 0
-        PARSER.filter_body(DST, RAW) and finalize_input
-        @tmp.write(DST)
+  class TeeInput < Struct.new(:socket, :req, :parser, :buf)
+
+    def initialize(*args)
+      super(*args)
+      @size = parser.content_length
+      @tmp = @size && @size < Const::MAX_BODY ? StringIO.new(Z.dup) : Util.tmpio
+      @buf2 = buf.dup
+      if buf.size > 0
+        parser.filter_body(@buf2, buf) and finalize_input
+        @tmp.write(@buf2)
         @tmp.seek(0)
       end
     end
@@ -39,45 +32,45 @@ module Unicorn
     def size
       @size and return @size
 
-      if @input
-        pos = @tmp.tell
-        while tee(Const::CHUNK_SIZE, DST)
+      if socket
+        pos = @tmp.pos
+        while tee(Const::CHUNK_SIZE, @buf2)
         end
         @tmp.seek(pos)
       end
 
-      @size = @tmp.stat.size
+      @size = tmp_size
     end
 
     def read(*args)
-      @input or return @tmp.read(*args)
+      socket or return @tmp.read(*args)
 
       length = args.shift
       if nil == length
         rv = @tmp.read || Z.dup
-        while tee(Const::CHUNK_SIZE, DST)
-          rv << DST
+        while tee(Const::CHUNK_SIZE, @buf2)
+          rv << @buf2
         end
         rv
       else
-        buf = args.shift || DST.dup
-        diff = @tmp.stat.size - @tmp.pos
+        rv = args.shift || @buf2.dup
+        diff = tmp_size - @tmp.pos
         if 0 == diff
-          tee(length, buf)
+          tee(length, rv)
         else
-          @tmp.read(diff > length ? length : diff, buf)
+          @tmp.read(diff > length ? length : diff, rv)
         end
       end
     end
 
     # takes zero arguments for strict Rack::Lint compatibility, unlike IO#gets
     def gets
-      @input or return @tmp.gets
+      socket or return @tmp.gets
       nil == $/ and return read
 
-      orig_size = @tmp.stat.size
+      orig_size = tmp_size
       if @tmp.pos == orig_size
-        tee(Const::CHUNK_SIZE, DST) or return nil
+        tee(Const::CHUNK_SIZE, @buf2) or return nil
         @tmp.seek(orig_size)
       end
 
@@ -87,7 +80,7 @@ module Unicorn
       # unlikely, if we got here, then @tmp is at EOF
       begin
         orig_size = @tmp.pos
-        tee(Const::CHUNK_SIZE, DST) or break
+        tee(Const::CHUNK_SIZE, @buf2) or break
         @tmp.seek(orig_size)
         line << @tmp.gets
         $/ == line[-$/.size, $/.size] and return line
@@ -102,11 +95,11 @@ module Unicorn
         yield line
       end
 
-      self # Rack does not specify what the return value here
+      self # Rack does not specify what the return value is here
     end
 
     def rewind
-      @tmp.rewind # Rack does not specify what the return value here
+      @tmp.rewind # Rack does not specify what the return value is here
     end
 
   private
@@ -114,12 +107,12 @@ module Unicorn
     # tees off a +length+ chunk of data from the input into the IO
     # backing store as well as returning it.  +buf+ must be specified.
     # returns nil if reading from the input returns nil
-    def tee(length, buf)
-      unless PARSER.body_eof?
+    def tee(length, dst)
+      unless parser.body_eof?
         begin
-          if PARSER.filter_body(buf, @input.readpartial(length, RAW)).nil?
-            @tmp.write(buf)
-            return buf
+          if parser.filter_body(dst, socket.readpartial(length, buf)).nil?
+            @tmp.write(dst)
+            return dst
           end
         rescue EOFError
         end
@@ -128,10 +121,14 @@ module Unicorn
     end
 
     def finalize_input
-      while PARSER.trailers(REQ, RAW).nil?
-        RAW << @input.readpartial(Const::CHUNK_SIZE, DST)
+      while parser.trailers(req, buf).nil?
+        buf << socket.readpartial(Const::CHUNK_SIZE, @buf2)
       end
-      @input = nil
+      self.socket = nil
+    end
+
+    def tmp_size
+      StringIO === @tmp ? @tmp.size : @tmp.stat.size
     end
 
   end
diff --git a/test/unit/test_tee_input.rb b/test/unit/test_tee_input.rb
index 3a6c998..0594d6d 100644
--- a/test/unit/test_tee_input.rb
+++ b/test/unit/test_tee_input.rb
@@ -27,7 +27,7 @@ class TestTeeInput < Test::Unit::TestCase
 
   def test_gets_long
     init_parser("hello", 5 + (4096 * 4 * 3) + "#$/foo#$/".size)
-    ti = Unicorn::TeeInput.new(@rd)
+    ti = Unicorn::TeeInput.new(@rd, @env, @parser, @buf)
     status = line = nil
     pid = fork {
       @rd.close
@@ -48,7 +48,7 @@ class TestTeeInput < Test::Unit::TestCase
 
   def test_gets_short
     init_parser("hello", 5 + "#$/foo".size)
-    ti = Unicorn::TeeInput.new(@rd)
+    ti = Unicorn::TeeInput.new(@rd, @env, @parser, @buf)
     status = line = nil
     pid = fork {
       @rd.close
@@ -65,19 +65,118 @@ class TestTeeInput < Test::Unit::TestCase
     assert status.success?
   end
 
+  def test_small_body
+    init_parser('hello')
+    ti = Unicorn::TeeInput.new(@rd, @env, @parser, @buf)
+    assert_equal 0, @parser.content_length
+    assert @parser.body_eof?
+    assert_equal StringIO, ti.instance_eval { @tmp.class }
+    assert_equal 0, ti.instance_eval { @tmp.pos }
+    assert_equal 5, ti.size
+    assert_equal 'hello', ti.read
+    assert_equal '', ti.read
+    assert_nil ti.read(4096)
+  end
+
+  def test_read_with_buffer
+    init_parser('hello')
+    ti = Unicorn::TeeInput.new(@rd, @env, @parser, @buf)
+    buf = ''
+    rv = ti.read(4, buf)
+    assert_equal 'hell', rv
+    assert_equal 'hell', buf
+    assert_equal rv.object_id, buf.object_id
+    assert_equal 'o', ti.read
+    assert_equal nil, ti.read(5, buf)
+    assert_equal 0, ti.rewind
+    assert_equal 'hello', ti.read(5, buf)
+    assert_equal 'hello', buf
+  end
+
+  def test_big_body
+    init_parser('.' * Unicorn::Const::MAX_BODY << 'a')
+    ti = Unicorn::TeeInput.new(@rd, @env, @parser, @buf)
+    assert_equal 0, @parser.content_length
+    assert @parser.body_eof?
+    assert_kind_of File, ti.instance_eval { @tmp }
+    assert_equal 0, ti.instance_eval { @tmp.pos }
+    assert_equal Unicorn::Const::MAX_BODY + 1, ti.size
+  end
+
+  def test_big_body_multi
+    init_parser('.', Unicorn::Const::MAX_BODY + 1)
+    ti = Unicorn::TeeInput.new(@rd, @env, @parser, @buf)
+    assert_equal Unicorn::Const::MAX_BODY, @parser.content_length
+    assert ! @parser.body_eof?
+    assert_kind_of File, ti.instance_eval { @tmp }
+    assert_equal 0, ti.instance_eval { @tmp.pos }
+    assert_equal 1, ti.instance_eval { tmp_size }
+    assert_equal Unicorn::Const::MAX_BODY + 1, ti.size
+    nr = Unicorn::Const::MAX_BODY / 4
+    pid = fork {
+      @rd.close
+      nr.times { @wr.write('....') }
+      @wr.close
+    }
+    @wr.close
+    assert_equal '.', ti.read(1)
+    assert_equal Unicorn::Const::MAX_BODY + 1, ti.size
+    nr.times {
+      assert_equal '....', ti.read(4)
+      assert_equal Unicorn::Const::MAX_BODY + 1, ti.size
+    }
+    assert_nil ti.read(1)
+    status = nil
+    assert_nothing_raised { pid, status = Process.waitpid2(pid) }
+    assert status.success?
+  end
+
+  def test_chunked
+    @parser = Unicorn::HttpParser.new
+    @buf = "POST / HTTP/1.1\r\n" \
+           "Host: localhost\r\n" \
+           "Transfer-Encoding: chunked\r\n" \
+           "\r\n"
+    assert_equal @env, @parser.headers(@env, @buf)
+    assert_equal "", @buf
+
+    pid = fork {
+      @rd.close
+      5.times { @wr.write("5\r\nabcde\r\n") }
+      @wr.write("0\r\n")
+    }
+    @wr.close
+    ti = Unicorn::TeeInput.new(@rd, @env, @parser, @buf)
+    assert_nil @parser.content_length
+    assert_nil ti.instance_eval { @size }
+    assert ! @parser.body_eof?
+    assert_equal 25, ti.size
+    assert @parser.body_eof?
+    assert_equal 25, ti.instance_eval { @size }
+    assert_equal 0, ti.instance_eval { @tmp.pos }
+    assert_nothing_raised { ti.rewind }
+    assert_equal 0, ti.instance_eval { @tmp.pos }
+    assert_equal 'abcdeabcdeabcdeabcde', ti.read(20)
+    assert_equal 20, ti.instance_eval { @tmp.pos }
+    assert_nothing_raised { ti.rewind }
+    assert_equal 0, ti.instance_eval { @tmp.pos }
+    assert_kind_of File, ti.instance_eval { @tmp }
+    status = nil
+    assert_nothing_raised { pid, status = Process.waitpid2(pid) }
+    assert status.success?
+  end
+
 private
 
   def init_parser(body, size = nil)
-    @parser = Unicorn::TeeInput::PARSER
-    @parser.reset
+    @parser = Unicorn::HttpParser.new
     body = body.to_s.freeze
-    buf = "POST / HTTP/1.1\r\n" \
-          "Host: localhost\r\n" \
-          "Content-Length: #{size || body.size}\r\n" \
-          "\r\n#{body}"
-    buf = Unicorn::TeeInput::RAW.replace(buf)
-    assert_equal @env, @parser.headers(@env, buf)
-    assert_equal body, buf
+    @buf = "POST / HTTP/1.1\r\n" \
+           "Host: localhost\r\n" \
+           "Content-Length: #{size || body.size}\r\n" \
+           "\r\n#{body}"
+    assert_equal @env, @parser.headers(@env, @buf)
+    assert_equal body, @buf
   end
 
 end