Repository: decathorpe/mitmproxy_wireguard (language: python)

Issue:
implement UDP equivalent of TcpStream
mitmproxy has its own `asyncio.StreamReader/StreamWriter` UDP analog wrapper implementation based on Python's `asyncio` module.

To integrate the UDP functionality from `mitmproxy_wireguard` into mitmproxy, a similar wrapper interface (`UdpStream`?) around the existing UDP functionality will be required.


Discussion / hints:
/cc @meitinger - a simple callback for incoming packets + an API to send could be good enough?

Fix patch:
diff --git a/benches/wg_echo_server.py b/benches/wg_echo_server.py
index 5fb46f8..2124cac 100644
--- a/benches/wg_echo_server.py
+++ b/benches/wg_echo_server.py
@@ -80,7 +80,7 @@ async def handle_connection(rw: mitmproxy_wireguard.TcpStream):
     logger.debug(f"Connection closed: {rw}")
 
 
-def receive_datagram(_data, _src_addr, _dst_addr):
+async def receive_datagram(_stream):
     pass
 
 
diff --git a/mitmproxy_wireguard.pyi b/mitmproxy_wireguard.pyi
index daf7554..a25b751 100644
--- a/mitmproxy_wireguard.pyi
+++ b/mitmproxy_wireguard.pyi
@@ -23,9 +23,18 @@ class TcpStream:
     def get_extra_info(self, name: str) -> tuple[str, int]: ...
     def __repr__(self) -> str: ...
 
+class UdpStream:
+    async def read(self, n: int) -> bytes: ...
+    def write(self, data: bytes): ...
+    async def drain(self) -> None: ...
+    def write_eof(self): ...
+    def close(self): ...
+    def get_extra_info(self, name: str) -> tuple[str, int]: ...
+    def __repr__(self) -> str: ...
+
 async def start_server(
     host: str,
     conf: WireGuardServerConf,
     handle_connection: Callable[[TcpStream], Awaitable[None]],
-    receive_datagram: Callable[[bytes, tuple[str, int], tuple[str, int]], None],
+    receive_datagram: Callable[[UdpStream], Awaitable[None]],
 ) -> WireGuardServer: ...
diff --git a/src/lib.rs b/src/lib.rs
index 99937d6..a6f5e66 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -23,7 +23,7 @@ mod wireguard;
 use conf::WireGuardServerConf;
 use messages::TransportCommand;
 use network::NetworkTask;
-use python::{event_queue_unavailable, py_to_socketaddr, socketaddr_to_py, PyInteropTask, TcpStream};
+use python::{event_queue_unavailable, py_to_socketaddr, socketaddr_to_py, PyInteropTask, TcpStream, UdpStream};
 use shutdown::ShutdownTask;
 use wireguard::WireGuardTaskBuilder;
 
@@ -221,11 +221,11 @@ fn start_server(
     py: Python<'_>,
     host: String,
     conf: WireGuardServerConf,
-    handle_connection: PyObject,
-    receive_datagram: PyObject,
+    tcp_handler: PyObject,
+    udp_handler: PyObject,
 ) -> PyResult<&PyAny> {
     pyo3_asyncio::tokio::future_into_py(py, async move {
-        let server = WireGuardServer::init(host, conf, handle_connection, receive_datagram).await?;
+        let server = WireGuardServer::init(host, conf, tcp_handler, udp_handler).await?;
         Ok(server)
     })
 }
@@ -247,6 +247,7 @@ fn mitmproxy_wireguard(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<WireGuardServer>()?;
     m.add_class::<WireGuardServerConf>()?;
     m.add_class::<TcpStream>()?;
+    m.add_class::<UdpStream>()?;
 
     Ok(())
 }
diff --git a/src/python.rs b/src/python.rs
index c3823c4..00429ab 100644
--- a/src/python.rs
+++ b/src/python.rs
@@ -126,6 +126,94 @@ impl Drop for TcpStream {
     }
 }
 
+
+#[pyclass]
+#[derive(Clone, Debug)]
+pub struct UdpStream {
+    data: Vec<u8>,
+    event_tx: mpsc::UnboundedSender<TransportCommand>,
+    peername: SocketAddr,
+    sockname: SocketAddr,
+    original_dst: SocketAddr,
+}
+
+#[pymethods]
+impl UdpStream {
+    /// Fake `read` method for UDP streams. This method will return the payload of the initially
+    /// received UDP packet, and zero bytes after this payload has been read.
+    fn read<'p>(&mut self, py: Python<'p>, n: u32) -> PyResult<&'p PyAny> {
+        let bytes = if self.data.is_empty() {
+            // entire payload has been read, return zero bytes
+            Vec::new()
+        } else {
+            // drain at most all remaining elements but never more
+            let l = self.data.len();
+            let r = if n as usize >= l { l } else { n as usize };
+            self.data.drain(0..r).collect()
+        };
+
+        pyo3_asyncio::tokio::future_into_py(py, async move {
+            let bytes: Py<PyBytes> = Python::with_gil(|py| PyBytes::new(py, &bytes).into_py(py));
+            Ok(bytes)
+        })
+    }
+
+    /// Send a UDP datagram with data as a response to a datagram that was previously received.
+    fn write(&self, data: Vec<u8>) -> PyResult<()> {
+        self.event_tx
+            .send(TransportCommand::SendDatagram {
+                data,
+                src_addr: self.original_dst,
+                dst_addr: self.peername,
+            })
+            .map_err(event_queue_unavailable)?;
+
+        Ok(())
+    }
+
+    /// Fake `drain` method for UDP streams. This method returns immediately, since UDP packets
+    /// are either sent immediately or are dropped if the send queue is full.
+    fn drain<'p>(&self, py: Python<'p>) -> PyResult<&'p PyAny> {
+        pyo3_asyncio::tokio::future_into_py(py, async move { Ok(()) })
+    }
+
+    /// Fake `write_eof` method for UDP streams. This method does nothing, since UDP is not
+    /// connection-based.
+    fn write_eof(&self) -> PyResult<()> {
+        Ok(())
+    }
+
+    /// Fake `close` method for UDP streams. This method returns immediately since there is no
+    /// connection to close.
+    fn close(&self) -> PyResult<()> {
+        Ok(())
+    }
+
+    /// Query the UDP stream for details of the underlying network connection.
+    ///
+    /// Supported values: `peername`, `sockname`, `original_dst`.
+    fn get_extra_info(&self, py: Python, name: String) -> PyResult<PyObject> {
+        match name.as_str() {
+            "peername" => Ok(socketaddr_to_py(py, self.peername)),
+            "sockname" => Ok(socketaddr_to_py(py, self.sockname)),
+            "original_dst" => Ok(socketaddr_to_py(py, self.original_dst)),
+            _ => Err(PyKeyError::new_err(name)),
+        }
+    }
+
+    fn __repr__(&self) -> String {
+        format!(
+            "UdpStream(peer={}, sock={}, dst={})",
+            self.peername, self.sockname, self.original_dst
+        )
+    }
+}
+
+impl Drop for UdpStream {
+    fn drop(&mut self) {}
+}
+
+
 pub fn socketaddr_to_py(py: Python, s: SocketAddr) -> PyObject {
     match s {
         SocketAddr::V4(addr) => (addr.ip().to_string(), addr.port()).into_py(py),
@@ -235,18 +323,28 @@ impl PyInteropTask {
                                 src_addr,
                                 dst_addr,
                             } => {
+                                let stream = UdpStream {
+                                    data,
+                                    event_tx: self.py_to_smol_tx.clone(),
+                                    peername: src_addr,
+                                    sockname: self.local_addr,
+                                    original_dst: dst_addr,
+                                };
+
                                 Python::with_gil(|py| {
-                                    let bytes: Py<PyBytes> = PyBytes::new(py, &data).into_py(py);
+                                    let stream = stream.into_py(py);
 
-                                    if let Err(err) = self.py_loop.call_method1(
+                                    let coro = match self.py_udp_handler.call1(py, (stream,)) {
+                                        Ok(coro) => coro,
+                                        Err(err) => {
+                                            err.print(py);
+                                            return;
+                                        },
+                                    };
+
+                                    if let Err(err) = self.run_coroutine_threadsafe.call1(
                                         py,
-                                        "call_soon_threadsafe",
-                                        (
-                                            self.py_udp_handler.as_ref(py),
-                                            bytes,
-                                            socketaddr_to_py(py, src_addr),
-                                            socketaddr_to_py(py, dst_addr),
-                                        ),
+                                        (coro, self.py_loop.as_ref(py))
                                     ) {
                                         err.print(py);
                                     }


Test patch:
diff --git a/echo_test_server.py b/echo_test_server.py
index 11f0db5..1ce345a 100644
--- a/echo_test_server.py
+++ b/echo_test_server.py
@@ -34,13 +34,8 @@
 
 
 async def main():
-    def receive_datagram(data, src_addr, dst_addr):
-        logger.debug(f"Received datagram: {data=} {src_addr=} {dst_addr=}")
-        server.send_datagram(data.upper(), dst_addr, src_addr)
-        logger.debug("Echoed datagram.")
-
     conf = mitmproxy_wireguard.WireGuardServerConf.custom(51820, server_keypair[0], [client_keypair[1]])
-    server = await mitmproxy_wireguard.start_server("0.0.0.0", conf, handle_connection, receive_datagram)
+    server = await mitmproxy_wireguard.start_server("0.0.0.0", conf, handle_tcp, handle_udp)
 
     print(
         textwrap.dedent(
@@ -72,7 +67,7 @@ def stop(*_):
     await server.wait_closed()
 
 
-async def handle_connection(rw: mitmproxy_wireguard.TcpStream):
+async def handle_tcp(rw: mitmproxy_wireguard.TcpStream):
     logger.debug(f"connection task {rw=}")
     logger.debug(f"{rw.get_extra_info('peername')=}")
 
@@ -107,5 +102,23 @@ async def handle_connection(rw: mitmproxy_wireguard.TcpStream):
     logger.debug("closed.")
 
 
+async def handle_udp(rw: mitmproxy_wireguard.UdpStream):
+    logger.debug(f"datagram task {rw=}")
+    logger.debug(f"{rw.get_extra_info('peername')=}")
+
+    try:
+        data = await rw.read(4096)
+    except Exception as exc:
+        logger.debug(f"read {exc=}")
+        data = b""
+    logger.debug(f"read complete. writing... {len(data)=} {data[:10]=} ")
+
+    try:
+        rw.write(data.upper())
+    except Exception as exc:
+        logger.debug(f"write {exc=}")
+    logger.debug("write complete.")
+
+
 if __name__ == "__main__":
     asyncio.run(main(), debug=True)