Merge branch 'add-maybenot-test'

3 files changed, 94 insertions, 31 deletions

diff --git a/talpid-wireguard/src/wireguard_go/mod.rs b/talpid-wireguard/src/wireguard_go/mod.rs
index 28470a0e46..f9a062f3c8 100644
--- a/talpid-wireguard/src/wireguard_go/mod.rs
+++ b/talpid-wireguard/src/wireguard_go/mod.rs
@@ -218,16 +218,8 @@ impl Tunnel for WgGoTunnel {
     #[cfg(daita)]
     fn start_daita(&mut self) -> Result<()> {
         static MAYBENOT_MACHINES: OnceCell<CString> = OnceCell::new();
-        let machines = MAYBENOT_MACHINES.get_or_try_init(|| {
-            let path = self.resource_dir.join("maybenot_machines_v2");
-            log::debug!("Reading maybenot machines from {}", path.display());
-
-            let machines =
-                fs::read_to_string(path).map_err(|e| TunnelError::StartDaita(Box::new(e)))?;
-            let machines =
-                CString::new(machines).map_err(|e| TunnelError::StartDaita(Box::new(e)))?;
-            Ok(machines)
-        })?;
+        let machines =
+            MAYBENOT_MACHINES.get_or_try_init(|| load_maybenot_machines(&self.resource_dir))?;
 
         log::info!("Initializing DAITA for wireguard device");
         let peer_public_key = &self.config.entry_peer.public_key;
@@ -244,6 +236,36 @@ impl Tunnel for WgGoTunnel {
     }
 }
 
+#[cfg(daita)]
+fn load_maybenot_machines(resource_dir: &Path) -> Result<CString> {
+    let path = resource_dir.join("maybenot_machines_v2");
+    log::debug!("Reading maybenot machines from {}", path.display());
+
+    let machines = fs::read_to_string(path).map_err(|e| TunnelError::StartDaita(Box::new(e)))?;
+    let machines = CString::new(machines).map_err(|e| TunnelError::StartDaita(Box::new(e)))?;
+    Ok(machines)
+}
+
+#[cfg(test)]
+mod test {
+    /// Test whether `maybenot_machines` in dist-assets contains valid machines.
+    /// TODO: Remove when switching to dynamic machines.
+    #[cfg(daita)]
+    #[test]
+    fn test_load_maybenot_machines() {
+        use super::load_maybenot_machines;
+        use std::path::PathBuf;
+
+        let dist_assets = std::env::var("CARGO_MANIFEST_DIR")
+            .map(PathBuf::from)
+            .expect("CARGO_MANIFEST_DIR env var not set")
+            .join("..")
+            .join("dist-assets");
+        let machines = load_maybenot_machines(&dist_assets).unwrap();
+        wireguard_go_rs::validate_maybenot_machines(&machines).unwrap();
+    }
+}
+
 mod stats {
     use super::{Stats, StatsMap};
 
diff --git a/talpid-wireguard/src/wireguard_nt/daita.rs b/talpid-wireguard/src/wireguard_nt/daita.rs
index 88c16c6a91..135ea9a2fd 100644
--- a/talpid-wireguard/src/wireguard_nt/daita.rs
+++ b/talpid-wireguard/src/wireguard_nt/daita.rs
@@ -250,26 +250,8 @@ impl Machinist {
 
         static MAYBENOT_MACHINES: OnceCell<Vec<maybenot::Machine>> = OnceCell::new();
 
-        let machines = MAYBENOT_MACHINES.get_or_try_init(|| {
-            let path = resource_dir.join("maybenot_machines_v2");
-            log::debug!("Reading maybenot machines from {}", path.display());
-
-            let mut machines = vec![];
-            let machines_str = fs::read_to_string(path).map_err(Error::EnumerateMachines)?;
-            for machine_str in machines_str.lines() {
-                let machine_str = machine_str.trim();
-                if matches!(machine_str.chars().next(), None | Some('#')) {
-                    continue;
-                }
-                log::debug!("Adding maybenot machine: {machine_str}");
-                machines.push(
-                    machine_str
-                        .parse::<maybenot::Machine>()
-                        .map_err(|_error| Error::InvalidMachine(machine_str.to_owned()))?,
-                );
-            }
-            Ok(machines)
-        })?;
+        let machines =
+            MAYBENOT_MACHINES.get_or_try_init(|| load_maybenot_machines(resource_dir))?;
 
         let quit_event =
             talpid_windows::sync::Event::new(true, false).map_err(Error::InitializeQuitEvent)?;
@@ -456,3 +438,43 @@ impl Machinist {
         }
     }
 }
+
+fn load_maybenot_machines(resource_dir: &Path) -> Result<Vec<maybenot::Machine>, Error> {
+    let path = resource_dir.join("maybenot_machines_v2");
+    log::debug!("Reading maybenot machines from {}", path.display());
+
+    let mut machines = vec![];
+    let machines_str = fs::read_to_string(path).map_err(Error::EnumerateMachines)?;
+    for machine_str in machines_str.lines() {
+        let machine_str = machine_str.trim();
+        if matches!(machine_str.chars().next(), None | Some('#')) {
+            continue;
+        }
+        log::debug!("Adding maybenot machine: {machine_str}");
+        machines.push(
+            machine_str
+                .parse::<maybenot::Machine>()
+                .map_err(|_error| Error::InvalidMachine(machine_str.to_owned()))?,
+        );
+    }
+    Ok(machines)
+}
+
+#[cfg(test)]
+mod test {
+    use super::load_maybenot_machines;
+    use std::path::PathBuf;
+
+    /// Test whether `maybenot_machines` in dist-assets contains valid machines.
+    /// TODO: Remove when switching to dynamic machines.
+    #[test]
+    fn test_load_maybenot_machines() {
+        let dist_assets = std::env::var("CARGO_MANIFEST_DIR")
+            .map(PathBuf::from)
+            .expect("CARGO_MANIFEST_DIR env var not set")
+            .join("..")
+            .join("dist-assets");
+
+        load_maybenot_machines(&dist_assets).unwrap();
+    }
+}
diff --git a/wireguard-go-rs/src/lib.rs b/wireguard-go-rs/src/lib.rs
index 6bacd325b5..a77b48c0bd 100644
--- a/wireguard-go-rs/src/lib.rs
+++ b/wireguard-go-rs/src/lib.rs
@@ -11,7 +11,7 @@
 use core::slice;
 use std::{
     ffi::{c_char, CStr},
-    mem::ManuallyDrop,
+    mem::{ManuallyDrop, MaybeUninit},
 };
 use util::OnDrop;
 use zeroize::Zeroize;
@@ -199,6 +199,25 @@ impl Drop for Tunnel {
     }
 }
 
+/// Check whether `machines` contains a valid, LF-separated maybenot machines. Return an error
+/// otherwise.
+pub fn validate_maybenot_machines(machines: &CStr) -> Result<(), Error> {
+    use maybenot_ffi::MaybenotResult;
+
+    let mut framework = MaybeUninit::uninit();
+    // SAFETY: `machines` is a null-terminated string, and `&mut framework` is a valid pointer
+    let result =
+        unsafe { maybenot_ffi::maybenot_start(machines.as_ptr(), 0.0, 0.0, &mut framework) };
+
+    if result as u32 == MaybenotResult::Ok as u32 {
+        // SAFETY: `maybenot_start` succeeded, so `framework` points to a valid framework
+        unsafe { maybenot_ffi::maybenot_stop(framework.assume_init()) };
+        Ok(())
+    } else {
+        Err(Error::Other)
+    }
+}
+
 fn result_from_code(code: i32) -> Result<(), Error> {
     // NOTE: must be kept in sync with enum definition
     Err(match code {