crystal-lang/crystal

require "http"

class Wakr::Handlers::ProxyWebSocket
  include HTTP::Handler

  def initialize(@backend_host : String, @backend_port : Int32)
  end

  def call(context) : Nil
    request = context.request
    response = context.response

    return call_next context unless websocket_upgrade_request? request

    version = request.headers["Sec-WebSocket-Version"]?
    unless version && version == HTTP::WebSocket::Protocol::VERSION
      response.status = :upgrade_required
      response.headers["Sec-WebSocket-Version"] = HTTP::WebSocket::Protocol::VERSION
      Log.error { "Invalid websocket version: `wss://#{request.headers["Host"]?}`#{request.path}" }
      return
    end

    key = request.headers["Sec-WebSocket-Key"]?

    unless key
      response.respond_with_status(:bad_request)
      Log.error { "Invalid request, missing key: `wss://#{request.headers["Host"]?}`#{request.path}" }
      return
    end

    Log.info { "Proxying websocket session: `wss://#{request.headers["Host"]?}`#{request.path}" }
    accept_code = HTTP::WebSocket::Protocol.key_challenge(key)
    response.status = :switching_protocols
    response.headers["Upgrade"] = "websocket"
    response.headers["Connection"] = "Upgrade"
    response.headers["Sec-WebSocket-Accept"] = accept_code
    response.upgrade do |io|
      clients = [
        get_backend_client(request),
        HTTP::WebSocket.new(io, sync_close: false),
      ]
      clients.permutations(2).map do |(source, dest)|
        Log.info { "bonding #{source} and #{dest} "}

        channel = Channel(Bool).new(1)

        # Proxy Ping
        source.on_ping do |message|
          Log.info { "PING" }
          dest.ping(message)
        end

        # Proxy Pong
        source.on_pong do |message|
          Log.info { "PONG" }
          dest.pong(message)
        end

        # Proxy Message
        source.on_message do |message|
          Log.info { "MESSAGE" }
          dest.send(message)
        end

        # Proxy Close
        source.on_close do |code, message|
          Log.info { "CLOSE" }
          dest.close(code, message)
          channel.send true
        end

        # Proxy Binary
        source.on_binary do |bytes|
          Log.info { "BINARY" }
          dest.stream do |stream_io|
            stream_io.write bytes
          end
        end

        # Run the source socket
        spawn {
          Log.info { "running #{source}" }
          source.run
        }

        # Return a waiting function
        ->{
          Fiber.yield
          Log.info { "waiting for close: #{source}" }
          channel.receive
          Log.info { "closed: #{source}" }
        }
      end.each(&.call)
    end
  end

  private def websocket_upgrade_request?(request)
    return false unless upgrade = request.headers["Upgrade"]?
    return false unless upgrade.compare("websocket", case_insensitive: true) == 0

    request.headers.includes_word?("Connection", "Upgrade")
  end

  private def get_backend_client(request)
    HTTP::WebSocket.new(
      host: @backend_host,
      path: request.resource,
      port: @backend_port,
      headers: request.headers
    )
  end
end