unicorn.git  about / heads / tags
Rack HTTP server for Unix and fast clients
blob bbc496b6f2ffb1283106ca5f97d8e1e4fe112cd7 3613 bytes (raw)
$ git show v0.9.2:lib/unicorn/tee_input.rb	# shows this blob on the CLI

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
 
# Copyright (c) 2009 Eric Wong
# You can redistribute it and/or modify it under the same terms as Ruby.

# acts like tee(1) on an input input to provide a input-like stream
# while providing rewindable semantics through a File/StringIO
# backing store.  On the first pass, the input is only read on demand
# so your Rack application can use input notification (upload progress
# and like).  This should fully conform to the Rack::InputWrapper
# specification on the public API.  This class is intended to be a
# strict interpretation of Rack::InputWrapper functionality and will
# not support any deviations from it.

module Unicorn
  class TeeInput

    def initialize(input, size, body)
      @tmp = Unicorn::Util.tmpio

      if body
        @tmp.write(body)
        @tmp.seek(0)
      end
      @input = input
      @size = size # nil if chunked
    end

    # returns the size of the input.  This is what the Content-Length
    # header value should be, and how large our input is expected to be.
    # For TE:chunked, this requires consuming all of the input stream
    # before returning since there's no other way
    def size
      @size and return @size

      if @input
        buf = Z.dup
        while tee(Const::CHUNK_SIZE, buf)
        end
        @tmp.rewind
      end

      @size = @tmp.stat.size
    end

    def read(*args)
      @input or return @tmp.read(*args)

      length = args.shift
      if nil == length
        rv = @tmp.read || Z.dup
        tmp = Z.dup
        while tee(Const::CHUNK_SIZE, tmp)
          rv << tmp
        end
        rv
      else
        buf = args.shift || Z.dup
        diff = @tmp.stat.size - @tmp.pos
        if 0 == diff
          tee(length, buf)
        else
          @tmp.read(diff > length ? length : diff, buf)
        end
      end
    end

    # takes zero arguments for strict Rack::Lint compatibility, unlike IO#gets
    def gets
      @input or return @tmp.gets
      nil == $/ and return read

      orig_size = @tmp.stat.size
      if @tmp.pos == orig_size
        tee(Const::CHUNK_SIZE, Z.dup) or return nil
        @tmp.seek(orig_size)
      end

      line = @tmp.gets # cannot be nil here since size > pos
      $/ == line[-$/.size, $/.size] and return line

      # unlikely, if we got here, then @tmp is at EOF
      begin
        orig_size = @tmp.stat.size
        tee(Const::CHUNK_SIZE, Z.dup) or break
        @tmp.seek(orig_size)
        line << @tmp.gets
        $/ == line[-$/.size, $/.size] and return line
        # @tmp is at EOF again here, retry the loop
      end while true

      line
    end

    def each(&block)
      while line = gets
        yield line
      end

      self # Rack does not specify what the return value here
    end

    def rewind
      @tmp.rewind # Rack does not specify what the return value here
    end

  private

    # tees off a +length+ chunk of data from the input into the IO
    # backing store as well as returning it.  +buf+ must be specified.
    # returns nil if reading from the input returns nil
    def tee(length, buf)
      begin
        if @size
          left = @size - @tmp.stat.size
          0 == left and return nil
          if length >= left
            @input.readpartial(left, buf) == left and @input = nil
          elsif @input.nil?
            return nil
          else
            @input.readpartial(length, buf)
          end
        else # ChunkedReader#readpartial just raises EOFError when done
          @input.readpartial(length, buf)
        end
      rescue EOFError
        return @input = nil
      end
      @tmp.write(buf)
      buf
    end

  end
end

git clone https://yhbt.net/unicorn.git