Optimize how HttpProtocolHandler delegates to the core plugins (#925)

* Add `protocols` abstract static method to `HttpProtocolHandlerBase` which defines which HTTP specification is followed by the core plugin

* lint

* Fix tests

* Lint fixes

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

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

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

Author: abhinavsingh

README.md

@@ -2232,7 +2232,7 @@ usage: -m [-h] [--enable-events] [--enable-conn-pool] [--threadless]
           [--filtered-url-regex-config FILTERED_URL_REGEX_CONFIG]
           [--cloudflare-dns-mode CLOUDFLARE_DNS_MODE]
 
-proxy.py v2.4.0rc5.dev11+ga872675.d20211225
+proxy.py v2.4.0rc5.dev26+gb2b1bdc.d20211230
 
 options:
   -h, --help            show this help message and exit

docs/conf.py

@@ -219,7 +219,11 @@
 # -- Options for linkcheck builder -------------------------------------------
 
 linkcheck_ignore = [
-    r'http://localhost:\d+/',  # local URLs
+    # local URLs
+    r'http://localhost:\d+/',
+    # GHA sees "403 Client Error: Forbidden for url:"
+    # while the URL actually works
+    r'https://developers.cloudflare.com/',
 ]
 linkcheck_workers = 25
 

proxy/http/handler.py

@@ -16,20 +16,21 @@
 import logging
 import selectors
 
-from typing import Tuple, List, Union, Optional, Dict, Any
+from typing import Tuple, List, Type, Union, Optional, Dict, Any
 
-from .plugin import HttpProtocolHandlerPlugin
-from .parser import HttpParser, httpParserStates, httpParserTypes
-from .exception import HttpProtocolException
-
-from ..common.types import Readables, Writables
+from ..common.flag import flags
 from ..common.utils import wrap_socket
 from ..core.base import BaseTcpServerHandler
+from ..common.types import Readables, Writables
 from ..core.connection import TcpClientConnection
-from ..common.flag import flags
 from ..common.constants import DEFAULT_CLIENT_RECVBUF_SIZE, DEFAULT_KEY_FILE
 from ..common.constants import DEFAULT_SELECTOR_SELECT_TIMEOUT, DEFAULT_TIMEOUT
 
+from .exception import HttpProtocolException
+from .plugin import HttpProtocolHandlerPlugin
+from .responses import BAD_REQUEST_RESPONSE_PKT
+from .parser import HttpParser, httpParserStates, httpParserTypes
+
 
 logger = logging.getLogger(__name__)
 
@@ -78,31 +79,24 @@ def __init__(self, *args: Any, **kwargs: Any):
         self.selector: Optional[selectors.DefaultSelector] = None
         if not self.flags.threadless:
             self.selector = selectors.DefaultSelector()
-        self.plugins: Dict[str, HttpProtocolHandlerPlugin] = {}
+        self.plugin: Optional[HttpProtocolHandlerPlugin] = None
 
     ##
     # initialize, is_inactive, shutdown, get_events, handle_events
     # overrides Work class definitions.
     ##
 
     def initialize(self) -> None:
-        """Optionally upgrades connection to HTTPS, set ``conn`` in non-blocking mode and initializes plugins."""
+        """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)
         # Update client connection reference if connection was wrapped
         if self._encryption_enabled():
             self.work = TcpClientConnection(conn=conn, addr=self.work.addr)
-        if b'HttpProtocolHandlerPlugin' in self.flags.plugins:
-            for klass in self.flags.plugins[b'HttpProtocolHandlerPlugin']:
-                instance: HttpProtocolHandlerPlugin = klass(
-                    self.uid,
-                    self.flags,
-                    self.work,
-                    self.request,
-                    self.event_queue,
-                    self.upstream_conn_pool,
-                )
-                self.plugins[instance.name()] = instance
+        # self._initialize_plugins()
         logger.debug('Handling connection %s' % self.work.address)
 
     def is_inactive(self) -> bool:
@@ -120,8 +114,8 @@ def shutdown(self) -> None:
             if self.selector and self.work.has_buffer():
                 self._flush()
             # Invoke plugin.on_client_connection_close
-            for plugin in self.plugins.values():
-                plugin.on_client_connection_close()
+            if self.plugin:
+                self.plugin.on_client_connection_close()
             logger.debug(
                 'Closing client connection %s has buffer %s' %
                 (self.work.address, self.work.has_buffer()),
@@ -153,8 +147,8 @@ async def get_events(self) -> Dict[int, int]:
         # Get default client events
         events: Dict[int, int] = await super().get_events()
         # HttpProtocolHandlerPlugin.get_descriptors
-        for plugin in self.plugins.values():
-            plugin_read_desc, plugin_write_desc = plugin.get_descriptors()
+        if self.plugin:
+            plugin_read_desc, plugin_write_desc = self.plugin.get_descriptors()
             for rfileno in plugin_read_desc:
                 if rfileno not in events:
                     events[rfileno] = selectors.EVENT_READ
@@ -179,17 +173,17 @@ async def handle_events(
         if teardown:
             return True
         # Invoke plugin.write_to_descriptors
-        for plugin in self.plugins.values():
-            teardown = await plugin.write_to_descriptors(writables)
+        if self.plugin:
+            teardown = await self.plugin.write_to_descriptors(writables)
             if teardown:
                 return True
         # Read from ready to read sockets
         teardown = await self.handle_readables(readables)
         if teardown:
             return True
         # Invoke plugin.read_from_descriptors
-        for plugin in self.plugins.values():
-            teardown = await plugin.read_from_descriptors(readables)
+        if self.plugin:
+            teardown = await self.plugin.read_from_descriptors(readables)
             if teardown:
                 return True
         return False
@@ -209,33 +203,13 @@ def handle_data(self, data: memoryview) -> Optional[bool]:
             # apply custom logic to handle request data sent after 1st
             # valid request.
             if self.request.state != httpParserStates.COMPLETE:
-                # Parse http request
-                #
-                # TODO(abhinavsingh): Remove .tobytes after parser is
-                # memoryview compliant
-                self.request.parse(data.tobytes())
-                if self.request.is_complete:
-                    # Invoke plugin.on_request_complete
-                    for plugin in self.plugins.values():
-                        upgraded_sock = plugin.on_request_complete()
-                        if isinstance(upgraded_sock, ssl.SSLSocket):
-                            logger.debug(
-                                'Updated client conn to %s', upgraded_sock,
-                            )
-                            self.work._conn = upgraded_sock
-                            for plugin_ in self.plugins.values():
-                                if plugin_ != plugin:
-                                    plugin_.client._conn = upgraded_sock
-                        elif isinstance(upgraded_sock, bool) and upgraded_sock is True:
-                            return True
+                if self._parse_first_request(data):
+                    return True
             else:
                 # HttpProtocolHandlerPlugin.on_client_data
                 # Can raise HttpProtocolException to tear down the connection
-                for plugin in self.plugins.values():
-                    optional_data = plugin.on_client_data(data)
-                    if optional_data is None:
-                        break
-                    data = optional_data
+                if self.plugin:
+                    data = self.plugin.on_client_data(data) or data
         except HttpProtocolException as e:
             logger.info('HttpProtocolException: %s' % e)
             response: Optional[memoryview] = e.response(self.request)
@@ -248,17 +222,13 @@ async def handle_writables(self, writables: Writables) -> bool:
         if self.work.connection.fileno() in writables and self.work.has_buffer():
             logger.debug('Client is write ready, flushing...')
             self.last_activity = time.time()
-
             # TODO(abhinavsingh): This hook could just reside within server recv block
             # instead of invoking when flushed to client.
             #
             # Invoke plugin.on_response_chunk
             chunk = self.work.buffer
-            for plugin in self.plugins.values():
-                chunk = plugin.on_response_chunk(chunk)
-                if chunk is None:
-                    break
-
+            if self.plugin:
+                chunk = self.plugin.on_response_chunk(chunk)
             try:
                 # Call super() for client flush
                 teardown = await super().handle_writables(writables)
@@ -305,6 +275,61 @@ async def handle_readables(self, readables: Readables) -> bool:
     # Internal methods
     ##
 
+    def _initialize_plugin(
+            self,
+            klass: Type['HttpProtocolHandlerPlugin'],
+    ) -> HttpProtocolHandlerPlugin:
+        """Initializes passed HTTP protocol handler plugin class."""
+        return klass(
+            self.uid,
+            self.flags,
+            self.work,
+            self.request,
+            self.event_queue,
+            self.upstream_conn_pool,
+        )
+
+    def _discover_plugin_klass(self, protocol: int) -> Optional[Type['HttpProtocolHandlerPlugin']]:
+        """Discovers and return matching HTTP handler plugin matching protocol."""
+        if b'HttpProtocolHandlerPlugin' in self.flags.plugins:
+            for klass in self.flags.plugins[b'HttpProtocolHandlerPlugin']:
+                k: Type['HttpProtocolHandlerPlugin'] = klass
+                if protocol in k.protocols():
+                    return k
+        return None
+
+    def _parse_first_request(self, data: memoryview) -> bool:
+        # Parse http request
+        #
+        # TODO(abhinavsingh): Remove .tobytes after parser is
+        # memoryview compliant
+        self.request.parse(data.tobytes())
+        if not self.request.is_complete:
+            return False
+        # Discover which HTTP handler plugin is capable of
+        # handling the current incoming request
+        klass = self._discover_plugin_klass(
+            self.request.http_handler_protocol,
+        )
+        if klass is None:
+            # No matching protocol class found.
+            # Return bad request response and
+            # close the connection.
+            self.work.queue(BAD_REQUEST_RESPONSE_PKT)
+            return True
+        assert klass is not None
+        self.plugin = self._initialize_plugin(klass)
+        # Invoke plugin.on_request_complete
+        output = self.plugin.on_request_complete()
+        if isinstance(output, bool):
+            return output
+        assert isinstance(output, ssl.SSLSocket)
+        logger.debug(
+            'Updated client conn to %s', output,
+        )
+        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

proxy/http/parser/parser.py

@@ -22,6 +22,7 @@
 
 from ..url import Url
 from ..methods import httpMethods
+from ..protocols import httpProtocols
 from ..exception import HttpProtocolException
 
 from .protocol import ProxyProtocol
@@ -153,11 +154,10 @@ def set_url(self, url: bytes) -> None:
         self._url = Url.from_bytes(url)
         self._set_line_attributes()
 
-    def has_host(self) -> bool:
-        """Returns whether host line attribute was parsed or set.
-
-        NOTE: Host field WILL be None for incoming local WebServer requests."""
-        return self.host is not None
+    @property
+    def http_handler_protocol(self) -> int:
+        """Returns `HttpProtocols` that this request belongs to."""
+        return httpProtocols.HTTP_PROXY if self.host is not None else httpProtocols.WEB_SERVER
 
     @property
     def is_complete(self) -> bool:

proxy/http/plugin.py

@@ -14,12 +14,12 @@
 from abc import ABC, abstractmethod
 from typing import Tuple, List, Union, Optional, TYPE_CHECKING
 
-from .parser import HttpParser
-
 from ..common.types import Readables, Writables
 from ..core.event import EventQueue
 from ..core.connection import TcpClientConnection
 
+from .parser import HttpParser
+
 if TYPE_CHECKING:
     from ..core.connection import UpstreamConnectionPool
 
@@ -52,7 +52,7 @@ def __init__(
             flags: argparse.Namespace,
             client: TcpClientConnection,
             request: HttpParser,
-            event_queue: EventQueue,
+            event_queue: Optional[EventQueue],
             upstream_conn_pool: Optional['UpstreamConnectionPool'] = None,
     ):
         self.uid: str = uid
@@ -63,12 +63,10 @@ def __init__(
         self.upstream_conn_pool = upstream_conn_pool
         super().__init__()
 
-    def name(self) -> str:
-        """A unique name for your plugin.
-
-        Defaults to name of the class. This helps plugin developers to directly
-        access a specific plugin by its name."""
-        return self.__class__.__name__
+    @staticmethod
+    @abstractmethod
+    def protocols() -> List[int]:
+        raise NotImplementedError()
 
     @abstractmethod
     def get_descriptors(self) -> Tuple[List[int], List[int]]:

proxy/http/protocols.py

@@ -0,0 +1,33 @@
+# -*- 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 typing import NamedTuple
+
+
+HttpProtocols = NamedTuple(
+    'HttpProtocols', [
+        # Web server handling HTTP/1.0, HTTP/1.1, HTTP/2, HTTP/3
+        # over plain Text or encrypted connection with clients
+        ('WEB_SERVER', int),
+        # Proxies handling HTTP/1.0, HTTP/1.1, HTTP/2 protocols
+        # over plain text connection or encrypted connection
+        # with clients
+        ('HTTP_PROXY', int),
+        # Proxies handling SOCKS4, SOCKS4a, SOCKS5 protocol
+        ('SOCKS_PROXY', int),
+    ],
+)
+
+httpProtocols = HttpProtocols(1, 2, 3)