diff --git a/CHANGELOG.md b/CHANGELOG.md index 2a5f7dbc..05df9c69 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ ## Unreleased: mitmproxy_rs next +- `Stream.wait_closed()` now waits for outgoing UDP client tasks to finish. - Preserve OS error details when UDP and WireGuard servers fail to start. - `Stream.is_closing()` now returns `True` once the server or UDP client task has shut down. diff --git a/mitmproxy-rs/src/stream.rs b/mitmproxy-rs/src/stream.rs index cfe32efe..eff3c837 100644 --- a/mitmproxy-rs/src/stream.rs +++ b/mitmproxy-rs/src/stream.rs @@ -6,6 +6,7 @@ use pyo3::{IntoPyObjectExt, exceptions::PyOSError, intern, prelude::*}; use tokio::sync::{ mpsc::{self}, oneshot::{self}, + watch, }; use mitmproxy::messages::{ConnectionId, TransportCommand, TunnelInfo}; @@ -31,6 +32,8 @@ pub struct Stream { pub peername: SocketAddr, pub sockname: SocketAddr, pub tunnel_info: TunnelInfo, + /// For UDP client streams, closed when the underlying task has finished. + pub task_done: Option>, } #[pymethods] @@ -135,9 +138,17 @@ impl Stream { } } - /// Wait until the stream is closed (currently a no-op). + /// Wait until the stream is closed. + /// For UDP client streams, this waits for the underlying task to finish. fn wait_closed<'py>(&self, py: Python<'py>) -> PyResult> { - pyo3_async_runtimes::tokio::future_into_py(py, std::future::ready(Ok(()))) + let task_done = self.task_done.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + if let Some(mut task_done) = task_done { + // Errors once the sender is dropped by the finished task. + task_done.changed().await.ok(); + } + Ok(()) + }) } /// Query the stream for details of the underlying network connection. @@ -211,6 +222,45 @@ impl Drop for Stream { #[cfg(test)] mod tests { use super::*; + use std::time::Duration; + + #[test] + fn wait_closed_waits_for_udp_task() -> PyResult<()> { + Python::initialize(); + Python::attach(|py| { + pyo3_async_runtimes::tokio::run(py, async { + let (command_tx, _command_rx) = mpsc::unbounded_channel(); + let (task_done_tx, task_done) = watch::channel(()); + let stream = Stream { + connection_id: ConnectionId::unassigned_udp(), + state: StreamState::Closed, + command_tx, + peername: "127.0.0.1:1234".parse()?, + sockname: "127.0.0.1:4321".parse()?, + tunnel_info: TunnelInfo::None, + task_done: Some(task_done), + }; + let wait_closed = Python::attach(|py| { + let stream = Py::new(py, stream)?; + pyo3_async_runtimes::tokio::into_future( + stream.bind(py).call_method0("wait_closed")?, + ) + })?; + tokio::pin!(wait_closed); + assert!( + tokio::time::timeout(Duration::from_millis(20), &mut wait_closed) + .await + .is_err(), + "wait_closed returned before the task finished" + ); + drop(task_done_tx); + tokio::time::timeout(Duration::from_secs(1), wait_closed) + .await + .expect("wait_closed did not return after the task finished")?; + Ok(()) + }) + }) + } #[test] fn is_closing_when_server_shut_down() { @@ -222,6 +272,7 @@ mod tests { peername: "127.0.0.1:1234".parse().unwrap(), sockname: "127.0.0.1:4321".parse().unwrap(), tunnel_info: TunnelInfo::None, + task_done: None, }; assert!(!stream.is_closing()); drop(command_rx); diff --git a/mitmproxy-rs/src/task.rs b/mitmproxy-rs/src/task.rs index ddf95bde..ae4bb5b5 100644 --- a/mitmproxy-rs/src/task.rs +++ b/mitmproxy-rs/src/task.rs @@ -78,6 +78,7 @@ impl PyInteropTask { peername: src_addr, sockname: dst_addr, tunnel_info, + task_done: None, }; let mut conns = active_streams.lock().await; diff --git a/mitmproxy-rs/src/udp_client.rs b/mitmproxy-rs/src/udp_client.rs index 6c04e35d..0e25987f 100644 --- a/mitmproxy-rs/src/udp_client.rs +++ b/mitmproxy-rs/src/udp_client.rs @@ -5,7 +5,7 @@ use anyhow::{Result, anyhow}; use pyo3::prelude::*; use tokio::net::{UdpSocket, lookup_host}; use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel}; -use tokio::sync::oneshot; +use tokio::sync::{oneshot, watch}; use crate::stream::{Stream, StreamState}; use mitmproxy::MAX_PACKET_SIZE; @@ -32,8 +32,9 @@ pub fn open_udp_connection( let sockname = socket.local_addr()?; let (command_tx, command_rx) = unbounded_channel(); + let (task_done_tx, task_done) = watch::channel(()); - tokio::spawn(async { + tokio::spawn(async move { let task = UdpClientTask { socket, transport_commands_rx: command_rx, @@ -41,6 +42,7 @@ pub fn open_udp_connection( if let Err(e) = task.run().await { log::error!("UDP client errored: {e}"); } + drop(task_done_tx); }); let stream = Stream { @@ -50,6 +52,7 @@ pub fn open_udp_connection( peername, sockname, tunnel_info: TunnelInfo::None, + task_done: Some(task_done), }; Ok(stream)