Add ping monitoring to `helpers.rs`

2 files changed, 131 insertions, 46 deletions

diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index 5b1c0a9903..fcdeb7866a 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -1,5 +1,7 @@
 use super::{config::TEST_CONFIG, Error, PING_TIMEOUT, WAIT_FOR_TUNNEL_STATE_TIMEOUT};
-use crate::network_monitor::{start_packet_monitor, MonitorOptions};
+use crate::network_monitor::{
+    start_packet_monitor, MonitorOptions, MonitorUnexpectedlyStopped, PacketMonitor,
+};
 use futures::StreamExt;
 use mullvad_management_interface::{types, ManagementServiceClient, MullvadProxyClient};
 use mullvad_types::{
@@ -523,3 +525,124 @@ pub fn into_constraint(relay: &Relay) -> Constraint<LocationConstraint> {
         .map(Constraint::Only)
         .expect("relay is missing location")
 }
+
+/// Ping monitoring made easy!
+///
+/// Continously ping some destination while monitoring to detect diverging
+/// packets.
+///
+/// To customize [`Pinger`] before the pinging and network monitoring starts,
+/// see [`PingerBuilder`]. Call [`start`](Pinger::start) to start pinging, and
+/// [`stop`](Pinger::stop) to get the leak test results.
+#[allow(dead_code)]
+pub struct Pinger {
+    // These values can be configured with [`PingerBuilder`].
+    destination: SocketAddr,
+    interval: tokio::time::Interval,
+    // Run-time specifics
+    pub guest_ip: IpAddr,
+    ping_task: AbortOnDrop<tokio::task::JoinHandle<()>>,
+    monitor: PacketMonitor,
+}
+
+impl Pinger {
+    /// Create a [`Pinger`] with a default configuration.
+    ///
+    /// See [`PingerBuilder`] for details.
+    pub async fn start(rpc: &test_rpc::ServiceClient) -> Pinger {
+        let defaults = PingerBuilder::new();
+        Self::start_with(defaults, rpc).await
+    }
+
+    /// Create a [`Pinger`] using the configuration of `builder`.
+    ///
+    /// See [`PingerBuilder`] for details on how to configure a [`Pinger`]
+    /// before starting it.
+    pub async fn start_with(builder: PingerBuilder, rpc: &test_rpc::ServiceClient) -> Pinger {
+        let guest_iface = rpc
+            .get_default_interface()
+            .await
+            .expect("failed to obtain default interface");
+        let guest_ip = rpc
+            .get_interface_ip(guest_iface)
+            .await
+            .expect("failed to obtain non-tun IP");
+        log::debug!("Guest IP: {guest_ip}");
+
+        // Start a network monitor
+        log::debug!("Monitoring outgoing traffic");
+        let monitor = start_packet_monitor(
+            move |packet| {
+                // NOTE: Many packets will likely be observed for API traffic. Rather than filtering all
+                // of those specifically, simply fail if our probes are observed.
+                packet.source.ip() == guest_ip
+                    && packet.destination.ip() == builder.destination.ip()
+            },
+            MonitorOptions::default(),
+        )
+        .await;
+        // Start pinging
+        let ping_rpc = rpc.clone();
+        let mut interval = tokio::time::interval(builder.interval.period());
+        #[allow(clippy::async_yields_async)]
+        let ping_task = AbortOnDrop::new(tokio::spawn(async move {
+            // Send a ping once every second.
+            loop {
+                send_guest_probes_without_monitor(ping_rpc.clone(), None, builder.destination)
+                    .await;
+                interval.tick().await;
+            }
+        }));
+
+        Pinger {
+            destination: builder.destination,
+            interval: builder.interval,
+            guest_ip,
+            ping_task,
+            monitor,
+        }
+    }
+
+    pub async fn stop(
+        self,
+    ) -> Result<crate::network_monitor::MonitorResult, MonitorUnexpectedlyStopped> {
+        // Abort the inner probe sender, which is accomplish by dropping the
+        // join handle to the running task.
+        drop(self.ping_task);
+        self.monitor.into_result().await
+    }
+
+    /// Return the time period determining the cadence of pings that are sent.
+    pub fn period(&self) -> tokio::time::Duration {
+        self.interval.period()
+    }
+}
+
+/// Configure a [`Pinger`] before starting it.
+pub struct PingerBuilder {
+    destination: SocketAddr,
+    interval: tokio::time::Interval,
+}
+
+#[allow(dead_code)]
+impl PingerBuilder {
+    //
+    pub fn new() -> PingerBuilder {
+        PingerBuilder {
+            destination: "1.1.1.1:1337".parse().unwrap(),
+            interval: tokio::time::interval(Duration::from_secs(1)),
+        }
+    }
+
+    /// Set the target to ping.
+    pub fn destination(mut self, destination: SocketAddr) -> Self {
+        self.destination = destination;
+        self
+    }
+
+    /// How often a ping should be sent.
+    pub fn interval(mut self, period: Duration) -> Self {
+        self.interval = tokio::time::interval(period);
+        self
+    }
+}
diff --git a/test/test-manager/src/tests/install.rs b/test/test-manager/src/tests/install.rs
index c809fd9ad2..78b97e3abd 100644
--- a/test/test-manager/src/tests/install.rs
+++ b/test/test-manager/src/tests/install.rs
@@ -343,42 +343,7 @@ pub async fn test_installation_idempotency(
         .expect("failed to enable auto-connect");
     // Start a tunnel monitor. No traffic should be observed going outside of
     // the tunnel during either installation process.
-    let inet_destination: SocketAddr = "1.1.1.1:1337".parse().unwrap();
-    let guest_iface = rpc
-        .get_default_interface()
-        .await
-        .expect("failed to obtain default interface");
-    let guest_ip = rpc
-        .get_interface_ip(guest_iface)
-        .await
-        .expect("failed to obtain non-tun IP");
-    log::debug!("Guest IP: {guest_ip}");
-
-    log::debug!("Monitoring outgoing traffic");
-    let monitor = start_packet_monitor(
-        move |packet| {
-            // NOTE: Many packets will likely be observed for API traffic. Rather than filtering all
-            // of those specifically, simply fail if our probes are observed.
-            packet.source.ip() == guest_ip && packet.destination.ip() == inet_destination.ip()
-        },
-        MonitorOptions::default(),
-    )
-    .await;
-    // TODO: Refactor this into a nice-looking API which returns both a pinger and an associated ping-monitor. => (Pinger, PingMonitor).
-    let ping_rpc = rpc.clone();
-    let probe_sender = AbortOnDrop::new(tokio::spawn(async move {
-        // Send a ping once every second.
-        loop {
-            super::helpers::send_guest_probes_without_monitor(
-                ping_rpc.clone(),
-                None,
-                inet_destination,
-            )
-            .await;
-            tokio::time::sleep(Duration::from_secs(1)).await;
-        }
-    }));
-
+    let pinger = Pinger::start(&rpc).await;
     for _ in 1..=2 {
         // install package
         log::debug!("Installing new app");
@@ -393,19 +358,16 @@ pub async fn test_installation_idempotency(
                 );
                 err
             })?;
-
         // Wait for an arbitrary amount of time. The point is that the pinger
-        // should be able to ping `inet_destination` while the newly installed
-        // app is running.
-        tokio::time::sleep(Duration::from_secs(3)).await;
+        // should be able to ping while the newly installed app is running.
+        if let Some(delay) = pinger.period().checked_mul(3) {
+            tokio::time::sleep(delay).await;
+        }
     }
 
     // Make sure that no traffic leak occured during any installation process.
-    let probe_sender = probe_sender.into_inner();
-    probe_sender.abort();
-    let _ = probe_sender.await;
-
-    let monitor_result = monitor.into_result().await.unwrap();
+    let guest_ip = pinger.guest_ip;
+    let monitor_result = pinger.stop().await.unwrap();
     assert_eq!(
         monitor_result.packets.len(),
         0,