A POC for redirecting traffic from one Sauce Connect proxy to multiple upstream proxies

capybara_driver.rb

Capybara.register_driver :remote do |app|
      capabilities = {
        browserName: "Chrome",
        platform: "Windows 10",
        version: "60"
      }# {Put your Sauce Labs Desired Capabilities here as a hash}
      capabilities[:name] = scenario.full_description    # Give the Sauce Labs test a decent name
      
      client = Selenium::WebDriver::Remote::Http::Default.new
      client.open_timeout = 1500  # Optional; Deals with long session start times
      
      server_url = "https://#{ENV['SAUCE_USERNAME']}:#{ENV['SAUCE_ACCESS_KEY']}" +'@ondemand.saucelabs.com:443/wd/hub'

      proxy = {
        proxyType: "manual",
        sslProxy: "#{Billy.proxy.host}:#{Billy.proxy.port}",
        httpProxy: "#{Billy.proxy.host}:#{Billy.proxy.port}"
      }

      driver = SauceCapybara::Driver.new(
        app,
        browser: :remote,
        url: server_url,
        desired_capabilities: capabilities,
        http_client: client
      )

      server_uri = URI('http://localhost:4567/mapping')
      res = Net::HTTP.post_form(server_uri, 'job' => driver.browser.session_id, 'port' => Billy.proxy.port)

      driver
    end

id_port_proxy.rb

require "sinatra"
require "thread"
require 'webrick/httpproxy'
require_relative 'web_marble'
require_relative 'port_map'

class WebrickPortProxy  
  def initialize map, proxy
    @port_map = map
    @proxy = WEBrick::HTTPProxyServer.new(:Port => 8004, :ProxyURI => proxy)
    @proxy.port_map = @port_map
  end

  trap 'INT'  do @proxy.shutdown end
  trap 'TERM' do @proxy.shutdown end

  def start
    # puts "Starting Proxy"
    begin
      @proxy.start
    rescue => e
      puts "ERROR in thread"
      puts e
      puts "#{caller}"
    end
    # puts "Proxy Started"
  end
end


semaphore = Mutex.new
port_map = PortMap.new semaphore

proxyURI = URI.parse "http://localhost:1111"

Thread.new do 
  begin

    proxy = WebrickPortProxy.new port_map, proxyURI
    proxy.start
    # puts "Control Returned"
  rescue => e
    # puts "Thread issue"
    puts e
  end
end


get '/mapping' do
  template = "<h1>Job/Port Map</h1>"
  template += "<table><tr><th>Job ID</th><th>Puffing Billy Port</th></tr>"
    
  port_map.each_port do |id, port|
    template += "<tr><td>#{id}</td><td>#{port}</td></tr>"
  end

  template += "</table></br></br>"
  
  template += "<b>Total Rows:</b> #{port_map.size}"

  template
end


post '/mapping' do
  # puts "Storing a map between #{params['job']} and #{params['port']}"
  port_map.add(params['job'], params['port'])
end

delete '/mapping' do
  port_map.remove params['job']
end

port_map.rb

require "thread"

class PortMap

  def initialize semaphore
    @map = {}
    @semaphore = semaphore
  end

  def find_port job_id
    port = nil
    @semaphore.synchronize { port = @map[job_id.to_sym] }
    port
  end

  def add job_id, port
    @semaphore.synchronize { @map[job_id.to_sym] = port }
  end

  def remove job_id
    @semaphore.synchronize { @map.delete job_id }
  end

  def sizes
    @semaphore.synchronize { @map.size }
  end

  def each_port &block
    @semaphore.synchronize do
      @map.each_pair { |k, v| yield k, v }
    end
  end
end

web_marble.rb

  class WEBrick::HTTPProxyServer
    def proxy_uri(req, res)
      proxy = @config[:ProxyURI].dup

      job_id = req.header['x-sl-job-id']
      if job_id && !job_id.empty?
        puts "Checking port map #{@port_map}"
        redirect_port = @port_map.find_port(job_id)
          puts "Redirecting traffic for #{job_id} to port #{redirect_port}"
        proxy.port = redirect_port
        return proxy
      end
      # puts "This was not matched; Returning nil"
      # should return upstream proxy server's URI
      return nil
    end

    attr_accessor :port_map

    def port_map= pmap
      @port_map = pmap
    end

    def port_map
      @port_map
    end
  end