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