leakage

Author: abhinavsingh

proxy/core/base/tcp_server.py

@@ -104,7 +104,7 @@ class BaseTcpServerHandler(Work[T]):
     is ready to accept new data before flushing data to it.
 
     Most importantly, BaseTcpServerHandler ensures that pending buffers
-    to the client are flushed before connection is closed.
+    to the client are flushed before connection is closed with the client.
 
     Implementations must provide::
 
@@ -170,9 +170,9 @@ async def handle_events(
     async def handle_writables(self, writables: Writables) -> bool:
         teardown = False
         if self.work.connection.fileno() in writables and self.work.has_buffer():
-            logger.debug(
-                'Flushing buffer to client {0}'.format(self.work.address),
-            )
+            # logger.debug(
+            #     'Flushing buffer to client {0}'.format(self.work.address),
+            # )
             self.work.flush(self.flags.max_sendbuf_size)
             if self.must_flush_before_shutdown is True and \
                     not self.work.has_buffer():

proxy/core/connection/connection.py

@@ -12,13 +12,15 @@
 from abc import ABC, abstractmethod
 from typing import List, Union, Optional
 
+from .leak import Leakage
 from .types import tcpConnectionTypes
 from ...common.types import TcpOrTlsSocket
 from ...common.constants import DEFAULT_BUFFER_SIZE, DEFAULT_MAX_SEND_SIZE
 
 
 logger = logging.getLogger(__name__)
 
+EMPTY_MV = memoryview(b"")
 
 class TcpConnectionUninitializedException(Exception):
     pass
@@ -34,12 +36,23 @@ class TcpConnection(ABC):
     a socket connection object.
     """
 
-    def __init__(self, tag: int) -> None:
+    def __init__(
+        self,
+        tag: int,
+        flush_bw_in_bps: int = 512,
+        recv_bw_in_bps: int = 512,
+    ) -> None:
         self.tag: str = 'server' if tag == tcpConnectionTypes.SERVER else 'client'
         self.buffer: List[memoryview] = []
         self.closed: bool = False
         self._reusable: bool = False
         self._num_buffer = 0
+        self._flush_leakage = (
+            Leakage(rate=flush_bw_in_bps) if flush_bw_in_bps > 0 else None
+        )
+        self._recv_leakage = (
+            Leakage(rate=recv_bw_in_bps) if recv_bw_in_bps > 0 else None
+        )
 
     @property
     @abstractmethod
@@ -55,14 +68,19 @@ def recv(
             self, buffer_size: int = DEFAULT_BUFFER_SIZE,
     ) -> Optional[memoryview]:
         """Users must handle socket.error exceptions"""
+        if self._recv_leakage is not None:
+            allowed_bytes = self._recv_leakage.consume(buffer_size)
+            if allowed_bytes == 0:
+                return EMPTY_MV
+            buffer_size = min(buffer_size, allowed_bytes)
         data: bytes = self.connection.recv(buffer_size)
-        if len(data) == 0:
+        size = len(data)
+        if self._recv_leakage is not None:
+            self._recv_leakage.putback(buffer_size - size)
+        if size == 0:
             return None
-        logger.debug(
-            'received %d bytes from %s' %
-            (len(data), self.tag),
-        )
-        # logger.info(data)
+        logger.debug("received %d bytes from %s" % (size, self.tag))
+        logger.info(data)
         return memoryview(data)
 
     def close(self) -> bool:
@@ -75,6 +93,8 @@ def has_buffer(self) -> bool:
         return self._num_buffer != 0
 
     def queue(self, mv: memoryview) -> None:
+        if len(mv) == 0:
+            return
         self.buffer.append(mv)
         self._num_buffer += 1
 
@@ -83,21 +103,38 @@ def flush(self, max_send_size: Optional[int] = None) -> int:
         if not self.has_buffer():
             return 0
         mv = self.buffer[0]
+        print(self.buffer)
+        print(mv.tobytes())
         # TODO: Assemble multiple packets if total
         # size remains below max send size.
         max_send_size = max_send_size or DEFAULT_MAX_SEND_SIZE
-        try:
-            sent: int = self.send(mv[:max_send_size])
-        except BlockingIOError:
-            logger.warning('BlockingIOError when trying send to {0}'.format(self.tag))
-            return 0
+        allowed_bytes = (
+            self._flush_leakage.consume(min(len(mv), max_send_size))
+            if self._flush_leakage is not None
+            else max_send_size
+        )
+        sent: int = 0
+        if allowed_bytes > 0:
+            try:
+                sent = self.send(mv[:allowed_bytes])
+                if self._flush_leakage is not None:
+                    self._flush_leakage.putback(allowed_bytes - sent)
+            except BlockingIOError:
+                logger.warning(
+                    "BlockingIOError when trying send to {0}".format(self.tag)
+                )
+                del mv
+                return 0
+        # if sent == 0:
+        #     return 0
         if sent == len(mv):
             self.buffer.pop(0)
             self._num_buffer -= 1
         else:
             self.buffer[0] = mv[sent:]
-        logger.debug('flushed %d bytes to %s' % (sent, self.tag))
-        # logger.info(mv[:sent].tobytes())
+        # if sent > 0:
+        logger.debug("flushed %d bytes to %s" % (sent, self.tag))
+        logger.info(mv[:sent].tobytes())
         del mv
         return sent
 

proxy/core/connection/leak.py

@@ -0,0 +1,44 @@
+# -*- 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:
+    def __init__(self, rate: int):
+        """Initialize the leaky bucket with a specified leak rate in bytes per second."""
+        self.rate = (
+            rate  # Maximum number of tokens the bucket can hold (bytes per second)
+        )
+        self.tokens = rate  # Initially start with a full bucket
+        self.last_check = time.time()  # Record the current time
+
+    def _refill(self):
+        """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  # Update the last check time
+
+    def putback(self, tokens) -> None:
+        self.tokens += tokens
+
+    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()  # Refill the tokens before consumption
+        allowed = min(amount, self.tokens)  # Allow up to the available tokens
+        self.tokens -= allowed  # Subtract the allowed amount from the available tokens
+        return allowed  # Return the number of bytes allowed to be consumed

proxy/core/work/fd/fd.py

@@ -41,14 +41,15 @@ def work(self, *args: Any) -> None:
             publisher_id=self.__class__.__qualname__,
         )
         try:
+            logger.debug("Initializing work#{0}".format(fileno))
             self.works[fileno].initialize()
             self._total += 1
         except Exception as e:
             logger.exception(   # pragma: no cover
                 'Exception occurred during initialization',
                 exc_info=e,
             )
-            self._cleanup(fileno)
+            self._cleanup(fileno, "error")
 
     @property
     @abstractmethod

proxy/core/work/threadless.py

@@ -301,16 +301,18 @@ def _cleanup_inactive(self) -> None:
             if self.works[work_id].is_inactive():
                 inactive_works.append(work_id)
         for work_id in inactive_works:
-            self._cleanup(work_id)
+            self._cleanup(work_id, "inactive")
 
     # TODO: HttpProtocolHandler.shutdown can call flush which may block
-    def _cleanup(self, work_id: int) -> None:
+    def _cleanup(self, work_id: int, reason: str) -> None:
         if work_id in self.registered_events_by_work_ids:
             assert self.selector
             for fileno in self.registered_events_by_work_ids[work_id]:
                 logger.debug(
-                    'fd#{0} unregistered by work#{1}'.format(
-                        fileno, work_id,
+                    "fd#{0} unregistered by work#{1}, reason: {2}".format(
+                        fileno,
+                        work_id,
+                        reason,
                     ),
                 )
                 self.selector.unregister(fileno)
@@ -360,7 +362,7 @@ async def _run_once(self) -> bool:
             return False
         # Invoke Threadless.handle_events
         self.unfinished.update(self._create_tasks(work_by_ids))
-        # logger.debug('Executing {0} works'.format(len(self.unfinished)))
+        # logger.debug("Executing {0} works".format(len(self.unfinished)))
         # Cleanup finished tasks
         for task in await self._wait_for_tasks():
             # Checking for result can raise exception e.g.
@@ -374,11 +376,12 @@ async def _run_once(self) -> bool:
                 teardown = True
             finally:
                 if teardown:
-                    self._cleanup(work_id)
+                    self._cleanup(work_id, "teardown")
                     # self.cleanup(int(task.get_name()))
         # logger.debug(
-        #     'Done executing works, {0} pending, {1} registered'.format(
-        #         len(self.unfinished), len(self.registered_events_by_work_ids),
+        #     "Done executing works, {0} pending, {1} registered".format(
+        #         len(self.unfinished),
+        #         len(self.registered_events_by_work_ids),
         #     ),
         # )
         return False

proxy/http/client.py

@@ -32,8 +32,8 @@ def client(
     conn_close: bool = True,
     scheme: bytes = HTTPS_PROTO,
     timeout: float = DEFAULT_TIMEOUT,
-    content_type: bytes = b'application/x-www-form-urlencoded',
-    verify: bool = True,
+    content_type: bytes = b"application/x-www-form-urlencoded",
+    verify: bool = False,
 ) -> Optional[HttpParser]:
     """HTTP Client"""
     request = build_http_request(

proxy/http/handler.py

@@ -190,7 +190,7 @@ def handle_data(self, data: memoryview) -> Optional[bool]:
 
     async def handle_writables(self, writables: Writables) -> bool:
         if self.work.connection.fileno() in writables and self.work.has_buffer():
-            logger.debug('Client is write ready, flushing...')
+            # logger.debug('Client is write ready, flushing...')
             self.last_activity = time.time()
             # TODO(abhinavsingh): This hook could just reside within server recv block
             # instead of invoking when flushed to client.
@@ -219,7 +219,7 @@ async def handle_writables(self, writables: Writables) -> bool:
 
     async def handle_readables(self, readables: Readables) -> bool:
         if self.work.connection.fileno() in readables:
-            logger.debug('Client is read ready, receiving...')
+            # logger.debug('Client is read ready, receiving...')
             self.last_activity = time.time()
             try:
                 teardown = await super().handle_readables(readables)

proxy/http/server/reverse.py

@@ -85,6 +85,9 @@ def handle_request(self, request: HttpParser) -> None:
                             random.choice(route[1]),
                         )
                         needs_upstream = True
+                        logger.debug(
+                            "Starting connection to upstream {0}".format(self.choice)
+                        )
                         break
                 # Dynamic routes
                 elif isinstance(route, str):
@@ -95,14 +98,23 @@ def handle_request(self, request: HttpParser) -> None:
                             self.choice = choice
                             needs_upstream = True
                             self._upstream_proxy_pass = str(self.choice)
+                            logger.debug(
+                                "Starting connection to upstream {0}".format(choice)
+                            )
                         elif isinstance(choice, memoryview):
                             self.client.queue(choice)
                             self._upstream_proxy_pass = '{0} bytes'.format(len(choice))
+                            logger.debug("Sending raw response to client")
                         else:
                             self.upstream = choice
                             self._upstream_proxy_pass = '{0}:{1}'.format(
                                 *self.upstream.addr,
                             )
+                            logger.debug(
+                                "Using existing connection to upstream {0}".format(
+                                    self.upstream.addr
+                                )
+                            )
                         break
                 else:
                     raise ValueError('Invalid route')

tests/core/test_leakage.py

@@ -0,0 +1,59 @@
+# -*- 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.core.connection.leak import Leakage
+
+
+class TestTcpConnectionLeakage(unittest.TestCase):
+    def test_initial_consume_no_tokens(self):
+        # 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):
+        # 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):
+        # 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):
+        # 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):
+        # 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

tests/http/parser/test_http_parser.py

@@ -908,3 +908,17 @@ def test_cannot_parse_sip_protocol(self) -> None:
                     b'\r\n',
                 ),
             )
+
+    def test_byte_by_byte(self) -> None:
+        response = HttpParser(httpParserTypes.RESPONSE_PARSER)
+        request = [
+            # pylint: disable=line-too-long
+            b'HTTP/1.1 200 OK\r\naccess-control-allow-credentials: true\r\naccess-control-allow-origin: *\r\ncontent-type: application/json; charset=utf-8\r\ndate: Thu, 14 Nov 2024 10:24:11 GMT\r\ncontent-length: 550\r\nserver: Fly/a40a59d0 (2024-11-12)\r\nvia: 1.1 fly.io\r\nfly-request-id: 01JCN37CEK4TB4DRWZDFFQYSD9-bom\r\n\r\n{\n  "args": {},\n  "headers": {\n    "Accept": [\n      "*/*"\n    ],\n    "Host": [\n      "httpbingo.org"\n    ],\n    "User-Agent": [\n      "curl/8.6.0"\n    ],\n    "Via": [\n      "1.1 fly.io"\n    ],\n    "X-Forwarded-For": [\n      "183.82.162.68, 66.241.125.232"\n    ],\n    "X-Forwarded-Port": [\n      "443"\n    ],\n    "X-Forwarded-Proto": [\n      "https"\n    ],\n    "X-Forwarded-Ssl',
+            b'": [\n      "on"\n    ],\n    "X-Request-Start": [\n      "t=1731579851219982"\n    ]\n  },\n  "method": "GET",\n  "origin": "183.82.162.68",\n  "url": "https://httpbingo.org/get"\n}\n',
+        ]
+        response.parse(memoryview(request[0]))
+        self.assertEqual(response.state, httpParserStates.RCVING_BODY)
+        self.assertEqual(response.code, b"200")
+        for byte in (bytes([b]) for b in request[1]):
+            response.parse(memoryview(byte))
+        self.assertEqual(response.state, httpParserStates.COMPLETE)