lib/unicorn/stream_input.rb

   1 # -*- encoding: binary -*-
   2
   3 # When processing uploads, Unicorn may expose a StreamInput object under
   4 # "rack.input" of the Rack (2.x) environment.
   5 class Unicorn::StreamInput
   6   # The I/O chunk size (in +bytes+) for I/O operations where
   7   # the size cannot be user-specified when a method is called.
   8   # The default is 16 kilobytes.
   9   @@io_chunk_size = Unicorn::Const::CHUNK_SIZE
  10
  11   # Initializes a new StreamInput object.  You normally do not have to call
  12   # this unless you are writing an HTTP server.
  13   def initialize(socket, request)
  14     @chunked = request.content_length.nil?
  15     @socket = socket
  16     @parser = request
  17     @buf = request.buf
  18     @rbuf = ''
  19     @bytes_read = 0
  20     filter_body(@rbuf, @buf) unless @buf.empty?
  21   end
  22
  23   # :call-seq:
  24   #   ios.read([length [, buffer ]]) => string, buffer, or nil
  25   #
  26   # Reads at most length bytes from the I/O stream, or to the end of
  27   # file if length is omitted or is nil. length must be a non-negative
  28   # integer or nil. If the optional buffer argument is present, it
  29   # must reference a String, which will receive the data.
  30   #
  31   # At end of file, it returns nil or '' depend on length.
  32   # ios.read() and ios.read(nil) returns ''.
  33   # ios.read(length [, buffer]) returns nil.
  34   #
  35   # If the Content-Length of the HTTP request is known (as is the common
  36   # case for POST requests), then ios.read(length [, buffer]) will block
  37   # until the specified length is read (or it is the last chunk).
  38   # Otherwise, for uncommon "Transfer-Encoding: chunked" requests,
  39   # ios.read(length [, buffer]) will return immediately if there is
  40   # any data and only block when nothing is available (providing
  41   # IO#readpartial semantics).
  42   def read(*args)
  43     length = args.shift
  44     rv = args.shift || ''
  45     if length.nil?
  46       read_all(rv)
  47     else
  48       if length <= @rbuf.size
  49         rv.replace(@rbuf.slice(0, length))
  50         @rbuf.replace(@rbuf.slice(length, @rbuf.size) || '')
  51       else
  52         rv.replace(@rbuf)
  53         length -= @rbuf.size
  54         @rbuf.replace('')
  55         until length == 0 || eof? || (rv.size > 0 && @chunked)
  56           @socket.kgio_read(length, @buf) or eof!
  57           filter_body(@rbuf, @buf)
  58           rv << @rbuf
  59           length -= @rbuf.size
  60           @rbuf.replace('')
  61         end
  62       end
  63       rv = nil if rv.empty? && length != 0
  64     end
  65     rv
  66   end
  67
  68   # :call-seq:
  69   #   ios.gets   => string or nil
  70   #
  71   # Reads the next ``line'' from the I/O stream; lines are separated
  72   # by the global record separator ($/, typically "\n"). A global
  73   # record separator of nil reads the entire unread contents of ios.
  74   # Returns nil if called at the end of file.
  75   # This takes zero arguments for strict Rack::Lint compatibility,
  76   # unlike IO#gets.
  77   def gets
  78     sep = $/
  79     if sep.nil?
  80       read_all(rv = '')
  81       return rv.empty? ? nil : rv
  82     end
  83     re = /\A(.*?#{Regexp.escape(sep)})/
  84
  85     begin
  86       @rbuf.gsub!(re, '') and return $1
  87       if eof?
  88         if @rbuf.empty?
  89           return nil
  90         else
  91           rv = @rbuf.dup
  92           @rbuf.replace('')
  93           return rv
  94         end
  95       end
  96       @socket.kgio_read(@@io_chunk_size, @buf) or eof!
  97       filter_body(once = '', @buf)
  98       @rbuf << once
  99     end while true
 100   end
 101
 102   # :call-seq:
 103   #   ios.each { |line| block }  => ios
 104   #
 105   # Executes the block for every ``line'' in *ios*, where lines are
 106   # separated by the global record separator ($/, typically "\n").
 107   def each(&block)
 108     while line = gets
 109       yield line
 110     end
 111
 112     self # Rack does not specify what the return value is here
 113   end
 114
 115 private
 116
 117   def eof?
 118     if @parser.body_eof?
 119       until @parser.parse
 120         once = @socket.kgio_read(@@io_chunk_size) or eof!
 121         @buf << once
 122       end
 123       @socket = nil
 124       true
 125     else
 126       false
 127     end
 128   end
 129
 130   def filter_body(dst, src)
 131     rv = @parser.filter_body(dst, src)
 132     @bytes_read += dst.size
 133     rv
 134   end
 135
 136   def read_all(dst)
 137     dst.replace(@rbuf)
 138     @socket or return
 139     until eof?
 140       @socket.kgio_read(@@io_chunk_size, @buf) or eof!
 141       filter_body(@rbuf, @buf)
 142       dst << @rbuf
 143     end
 144     ensure
 145       @rbuf.replace('')
 146   end
 147
 148   def eof!
 149     # in case client only did a premature shutdown(SHUT_WR)
 150     # we do support clients that shutdown(SHUT_WR) after the
 151     # _entire_ request has been sent, and those will not have
 152     # raised EOFError on us.
 153     @socket.close if @socket
 154     raise Unicorn::ClientShutdown, "bytes_read=#{@bytes_read}", []
 155   end
 156 end