lib/unicorn/tee_input.rb

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