improve shutdown dance

mitmproxy-linux/src/main2.rs

@@ -16,6 +16,7 @@ use mitmproxy::packet_sources::tun::create_tun_device;
 use tun::AbstractDevice;
 use prost::Message;
 use tokio::io::AsyncReadExt;
+use tokio::signal::unix::{signal, SignalKind};
 use mitmproxy::ipc::{PacketWithMeta, from_proxy};
 use mitmproxy::ipc::FromProxy;
 use mitmproxy::packet_sources::IPC_BUF_SIZE;
@@ -89,6 +90,17 @@ async fn main() -> anyhow::Result<()> {
     fs::set_permissions(&redirector_addr, Permissions::from_mode(0o777))?;
     println!("{}", redirector_addr.to_string_lossy());
 
+    // Exit cleanly on SIGINT/SIGTERM
+    tokio::spawn(async {
+        let mut sigint = signal(SignalKind::interrupt()).context("failed to register SIGINT listener").unwrap();
+        let mut sigterm = signal(SignalKind::terminate()).context("failed to register SIGTERM listener").unwrap();
+        select! {
+            _ = sigint.recv() => (),
+            _ = sigterm.recv() => (),
+        }
+        std::process::exit(0);
+    });
+
     let mut ipc_buf = BytesMut::with_capacity(IPC_BUF_SIZE);
     let mut dev_buf = BytesMut::with_capacity(IPC_BUF_SIZE);
 

src/packet_sources/linux.rs

@@ -21,7 +21,11 @@ use tokio::process::Command;
 use tokio::time::timeout;
 use crate::shutdown;
 
-async fn start_redirector(executable: &Path, listener_addr: &Path) -> Result<PathBuf> {
+async fn start_redirector(
+    executable: &Path,
+    listener_addr: &Path,
+    shutdown: shutdown::Receiver,
+) -> Result<PathBuf> {
     debug!("Elevating privileges...");
     // Try to elevate privileges using a dummy sudo invocation.
     // The idea here is to block execution and give the user time to enter their password.
@@ -50,10 +54,17 @@ async fn start_redirector(executable: &Path, listener_addr: &Path) -> Result<Pat
 
     let stdout = redirector_process.stdout.take().unwrap();
     let stderr = redirector_process.stderr.take().unwrap();
+    let shutdown2 = shutdown.clone();
     tokio::spawn(async move {
         let mut stderr = BufReader::new(stderr).lines();
         let mut level = Level::Error;
         while let Ok(Some(line)) = stderr.next_line().await {
+            if shutdown2.is_shutting_down() {
+                // We don't want to log during exit, https://github.com/vorner/pyo3-log/issues/30
+                eprintln!("{}", line);
+                continue
+            }
+
             let new_level = line
                 .strip_prefix("[")
                 .and_then(|s| s.split_once(" "))
@@ -73,10 +84,18 @@ async fn start_redirector(executable: &Path, listener_addr: &Path) -> Result<Pat
     tokio::spawn(async move {
         match redirector_process.wait().await {
             Ok(status) if status.success() => {
-                debug!("[linux-redirector] exited successfully.")
+                if shutdown.is_shutting_down() {
+                    // We don't want to log during exit, https://github.com/vorner/pyo3-log/issues/30
+                } else {
+                    debug!("[linux-redirector] exited successfully.")
+                }
             }
             other => {
-                error!("[linux-redirector] exited: {:?}", other)
+                if shutdown.is_shutting_down() {
+                    eprintln!("[linux-redirector] exited during shutdown: {:?}", other)
+                } else {
+                    error!("[linux-redirector] exited: {:?}", other)
+                }
             }
         }
     });
@@ -150,7 +169,7 @@ impl PacketSourceConf for LinuxConf {
         let datagram_dir = tempdir().context("failed to create temp dir")?;
 
         let channel = UnixDatagram::bind(datagram_dir.path().join("mitmproxy"))?;
-        let dst = start_redirector(&self.executable_path, datagram_dir.path()).await?;
+        let dst = start_redirector(&self.executable_path, datagram_dir.path(), shutdown.clone()).await?;
 
         channel
             .connect(&dst)