From 6f23e132780c45a4d15b253da65040ae84945ac1 Mon Sep 17 00:00:00 2001 From: James Moriarty Date: Sat, 17 Jul 2021 11:02:31 +1000 Subject: [PATCH 1/2] feature: add connection timeout. --- .github/workflows/ci.yaml | 2 +- .github/workflows/cli.yaml | 2 +- exe/forward-proxy | 7 +++- .../errors/connection_timeout_error.rb | 5 +++ lib/forward_proxy/server.rb | 37 ++++++++++++------- test/forward_proxy_test.rb | 10 +++++ test/test_helper.rb | 6 ++- 7 files changed, 51 insertions(+), 18 deletions(-) create mode 100644 lib/forward_proxy/errors/connection_timeout_error.rb diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 4ed8434..3e76230 100644 --- a/.github/workflows/ci.yaml +++ b/.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 diff --git a/.github/workflows/cli.yaml b/.github/workflows/cli.yaml index aa0e7c8..1606c1b 100644 --- a/.github/workflows/cli.yaml +++ b/.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 diff --git a/exe/forward-proxy b/exe/forward-proxy index 61eb881..a37bf07 100755 --- a/exe/forward-proxy +++ b/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 diff --git a/lib/forward_proxy/errors/connection_timeout_error.rb b/lib/forward_proxy/errors/connection_timeout_error.rb new file mode 100644 index 0000000..4363c02 --- /dev/null +++ b/lib/forward_proxy/errors/connection_timeout_error.rb @@ -0,0 +1,5 @@ +module ForwardProxy + module Errors + class ConnectionTimeoutError < StandardError; end + end +end diff --git a/lib/forward_proxy/server.rb b/lib/forward_proxy/server.rb index 75ac70a..c28fd2f 100644 --- a/lib/forward_proxy/server.rb +++ b/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 diff --git a/test/forward_proxy_test.rb b/test/forward_proxy_test.rb index 0916f1f..dff413a 100644 --- a/test/forward_proxy_test.rb +++ b/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| diff --git a/test/test_helper.rb b/test/test_helper.rb index 2846a40..c17fa79 100644 --- a/test/test_helper.rb +++ b/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 } From b23eab51d034e88ead6dd01c091ade39e0ea5a9f Mon Sep 17 00:00:00 2001 From: James Moriarty Date: Sat, 17 Jul 2021 11:16:05 +1000 Subject: [PATCH 2/2] refinement: update change log. --- CHANGELOG.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index bdf8c00..bf395e5 100644 --- a/CHANGELOG.md +++ b/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.