yahns.git  about / heads / tags
sleepy, multi-threaded, non-blocking application server for Ruby
blob 40bf19a2af7cf55177ff73a70e9faf89502f9e0c 6145 bytes (raw)
$ git show HEAD:extras/proxy_pass.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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
 
# -*- encoding: binary -*-
# Copyright (C) 2013-2016 all contributors <yahns-public@yhbt.net>
# License: GPL-3.0+ (https://www.gnu.org/licenses/gpl-3.0.txt)
# frozen_string_literal: true
require 'time'
require 'socket'
require 'kgio'
require 'kcar' # gem install kcar
require 'rack/request'
require 'thread'
require 'timeout'

# Totally synchronous and Rack 1.1-compatible.  See Yahns::ProxyPass for
# the rewritten version which takes advantage of rack.hijack and uses
# the internal non-blocking I/O facilities in yahns.  yahns may have to
# grow a supported API for that...
#
# For now, we this blocks a worker thread; fortunately threads are reasonably
# cheap on GNU/Linux...
class ProxyPass # :nodoc:
  class ConnPool
    def initialize
      @mtx = Mutex.new
      @objs = []
    end

    def get
      @mtx.synchronize { @objs.pop }
    end

    def put(obj)
      @mtx.synchronize { @objs << obj }
    end
  end

  class UpstreamSocket < Kgio::Socket # :nodoc:
    attr_writer :expiry

    # called automatically by kgio_read!
    def wait_readable(timeout = nil)
      super(timeout || wait_time)
    end

    def wait_time
      tout = @expiry ? @expiry - Time.now : @timeout
      raise Timeout::Error, "request timed out", [] if tout < 0
      tout
    end

    def readpartial(bytes, buf = Thread.current[:proxy_pass_buf] ||= ''.dup)
      case rv = kgio_read!(bytes, buf)
      when String
        @expiry += @timeout # bump expiry when we succeed
      end
      rv
    end

    def req_write(buf, timeout)
      @timeout = timeout
      @expiry = Time.now + timeout
      case rv = kgio_trywrite(buf)
      when :wait_writable
        wait_writable(wait_time)
      when nil
        return
      when String
        buf = rv
      end while true
    end
  end # class UpstreamSocket

  class UpstreamResponse < Kcar::Response # :nodoc:
    # Called by the Rack server at the end of a successful response
    def close
      reusable = @parser.keepalive? && @parser.body_eof?
      super
      @pool.put(self) if reusable
      nil
    end

    # req is just a string buffer of HTTP headers
    def req_write(req, timeout)
      @sock.req_write(req, timeout)
    end

    # returns true if the socket is still alive, nil if dead
    def sock_alive?
      @reused = (:wait_readable == (@sock.kgio_tryread(1) rescue nil)) ?
                true : @sock.close
    end

    # returns true if the socket was reused and thus retryable
    def fail_retryable?
      @sock.close
      @reused
    end

    def initialize(sock, pool)
      super(sock)
      @reused = false
      @pool = pool
    end
  end # class UpstreamResponse

  # take a responder from the pool, we'll add the object back to the
  # pool in UpstreamResponse#close
  def responder_get
    while obj = @pool.get
      return obj if obj.sock_alive?
    end

    UpstreamResponse.new(UpstreamSocket.start(@sockaddr), @pool)
  end

  def initialize(dest, timeout = 5)
    case dest
    when %r{\Aunix:([^:]+)(?::(/.*))?\z}
      path = $2
      @sockaddr = Socket.sockaddr_un($1)
    when %r{\Ahttp://([^/]+)(/.*)?\z}
      path = $2
      host, port = $1.split(':')
      @sockaddr = Socket.sockaddr_in(port || 80, host)
    else
      raise ArgumentError, "destination must be an HTTP URL or unix: path"
    end
    init_path_vars(path)
    @pool = ConnPool.new
    @timeout = timeout
  end

  def init_path_vars(path)
    path ||= '$fullpath'
    # methods from Rack::Request we want:
    allow = %w(fullpath host_with_port host port url path)
    want = path.scan(/\$(\w+)/).flatten! || []
    diff = want - allow
    diff.empty? or
             raise ArgumentError, "vars not allowed: #{diff.uniq.join(' ')}"

    # kill leading slash just in case...
    @path = path.gsub(%r{\A/(\$(?:fullpath|path))}, '\1')
  end

  def call(env)
    request_method = env['REQUEST_METHOD']
    req = Rack::Request.new(env)
    path = @path.gsub(/\$(\w+)/) { req.__send__($1) }
    req = "#{request_method} #{path} HTTP/1.1\r\n" \
          "X-Forwarded-For: #{env["REMOTE_ADDR"]}\r\n".dup

    # pass most HTTP_* headers through as-is
    chunked = false
    env.each do |key, val|
      %r{\AHTTP_(\w+)\z} =~ key or next
      key = $1
      next if %r{\A(?:VERSION|CONNECTION|KEEP_ALIVE|X_FORWARDED_FOR)} =~ key
      chunked = true if %r{\ATRANSFER_ENCODING} =~ key && val =~ /\bchunked\b/i
      key.tr!("_", "-")
      req << "#{key}: #{val}\r\n"
    end

    # special cases which Rack does not prefix:
    ctype = env["CONTENT_TYPE"] and req << "Content-Type: #{ctype}\r\n"
    clen = env["CONTENT_LENGTH"] and req << "Content-Length: #{clen}\r\n"
    req << "\r\n"

    # get an open socket and send the headers
    ures = responder_get
    ures.req_write(req, @timeout)

    # send the request body if there was one
    send_body(env["rack.input"], ures, chunked) if chunked || clen

    # wait for the response here
    _, header, body = res = ures.rack

    # don't let the upstream Connection and Keep-Alive headers leak through
    header.delete_if do |k,_|
      k =~ /\A(?:Connection|Keep-Alive)\z/i
    end

    case request_method
    when "HEAD"
      # kcar doesn't know if it's a HEAD or GET response, and HEAD
      # responses have Content-Length in it which fools kcar...
      body.parser.body_bytes_left = 0
      res[1] = header.dup
      body.close # clobbers original header
      res[2] = body = []
    end
    res
  rescue => e
    retry if ures && ures.fail_retryable? && request_method != "POST"
    if defined?(Yahns::Log)
      logger = env['rack.logger'] and
        Yahns::Log.exception(logger, 'proxy_pass', e)
    end
    [ 502, { 'Content-Length' => '0', 'Content-Type' => 'text/plain' }, [] ]
  end

  def send_body(input, ures, chunked)
    buf = Thread.current[:proxy_pass_buf] ||= ''.dup

    if chunked # unlikely
      while input.read(16384, buf)
        buf.replace("#{buf.size.to_s(16)}\r\n#{buf}\r\n")
        ures.req_write(buf, @timeout)
      end
      ures.req_write("0\r\n\r\n", @timeout)
    else # common if we hit uploads
      while input.read(16384, buf)
        ures.req_write(buf, @timeout)
      end
    end
  end
end

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