Merge pull request #31 from jamesmoriarty/feature/add-connection-timeout

feature: add connection timeout.

.github/workflows/ci.yaml

@@ -8,7 +8,7 @@ jobs:
 
     strategy:
       matrix:
-        ruby: [ '2.3', '2.7' ]
+        ruby: [ '2.5', '2.7' ]
 
     steps:
     - uses: actions/checkout@v2

.github/workflows/cli.yaml

@@ -8,7 +8,7 @@ jobs:
 
     strategy:
       matrix:
-        ruby: [ '2.3', '2.7' ]
+        ruby: [ '2.5', '2.7' ]
 
     steps:
     - uses: actions/checkout@v2

CHANGELOG.md

@@ -1,5 +1,13 @@
 # CHANGELOG
 
+- add connection timeout to stop tracking connection from saturating client threads.
+- add cli flats for connection timeout `-t` and `--timeout`.
+- change cli short flag `-t` to `-c` for `--threads`.
+
+## 0.5.0
+
+- increase default threads from `32` to `128`.
+
 ## 0.2.0
 
 - Extract errors into module.

exe/forward-proxy

@@ -15,7 +15,12 @@ OptionParser.new do |parser|
     options[:bind_address] = bind_address
   end
 
-  parser.on("-tTHREADS", "--threads=THREADS", Integer,
+  parser.on("-tTIMEOUT", "--timeout=TIMEOUT", Integer,
+            "Specify the connection timout in seconds. Default: 300") do |threads|
+    options[:timeout] = threads
+  end
+
+  parser.on("-cTHREADS", "--threads=THREADS", Integer,
             "Specify the number of client threads. Default: 128") do |threads|
     options[:threads] = threads
   end

lib/forward_proxy/errors/connection_timeout_error.rb

@@ -0,0 +1,5 @@
+module ForwardProxy
+  module Errors
+    class ConnectionTimeoutError < StandardError; end
+  end
+end

lib/forward_proxy/server.rb

@@ -1,20 +1,23 @@
 require 'logger'
 require 'socket'
-require 'webrick'
+require 'timeout'
 require 'net/http'
+require 'webrick'
+require 'forward_proxy/errors/connection_timeout_error'
 require 'forward_proxy/errors/http_method_not_implemented'
 require 'forward_proxy/errors/http_parse_error'
 require 'forward_proxy/thread_pool'
 
 module ForwardProxy
   class Server
-    attr_reader :bind_address, :bind_port, :logger
+    attr_reader :bind_address, :bind_port, :logger, :timeout
 
-    def initialize(bind_address: "127.0.0.1", bind_port: 9292, threads: 128, logger: default_logger)
+    def initialize(bind_address: "127.0.0.1", bind_port: 9292, threads: 4, timeout: 1, logger: default_logger)
       @bind_address = bind_address
       @bind_port = bind_port
       @logger = logger
       @thread_pool = ThreadPool.new(threads)
+      @timeout = timeout
     end
 
     def start
@@ -27,18 +30,20 @@ def start
       loop do
         thread_pool.schedule(socket.accept) do |client_conn|
           begin
-            req = parse_req(client_conn)
+            Timeout::timeout(timeout, Errors::ConnectionTimeoutError, "connection exceeded #{timeout} seconds") do
+              req = parse_req(client_conn)
 
-            logger.info(req.request_line.strip)
+              logger.info(req.request_line.strip)
 
-            case req.request_method
-            when METHOD_CONNECT then handle_tunnel(client_conn, req)
-            when METHOD_GET, METHOD_POST then handle(client_conn, req)
-            else
-              raise Errors::HTTPMethodNotImplemented
+              case req.request_method
+              when METHOD_CONNECT then handle_tunnel(client_conn, req)
+              when METHOD_GET, METHOD_POST then handle(client_conn, req)
+              else
+                raise Errors::HTTPMethodNotImplemented
+              end
             end
           rescue => e
-            handle_error(e, client_conn)
+            handle_error(client_conn, e)
           ensure
             client_conn.close
           end
@@ -165,9 +170,15 @@ def handle(client_conn, req)
       end
     end
 
-    def handle_error(err, client_conn)
+    def handle_error(client_conn, err)
+      status_code = case err
+                    when Errors::ConnectionTimeoutError then 504
+                    else
+                      502
+                    end
+
       client_conn.puts <<~eos.chomp
-        HTTP/1.1 502
+        HTTP/1.1 #{status_code}
         Via: #{HEADER_VIA}
         #{HEADER_EOP}
       eos

test/forward_proxy_test.rb

@@ -39,6 +39,16 @@ def test_handle_with_unsupported_method
     end
   end
 
+  def test_handle_error_with_timeout
+    with_dest(uri = URI('http://127.0.0.1:8000/test/index.txt')) do
+      with_proxy(uri, timeout: 1/1000r) do |http|
+        resp = http.request Net::HTTP::Get.new(uri)
+
+        assert_equal "504", resp.code
+      end
+    end
+  end
+
   def test_handle_tunnel
     with_dest(uri = URI('https://127.0.0.1:8000/test/index.txt')) do
       with_proxy(uri) do |http|

test/test_helper.rb

@@ -113,11 +113,13 @@ def with_dest(uri)
   server_thread.join
 end
 
-def with_proxy(uri, bind_address: "127.0.0.1", bind_port: 3000, threads: 32, logger: Logger.new(STDOUT, level: :error))
+def with_proxy(uri, bind_address: "127.0.0.1", bind_port: 3000, timeout: 1, threads: 32, logger: Logger.new(STDOUT,
+                                                                                                            level: :error))
   proxy = ForwardProxy::Server.new(
     bind_address: bind_address,
     bind_port: bind_port,
-    logger: logger
+    logger: logger,
+    timeout: timeout
   )
 
   proxy_thread = Thread.new { proxy.start }