Make payload of connection checker configurable

4 files changed, 38 insertions, 19 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 ff45993408..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(PAYLOAD)
+        .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(PAYLOAD, 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/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.