about summary refs log tree commit homepage
diff options
context:
space:
mode:
-rw-r--r--lib/unicorn/http_request.rb2
-rw-r--r--lib/unicorn/tee_input.rb81
-rw-r--r--test/unit/test_tee_input.rb121
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