Refactor how we handle UDP

[mhils]

Our current UDP implementation roughly looks as follows:

1. On the Python side, we use asyncio's UDP functionality, and them some custom transport/protocol wrapper classes that provide similar APIs to asyncio TCP streams.
2. On the Rust side, we currently create DatagramTransport objects and then pass them to Python.

This has two major problems:

1. With asyncio's UDP protocols, we only get the remote address and packet bytes for each packet. We would really like to have the local address as well (to generate matching certificates), but that's not available via the asyncio UDP APIs [^1].
2. For UDP packets from Rust (WireGuard mode, OS proxy mode), the instantiation of a new DatagramTransport object for each packet is quite inefficient.

Now, how to fix it? I have long thought that a packet-based handle_udp_datagram(data, address) API is the right abstraction for UDP (in contrast to handle_tcp_connection, which is called once for an entire connection, this is called for each packet). But seeing how Apple's Network Extension API abstractions work – handleNewFlow and handleNewUDPFlow –, I think we should evaluate a similar approach as well. More concretely, in a first step, we could add a create_udp_server(host, port, handle_udp_flow) method over in Rust, which takes care of grouping packets by 4-tuple and calling handle_udp_flow with one DatagramTransport object per flow. If this works, we can reuse our flow grouping logic for the transparent modes as well. Over in Python land, UDP flows would then basically be handled the same way as TCP connections (right now we do a similar type of grouping in Python).

TL;DR: More Rust to work around asyncio limitations, and also improve our existing UDP Rust story on the way?

(cc @meitinger @decathorpe)

[^1]: We can obviously do .get_extra_info("sockname"), but that does not help for sockets binding to :: or 0.0.0.0.

[decathorpe]

So you're saying the "fake" UDP connections I had originally implemented in mitmproxy_wireguard are coming back? 😉

[meitinger]

The suggested approach makes a lot of sense to me, especially if the Python-side sees it as a connection and Rust does all the mapping from (local-endpoint,remote-endpoint) => connection. There is only one caveat where things might get trickier this way down the line, namely if we ever want to support full QUIC roaming.

[As food for though, maybe - in the long run - it would make sense to do the same for TCP as well, i.e. being protocol-agnostic in Python (cleaning out some OS-dependent code) and support pluggable Rust endpoints. So future modes would be a combination of what (regular, reverse, transparent, etc.) and from where (local socket, wireguard, ...) intead of having modes for where with an implicit assumption of what.]

[mhils]

So you're saying the "fake" UDP connections I had originally implemented in mitmproxy_wireguard are coming back? 😉

Yes, you get your "told you so" card. 😄 Do we still have that code lying around somewhere by chance?

There is only one caveat where things might get trickier this way down the line, namely if we ever want to support full QUIC roaming.

Yes. That's a special case we can handle within the QUIC layer I believe.

[As food for though, maybe - in the long run - it would make sense to do the same for TCP as well, i.e. being protocol-agnostic in Python (cleaning out some OS-dependent code) and support pluggable Rust endpoints.

I think this is not a completely bad idea... although the asyncio TCP story is not too bad from our end. Let's do UDP first and see how that goes!

---

Thank you two for the quick feedback. 🍰 https://github.com/mitmproxy/mitmproxy/commit/d587353ba0fd485d88161a042f2774f5de804234 is hopefully good enough for a first HTTP/3 release, but I'll give this a spin afterwards. :)

[decathorpe]

So you're saying the "fake" UDP connections I had originally implemented in mitmproxy_wireguard are coming back? wink

Yes, you get your "told you so" card. smile Do we still have that code lying around somewhere by chance?

I'm not even sure if this used to work the way you'd want it to work now, but here's the original issue / PR: https://github.com/decathorpe/mitmproxy_wireguard/issues/7 https://github.com/decathorpe/mitmproxy_wireguard/pull/8

If this would work the way you want it to, we might be able to re-purpose parts of the code here.

[mhils]

Awesome, thanks! I'll take a closer look when working on this. 😃 🍰

[mhils]

I've made considerable progress on this over at https://github.com/mitmproxy/mitmproxy_rs/pull/127. WireGuard, Windows Redirector, and plain UDP servers all work. I still need to fix macOS local redirect mode and then rip out all the old cruft. :)

[mhils]

And done! Fake UDP connections are back. 😄 I'll look over things tomorrow again with a fresh mind, but testing has been super successful so far. UDP clients and servers are now implemented entirely in Rust (which is much more sane than DatagramTransport). WireGuard and local redirect mode are ported as well. 🙌