feat: allow core classes to be extended

shard.yml

@@ -1,5 +1,5 @@
 name: connect-proxy
-version: 1.4.1
+version: 2.0.0
 license: MIT
 crystal: ">= 0.36.1"
 

spec/connect_spec.cr

@@ -12,13 +12,15 @@ describe ConnectProxy do
   end
 
   it "connect to a website and get a response using explicit proxy" do
+    expected_count = CONNECTION_COUNT[0] + 1
     host = URI.parse("https://github.com/")
     client = ConnectProxy::HTTPClient.new(host, ignore_env: true)
     proxy = ConnectProxy.new("localhost", 22222)
     client.set_proxy(proxy)
     response = client.exec("GET", "/")
-    response.success?.should eq(true)
     client.close
+    response.success?.should eq(true)
+    expected_count.should eq(CONNECTION_COUNT[0])
   end
 
   it "connect to a website with CRL checks disabled" do
@@ -30,8 +32,8 @@ describe ConnectProxy do
     proxy = ConnectProxy.new("localhost", 22222)
     client.set_proxy(proxy)
     response = client.exec("GET", "/")
-    response.success?.should eq(true)
     client.close
+    response.success?.should eq(true)
   end
 
   it "connect to a website with TLS disabled" do
@@ -42,13 +44,14 @@ describe ConnectProxy do
     proxy = ConnectProxy.new("localhost", 22222)
     client.set_proxy(proxy)
     response = client.exec("GET", "/")
-    response.success?.should eq(true)
     client.close
+    response.success?.should eq(true)
   end
 
   it "connect to a websocket using explicit proxy" do
+    expected_count = CONNECTION_COUNT[0] + 1
     received = ""
-    host = URI.parse("wss://echo.websocket.org/")
+    host = URI.parse("wss://ws.postman-echo.com/raw")
     proxy = ConnectProxy.new("localhost", 22222)
 
     ws = ConnectProxy::WebSocket.new(host, proxy: proxy)
@@ -61,5 +64,6 @@ describe ConnectProxy do
     ws.run
 
     received.should eq("test")
+    expected_count.should eq(CONNECTION_COUNT[0])
   end
 end

spec/ext_spec.cr

@@ -0,0 +1,72 @@
+require "spec"
+require "../src/connect-proxy"
+require "../src/connect-proxy/ext/*"
+require "./proxy_server"
+
+describe "proxy ext" do
+  it "connect to a website and get a response" do
+    host = URI.parse("https://github.com/")
+    response = ::HTTP::Client.new(host) do |client|
+      client.exec("GET", "/")
+    end
+    response.success?.should eq(true)
+  end
+
+  it "connect to a website and get a response using explicit proxy" do
+    host = URI.parse("https://github.com/")
+    expected_count = CONNECTION_COUNT[0] + 1
+    client = ::HTTP::Client.new(host, ignore_env: true)
+    proxy = ConnectProxy.new("localhost", 22222)
+    client.set_proxy(proxy)
+    response = client.exec("GET", "/")
+    client.close
+    response.success?.should eq(true)
+    expected_count.should eq(CONNECTION_COUNT[0])
+  end
+
+  it "connect to a website with CRL checks disabled" do
+    ConnectProxy.verify_tls = true
+    ConnectProxy.disable_crl_checks = true
+
+    host = URI.parse("https://github.com/")
+    client = ::HTTP::Client.new(host, ignore_env: true)
+    proxy = ConnectProxy.new("localhost", 22222)
+    client.set_proxy(proxy)
+    response = client.exec("GET", "/")
+    client.close
+    response.success?.should eq(true)
+  end
+
+  it "connect to a website with TLS disabled" do
+    ConnectProxy.verify_tls = false
+
+    host = URI.parse("https://github.com/")
+    client = ::HTTP::Client.new(host, ignore_env: true)
+    proxy = ConnectProxy.new("localhost", 22222)
+    client.set_proxy(proxy)
+    response = client.exec("GET", "/")
+    client.close
+    response.success?.should eq(true)
+  end
+
+  it "connect to a websocket using explicit proxy" do
+    received = ""
+
+    expected_count = CONNECTION_COUNT[0] + 1
+
+    host = URI.parse("wss://ws.postman-echo.com/raw")
+    proxy = ConnectProxy.new("localhost", 22222)
+
+    ws = ::HTTP::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")
+    expected_count.should eq(CONNECTION_COUNT[0])
+  end
+end

spec/proxy_server.cr

@@ -1,5 +1,7 @@
 require "socket"
 
+CONNECTION_COUNT = [0]
+
 spawn do
   server = TCPServer.new(22222)
   loop do
@@ -15,6 +17,8 @@ spawn do
 
     socket << "HTTP/1.1 200 OK\r\n\r\n"
 
+    CONNECTION_COUNT[0] = CONNECTION_COUNT[0] + 1
+
     spawn do
       begin
         raw_data = Bytes.new(2048)

src/connect-proxy.cr

@@ -100,7 +100,10 @@ class ConnectProxy
         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
+          begin
+            context.add_x509_verify_flags OpenSSL::SSL::X509VerifyFlags::IGNORE_CRITICAL
+          rescue NotImplementedError
+          end
         end
 
         socket = OpenSSL::SSL::Socket::Client.new(socket, context: context, sync_close: true, hostname: host)

src/connect-proxy/ext/http-client.cr

@@ -0,0 +1,5 @@
+require "../../connect-proxy"
+
+class ::HTTP::Client
+  include ConnectProxy::ProxyHTTP
+end

src/connect-proxy/ext/websocket.cr

@@ -0,0 +1,5 @@
+require "../../connect-proxy"
+
+class ::HTTP::WebSocket
+  include ConnectProxy::ProxyWebSocket
+end

src/connect-proxy/http_client.cr

@@ -1,17 +1,24 @@
 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
+module ConnectProxy::ProxyHTTP
+  macro included
+    def self.new(uri : URI, tls : TLSContext = nil, ignore_env = false)
+      {% if @type.stringify == "HTTP::Client" %}
+        inst = previous_def(uri, tls)
+      {% else %}
+        inst = super(uri, tls)
+      {% end %}
 
-    inst
-  end
+      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)
+    def self.new(uri : URI, tls : TLSContext = nil, ignore_env = false)
+      yield new(uri, tls, ignore_env)
+    end
   end
 
   def set_proxy(proxy : ConnectProxy = nil)
@@ -41,3 +48,7 @@ class ConnectProxy::HTTPClient < ::HTTP::Client
     }
   end
 end
+
+class ConnectProxy::HTTPClient < ::HTTP::Client
+  include ConnectProxy::ProxyHTTP
+end

src/connect-proxy/websocket.cr

@@ -35,42 +35,48 @@ class HTTP::WebSocket::Protocol
   end
 end
 
-class ConnectProxy::WebSocket < ::HTTP::WebSocket
-  def self.new(
-    uri : URI | String,
-    headers = HTTP::Headers.new,
-    proxy : ConnectProxy? = nil,
-    ignore_env : Bool = false
-  )
-    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, ignore_env)
-    end
+module ConnectProxy::ProxyWebSocket
+  macro included
+    def self.new(
+      uri : URI | String,
+      headers = HTTP::Headers.new,
+      proxy : ConnectProxy? = nil,
+      ignore_env : Bool = false
+    )
+      uri = URI.parse(uri) if uri.is_a?(String)
 
-    raise ArgumentError.new("No host or path specified which are required.")
-  end
+      if (host = uri.hostname) && (path = uri.request_target)
+        tls = uri.scheme.in?("https", "wss")
+        return new(host, path, uri.port, tls, headers, proxy, ignore_env)
+      end
 
-  def self.new(
-    host : String,
-    path : String,
-    port = nil,
-    tls : HTTP::Client::TLSContext = nil,
-    headers = HTTP::Headers.new,
-    proxy : ConnectProxy? = nil,
-    ignore_env : Bool = false
-  )
-    if proxy.nil? && !ignore_env && ConnectProxy.behind_proxy?
-      proxy = ConnectProxy.new(*ConnectProxy.parse_proxy_url)
+      raise ArgumentError.new("No host or path specified which are required.")
     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))
+    def self.new(
+      host : String,
+      path : String,
+      port = nil,
+      tls : HTTP::Client::TLSContext = nil,
+      headers = HTTP::Headers.new,
+      proxy : ConnectProxy? = nil,
+      ignore_env : Bool = false
+    )
+      if proxy.nil? && !ignore_env && 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
 end
+
+class ConnectProxy::WebSocket < ::HTTP::WebSocket
+  include ConnectProxy::ProxyWebSocket
+end