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