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 by default.
  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 # :nodoc:
  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) # :nodoc:
  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 # :nodoc:
  28     @@client_body_buffer_size
  29   end
  30
  31   # for Rack::TempfileReaper in rack 1.6+
  32   def new_tmpio # :nodoc:
  33     tmpio = Unicorn::TmpIO.new
  34     (@parser.env['rack.tempfiles'] ||= []) << tmpio
  35     tmpio
  36   end
  37
  38   # Initializes a new TeeInput object.  You normally do not have to call
  39   # this unless you are writing an HTTP server.
  40   def initialize(socket, request) # :nodoc:
  41     @len = request.content_length
  42     super
  43     @tmp = @len && @len <= @@client_body_buffer_size ?
  44            StringIO.new("") : new_tmpio
  45   end
  46
  47   # :call-seq:
  48   #   ios.size  => Integer
  49   #
  50   # Returns the size of the input.  For requests with a Content-Length
  51   # header value, this will not read data off the socket and just return
  52   # the value of the Content-Length header as an Integer.
  53   #
  54   # For Transfer-Encoding:chunked requests, this requires consuming
  55   # all of the input stream before returning since there's no other
  56   # way to determine the size of the request body beforehand.
  57   #
  58   # This method is no longer part of the Rack specification as of
  59   # Rack 1.2, so its use is not recommended.  This method only exists
  60   # for compatibility with Rack applications designed for Rack 1.1 and
  61   # earlier.  Most applications should only need to call +read+ with a
  62   # specified +length+ in a loop until it returns +nil+.
  63   def size
  64     @len and return @len
  65     pos = @tmp.pos
  66     consume!
  67     @tmp.pos = pos
  68     @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 ? tee(super) : @tmp.read(*args)
  92   end
  93
  94   # :call-seq:
  95   #   ios.gets   => string or nil
  96   #
  97   # Reads the next ``line'' from the I/O stream; lines are separated
  98   # by the global record separator ($/, typically "\n"). A global
  99   # record separator of nil reads the entire unread contents of ios.
 100   # Returns nil if called at the end of file.
 101   # This takes zero arguments for strict Rack::Lint compatibility,
 102   # unlike IO#gets.
 103   def gets
 104     @socket ? tee(super) : @tmp.gets
 105   end
 106
 107   # :call-seq:
 108   #   ios.rewind    => 0
 109   #
 110   # Positions the *ios* pointer to the beginning of input, returns
 111   # the offset (zero) of the +ios+ pointer.  Subsequent reads will
 112   # start from the beginning of the previously-buffered input.
 113   def rewind
 114     return 0 if 0 == @tmp.size
 115     consume! if @socket
 116     @tmp.rewind # Rack does not specify what the return value is here
 117   end
 118
 119 private
 120
 121   # consumes the stream of the socket
 122   def consume!
 123     junk = ""
 124     nil while read(@@io_chunk_size, junk)
 125   end
 126
 127   def tee(buffer)
 128     @tmp.write(buffer) if buffer
 129     buffer
 130   end
 131 end