diff options
-rw-r--r-- | lib/unicorn/http_request.rb | 2 | ||||
-rw-r--r-- | lib/unicorn/tee_input.rb | 81 | ||||
-rw-r--r-- | test/unit/test_tee_input.rb | 121 |
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 |