lib/unicorn/tee_input.rb

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