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 (future) 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(length = nil, rv = '')
  43     if length
  44       if length <= @rbuf.size
  45         length < 0 and raise ArgumentError, "negative length #{length} given"
  46         rv.replace(@rbuf.slice!(0, length))
  47       else
  48         to_read = length - @rbuf.size
  49         rv.replace(@rbuf.slice!(0, @rbuf.size))
  50         until to_read == 0 || eof? || (rv.size > 0 && @chunked)
  51           @socket.kgio_read(to_read, @buf) or eof!
  52           filter_body(@rbuf, @buf)
  53           rv << @rbuf
  54           to_read -= @rbuf.size
  55         end
  56         @rbuf.replace('')
  57       end
  58       rv = nil if rv.empty? && length != 0
  59     else
  60       read_all(rv)
  61     end
  62     rv
  63   end
  64
  65   # :call-seq:
  66   #   ios.gets   => string or nil
  67   #
  68   # Reads the next ``line'' from the I/O stream; lines are separated
  69   # by the global record separator ($/, typically "\n"). A global
  70   # record separator of nil reads the entire unread contents of ios.
  71   # Returns nil if called at the end of file.
  72   # This takes zero arguments for strict Rack::Lint compatibility,
  73   # unlike IO#gets.
  74   def gets
  75     sep = $/
  76     if sep.nil?
  77       read_all(rv = '')
  78       return rv.empty? ? nil : rv
  79     end
  80     re = /\A(.*?#{Regexp.escape(sep)})/
  81
  82     begin
  83       @rbuf.sub!(re, '') and return $1
  84       return @rbuf.empty? ? nil : @rbuf.slice!(0, @rbuf.size) if eof?
  85       @socket.kgio_read(@@io_chunk_size, @buf) or eof!
  86       filter_body(once = '', @buf)
  87       @rbuf << once
  88     end while true
  89   end
  90
  91   # :call-seq:
  92   #   ios.each { |line| block }  => ios
  93   #
  94   # Executes the block for every ``line'' in *ios*, where lines are
  95   # separated by the global record separator ($/, typically "\n").
  96   def each
  97     while line = gets
  98       yield line
  99     end
 100
 101     self # Rack does not specify what the return value is here
 102   end
 103
 104 private
 105
 106   def eof?
 107     if @parser.body_eof?
 108       while @chunked && ! @parser.parse
 109         once = @socket.kgio_read(@@io_chunk_size) or eof!
 110         @buf << once
 111       end
 112       @socket = nil
 113       true
 114     else
 115       false
 116     end
 117   end
 118
 119   def filter_body(dst, src)
 120     rv = @parser.filter_body(dst, src)
 121     @bytes_read += dst.size
 122     rv
 123   end
 124
 125   def read_all(dst)
 126     dst.replace(@rbuf)
 127     @socket or return
 128     until eof?
 129       @socket.kgio_read(@@io_chunk_size, @buf) or eof!
 130       filter_body(@rbuf, @buf)
 131       dst << @rbuf
 132     end
 133     ensure
 134       @rbuf.replace('')
 135   end
 136
 137   def eof!
 138     # in case client only did a premature shutdown(SHUT_WR)
 139     # we do support clients that shutdown(SHUT_WR) after the
 140     # _entire_ request has been sent, and those will not have
 141     # raised EOFError on us.
 142     @socket.shutdown if @socket
 143   ensure
 144     raise Unicorn::ClientShutdown, "bytes_read=#{@bytes_read}", []
 145   end
 146 end