fixup UDP support

mitmproxy · mhils · Sep 18, 2022 · Sep 4, 2022 · Sep 6, 2022 · Sep 11, 2022
commit a5f75a02a597a7fb51e8f14a12e0ca49d33e9588
diff --git a/mitmproxy/net/udp.py b/mitmproxy/net/udp.py
@@ -82,7 +82,7 @@ async def wait_closed(self) -> None:
 
 class WireGuardDatagramTransport(asyncio.DatagramTransport):
     def __init__(self, server: wg.Server, local_addr: Address, remote_addr: Address):
-        self._server = server
+        self._server: wg.Server = server
         self._local_addr: Address = local_addr
         self._remote_addr: Address = remote_addr
         self._protocol: DrainableDatagramProtocol | None = None
@@ -91,12 +91,17 @@ def __init__(self, server: wg.Server, local_addr: Address, remote_addr: Address)
     def sendto(self, data, addr=None):
         self._server.send_datagram(data, self._local_addr, addr or self._remote_addr)
 
-    def set_protocol(self, protocol):
-        assert isinstance(protocol, DrainableDatagramProtocol)
-        self._protocol = cast(DrainableDatagramProtocol, protocol)
+    def get_extra_info(self, name: str, default: Any = None) -> Any:
+        if name == "sockname":
+            return self._server.getsockname()
+        else:
+            raise NotImplementedError
 
     def get_protocol(self):
-        return self._protocol
+        return self
+
+    async def drain(self) -> None:
+        pass
 
 
 DatagramTransport = Union[asyncio.DatagramTransport, WireGuardDatagramTransport]
@@ -201,7 +206,7 @@ def __init__(
         self._transport = transport
         self._remote_addr = remote_addr
         proto = transport.get_protocol()
-        assert isinstance(proto, DrainableDatagramProtocol)
+        assert isinstance(proto, (DrainableDatagramProtocol, WireGuardDatagramTransport))
         self._reader = reader
         self._closed = asyncio.Event() if reader is not None else None
 

diff --git a/mitmproxy/proxy/mode_servers.py b/mitmproxy/proxy/mode_servers.py
@@ -392,6 +392,7 @@ def handle_udp_datagram(self, data: bytes, remote_addr: Address, local_addr: Add
             handler.layer = self.make_top_layer(handler.layer.context)
             handler.layer.context.client.transport_protocol = "udp"
             handler.layer.context.server.transport_protocol = "udp"
+            handler.layer.context.server.address = local_addr
 
             # pre-register here - we may get datagrams before the task is executed.
             self.manager.connections[connection_id] = handler

diff --git a/mitmproxy/proxy/server.py b/mitmproxy/proxy/server.py
@@ -129,6 +129,8 @@ async def handle_client(self) -> None:
             self.transports[self.client].handler = handler
             self.server_event(events.Start())
             await asyncio.wait([handler])
+            if e := handler.exception():
+                self.log(f"mitmproxy has crashed!\n{traceback.format_exception(e)}", level="error")
 
         watch.cancel()
         while self.wakeup_timer: