a

host.rb

require 'socket'
require 'securerandom'
require 'digest'
# require 'io/console'
require_relative './util'
# $> << '>'
password = 'password' #STDIN.noecho { STDIN.gets }.strip
server_host, server_port = ARGV[0].split(':')
local_port = ARGV[1]
socket = TCPSocket.new server_host, server_port.to_i
challenge = socket.gets.strip
socket.puts Digest::SHA256.hexdigest(challenge + password)

class Proxy
  def initialize(socket, host: 'localhost', port:)
    @socket = socket
    @host = host
    @port = port
    @sockets = {}
  end

  def start
    loop do
      op = @socket.gets
      next if op == "\n"
      id = op.to_i
      cmd = @socket.gets.strip
      if cmd == '+'
        connect id
      elsif cmd == '-'
        disconnect id
      else
        size = cmd.to_i
        data = read_byte @socket, size
        send_data id, data
      end
    end
  end

  def connect(id)
    return if @sockets.has_key? id
    socket = TCPSocket.new @host, @port
    @sockets[id] = socket
    Thread.new do
      loop do
        data = socket.readpartial 65536
        raise if data.nil? || data.empty?
        @socket.write "#{id}\n#{data.size}\n#{data}"
        @socket.flush
      end
    rescue => e
      @socket.puts "#{id}\n-"
      @sockets.delete id
    end
  rescue => e
    @socket.puts "#{id}\n-"
    @sockets.delete id
  end

  def send_data(id, data)
    socket = @sockets[id]
    socket.write data
    socket.flush
  rescue => e
    @socket.puts "#{id}\n-"
    @sockets[id]&.close rescue nil
    @sockets.delete id
  end

  def disconnect(id)
    @sockets[id]&.close rescue nil
    @sockets.delete id
  end
end

Proxy.new(socket, port: local_port).start

server.rb

require 'socket'
require 'securerandom'
require 'digest'
require_relative './util'
host_port = ENV['HOST_PORT']&.to_i || 9998
port = ENV['PORT']&.to_i || 9999
password = ENV['PASSWORD'] || 'password'

class HostProxy
  def initialize(port, password)
    @server = TCPServer.new port
    @password = password
    @sockets = {}
    @cnt = 0
  end

  def start
    loop do
      socket = @server.accept
      Thread.new do
        authorize socket
      end
    end
  end

  def connect(socket)
    unless @host_socket
      socket.close
      return nil
    end
    id = @cnt += 1
    @sockets[id] = socket
    @host_socket.puts "#{id}\n+"
    id
  end

  def disconnect(id)
    return unless @sockets.has_key?(id)
    @sockets[id].close rescue nil
    @sockets.delete id
    @host_socket.puts "#{id}\n-"
  end

  def send_to_host(id, data)
    return unless @sockets.has_key?(id)
    @host_socket.write "#{id}\n#{data.bytesize}\n#{data}"
    @host_socket.flush
  end

  def send_from_host(id, data)
    socket = @sockets[id]
    return unless socket
    socket.write data
    socket.flush
  rescue => e
    disconnect id
  end

  def authorize(socket)
    raise 'dup' if @host_socket
    challenge = SecureRandom.hex
    socket.puts challenge
    result = socket.gets.strip
    raise 'auth failed' unless result == Digest::SHA256.hexdigest(challenge + @password)
    handle socket
  rescue => e
    socket.close
  end

  def handle(socket)
    @host_socket = socket
    Thread.new do
      loop do
        sleep 10
        socket.puts
      end
    rescue => e
      socket.close
    end
    loop do
      dst = socket.gets.to_i
      cmd = socket.gets
      if cmd.strip == '-'
        disconnect dst
      else
        size = cmd.to_i
        data = read_byte socket, size
        send_from_host dst, data
      end
    end
  rescue => e
    close_all
    @host_socket = nil
  end

  def close_all
    @host_socket.close rescue nil
    @sockets.each do |_k, socket|
      socket.close rescue nil
    end
    @sockets = {}
  end
end

class Proxy
  def initialize(port, host_proxy)
    @server = TCPServer.new port
    @host = host_proxy
  end

  def start
    loop do
      socket = @server.accept
      Thread.new do
        handle socket
      end
    end
  end

  def handle socket
    id = @host.connect socket
    loop do
      data = socket.readpartial 65536
      raise if data.nil? || data.empty?
      @host.send_to_host id, data
    end
  rescue => e
    @host.disconnect id
  end
end

host_proxy = HostProxy.new host_port, password
Thread.new { host_proxy.start }
proxy = Proxy.new port, host_proxy
proxy.start

util.rb

def read_byte(io, size)
  data = ''
  while data.bytesize < size
    s = io.read(size - data.bytesize)
    raise if s.nil? || s.bytesize.zero?
    data << s
  end
  data
end