Merge branch 'add-test-for-2020-app-audit-des-419'

9 files changed, 501 insertions, 389 deletions

diff --git a/test/connection-checker/src/cli.rs b/test/connection-checker/src/cli.rs
index dddb348b25..ecc33dd7e9 100644
--- a/test/connection-checker/src/cli.rs
+++ b/test/connection-checker/src/cli.rs
@@ -33,4 +33,8 @@ pub struct Opt {
     /// Timeout for leak check network connections (in millis).
     #[clap(long, default_value = "1000")]
     pub leak_timeout: u64,
+
+    /// Junk data for each UDP and TCP packet
+    #[clap(long, requires = "leak", default_value = "Hello there!")]
+    pub payload: String,
 }
diff --git a/test/connection-checker/src/net.rs b/test/connection-checker/src/net.rs
index 2e17423933..40db99e8b5 100644
--- a/test/connection-checker/src/net.rs
+++ b/test/connection-checker/src/net.rs
@@ -31,13 +31,13 @@ pub fn send_tcp(opt: &Opt, destination: SocketAddr) -> eyre::Result<()> {
 
     let mut stream = std::net::TcpStream::from(sock);
     stream
-        .write_all(b"hello there")
+        .write_all(opt.payload.as_bytes())
         .wrap_err(eyre!("Failed to send message to {destination}"))?;
 
     Ok(())
 }
 
-pub fn send_udp(_opt: &Opt, destination: SocketAddr) -> Result<(), eyre::Error> {
+pub fn send_udp(opt: &Opt, destination: SocketAddr) -> Result<(), eyre::Error> {
     let bind_addr: SocketAddr = SocketAddr::new(Ipv4Addr::new(0, 0, 0, 0).into(), 0);
 
     eprintln!("Leaking UDP packets to {destination}");
@@ -52,11 +52,9 @@ pub fn send_udp(_opt: &Opt, destination: SocketAddr) -> Result<(), eyre::Error>
     sock.bind(&socket2::SockAddr::from(bind_addr))
         .wrap_err(eyre!("Failed to bind UDP socket to {bind_addr}"))?;
 
-    // log::debug!("Send message from {bind_addr} to {destination}/UDP");
-
     let std_socket = std::net::UdpSocket::from(sock);
     std_socket
-        .send_to(b"Hello there!", destination)
+        .send_to(opt.payload.as_bytes(), destination)
         .wrap_err(eyre!("Failed to send message to {destination}"))?;
 
     Ok(())
diff --git a/test/test-manager/src/network_monitor.rs b/test/test-manager/src/network_monitor.rs
index 87a8193e29..ddd35edc7e 100644
--- a/test/test-manager/src/network_monitor.rs
+++ b/test/test-manager/src/network_monitor.rs
@@ -25,6 +25,7 @@ pub struct ParsedPacket {
     pub source: SocketAddr,
     pub destination: SocketAddr,
     pub protocol: IpNextHeaderProtocol,
+    pub payload: Vec<u8>,
 }
 
 impl PacketCodec for Codec {
@@ -74,9 +75,9 @@ impl Codec {
 
         let mut source = SocketAddr::new(IpAddr::V4(packet.get_source()), 0);
         let mut destination = SocketAddr::new(IpAddr::V4(packet.get_destination()), 0);
+        let mut payload = vec![];
 
         let protocol = packet.get_next_level_protocol();
-
         match protocol {
             IpHeaderProtocols::Tcp => {
                 let seg = TcpPacket::new(packet.payload()).or_else(|| {
@@ -85,6 +86,7 @@ impl Codec {
                 })?;
                 source.set_port(seg.get_source());
                 destination.set_port(seg.get_destination());
+                payload = seg.payload().to_vec();
             }
             IpHeaderProtocols::Udp => {
                 let seg = UdpPacket::new(packet.payload()).or_else(|| {
@@ -93,6 +95,7 @@ impl Codec {
                 })?;
                 source.set_port(seg.get_source());
                 destination.set_port(seg.get_destination());
+                payload = seg.payload().to_vec();
             }
             IpHeaderProtocols::Icmp => {}
             proto => log::debug!("ignoring v4 packet, transport/protocol type {proto}"),
@@ -102,6 +105,7 @@ impl Codec {
             source,
             destination,
             protocol,
+            payload,
         })
     }
 
@@ -113,6 +117,7 @@ impl Codec {
 
         let mut source = SocketAddr::new(IpAddr::V6(packet.get_source()), 0);
         let mut destination = SocketAddr::new(IpAddr::V6(packet.get_destination()), 0);
+        let mut payload = vec![];
 
         let protocol = packet.get_next_header();
         match protocol {
@@ -123,6 +128,7 @@ impl Codec {
                 })?;
                 source.set_port(seg.get_source());
                 destination.set_port(seg.get_destination());
+                payload = seg.payload().to_vec();
             }
             IpHeaderProtocols::Udp => {
                 let seg = UdpPacket::new(packet.payload()).or_else(|| {
@@ -131,6 +137,7 @@ impl Codec {
                 })?;
                 source.set_port(seg.get_source());
                 destination.set_port(seg.get_destination());
+                payload = seg.payload().to_vec();
             }
             IpHeaderProtocols::Icmpv6 => {}
             proto => log::debug!("ignoring v6 packet, transport/protocol type {proto}"),
@@ -140,12 +147,14 @@ impl Codec {
             source,
             destination,
             protocol,
+            payload,
         })
     }
 }
 
-#[derive(Debug)]
-pub struct MonitorUnexpectedlyStopped(());
+#[derive(Debug, thiserror::Error)]
+#[error("Packet monitor stopped unexpectedly")]
+pub struct MonitorUnexpectedlyStopped;
 
 pub struct PacketMonitor {
     handle: tokio::task::JoinHandle<Result<MonitorResult, MonitorUnexpectedlyStopped>>,
@@ -297,7 +306,7 @@ async fn start_packet_monitor_for_interface(
                         }
                         _ => {
                             log::error!("lost packet stream");
-                            break Err(MonitorUnexpectedlyStopped(()));
+                            break Err(MonitorUnexpectedlyStopped);
                         }
                     }
                 }
diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index 3ef42751b2..86cbecb50c 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -2,17 +2,23 @@ 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_relay_selector::{
+    query::RelayQuery, GetRelay, RelaySelector, SelectorConfig, WireguardConfig,
+};
 use mullvad_types::{
     constraints::Constraint,
     location::Location,
     relay_constraints::{
-        BridgeSettings, GeographicLocationConstraint, LocationConstraint, RelaySettings,
+        BridgeSettings, GeographicLocationConstraint, LocationConstraint, RelayConstraints,
+        RelaySettings,
     },
-    relay_list::{Relay, RelayList},
+    relay_list::Relay,
     states::TunnelState,
 };
+use pcap::Direction;
 use pnet_packet::ip::IpNextHeaderProtocols;
 use std::{
     collections::HashMap,
@@ -21,10 +27,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 +252,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 +277,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> {
@@ -480,27 +503,49 @@ pub fn get_app_env() -> HashMap<String, String> {
     ])
 }
 
-/// Return a filtered version of the daemon's relay list.
+/// Constrain the daemon to only select the relay selected with `query` when establishing all
+/// future tunnels (until relay settings are updated, see [`set_relay_settings`]). Returns the
+/// selected [`Relay`] for future reference.
 ///
-/// * `mullvad_client` - An interface to the Mullvad daemon.
-/// * `critera` - A function used to determine which relays to return.
-pub async fn filter_relays<Filter>(
+/// # Note
+/// This function does not handle bridges and multihop configurations (currently). There is no
+/// particular reason for this other than it not being needed at the time, so feel free to extend this
+/// function :).
+pub async fn constrain_to_relay(
     mullvad_client: &mut MullvadProxyClient,
-    criteria: Filter,
-) -> Result<Vec<Relay>, Error>
-where
-    Filter: Fn(&Relay) -> bool,
-{
-    let relay_list: RelayList = mullvad_client
-        .get_relay_locations()
-        .await
-        .map_err(|error| Error::Daemon(format!("Failed to obtain relay list: {}", error)))?;
+    query: RelayQuery,
+) -> anyhow::Result<Relay> {
+    /// Convert the result of invoking the relay selector to a relay constraint.
+    fn convert_to_relay_constraints(
+        selected_relay: GetRelay,
+    ) -> anyhow::Result<(Relay, RelayConstraints)> {
+        match selected_relay {
+            GetRelay::Wireguard {
+                inner: WireguardConfig::Singlehop { exit },
+                ..
+            }
+            | GetRelay::OpenVpn { exit, .. } => {
+                let location = into_constraint(&exit)?;
+                let relay_constraints = RelayConstraints {
+                    location,
+                    ..Default::default()
+                };
+                Ok((exit, relay_constraints))
+            }
+            unsupported => bail!("Can not constrain to a {unsupported:?}"),
+        }
+    }
 
-    Ok(relay_list
-        .relays()
-        .filter(|relay| criteria(relay))
-        .cloned()
-        .collect())
+    // Construct a relay selector with up-to-date information from the runnin daemon's relay list
+    let relay_list = mullvad_client.get_relay_locations().await?;
+    let relay_selector = RelaySelector::from_list(SelectorConfig::default(), relay_list);
+    // Select an(y) appropriate relay for the given query and constrain the daemon to only connect
+    // to that specific relay (when connecting).
+    let relay = relay_selector.get_relay_by_query(query)?;
+    let (exit, relay_constraints) = convert_to_relay_constraints(relay)?;
+    set_relay_settings(mullvad_client, RelaySettings::Normal(relay_constraints)).await?;
+
+    Ok(exit)
 }
 
 /// Convenience function for constructing a constraint from a given [`Relay`].
@@ -508,7 +553,7 @@ where
 /// # Panics
 ///
 /// The relay must have a location set.
-pub fn into_constraint(relay: &Relay) -> Constraint<LocationConstraint> {
+pub fn into_constraint(relay: &Relay) -> anyhow::Result<Constraint<LocationConstraint>> {
     relay
         .location
         .as_ref()
@@ -518,16 +563,12 @@ pub fn into_constraint(relay: &Relay) -> Constraint<LocationConstraint> {
                  city_code,
                  ..
              }| {
-                GeographicLocationConstraint::Hostname(
-                    country_code.to_string(),
-                    city_code.to_string(),
-                    relay.hostname.to_string(),
-                )
+                GeographicLocationConstraint::hostname(country_code, city_code, &relay.hostname)
             },
         )
         .map(LocationConstraint::Location)
         .map(Constraint::Only)
-        .expect("relay is missing location")
+        .ok_or(anyhow!("relay is missing location"))
 }
 
 /// Ping monitoring made easy!
@@ -661,3 +702,301 @@ 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,
+
+    /// Some arbitrary payload
+    payload: Option<String>,
+}
+
+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,
+            payload: None,
+        }
+    }
+
+    /// Set a custom magic payload that the connection checker binary should use when leak-testing.
+    pub fn payload(&mut self, payload: impl Into<String>) {
+        self.payload = Some(payload.into())
+    }
+
+    /// 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 = {
+            let mut args = [
+                "--interactive",
+                "--timeout",
+                &AM_I_MULLVAD_TIMEOUT_MS.to_string(),
+                // try to leak traffic to LEAK_DESTINATION
+                "--leak",
+                &self.leak_destination.to_string(),
+                "--leak-timeout",
+                &LEAK_TIMEOUT_MS.to_string(),
+                "--leak-tcp",
+                "--leak-udp",
+                "--leak-icmp",
+            ]
+            .map(String::from)
+            .to_vec();
+
+            if let Some(payload) = &self.payload {
+                args.push("--payload".to_string());
+                args.push(payload.clone());
+            };
+
+            SpawnOpts {
+                attach_stdin: true,
+                attach_stdout: true,
+                args,
+                ..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:?}")
+    })
+}
diff --git a/test/test-manager/src/tests/install.rs b/test/test-manager/src/tests/install.rs
index cb1ddf58c5..f0e0942af7 100644
--- a/test/test-manager/src/tests/install.rs
+++ b/test/test-manager/src/tests/install.rs
@@ -281,6 +281,7 @@ pub async fn test_installation_idempotency(
     // Connect to any relay. This forces the daemon to enter a secured target state
     connect_and_wait(&mut mullvad_client)
         .await
+        .map(|_| ()) // Discard the new tunnel state
         .or_else(|error| match error {
             Error::UnexpectedErrorState(_) => Ok(()),
             err => Err(err),
diff --git a/test/test-manager/src/tests/split_tunnel.rs b/test/test-manager/src/tests/split_tunnel.rs
index be2379de42..609acf7ac8 100644
--- a/test/test-manager/src/tests/split_tunnel.rs
+++ b/test/test-manager/src/tests/split_tunnel.rs
@@ -1,35 +1,16 @@
-use anyhow::{anyhow, bail, ensure, Context};
+use anyhow::Context;
 use mullvad_management_interface::MullvadProxyClient;
-use pcap::Direction;
-use pnet_packet::ip::IpNextHeaderProtocols;
-use std::{
-    net::{IpAddr, Ipv4Addr, SocketAddr},
-    str,
-    time::Duration,
-};
+use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use test_macro::test_function;
-use test_rpc::{meta::Os, ServiceClient, SpawnOpts};
-use tokio::time::{sleep, timeout};
-
-use crate::network_monitor::{start_packet_monitor, MonitorOptions};
+use test_rpc::ServiceClient;
 
-use super::{config::TEST_CONFIG, helpers, TestContext};
+use super::{
+    helpers::{self, ConnChecker},
+    TestContext,
+};
 
-const CHECKER_FILENAME_WINDOWS: &str = "connection-checker.exe";
-const CHECKER_FILENAME_UNIX: &str = "connection-checker";
 const LEAK_DESTINATION: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)), 1337);
 
-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
-);
-
 /// Test that split tunneling works by asserting the following:
 /// - Splitting a process shouldn't do anything if tunnel is not connected.
 /// - A split process should never push traffic through the tunnel.
@@ -40,7 +21,7 @@ pub async fn test_split_tunnel(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> anyhow::Result<()> {
-    let mut checker = ConnChecker::new(rpc.clone(), mullvad_client.clone());
+    let mut checker = ConnChecker::new(rpc.clone(), mullvad_client.clone(), LEAK_DESTINATION);
 
     // Test that program is behaving when we are disconnected
     (checker.spawn().await?.assert_insecure().await)
@@ -93,276 +74,3 @@ pub async fn test_split_tunnel(
 
     Ok(())
 }
-
-/// 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].
-struct ConnChecker {
-    rpc: ServiceClient,
-    mullvad_client: MullvadProxyClient,
-
-    /// Path to the process binary.
-    executable_path: String,
-
-    /// Whether the process should be split when spawned. Needed on Linux.
-    split: bool,
-}
-
-struct ConnCheckerHandle<'a> {
-    checker: &'a mut ConnChecker,
-
-    /// ID of the spawned process.
-    pid: u32,
-}
-
-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) -> 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,
-            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",
-                &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(())
-    }
-
-    async fn check_connection(&mut self) -> anyhow::Result<ConnectionStatus> {
-        // Monitor all pakets going to LEAK_DESTINATION during the check.
-        let monitor = start_packet_monitor(
-            |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.
-        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:?}")
-    })
-}
diff --git a/test/test-manager/src/tests/tunnel.rs b/test/test-manager/src/tests/tunnel.rs
index a0acbf2114..8c6482d483 100644
--- a/test/test-manager/src/tests/tunnel.rs
+++ b/test/test-manager/src/tests/tunnel.rs
@@ -6,10 +6,11 @@ use super::{
     Error, TestContext,
 };
 use crate::{
-    network_monitor::{start_packet_monitor, MonitorOptions},
-    tests::helpers::login_with_retries,
+    network_monitor::{start_packet_monitor, MonitorOptions, ParsedPacket},
+    tests::helpers::{login_with_retries, ConnChecker},
 };
 
+use anyhow::{bail, ensure};
 use mullvad_management_interface::MullvadProxyClient;
 use mullvad_relay_selector::query::builder::RelayQueryBuilder;
 use mullvad_types::{
@@ -19,6 +20,7 @@ use mullvad_types::{
         RelaySettings, SelectedObfuscation, TransportPort, Udp2TcpObfuscationSettings,
         WireguardConstraints,
     },
+    states::TunnelState,
     wireguard,
 };
 use std::net::SocketAddr;
@@ -787,3 +789,87 @@ pub async fn test_establish_tunnel_without_api(
     // Profit
     Ok(())
 }
+
+/// Fail to leak traffic to verify that mitigation for MUL-02-002-WP2
+/// ("Firewall allows deanonymization by eavesdropper") works.
+///
+/// # Vulnerability
+/// 1. Connect to a relay on port 443. Record this relay's IP address (the new gateway of the
+///    client)
+/// 2. Start listening for unencrypted traffic on the outbound network interface
+/// (Choose some human-readable, identifiable payload to look for in the outgoing TCP packets)
+/// 3. Start a rogue program which performs a GET request* containing the payload defined in step 2
+/// 4. The network snooper started in step 2 should now be able to observe the network request
+///    containing the identifiable payload being sent unencrypted over the wire
+///
+/// * or something similiar, as long as it generates some traffic containing UDP and/or TCP packets
+/// with the correct payload.
+#[test_function]
+pub async fn test_mul_02_002(
+    _: TestContext,
+    rpc: ServiceClient,
+    mut mullvad_client: MullvadProxyClient,
+) -> anyhow::Result<()> {
+    // Step 1 - Choose a relay
+    helpers::constrain_to_relay(
+        &mut mullvad_client,
+        RelayQueryBuilder::new()
+            .openvpn()
+            .transport_protocol(TransportProtocol::Tcp)
+            .port(443)
+            .build(),
+    )
+    .await?;
+
+    // Step 1.5 - Temporarily connect to the relay to get the target endpoint
+    let tunnel_state = helpers::connect_and_wait(&mut mullvad_client).await?;
+    let TunnelState::Connected { endpoint, .. } = tunnel_state else {
+        bail!("Expected tunnel state to be `Connected` - instead it was {tunnel_state:?}");
+    };
+    helpers::disconnect_and_wait(&mut mullvad_client).await?;
+    let target_endpoint = endpoint.endpoint.address;
+
+    // Step 2 - Start a network monitor snooping the outbound network interface for some
+    // identifiable payload
+    let unique_identifier = "Hello there!";
+    let identify_rogue_packet = move |packet: &ParsedPacket| {
+        packet
+            .payload
+            .windows(unique_identifier.len())
+            .any(|window| window == unique_identifier.as_bytes())
+    };
+    let rogue_packet_monitor =
+        start_packet_monitor(identify_rogue_packet, MonitorOptions::default()).await;
+
+    // Step 3 - Start the rogue program which will try to leak the unique identifier payload
+    // to the chosen relay endpoint
+    let mut checker = ConnChecker::new(rpc.clone(), mullvad_client.clone(), target_endpoint);
+    checker.payload(unique_identifier);
+    let mut conn_artist = checker.spawn().await?;
+    // Before proceeding, assert that the method of detecting identifiable packets work.
+    conn_artist.check_connection().await?;
+    let monitor_result = rogue_packet_monitor.into_result().await?;
+
+    log::info!("Checking that the identifiable payload was detectable without encryption");
+    ensure!(
+        !monitor_result.packets.is_empty(),
+        "Did not observe rogue packets! The method seems to be broken"
+    );
+    log::info!("The identifiable payload was detected! (that's good)");
+
+    // Step 4 - Finally, connect to a tunnel and assert that no outgoing traffic contains the
+    // payload in plain text.
+    helpers::connect_and_wait(&mut mullvad_client).await?;
+    let rogue_packet_monitor =
+        start_packet_monitor(identify_rogue_packet, MonitorOptions::default()).await;
+    conn_artist.check_connection().await?;
+    let monitor_result = rogue_packet_monitor.into_result().await?;
+
+    log::info!("Checking that the identifiable payload was not detected");
+    ensure!(
+        monitor_result.packets.is_empty(),
+        "Observed rogue packets! The tunnel seems to be leaking traffic"
+    );
+
+    Ok(())
+}
diff --git a/test/test-manager/src/tests/tunnel_state.rs b/test/test-manager/src/tests/tunnel_state.rs
index 5e69b33b7a..edbe595092 100644
--- a/test/test-manager/src/tests/tunnel_state.rs
+++ b/test/test-manager/src/tests/tunnel_state.rs
@@ -11,12 +11,12 @@ use crate::{
 };
 
 use mullvad_management_interface::MullvadProxyClient;
+use mullvad_relay_selector::query::builder::RelayQueryBuilder;
 use mullvad_types::{
     constraints::Constraint,
     relay_constraints::{
         GeographicLocationConstraint, LocationConstraint, RelayConstraints, RelaySettings,
     },
-    relay_list::{Relay, RelayEndpointData},
     states::TunnelState,
     CustomTunnelEndpoint,
 };
@@ -340,40 +340,21 @@ pub async fn test_connected_state(
     _: TestContext,
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
-) -> Result<(), Error> {
+) -> anyhow::Result<()> {
     let inet_destination = "1.1.1.1:1337".parse().unwrap();
 
     // Set relay to use
-    //
-
     log::info!("Select relay");
-
-    let relay_filter = |relay: &Relay| {
-        relay.active && matches!(relay.endpoint_data, RelayEndpointData::Wireguard(_))
-    };
-
-    let relay = helpers::filter_relays(&mut mullvad_client, relay_filter)
-        .await?
-        .pop()
-        .unwrap();
-
-    let relay_settings = RelaySettings::Normal(RelayConstraints {
-        location: helpers::into_constraint(&relay),
-        ..Default::default()
-    });
-
-    set_relay_settings(&mut mullvad_client, relay_settings)
-        .await
-        .expect("failed to update relay settings");
+    let relay = helpers::constrain_to_relay(
+        &mut mullvad_client,
+        RelayQueryBuilder::new().wireguard().build(),
+    )
+    .await?;
 
     // Connect
-    //
-
     connect_and_wait(&mut mullvad_client).await?;
 
     // Verify that endpoint was selected
-    //
-
     match mullvad_client.get_tunnel_state().await? {
         TunnelState::Connected {
             endpoint:
diff --git a/test/test-manager/src/tests/ui.rs b/test/test-manager/src/tests/ui.rs
index ebd6fb3ee9..cce7cdd990 100644
--- a/test/test-manager/src/tests/ui.rs
+++ b/test/test-manager/src/tests/ui.rs
@@ -1,9 +1,6 @@
 use super::{config::TEST_CONFIG, helpers, Error, TestContext};
 use mullvad_management_interface::MullvadProxyClient;
-use mullvad_types::{
-    relay_constraints::{RelayConstraints, RelaySettings},
-    relay_list::{Relay, RelayEndpointData},
-};
+use mullvad_relay_selector::query::builder::RelayQueryBuilder;
 use std::{
     collections::BTreeMap,
     fmt::Debug,
@@ -85,25 +82,14 @@ pub async fn test_ui_tunnel_settings(
     _: TestContext,
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
-) -> Result<(), Error> {
+) -> anyhow::Result<()> {
     // tunnel-state.spec precondition: a single WireGuard relay should be selected
     log::info!("Select WireGuard relay");
-    let entry = helpers::filter_relays(&mut mullvad_client, |relay: &Relay| {
-        relay.active && matches!(relay.endpoint_data, RelayEndpointData::Wireguard(_))
-    })
-    .await?
-    .pop()
-    .unwrap();
-
-    // The test expects us to be disconnected and logged in but to have a specific relay selected
-    let relay_settings = RelaySettings::Normal(RelayConstraints {
-        location: helpers::into_constraint(&entry),
-        ..Default::default()
-    });
-
-    helpers::set_relay_settings(&mut mullvad_client, relay_settings)
-        .await
-        .expect("failed to update relay settings");
+    let entry = helpers::constrain_to_relay(
+        &mut mullvad_client,
+        RelayQueryBuilder::new().wireguard().build(),
+    )
+    .await?;
 
     let ui_result = run_test_env(
         &rpc,