[Work] Refactor for generic work type (#1048)

* Add `TcpOrTlsSocket` type

* isort

* Update to use Fd executor

* Define `HostPort` type

* Fix fileno

* spellfix

Author: abhinavsingh

docs/conf.py

@@ -317,7 +317,10 @@
     (_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'),
 ]

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 @@
 Readables = Selectables
 Writables = Selectables
 Descriptors = Tuple[Readables, Writables]
-
 IpAddress = Union[ipaddress.IPv4Address, ipaddress.IPv6Address]
+TcpOrTlsSocket = Union[ssl.SSLSocket, socket.socket]
+HostPort = Tuple[str, int]

proxy/common/utils.py

@@ -23,6 +23,7 @@
 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__()
 

proxy/core/acceptor/acceptor.py

@@ -26,6 +26,7 @@
 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 @@ def __init__(
     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 @@ def run(self) -> None:
             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 @@ def _stop_local(self) -> None:
             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

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 @@ def _encryption_enabled(self) -> bool:
         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.

proxy/core/connection/client.py

@@ -9,32 +9,32 @@
     :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):
     """A buffered client connection object."""
 
     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

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 @@ def __init__(self, tag: int) -> None:
 
     @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
 

proxy/core/connection/pool.py

@@ -20,7 +20,7 @@
 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 @@ async def handle_events(self, readables: Readables, _writables: Writables) -> bo
             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.

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(

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 @@ def __init__(self, flags: argparse.Namespace) -> None:
     def on_connection(
             self,
             chan: 'Channel',
-            origin: Tuple[str, int],
-            server: Tuple[str, int],
+            origin: 'HostPort',
+            server: 'HostPort',
     ) -> None:
         pass