Make payload of connection checker configurable

2 files changed, 31 insertions, 14 deletions

diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index 2dff6eccdd..86cbecb50c 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -715,6 +715,9 @@ pub struct ConnChecker {
 
     /// Whether the process should be split when spawned. Needed on Linux.
     split: bool,
+
+    /// Some arbitrary payload
+    payload: Option<String>,
 }
 
 pub struct ConnCheckerHandle<'a> {
@@ -756,9 +759,15 @@ impl ConnChecker {
             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.
@@ -766,18 +775,14 @@ impl ConnChecker {
     pub async fn spawn(&mut self) -> anyhow::Result<ConnCheckerHandle<'_>> {
         log::debug!("spawning connection checker");
 
-        let opts = SpawnOpts {
-            attach_stdin: true,
-            attach_stdout: true,
-            args: [
+        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(),
-                //TODO(markus): Remove
-                //&LEAK_DESTINATION.to_string(),
                 "--leak-timeout",
                 &LEAK_TIMEOUT_MS.to_string(),
                 "--leak-tcp",
@@ -785,8 +790,19 @@ impl ConnChecker {
                 "--leak-icmp",
             ]
             .map(String::from)
-            .to_vec(),
-            ..SpawnOpts::new(&self.executable_path)
+            .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?;
diff --git a/test/test-manager/src/tests/tunnel.rs b/test/test-manager/src/tests/tunnel.rs
index 594c930ef0..8c6482d483 100644
--- a/test/test-manager/src/tests/tunnel.rs
+++ b/test/test-manager/src/tests/tunnel.rs
@@ -831,30 +831,31 @@ pub async fn test_mul_02_002(
 
     // Step 2 - Start a network monitor snooping the outbound network interface for some
     // identifiable payload
-    // FIXME: This needs to be kept in sync with the magic payload string defined in `connection_cheker::net`.
-    // An easy fix would be to make the payload for `ConnCheck` configurable.
-    let unique_identifier = b"Hello there!";
+    let unique_identifier = "Hello there!";
     let identify_rogue_packet = move |packet: &ParsedPacket| {
         packet
             .payload
             .windows(unique_identifier.len())
-            .any(|window| window == unique_identifier)
+            .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 traffic to the chosen relay endpoint
+    // 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.unwrap();
+    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.