lib/unicorn/tee_input.rb

   1 # -*- encoding: binary -*-
   2
   3 module Unicorn
   4
   5   # acts like tee(1) on an input input to provide a input-like stream
   6   # while providing rewindable semantics through a File/StringIO
   7   # backing store.  On the first pass, the input is only read on demand
   8   # so your Rack application can use input notification (upload progress
   9   # and like).  This should fully conform to the Rack::InputWrapper
  10   # specification on the public API.  This class is intended to be a
  11   # strict interpretation of Rack::InputWrapper functionality and will
  12   # not support any deviations from it.
  13   class TeeInput < Struct.new(:socket, :req, :parser, :buf)
  14
  15     def initialize(*args)
  16       super(*args)
  17       @size = parser.content_length
  18       @tmp = @size && @size < Const::MAX_BODY ? StringIO.new("") : Util.tmpio
  19       @buf2 = buf.dup
  20       if buf.size > 0
  21         parser.filter_body(@buf2, buf) and finalize_input
  22         @tmp.write(@buf2)
  23         @tmp.seek(0)
  24       end
  25     end
  26
  27     # returns the size of the input.  This is what the Content-Length
  28     # header value should be, and how large our input is expected to be.
  29     # For TE:chunked, this requires consuming all of the input stream
  30     # before returning since there's no other way
  31     def size
  32       @size and return @size
  33
  34       if socket
  35         pos = @tmp.pos
  36         while tee(Const::CHUNK_SIZE, @buf2)
  37         end
  38         @tmp.seek(pos)
  39       end
  40
  41       @size = @tmp.size
  42     end
  43
  44     # call-seq:
  45     #   ios = env['rack.input']
  46     #   ios.read([length [, buffer ]]) => string, buffer, or nil
  47     #
  48     # Reads at most length bytes from the I/O stream, or to the end of
  49     # file if length is omitted or is nil. length must be a non-negative
  50     # integer or nil. If the optional buffer argument is present, it
  51     # must reference a String, which will receive the data.
  52     #
  53     # At end of file, it returns nil or "" depend on length.
  54     # ios.read() and ios.read(nil) returns "".
  55     # ios.read(length [, buffer]) returns nil.
  56     #
  57     # If the Content-Length of the HTTP request is known (as is the common
  58     # case for POST requests), then ios.read(length [, buffer]) will block
  59     # until the specified length is read (or it is the last chunk).
  60     # Otherwise, for uncommon "Transfer-Encoding: chunked" requests,
  61     # ios.read(length [, buffer]) will return immediately if there is
  62     # any data and only block when nothing is available (providing
  63     # IO#readpartial semantics).
  64     def read(*args)
  65       socket or return @tmp.read(*args)
  66
  67       length = args.shift
  68       if nil == length
  69         rv = @tmp.read || ""
  70         while tee(Const::CHUNK_SIZE, @buf2)
  71           rv << @buf2
  72         end
  73         rv
  74       else
  75         rv = args.shift || @buf2.dup
  76         diff = @tmp.size - @tmp.pos
  77         if 0 == diff
  78           ensure_length(tee(length, rv), length)
  79         else
  80           ensure_length(@tmp.read(diff > length ? length : diff, rv), length)
  81         end
  82       end
  83     end
  84
  85     # takes zero arguments for strict Rack::Lint compatibility, unlike IO#gets
  86     def gets
  87       socket or return @tmp.gets
  88       nil == $/ and return read
  89
  90       orig_size = @tmp.size
  91       if @tmp.pos == orig_size
  92         tee(Const::CHUNK_SIZE, @buf2) or return nil
  93         @tmp.seek(orig_size)
  94       end
  95
  96       line = @tmp.gets # cannot be nil here since size > pos
  97       $/ == line[-$/.size, $/.size] and return line
  98
  99       # unlikely, if we got here, then @tmp is at EOF
 100       begin
 101         orig_size = @tmp.pos
 102         tee(Const::CHUNK_SIZE, @buf2) or break
 103         @tmp.seek(orig_size)
 104         line << @tmp.gets
 105         $/ == line[-$/.size, $/.size] and return line
 106         # @tmp is at EOF again here, retry the loop
 107       end while true
 108
 109       line
 110     end
 111
 112     def each(&block)
 113       while line = gets
 114         yield line
 115       end
 116
 117       self # Rack does not specify what the return value is here
 118     end
 119
 120     def rewind
 121       @tmp.rewind # Rack does not specify what the return value is here
 122     end
 123
 124   private
 125
 126     def client_error(e)
 127       case e
 128       when EOFError
 129         # in case client only did a premature shutdown(SHUT_WR)
 130         # we do support clients that shutdown(SHUT_WR) after the
 131         # _entire_ request has been sent, and those will not have
 132         # raised EOFError on us.
 133         socket.close if socket
 134         raise ClientShutdown, "bytes_read=#{@tmp.size}", []
 135       when HttpParserError
 136         e.set_backtrace([])
 137         raise e
 138       end
 139     end
 140
 141     # tees off a +length+ chunk of data from the input into the IO
 142     # backing store as well as returning it.  +dst+ must be specified.
 143     # returns nil if reading from the input returns nil
 144     def tee(length, dst)
 145       unless parser.body_eof?
 146         if parser.filter_body(dst, socket.readpartial(length, buf)).nil?
 147           @tmp.write(dst)
 148           @tmp.seek(0, IO::SEEK_END) # workaround FreeBSD/OSX + MRI 1.8.x bug
 149           return dst
 150         end
 151       end
 152       finalize_input
 153       rescue => e
 154         client_error(e)
 155     end
 156
 157     def finalize_input
 158       while parser.trailers(req, buf).nil?
 159         # Don't worry about raising ClientShutdown here on EOFError, tee()
 160         # will catch EOFError when app is processing it, otherwise in
 161         # initialize we never get any chance to enter the app so the
 162         # EOFError will just get trapped by Unicorn and not the Rack app
 163         buf << socket.readpartial(Const::CHUNK_SIZE)
 164       end
 165       self.socket = nil
 166     end
 167
 168     # tee()s into +dst+ until it is of +length+ bytes (or until
 169     # we've reached the Content-Length of the request body).
 170     # Returns +dst+ (the exact object, not a duplicate)
 171     # To continue supporting applications that need near-real-time
 172     # streaming input bodies, this is a no-op for
 173     # "Transfer-Encoding: chunked" requests.
 174     def ensure_length(dst, length)
 175       # @size is nil for chunked bodies, so we can't ensure length for those
 176       # since they could be streaming bidirectionally and we don't want to
 177       # block the caller in that case.
 178       return dst if dst.nil? || @size.nil?
 179
 180       while dst.size < length && tee(length - dst.size, @buf2)
 181         dst << @buf2
 182       end
 183
 184       dst
 185     end
 186
 187   end
 188 end