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
  18     def initialize(input, size = nil, buffer = nil)
  19       @wr = Tempfile.new(nil)
  20       @wr.binmode
  21       @rd = File.open(@wr.path, 'rb')
  22       @wr.unlink
  23       @rd.sync = @wr.sync = true
  24
  25       @wr.write(buffer) if buffer
  26       @input = input
  27       @size = size # nil if chunked
  28     end
  29
  30     def consume
  31       @input or return
  32       buf = Z.dup
  33       while tee(Const::CHUNK_SIZE, buf)
  34       end
  35       self
  36     end
  37
  38     # returns the size of the input.  This is what the Content-Length
  39     # header value should be, and how large our input is expected to be.
  40     # For TE:chunked, this requires consuming all of the input stream
  41     # before returning since there's no other way
  42     def size
  43       @size and return @size
  44       @input and consume
  45       @size = @wr.stat.size
  46     end
  47
  48     def read(*args)
  49       @input or return @rd.read(*args)
  50
  51       length = args.shift
  52       if nil == length
  53         rv = @rd.read || Z.dup
  54         tmp = Z.dup
  55         while tee(Const::CHUNK_SIZE, tmp)
  56           rv << tmp
  57         end
  58         rv
  59       else
  60         buf = args.shift || Z.dup
  61         @rd.read(length, buf) || tee(length, buf)
  62       end
  63     end
  64
  65     # takes zero arguments for strict Rack::Lint compatibility, unlike IO#gets
  66     def gets
  67       @input or return @rd.gets
  68       nil == $/ and return read
  69
  70       line = nil
  71       if @rd.pos < @wr.stat.size
  72         line = @rd.gets # cannot be nil here
  73         $/ == line[-$/.size, $/.size] and return line
  74
  75         # half the line was already read, and the rest of has not been read
  76         if buf = @input.gets
  77           @wr.write(buf)
  78           line << buf
  79         else
  80           @input = nil
  81         end
  82       elsif line = @input.gets
  83         @wr.write(line)
  84       end
  85
  86       line
  87     end
  88
  89     def each(&block)
  90       while line = gets
  91         yield line
  92       end
  93
  94       self # Rack does not specify what the return value here
  95     end
  96
  97     def rewind
  98       @rd.rewind # Rack does not specify what the return value here
  99     end
 100
 101   private
 102
 103     # tees off a +length+ chunk of data from the input into the IO
 104     # backing store as well as returning it.  +buf+ must be specified.
 105     # returns nil if reading from the input returns nil
 106     def tee(length, buf)
 107       begin
 108         if @size
 109           left = @size - @rd.stat.size
 110           0 == left and return nil
 111           if length >= left
 112             @input.readpartial(left, buf) == left and @input = nil
 113           elsif @input.nil?
 114             return nil
 115           else
 116             @input.readpartial(length, buf)
 117           end
 118         else # ChunkedReader#readpartial just raises EOFError when done
 119           @input.readpartial(length, buf)
 120         end
 121       rescue EOFError
 122         return @input = nil
 123       end
 124       @wr.write(buf)
 125       buf
 126     end
 127
 128   end
 129 end