[Middleware] HttpClientConnection preparation

docs/conf.py

@@ -308,12 +308,14 @@
     (_py_class_role, 'unittest.result.TestResult'),
     (_py_class_role, 'UUID'),
     (_py_class_role, 'UpstreamConnectionPool'),
+    (_py_class_role, 'HttpClientConnection'),
     (_py_class_role, 'Url'),
     (_py_class_role, 'WebsocketFrame'),
     (_py_class_role, 'Work'),
     (_py_class_role, 'proxy.core.acceptor.work.Work'),
     (_py_class_role, 'connection.Connection'),
     (_py_class_role, 'EventQueue'),
+    (_py_class_role, 'T'),
     (_py_obj_role, 'proxy.core.work.threadless.T'),
     (_py_obj_role, 'proxy.core.work.work.T'),
     (_py_obj_role, 'proxy.core.base.tcp_server.T'),

examples/ssl_echo_server.py

@@ -9,7 +9,7 @@
     :license: BSD, see LICENSE for more details.
 """
 import time
-from typing import Optional
+from typing import Any, Optional
 
 from proxy import Proxy
 from proxy.core.base import BaseTcpServerHandler
@@ -20,6 +20,10 @@
 class EchoSSLServerHandler(BaseTcpServerHandler[TcpClientConnection]):
     """Wraps client socket during initialization."""
 
+    @staticmethod
+    def create(**kwargs: Any) -> TcpClientConnection:
+        return TcpClientConnection(**kwargs)
+
     def initialize(self) -> None:
         # Acceptors don't perform TLS handshake.  Perform the same
         # here using wrap_socket() utility.

examples/tcp_echo_server.py

@@ -9,7 +9,7 @@
     :license: BSD, see LICENSE for more details.
 """
 import time
-from typing import Optional
+from typing import Any, Optional
 
 from proxy import Proxy
 from proxy.core.base import BaseTcpServerHandler
@@ -19,6 +19,10 @@
 class EchoServerHandler(BaseTcpServerHandler[TcpClientConnection]):
     """Sets client socket to non-blocking during initialization."""
 
+    @staticmethod
+    def create(**kwargs: Any) -> TcpClientConnection:
+        return TcpClientConnection(**kwargs)
+
     def initialize(self) -> None:
         self.work.connection.setblocking(False)
 

proxy/core/acceptor/acceptor.py

@@ -224,7 +224,8 @@ def _work(self, conn: socket.socket, addr: Optional[Tuple[str, int]]) -> None:
                 ),
             )
             thread.start()
-            logger.debug(
+            # TODO: Move me into target method
+            logger.debug(   # pragma: no cover
                 'Dispatched work#{0}.{1}.{2} to worker#{3}'.format(
                     conn.fileno(), self.idd, self._total, index,
                 ),
@@ -237,6 +238,7 @@ def _work(self, conn: socket.socket, addr: Optional[Tuple[str, int]]) -> None:
                 event_queue=self.event_queue,
                 publisher_id=self.__class__.__name__,
             )
+            # TODO: Move me into target method
             logger.debug(   # pragma: no cover
                 'Started work#{0}.{1}.{2} in thread#{3}'.format(
                     conn.fileno(), self.idd, self._total, thread.ident,

proxy/core/base/tcp_tunnel.py

@@ -47,6 +47,10 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
     def handle_data(self, data: memoryview) -> Optional[bool]:
         pass    # pragma: no cover
 
+    @staticmethod
+    def create(**kwargs: Any) -> TcpClientConnection:
+        return TcpClientConnection(**kwargs)
+
     def initialize(self) -> None:
         self.work.connection.setblocking(False)
 

proxy/core/connection/pool.py

@@ -16,7 +16,7 @@
 import logging
 import selectors
 
-from typing import TYPE_CHECKING, Set, Dict, Tuple
+from typing import TYPE_CHECKING, Any, Set, Dict, Tuple
 
 from ...common.flag import flags
 from ...common.types import Readables, SelectableEvents, Writables
@@ -77,6 +77,10 @@ def __init__(self) -> None:
         self.connections: Dict[int, TcpServerConnection] = {}
         self.pools: Dict[Tuple[str, int], Set[TcpServerConnection]] = {}
 
+    @staticmethod
+    def create(**kwargs: Any) -> TcpServerConnection:
+        return TcpServerConnection(**kwargs)
+
     def acquire(self, addr: Tuple[str, int]) -> Tuple[bool, TcpServerConnection]:
         """Returns a reusable connection from the pool.
 
@@ -152,7 +156,7 @@ def add(self, addr: Tuple[str, int]) -> TcpServerConnection:
 
         NOTE: You must not use the returned connection, instead use `acquire`.
         """
-        new_conn = TcpServerConnection(addr[0], addr[1])
+        new_conn = self.create(host=addr[0], port=addr[1])
         new_conn.connect()
         self._add(new_conn)
         logger.debug(

proxy/core/ssh/handler.py

@@ -12,7 +12,7 @@
 
 from typing import TYPE_CHECKING, Tuple
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:   # pragma: no cover
     try:
         from paramiko.channel import Channel
     except ImportError:

proxy/core/ssh/listener.py

@@ -16,7 +16,7 @@
 try:
     from paramiko import SSHClient, AutoAddPolicy
     from paramiko.transport import Transport
-    if TYPE_CHECKING:
+    if TYPE_CHECKING:   # pragma: no cover
         from paramiko.channel import Channel
 except ImportError:
     pass

proxy/core/work/threaded.py

@@ -12,25 +12,27 @@
 import argparse
 import threading
 
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import TYPE_CHECKING, Optional, Tuple, TypeVar
 
-from ..connection import TcpClientConnection
 from ..event import EventQueue, eventNames
 
 if TYPE_CHECKING:   # pragma: no cover
     from .work import Work
 
+T = TypeVar('T')
 
+
+# TODO: Add generic T
 def start_threaded_work(
         flags: argparse.Namespace,
         conn: socket.socket,
         addr: Optional[Tuple[str, int]],
         event_queue: Optional[EventQueue] = None,
         publisher_id: Optional[str] = None,
-) -> Tuple['Work[TcpClientConnection]', threading.Thread]:
+) -> Tuple['Work[T]', threading.Thread]:
     """Utility method to start a work in a new thread."""
     work = flags.work_klass(
-        TcpClientConnection(conn, addr),
+        flags.work_klass.create(conn=conn, addr=addr),
         flags=flags,
         event_queue=event_queue,
         upstream_conn_pool=None,

proxy/core/work/threadless.py

@@ -25,7 +25,6 @@
 from ...common.constants import DEFAULT_INACTIVE_CONN_CLEANUP_TIMEOUT, DEFAULT_SELECTOR_SELECT_TIMEOUT
 from ...common.constants import DEFAULT_WAIT_FOR_TASKS_TIMEOUT
 
-from ..connection import TcpClientConnection
 from ..event import eventNames
 
 if TYPE_CHECKING:   # pragma: no cover
@@ -138,7 +137,7 @@ def work_on_tcp_conn(
         )
         uid = '%s-%s-%s' % (self.iid, self._total, fileno)
         self.works[fileno] = self.flags.work_klass(
-            TcpClientConnection(
+            self.flags.work_klass.create(
                 conn=conn,
                 addr=addr,
             ),

proxy/core/work/work.py

@@ -47,6 +47,14 @@ def __init__(
         self.work = work
         self.upstream_conn_pool = upstream_conn_pool
 
+    @staticmethod
+    @abstractmethod
+    def create(**kwargs: Any) -> T:
+        """Implementations are responsible for creation of work objects
+        from incoming args.  This helps keep work core agnostic to
+        creation of externally defined work class objects."""
+        raise NotImplementedError()
+
     @abstractmethod
     async def get_events(self) -> SelectableEvents:
         """Return sockets and events (read or write) that we are interested in."""

proxy/http/__init__.py

@@ -9,6 +9,7 @@
     :license: BSD, see LICENSE for more details.
 """
 from .handler import HttpProtocolHandler
+from .connection import HttpClientConnection
 from .plugin import HttpProtocolHandlerPlugin
 from .codes import httpStatusCodes
 from .methods import httpMethods
@@ -17,6 +18,7 @@
 
 __all__ = [
     'HttpProtocolHandler',
+    'HttpClientConnection',
     'HttpProtocolHandlerPlugin',
     'httpStatusCodes',
     'httpMethods',

proxy/http/connection.py

@@ -0,0 +1,20 @@
+# -*- 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.
+
+    .. spelling::
+
+       http
+       iterable
+"""
+from ..core.connection import TcpClientConnection
+
+
+class HttpClientConnection(TcpClientConnection):
+    pass

proxy/http/handler.py

@@ -19,10 +19,10 @@
 from typing import Tuple, List, Type, Optional, Any
 
 from ..core.base import BaseTcpServerHandler
-from ..core.connection import TcpClientConnection
 from ..common.types import Readables, SelectableEvents, Writables
 from ..common.constants import DEFAULT_SELECTOR_SELECT_TIMEOUT
 
+from .connection import HttpClientConnection
 from .exception import HttpProtocolException
 from .plugin import HttpProtocolHandlerPlugin
 from .responses import BAD_REQUEST_RESPONSE_PKT
@@ -32,7 +32,7 @@
 logger = logging.getLogger(__name__)
 
 
-class HttpProtocolHandler(BaseTcpServerHandler[TcpClientConnection]):
+class HttpProtocolHandler(BaseTcpServerHandler[HttpClientConnection]):
     """HTTP, HTTPS, HTTP2, WebSockets protocol handler.
 
     Accepts `Client` connection and delegates to HttpProtocolHandlerPlugin.
@@ -56,14 +56,14 @@ def __init__(self, *args: Any, **kwargs: Any):
     # overrides Work class definitions.
     ##
 
+    @staticmethod
+    def create(**kwargs: Any) -> HttpClientConnection:
+        return HttpClientConnection(**kwargs)
+
     def initialize(self) -> None:
         super().initialize()
-        # Update client connection reference if connection was wrapped
-        # This is here in `handler` and not `tcp_server` because
-        # `tcp_server` is agnostic to constructing TcpClientConnection
-        # objects.
         if self._encryption_enabled():
-            self.work = TcpClientConnection(
+            self.work = HttpClientConnection(
                 conn=self.work.connection,
                 addr=self.work.addr,
             )
@@ -204,12 +204,12 @@ async def handle_writables(self, writables: Writables) -> bool:
                 if teardown:
                     return True
             except BrokenPipeError:
-                logger.error(
+                logger.warning(
                     'BrokenPipeError when flushing buffer for client',
                 )
                 return True
             except OSError:
-                logger.error('OSError when flushing buffer to client')
+                logger.warning('OSError when flushing buffer to client')
                 return True
         return False
 

proxy/http/inspector/inspect_traffic.py

@@ -9,13 +9,16 @@
     :license: BSD, see LICENSE for more details.
 """
 import json
-from typing import List, Dict, Any
+from typing import TYPE_CHECKING, List, Dict, Any
 
 from ...common.utils import bytes_
 from ...core.event import EventSubscriber
-from ...core.connection import TcpClientConnection
+
 from ..websocket import WebsocketFrame, WebSocketTransportBasePlugin
 
+if TYPE_CHECKING:   # pragma: no cover
+    from ..connection import HttpClientConnection
+
 
 class InspectTrafficPlugin(WebSocketTransportBasePlugin):
     """Websocket API for inspect_traffic.ts frontend plugin."""
@@ -58,7 +61,7 @@ def handle_message(self, message: Dict[str, Any]) -> None:
             raise NotImplementedError()
 
     @staticmethod
-    def callback(client: TcpClientConnection, event: Dict[str, Any]) -> None:
+    def callback(client: 'HttpClientConnection', event: Dict[str, Any]) -> None:
         event['push'] = 'inspect_traffic'
         client.queue(
             memoryview(

proxy/http/inspector/transformer.py

@@ -10,15 +10,18 @@
 """
 import json
 import time
-from typing import Any, Dict
+from typing import TYPE_CHECKING, Any, Dict
 
-from ..websocket import WebsocketFrame
 from ...common.constants import PROXY_PY_START_TIME, DEFAULT_DEVTOOLS_DOC_URL
 from ...common.constants import DEFAULT_DEVTOOLS_FRAME_ID, DEFAULT_DEVTOOLS_LOADER_ID
 from ...common.utils import bytes_
-from ...core.connection import TcpClientConnection
 from ...core.event import eventNames
 
+from ..websocket import WebsocketFrame
+
+if TYPE_CHECKING:   # pragma: no cover
+    from ..connection import HttpClientConnection
+
 
 class CoreEventsToDevtoolsProtocol:
     """Open in Chrome
@@ -30,7 +33,7 @@ class CoreEventsToDevtoolsProtocol:
 
     @staticmethod
     def transformer(
-        client: TcpClientConnection,
+        client: 'HttpClientConnection',
         event: Dict[str, Any],
     ) -> None:
         event_name = event['event_name']

proxy/http/plugin.py

@@ -15,9 +15,9 @@
 from typing import List, Union, Optional, TYPE_CHECKING
 
 from ..core.event import EventQueue
-from ..core.connection import TcpClientConnection
 
 from .parser import HttpParser
+from .connection import HttpClientConnection
 from .descriptors import DescriptorsHandlerMixin
 from .mixins import TlsInterceptionPropertyMixin
 
@@ -55,15 +55,15 @@ def __init__(
             self,
             uid: str,
             flags: argparse.Namespace,
-            client: TcpClientConnection,
+            client: HttpClientConnection,
             request: HttpParser,
             event_queue: Optional[EventQueue] = None,
             upstream_conn_pool: Optional['UpstreamConnectionPool'] = None,
     ):
         super().__init__(uid, flags, client, event_queue, upstream_conn_pool)
         self.uid: str = uid
         self.flags: argparse.Namespace = flags
-        self.client: TcpClientConnection = client
+        self.client: HttpClientConnection = client
         self.request: HttpParser = request
         self.event_queue = event_queue
         self.upstream_conn_pool = upstream_conn_pool

proxy/http/proxy/plugin.py

@@ -17,9 +17,9 @@
 
 from ..parser import HttpParser
 from ..descriptors import DescriptorsHandlerMixin
+from ..connection import HttpClientConnection
 
 from ...core.event import EventQueue
-from ...core.connection import TcpClientConnection
 
 if TYPE_CHECKING:   # pragma: no cover
     from ...core.connection import UpstreamConnectionPool
@@ -38,7 +38,7 @@ def __init__(
             self,
             uid: str,
             flags: argparse.Namespace,
-            client: TcpClientConnection,
+            client: HttpClientConnection,
             event_queue: EventQueue,
             upstream_conn_pool: Optional['UpstreamConnectionPool'] = None,
     ) -> None: