Repository: abhinavsingh/proxy.py
Pull request #1033 (2022-01-21 15:26:57)

[ReverseProxy] Move within core lib with ability to write its plugin
Added a way to create `ReverseProxy` plugins instead :D


Patch:
diff --git a/README.md b/README.md
index b85d1d7..ad7de4e 100644
--- a/README.md
+++ b/README.md
@@ -63,8 +63,9 @@
     - [Custom Network Interface](#customnetworkinterface)
     - [Program Name Plugin](#programnameplugin)
   - [HTTP Web Server Plugins](#http-web-server-plugins)
-    - [Reverse Proxy](#reverse-proxy)
     - [Web Server Route](#web-server-route)
+  - [Reverse Proxy Plugins](#reverse-proxy-plugins)
+    - [Reverse Proxy](#reverse-proxy)
   - [Plugin Ordering](#plugin-ordering)
 - [End-to-End Encryption](#end-to-end-encryption)
 - [TLS Interception](#tls-interception)
@@ -202,8 +203,10 @@
 - Programmable
   - Customize proxy behavior using [Proxy Server Plugins](#http-proxy-plugins). Example:
     - `--plugins proxy.plugin.ProxyPoolPlugin`
-  - Optionally, enable builtin [Web Server](#http-web-server-plugins). Example:
-    - `--enable-web-server --plugins proxy.plugin.ReverseProxyPlugin`
+  - Enable builtin [Web Server](#http-web-server-plugins). Example:
+    - `--enable-web-server --plugins proxy.plugin.WebServerPlugin`
+  - Enable builtin [Reverse Proxy Server](#reverse-proxy-plugins). Example:
+    - `--enable-reverse-proxy --plugins proxy.plugin.ReverseProxyPlugin`
   - Plugin API is currently in *development phase*. Expect breaking changes. See [Deploying proxy.py in production](#deploying-proxypy-in-production) on how to ensure reliability across code changes.
 - Real-time Dashboard
   - Optionally, enable [proxy.py dashboard](#run-dashboard).
@@ -940,14 +943,33 @@ Notice `curl` in-place of `::1` or `127.0.0.1` as client IP.
 
 ## HTTP Web Server Plugins
 
-### Reverse Proxy
+### Web Server Route
 
-Extend in-built Web Server to add Reverse Proxy capabilities.
+Demonstrates inbuilt web server routing using plugin.
 
 Start `proxy.py` as:
 
 ```console
 ❯ proxy --enable-web-server \
+    --plugins proxy.plugin.WebServerPlugin
+```
+
+Verify using `curl -v localhost:8899/http-route-example`, should return:
+
+```console
+HTTP route response
+```
+
+## Reverse Proxy Plugins
+
+Extends in-built Web Server to add Reverse Proxy capabilities.
+
+### Reverse Proxy
+
+Start `proxy.py` as:
+
+```console
+❯ proxy --enable-reverse-proxy \
     --plugins proxy.plugin.ReverseProxyPlugin
 ```
 
@@ -975,23 +997,6 @@ Verify using `curl -v localhost:8899/get`:
 }
 ```
 
-### Web Server Route
-
-Demonstrates inbuilt web server routing using plugin.
-
-Start `proxy.py` as:
-
-```console
-❯ proxy --enable-web-server \
-    --plugins proxy.plugin.WebServerPlugin
-```
-
-Verify using `curl -v localhost:8899/http-route-example`, should return:
-
-```console
-HTTP route response
-```
-
 ## Plugin Ordering
 
 When using multiple plugins, depending upon plugin functionality,
diff --git a/proxy/common/constants.py b/proxy/common/constants.py
index e21586a..3f58c51 100644
--- a/proxy/common/constants.py
+++ b/proxy/common/constants.py
@@ -94,6 +94,7 @@ DEFAULT_ENABLE_EVENTS = False
 DEFAULT_EVENTS_QUEUE = None
 DEFAULT_ENABLE_STATIC_SERVER = False
 DEFAULT_ENABLE_WEB_SERVER = False
+DEFAULT_ENABLE_REVERSE_PROXY = False
 DEFAULT_ALLOWED_URL_SCHEMES = [HTTP_PROTO, HTTPS_PROTO]
 DEFAULT_IPV4_HOSTNAME = ipaddress.IPv4Address('127.0.0.1')
 DEFAULT_IPV6_HOSTNAME = ipaddress.IPv6Address('::1')
@@ -110,6 +111,8 @@ DEFAULT_HTTP_PROXY_ACCESS_LOG_FORMAT = '{client_ip}:{client_port} - ' + \
 DEFAULT_HTTPS_PROXY_ACCESS_LOG_FORMAT = '{client_ip}:{client_port} - ' + \
     '{request_method} {server_host}:{server_port} - ' + \
     '{response_bytes} bytes - {connection_time_ms}ms'
+DEFAULT_REVERSE_PROXY_ACCESS_LOG_FORMAT = '{client_ip}:{client_port} - ' + \
+    '{request_method} {request_path} -> {upstream_proxy_pass} - {connection_time_ms}ms'
 DEFAULT_NUM_ACCEPTORS = 0
 DEFAULT_NUM_WORKERS = 0
 DEFAULT_OPEN_FILE_LIMIT = 1024
@@ -149,11 +152,13 @@ DEFAULT_ABC_PLUGINS = [
     'HttpProxyBasePlugin',
     'HttpWebServerBasePlugin',
     'WebSocketTransportBasePlugin',
+    'ReverseProxyBasePlugin',
 ]
 PLUGIN_DASHBOARD = 'proxy.dashboard.ProxyDashboard'
 PLUGIN_HTTP_PROXY = 'proxy.http.proxy.HttpProxyPlugin'
 PLUGIN_PROXY_AUTH = 'proxy.http.proxy.auth.AuthPlugin'
 PLUGIN_WEB_SERVER = 'proxy.http.server.HttpWebServerPlugin'
+PLUGIN_REVERSE_PROXY = 'proxy.http.server.reverse.ReverseProxy'
 PLUGIN_PAC_FILE = 'proxy.http.server.HttpWebServerPacFilePlugin'
 PLUGIN_DEVTOOLS_PROTOCOL = 'proxy.http.inspector.devtools.DevtoolsProtocolPlugin'
 PLUGIN_INSPECT_TRAFFIC = 'proxy.http.inspector.inspect_traffic.InspectTrafficPlugin'
diff --git a/proxy/common/flag.py b/proxy/common/flag.py
index 6786456..b9f420f 100644
--- a/proxy/common/flag.py
+++ b/proxy/common/flag.py
@@ -26,8 +26,8 @@ from .version import __version__
 from .constants import (
     COMMA, IS_WINDOWS, PLUGIN_PAC_FILE, PLUGIN_DASHBOARD, PLUGIN_HTTP_PROXY,
     PLUGIN_PROXY_AUTH, PLUGIN_WEB_SERVER, DEFAULT_NUM_WORKERS,
-    DEFAULT_NUM_ACCEPTORS, PLUGIN_INSPECT_TRAFFIC, DEFAULT_DISABLE_HEADERS,
-    PY2_DEPRECATION_MESSAGE, DEFAULT_DEVTOOLS_WS_PATH,
+    PLUGIN_REVERSE_PROXY, DEFAULT_NUM_ACCEPTORS, PLUGIN_INSPECT_TRAFFIC,
+    DEFAULT_DISABLE_HEADERS, PY2_DEPRECATION_MESSAGE, DEFAULT_DEVTOOLS_WS_PATH,
     PLUGIN_DEVTOOLS_PROTOCOL, PLUGIN_WEBSOCKET_TRANSPORT,
     DEFAULT_DATA_DIRECTORY_PATH, DEFAULT_MIN_COMPRESSION_LIMIT,
 )
@@ -415,6 +415,9 @@ class FlagParser:
                 args.pac_file is not None or \
                 args.enable_static_server:
             default_plugins.append(PLUGIN_WEB_SERVER)
+        if args.enable_reverse_proxy:
+            default_plugins.append(PLUGIN_WEB_SERVER)
+            default_plugins.append(PLUGIN_REVERSE_PROXY)
         if args.pac_file is not None:
             default_plugins.append(PLUGIN_PAC_FILE)
         return list(collections.OrderedDict.fromkeys(default_plugins).keys())
diff --git a/proxy/core/base/tcp_upstream.py b/proxy/core/base/tcp_upstream.py
index 564f9a5..f045f1f 100644
--- a/proxy/core/base/tcp_upstream.py
+++ b/proxy/core/base/tcp_upstream.py
@@ -36,7 +36,7 @@ class TcpUpstreamConnectionHandler(ABC):
         super().__init__(*args, **kwargs)   # type: ignore
         self.upstream: Optional[TcpServerConnection] = None
         # TODO: Currently, :class:`~proxy.core.base.TcpUpstreamConnectionHandler`
-        # is used within :class:`~proxy.plugin.ReverseProxyPlugin` and
+        # is used within :class:`~proxy.http.server.ReverseProxy` and
         # :class:`~proxy.plugin.ProxyPoolPlugin`.
         #
         # For both of which we expect a 4-tuple as arguments
@@ -47,7 +47,7 @@ class TcpUpstreamConnectionHandler(ABC):
         # A separate tunnel class must be created which handles
         # client connection too.
         #
-        # Both :class:`~proxy.plugin.ReverseProxyPlugin` and
+        # Both :class:`~proxy.http.server.ReverseProxy` and
         # :class:`~proxy.plugin.ProxyPoolPlugin` are currently
         # calling client queue within `handle_upstream_data` callback.
         #
diff --git a/proxy/http/server/__init__.py b/proxy/http/server/__init__.py
index dfbaa02..dd4bf30 100644
--- a/proxy/http/server/__init__.py
+++ b/proxy/http/server/__init__.py
@@ -9,7 +9,7 @@
     :license: BSD, see LICENSE for more details.
 """
 from .web import HttpWebServerPlugin
-from .plugin import HttpWebServerBasePlugin
+from .plugin import ReverseProxyBasePlugin, HttpWebServerBasePlugin
 from .protocols import httpProtocolTypes
 from .pac_plugin import HttpWebServerPacFilePlugin
 
@@ -19,4 +19,5 @@ __all__ = [
     'HttpWebServerPacFilePlugin',
     'HttpWebServerBasePlugin',
     'httpProtocolTypes',
+    'ReverseProxyBasePlugin',
 ]
diff --git a/proxy/http/server/plugin.py b/proxy/http/server/plugin.py
index 62ebd3e..7870620 100644
--- a/proxy/http/server/plugin.py
+++ b/proxy/http/server/plugin.py
@@ -96,3 +96,12 @@ class HttpWebServerBasePlugin(DescriptorsHandlerMixin, ABC):
         Return None if plugin has logged the request.
         """
         return context
+
+
+class ReverseProxyBasePlugin(ABC):
+    """ReverseProxy base plugin class."""
+
+    @abstractmethod
+    def routes(self) -> List[Tuple[str, List[bytes]]]:
+        """Return List(path, List(upstream)) reverse proxy config."""
+        raise NotImplementedError()     # pragma: no cover
diff --git a/proxy/http/server/reverse.py b/proxy/http/server/reverse.py
new file mode 100644
index 0000000..326d8ba
--- /dev/null
+++ b/proxy/http/server/reverse.py
@@ -0,0 +1,106 @@
+# -*- 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 re
+import random
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Optional
+
+from . import HttpWebServerBasePlugin, httpProtocolTypes
+from .. import Url
+from ..parser import HttpParser
+from ..exception import HttpProtocolException
+from ...core.base import TcpUpstreamConnectionHandler
+from ...common.utils import text_
+from ...common.constants import (
+    DEFAULT_HTTP_PORT, DEFAULT_HTTPS_PORT,
+    DEFAULT_REVERSE_PROXY_ACCESS_LOG_FORMAT,
+)
+
+
+if TYPE_CHECKING:
+    from .plugin import ReverseProxyBasePlugin
+
+
+logger = logging.getLogger(__name__)
+
+
+class ReverseProxy(TcpUpstreamConnectionHandler, HttpWebServerBasePlugin):
+    """Extend in-built Web Server to add Reverse Proxy capabilities."""
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self.choice: Optional[Url] = None
+        self.reverse: Dict[str, List[bytes]] = {}
+
+    def handle_upstream_data(self, raw: memoryview) -> None:
+        self.client.queue(raw)
+
+    def routes(self) -> List[Tuple[int, str]]:
+        reverse: List[Tuple[str, List[bytes]]] = []
+        for klass in self.flags.plugins[b'ReverseProxyBasePlugin']:
+            instance: 'ReverseProxyBasePlugin' = klass()
+            reverse.extend(instance.routes())
+        r = []
+        for (route, upstreams) in reverse:
+            r.append((httpProtocolTypes.HTTP, route))
+            r.append((httpProtocolTypes.HTTPS, route))
+            self.reverse[route] = upstreams
+        return r
+
+    def handle_request(self, request: HttpParser) -> None:
+        # TODO: Core must be capable of dispatching a context
+        # with each invocation of handle request callback.
+        #
+        # Example, here we don't know which of our registered
+        # route actually matched.
+        #
+        for route in self.reverse:
+            pattern = re.compile(route)
+            if pattern.match(text_(request.path)):
+                self.choice = Url.from_bytes(
+                    random.choice(self.reverse[route]),
+                )
+                break
+        assert self.choice and self.choice.hostname
+        port = self.choice.port or \
+            DEFAULT_HTTP_PORT \
+            if self.choice.scheme == b'http' \
+            else DEFAULT_HTTPS_PORT
+        self.initialize_upstream(text_(self.choice.hostname), port)
+        assert self.upstream
+        try:
+            self.upstream.connect()
+            if self.choice.scheme == b'https':
+                self.upstream.wrap(
+                    text_(
+                        self.choice.hostname,
+                    ), ca_file=str(self.flags.ca_file),
+                )
+            self.upstream.queue(memoryview(request.build()))
+        except ConnectionRefusedError:
+            raise HttpProtocolException(
+                'Connection refused by upstream server {0}:{1}'.format(
+                    text_(self.choice.hostname), port,
+                ),
+            )
+
+    def on_client_connection_close(self) -> None:
+        if self.upstream and not self.upstream.closed:
+            logger.debug('Closing upstream server connection')
+            self.upstream.close()
+            self.upstream = None
+
+    def on_access_log(self, context: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        context.update({
+            'upstream_proxy_pass': str(self.choice) if self.choice else None,
+        })
+        logger.info(DEFAULT_REVERSE_PROXY_ACCESS_LOG_FORMAT.format_map(context))
+        return None
diff --git a/proxy/http/server/web.py b/proxy/http/server/web.py
index efc734c..f6b0a12 100644
--- a/proxy/http/server/web.py
+++ b/proxy/http/server/web.py
@@ -28,8 +28,8 @@ from ...common.types import Readables, Writables, Descriptors
 from ...common.utils import text_, bytes_, build_websocket_handshake_response
 from ...common.constants import (
     DEFAULT_ENABLE_WEB_SERVER, DEFAULT_STATIC_SERVER_DIR,
-    DEFAULT_ENABLE_STATIC_SERVER, DEFAULT_MIN_COMPRESSION_LIMIT,
-    DEFAULT_WEB_ACCESS_LOG_FORMAT,
+    DEFAULT_ENABLE_REVERSE_PROXY, DEFAULT_ENABLE_STATIC_SERVER,
+    DEFAULT_MIN_COMPRESSION_LIMIT, DEFAULT_WEB_ACCESS_LOG_FORMAT,
 )
 
 
@@ -70,6 +70,13 @@ flags.add_argument(
     'Sets the minimum length of a response that will be compressed (gzipped).',
 )
 
+flags.add_argument(
+    '--enable-reverse-proxy',
+    action='store_true',
+    default=DEFAULT_ENABLE_REVERSE_PROXY,
+    help='Default: False.  Whether to enable reverse proxy core.',
+)
+
 
 class HttpWebServerPlugin(HttpProtocolHandlerPlugin):
     """HttpProtocolHandler plugin which handles incoming requests to local web server."""
diff --git a/proxy/plugin/reverse_proxy.py b/proxy/plugin/reverse_proxy.py
index e21442e..01eca2b 100644
--- a/proxy/plugin/reverse_proxy.py
+++ b/proxy/plugin/reverse_proxy.py
@@ -8,26 +8,26 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
-import random
-import logging
-from typing import Any, Dict, List, Tuple, Optional
+from typing import List, Tuple
 
-from ..http import Url
-from ..core.base import TcpUpstreamConnectionHandler
-from ..http.parser import HttpParser
-from ..http.server import HttpWebServerBasePlugin, httpProtocolTypes
-from ..common.utils import text_
-from ..http.exception import HttpProtocolException
-from ..common.constants import DEFAULT_HTTP_PORT, DEFAULT_HTTPS_PORT
+from ..http.server import ReverseProxyBasePlugin
 
 
-logger = logging.getLogger(__name__)
+# TODO: We must use nginx python parser and
+# make this plugin nginx.conf complaint.
+REVERSE_PROXY_LOCATION: str = r'/get$'
+# Randomly choose either http or https upstream endpoint.
+#
+# This is just to demonstrate that both http and https upstream
+# reverse proxy works.
+REVERSE_PROXY_PASS = [
+    b'http://httpbin.org/get',
+    b'https://httpbin.org/get',
+]
 
 
-class ReverseProxyPlugin(TcpUpstreamConnectionHandler, HttpWebServerBasePlugin):
-    """Extend in-built Web Server to add Reverse Proxy capabilities.
-
-    This example plugin is equivalent to following Nginx configuration::
+class ReverseProxyPlugin(ReverseProxyBasePlugin):
+    """This example plugin is equivalent to following Nginx configuration::
 
         ```text
         location /get {
@@ -35,86 +35,8 @@ class ReverseProxyPlugin(TcpUpstreamConnectionHandler, HttpWebServerBasePlugin):
         }
         ```
 
-    Example::
-
-        ```console
-        $ curl http://localhost:9000/get
-        {
-          "args": {},
-          "headers": {
-            "Accept": "*/*",
-            "Host": "localhost",
-            "User-Agent": "curl/7.64.1"
-          },
-          "origin": "1.2.3.4, 5.6.7.8",
-          "url": "http://localhost/get"
-        }
-        ```
+    Update the routes config before.
     """
 
-    # TODO: We must use nginx python parser and
-    # make this plugin nginx.conf complaint.
-    REVERSE_PROXY_LOCATION: str = r'/get$'
-    # Randomly choose either http or https upstream endpoint.
-    #
-    # This is just to demonstrate that both http and https upstream
-    # reverse proxy works.
-    REVERSE_PROXY_PASS = [
-        b'http://httpbin.org/get',
-        b'https://httpbin.org/get',
-    ]
-
-    def __init__(self, *args: Any, **kwargs: Any):
-        super().__init__(*args, **kwargs)
-        self.choice: Optional[Url] = None
-
-    def handle_upstream_data(self, raw: memoryview) -> None:
-        self.client.queue(raw)
-
-    def routes(self) -> List[Tuple[int, str]]:
-        return [
-            (httpProtocolTypes.HTTP, ReverseProxyPlugin.REVERSE_PROXY_LOCATION),
-            (httpProtocolTypes.HTTPS, ReverseProxyPlugin.REVERSE_PROXY_LOCATION),
-        ]
-
-    def handle_request(self, request: HttpParser) -> None:
-        self.choice = Url.from_bytes(
-            random.choice(ReverseProxyPlugin.REVERSE_PROXY_PASS),
-        )
-        assert self.choice.hostname
-        port = self.choice.port or \
-            DEFAULT_HTTP_PORT \
-            if self.choice.scheme == b'http' \
-            else DEFAULT_HTTPS_PORT
-
-        self.initialize_upstream(text_(self.choice.hostname), port)
-        assert self.upstream
-        try:
-            self.upstream.connect()
-            if self.choice.scheme == b'https':
-                self.upstream.wrap(
-                    text_(
-                        self.choice.hostname,
-                    ), ca_file=str(self.flags.ca_file),
-                )
-            self.upstream.queue(memoryview(request.build()))
-        except ConnectionRefusedError:
-            raise HttpProtocolException(
-                'Connection refused by upstream server {0}:{1}'.format(
-                    text_(self.choice.hostname), port,
-                ),
-            )
-
-    def on_client_connection_close(self) -> None:
-        if self.upstream and not self.upstream.closed:
-            logger.debug('Closing upstream server connection')
-            self.upstream.close()
-            self.upstream = None
-
-    def on_access_log(self, context: Dict[str, Any]) -> Optional[Dict[str, Any]]:
-        log_format = '{client_addr} - {request_method} {request_path} -> {upstream_proxy_pass} - {connection_time_ms}ms'
-        context.update({
-            'upstream_proxy_pass': str(self.choice) if self.choice else None,
-        })
-        logger.info(log_format.format_map(context))
-        return None
+    def routes(self) -> List[Tuple[str, List[bytes]]]:
+        return [(REVERSE_PROXY_LOCATION, REVERSE_PROXY_PASS)]


Test patch:
diff --git a/tests/test_main.py b/tests/test_main.py
index f3b7a85..718fa17 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -30,8 +30,8 @@ from proxy.common.constants import (  # noqa: WPS450
     DEFAULT_ENABLE_SSH_TUNNEL, DEFAULT_ENABLE_WEB_SERVER,
     DEFAULT_DISABLE_HTTP_PROXY, PLUGIN_WEBSOCKET_TRANSPORT,
     DEFAULT_CA_SIGNING_KEY_FILE, DEFAULT_CLIENT_RECVBUF_SIZE,
-    DEFAULT_SERVER_RECVBUF_SIZE, DEFAULT_ENABLE_STATIC_SERVER,
-    _env_threadless_compliant,
+    DEFAULT_SERVER_RECVBUF_SIZE, DEFAULT_ENABLE_REVERSE_PROXY,
+    DEFAULT_ENABLE_STATIC_SERVER, _env_threadless_compliant,
 )
 
 
@@ -77,6 +77,7 @@ class TestMain(unittest.TestCase):
         mock_args.local_executor = int(DEFAULT_LOCAL_EXECUTOR)
         mock_args.port_file = DEFAULT_PORT_FILE
         mock_args.enable_ssh_tunnel = DEFAULT_ENABLE_SSH_TUNNEL
+        mock_args.enable_reverse_proxy = DEFAULT_ENABLE_REVERSE_PROXY
 
     @mock.patch('os.remove')
     @mock.patch('os.path.exists')
@@ -243,7 +244,7 @@ class TestMain(unittest.TestCase):
         self.mock_default_args(mock_args)
         mock_args.enable_dashboard = True
         mock_args.local_executor = 0
-        main(enable_dashboard=True, local_executor=0)
+        main()
         mock_load_plugins.assert_called()
         self.assertEqual(
             mock_load_plugins.call_args_list[0][0][0], [
@@ -290,7 +291,7 @@ class TestMain(unittest.TestCase):
         self.mock_default_args(mock_args)
         mock_args.enable_devtools = True
         mock_args.local_executor = 0
-        main(enable_devtools=True, local_executor=0)
+        main()
         mock_load_plugins.assert_called()
         self.assertEqual(
             mock_load_plugins.call_args_list[0][0][0], [
@@ -335,7 +336,7 @@ class TestMain(unittest.TestCase):
         self.mock_default_args(mock_args)
         mock_args.enable_ssh_tunnel = True
         mock_args.local_executor = 0
-        main(enable_ssh_tunnel=True, local_executor=0)
+        main()
         mock_load_plugins.assert_called()
         self.assertEqual(
             mock_load_plugins.call_args_list[0][0][0], [