feat: add websocket proxy support

stakach · stakach · commit fe7097c82aa9 · 2021-06-08T00:40:48.000+10:00
diff --git a/shard.yml b/shard.yml
@@ -1,5 +1,5 @@
 name: connect-proxy
-version: 1.3.0
+version: 1.4.0
 license: MIT
 crystal: ">= 0.36.1"
 
diff --git a/spec/connect_spec.cr b/spec/connect_spec.cr
@@ -1,5 +1,6 @@
 require "spec"
 require "../src/connect-proxy"
+require "./proxy_server"
 
 describe ConnectProxy do
   it "connect to a website and get a response" do
@@ -13,10 +14,11 @@ describe ConnectProxy do
   it "connect to a website and get a response using explicit proxy" do
     host = URI.parse("https://github.com/")
     client = ConnectProxy::HTTPClient.new(host, ignore_env: true)
-    proxy = ConnectProxy.new("51.38.71.101", 8080)
+    proxy = ConnectProxy.new("localhost", 22222)
     client.set_proxy(proxy)
     response = client.exec("GET", "/")
     response.success?.should eq(true)
+    client.close
   end
 
   it "connect to a website with CRL checks disabled" do
@@ -25,20 +27,39 @@ describe ConnectProxy do
 
     host = URI.parse("https://github.com/")
     client = ConnectProxy::HTTPClient.new(host, ignore_env: true)
-    proxy = ConnectProxy.new("51.38.71.101", 8080)
+    proxy = ConnectProxy.new("localhost", 22222)
     client.set_proxy(proxy)
     response = client.exec("GET", "/")
     response.success?.should eq(true)
+    client.close
   end
 
   it "connect to a website with TLS disabled" do
     ConnectProxy.verify_tls = false
 
     host = URI.parse("https://github.com/")
     client = ConnectProxy::HTTPClient.new(host, ignore_env: true)
-    proxy = ConnectProxy.new("51.38.71.101", 8080)
+    proxy = ConnectProxy.new("localhost", 22222)
     client.set_proxy(proxy)
     response = client.exec("GET", "/")
     response.success?.should eq(true)
+    client.close
+  end
+
+  it "connect to a websocket using explicit proxy" do
+    received = ""
+    host = URI.parse("wss://echo.websocket.org/")
+    proxy = ConnectProxy.new("localhost", 22222)
+
+    ws = ConnectProxy::WebSocket.new(host, proxy: proxy)
+    ws.on_message do |msg|
+      puts msg
+      received = msg
+      ws.close
+    end
+    ws.send "test"
+    ws.run
+
+    received.should eq("test")
   end
 end
diff --git a/spec/proxy_server.cr b/spec/proxy_server.cr
@@ -0,0 +1,44 @@
+require "socket"
+
+spawn do
+  server = TCPServer.new(22222)
+  loop do
+    socket = server.accept
+    request = socket.gets("\r\n\r\n").not_nil!
+    details = request.split(" ", 3)[1].split(":")
+    host = details[0]
+    port = details[1].to_i
+
+    puts "connecting to #{host} #{port}..."
+    client = TCPSocket.new(host, port)
+    puts "proxy connection established"
+
+    socket << "HTTP/1.1 200 OK\r\n\r\n"
+
+    spawn do
+      begin
+        raw_data = Bytes.new(2048)
+        while !client.closed?
+          bytes_read = client.read(raw_data)
+          break if bytes_read.zero? # IO was closed
+          socket.write raw_data[0, bytes_read].dup
+        end
+      rescue IO::Error
+      ensure
+        socket.close
+      end
+    end
+
+    begin
+      out_data = Bytes.new(2048)
+      while !socket.closed?
+        read = socket.read(out_data)
+        break if read.zero? # IO was closed
+        client.write out_data[0, read].dup
+      end
+    rescue IO::Error
+    ensure
+      client.close
+    end
+  end
+end
diff --git a/src/connect-proxy.cr b/src/connect-proxy.cr
@@ -91,8 +91,19 @@ class ConnectProxy
 
     if resp[:code]? == 200
       if tls
-        tls_socket = OpenSSL::SSL::Socket::Client.new(socket, context: tls, sync_close: true, hostname: host)
-        socket = tls_socket
+        if tls.is_a?(Bool) # true, but we want to get rid of the union
+          context = OpenSSL::SSL::Context::Client.new
+        else
+          context = tls
+        end
+
+        if !ConnectProxy.verify_tls
+          context.verify_mode = OpenSSL::SSL::VerifyMode::NONE
+        elsif ConnectProxy.disable_crl_checks
+          context.add_x509_verify_flags OpenSSL::SSL::X509VerifyFlags::IGNORE_CRITICAL
+        end
+
+        socket = OpenSSL::SSL::Socket::Client.new(socket, context: context, sync_close: true, hostname: host)
       end
 
       socket
@@ -119,59 +130,12 @@ class ConnectProxy
       resp[:code] = code.to_i
       resp[:reason] = reason
       resp[:headers] = headers
-    rescue
+    rescue error
+      raise IO::Error.new("parsing proxy initialization", cause: error)
     end
 
     resp
   end
-
-  class HTTPClient < ::HTTP::Client
-    def self.new(uri : URI, tls = nil, ignore_env = false)
-      inst = super(uri, tls)
-      if !ignore_env && ConnectProxy.behind_proxy?
-        inst.set_proxy ConnectProxy.new(*ConnectProxy.parse_proxy_url)
-      end
-
-      inst
-    end
-
-    def self.new(uri : URI, tls = nil, ignore_env = false)
-      yield new(uri, tls, ignore_env)
-    end
-
-    def set_proxy(proxy : ConnectProxy = nil)
-      socket = {% if compare_versions(Crystal::VERSION, "0.36.0") < 0 %} @socket {% else %} @io {% end %}
-      return if socket && !socket.closed?
-
-      if tls = @tls
-        if !ConnectProxy.verify_tls
-          tls.verify_mode = OpenSSL::SSL::VerifyMode::NONE
-        elsif ConnectProxy.disable_crl_checks
-          tls.add_x509_verify_flags OpenSSL::SSL::X509VerifyFlags::IGNORE_CRITICAL
-        end
-      end
-
-      {% if compare_versions(Crystal::VERSION, "0.36.0") < 0 %}
-        begin
-          @socket = proxy.open(@host, @port, @tls, **proxy_connection_options)
-        rescue IO::Error
-          @socket = nil
-        end
-      {% else %}
-        begin
-          @io = proxy.open(@host, @port, @tls, **proxy_connection_options)
-        rescue IO::Error
-          @io = nil
-        end
-      {% end %}
-    end
-
-    def proxy_connection_options
-      {
-        dns_timeout:     @dns_timeout,
-        connect_timeout: @connect_timeout,
-        read_timeout:    @read_timeout,
-      }
-    end
-  end
 end
+
+require "./connect-proxy/*"
diff --git a/src/connect-proxy/http_client.cr b/src/connect-proxy/http_client.cr
@@ -0,0 +1,43 @@
+require "../connect-proxy"
+
+class ConnectProxy::HTTPClient < ::HTTP::Client
+  def self.new(uri : URI, tls = nil, ignore_env = false)
+    inst = super(uri, tls)
+    if !ignore_env && ConnectProxy.behind_proxy?
+      inst.set_proxy ConnectProxy.new(*ConnectProxy.parse_proxy_url)
+    end
+
+    inst
+  end
+
+  def self.new(uri : URI, tls = nil, ignore_env = false)
+    yield new(uri, tls, ignore_env)
+  end
+
+  def set_proxy(proxy : ConnectProxy = nil)
+    socket = {% if compare_versions(Crystal::VERSION, "0.36.0") < 0 %} @socket {% else %} @io {% end %}
+    return if socket && !socket.closed?
+
+    {% if compare_versions(Crystal::VERSION, "0.36.0") < 0 %}
+      begin
+        @socket = proxy.open(@host, @port, @tls, **proxy_connection_options)
+      rescue IO::Error
+        @socket = nil
+      end
+    {% else %}
+      begin
+        @io = proxy.open(@host, @port, @tls, **proxy_connection_options)
+      rescue IO::Error
+        @io = nil
+      end
+    {% end %}
+  end
+
+  def proxy_connection_options
+    {
+      dns_timeout:     @dns_timeout,
+      connect_timeout: @connect_timeout,
+      read_timeout:    @read_timeout,
+    }
+  end
+end
diff --git a/src/connect-proxy/websocket.cr b/src/connect-proxy/websocket.cr
@@ -0,0 +1,74 @@
+require "../connect-proxy"
+
+# alternative initialization
+class HTTP::WebSocket::Protocol
+  def self.new(socket, host : String, path : String, port, headers = HTTP::Headers.new)
+    begin
+      random_key = Base64.strict_encode(StaticArray(UInt8, 16).new { rand(256).to_u8 })
+
+      headers["Host"] = "#{host}:#{port}"
+      headers["Connection"] = "Upgrade"
+      headers["Upgrade"] = "websocket"
+      headers["Sec-WebSocket-Version"] = VERSION
+      headers["Sec-WebSocket-Key"] = random_key
+
+      path = "/" if path.empty?
+      handshake = HTTP::Request.new("GET", path, headers)
+      handshake.to_io(socket)
+      socket.flush
+
+      handshake_response = HTTP::Client::Response.from_io(socket, ignore_body: true)
+      unless handshake_response.status.switching_protocols?
+        raise Socket::Error.new("Handshake got denied. Status code was #{handshake_response.status.code}.")
+      end
+
+      challenge_response = Protocol.key_challenge(random_key)
+      unless handshake_response.headers["Sec-WebSocket-Accept"]? == challenge_response
+        raise Socket::Error.new("Handshake got denied. Server did not verify WebSocket challenge.")
+      end
+    rescue exc
+      socket.close
+      raise exc
+    end
+
+    new(socket, masked: true)
+  end
+end
+
+class ConnectProxy::WebSocket < ::HTTP::WebSocket
+  def self.new(
+    uri : URI | String,
+    headers = HTTP::Headers.new,
+    proxy : ConnectProxy? = nil
+  )
+    uri = URI.parse(uri) if uri.is_a?(String)
+
+    if (host = uri.hostname) && (path = uri.request_target)
+      tls = uri.scheme.in?("https", "wss")
+      return new(host, path, uri.port, tls, headers, proxy)
+    end
+
+    raise ArgumentError.new("No host or path specified which are required.")
+  end
+
+  def self.new(
+    host : String,
+    path : String,
+    port = nil,
+    tls : HTTP::Client::TLSContext = nil,
+    headers = HTTP::Headers.new,
+    proxy : ConnectProxy? = nil
+  )
+    if proxy.nil? && ConnectProxy.behind_proxy?
+      proxy = ConnectProxy.new(*ConnectProxy.parse_proxy_url)
+    end
+
+    if proxy
+      port ||= tls ? 443 : 80
+      socket = proxy.open(host, port, tls)
+      new(HTTP::WebSocket::Protocol.new(socket, host, path, port, headers))
+    else
+      new(HTTP::WebSocket::Protocol.new(host, path, port, tls, headers))
+    end
+  end
+end
