[Middleware] Capability in the core to allow custom client connection classes (#993)

* Move all TCP server related flags within `tcp_server.py` and also move the encryption functionality within TCP base server

* Templatize `BaseTcpServerHandler` which now expects a client connection object bound to `TcpClientConnection`.  This will allow for custom `HttpClientConnection` object in future to be used by `HttpProtocolHandler`

* Pass necessary flags to allow self-signed certificates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix https integration tests

* Affected by #994

* Fix docs

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: abhinavsingh

Makefile

@@ -57,7 +57,7 @@ sign-https-certificates:
 	python -m proxy.common.pki sign_csr \
 		--csr-path $(HTTPS_CSR_FILE_PATH) \
 		--crt-path $(HTTPS_SIGNED_CERT_FILE_PATH) \
-		--hostname example.com \
+		--hostname localhost \
 		--private-key-path $(CA_KEY_FILE_PATH) \
 		--public-key-path $(CA_CERT_FILE_PATH)
 

docs/conf.py

@@ -316,4 +316,5 @@
     (_py_class_role, 'EventQueue'),
     (_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

@@ -17,7 +17,7 @@
 from proxy.core.connection import TcpClientConnection
 
 
-class EchoSSLServerHandler(BaseTcpServerHandler):
+class EchoSSLServerHandler(BaseTcpServerHandler[TcpClientConnection]):
     """Wraps client socket during initialization."""
 
     def initialize(self) -> None:

examples/tcp_echo_server.py

@@ -13,9 +13,10 @@
 
 from proxy import Proxy
 from proxy.core.base import BaseTcpServerHandler
+from proxy.core.connection import TcpClientConnection
 
 
-class EchoServerHandler(BaseTcpServerHandler):
+class EchoServerHandler(BaseTcpServerHandler[TcpClientConnection]):
     """Sets client socket to non-blocking during initialization."""
 
     def initialize(self) -> None:

proxy/common/pki.py

@@ -268,8 +268,8 @@ def run_openssl_command(command: List[str], timeout: int) -> bool:
     parser.add_argument(
         '--subject',
         type=str,
-        default='/CN=example.com',
-        help='Subject to use for public key generation. Default: /CN=example.com',
+        default='/CN=localhost',
+        help='Subject to use for public key generation. Default: /CN=localhost',
     )
     parser.add_argument(
         '--csr-path',

proxy/core/base/tcp_server.py

@@ -12,20 +12,75 @@
 
        tcp
 """
+import ssl
+import socket
 import logging
 import selectors
 
 from abc import abstractmethod
-from typing import Any, Optional
+from typing import Any, Optional, TypeVar, Union
+
+from ...common.flag import flags
+from ...common.utils import wrap_socket
+from ...common.types import Readables, SelectableEvents, Writables
+from ...common.constants import DEFAULT_CERT_FILE, DEFAULT_CLIENT_RECVBUF_SIZE
+from ...common.constants import DEFAULT_KEY_FILE, DEFAULT_SERVER_RECVBUF_SIZE, DEFAULT_TIMEOUT
 
 from ...core.work import Work
 from ...core.connection import TcpClientConnection
-from ...common.types import Readables, SelectableEvents, Writables
 
 logger = logging.getLogger(__name__)
 
 
-class BaseTcpServerHandler(Work[TcpClientConnection]):
+flags.add_argument(
+    '--key-file',
+    type=str,
+    default=DEFAULT_KEY_FILE,
+    help='Default: None. Server key file to enable end-to-end TLS encryption with clients. '
+    'If used, must also pass --cert-file.',
+)
+
+flags.add_argument(
+    '--cert-file',
+    type=str,
+    default=DEFAULT_CERT_FILE,
+    help='Default: None. Server certificate to enable end-to-end TLS encryption with clients. '
+    'If used, must also pass --key-file.',
+)
+
+flags.add_argument(
+    '--client-recvbuf-size',
+    type=int,
+    default=DEFAULT_CLIENT_RECVBUF_SIZE,
+    help='Default: ' + str(int(DEFAULT_CLIENT_RECVBUF_SIZE / 1024)) +
+    ' KB. Maximum amount of data received from the '
+    'client in a single recv() operation.',
+)
+
+flags.add_argument(
+    '--server-recvbuf-size',
+    type=int,
+    default=DEFAULT_SERVER_RECVBUF_SIZE,
+    help='Default: ' + str(int(DEFAULT_SERVER_RECVBUF_SIZE / 1024)) +
+    ' KB. Maximum amount of data received from the '
+    'server in a single recv() operation.',
+)
+
+flags.add_argument(
+    '--timeout',
+    type=int,
+    default=DEFAULT_TIMEOUT,
+    help='Default: ' + str(DEFAULT_TIMEOUT) +
+    '.  Number of seconds after which '
+    'an inactive connection must be dropped.  Inactivity is defined by no '
+    'data sent or received by the client.',
+)
+
+
+T = TypeVar('T', bound=TcpClientConnection)
+
+
+class BaseTcpServerHandler(Work[T]):
     """BaseTcpServerHandler implements Work interface.
 
     BaseTcpServerHandler lifecycle is controlled by Threadless core
@@ -56,6 +111,14 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             self.work.address,
         )
 
+    def initialize(self) -> None:
+        """Optionally upgrades connection to HTTPS,
+        sets ``conn`` in non-blocking mode and initializes
+        HTTP protocol plugins."""
+        conn = self._optionally_wrap_socket(self.work.connection)
+        conn.setblocking(False)
+        logger.debug('Handling connection %s' % self.work.address)
+
     @abstractmethod
     def handle_data(self, data: memoryview) -> Optional[bool]:
         """Optionally return True to close client connection."""
@@ -139,3 +202,21 @@ async def handle_readables(self, readables: Readables) -> bool:
                     else:
                         teardown = True
         return teardown
+
+    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]:
+        """Attempts to wrap accepted client connection using provided certificates.
+
+        Shutdown and closes client connection upon error.
+        """
+        if self._encryption_enabled():
+            assert self.flags.keyfile and self.flags.certfile
+            # TODO(abhinavsingh): Insecure TLS versions must not be accepted by default
+            conn = wrap_socket(conn, self.flags.keyfile, self.flags.certfile)
+            self.work._conn = conn
+        return conn

proxy/core/base/tcp_tunnel.py

@@ -18,13 +18,13 @@
 from ...common.types import Readables, SelectableEvents, Writables
 from ...common.utils import text_
 
-from ..connection import TcpServerConnection
+from ..connection import TcpServerConnection, TcpClientConnection
 from .tcp_server import BaseTcpServerHandler
 
 logger = logging.getLogger(__name__)
 
 
-class BaseTcpTunnelHandler(BaseTcpServerHandler):
+class BaseTcpTunnelHandler(BaseTcpServerHandler[TcpClientConnection]):
     """BaseTcpTunnelHandler build on-top of BaseTcpServerHandler work class.
 
     On-top of BaseTcpServerHandler implementation,

proxy/http/handler.py

@@ -16,15 +16,12 @@
 import logging
 import selectors
 
-from typing import Tuple, List, Type, Union, Optional, Any
+from typing import Tuple, List, Type, Optional, Any
 
-from ..common.flag import flags
-from ..common.utils import wrap_socket
 from ..core.base import BaseTcpServerHandler
 from ..core.connection import TcpClientConnection
 from ..common.types import Readables, SelectableEvents, Writables
-from ..common.constants import DEFAULT_CLIENT_RECVBUF_SIZE, DEFAULT_KEY_FILE
-from ..common.constants import DEFAULT_SELECTOR_SELECT_TIMEOUT, DEFAULT_TIMEOUT
+from ..common.constants import DEFAULT_SELECTOR_SELECT_TIMEOUT
 
 from .exception import HttpProtocolException
 from .plugin import HttpProtocolHandlerPlugin
@@ -35,33 +32,7 @@
 logger = logging.getLogger(__name__)
 
 
-flags.add_argument(
-    '--client-recvbuf-size',
-    type=int,
-    default=DEFAULT_CLIENT_RECVBUF_SIZE,
-    help='Default: ' + str(int(DEFAULT_CLIENT_RECVBUF_SIZE / 1024)) +
-    ' KB. Maximum amount of data received from the '
-    'client in a single recv() operation.',
-)
-flags.add_argument(
-    '--key-file',
-    type=str,
-    default=DEFAULT_KEY_FILE,
-    help='Default: None. Server key file to enable end-to-end TLS encryption with clients. '
-    'If used, must also pass --cert-file.',
-)
-flags.add_argument(
-    '--timeout',
-    type=int,
-    default=DEFAULT_TIMEOUT,
-    help='Default: ' + str(DEFAULT_TIMEOUT) +
-    '.  Number of seconds after which '
-    'an inactive connection must be dropped.  Inactivity is defined by no '
-    'data sent or received by the client.',
-)
-
-
-class HttpProtocolHandler(BaseTcpServerHandler):
+class HttpProtocolHandler(BaseTcpServerHandler[TcpClientConnection]):
     """HTTP, HTTPS, HTTP2, WebSockets protocol handler.
 
     Accepts `Client` connection and delegates to HttpProtocolHandlerPlugin.
@@ -86,17 +57,16 @@ def __init__(self, *args: Any, **kwargs: Any):
     ##
 
     def initialize(self) -> None:
-        """Optionally upgrades connection to HTTPS,
-        sets ``conn`` in non-blocking mode and initializes
-        HTTP protocol plugins.
-        """
-        conn = self._optionally_wrap_socket(self.work.connection)
-        conn.setblocking(False)
+        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(conn=conn, addr=self.work.addr)
-        # self._initialize_plugins()
-        logger.debug('Handling connection %s' % self.work.address)
+            self.work = TcpClientConnection(
+                conn=self.work.connection,
+                addr=self.work.addr,
+            )
 
     def is_inactive(self) -> bool:
         if not self.work.has_buffer() and \
@@ -334,23 +304,6 @@ def _parse_first_request(self, data: memoryview) -> bool:
         self.work._conn = output
         return False
 
-    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]:
-        """Attempts to wrap accepted client connection using provided certificates.
-
-        Shutdown and closes client connection upon error.
-        """
-        if self._encryption_enabled():
-            assert self.flags.keyfile and self.flags.certfile
-            # TODO(abhinavsingh): Insecure TLS versions must not be accepted by default
-            conn = wrap_socket(conn, self.flags.keyfile, self.flags.certfile)
-        return conn
-
     def _connection_inactive_for(self) -> float:
         return time.time() - self.last_activity
 

proxy/http/proxy/server.py

@@ -37,9 +37,9 @@
 from ...common.types import Readables, Writables, Descriptors
 from ...common.constants import DEFAULT_CA_CERT_DIR, DEFAULT_CA_CERT_FILE, DEFAULT_CA_FILE
 from ...common.constants import DEFAULT_CA_KEY_FILE, DEFAULT_CA_SIGNING_KEY_FILE
-from ...common.constants import COMMA, DEFAULT_SERVER_RECVBUF_SIZE, DEFAULT_CERT_FILE
+from ...common.constants import COMMA, DEFAULT_HTTPS_PROXY_ACCESS_LOG_FORMAT
 from ...common.constants import PROXY_AGENT_HEADER_VALUE, DEFAULT_DISABLE_HEADERS
-from ...common.constants import DEFAULT_HTTP_PROXY_ACCESS_LOG_FORMAT, DEFAULT_HTTPS_PROXY_ACCESS_LOG_FORMAT
+from ...common.constants import DEFAULT_HTTP_PROXY_ACCESS_LOG_FORMAT
 from ...common.constants import DEFAULT_DISABLE_HTTP_PROXY, PLUGIN_PROXY_AUTH
 from ...common.utils import text_
 from ...common.pki import gen_public_key, gen_csr, sign_csr
@@ -52,15 +52,6 @@
 logger = logging.getLogger(__name__)
 
 
-flags.add_argument(
-    '--server-recvbuf-size',
-    type=int,
-    default=DEFAULT_SERVER_RECVBUF_SIZE,
-    help='Default: ' + str(int(DEFAULT_SERVER_RECVBUF_SIZE / 1024)) +
-    ' KB. Maximum amount of data received from the '
-    'server in a single recv() operation.',
-)
-
 flags.add_argument(
     '--disable-http-proxy',
     action='store_true',
@@ -116,14 +107,6 @@
     'HTTPS certificates.  If used, must also pass --ca-key-file and --ca-cert-file',
 )
 
-flags.add_argument(
-    '--cert-file',
-    type=str,
-    default=DEFAULT_CERT_FILE,
-    help='Default: None. Server certificate to enable end-to-end TLS encryption with clients. '
-    'If used, must also pass --key-file.',
-)
-
 flags.add_argument(
     '--auth-plugin',
     type=str,

proxy/plugin/web_server_route.py

@@ -18,7 +18,7 @@
 logger = logging.getLogger(__name__)
 
 HTTP_RESPONSE = okResponse(content=b'HTTP route response')
-HTTPS_RESPONSE = okResponse(content=b'HTTP route response')
+HTTPS_RESPONSE = okResponse(content=b'HTTPS route response')
 
 
 class WebServerPlugin(HttpWebServerBasePlugin):