Repository: abhinavsingh/proxy.py
Pull request #1048 (2022-01-23 17:17:39)

[Work] Refactor for generic work type
- Currently work framework assumes a client connection object as work
- Refactor to allow custom definition of work objects

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

[![Impacted file tree graph](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1048/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/1048?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    #1048      +/-   ##
===========================================
- Coverage    87.11%   80.05%   -7.07%     
===========================================
  Files          156      157       +1     
  Lines         6591     6603      +12     
  Branches       668      668              
===========================================
- Hits          5742     5286     -456     
- Misses         738     1191     +453     
- Partials       111      126      +15     
```

| Flag | Coverage Δ | |
|---|---|---|
| pytest | `80.03% <73.21%> (-6.90%)` | :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/1048?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | Coverage Δ | |
|---|---|---|
| [proxy/core/work/fd.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1048/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29yZS93b3JrL2ZkLnB5) | `45.83% <45.83%> (ø)` | |
| [proxy/core/work/local.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1048/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29yZS93b3JrL2xvY2FsLnB5) | `50.00% <66.66%> (-50.00%)` | :arrow_down: |
| [proxy/core/work/remote.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1048/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29yZS93b3JrL3JlbW90ZS5weQ==) | `53.57% <66.66%> (-42.86%)` | :arrow_down: |
| [proxy/common/types.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1048/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29tbW9uL3R5cGVzLnB5) | `100.00% <100.00%> (ø)` | |
| [proxy/core/acceptor/acceptor.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1048/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29yZS9hY2NlcHRvci9hY2NlcHRvci5weQ==) | `76.36% <100.00%> (-20.91%)` | :arrow_down: |
| [proxy/core/base/tcp\_server.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1048/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29yZS9iYXNlL3RjcF9zZXJ2ZXIucHk=) | `83.75% <100.00%> (-5.14%)` | :arrow_down: |
| [proxy/core/connection/client.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1048/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29yZS9jb25uZWN0aW9uL2NsaWVudC5weQ==) | `100.00% <100.00%> (ø)` | |
| [proxy/core/connection/connection.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1048/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29yZS9jb25uZWN0aW9uL2Nvbm5lY3Rpb24ucHk=) | `89.83% <100.00%> (-0.17%)` | :arrow_down: |
| [proxy/core/connection/server.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1048/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29yZS9jb25uZWN0aW9uL3NlcnZlci5weQ==) | `100.00% <100.00%> (ø)` | |
| [proxy/core/work/threadless.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1048/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29yZS93b3JrL3RocmVhZGxlc3MucHk=) | `18.50% <100.00%> (-64.05%)` | :arrow_down: |
| ... and [32 more](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1048/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/1048?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/1048?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh). Last update [d7a568e...d305ebe](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1048?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 452a09a..a370246 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -317,7 +317,10 @@ nitpick_ignore = [
     (_py_class_role, 'connection.Connection'),
     (_py_class_role, 'EventQueue'),
     (_py_class_role, 'T'),
+    (_py_class_role, 'HostPort'),
+    (_py_class_role, 'TcpOrTlsSocket'),
     (_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'),
+    (_py_obj_role, 'proxy.core.work.fd.T'),
 ]
diff --git a/proxy/common/types.py b/proxy/common/types.py
index bc230c9..603dd95 100644
--- a/proxy/common/types.py
+++ b/proxy/common/types.py
@@ -8,7 +8,9 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
+import ssl
 import queue
+import socket
 import ipaddress
 from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union
 
@@ -25,5 +27,6 @@ SelectableEvents = Dict[Selectable, int]    # Values are event masks
 Readables = Selectables
 Writables = Selectables
 Descriptors = Tuple[Readables, Writables]
-
 IpAddress = Union[ipaddress.IPv4Address, ipaddress.IPv6Address]
+TcpOrTlsSocket = Union[ssl.SSLSocket, socket.socket]
+HostPort = Tuple[str, int]
diff --git a/proxy/common/utils.py b/proxy/common/utils.py
index 2141e2c..324dc16 100644
--- a/proxy/common/utils.py
+++ b/proxy/common/utils.py
@@ -23,6 +23,7 @@ import contextlib
 from types import TracebackType
 from typing import Any, Dict, List, Type, Tuple, Callable, Optional
 
+from .types import HostPort
 from .constants import (
     CRLF, COLON, HTTP_1_1, IS_WINDOWS, WHITESPACE, DEFAULT_TIMEOUT,
     DEFAULT_THREADLESS,
@@ -220,9 +221,9 @@ def wrap_socket(
 
 
 def new_socket_connection(
-        addr: Tuple[str, int],
+        addr: HostPort,
         timeout: float = DEFAULT_TIMEOUT,
-        source_address: Optional[Tuple[str, int]] = None,
+        source_address: Optional[HostPort] = None,
 ) -> socket.socket:
     conn = None
     try:
@@ -252,8 +253,8 @@ def new_socket_connection(
 class socket_connection(contextlib.ContextDecorator):
     """Same as new_socket_connection but as a context manager and decorator."""
 
-    def __init__(self, addr: Tuple[str, int]):
-        self.addr: Tuple[str, int] = addr
+    def __init__(self, addr: HostPort):
+        self.addr: HostPort = addr
         self.conn: Optional[socket.socket] = None
         super().__init__()
 
diff --git a/proxy/core/acceptor/acceptor.py b/proxy/core/acceptor/acceptor.py
index 74721de..640595f 100644
--- a/proxy/core/acceptor/acceptor.py
+++ b/proxy/core/acceptor/acceptor.py
@@ -26,6 +26,7 @@ from multiprocessing.reduction import recv_handle
 from ..work import LocalExecutor, start_threaded_work, delegate_work_to_pool
 from ..event import EventQueue
 from ...common.flag import flags
+from ...common.types import HostPort
 from ...common.logger import Logger
 from ...common.backports import NonBlockingQueue
 from ...common.constants import DEFAULT_LOCAL_EXECUTOR
@@ -104,7 +105,7 @@ class Acceptor(multiprocessing.Process):
     def accept(
             self,
             events: List[Tuple[selectors.SelectorKey, int]],
-    ) -> List[Tuple[socket.socket, Optional[Tuple[str, int]]]]:
+    ) -> List[Tuple[socket.socket, Optional[HostPort]]]:
         works = []
         for key, mask in events:
             if mask & selectors.EVENT_READ:
@@ -156,8 +157,8 @@ class Acceptor(multiprocessing.Process):
             self.flags.log_format,
         )
         self.selector = selectors.DefaultSelector()
-        self._recv_and_setup_socks()
         try:
+            self._recv_and_setup_socks()
             if self.flags.threadless and self.flags.local_executor:
                 self._start_local()
             for fileno in self.socks:
@@ -209,7 +210,7 @@ class Acceptor(multiprocessing.Process):
             self._local_work_queue.put(False)
             self._lthread.join()
 
-    def _work(self, conn: socket.socket, addr: Optional[Tuple[str, int]]) -> None:
+    def _work(self, conn: socket.socket, addr: Optional[HostPort]) -> None:
         self._total = self._total or 0
         if self.flags.threadless:
             # Index of worker to which this work should be dispatched
diff --git a/proxy/core/base/tcp_server.py b/proxy/core/base/tcp_server.py
index 107e1ac..dfc368f 100644
--- a/proxy/core/base/tcp_server.py
+++ b/proxy/core/base/tcp_server.py
@@ -12,16 +12,17 @@
 
        tcp
 """
-import ssl
 import socket
 import logging
 import selectors
 from abc import abstractmethod
-from typing import Any, Union, TypeVar, Optional
+from typing import Any, TypeVar, Optional
 
 from ...core.work import Work
 from ...common.flag import flags
-from ...common.types import Readables, Writables, SelectableEvents
+from ...common.types import (
+    Readables, Writables, TcpOrTlsSocket, SelectableEvents,
+)
 from ...common.utils import wrap_socket
 from ...core.connection import TcpClientConnection
 from ...common.constants import (
@@ -208,9 +209,7 @@ class BaseTcpServerHandler(Work[T]):
         return self.flags.keyfile is not None and \
             self.flags.certfile is not None
 
-    def _optionally_wrap_socket(
-            self, conn: socket.socket,
-    ) -> Union[ssl.SSLSocket, socket.socket]:
+    def _optionally_wrap_socket(self, conn: socket.socket) -> TcpOrTlsSocket:
         """Attempts to wrap accepted client connection using provided certificates.
 
         Shutdown and closes client connection upon error.
diff --git a/proxy/core/connection/client.py b/proxy/core/connection/client.py
index 0ba6388..f241c56 100644
--- a/proxy/core/connection/client.py
+++ b/proxy/core/connection/client.py
@@ -9,11 +9,11 @@
     :license: BSD, see LICENSE for more details.
 """
 import ssl
-import socket
-from typing import Tuple, Union, Optional
+from typing import Optional
 
 from .types import tcpConnectionTypes
 from .connection import TcpConnection, TcpConnectionUninitializedException
+from ...common.types import HostPort, TcpOrTlsSocket
 
 
 class TcpClientConnection(TcpConnection):
@@ -21,20 +21,20 @@ class TcpClientConnection(TcpConnection):
 
     def __init__(
         self,
-        conn: Union[ssl.SSLSocket, socket.socket],
+        conn: TcpOrTlsSocket,
         # optional for unix socket servers
-        addr: Optional[Tuple[str, int]] = None,
+        addr: Optional[HostPort] = None,
     ) -> None:
         super().__init__(tcpConnectionTypes.CLIENT)
-        self._conn: Optional[Union[ssl.SSLSocket, socket.socket]] = conn
-        self.addr: Optional[Tuple[str, int]] = addr
+        self._conn: Optional[TcpOrTlsSocket] = conn
+        self.addr: Optional[HostPort] = addr
 
     @property
     def address(self) -> str:
         return 'unix:client' if not self.addr else '{0}:{1}'.format(self.addr[0], self.addr[1])
 
     @property
-    def connection(self) -> Union[ssl.SSLSocket, socket.socket]:
+    def connection(self) -> TcpOrTlsSocket:
         if self._conn is None:
             raise TcpConnectionUninitializedException()
         return self._conn
diff --git a/proxy/core/connection/connection.py b/proxy/core/connection/connection.py
index 9aa0d77..cddf585 100644
--- a/proxy/core/connection/connection.py
+++ b/proxy/core/connection/connection.py
@@ -8,13 +8,12 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
-import ssl
-import socket
 import logging
 from abc import ABC, abstractmethod
-from typing import List, Union, Optional
+from typing import List, Optional
 
 from .types import tcpConnectionTypes
+from ...common.types import TcpOrTlsSocket
 from ...common.constants import DEFAULT_BUFFER_SIZE, DEFAULT_MAX_SEND_SIZE
 
 
@@ -44,7 +43,7 @@ class TcpConnection(ABC):
 
     @property
     @abstractmethod
-    def connection(self) -> Union[ssl.SSLSocket, socket.socket]:
+    def connection(self) -> TcpOrTlsSocket:
         """Must return the socket connection to use in this class."""
         raise TcpConnectionUninitializedException()     # pragma: no cover
 
diff --git a/proxy/core/connection/pool.py b/proxy/core/connection/pool.py
index 86c54a3..2536c5b 100644
--- a/proxy/core/connection/pool.py
+++ b/proxy/core/connection/pool.py
@@ -20,7 +20,7 @@ from typing import TYPE_CHECKING, Any, Set, Dict, Tuple
 from ..work import Work
 from .server import TcpServerConnection
 from ...common.flag import flags
-from ...common.types import Readables, Writables, SelectableEvents
+from ...common.types import HostPort, Readables, Writables, SelectableEvents
 
 
 logger = logging.getLogger(__name__)
@@ -73,13 +73,13 @@ class UpstreamConnectionPool(Work[TcpServerConnection]):
 
     def __init__(self) -> None:
         self.connections: Dict[int, TcpServerConnection] = {}
-        self.pools: Dict[Tuple[str, int], Set[TcpServerConnection]] = {}
+        self.pools: Dict[HostPort, Set[TcpServerConnection]] = {}
 
     @staticmethod
     def create(**kwargs: Any) -> TcpServerConnection:   # pragma: no cover
         return TcpServerConnection(**kwargs)
 
-    def acquire(self, addr: Tuple[str, int]) -> Tuple[bool, TcpServerConnection]:
+    def acquire(self, addr: HostPort) -> Tuple[bool, TcpServerConnection]:
         """Returns a reusable connection from the pool.
 
         If none exists, will create and return a new connection."""
@@ -147,7 +147,7 @@ class UpstreamConnectionPool(Work[TcpServerConnection]):
             self._remove(fileno)
         return False
 
-    def add(self, addr: Tuple[str, int]) -> TcpServerConnection:
+    def add(self, addr: HostPort) -> TcpServerConnection:
         """Creates, connects and adds a new connection to the pool.
 
         Returns newly created connection.
diff --git a/proxy/core/connection/server.py b/proxy/core/connection/server.py
index 15fae34..83ad0ae 100644
--- a/proxy/core/connection/server.py
+++ b/proxy/core/connection/server.py
@@ -9,11 +9,11 @@
     :license: BSD, see LICENSE for more details.
 """
 import ssl
-import socket
-from typing import Tuple, Union, Optional
+from typing import Optional
 
 from .types import tcpConnectionTypes
 from .connection import TcpConnection, TcpConnectionUninitializedException
+from ...common.types import HostPort, TcpOrTlsSocket
 from ...common.utils import new_socket_connection
 
 
@@ -22,20 +22,20 @@ class TcpServerConnection(TcpConnection):
 
     def __init__(self, host: str, port: int) -> None:
         super().__init__(tcpConnectionTypes.SERVER)
-        self._conn: Optional[Union[ssl.SSLSocket, socket.socket]] = None
-        self.addr: Tuple[str, int] = (host, port)
+        self._conn: Optional[TcpOrTlsSocket] = None
+        self.addr: HostPort = (host, port)
         self.closed = True
 
     @property
-    def connection(self) -> Union[ssl.SSLSocket, socket.socket]:
+    def connection(self) -> TcpOrTlsSocket:
         if self._conn is None:
             raise TcpConnectionUninitializedException()
         return self._conn
 
     def connect(
             self,
-            addr: Optional[Tuple[str, int]] = None,
-            source_address: Optional[Tuple[str, int]] = None,
+            addr: Optional[HostPort] = None,
+            source_address: Optional[HostPort] = None,
     ) -> None:
         assert self._conn is None
         self._conn = new_socket_connection(
diff --git a/proxy/core/ssh/handler.py b/proxy/core/ssh/handler.py
index 12557ed..ed6ea78 100644
--- a/proxy/core/ssh/handler.py
+++ b/proxy/core/ssh/handler.py
@@ -9,10 +9,11 @@
     :license: BSD, see LICENSE for more details.
 """
 import argparse
-from typing import TYPE_CHECKING, Tuple
+from typing import TYPE_CHECKING
 
 
 if TYPE_CHECKING:   # pragma: no cover
+    from ...common.types import HostPort
     try:
         from paramiko.channel import Channel
     except ImportError:
@@ -28,7 +29,7 @@ class SshHttpProtocolHandler:
     def on_connection(
             self,
             chan: 'Channel',
-            origin: Tuple[str, int],
-            server: Tuple[str, int],
+            origin: 'HostPort',
+            server: 'HostPort',
     ) -> None:
         pass
diff --git a/proxy/core/ssh/listener.py b/proxy/core/ssh/listener.py
index 1b17120..d851600 100644
--- a/proxy/core/ssh/listener.py
+++ b/proxy/core/ssh/listener.py
@@ -10,7 +10,7 @@
 """
 import logging
 import argparse
-from typing import TYPE_CHECKING, Any, Set, Tuple, Callable, Optional
+from typing import TYPE_CHECKING, Any, Set, Callable, Optional
 
 
 try:
@@ -18,6 +18,8 @@ try:
     from paramiko.transport import Transport
     if TYPE_CHECKING:   # pragma: no cover
         from paramiko.channel import Channel
+
+        from ...common.types import HostPort
 except ImportError:     # pragma: no cover
     pass
 
@@ -78,15 +80,15 @@ class SshTunnelListener:
     def __init__(
             self,
             flags: argparse.Namespace,
-            on_connection_callback: Callable[['Channel', Tuple[str, int], Tuple[str, int]], None],
+            on_connection_callback: Callable[['Channel', 'HostPort', 'HostPort'], None],
     ) -> None:
         self.flags = flags
         self.on_connection_callback = on_connection_callback
         self.ssh: Optional[SSHClient] = None
         self.transport: Optional[Transport] = None
-        self.forwarded: Set[Tuple[str, int]] = set()
+        self.forwarded: Set['HostPort'] = set()
 
-    def start_port_forward(self, remote_addr: Tuple[str, int]) -> None:
+    def start_port_forward(self, remote_addr: 'HostPort') -> None:
         assert self.transport is not None
         self.transport.request_port_forward(
             *remote_addr,
@@ -95,7 +97,7 @@ class SshTunnelListener:
         self.forwarded.add(remote_addr)
         logger.info('%s:%d forwarding successful...' % remote_addr)
 
-    def stop_port_forward(self, remote_addr: Tuple[str, int]) -> None:
+    def stop_port_forward(self, remote_addr: 'HostPort') -> None:
         assert self.transport is not None
         self.transport.cancel_port_forward(*remote_addr)
         self.forwarded.remove(remote_addr)
diff --git a/proxy/core/work/delegate.py b/proxy/core/work/delegate.py
index 76207d9..76f7e71 100644
--- a/proxy/core/work/delegate.py
+++ b/proxy/core/work/delegate.py
@@ -8,7 +8,7 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
-from typing import TYPE_CHECKING, Tuple, Optional
+from typing import TYPE_CHECKING, Optional
 from multiprocessing.reduction import send_handle
 
 
@@ -17,13 +17,15 @@ if TYPE_CHECKING:   # pragma: no cover
     import multiprocessing
     from multiprocessing import connection
 
+    from ...common.types import HostPort
+
 
 def delegate_work_to_pool(
         worker_pid: int,
         work_queue: 'connection.Connection',
         work_lock: 'multiprocessing.synchronize.Lock',
         conn: 'socket.socket',
-        addr: Optional[Tuple[str, int]],
+        addr: Optional['HostPort'],
         unix_socket_path: Optional[str] = None,
 ) -> None:
     """Utility method to delegate a work to threadless executor pool."""
diff --git a/proxy/core/work/fd.py b/proxy/core/work/fd.py
new file mode 100644
index 0000000..797d478
--- /dev/null
+++ b/proxy/core/work/fd.py
@@ -0,0 +1,64 @@
+# -*- 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 socket
+import logging
+from typing import Any, TypeVar, Optional
+
+from ..event import eventNames
+from .threadless import Threadless
+from ...common.types import HostPort, TcpOrTlsSocket
+
+
+T = TypeVar('T')
+
+logger = logging.getLogger(__name__)
+
+
+class ThreadlessFdExecutor(Threadless[T]):
+
+    def work(
+            self,
+            *args: Any,
+            **kwargs: Any,
+    ) -> None:
+        fileno: int = kwargs['fileno']
+        addr: Optional[HostPort] = kwargs.get('addr', None)
+        conn: Optional[TcpOrTlsSocket] = \
+            kwargs.get('conn', None)
+        conn = conn or socket.fromfd(
+            fileno, family=socket.AF_INET if self.flags.hostname.version == 4 else socket.AF_INET6,
+            type=socket.SOCK_STREAM,
+        )
+        uid = '%s-%s-%s' % (self.iid, self._total, fileno)
+        self.works[fileno] = self.flags.work_klass(
+            self.flags.work_klass.create(
+                conn=conn,
+                addr=addr,
+            ),
+            flags=self.flags,
+            event_queue=self.event_queue,
+            uid=uid,
+            upstream_conn_pool=self._upstream_conn_pool,
+        )
+        self.works[fileno].publish_event(
+            event_name=eventNames.WORK_STARTED,
+            event_payload={'fileno': fileno, 'addr': addr},
+            publisher_id=self.__class__.__name__,
+        )
+        try:
+            self.works[fileno].initialize()
+            self._total += 1
+        except Exception as e:
+            logger.exception(
+                'Exception occurred during initialization',
+                exc_info=e,
+            )
+            self._cleanup(fileno)
diff --git a/proxy/core/work/local.py b/proxy/core/work/local.py
index e081b2d..cdb26bd 100644
--- a/proxy/core/work/local.py
+++ b/proxy/core/work/local.py
@@ -14,7 +14,7 @@ import logging
 import contextlib
 from typing import Any, Optional
 
-from .threadless import Threadless
+from .fd import ThreadlessFdExecutor
 from ...common.backports import (  # noqa: W0611, F401   pylint: disable=unused-import
     NonBlockingQueue,
 )
@@ -23,7 +23,7 @@ from ...common.backports import (  # noqa: W0611, F401   pylint: disable=unused-
 logger = logging.getLogger(__name__)
 
 
-class LocalExecutor(Threadless['NonBlockingQueue']):
+class LocalExecutor(ThreadlessFdExecutor['NonBlockingQueue']):
     """A threadless executor implementation which uses a queue to receive new work."""
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
@@ -49,5 +49,5 @@ class LocalExecutor(Threadless['NonBlockingQueue']):
             # NOTE: Here we are assuming to receive a connection object
             # and not a fileno because we are a LocalExecutor.
             fileno = conn.fileno()
-            self.work_on_tcp_conn(fileno=fileno, addr=addr, conn=conn)
+            self.work(fileno=fileno, addr=addr, conn=conn)
         return False
diff --git a/proxy/core/work/remote.py b/proxy/core/work/remote.py
index 6dd91ef..6168bb3 100644
--- a/proxy/core/work/remote.py
+++ b/proxy/core/work/remote.py
@@ -14,13 +14,13 @@ from typing import Any, Optional
 from multiprocessing import connection
 from multiprocessing.reduction import recv_handle
 
-from .threadless import Threadless
+from .fd import ThreadlessFdExecutor
 
 
 logger = logging.getLogger(__name__)
 
 
-class RemoteExecutor(Threadless[connection.Connection]):
+class RemoteExecutor(ThreadlessFdExecutor[connection.Connection]):
     """A threadless executor implementation which receives work over a connection.
 
     NOTE: RemoteExecutor uses ``recv_handle`` to accept file descriptors.
@@ -53,5 +53,5 @@ class RemoteExecutor(Threadless[connection.Connection]):
         if not self.flags.unix_socket_path:
             addr = self.work_queue.recv()
         fileno = recv_handle(self.work_queue)
-        self.work_on_tcp_conn(fileno=fileno, addr=addr)
+        self.work(fileno=fileno, addr=addr)
         return False
diff --git a/proxy/core/work/threaded.py b/proxy/core/work/threaded.py
index 6f2569a..4437424 100644
--- a/proxy/core/work/threaded.py
+++ b/proxy/core/work/threaded.py
@@ -18,6 +18,7 @@ from ..event import EventQueue, eventNames
 
 if TYPE_CHECKING:   # pragma: no cover
     from .work import Work
+    from ...common.types import HostPort
 
 T = TypeVar('T')
 
@@ -26,7 +27,7 @@ T = TypeVar('T')
 def start_threaded_work(
         flags: argparse.Namespace,
         conn: socket.socket,
-        addr: Optional[Tuple[str, int]],
+        addr: Optional['HostPort'],
         event_queue: Optional[EventQueue] = None,
         publisher_id: Optional[str] = None,
 ) -> Tuple['Work[T]', threading.Thread]:
diff --git a/proxy/core/work/threadless.py b/proxy/core/work/threadless.py
index 57eca4d..ac76c7e 100644
--- a/proxy/core/work/threadless.py
+++ b/proxy/core/work/threadless.py
@@ -9,8 +9,6 @@
     :license: BSD, see LICENSE for more details.
 """
 import os
-import ssl
-import socket
 import asyncio
 import logging
 import argparse
@@ -18,10 +16,9 @@ import selectors
 import multiprocessing
 from abc import ABC, abstractmethod
 from typing import (
-    TYPE_CHECKING, Set, Dict, List, Tuple, Union, Generic, TypeVar, Optional,
+    TYPE_CHECKING, Any, Set, Dict, List, Tuple, Generic, TypeVar, Optional,
 )
 
-from ..event import eventNames
 from ...common.types import Readables, Writables, SelectableEvents
 from ...common.logger import Logger
 from ...common.constants import (
@@ -31,8 +28,6 @@ from ...common.constants import (
 
 
 if TYPE_CHECKING:   # pragma: no cover
-    from typing import Any
-
     from .work import Work
     from ..event import EventQueue
 
@@ -124,48 +119,16 @@ class Threadless(ABC, Generic[T]):
         return work queue fd."""
         raise NotImplementedError()
 
+    @abstractmethod
+    def work(self, *args: Any, **kwargs: Any) -> None:
+        raise NotImplementedError()
+
     def close_work_queue(self) -> None:
         """Only called if ``work_queue_fileno`` returns an integer.
         If an fd is select-able for work queue, make sure
         to close the work queue fd now."""
         pass    # pragma: no cover
 
-    def work_on_tcp_conn(
-            self,
-            fileno: int,
-            addr: Optional[Tuple[str, int]] = None,
-            conn: Optional[Union[ssl.SSLSocket, socket.socket]] = None,
-    ) -> None:
-        conn = conn or socket.fromfd(
-            fileno, family=socket.AF_INET if self.flags.hostname.version == 4 else socket.AF_INET6,
-            type=socket.SOCK_STREAM,
-        )
-        uid = '%s-%s-%s' % (self.iid, self._total, fileno)
-        self.works[fileno] = self.flags.work_klass(
-            self.flags.work_klass.create(
-                conn=conn,
-                addr=addr,
-            ),
-            flags=self.flags,
-            event_queue=self.event_queue,
-            uid=uid,
-            upstream_conn_pool=self._upstream_conn_pool,
-        )
-        self.works[fileno].publish_event(
-            event_name=eventNames.WORK_STARTED,
-            event_payload={'fileno': fileno, 'addr': addr},
-            publisher_id=self.__class__.__name__,
-        )
-        try:
-            self.works[fileno].initialize()
-            self._total += 1
-        except Exception as e:
-            logger.exception(
-                'Exception occurred during initialization',
-                exc_info=e,
-            )
-            self._cleanup(fileno)
-
     async def _update_work_events(self, work_id: int) -> None:
         assert self.selector is not None
         worker_events = await self.works[work_id].get_events()
diff --git a/proxy/http/proxy/plugin.py b/proxy/http/proxy/plugin.py
index 0fe1028..a71e553 100644
--- a/proxy/http/proxy/plugin.py
+++ b/proxy/http/proxy/plugin.py
@@ -20,6 +20,7 @@ from ...common.utils import tls_interception_enabled
 
 
 if TYPE_CHECKING:   # pragma: no cover
+    from ...common.types import HostPort
     from ...core.connection import UpstreamConnectionPool
 
 
@@ -52,7 +53,7 @@ class HttpProxyBasePlugin(
         access a specific plugin by its name."""
         return self.__class__.__name__      # pragma: no cover
 
-    def resolve_dns(self, host: str, port: int) -> Tuple[Optional[str], Optional[Tuple[str, int]]]:
+    def resolve_dns(self, host: str, port: int) -> Tuple[Optional[str], Optional['HostPort']]:
         """Resolve upstream server host to an IP address.
 
         Optionally also override the source address to use for
diff --git a/proxy/http/websocket/client.py b/proxy/http/websocket/client.py
index f398b5d..c8c53eb 100644
--- a/proxy/http/websocket/client.py
+++ b/proxy/http/websocket/client.py
@@ -8,15 +8,15 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
-import ssl
 import base64
 import socket
 import secrets
 import selectors
-from typing import Union, Callable, Optional
+from typing import Callable, Optional
 
 from .frame import WebsocketFrame
 from ..parser import HttpParser, httpParserTypes
+from ...common.types import TcpOrTlsSocket
 from ...common.utils import (
     text_, new_socket_connection, build_websocket_handshake_request,
 )
@@ -53,7 +53,7 @@ class WebsocketClient(TcpConnection):
         self.selector: selectors.DefaultSelector = selectors.DefaultSelector()
 
     @property
-    def connection(self) -> Union[ssl.SSLSocket, socket.socket]:
+    def connection(self) -> TcpOrTlsSocket:
         return self.sock
 
     def handshake(self) -> None:
diff --git a/proxy/plugin/cloudflare_dns.py b/proxy/plugin/cloudflare_dns.py
index 9c459b5..ee47900 100644
--- a/proxy/plugin/cloudflare_dns.py
+++ b/proxy/plugin/cloudflare_dns.py
@@ -26,6 +26,7 @@ from typing import Tuple, Optional
 
 from ..http.proxy import HttpProxyBasePlugin
 from ..common.flag import flags
+from ..common.types import HostPort
 
 
 logger = logging.getLogger(__name__)
@@ -61,7 +62,7 @@ class CloudflareDnsResolverPlugin(HttpProxyBasePlugin):
            pip install "httpx[http2]"
     """
 
-    def resolve_dns(self, host: str, port: int) -> Tuple[Optional[str], Optional[Tuple[str, int]]]:
+    def resolve_dns(self, host: str, port: int) -> Tuple[Optional[str], Optional[HostPort]]:
         try:
             context = httpx.create_ssl_context(http2=True)
             # TODO: Support resolution via Authority (SOA) to add support for
diff --git a/proxy/plugin/custom_dns_resolver.py b/proxy/plugin/custom_dns_resolver.py
index 5ee7ade..8a1ee7c 100644
--- a/proxy/plugin/custom_dns_resolver.py
+++ b/proxy/plugin/custom_dns_resolver.py
@@ -16,12 +16,13 @@ import socket
 from typing import Tuple, Optional
 
 from ..http.proxy import HttpProxyBasePlugin
+from ..common.types import HostPort
 
 
 class CustomDnsResolverPlugin(HttpProxyBasePlugin):
     """This plugin demonstrate how to use your own custom DNS resolver."""
 
-    def resolve_dns(self, host: str, port: int) -> Tuple[Optional[str], Optional[Tuple[str, int]]]:
+    def resolve_dns(self, host: str, port: int) -> Tuple[Optional[str], Optional[HostPort]]:
         """Here we are using in-built python resolver for demonstration.
 
         Ideally you would like to query your custom DNS server or even


Test patch:
diff --git a/tests/core/test_connection.py b/tests/core/test_connection.py
index b3c00ad..275a0d9 100644
--- a/tests/core/test_connection.py
+++ b/tests/core/test_connection.py
@@ -8,13 +8,12 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
-import ssl
-import socket
-from typing import Union, Optional
+from typing import Optional
 
 import unittest
 from unittest import mock
 
+from proxy.common.types import TcpOrTlsSocket
 from proxy.core.connection import (
     TcpConnection, TcpClientConnection, TcpServerConnection,
     TcpConnectionUninitializedException, tcpConnectionTypes,
@@ -28,14 +27,14 @@ class TestTcpConnection(unittest.TestCase):
     class TcpConnectionToTest(TcpConnection):
 
         def __init__(
-            self, conn: Optional[Union[ssl.SSLSocket, socket.socket]] = None,
+            self, conn: Optional[TcpOrTlsSocket] = None,
             tag: int = tcpConnectionTypes.CLIENT,
         ) -> None:
             super().__init__(tag)
             self._conn = conn
 
         @property
-        def connection(self) -> Union[ssl.SSLSocket, socket.socket]:
+        def connection(self) -> TcpOrTlsSocket:
             if self._conn is None:
                 raise TcpConnectionUninitializedException()
             return self._conn