1 files changed, 145 insertions, 0 deletions
diff --git a/lib/unicorn/tee_input.rb b/lib/unicorn/tee_input.rb
new file mode 100644
index 0000000..1e281b4
--- /dev/null
+++ b/lib/unicorn/tee_input.rb
@@ -0,0 +1,145 @@
+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
+    Z = ''
+    Z.force_encoding(Encoding::BINARY) if Z.respond_to?(:force_encoding)
+
+    def initialize
+      @rd = @wr = @size = @input = nil
+      setup
+    end
+
+    def setup
+      @tmp = tmp = Tempfile.new(nil)
+      @rd.close if @rd
+      @rd = File.open(tmp.path, 'wb+')
+      @wr.close if @wr
+      @wr = File.open(tmp.path, 'wb')
+      @rd.sync = @wr.sync = true
+
+      def @rd.size
+        stat.size
+      end
+      tmp.close!
+    end
+
+    def reopen(input, size = nil, buffer = nil)
+      @rd.seek(0)
+      @wr.seek(0)
+      @rd.truncate(0) # truncate read to flush luserspace read buffers
+      @wr.write(buffer) if buffer
+      @input = input
+      @size = size # nil if chunked
+      self
+    end
+
+    def consume
+      @input or return
+      buf = Z.dup
+      while tee(Const::CHUNK_SIZE, buf)
+      end
+      @rd
+    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
+      @input and consume
+      @size = @wr.stat.size
+    end
+
+    def read(*args)
+      @input or return @rd.read(*args)
+
+      length = args.shift
+      if nil == length
+        rv = @rd.read || Z.dup
+        tmp = Z.dup
+        while tee(Const::CHUNK_SIZE, tmp)
+          rv << tmp
+        end
+        rv
+      else
+        buf = args.shift || Z.dup
+        @rd.read(length, buf) || tee(length, buf)
+      end
+    end
+
+    # takes zero arguments for strict Rack::Lint compatibility, unlike IO#gets
+    def gets
+      @input or return @rd.gets
+      nil == $/ and return read
+
+      line = nil
+      if @rd.pos < @wr.stat.size
+        line = @rd.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
+          @wr.write(buf)
+          line << buf
+        else
+          @input = nil
+        end
+      elsif line = @input.gets
+        @wr.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
+      @rd.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 - @rd.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
+      @wr.write(buf)
+      buf
+    end
+
+  end
+end