Fully abort all tasks before leaving a test

3 files changed, 119 insertions, 79 deletions

diff --git a/test/test-manager/src/tests/dns.rs b/test/test-manager/src/tests/dns.rs
index e87da24db0..69de367b88 100644
--- a/test/test-manager/src/tests/dns.rs
+++ b/test/test-manager/src/tests/dns.rs
@@ -212,44 +212,50 @@ async fn leak_test_dns(
     // We should observe 2 outgoing packets to the whitelisted destination
     // on port 53, and only inside the desired interface.
 
-    spoof_packets(
-        rpc,
-        Some(Interface::Tunnel),
-        tun_bind_addr,
-        whitelisted_dest,
-    );
-    spoof_packets(
-        rpc,
-        Some(Interface::NonTunnel),
-        guest_bind_addr,
-        whitelisted_dest,
-    );
-
-    spoof_packets(
-        rpc,
-        Some(Interface::Tunnel),
-        tun_bind_addr,
-        blocked_dest_local,
-    );
-    spoof_packets(
-        rpc,
-        Some(Interface::NonTunnel),
-        guest_bind_addr,
-        blocked_dest_local,
-    );
-
-    spoof_packets(
-        rpc,
-        Some(Interface::Tunnel),
-        tun_bind_addr,
-        blocked_dest_public,
-    );
-    spoof_packets(
-        rpc,
-        Some(Interface::NonTunnel),
-        guest_bind_addr,
-        blocked_dest_public,
-    );
+    let rpc = rpc.clone();
+    let probes = tokio::spawn(async move {
+        tokio::join!(
+            // send to allowed dest
+            spoof_packets(
+                &rpc,
+                Some(Interface::Tunnel),
+                tun_bind_addr,
+                whitelisted_dest,
+            ),
+            spoof_packets(
+                &rpc,
+                Some(Interface::NonTunnel),
+                guest_bind_addr,
+                whitelisted_dest,
+            ),
+            // send to blocked local dest
+            spoof_packets(
+                &rpc,
+                Some(Interface::Tunnel),
+                tun_bind_addr,
+                blocked_dest_local,
+            ),
+            spoof_packets(
+                &rpc,
+                Some(Interface::NonTunnel),
+                guest_bind_addr,
+                blocked_dest_local,
+            ),
+            // send to blocked public dest
+            spoof_packets(
+                &rpc,
+                Some(Interface::Tunnel),
+                tun_bind_addr,
+                blocked_dest_public,
+            ),
+            spoof_packets(
+                &rpc,
+                Some(Interface::NonTunnel),
+                guest_bind_addr,
+                blocked_dest_public,
+            ),
+        )
+    });
 
     if use_tun {
         //
@@ -257,6 +263,10 @@ async fn leak_test_dns(
         //
 
         let tunnel_result = tunnel_monitor.wait().await.unwrap();
+
+        probes.abort();
+        let _ = probes.await;
+
         assert!(
             tunnel_result.packets.len() >= 2,
             "expected at least 2 in-tunnel packets to allowed destination only"
@@ -282,6 +292,9 @@ async fn leak_test_dns(
     } else {
         let non_tunnel_result = non_tunnel_monitor.wait().await.unwrap();
 
+        probes.abort();
+        let _ = probes.await;
+
         //
         // Examine tunnel traffic
         //
@@ -605,6 +618,7 @@ async fn run_dns_config_test<
     );
 
     handle.abort();
+    let _ = handle.await;
 
     Ok(())
 }
@@ -647,22 +661,23 @@ async fn connect_local_wg_relay(mullvad_client: &mut ManagementServiceClient) ->
     Ok(())
 }
 
-fn spoof_packets(
+async fn spoof_packets(
     rpc: &ServiceClient,
     interface: Option<Interface>,
     bind_addr: SocketAddr,
     dest: SocketAddr,
 ) {
-    let rpc1 = rpc.clone();
-    let rpc2 = rpc.clone();
-    tokio::spawn(async move {
+    let tcp_rpc = rpc.clone();
+    let tcp_send = async move {
         log::debug!("sending to {}/tcp from {}", dest, bind_addr);
-        let _ = rpc1.send_tcp(interface, bind_addr, dest).await;
-    });
-    tokio::spawn(async move {
+        let _ = tcp_rpc.send_tcp(interface, bind_addr, dest).await;
+    };
+    let udp_rpc = rpc.clone();
+    let udp_send = async move {
         log::debug!("sending to {}/udp from {}", dest, bind_addr);
-        let _ = rpc2.send_udp(interface, bind_addr, dest).await;
-    });
+        let _ = udp_rpc.send_udp(interface, bind_addr, dest).await;
+    };
+    let _ = tokio::join!(tcp_send, udp_send);
 }
 
 type ShouldContinue = bool;
diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index daef783247..f73e6f6e79 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -79,33 +79,16 @@ pub async fn send_guest_probes(
     )
     .await;
 
-    let bind_addr = if let Some(interface) = interface {
-        SocketAddr::new(
-            rpc.get_interface_ip(interface)
-                .await
-                .expect("failed to obtain interface IP"),
-            0,
-        )
-    } else {
-        "0.0.0.0:0".parse().unwrap()
-    };
-
-    let send_handle = tokio::spawn(async move {
-        let tcp_rpc = rpc.clone();
-        let udp_rpc = rpc.clone();
-        tokio::spawn(async move {
-            let _ = tcp_rpc.send_tcp(interface, bind_addr, destination).await;
-        });
-        tokio::spawn(async move {
-            let _ = udp_rpc.send_udp(interface, bind_addr, destination).await;
-        });
-        ping_with_timeout(&rpc, destination.ip(), interface).await?;
-        Ok::<(), Error>(())
-    });
+    let send_handle = tokio::spawn(send_guest_probes_without_monitor(
+        rpc,
+        interface,
+        destination,
+    ));
 
     let monitor_result = pktmon.wait().await.unwrap();
 
     send_handle.abort();
+    let _ = send_handle.await;
 
     let mut result = ProbeResult::default();
 
@@ -127,6 +110,31 @@ pub async fn send_guest_probes(
     Ok(result)
 }
 
+/// Send one probe per transport protocol to `destination` without running a packet monitor
+pub async fn send_guest_probes_without_monitor(
+    rpc: ServiceClient,
+    interface: Option<Interface>,
+    destination: SocketAddr,
+) {
+    let bind_addr = if let Some(interface) = interface {
+        SocketAddr::new(
+            rpc.get_interface_ip(interface)
+                .await
+                .expect("failed to obtain interface IP"),
+            0,
+        )
+    } else {
+        "0.0.0.0:0".parse().unwrap()
+    };
+
+    let tcp_rpc = rpc.clone();
+    let tcp_send = async move { tcp_rpc.send_tcp(interface, bind_addr, destination).await };
+    let udp_rpc = rpc.clone();
+    let udp_send = async move { udp_rpc.send_udp(interface, bind_addr, destination).await };
+    let icmp = async move { ping_with_timeout(&rpc, destination.ip(), interface).await };
+    let _ = tokio::join!(tcp_send, udp_send, icmp);
+}
+
 pub async fn ping_with_timeout(
     rpc: &ServiceClient,
     dest: IpAddr,
@@ -274,11 +282,23 @@ pub async fn geoip_lookup_with_retries(rpc: &ServiceClient) -> Result<AmIMullvad
     }
 }
 
-pub struct AbortOnDrop<T>(pub tokio::task::JoinHandle<T>);
+pub struct AbortOnDrop<T>(Option<tokio::task::JoinHandle<T>>);
+
+impl<T> AbortOnDrop<T> {
+    pub fn new(inner: tokio::task::JoinHandle<T>) -> AbortOnDrop<T> {
+        AbortOnDrop(Some(inner))
+    }
+
+    pub fn into_inner(mut self) -> tokio::task::JoinHandle<T> {
+        self.0.take().unwrap()
+    }
+}
 
 impl<T> Drop for AbortOnDrop<T> {
     fn drop(&mut self) {
-        self.0.abort();
+        if let Some(task) = self.0.take() {
+            task.abort();
+        }
     }
 }
 
diff --git a/test/test-manager/src/tests/install.rs b/test/test-manager/src/tests/install.rs
index abd105d4dd..04d401a22d 100644
--- a/test/test-manager/src/tests/install.rs
+++ b/test/test-manager/src/tests/install.rs
@@ -1,4 +1,4 @@
-use super::helpers::{get_package_desc, ping_with_timeout, AbortOnDrop};
+use super::helpers::{get_package_desc, AbortOnDrop};
 use super::{Error, TestContext};
 
 use super::config::TEST_CONFIG;
@@ -48,7 +48,6 @@ pub async fn test_install_previous_app(_: TestContext, rpc: ServiceClient) -> Re
 #[test_function(priority = -190)]
 pub async fn test_upgrade_app(ctx: TestContext, rpc: ServiceClient) -> Result<(), Error> {
     let inet_destination: SocketAddr = "1.1.1.1:1337".parse().unwrap();
-    let bind_addr: SocketAddr = "0.0.0.0:0".parse().unwrap();
 
     // Verify that daemon is running
     if rpc.mullvad_daemon_get_status().await? != ServiceStatus::Running {
@@ -122,11 +121,14 @@ pub async fn test_upgrade_app(ctx: TestContext, rpc: ServiceClient) -> Result<()
     .await;
 
     let ping_rpc = rpc.clone();
-    let abort_on_drop = AbortOnDrop(tokio::spawn(async move {
+    let probe_sender = AbortOnDrop::new(tokio::spawn(async move {
         loop {
-            let _ = ping_rpc.send_tcp(None, bind_addr, inet_destination).await;
-            let _ = ping_rpc.send_udp(None, bind_addr, inet_destination).await;
-            let _ = ping_with_timeout(&ping_rpc, inet_destination.ip(), None).await;
+            super::helpers::send_guest_probes_without_monitor(
+                ping_rpc.clone(),
+                None,
+                inet_destination,
+            )
+            .await;
             tokio::time::sleep(Duration::from_secs(1)).await;
         }
     }));
@@ -147,7 +149,10 @@ pub async fn test_upgrade_app(ctx: TestContext, rpc: ServiceClient) -> Result<()
     //
     // Check if any traffic was observed
     //
-    drop(abort_on_drop);
+    let probe_sender = probe_sender.into_inner();
+    probe_sender.abort();
+    let _ = probe_sender.await;
+
     let monitor_result = monitor.into_result().await.unwrap();
     assert_eq!(
         monitor_result.packets.len(),