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