Merge branch 'develop' into leakage

Author: abhinavsingh

README.md

@@ -2619,6 +2619,7 @@ usage: -m [-h] [--tunnel-hostname TUNNEL_HOSTNAME] [--tunnel-port TUNNEL_PORT]
           [--tunnel-ssh-key-passphrase TUNNEL_SSH_KEY_PASSPHRASE]
           [--tunnel-remote-port TUNNEL_REMOTE_PORT] [--threadless]
           [--threaded] [--num-workers NUM_WORKERS] [--enable-events]
+          [--inactive-conn-cleanup-timeout INACTIVE_CONN_CLEANUP_TIMEOUT]
           [--enable-proxy-protocol] [--enable-conn-pool] [--key-file KEY_FILE]
           [--cert-file CERT_FILE] [--client-recvbuf-size CLIENT_RECVBUF_SIZE]
           [--server-recvbuf-size SERVER_RECVBUF_SIZE]
@@ -2682,6 +2683,16 @@ options:
   --enable-events       Default: False. Enables core to dispatch lifecycle
                         events. Plugins can be used to subscribe for core
                         events.
+  --inactive-conn-cleanup-timeout INACTIVE_CONN_CLEANUP_TIMEOUT
+                        Time after which inactive works must be cleaned up.
+                        Increase this value if your backend services are slow
+                        to response or when proxy.py is handling a high
+                        volume. When running proxy.py on Google Cloud (GCP)
+                        you may see 'backend_connection_closed_before_data_sen
+                        t_to_client', with curl clients you may see 'Empty
+                        reply from server' error when '--inactive-conn-
+                        cleanup-timeout' value is low for your use-case.
+                        Default 1 seconds
   --enable-proxy-protocol
                         Default: False. If used, will enable proxy protocol.
                         Only version 1 is currently supported.

proxy/common/leakage.py

@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+"""
+    proxy.py
+    ~~~~~~~~
+    ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
+    Network monitoring, controls & Application development, testing, debugging.
+
+    :copyright: (c) 2013-present by Abhinav Singh and contributors.
+    :license: BSD, see LICENSE for more details.
+"""
+import time
+
+
+class Leakage:
+    """Leaky Bucket algorithm."""
+
+    def __init__(self, rate: int) -> None:
+        """Initialize the leaky bucket with a specified leak rate in bytes per second."""
+        # Maximum number of tokens the bucket can hold (bytes per second)
+        self.rate = rate
+        self.tokens = rate
+        self.last_check = time.time()
+
+    def _refill(self) -> None:
+        """Refill tokens based on the elapsed time since the last check."""
+        now = time.time()
+        elapsed = now - self.last_check
+        # Add tokens proportional to elapsed time, up to the rate
+        self.tokens += int(elapsed * self.rate)
+        # Cap tokens at the maximum rate to enforce the rate limit
+        self.tokens = min(self.tokens, self.rate)
+        self.last_check = now
+
+    def release(self, tokens: int) -> None:
+        """When you are unable to consume amount units of token, release them into the bucket.
+
+        E.g. say you wanted to read 1024 units, but only 24 units were read, then put
+        back unconsumed 1000 tokens back in the bucket."""
+        if tokens < 0:
+            raise ValueError('Cannot release a negative number of tokens')
+        self.tokens += tokens
+        self.tokens = min(self.tokens, self.rate)
+
+    def consume(self, amount: int) -> int:
+        """Attempt to consume the amount from the bucket.
+
+        Returns the amount allowed to be sent, up to the available tokens (rate).
+        """
+        self._refill()
+        allowed = min(amount, self.tokens)
+        self.tokens -= allowed
+        return allowed

proxy/core/work/threadless.py

@@ -20,6 +20,7 @@
     cast,
 )
 
+from ...common.flag import flags as proxy_flags
 from ...common.types import Readables, Writables, SelectableEvents
 from ...common.logger import Logger
 from ...common.constants import (
@@ -37,6 +38,20 @@
 logger = logging.getLogger(__name__)
 
 
+proxy_flags.add_argument(
+    '--inactive-conn-cleanup-timeout',
+    default=DEFAULT_INACTIVE_CONN_CLEANUP_TIMEOUT,
+    help='Time after which inactive works must be cleaned up. '
+    + 'Increase this value if your backend services are slow to response '
+    + 'or when proxy.py is handling a high volume. When running proxy.py on Google Cloud (GCP) '
+    + "you may see 'backend_connection_closed_before_data_sent_to_client', with curl clients "
+    + "you may see 'Empty reply from server' error when '--inactive-conn-cleanup-timeout' "
+    + 'value is low for your use-case. Default {0} seconds'.format(
+        DEFAULT_INACTIVE_CONN_CLEANUP_TIMEOUT,
+    ),
+)
+
+
 class Threadless(ABC, Generic[T]):
     """Work executor base class.
 
@@ -87,7 +102,9 @@ def __init__(
             SelectableEvents,
         ] = {}
         self.wait_timeout: float = DEFAULT_WAIT_FOR_TASKS_TIMEOUT
-        self.cleanup_inactive_timeout: float = DEFAULT_INACTIVE_CONN_CLEANUP_TIMEOUT
+        self.cleanup_inactive_timeout: float = float(
+            self.flags.inactive_conn_cleanup_timeout,
+        )
         self._total: int = 0
         # When put at the top, causes circular import error
         # since integrated ssh tunnel was introduced.
@@ -318,10 +335,17 @@ def _cleanup(self, work_id: int, reason: str) -> None:
                 self.selector.unregister(fileno)
             self.registered_events_by_work_ids[work_id].clear()
             del self.registered_events_by_work_ids[work_id]
-        self.works[work_id].shutdown()
-        del self.works[work_id]
-        if self.work_queue_fileno() is not None:
-            os.close(work_id)
+        try:
+            self.works[work_id].shutdown()
+        except Exception as exc:
+            logger.exception(
+                'Error when shutting down work#{0}'.format(work_id),
+                exc_info=exc,
+            )
+        finally:
+            del self.works[work_id]
+            if self.work_queue_fileno() is not None:
+                os.close(work_id)
 
     def _create_tasks(
             self,

tests/common/test_leakage.py

@@ -0,0 +1,61 @@
+# -*- coding: utf-8 -*-
+"""
+    proxy.py
+    ~~~~~~~~
+    ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
+    Network monitoring, controls & Application development, testing, debugging.
+
+    :copyright: (c) 2013-present by Abhinav Singh and contributors.
+    :license: BSD, see LICENSE for more details.
+"""
+import time
+
+import unittest
+
+from proxy.common.leakage import Leakage
+
+
+class TestLeakage(unittest.TestCase):
+
+    def test_initial_consume_no_tokens(self) -> None:
+        # Test consuming with no tokens available initially
+        rate = 100  # bytes per second
+        bucket = Leakage(rate)
+        self.assertEqual(
+            bucket.consume(150),
+            100,
+        )  # No tokens yet, so expect 0 bytes to be sent
+
+    def test_consume_with_refill(self) -> None:
+        # Test consuming with refill after waiting
+        rate = 100  # bytes per second
+        bucket = Leakage(rate)
+        time.sleep(1)  # Wait for a second to allow refill
+        self.assertEqual(bucket.consume(50), 50)  # 50 bytes should be available
+
+    def test_consume_above_leak_rate(self) -> None:
+        # Test attempting to consume more than the leak rate after a refill
+        rate = 100  # bytes per second
+        bucket = Leakage(rate)
+        time.sleep(1)  # Wait for a second to allow refill
+        self.assertEqual(bucket.consume(150), 100)  # Only 100 bytes should be allowed
+
+    def test_repeated_consume_with_partial_refill(self) -> None:
+        # Test repeated consumption with partial refill
+        rate = 100  # bytes per second
+        bucket = Leakage(rate)
+
+        time.sleep(1)  # Allow tokens to accumulate
+        bucket.consume(80)  # Consume 80 bytes, should leave 20
+        time.sleep(0.5)  # Wait half a second to refill by 50 bytes
+
+        self.assertEqual(bucket.consume(50), 50)  # 50 bytes should be available now
+
+    def test_negative_token_guard(self) -> None:
+        # Ensure tokens do not go negative
+        rate = 100  # bytes per second
+        bucket = Leakage(rate)
+        time.sleep(1)  # Allow tokens to accumulate
+        bucket.consume(150)  # Consume all available tokens
+        self.assertEqual(bucket.consume(10), 0)  # Should return 0 as no tokens are left
+        self.assertEqual(bucket.tokens, 0)  # Tokens should not be negative