Move ConnChecker to helpers.rs

test/test-manager/src/tests/helpers.rs

@@ -2,6 +2,7 @@ use super::{config::TEST_CONFIG, Error, TestContext, WAIT_FOR_TUNNEL_STATE_TIMEO
 use crate::network_monitor::{
     self, start_packet_monitor, MonitorOptions, MonitorUnexpectedlyStopped, PacketMonitor,
 };
+use anyhow::{anyhow, bail, ensure, Context};
 use futures::StreamExt;
 use mullvad_management_interface::{client::DaemonEvent, MullvadProxyClient};
 use mullvad_types::{
@@ -13,6 +14,7 @@ use mullvad_types::{
     relay_list::{Relay, RelayList},
     states::TunnelState,
 };
+use pcap::Direction;
 use pnet_packet::ip::IpNextHeaderProtocols;
 use std::{
     collections::HashMap,
@@ -21,10 +23,25 @@ use std::{
     time::Duration,
 };
 use talpid_types::net::wireguard::{PeerConfig, PrivateKey, TunnelConfig};
-use test_rpc::{package::Package, AmIMullvad, ServiceClient};
+use test_rpc::{meta::Os, package::Package, AmIMullvad, ServiceClient, SpawnOpts};
+use tokio::time::sleep;
 
 pub const THROTTLE_RETRY_DELAY: Duration = Duration::from_secs(120);
 
+const CHECKER_FILENAME_WINDOWS: &str = "connection-checker.exe";
+const CHECKER_FILENAME_UNIX: &str = "connection-checker";
+
+const AM_I_MULLVAD_TIMEOUT_MS: u64 = 10000;
+const LEAK_TIMEOUT_MS: u64 = 500;
+
+/// Timeout of [ConnCheckerHandle::check_connection].
+const CONN_CHECKER_TIMEOUT: Duration = Duration::from_millis(
+    AM_I_MULLVAD_TIMEOUT_MS // https://am.i.mullvad.net timeout
+    + LEAK_TIMEOUT_MS // leak-tcp timeout
+    + LEAK_TIMEOUT_MS // leak-icmp timeout
+    + 1000, // plus some extra grace time
+);
+
 #[macro_export]
 macro_rules! assert_tunnel_state {
     ($mullvad_client:expr, $pattern:pat) => {{
@@ -231,12 +248,14 @@ pub async fn login_with_retries(
 /// Try to connect to a Mullvad Tunnel.
 ///
 /// # Returns
-/// - `Result::Ok` if the daemon successfully connected to a tunnel
+/// - `Result::Ok(new_state)` if the daemon successfully connected to a tunnel
 /// - `Result::Err` if:
 ///     - The daemon failed to even begin connecting. Then [`Error::Rpc`] is returned.
 ///     - The daemon started to connect but ended up in the [`TunnelState::Error`] state.
 ///     Then [`Error::UnexpectedErrorState`] is returned
-pub async fn connect_and_wait(mullvad_client: &mut MullvadProxyClient) -> Result<(), Error> {
+pub async fn connect_and_wait(
+    mullvad_client: &mut MullvadProxyClient,
+) -> Result<TunnelState, Error> {
     log::info!("Connecting");
 
     mullvad_client.connect_tunnel().await?;
@@ -254,7 +273,7 @@ pub async fn connect_and_wait(mullvad_client: &mut MullvadProxyClient) -> Result
 
     log::info!("Connected");
 
-    Ok(())
+    Ok(new_state)
 }
 
 pub async fn disconnect_and_wait(mullvad_client: &mut MullvadProxyClient) -> Result<(), Error> {
@@ -661,3 +680,285 @@ impl PingerBuilder {
         self
     }
 }
+
+/// This helper spawns a seperate process which checks if we are connected to Mullvad, and tries to
+/// leak traffic outside the tunnel by sending TCP, UDP, and ICMP packets to [LEAK_DESTINATION].
+pub struct ConnChecker {
+    rpc: ServiceClient,
+    mullvad_client: MullvadProxyClient,
+    leak_destination: SocketAddr,
+
+    /// Path to the process binary.
+    executable_path: String,
+
+    /// Whether the process should be split when spawned. Needed on Linux.
+    split: bool,
+}
+
+pub struct ConnCheckerHandle<'a> {
+    checker: &'a mut ConnChecker,
+
+    /// ID of the spawned process.
+    pid: u32,
+}
+
+pub struct ConnectionStatus {
+    /// True if <https://am.i.mullvad.net/> reported we are connected.
+    am_i_mullvad: bool,
+
+    /// True if we sniffed TCP packets going outside the tunnel.
+    leaked_tcp: bool,
+
+    /// True if we sniffed UDP packets going outside the tunnel.
+    leaked_udp: bool,
+
+    /// True if we sniffed ICMP packets going outside the tunnel.
+    leaked_icmp: bool,
+}
+
+impl ConnChecker {
+    pub fn new(
+        rpc: ServiceClient,
+        mullvad_client: MullvadProxyClient,
+        leak_destination: SocketAddr,
+    ) -> Self {
+        let artifacts_dir = &TEST_CONFIG.artifacts_dir;
+        let executable_path = match TEST_CONFIG.os {
+            Os::Linux | Os::Macos => format!("{artifacts_dir}/{CHECKER_FILENAME_UNIX}"),
+            Os::Windows => format!("{artifacts_dir}\\{CHECKER_FILENAME_WINDOWS}"),
+        };
+
+        Self {
+            rpc,
+            mullvad_client,
+            leak_destination,
+            split: false,
+            executable_path,
+        }
+    }
+
+    /// Spawn the connecton checker process and return a handle to it.
+    ///
+    /// Dropping the handle will stop the process.
+    /// **NOTE**: The handle must be dropped from a tokio runtime context.
+    pub async fn spawn(&mut self) -> anyhow::Result<ConnCheckerHandle<'_>> {
+        log::debug!("spawning connection checker");
+
+        let opts = SpawnOpts {
+            attach_stdin: true,
+            attach_stdout: true,
+            args: [
+                "--interactive",
+                "--timeout",
+                &AM_I_MULLVAD_TIMEOUT_MS.to_string(),
+                // try to leak traffic to LEAK_DESTINATION
+                "--leak",
+                &self.leak_destination.to_string(),
+                //TODO(markus): Remove
+                //&LEAK_DESTINATION.to_string(),
+                "--leak-timeout",
+                &LEAK_TIMEOUT_MS.to_string(),
+                "--leak-tcp",
+                "--leak-udp",
+                "--leak-icmp",
+            ]
+            .map(String::from)
+            .to_vec(),
+            ..SpawnOpts::new(&self.executable_path)
+        };
+
+        let pid = self.rpc.spawn(opts).await?;
+
+        if self.split && TEST_CONFIG.os == Os::Linux {
+            self.mullvad_client
+                .add_split_tunnel_process(pid as i32)
+                .await?;
+        }
+
+        Ok(ConnCheckerHandle { pid, checker: self })
+    }
+
+    /// Enable split tunneling for the connection checker.
+    pub async fn split(&mut self) -> anyhow::Result<()> {
+        log::debug!("enable split tunnel");
+        self.split = true;
+
+        match TEST_CONFIG.os {
+            Os::Linux => { /* linux programs can't be split until they are spawned */ }
+            Os::Windows => {
+                self.mullvad_client
+                    .add_split_tunnel_app(&self.executable_path)
+                    .await?;
+                self.mullvad_client.set_split_tunnel_state(true).await?;
+            }
+            Os::Macos => unimplemented!("MacOS"),
+        }
+
+        Ok(())
+    }
+
+    /// Disable split tunneling for the connection checker.
+    pub async fn unsplit(&mut self) -> anyhow::Result<()> {
+        log::debug!("disable split tunnel");
+        self.split = false;
+
+        match TEST_CONFIG.os {
+            Os::Linux => {}
+            Os::Windows => {
+                self.mullvad_client.set_split_tunnel_state(false).await?;
+                self.mullvad_client
+                    .remove_split_tunnel_app(&self.executable_path)
+                    .await?;
+            }
+            Os::Macos => unimplemented!("MacOS"),
+        }
+
+        Ok(())
+    }
+}
+
+impl ConnCheckerHandle<'_> {
+    pub async fn split(&mut self) -> anyhow::Result<()> {
+        if TEST_CONFIG.os == Os::Linux {
+            self.checker
+                .mullvad_client
+                .add_split_tunnel_process(self.pid as i32)
+                .await?;
+        }
+
+        self.checker.split().await
+    }
+
+    pub async fn unsplit(&mut self) -> anyhow::Result<()> {
+        if TEST_CONFIG.os == Os::Linux {
+            self.checker
+                .mullvad_client
+                .remove_split_tunnel_process(self.pid as i32)
+                .await?;
+        }
+
+        self.checker.unsplit().await
+    }
+
+    /// Assert that traffic is flowing through the Mullvad tunnel and that no packets are leaked.
+    pub async fn assert_secure(&mut self) -> anyhow::Result<()> {
+        log::info!("checking that connection is secure");
+        let status = self.check_connection().await?;
+        ensure!(status.am_i_mullvad);
+        ensure!(!status.leaked_tcp);
+        ensure!(!status.leaked_udp);
+        ensure!(!status.leaked_icmp);
+
+        Ok(())
+    }
+
+    /// Assert that traffic is NOT flowing through the Mullvad tunnel and that packets ARE leaked.
+    pub async fn assert_insecure(&mut self) -> anyhow::Result<()> {
+        log::info!("checking that connection is not secure");
+        let status = self.check_connection().await?;
+        ensure!(!status.am_i_mullvad);
+        ensure!(status.leaked_tcp);
+        ensure!(status.leaked_udp);
+        ensure!(status.leaked_icmp);
+
+        Ok(())
+    }
+
+    pub async fn check_connection(&mut self) -> anyhow::Result<ConnectionStatus> {
+        // Monitor all pakets going to LEAK_DESTINATION during the check.
+        let leak_destination = self.checker.leak_destination;
+        let monitor = start_packet_monitor(
+            move |packet| packet.destination.ip() == leak_destination.ip(),
+            MonitorOptions {
+                direction: Some(Direction::In),
+                ..MonitorOptions::default()
+            },
+        )
+        .await;
+
+        // Write a newline to the connection checker to prompt it to perform the check.
+        self.checker
+            .rpc
+            .write_child_stdin(self.pid, "Say the line, Bart!\r\n".into())
+            .await?;
+
+        // The checker responds when the check is complete.
+        let line = self.read_stdout_line().await?;
+
+        let monitor_result = monitor
+            .into_result()
+            .await
+            .map_err(|_e| anyhow!("Packet monitor unexpectedly stopped"))?;
+
+        Ok(ConnectionStatus {
+            am_i_mullvad: parse_am_i_mullvad(line)?,
+
+            leaked_tcp: (monitor_result.packets.iter())
+                .any(|pkt| pkt.protocol == IpNextHeaderProtocols::Tcp),
+
+            leaked_udp: (monitor_result.packets.iter())
+                .any(|pkt| pkt.protocol == IpNextHeaderProtocols::Udp),
+
+            leaked_icmp: (monitor_result.packets.iter())
+                .any(|pkt| pkt.protocol == IpNextHeaderProtocols::Icmp),
+        })
+    }
+
+    /// Try to a single line of output from the spawned process
+    async fn read_stdout_line(&mut self) -> anyhow::Result<String> {
+        // Add a timeout to avoid waiting forever.
+        tokio::time::timeout(CONN_CHECKER_TIMEOUT, async {
+            let mut line = String::new();
+
+            // tarpc doesn't support streams, so we poll the checker process in a loop instead
+            loop {
+                let Some(output) = self.checker.rpc.read_child_stdout(self.pid).await? else {
+                    bail!("got EOF from connection checker process");
+                };
+
+                if output.is_empty() {
+                    sleep(Duration::from_millis(500)).await;
+                    continue;
+                }
+
+                line.push_str(&output);
+
+                if line.contains('\n') {
+                    log::info!("output from child process: {output:?}");
+                    return Ok(line);
+                }
+            }
+        })
+        .await
+        .with_context(|| "Timeout reading stdout from connection checker")?
+    }
+}
+
+impl Drop for ConnCheckerHandle<'_> {
+    fn drop(&mut self) {
+        let rpc = self.checker.rpc.clone();
+        let pid = self.pid;
+
+        let Ok(runtime_handle) = tokio::runtime::Handle::try_current() else {
+            log::error!("ConnCheckerHandle dropped outside of a tokio runtime.");
+            return;
+        };
+
+        runtime_handle.spawn(async move {
+            // Make sure child process is stopped when this handle is dropped.
+            // Closing stdin does the trick.
+            let _ = rpc.close_child_stdin(pid).await;
+        });
+    }
+}
+
+/// Parse output from connection-checker. Returns true if connected to Mullvad.
+fn parse_am_i_mullvad(result: String) -> anyhow::Result<bool> {
+    Ok(if result.contains("You are connected") {
+        true
+    } else if result.contains("You are not connected") {
+        false
+    } else {
+        bail!("Unexpected output from connection-checker: {result:?}")
+    })
+}