Add CustomDnsResolver plugin. Addresses #535 and #664

Author: abhinavsingh

LICENSE

@@ -1,4 +1,4 @@
-Copyright (c) 2013-2022 by Abhinav Singh and contributors.
+Copyright (c) 2013-present by Abhinav Singh and contributors.
 All rights reserved.
 
 Redistribution and use in source and binary forms, with or without modification,

README.md

@@ -56,6 +56,7 @@
     - [Proxy Pool Plugin](#proxypoolplugin)
     - [FilterByClientIpPlugin](#filterbyclientipplugin)
     - [ModifyChunkResponsePlugin](#modifychunkresponseplugin)
+    - [CustomDnsResolverPlugin](#customdnsresolverplugin)
   - [HTTP Web Server Plugins](#http-web-server-plugins)
     - [Reverse Proxy](#reverse-proxy)
     - [Web Server Route](#web-server-route)
@@ -720,6 +721,22 @@ plugin
 
 Modify `ModifyChunkResponsePlugin` to your taste. Example, instead of sending hardcoded chunks, parse and modify the original `JSON` chunks received from the upstream server.
 
+### CustomDnsResolverPlugin
+
+This plugin demonstrate how to use a custom DNS resolution implementation with `proxy.py`.
+This example plugin currently uses Python's in-built resolution mechanism.  Customize code
+to your taste.  Example, query your custom DNS server, implement DoH or other mechanisms.
+
+Start `proxy.py` as:
+
+```bash
+❯ proxy \
+    --plugins proxy.plugin.CustomDnsResolverPlugin
+```
+
+`HttpProxyBasePlugin.resolve_dns` can also be used to configure `network interface` which
+must be used as the `source_address` for connection to the upstream server.
+
 ## HTTP Web Server Plugins
 
 ### Reverse Proxy

proxy/common/utils.py

@@ -181,7 +181,7 @@ def wrap_socket(
 
 
 def new_socket_connection(
-        addr: Tuple[str, int], timeout: int = DEFAULT_TIMEOUT,
+        addr: Tuple[str, int], timeout: int = DEFAULT_TIMEOUT, source_address: Optional[Tuple[str, int]] = None,
 ) -> socket.socket:
     conn = None
     try:
@@ -205,7 +205,7 @@ def new_socket_connection(
         return conn
 
     # try to establish dual stack IPv4/IPv6 connection.
-    return socket.create_connection(addr, timeout=timeout)
+    return socket.create_connection(addr, timeout=timeout, source_address=source_address)
 
 
 class socket_connection(contextlib.ContextDecorator):

proxy/core/connection/server.py

@@ -33,9 +33,11 @@ def connection(self) -> Union[ssl.SSLSocket, socket.socket]:
             raise TcpConnectionUninitializedException()
         return self._conn
 
-    def connect(self) -> None:
+    def connect(self, addr: Optional[Tuple[str, int]] = None, source_address: Optional[Tuple[str, int]] = None) -> None:
         if self._conn is None:
-            self._conn = new_socket_connection(self.addr)
+            self._conn = new_socket_connection(
+                addr or self.addr, source_address=source_address,
+            )
         self.closed = False
 
     def wrap(self, hostname: str, ca_file: Optional[str]) -> None:

proxy/http/proxy/plugin.py

@@ -77,6 +77,22 @@ def read_from_descriptors(self, r: Readables) -> bool:
         """Implementations must now read data over the socket."""
         return False  # pragma: no cover
 
+    def resolve_dns(self, host: str, port: int) -> Tuple[Optional[str], Optional[Tuple[str, int]]]:
+        """Resolve upstream server host to an IP address.
+
+        Optionally also override the source address to use for
+        connection with upstream server.
+
+        For upstream IP:
+        Return None to use default resolver available to the system.
+        Return ip address as string to use your custom resolver.
+
+        For source address:
+        Return None to use default source address
+        Return 2-tuple representing (host, port) to use as source address
+        """
+        return None, None
+
     @abstractmethod
     def before_upstream_connection(
             self, request: HttpParser,

proxy/http/proxy/server.py

@@ -509,7 +509,21 @@ def connect_upstream(self) -> None:
                     'Connecting to upstream %s:%s' %
                     (text_(host), port),
                 )
-                self.server.connect()
+                # Invoke plugin.resolve_dns
+                upstream_ip, source_addr = None, None
+                for plugin in self.plugins.values():
+                    upstream_ip, source_addr = plugin.resolve_dns(
+                        text_(host), port,
+                    )
+                    if upstream_ip or source_addr:
+                        break
+                # Connect with overridden upstream IP and source address
+                # if any of the plugin returned a non-null value.
+                self.server.connect(
+                    addr=None if not upstream_ip else (
+                        upstream_ip, port,
+                    ), source_address=source_addr,
+                )
                 self.server.connection.setblocking(False)
                 logger.debug(
                     'Connected to upstream %s:%s' %

proxy/plugin/__init__.py

@@ -21,6 +21,7 @@
 from .filter_by_client_ip import FilterByClientIpPlugin
 from .filter_by_url_regex import FilterByURLRegexPlugin
 from .modify_chunk_response import ModifyChunkResponsePlugin
+from .custom_dns_resolver import CustomDnsResolverPlugin
 
 __all__ = [
     'CacheResponsesPlugin',
@@ -37,4 +38,5 @@
     'FilterByClientIpPlugin',
     'ModifyChunkResponsePlugin',
     'FilterByURLRegexPlugin',
+    'CustomDnsResolverPlugin',
 ]

proxy/plugin/custom_dns_resolver.py

@@ -0,0 +1,52 @@
+# -*- 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
+
+from typing import Optional, Tuple
+
+from ..http.parser import HttpParser
+from ..http.proxy import HttpProxyBasePlugin
+
+
+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]]]:
+        """Here we are using in-built python resolver for demonstration.
+
+        Ideally you would like to query your custom DNS server or even use DoH to make
+        real sense out of this plugin.
+
+        2nd parameter returned is None.  Return a 2-tuple to configure underlying interface
+        to use for connection to the upstream server.
+        """
+        try:
+            return socket.getaddrinfo(host, port, proto=socket.IPPROTO_TCP)[0][4][0], None
+        except socket.gaierror:
+            # Ideally we can also thrown HttpRequestRejected or HttpProtocolException here
+            # Returning None simply fallback to core generated exceptions.
+            return None, None
+
+    def before_upstream_connection(
+            self, request: HttpParser,
+    ) -> Optional[HttpParser]:
+        return request
+
+    def handle_client_request(
+            self, request: HttpParser,
+    ) -> Optional[HttpParser]:
+        return request
+
+    def handle_upstream_chunk(self, chunk: memoryview) -> memoryview:
+        return chunk
+
+    def on_upstream_connection_close(self) -> None:
+        pass

tests/common/test_utils.py

@@ -43,7 +43,9 @@ def test_new_socket_connection_ipv6(self, mock_socket: mock.Mock) -> None:
     @mock.patch('socket.create_connection')
     def test_new_socket_connection_dual(self, mock_socket: mock.Mock) -> None:
         conn = new_socket_connection(self.addr_dual)
-        mock_socket.assert_called_with(self.addr_dual, timeout=DEFAULT_TIMEOUT)
+        mock_socket.assert_called_with(
+            self.addr_dual, timeout=DEFAULT_TIMEOUT, source_address=None,
+        )
         self.assertEqual(conn, mock_socket.return_value)
 
     @mock.patch('proxy.common.utils.new_socket_connection')

tests/http/test_http_proxy.py

@@ -60,6 +60,7 @@ def test_proxy_plugin_on_and_before_upstream_connection(
         self.plugin.return_value.read_from_descriptors.return_value = False
         self.plugin.return_value.before_upstream_connection.side_effect = lambda r: r
         self.plugin.return_value.handle_client_request.side_effect = lambda r: r
+        self.plugin.return_value.resolve_dns.return_value = None, None
 
         self._conn.recv.return_value = build_http_request(
             b'GET', b'http://upstream.host/not-found.html',

tests/http/test_http_proxy_tls_interception.py

@@ -133,6 +133,7 @@ def mock_connection() -> Any:
         self.proxy_plugin.return_value.read_from_descriptors.return_value = False
         self.proxy_plugin.return_value.before_upstream_connection.side_effect = lambda r: r
         self.proxy_plugin.return_value.handle_client_request.side_effect = lambda r: r
+        self.proxy_plugin.return_value.resolve_dns.return_value = None, None
 
         self.mock_selector.return_value.select.side_effect = [
             [(