Add unit test for maybenot machines for wg-go

mullvad · Oct 30, 2024 · 7fedf8a · 7fedf8a
1 parent 1c224df
commit 7fedf8a
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 11 deletions.
diff --git 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");
-            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,34 @@ impl Tunnel for WgGoTunnel {
     }
 }
 
+#[cfg(daita)]
+fn load_maybenot_machines(resource_dir: &Path) -> Result<CString> {
+    let path = resource_dir.join("maybenot_machines");
+    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 {
+    #[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/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,21 @@ impl Drop for Tunnel {
     }
 }
 
+pub fn validate_maybenot_machines(machines: &CStr) -> Result<(), Error> {
+    use maybenot_ffi::MaybenotResult;
+
+    let mut framework = MaybeUninit::uninit();
+    let result =
+        unsafe { maybenot_ffi::maybenot_start(machines.as_ptr(), 0.0, 0.0, &mut framework) };
+
+    if result as u32 == MaybenotResult::Ok as u32 {
+        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 {