From 60f9bcd4ba54d84a4b2e2cd621d5995ddf206ace Mon Sep 17 00:00:00 2001 From: Eric Wong Date: Fri, 14 Aug 2009 21:03:40 -0700 Subject: 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. --- lib/unicorn/http_request.rb | 2 +- lib/unicorn/tee_input.rb | 81 ++++++++++++++--------------- 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 -- cgit v1.2.3-24-ge0c7