lib/unicorn/tee_input.rb

   1 # Copyright (c) 2009 Eric Wong
   2 # You can redistribute it and/or modify it under the same terms as Ruby.
   3
   4 require 'tempfile'
   5
   6 # acts like tee(1) on an input input to provide a input-like stream
   7 # while providing rewindable semantics through a Tempfile/StringIO
   8 # backing store.  On the first pass, the input is only read on demand
   9 # so your Rack application can use input notification (upload progress
  10 # and like).  This should fully conform to the Rack::InputWrapper
  11 # specification on the public API.  This class is intended to be a
  12 # strict interpretation of Rack::InputWrapper functionality and will
  13 # not support any deviations from it.
  14
  15 module Unicorn
  16   class TeeInput
  17     Z = ''
  18     Z.force_encoding(Encoding::BINARY) if Z.respond_to?(:force_encoding)
  19
  20     def initialize
  21       @rd = @wr = @size = @input = nil
  22       setup
  23     end
  24
  25     def setup
  26       @tmp = tmp = Tempfile.new(nil)
  27       @rd.close if @rd
  28       @rd = File.open(tmp.path, 'wb+')
  29       @wr.close if @wr
  30       @wr = File.open(tmp.path, 'wb')
  31       @rd.sync = @wr.sync = true
  32
  33       def @rd.size
  34         stat.size
  35       end
  36       tmp.close!
  37     end
  38
  39     def reopen(input, size = nil, buffer = nil)
  40       @rd.seek(0)
  41       @wr.seek(0)
  42       @rd.truncate(0) # truncate read to flush luserspace read buffers
  43       @wr.write(buffer) if buffer
  44       @input = input
  45       @size = size # nil if chunked
  46       self
  47     end
  48
  49     def consume
  50       @input or return
  51       buf = Z.dup
  52       while tee(Const::CHUNK_SIZE, buf)
  53       end
  54       @rd
  55     end
  56
  57     # returns the size of the input.  This is what the Content-Length
  58     # header value should be, and how large our input is expected to be.
  59     # For TE:chunked, this requires consuming all of the input stream
  60     # before returning since there's no other way
  61     def size
  62       @size and return @size
  63       @input and consume
  64       @size = @wr.stat.size
  65     end
  66
  67     def read(*args)
  68       @input or return @rd.read(*args)
  69
  70       length = args.shift
  71       if nil == length
  72         rv = @rd.read || Z.dup
  73         tmp = Z.dup
  74         while tee(Const::CHUNK_SIZE, tmp)
  75           rv << tmp
  76         end
  77         rv
  78       else
  79         buf = args.shift || Z.dup
  80         @rd.read(length, buf) || tee(length, buf)
  81       end
  82     end
  83
  84     # takes zero arguments for strict Rack::Lint compatibility, unlike IO#gets
  85     def gets
  86       @input or return @rd.gets
  87       nil == $/ and return read
  88
  89       line = nil
  90       if @rd.pos < @wr.stat.size
  91         line = @rd.gets # cannot be nil here
  92         $/ == line[-$/.size, $/.size] and return line
  93
  94         # half the line was already read, and the rest of has not been read
  95         if buf = @input.gets
  96           @wr.write(buf)
  97           line << buf
  98         else
  99           @input = nil
 100         end
 101       elsif line = @input.gets
 102         @wr.write(line)
 103       end
 104
 105       line
 106     end
 107
 108     def each(&block)
 109       while line = gets
 110         yield line
 111       end
 112
 113       self # Rack does not specify what the return value here
 114     end
 115
 116     def rewind
 117       @rd.rewind # Rack does not specify what the return value here
 118     end
 119
 120   private
 121
 122     # tees off a +length+ chunk of data from the input into the IO
 123     # backing store as well as returning it.  +buf+ must be specified.
 124     # returns nil if reading from the input returns nil
 125     def tee(length, buf)
 126       begin
 127         if @size
 128           left = @size - @rd.stat.size
 129           0 == left and return nil
 130           if length >= left
 131             @input.readpartial(left, buf) == left and @input = nil
 132           elsif @input.nil?
 133             return nil
 134           else
 135             @input.readpartial(length, buf)
 136           end
 137         else # ChunkedReader#readpartial just raises EOFError when done
 138           @input.readpartial(length, buf)
 139         end
 140       rescue EOFError
 141         return @input = nil
 142       end
 143       @wr.write(buf)
 144       buf
 145     end
 146
 147   end
 148 end