From 64da39eb9b27b8a3afee6fdf6f06c81be08f4d75 Mon Sep 17 00:00:00 2001
From: Anton Maminov <anton.linux@gmail.com>
Date: Sun, 8 Jan 2023 19:45:22 +0200
Subject: [PATCH] add proxy support

---
 shard.yml                   |  2 ++
 spec/proxy_spec.cr          | 17 +++++++++++++++++
 spec/spec_helper.cr         | 20 ++++++++++++++++++++
 src/mechanize.cr            |  6 ++++++
 src/mechanize/http/agent.cr |  7 +++++++
 5 files changed, 52 insertions(+)
 create mode 100644 spec/proxy_spec.cr

diff --git a/shard.yml b/shard.yml
index 718580f..4a30b7d 100644
--- a/shard.yml
+++ b/shard.yml
@@ -11,6 +11,8 @@ license: MIT
 dependencies:
   lexbor:
     github: kostya/lexbor
+  http_proxy:
+    github: mamantoha/http_proxy
 
 development_dependencies:
   webmock:
diff --git a/spec/proxy_spec.cr b/spec/proxy_spec.cr
new file mode 100644
index 0000000..987e49e
--- /dev/null
+++ b/spec/proxy_spec.cr
@@ -0,0 +1,17 @@
+require "./spec_helper"
+
+WebMock.stub(:get, "http://example.com/with_proxy").to_return(body: "success")
+
+describe "Mechanize proxy test" do
+  it "set proxy" do
+    with_proxy_server do |host, port, wants_close|
+      agent = Mechanize.new
+      agent.set_proxy("127.0.0.1", 8080)
+      page = agent.get("http://example.com/with_proxy")
+      page.body.should eq("success")
+      page.code.should eq(200)
+    ensure
+      wants_close.send(nil)
+    end
+  end
+end
diff --git a/spec/spec_helper.cr b/spec/spec_helper.cr
index f3394a2..ed422b2 100644
--- a/spec/spec_helper.cr
+++ b/spec/spec_helper.cr
@@ -1,5 +1,6 @@
 require "spec"
 require "webmock"
+require "http_proxy"
 require "../src/mechanize"
 
 WebMock.stub(:get, "example.com")
@@ -46,3 +47,22 @@ WebMock.stub(:post, "example.com/post_path")
 WebMock.stub(:post, "example.com/post_path")
   .with(body: "name=foo&email=bar&commit=submit", headers: {"Content-Type" => "application/x-www-form-urlencoded"})
   .to_return(body: "success with button")
+
+def with_proxy_server(host = "127.0.0.1", port = 8080, &)
+  wants_close = Channel(Nil).new
+  server = HTTP::Proxy::Server.new
+
+  spawn do
+    server.bind_tcp(host, port)
+    server.listen
+  end
+
+  spawn do
+    wants_close.receive
+    server.close
+  end
+
+  Fiber.yield
+
+  yield host, port, wants_close
+end
diff --git a/src/mechanize.cr b/src/mechanize.cr
index 05d29ee..14721dc 100644
--- a/src/mechanize.cr
+++ b/src/mechanize.cr
@@ -2,6 +2,7 @@ require "log"
 require "uri"
 require "http/client"
 require "lexbor"
+require "http_proxy"
 require "./mechanize/http/agent"
 require "./mechanize/form"
 require "./mechanize/node"
@@ -305,6 +306,11 @@ class Mechanize
     @agent.add_auth(uri, user, pass)
   end
 
+  # Sets the proxy +address+ at +port+ with an optional +user+ and +password+
+  def set_proxy(address : String, port : Int32, user : String? = nil, password : String? = nil)
+    @agent.set_proxy(address, port, user, password)
+  end
+
   # Runs given block, then resets the page history as it was before.
   private def transact
     # save the previous history status.
diff --git a/src/mechanize/http/agent.cr b/src/mechanize/http/agent.cr
index 89abec2..1f227e4 100644
--- a/src/mechanize/http/agent.cr
+++ b/src/mechanize/http/agent.cr
@@ -16,6 +16,8 @@ class Mechanize
       getter authenticate_methods : Hash(URI, Hash(String, Array(AuthRealm)))
       getter authenticate_parser : WWWAuthenticateParser
 
+      @proxy : ::HTTP::Proxy::Client?
+
       def initialize(@context : Mechanize? = nil)
         @history = History.new
         @request_headers = ::HTTP::Headers.new
@@ -43,6 +45,7 @@ class Mechanize
         uri, params = resolve_parameters(uri, method, params)
         client = ::HTTP::Client.new(uri)
         request_auth client, uri
+        client.set_proxy(@proxy) if @proxy
         response = http_request(client, uri, method, params, body)
         body = response.not_nil!.body
         page = response_parse(response, body, uri)
@@ -170,6 +173,10 @@ class Mechanize
         @auth_store.add_auth(uri, user, pass)
       end
 
+      def set_proxy(address : String, port : Int32, user : String? = nil, password : String? = nil)
+        @proxy = ::HTTP::Proxy::Client.new(address, port, username: user, password: password)
+      end
+
       private def set_request_headers(uri, headers)
         reset_request_header_cookies
         headers.each do |k, v|