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 < Unicorn::StreamInput
  15   # The maximum size (in +bytes+) to buffer in memory before
  16   # resorting to a temporary file.  Default is 112 kilobytes.
  17   @@client_body_buffer_size = Unicorn::Const::MAX_BODY
  18
  19   # sets the maximum size of request bodies to buffer in memory,
  20   # amounts larger than this are buffered to the filesystem
  21   def self.client_body_buffer_size=(bytes)
  22     @@client_body_buffer_size = bytes
  23   end
  24
  25   # returns the maximum size of request bodies to buffer in memory,
  26   # amounts larger than this are buffered to the filesystem
  27   def self.client_body_buffer_size
  28     @@client_body_buffer_size
  29   end
  30
  31   # Initializes a new TeeInput object.  You normally do not have to call
  32   # this unless you are writing an HTTP server.
  33   def initialize(socket, request)
  34     @len = request.content_length
  35     super
  36     @tmp = @len && @len <= @@client_body_buffer_size ?
  37            StringIO.new("") : Unicorn::TmpIO.new
  38   end
  39
  40   # :call-seq:
  41   #   ios.size  => Integer
  42   #
  43   # Returns the size of the input.  For requests with a Content-Length
  44   # header value, this will not read data off the socket and just return
  45   # the value of the Content-Length header as an Integer.
  46   #
  47   # For Transfer-Encoding:chunked requests, this requires consuming
  48   # all of the input stream before returning since there's no other
  49   # way to determine the size of the request body beforehand.
  50   #
  51   # This method is no longer part of the Rack specification as of
  52   # Rack 1.2, so its use is not recommended.  This method only exists
  53   # for compatibility with Rack applications designed for Rack 1.1 and
  54   # earlier.  Most applications should only need to call +read+ with a
  55   # specified +length+ in a loop until it returns +nil+.
  56   def size
  57     @len and return @len
  58     pos = @tmp.pos
  59     consume!
  60     @tmp.pos = pos
  61     @len = @tmp.size
  62   end
  63
  64   # :call-seq:
  65   #   ios.read([length [, buffer ]]) => string, buffer, or nil
  66   #
  67   # Reads at most length bytes from the I/O stream, or to the end of
  68   # file if length is omitted or is nil. length must be a non-negative
  69   # integer or nil. If the optional buffer argument is present, it
  70   # must reference a String, which will receive the data.
  71   #
  72   # At end of file, it returns nil or "" depend on length.
  73   # ios.read() and ios.read(nil) returns "".
  74   # ios.read(length [, buffer]) returns nil.
  75   #
  76   # If the Content-Length of the HTTP request is known (as is the common
  77   # case for POST requests), then ios.read(length [, buffer]) will block
  78   # until the specified length is read (or it is the last chunk).
  79   # Otherwise, for uncommon "Transfer-Encoding: chunked" requests,
  80   # ios.read(length [, buffer]) will return immediately if there is
  81   # any data and only block when nothing is available (providing
  82   # IO#readpartial semantics).
  83   def read(*args)
  84     @socket ? tee(super) : @tmp.read(*args)
  85   end
  86
  87   # :call-seq:
  88   #   ios.gets   => string or nil
  89   #
  90   # Reads the next ``line'' from the I/O stream; lines are separated
  91   # by the global record separator ($/, typically "\n"). A global
  92   # record separator of nil reads the entire unread contents of ios.
  93   # Returns nil if called at the end of file.
  94   # This takes zero arguments for strict Rack::Lint compatibility,
  95   # unlike IO#gets.
  96   def gets
  97     @socket ? tee(super) : @tmp.gets
  98   end
  99
 100   # :call-seq:
 101   #   ios.rewind    => 0
 102   #
 103   # Positions the *ios* pointer to the beginning of input, returns
 104   # the offset (zero) of the +ios+ pointer.  Subsequent reads will
 105   # start from the beginning of the previously-buffered input.
 106   def rewind
 107     return 0 if 0 == @tmp.size
 108     consume! if @socket
 109     @tmp.rewind # Rack does not specify what the return value is here
 110   end
 111
 112 private
 113
 114   # consumes the stream of the socket
 115   def consume!
 116     junk = ""
 117     nil while read(@@io_chunk_size, junk)
 118   end
 119
 120   def tee(buffer)
 121     if buffer && buffer.size > 0
 122       @tmp.write(buffer)
 123       @tmp.seek(0, IO::SEEK_END) # workaround FreeBSD/OSX + MRI 1.8.x bug
 124     end
 125     buffer
 126   end
 127 end