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

@@ -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]):

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)

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}"
         ):