lib/unicorn/tee_input.rb

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