diff options
Diffstat (limited to 'lib/unicorn/tee_input.rb')
-rw-r--r-- | lib/unicorn/tee_input.rb | 135 |
1 files changed, 135 insertions, 0 deletions
diff --git a/lib/unicorn/tee_input.rb b/lib/unicorn/tee_input.rb new file mode 100644 index 0000000..06028a6 --- /dev/null +++ b/lib/unicorn/tee_input.rb @@ -0,0 +1,135 @@ +# Copyright (c) 2009 Eric Wong +# You can redistribute it and/or modify it under the same terms as Ruby. + +require 'tempfile' + +# acts like tee(1) on an input input to provide a input-like stream +# while providing rewindable semantics through a Tempfile/StringIO +# backing store. On the first pass, the input is only read on demand +# so your Rack application can use input notification (upload progress +# and like). This should fully conform to the Rack::InputWrapper +# specification on the public API. This class is intended to be a +# strict interpretation of Rack::InputWrapper functionality and will +# not support any deviations from it. + +module Unicorn + class TeeInput + + def initialize(input, size, body) + @tmp = Tempfile.new(nil) + @tmp.unlink + @tmp.binmode + @tmp.sync = true + + if body + @tmp.write(body) + @tmp.seek(0) + end + @input = input + @size = size # nil if chunked + end + + # returns the size of the input. This is what the Content-Length + # header value should be, and how large our input is expected to be. + # For TE:chunked, this requires consuming all of the input stream + # before returning since there's no other way + def size + @size and return @size + + if @input + buf = Z.dup + while tee(Const::CHUNK_SIZE, buf) + end + @tmp.rewind + end + + @size = @tmp.stat.size + end + + def read(*args) + @input or return @tmp.read(*args) + + length = args.shift + if nil == length + rv = @tmp.read || Z.dup + tmp = Z.dup + while tee(Const::CHUNK_SIZE, tmp) + rv << tmp + end + rv + else + buf = args.shift || Z.dup + diff = @tmp.stat.size - @tmp.pos + if 0 == diff + tee(length, buf) + else + @tmp.read(diff > length ? length : diff, buf) + end + end + end + + # takes zero arguments for strict Rack::Lint compatibility, unlike IO#gets + def gets + @input or return @tmp.gets + nil == $/ and return read + + line = nil + if @tmp.pos < @tmp.stat.size + line = @tmp.gets # cannot be nil here + $/ == line[-$/.size, $/.size] and return line + + # half the line was already read, and the rest of has not been read + if buf = @input.gets + @tmp.write(buf) + line << buf + else + @input = nil + end + elsif line = @input.gets + @tmp.write(line) + end + + line + end + + def each(&block) + while line = gets + yield line + end + + self # Rack does not specify what the return value here + end + + def rewind + @tmp.rewind # Rack does not specify what the return value here + end + + private + + # 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) + begin + if @size + left = @size - @tmp.stat.size + 0 == left and return nil + if length >= left + @input.readpartial(left, buf) == left and @input = nil + elsif @input.nil? + return nil + else + @input.readpartial(length, buf) + end + else # ChunkedReader#readpartial just raises EOFError when done + @input.readpartial(length, buf) + end + rescue EOFError + return @input = nil + end + @tmp.write(buf) + buf + end + + end +end |