From 95bf441bc3a1dfcac2cb313cac8e72d3716a4bed Mon Sep 17 00:00:00 2001 From: Gaurav Jain <64748057+errorxyz@users.noreply.github.com> Date: Mon, 17 Jun 2024 13:03:25 +0530 Subject: [PATCH] make DNS mode listen for both UDP and TCP (#6912) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- mitmproxy/proxy/mode_servers.py | 47 ++++++++++++++--------- mitmproxy/proxy/mode_specs.py | 13 ++++--- test/mitmproxy/proxy/test_mode_servers.py | 6 ++- 3 files changed, 41 insertions(+), 25 deletions(-) diff --git a/mitmproxy/proxy/mode_servers.py b/mitmproxy/proxy/mode_servers.py index 921789a0a..a8ff78d70 100644 --- a/mitmproxy/proxy/mode_servers.py +++ b/mitmproxy/proxy/mode_servers.py @@ -271,7 +271,11 @@ class AsyncioServerInstance(ServerInstance[M], metaclass=ABCMeta): async def listen( self, host: str, port: int ) -> list[asyncio.Server | mitmproxy_rs.UdpServer]: - if self.mode.transport_protocol == "tcp": + if self.mode.transport_protocol not in ("tcp", "udp", "both"): + raise AssertionError(self.mode.transport_protocol) + + servers: list[asyncio.Server | mitmproxy_rs.UdpServer] = [] + if self.mode.transport_protocol in ("tcp", "both"): # workaround for https://github.com/python/cpython/issues/89856: # We want both IPv4 and IPv6 sockets to bind to the same port. # This may fail (https://github.com/mitmproxy/mitmproxy/pull/5542#issuecomment-1222803291), @@ -280,17 +284,24 @@ class AsyncioServerInstance(ServerInstance[M], metaclass=ABCMeta): try: s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) s.bind(("", 0)) - fixed_port = s.getsockname()[1] + port = s.getsockname()[1] s.close() - return [ - await asyncio.start_server(self.handle_stream, host, fixed_port) - ] + servers.append( + await asyncio.start_server(self.handle_stream, host, port) + ) except Exception as e: logger.debug( f"Failed to listen on a single port ({e!r}), falling back to default behavior." ) - return [await asyncio.start_server(self.handle_stream, host, port)] - elif self.mode.transport_protocol == "udp": + port = 0 + servers.append( + await asyncio.start_server(self.handle_stream, host, port) + ) + else: + servers.append( + await asyncio.start_server(self.handle_stream, host, port) + ) + if self.mode.transport_protocol in ("udp", "both"): # we start two servers for dual-stack support. # On Linux, this would also be achievable by toggling IPV6_V6ONLY off, but this here works cross-platform. if host == "": @@ -299,26 +310,26 @@ class AsyncioServerInstance(ServerInstance[M], metaclass=ABCMeta): port, self.handle_udp_stream, ) + servers.append(ipv4) try: ipv6 = await mitmproxy_rs.start_udp_server( "::", ipv4.getsockname()[1], self.handle_udp_stream, ) + servers.append(ipv6) # pragma: no cover except Exception: # pragma: no cover logger.debug("Failed to listen on '::', listening on IPv4 only.") - return [ipv4] - else: # pragma: no cover - return [ipv4, ipv6] - return [ - await mitmproxy_rs.start_udp_server( - host, - port, - self.handle_udp_stream, + else: + servers.append( + await mitmproxy_rs.start_udp_server( + host, + port, + self.handle_udp_stream, + ) ) - ] - else: - raise AssertionError(self.mode.transport_protocol) + + return servers class WireGuardServerInstance(ServerInstance[mode_specs.WireGuardMode]): diff --git a/mitmproxy/proxy/mode_specs.py b/mitmproxy/proxy/mode_specs.py index 7d189ddf3..4aed72856 100644 --- a/mitmproxy/proxy/mode_specs.py +++ b/mitmproxy/proxy/mode_specs.py @@ -90,7 +90,7 @@ class ProxyMode(Serializable, metaclass=ABCMeta): @property @abstractmethod - def transport_protocol(self) -> Literal["tcp", "udp"] | None: + def transport_protocol(self) -> Literal["tcp", "udp", "both"] | None: """The transport protocol used by this mode's server.""" @classmethod @@ -172,8 +172,9 @@ class ProxyMode(Serializable, metaclass=ABCMeta): raise dataclasses.FrozenInstanceError("Proxy modes are immutable.") -TCP: Literal["tcp", "udp"] = "tcp" -UDP: Literal["tcp", "udp"] = "udp" +TCP: Literal["tcp", "udp", "both"] = "tcp" +UDP: Literal["tcp", "udp", "both"] = "udp" +BOTH: Literal["tcp", "udp", "both"] = "both" def _check_empty(data): @@ -230,8 +231,10 @@ class ReverseMode(ProxyMode): # noinspection PyDataclass def __post_init__(self) -> None: self.scheme, self.address = server_spec.parse(self.data, default_scheme="https") - if self.scheme in ("http3", "dtls", "udp", "dns", "quic"): + if self.scheme in ("http3", "dtls", "udp", "quic"): self.transport_protocol = UDP + elif self.scheme == "dns": + self.transport_protocol = BOTH self.description = f"{self.description} to {self.data}" @property @@ -257,7 +260,7 @@ class DnsMode(ProxyMode): description = "DNS server" default_port = 53 - transport_protocol = UDP + transport_protocol = BOTH def __post_init__(self) -> None: _check_empty(self.data) diff --git a/test/mitmproxy/proxy/test_mode_servers.py b/test/mitmproxy/proxy/test_mode_servers.py index 26e53345e..677fa3af0 100644 --- a/test/mitmproxy/proxy/test_mode_servers.py +++ b/test/mitmproxy/proxy/test_mode_servers.py @@ -276,10 +276,12 @@ async def test_udp_start_error(): manager = MagicMock() with taddons.context(): - inst = ServerInstance.make("dns@127.0.0.1:0", manager) + inst = ServerInstance.make("reverse:udp://127.0.0.1:1234@127.0.0.1:0", manager) await inst.start() port = inst.listen_addrs[0][1] - inst2 = ServerInstance.make(f"dns@127.0.0.1:{port}", manager) + inst2 = ServerInstance.make( + f"reverse:udp://127.0.0.1:1234@127.0.0.1:{port}", manager + ) with pytest.raises( Exception, match=f"Failed to bind UDP socket to 127.0.0.1:{port}" ):