From 7b7afffc4621ee87528d10a9cc4466074e8c3226 Mon Sep 17 00:00:00 2001
From: Vicki Pfau <vi@endrift.com>
Date: Tue, 14 Jan 2025 20:53:12 -0800
Subject: [PATCH] manager: Expose new CaptureDebugTraceOutput method

---
 Cargo.toml                           |  4 +-
 com.steampowered.SteamOSManager1.xml | 13 +++++++
 src/bin/steamosctl.rs                |  8 ++++
 src/manager/root.rs                  | 14 ++++++-
 src/manager/user.rs                  |  4 ++
 src/process.rs                       | 36 ++++++++++++++++++
 src/proxy/wifi_debug1.rs             |  3 ++
 src/wifi.rs                          | 55 +++++++++++++++++++++++++---
 8 files changed, 126 insertions(+), 11 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 9ab3d38..96f66aa 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -3,9 +3,6 @@ name = "steamos-manager"
 version = "24.5.1"
 edition = "2021"
 
-[dev-dependencies]
-tempfile = "3"
-
 [profile.release]
 strip="symbols"
 
@@ -22,6 +19,7 @@ nix = { version = "0.29", default-features = false, features = ["fs", "poll", "s
 num_enum = "0.7"
 regex = "1"
 serde = { version = "1.0", default-features = false, features = ["derive"] }
+tempfile = "3"
 tokio = { version = "1", default-features = false, features = ["fs", "io-std", "io-util", "macros", "process", "rt-multi-thread", "signal", "sync"] }
 tokio-stream = { version = "0.1", default-features = false }
 tokio-util = { version = "0.7", default-features = false }
diff --git a/com.steampowered.SteamOSManager1.xml b/com.steampowered.SteamOSManager1.xml
index 688bdec..2cfcbba 100644
--- a/com.steampowered.SteamOSManager1.xml
+++ b/com.steampowered.SteamOSManager1.xml
@@ -361,6 +361,19 @@
     -->
     <property name="WifiBackend" type="s" access="readwrite"/>
 
+    <!--
+        CaptureDebugTraceOutput:
+
+        @path: The path to the captured file. This will be world-readable and
+        in a temporary directory, so make sure to move it to a permanent
+        location if keeping it is desired.
+
+        Extract the Wi-Fi trace to a file, if present.
+    -->
+    <method name="CaptureDebugTraceOutput">
+      <arg type="s" name="path" direction="out"/>
+    </method>
+
   </interface>
 
   <!--
diff --git a/src/bin/steamosctl.rs b/src/bin/steamosctl.rs
index bd2217d..809c6cb 100644
--- a/src/bin/steamosctl.rs
+++ b/src/bin/steamosctl.rs
@@ -133,6 +133,9 @@ enum Commands {
     /// Get Wi-Fi debug mode
     GetWifiDebugMode,
 
+    /// Capture the current Wi-Fi debug trace
+    CaptureWifiDebugTraceOutput,
+
     /// Set the Wi-Fi power management state
     SetWifiPowerManagementState {
         /// Valid modes are `enabled`, `disabled`
@@ -380,6 +383,11 @@ async fn main() -> Result<()> {
                 Err(_) => println!("Got unknown value {mode} from backend"),
             }
         }
+        Commands::CaptureWifiDebugTraceOutput => {
+            let proxy = WifiDebugDump1Proxy::new(&conn).await?;
+            let path = proxy.generate_debug_dump().await?;
+            println!("{path}");
+        }
         Commands::SetWifiPowerManagementState { state } => {
             let proxy = WifiPowerManagement1Proxy::new(&conn).await?;
             proxy.set_wifi_power_management_state(*state as u32).await?;
diff --git a/src/manager/root.rs b/src/manager/root.rs
index d892ed2..de0d23a 100644
--- a/src/manager/root.rs
+++ b/src/manager/root.rs
@@ -29,8 +29,8 @@ use crate::power::{
 };
 use crate::process::{run_script, script_output};
 use crate::wifi::{
-    set_wifi_backend, set_wifi_debug_mode, set_wifi_power_management_state, WifiBackend,
-    WifiDebugMode, WifiPowerManagement,
+    extract_wifi_trace, set_wifi_backend, set_wifi_debug_mode, set_wifi_power_management_state,
+    WifiBackend, WifiDebugMode, WifiPowerManagement,
 };
 use crate::{path, API_VERSION};
 
@@ -369,6 +369,16 @@ impl SteamOSManager {
             .map_err(to_zbus_fdo_error)
     }
 
+    async fn capture_debug_trace_output(&self) -> fdo::Result<String> {
+        Ok(extract_wifi_trace()
+            .await
+            .inspect_err(|message| error!("Error capturing trace output: {message}"))
+            .map_err(to_zbus_fdo_error)?
+            .into_os_string()
+            .to_string_lossy()
+            .into())
+    }
+
     #[zbus(property)]
     async fn inhibit_ds(&self) -> fdo::Result<bool> {
         let (tx, rx) = oneshot::channel();
diff --git a/src/manager/user.rs b/src/manager/user.rs
index 69f5d75..9aa1a74 100644
--- a/src/manager/user.rs
+++ b/src/manager/user.rs
@@ -486,6 +486,10 @@ impl WifiDebug1 {
     async fn set_wifi_backend(&self, backend: &str) -> zbus::Result<()> {
         self.proxy.call("SetWifiBackend", &(backend)).await
     }
+
+    async fn capture_debug_trace_output(&self) -> fdo::Result<String> {
+        method!(self, "CaptureDebugTraceOutput")
+    }
 }
 
 #[interface(name = "com.steampowered.SteamOSManager1.WifiPowerManagement1")]
diff --git a/src/process.rs b/src/process.rs
index 96377a2..9570d9f 100644
--- a/src/process.rs
+++ b/src/process.rs
@@ -7,11 +7,18 @@
 
 use anyhow::{anyhow, Result};
 use std::ffi::OsStr;
+use tokio::process::ChildStdout;
+
 #[cfg(not(test))]
 use std::process::Stdio;
 #[cfg(not(test))]
 use tokio::process::Command;
 
+#[cfg(test)]
+use nix::{fcntl::OFlag, unistd::pipe2};
+#[cfg(test)]
+use std::{fs::File, io::Write};
+
 #[cfg(not(test))]
 pub async fn script_exit_code(
     executable: impl AsRef<OsStr>,
@@ -73,6 +80,35 @@ pub async fn script_output(
     cb(executable.as_ref(), args.as_ref()).map(|(_, res)| res)
 }
 
+#[cfg(not(test))]
+pub async fn script_pipe_output(
+    executable: impl AsRef<OsStr>,
+    args: &[impl AsRef<OsStr>],
+) -> Result<ChildStdout> {
+    // Run given command and return the output given, as an fd instead of a String
+    let child = Command::new(executable)
+        .args(args)
+        .stdout(Stdio::piped())
+        .stderr(Stdio::null())
+        .spawn()?;
+
+    child.stdout.ok_or(anyhow!("Failed to get stdout"))
+}
+
+#[cfg(test)]
+pub async fn script_pipe_output(
+    executable: impl AsRef<OsStr>,
+    args: &[impl AsRef<OsStr>],
+) -> Result<ChildStdout> {
+    let test = crate::testing::current();
+    let args: Vec<&OsStr> = args.iter().map(std::convert::AsRef::as_ref).collect();
+    let cb = test.process_cb.get();
+    let string = cb(executable.as_ref(), args.as_ref()).map(|(_, res)| res)?;
+    let (rx, tx) = pipe2(OFlag::O_CLOEXEC)?;
+    File::from(tx).write(string.as_bytes())?;
+    Ok(ChildStdout::from_std(std::process::ChildStdout::from(rx))?)
+}
+
 #[cfg(test)]
 pub(crate) mod test {
     use super::*;
diff --git a/src/proxy/wifi_debug1.rs b/src/proxy/wifi_debug1.rs
index da02dbe..d61f59a 100644
--- a/src/proxy/wifi_debug1.rs
+++ b/src/proxy/wifi_debug1.rs
@@ -19,6 +19,9 @@ use zbus::proxy;
     assume_defaults = true
 )]
 pub trait WifiDebug1 {
+    /// CaptureDebugTraceOutput method
+    fn capture_debug_trace_output(&self) -> zbus::Result<String>;
+
     /// SetWifiDebugMode method
     fn set_wifi_debug_mode(
         &self,
diff --git a/src/wifi.rs b/src/wifi.rs
index e59bd2e..4050211 100644
--- a/src/wifi.rs
+++ b/src/wifi.rs
@@ -8,14 +8,20 @@
 use anyhow::{bail, ensure, Result};
 use config::builder::AsyncState;
 use config::{ConfigBuilder, FileFormat};
+use nix::sys::stat::{self, Mode};
 use num_enum::TryFromPrimitive;
+use std::fs::Permissions;
+use std::os::unix::fs::PermissionsExt;
+use std::path::PathBuf;
 use std::str::FromStr;
 use strum::{Display, EnumString};
+use tempfile::Builder as TempFileBuilder;
 use tokio::fs;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tracing::error;
 use zbus::Connection;
 
-use crate::process::{run_script, script_output};
+use crate::process::{run_script, script_output, script_pipe_output};
 use crate::systemd::{daemon_reload, SystemdUnit};
 use crate::{path, read_config_directory};
 
@@ -28,7 +34,6 @@ const OVERRIDE_PATH: &str = "/etc/systemd/system/iwd.service.d/99-valve-override
 
 // Only use one path for output for now. If needed we can add a timestamp later
 // to have multiple files, etc.
-const OUTPUT_FILE: &str = "/var/log/wifitrace.dat";
 const TRACE_CMD_PATH: &str = "/usr/bin/trace-cmd";
 
 const MIN_BUFFER_SIZE: u32 = 100;
@@ -114,10 +119,7 @@ async fn restart_iwd(connection: Connection) -> Result<()> {
 }
 
 async fn stop_tracing() -> Result<()> {
-    // Stop tracing and extract ring buffer to disk for capture
-    run_script(TRACE_CMD_PATH, &["stop"]).await?;
-    // stop tracing worked
-    run_script(TRACE_CMD_PATH, &["extract", "-o", OUTPUT_FILE]).await
+    run_script(TRACE_CMD_PATH, &["stop"]).await
 }
 
 async fn start_tracing(buffer_size: u32) -> Result<()> {
@@ -130,6 +132,32 @@ async fn start_tracing(buffer_size: u32) -> Result<()> {
     .await
 }
 
+fn make_tempfile(prefix: &str) -> Result<(fs::File, PathBuf)> {
+    let umask = stat::umask(Mode::from_bits_truncate(0));
+    let output = TempFileBuilder::new()
+        .prefix(prefix)
+        .permissions(Permissions::from_mode(0o666))
+        .tempfile()?;
+    let (output, path) = output.keep()?;
+    let output = fs::File::from_std(output);
+    stat::umask(umask);
+
+    Ok((output, path))
+}
+
+pub async fn extract_wifi_trace() -> Result<PathBuf> {
+    let (mut output, path) = make_tempfile("wifi-trace-")?;
+    let mut pipe = script_pipe_output("trace-cmd", &["extract"]).await?;
+    let mut buf = [0; 4096];
+    loop {
+        let read = pipe.read(&mut buf).await?;
+        if read == 0 {
+            break Ok(path);
+        }
+        output.write_all(&buf[..read]).await?;
+    }
+}
+
 pub(crate) async fn set_wifi_debug_mode(
     mode: WifiDebugMode,
     buffer_size: u32,
@@ -469,4 +497,19 @@ mod test {
         assert!(WifiBackend::try_from(2).is_err());
         assert!(WifiBackend::from_str("iwl").is_err());
     }
+
+    #[tokio::test]
+    async fn trace_extract() {
+        let h = testing::start();
+
+        fn process_output(_: &OsStr, _: &[&OsStr]) -> Result<(i32, String)> {
+            Ok((0, String::from("output")))
+        }
+        h.test.process_cb.set(process_output);
+
+        let pathbuf = extract_wifi_trace().await.unwrap();
+
+        assert_eq!(fs::read_to_string(&pathbuf).await.unwrap(), "output");
+        fs::remove_file(pathbuf).await.unwrap();
+    }
 }