Repository: abhinavsingh/proxy.py
Pull request #995 (2022-01-16 10:19:09)

[Middleware] `HttpClientConnection` preparation
## This PR
- Adds `HttpClientConnection` skeleton, now used by HTTP layer core and plugins
- Make `work` core agnostic to work object construction e.g. `TcpClientConnection`, `HttpClientConnection` and other work objects must now be constructed by the implementation themselves using incoming `**kwargs`.

## Related PR
- https://github.com/abhinavsingh/proxy.py/pull/993

## Agenda: All plugins must use client.response not client.queue

Using `client.queue` works but imposes challenges for outgoing response handling/customization/inspection.  Instead all plugins must now call `client.response`.  Internally, it works the same as `client.queue`.  Just that now, plugins must queue a `HttpParser`, `WebsocketFrame` etc objects instead of `memoryview` or raw `bytes`.

By doing so it becomes possible for the core to provide a response based callback eco-system e.g. proxy plugins can modify/update responses queued by other plugins (currently only a single plugin can technically send response).  Also, we'll be able to add `middleware` support for proxy and web server plugins.

Note that `client.queue` itself will continue to work.  Just that such plugins will not be able to take advantage of core capabilities that will depend upon plugins use of `client.response`.   At some point in far future, `client.queue` method will be made private.


codecov[bot]: # [Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) Report
> Merging [#995](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) (3ba6847) into [develop](https://codecov.io/gh/abhinavsingh/proxy.py/commit/c6fceb639a5925994f04b15b333413693eb434eb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) (c6fceb6) will **decrease** coverage by `0.24%`.
> The diff coverage is `68.00%`.

[![Impacted file tree graph](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995/graphs/tree.svg?width=650&height=150&src=pr&token=Zh9J7b4la2&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)

```diff
@@             Coverage Diff             @@
##           develop     #995      +/-   ##
===========================================
- Coverage    87.11%   86.86%   -0.25%     
===========================================
  Files          145      146       +1     
  Lines         6284     6290       +6     
  Branches       626      629       +3     
===========================================
- Hits          5474     5464      -10     
- Misses         705      715      +10     
- Partials       105      111       +6     
```

| Flag | Coverage Δ | |
|---|---|---|
| pytest | `86.75% <68.00%> (-0.25%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | Coverage Δ | |
|---|---|---|
| [proxy/http/inspector/inspect\_traffic.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9pbnNwZWN0b3IvaW5zcGVjdF90cmFmZmljLnB5) | `44.44% <50.00%> (-5.56%)` | :arrow_down: |
| [proxy/http/inspector/transformer.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9pbnNwZWN0b3IvdHJhbnNmb3JtZXIucHk=) | `53.84% <50.00%> (-4.05%)` | :arrow_down: |
| [proxy/http/websocket/plugin.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC93ZWJzb2NrZXQvcGx1Z2luLnB5) | `73.91% <50.00%> (-7.91%)` | :arrow_down: |
| [proxy/http/handler.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9oYW5kbGVyLnB5) | `74.19% <60.00%> (-0.93%)` | :arrow_down: |
| [proxy/http/connection.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9jb25uZWN0aW9uLnB5) | `100.00% <100.00%> (ø)` | |
| [proxy/http/plugin.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9wbHVnaW4ucHk=) | `100.00% <100.00%> (ø)` | |
| [proxy/http/server/plugin.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9zZXJ2ZXIvcGx1Z2luLnB5) | `100.00% <100.00%> (ø)` | |
| [tests/http/test\_protocol\_handler.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-dGVzdHMvaHR0cC90ZXN0X3Byb3RvY29sX2hhbmRsZXIucHk=) | `100.00% <100.00%> (ø)` | |
| [tests/http/parser/test\_proxy\_protocol.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-dGVzdHMvaHR0cC9wYXJzZXIvdGVzdF9wcm94eV9wcm90b2NvbC5weQ==) | `96.07% <0.00%> (-3.93%)` | :arrow_down: |
| [tests/testing/test\_embed.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-dGVzdHMvdGVzdGluZy90ZXN0X2VtYmVkLnB5) | `90.90% <0.00%> (-3.04%)` | :arrow_down: |
| ... and [3 more](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh). Last update [c6fceb6...3ba6847](https://codecov.io/gh/abhinavsingh/proxy.py/pull/995?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh).


Patch:
diff --git a/docs/conf.py b/docs/conf.py
index adabd03..350dd21 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -308,12 +308,14 @@ nitpick_ignore = [
     (_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'),
diff --git a/examples/ssl_echo_server.py b/examples/ssl_echo_server.py
index 8103e01..c2b25a1 100644
--- a/examples/ssl_echo_server.py
+++ b/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 @@ from proxy.core.connection import TcpClientConnection
 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.
diff --git a/examples/tcp_echo_server.py b/examples/tcp_echo_server.py
index 309ac79..22030ce 100644
--- a/examples/tcp_echo_server.py
+++ b/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 @@ from proxy.core.connection import TcpClientConnection
 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)
 
diff --git a/proxy/core/acceptor/acceptor.py b/proxy/core/acceptor/acceptor.py
index f578e0a..d83c5e0 100644
--- a/proxy/core/acceptor/acceptor.py
+++ b/proxy/core/acceptor/acceptor.py
@@ -224,7 +224,8 @@ class Acceptor(multiprocessing.Process):
                 ),
             )
             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 @@ class Acceptor(multiprocessing.Process):
                 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,
diff --git a/proxy/core/base/tcp_tunnel.py b/proxy/core/base/tcp_tunnel.py
index 20f0bbe..50320aa 100644
--- a/proxy/core/base/tcp_tunnel.py
+++ b/proxy/core/base/tcp_tunnel.py
@@ -47,6 +47,10 @@ class BaseTcpTunnelHandler(BaseTcpServerHandler[TcpClientConnection]):
     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)
 
diff --git a/proxy/core/connection/pool.py b/proxy/core/connection/pool.py
index 399aa59..379875e 100644
--- a/proxy/core/connection/pool.py
+++ b/proxy/core/connection/pool.py
@@ -16,7 +16,7 @@ import socket
 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 @@ class UpstreamConnectionPool(Work[TcpServerConnection]):
         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 @@ class UpstreamConnectionPool(Work[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(
diff --git a/proxy/core/ssh/handler.py b/proxy/core/ssh/handler.py
index fdf6f3a..1c021f8 100644
--- a/proxy/core/ssh/handler.py
+++ b/proxy/core/ssh/handler.py
@@ -12,7 +12,7 @@ import argparse
 
 from typing import TYPE_CHECKING, Tuple
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:   # pragma: no cover
     try:
         from paramiko.channel import Channel
     except ImportError:
diff --git a/proxy/core/ssh/listener.py b/proxy/core/ssh/listener.py
index 47c9b41..352b60d 100644
--- a/proxy/core/ssh/listener.py
+++ b/proxy/core/ssh/listener.py
@@ -16,7 +16,7 @@ from typing import TYPE_CHECKING, Any, Callable, Optional, Set, Tuple
 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
diff --git a/proxy/core/work/threaded.py b/proxy/core/work/threaded.py
index 4e583a0..7366f8a 100644
--- a/proxy/core/work/threaded.py
+++ b/proxy/core/work/threaded.py
@@ -12,25 +12,27 @@ import socket
 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,
diff --git a/proxy/core/work/threadless.py b/proxy/core/work/threadless.py
index bf50507..5ff7713 100644
--- a/proxy/core/work/threadless.py
+++ b/proxy/core/work/threadless.py
@@ -25,7 +25,6 @@ from ...common.types import Readables, SelectableEvents, Writables
 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 @@ class Threadless(ABC, Generic[T]):
         )
         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,
             ),
diff --git a/proxy/core/work/work.py b/proxy/core/work/work.py
index 5c94d2b..7747d26 100644
--- a/proxy/core/work/work.py
+++ b/proxy/core/work/work.py
@@ -47,6 +47,14 @@ class Work(ABC, Generic[T]):
         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."""
diff --git a/proxy/http/__init__.py b/proxy/http/__init__.py
index b918c3e..9c800b5 100644
--- a/proxy/http/__init__.py
+++ b/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 @@ from .url import Url
 
 __all__ = [
     'HttpProtocolHandler',
+    'HttpClientConnection',
     'HttpProtocolHandlerPlugin',
     'httpStatusCodes',
     'httpMethods',
diff --git a/proxy/http/connection.py b/proxy/http/connection.py
new file mode 100644
index 0000000..d31d340
--- /dev/null
+++ b/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
diff --git a/proxy/http/handler.py b/proxy/http/handler.py
index 1d80b4e..7b4a306 100644
--- a/proxy/http/handler.py
+++ b/proxy/http/handler.py
@@ -19,10 +19,10 @@ import selectors
 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 @@ from .parser import HttpParser, httpParserStates, httpParserTypes
 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 @@ class HttpProtocolHandler(BaseTcpServerHandler[TcpClientConnection]):
     # 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 @@ class HttpProtocolHandler(BaseTcpServerHandler[TcpClientConnection]):
                 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
 
diff --git a/proxy/http/inspector/inspect_traffic.py b/proxy/http/inspector/inspect_traffic.py
index 456626d..60225f5 100644
--- a/proxy/http/inspector/inspect_traffic.py
+++ b/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 @@ class InspectTrafficPlugin(WebSocketTransportBasePlugin):
             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(
diff --git a/proxy/http/inspector/transformer.py b/proxy/http/inspector/transformer.py
index da7e1d6..ea9276d 100644
--- a/proxy/http/inspector/transformer.py
+++ b/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']
diff --git a/proxy/http/plugin.py b/proxy/http/plugin.py
index fe44999..3b587a1 100644
--- a/proxy/http/plugin.py
+++ b/proxy/http/plugin.py
@@ -15,9 +15,9 @@ from abc import ABC, abstractmethod
 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,7 +55,7 @@ class HttpProtocolHandlerPlugin(
             self,
             uid: str,
             flags: argparse.Namespace,
-            client: TcpClientConnection,
+            client: HttpClientConnection,
             request: HttpParser,
             event_queue: Optional[EventQueue] = None,
             upstream_conn_pool: Optional['UpstreamConnectionPool'] = None,
@@ -63,7 +63,7 @@ class HttpProtocolHandlerPlugin(
         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
diff --git a/proxy/http/proxy/plugin.py b/proxy/http/proxy/plugin.py
index 2f893d2..9f6f114 100644
--- a/proxy/http/proxy/plugin.py
+++ b/proxy/http/proxy/plugin.py
@@ -17,9 +17,9 @@ from ..mixins import TlsInterceptionPropertyMixin
 
 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 @@ class HttpProxyBasePlugin(
             self,
             uid: str,
             flags: argparse.Namespace,
-            client: TcpClientConnection,
+            client: HttpClientConnection,
             event_queue: EventQueue,
             upstream_conn_pool: Optional['UpstreamConnectionPool'] = None,
     ) -> None:
diff --git a/proxy/http/proxy/server.py b/proxy/http/proxy/server.py
index f0ccf11..3008898 100644
--- a/proxy/http/proxy/server.py
+++ b/proxy/http/proxy/server.py
@@ -196,7 +196,7 @@ class HttpProxyPlugin(HttpProtocolHandlerPlugin):
                 )
                 return False
             except BrokenPipeError:
-                logger.error(
+                logger.warning(
                     'BrokenPipeError when flushing buffer for server',
                 )
                 return self._close_and_release()
@@ -840,7 +840,7 @@ class HttpProxyPlugin(HttpProtocolHandlerPlugin):
                 )
             do_close = True
         except BrokenPipeError:
-            logger.error(
+            logger.warning(
                 'BrokenPipeError when wrapping client for upstream: {0}'.format(
                     self.upstream.addr[0],
                 ),
diff --git a/proxy/http/server/plugin.py b/proxy/http/server/plugin.py
index dde91a9..069b0df 100644
--- a/proxy/http/server/plugin.py
+++ b/proxy/http/server/plugin.py
@@ -16,8 +16,8 @@ from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING
 from ..websocket import WebsocketFrame
 from ..parser import HttpParser
 from ..descriptors import DescriptorsHandlerMixin
+from ..connection import HttpClientConnection
 
-from ...core.connection import TcpClientConnection
 from ...core.event import EventQueue
 
 if TYPE_CHECKING:   # pragma: no cover
@@ -31,7 +31,7 @@ class HttpWebServerBasePlugin(DescriptorsHandlerMixin, ABC):
             self,
             uid: str,
             flags: argparse.Namespace,
-            client: TcpClientConnection,
+            client: HttpClientConnection,
             event_queue: EventQueue,
             upstream_conn_pool: Optional['UpstreamConnectionPool'] = None,
     ):
diff --git a/proxy/http/websocket/plugin.py b/proxy/http/websocket/plugin.py
index f688163..f7f2748 100644
--- a/proxy/http/websocket/plugin.py
+++ b/proxy/http/websocket/plugin.py
@@ -11,13 +11,16 @@
 import argparse
 import json
 from abc import ABC, abstractmethod
-from typing import List, Dict, Any
+from typing import TYPE_CHECKING, List, Dict, Any
 
 from ...common.utils import bytes_
-from . import WebsocketFrame
-from ...core.connection import TcpClientConnection
 from ...core.event import EventQueue
 
+from . import WebsocketFrame
+
+if TYPE_CHECKING:   # pragma: no cover
+    from ..connection import HttpClientConnection
+
 
 class WebSocketTransportBasePlugin(ABC):
     """Abstract class for plugins extending dashboard websocket API."""
@@ -25,7 +28,7 @@ class WebSocketTransportBasePlugin(ABC):
     def __init__(
             self,
             flags: argparse.Namespace,
-            client: TcpClientConnection,
+            client: 'HttpClientConnection',
             event_queue: EventQueue,
     ) -> None:
         self.flags = flags


Test patch:
diff --git a/tests/core/test_acceptor.py b/tests/core/test_acceptor.py
index b18f640..376186c 100644
--- a/tests/core/test_acceptor.py
+++ b/tests/core/test_acceptor.py
@@ -24,9 +24,10 @@ class TestAcceptor(unittest.TestCase):
     def setUp(self) -> None:
         self.acceptor_id = 1
         self.pipe = multiprocessing.Pipe()
+        self.work_klass = mock.MagicMock()
         self.flags = FlagParser.initialize(
             threaded=True,
-            work_klass=mock.MagicMock(),
+            work_klass=self.work_klass,
             local_executor=0,
         )
         self.acceptor = Acceptor(
@@ -63,7 +64,6 @@ class TestAcceptor(unittest.TestCase):
         sock.accept.assert_not_called()
         self.flags.work_klass.assert_not_called()
 
-    @mock.patch('proxy.core.work.threaded.TcpClientConnection')
     @mock.patch('threading.Thread')
     @mock.patch('selectors.DefaultSelector')
     @mock.patch('socket.fromfd')
@@ -74,7 +74,6 @@ class TestAcceptor(unittest.TestCase):
             mock_fromfd: mock.Mock,
             mock_selector: mock.Mock,
             mock_thread: mock.Mock,
-            mock_client: mock.Mock,
     ) -> None:
         fileno = 10
         conn = mock.MagicMock()
@@ -99,7 +98,7 @@ class TestAcceptor(unittest.TestCase):
             type=socket.SOCK_STREAM,
         )
         self.flags.work_klass.assert_called_with(
-            mock_client.return_value,
+            self.work_klass.create.return_value,
             flags=self.flags,
             event_queue=None,
             upstream_conn_pool=None,
diff --git a/tests/core/test_conn_pool.py b/tests/core/test_conn_pool.py
index 6547a0d..3918372 100644
--- a/tests/core/test_conn_pool.py
+++ b/tests/core/test_conn_pool.py
@@ -33,7 +33,10 @@ class TestConnectionPool(unittest.TestCase):
         mock_conn.closed = False
         # Acquire
         created, conn = pool.acquire(addr)
-        mock_tcp_server_connection.assert_called_once_with(addr[0], addr[1])
+        mock_tcp_server_connection.assert_called_once_with(
+            host=addr[0],
+            port=addr[1],
+        )
         mock_conn.mark_inuse.assert_called_once()
         mock_conn.reset.assert_not_called()
         self.assertTrue(created)
@@ -66,7 +69,10 @@ class TestConnectionPool(unittest.TestCase):
         # Acquire
         created, conn = pool.acquire(addr)
         self.assertTrue(created)
-        mock_tcp_server_connection.assert_called_once_with(addr[0], addr[1])
+        mock_tcp_server_connection.assert_called_once_with(
+            host=addr[0],
+            port=addr[1],
+        )
         self.assertEqual(conn, mock_conn)
         self.assertEqual(len(pool.pools[addr]), 1)
         self.assertTrue(conn in pool.pools[addr])
@@ -91,7 +97,10 @@ class TestConnectionPoolAsync:
         mock_conn = mock_tcp_server_connection.return_value
         addr = mock_conn.addr
         pool.add(addr)
-        mock_tcp_server_connection.assert_called_once_with(addr[0], addr[1])
+        mock_tcp_server_connection.assert_called_once_with(
+            host=addr[0],
+            port=addr[1],
+        )
         mock_conn.connect.assert_called_once()
         events = await pool.get_events()
         print(events)
diff --git a/tests/http/exceptions/test_http_proxy_auth_failed.py b/tests/http/exceptions/test_http_proxy_auth_failed.py
index e0695bc..5d74aec 100644
--- a/tests/http/exceptions/test_http_proxy_auth_failed.py
+++ b/tests/http/exceptions/test_http_proxy_auth_failed.py
@@ -14,11 +14,10 @@ import pytest
 
 from pytest_mock import MockerFixture
 
-from proxy.http import HttpProtocolHandler, httpHeaders
+from proxy.http import HttpProtocolHandler, httpHeaders, HttpClientConnection
 from proxy.common.flag import FlagParser
 from proxy.common.utils import build_http_request
 from proxy.http.responses import PROXY_AUTH_FAILED_RESPONSE_PKT
-from proxy.core.connection import TcpClientConnection
 from ...test_assertions import Assertions
 
 
@@ -39,7 +38,7 @@ class TestHttpProxyAuthFailed(Assertions):
         )
         self._conn = self.mock_fromfd.return_value
         self.protocol_handler = HttpProtocolHandler(
-            TcpClientConnection(self._conn, self._addr),
+            HttpClientConnection(self._conn, self._addr),
             flags=self.flags,
         )
         self.protocol_handler.initialize()
diff --git a/tests/http/proxy/test_http_proxy.py b/tests/http/proxy/test_http_proxy.py
index 54789ff..68f6bfd 100644
--- a/tests/http/proxy/test_http_proxy.py
+++ b/tests/http/proxy/test_http_proxy.py
@@ -14,12 +14,11 @@ import pytest
 
 from pytest_mock import MockerFixture
 
-from proxy.http import HttpProtocolHandler
+from proxy.http import HttpProtocolHandler, HttpClientConnection
 from proxy.http.proxy import HttpProxyPlugin
 from proxy.common.flag import FlagParser
 from proxy.common.utils import build_http_request
 from proxy.http.exception import HttpProtocolException
-from proxy.core.connection import TcpClientConnection
 from proxy.common.constants import DEFAULT_HTTP_PORT
 
 
@@ -43,7 +42,7 @@ class TestHttpProxyPlugin:
         }
         self._conn = self.mock_fromfd.return_value
         self.protocol_handler = HttpProtocolHandler(
-            TcpClientConnection(self._conn, self._addr),
+            HttpClientConnection(self._conn, self._addr),
             flags=self.flags,
         )
         self.protocol_handler.initialize()
diff --git a/tests/http/proxy/test_http_proxy_tls_interception.py b/tests/http/proxy/test_http_proxy_tls_interception.py
index 4acff4d..b13c748 100644
--- a/tests/http/proxy/test_http_proxy_tls_interception.py
+++ b/tests/http/proxy/test_http_proxy_tls_interception.py
@@ -19,12 +19,12 @@ from unittest import mock
 
 from pytest_mock import MockerFixture
 
-from proxy.http import HttpProtocolHandler, httpMethods
+from proxy.http import HttpProtocolHandler, httpMethods, HttpClientConnection
 from proxy.http.proxy import HttpProxyPlugin
 from proxy.common.flag import FlagParser
 from proxy.common.utils import bytes_, build_http_request
 from proxy.http.responses import PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT
-from proxy.core.connection import TcpClientConnection, TcpServerConnection
+from proxy.core.connection import TcpServerConnection
 from proxy.common.constants import DEFAULT_CA_FILE
 from ...test_assertions import Assertions
 
@@ -88,7 +88,7 @@ class TestHttpProxyTlsInterception(Assertions):
         }
         self._conn = self.mock_fromfd.return_value
         self.protocol_handler = HttpProtocolHandler(
-            TcpClientConnection(self._conn, self._addr),
+            HttpClientConnection(self._conn, self._addr),
             flags=self.flags,
         )
         self.protocol_handler.initialize()
diff --git a/tests/http/test_protocol_handler.py b/tests/http/test_protocol_handler.py
index 6faf99f..ade2101 100644
--- a/tests/http/test_protocol_handler.py
+++ b/tests/http/test_protocol_handler.py
@@ -17,7 +17,7 @@ from unittest import mock
 
 from pytest_mock import MockerFixture
 
-from proxy.http import HttpProtocolHandler, httpHeaders
+from proxy.http import HttpProtocolHandler, httpHeaders, HttpClientConnection
 from proxy.http.proxy import HttpProxyPlugin
 from proxy.common.flag import FlagParser
 from proxy.http.parser import HttpParser, httpParserTypes, httpParserStates
@@ -28,7 +28,6 @@ from proxy.http.responses import (
     BAD_GATEWAY_RESPONSE_PKT, PROXY_AUTH_FAILED_RESPONSE_PKT,
     PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT,
 )
-from proxy.core.connection import TcpClientConnection
 from proxy.common.constants import (
     CRLF, PLUGIN_HTTP_PROXY, PLUGIN_PROXY_AUTH, PLUGIN_WEB_SERVER,
 )
@@ -68,7 +67,7 @@ class TestHttpProtocolHandlerWithoutServerMock(Assertions):
         ])
 
         self.protocol_handler = HttpProtocolHandler(
-            TcpClientConnection(self._conn, self._addr),
+            HttpClientConnection(self._conn, self._addr),
             flags=self.flags,
         )
         self.protocol_handler.initialize()
@@ -101,7 +100,7 @@ class TestHttpProtocolHandlerWithoutServerMock(Assertions):
             bytes_(PLUGIN_PROXY_AUTH),
         ])
         self.protocol_handler = HttpProtocolHandler(
-            TcpClientConnection(self._conn, self._addr), flags=flags,
+            HttpClientConnection(self._conn, self._addr), flags=flags,
         )
         self.protocol_handler.initialize()
         self._conn.recv.return_value = CRLF.join([
@@ -138,7 +137,7 @@ class TestHttpProtocolHandler(Assertions):
         ])
 
         self.protocol_handler = HttpProtocolHandler(
-            TcpClientConnection(self._conn, self._addr),
+            HttpClientConnection(self._conn, self._addr),
             flags=self.flags,
         )
         self.protocol_handler.initialize()
@@ -301,7 +300,7 @@ class TestHttpProtocolHandler(Assertions):
         ])
 
         self.protocol_handler = HttpProtocolHandler(
-            TcpClientConnection(self._conn, self._addr), flags=flags,
+            HttpClientConnection(self._conn, self._addr), flags=flags,
         )
         self.protocol_handler.initialize()
         assert self.http_server_port is not None
@@ -349,7 +348,7 @@ class TestHttpProtocolHandler(Assertions):
         ])
 
         self.protocol_handler = HttpProtocolHandler(
-            TcpClientConnection(self._conn, self._addr), flags=flags,
+            HttpClientConnection(self._conn, self._addr), flags=flags,
         )
         self.protocol_handler.initialize()
 
diff --git a/tests/http/web/test_web_server.py b/tests/http/web/test_web_server.py
index bae92dd..e5bbabb 100644
--- a/tests/http/web/test_web_server.py
+++ b/tests/http/web/test_web_server.py
@@ -18,13 +18,12 @@ import pytest
 
 from pytest_mock import MockerFixture
 
-from proxy.http import HttpProtocolHandler
+from proxy.http import HttpProtocolHandler, HttpClientConnection
 from proxy.common.flag import FlagParser
 from proxy.http.parser import HttpParser, httpParserTypes, httpParserStates
 from proxy.common.utils import bytes_, build_http_request, build_http_response
 from proxy.common.plugins import Plugins
 from proxy.http.responses import NOT_FOUND_RESPONSE_PKT
-from proxy.core.connection import TcpClientConnection
 from proxy.common.constants import (
     CRLF, PROXY_PY_DIR, PLUGIN_PAC_FILE, PLUGIN_HTTP_PROXY, PLUGIN_WEB_SERVER,
 )
@@ -48,7 +47,7 @@ def test_on_client_connection_called_on_teardown(mocker: MockerFixture) -> None:
     _conn = mock_fromfd.return_value
     _addr = ('127.0.0.1', 54382)
     protocol_handler = HttpProtocolHandler(
-        TcpClientConnection(_conn, _addr),
+        HttpClientConnection(_conn, _addr),
         flags=flags,
     )
     protocol_handler.initialize()
@@ -81,7 +80,7 @@ def mock_selector_for_client_read(self: Any) -> None:
 #     flags.plugins = {b'HttpProtocolHandlerPlugin': [plugin]}
 #     self._conn = mock_fromfd.return_value
 #     self.protocol_handler = HttpProtocolHandler(
-#         TcpClientConnection(self._conn, self._addr),
+#         HttpClientConnection(self._conn, self._addr),
 #         flags=flags,
 #     )
 #     self.protocol_handler.initialize()
@@ -101,7 +100,7 @@ def mock_selector_for_client_read(self: Any) -> None:
 #     flags.plugins = {b'HttpProtocolHandlerPlugin': [plugin]}
 #     self._conn = mock_fromfd.return_value
 #     self.protocol_handler = HttpProtocolHandler(
-#         TcpClientConnection(self._conn, self._addr),
+#         HttpClientConnection(self._conn, self._addr),
 #         flags=flags,
 #     )
 #     self.protocol_handler.initialize()
@@ -121,7 +120,7 @@ def mock_selector_for_client_read(self: Any) -> None:
 #     flags.plugins = {b'HttpProtocolHandlerPlugin': [plugin]}
 #     self._conn = mock_fromfd.return_value
 #     self.protocol_handler = HttpProtocolHandler(
-#         TcpClientConnection(self._conn, self._addr),
+#         HttpClientConnection(self._conn, self._addr),
 #         flags=flags,
 #     )
 #     self.protocol_handler.initialize()
@@ -162,7 +161,7 @@ class TestWebServerPluginWithPacFilePlugin(Assertions):
             bytes_(PLUGIN_PAC_FILE),
         ])
         self.protocol_handler = HttpProtocolHandler(
-            TcpClientConnection(self._conn, self._addr),
+            HttpClientConnection(self._conn, self._addr),
             flags=self.flags,
         )
         self.protocol_handler.initialize()
@@ -221,7 +220,7 @@ class TestStaticWebServerPlugin(Assertions):
             bytes_(PLUGIN_WEB_SERVER),
         ])
         self.protocol_handler = HttpProtocolHandler(
-            TcpClientConnection(self._conn, self._addr),
+            HttpClientConnection(self._conn, self._addr),
             flags=flags,
         )
         self.protocol_handler.initialize()
@@ -328,7 +327,7 @@ class TestWebServerPlugin(Assertions):
             bytes_(PLUGIN_WEB_SERVER),
         ])
         self.protocol_handler = HttpProtocolHandler(
-            TcpClientConnection(self._conn, self._addr),
+            HttpClientConnection(self._conn, self._addr),
             flags=self.flags,
         )
         self.protocol_handler.initialize()
@@ -353,7 +352,7 @@ class TestWebServerPlugin(Assertions):
             bytes_(PLUGIN_WEB_SERVER),
         ])
         self.protocol_handler = HttpProtocolHandler(
-            TcpClientConnection(self._conn, self._addr),
+            HttpClientConnection(self._conn, self._addr),
             flags=flags,
         )
         self.protocol_handler.initialize()
diff --git a/tests/plugin/test_http_proxy_plugins.py b/tests/plugin/test_http_proxy_plugins.py
index 3550c32..742b7fd 100644
--- a/tests/plugin/test_http_proxy_plugins.py
+++ b/tests/plugin/test_http_proxy_plugins.py
@@ -20,13 +20,12 @@ from unittest import mock
 
 from pytest_mock import MockerFixture
 
-from proxy.http import HttpProtocolHandler, httpStatusCodes
+from proxy.http import HttpProtocolHandler, httpStatusCodes, HttpClientConnection
 from proxy.plugin import ProposedRestApiPlugin, RedirectToCustomServerPlugin
 from proxy.http.proxy import HttpProxyPlugin
 from proxy.common.flag import FlagParser
 from proxy.http.parser import HttpParser, httpParserTypes
 from proxy.common.utils import bytes_, build_http_request, build_http_response
-from proxy.core.connection import TcpClientConnection
 from proxy.common.constants import DEFAULT_HTTP_PORT, PROXY_AGENT_HEADER_VALUE
 from .utils import get_plugin_by_test_name
 from ..test_assertions import Assertions
@@ -64,7 +63,7 @@ class TestHttpProxyPluginExamples(Assertions):
         }
         self._conn = self.mock_fromfd.return_value
         self.protocol_handler = HttpProtocolHandler(
-            TcpClientConnection(self._conn, self._addr),
+            HttpClientConnection(self._conn, self._addr),
             flags=self.flags,
         )
         self.protocol_handler.initialize()
diff --git a/tests/plugin/test_http_proxy_plugins_with_tls_interception.py b/tests/plugin/test_http_proxy_plugins_with_tls_interception.py
index 55d09b6..a985f0c 100644
--- a/tests/plugin/test_http_proxy_plugins_with_tls_interception.py
+++ b/tests/plugin/test_http_proxy_plugins_with_tls_interception.py
@@ -18,7 +18,7 @@ import pytest
 
 from pytest_mock import MockerFixture
 
-from proxy.http import HttpProtocolHandler, httpMethods
+from proxy.http import HttpProtocolHandler, httpMethods, HttpClientConnection
 from proxy.http.proxy import HttpProxyPlugin
 from proxy.common.flag import FlagParser
 from proxy.http.parser import HttpParser, httpParserTypes
@@ -26,7 +26,7 @@ from proxy.common.utils import bytes_, build_http_request
 from proxy.http.responses import (
     PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, okResponse,
 )
-from proxy.core.connection import TcpClientConnection, TcpServerConnection
+from proxy.core.connection import TcpServerConnection
 from .utils import get_plugin_by_test_name
 from ..test_assertions import Assertions
 
@@ -71,7 +71,7 @@ class TestHttpProxyPluginExamplesWithTlsInterception(Assertions):
         self._conn = mocker.MagicMock(spec=socket.socket)
         self.mock_fromfd.return_value = self._conn
         self.protocol_handler = HttpProtocolHandler(
-            TcpClientConnection(self._conn, self._addr), flags=self.flags,
+            HttpClientConnection(self._conn, self._addr), flags=self.flags,
         )
         self.protocol_handler.initialize()