rack-streaming-proxy example code

rack-streaming-proxy.rb

require "servolux"
require "net/http"
require "uri"

# see: http://github.com/aniero/rack-streaming-proxy for the latest code
# or: sudo gem install rack-streaming-proxy --source http://gemcutter.org

module Rack
  class StreamingProxy

    class Error < StandardError; end

    # The block provided to the initializer is given a Rack::Request
    # and should return:
    #
    #   * nil/false to skip the proxy and continue down the stack
    #   * a complete uri (with query string if applicable) to proxy to
    #
    # E.g.
    #
    #   use Rack::StreamingProxy do |req|
    #     if req.path.start_with?("/search")
    #       "http://some_other_service/search?#{req.query}"
    #     end
    #   end
    #
    # Most headers, request body, and HTTP method are preserved.
    #
    def initialize(app, &block)
      @request_uri = block
      @app = app
    end

    def call(env)
      req = Rack::Request.new(env)
      return app.call(env) unless uri = request_uri.call(req)
      proxy = ProxyRequest.new(req, uri)
      [proxy.status, proxy.headers, proxy]
    rescue => e
      msg = "Proxy error when proxying to #{uri}: #{e.class}: #{e.message}"
      env["rack.errors"].puts msg
      env["rack.errors"].puts e.backtrace.map { |l| "\t" + l }
      env["rack.errors"].flush
      raise Error, msg
    end

    protected

    attr_reader :request_uri, :app

    public

    class ProxyRequest
      include Rack::Utils

      attr_reader :status, :headers

      def initialize(request, uri)
        uri = URI.parse(uri)

        method = request.request_method.downcase
        method[0..0] = method[0..0].upcase

        proxy_request = Net::HTTP.const_get(method).new("#{uri.path}#{"?" if uri.query}#{uri.query}")

        if proxy_request.request_body_permitted? and request.body
          proxy_request.body_stream = request.body
          proxy_request.content_length = request.content_length
          proxy_request.content_type = request.content_type
        end

      %w(Accept Accept-Encoding Accept-Charset
        X-Requested-With Referer User-Agent Cookie).each do |header|
        key = "HTTP_#{header.upcase.gsub('-', '_')}"
        proxy_request[header] = request.env[key] if request.env[key]
        end
      proxy_request["X-Forwarded-For"] =
        (request.env["X-Forwarded-For"].to_s.split(/, +/) + [request.env["REMOTE_ADDR"]]).join(", ")
      proxy_request.basic_auth(*uri.userinfo.split(':')) if (uri.userinfo && uri.userinfo.index(':'))

      @piper = Servolux::Piper.new 'r', :timeout => 30

      @piper.child do
        Net::HTTP.start(uri.host, uri.port) do |http|
          http.request(proxy_request) do |response|
            # at this point the headers and status are available, but the body
            # has not yet been read. start reading it and putting it in the parent's pipe.
            response_headers = {}
            response.each_header {|k,v| response_headers[k] = v}
            @piper.puts [response.code.to_i, response_headers]

            response.read_body do |chunk|
              @piper.puts chunk
            end
            @piper.puts :done
          end
        end

        exit!
      end

      @piper.parent do
        # wait for the status and headers to come back from the child
        @status, @headers = @piper.gets
        @headers = HeaderHash.new(@headers)
      end
      end

      def each
        chunked = @headers["Transfer-Encoding"] == "chunked"
        term = "\r\n"

        while chunk = @piper.gets
          break if chunk == :done
          if chunked
            size = bytesize(chunk)
            next if size == 0
            yield [size.to_s(16), term, chunk, term].join
          else
            yield chunk
          end
        end

        yield ["0", term, "", term].join if chunked
      end

    end # ProxyRequest

  end

end