steamos-manager/src/sls/ftrace.rs

/* SPDX-License-Identifier: BSD-2-Clause */
use anyhow::{Error, Result};
use std::collections::HashMap;
use std::fmt::Debug;
use std::path::Path;
use tokio::fs;
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::net::unix::pipe;
use tracing::{error, info};
use zbus::connection::Connection;
use zbus::zvariant;

use crate::{get_appid, read_comm, sysbase, Service};

#[zbus::proxy(
    interface = "com.steampowered.SteamOSLogSubmitter.Trace",
    default_service = "com.steampowered.SteamOSLogSubmitter",
    default_path = "/com/steampowered/SteamOSLogSubmitter/helpers/Trace"
)]
trait TraceHelper {
    async fn log_event(
        &self,
        trace: &str,
        data: HashMap<&str, zvariant::Value<'_>>,
    ) -> zbus::Result<()>;
}

pub struct Ftrace
where
    Self: 'static,
{
    pipe: Option<BufReader<pipe::Receiver>>,
    proxy: TraceHelperProxy<'static>,
}

async fn setup_traces(path: &Path) -> Result<()> {
    fs::write(path.join("events/oom/mark_victim/enable"), "1").await?;
    fs::write(path.join("set_ftrace_filter"), "split_lock_warn").await?;
    fs::write(path.join("current_tracer"), "function").await?;
    Ok(())
}

impl Ftrace {
    pub async fn init(connection: Connection) -> Result<Ftrace> {
        let base = Self::base();
        let path = Path::new(base.as_str());
        fs::create_dir_all(path).await?;
        setup_traces(path).await?;
        let file = pipe::OpenOptions::new()
            .unchecked(true) // Thanks tracefs for making trace_pipe a "regular" file
            .open_receiver(path.join("trace_pipe"))?;
        Ok(Ftrace {
            pipe: Some(BufReader::new(file)),
            proxy: TraceHelperProxy::new(&connection).await?,
        })
    }

    fn base() -> String {
        sysbase() + "/sys/kernel/tracing/instances/steamos-log-submitter"
    }

    async fn handle_pid(data: &mut HashMap<&str, zvariant::Value<'_>>, pid: u32) -> Result<()> {
        if let Ok(comm) = read_comm(pid) {
            info!("├─ comm: {}", comm);
            data.insert("comm", zvariant::Value::new(comm));
        } else {
            info!("├─ comm not found");
        }
        if let Ok(Some(appid)) = get_appid(pid) {
            info!("└─ appid: {}", appid);
            data.insert("appid", zvariant::Value::new(appid));
        } else {
            info!("└─ appid not found");
        }
        Ok(())
    }

    async fn handle_event(&mut self, line: &str) -> Result<()> {
        info!("Forwarding line {}", line);
        let mut data = HashMap::new();
        let mut split = line.rsplit(' ');
        if let Some(("pid", pid)) = split.next().and_then(|arg| arg.split_once('=')) {
            let pid = pid.parse()?;
            Ftrace::handle_pid(&mut data, pid).await?;
        }
        self.proxy.log_event(line, data).await?;
        Ok(())
    }
}

impl Service for Ftrace {
    const NAME: &'static str = "ftrace";

    async fn run(&mut self) -> Result<()> {
        loop {
            let mut string = String::new();
            self.pipe
                .as_mut()
                .ok_or(Error::msg("BUG: trace_pipe missing"))?
                .read_line(&mut string)
                .await?;
            if let Err(e) = self.handle_event(string.trim_end()).await {
                error!("Encountered an error handling event: {}", e);
            }
        }
    }

    async fn shutdown(&mut self) -> Result<()> {
        self.pipe.take();
        fs::remove_dir(Self::base()).await?;
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::testing;
    use nix::sys::stat::Mode;
    use nix::unistd;
    use std::cell::Cell;
    use std::fs;
    use std::path::PathBuf;
    use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender};

    struct MockTrace {
        traces: UnboundedSender<(String, HashMap<String, zvariant::OwnedValue>)>,
    }

    #[zbus::interface(name = "com.steampowered.SteamOSLogSubmitter.Trace")]
    impl MockTrace {
        fn log_event(
            &mut self,
            trace: &str,
            data: HashMap<&str, zvariant::Value<'_>>,
        ) -> zbus::fdo::Result<()> {
            self.traces.send((
                String::from(trace),
                HashMap::from_iter(
                    data.iter()
                        .map(|(k, v)| (String::from(*k), v.try_to_owned().unwrap())),
                ),
            ));
            Ok(())
        }
    }

    #[tokio::test]
    async fn handle_pid() {
        let h = testing::start();
        let path = h.test.path();

        fs::create_dir_all(path.join("proc/1234")).expect("create_dir_all");
        fs::write(path.join("proc/1234/comm"), "ftrace\n").expect("write comm");
        fs::write(path.join("proc/1234/environ"), "SteamGameId=5678").expect("write environ");

        fs::create_dir_all(path.join("proc/1235")).expect("create_dir_all");
        fs::write(path.join("proc/1235/comm"), "ftrace\n").expect("write comm");

        fs::create_dir_all(path.join("proc/1236")).expect("create_dir_all");
        fs::write(path.join("proc/1236/environ"), "SteamGameId=5678").expect("write environ");

        let mut map = HashMap::new();
        assert!(Ftrace::handle_pid(&mut map, 1234).await.is_ok());
        assert_eq!(
            *map.get("comm").expect("comm"),
            zvariant::Value::new("ftrace")
        );
        assert_eq!(
            *map.get("appid").expect("appid"),
            zvariant::Value::new(5678 as u64)
        );

        let mut map = HashMap::new();
        assert!(Ftrace::handle_pid(&mut map, 1235).await.is_ok());
        assert_eq!(
            *map.get("comm").expect("comm"),
            zvariant::Value::new("ftrace")
        );
        assert!(map.get("appid").is_none());

        let mut map = HashMap::new();
        assert!(Ftrace::handle_pid(&mut map, 1236).await.is_ok());
        assert!(map.get("comm").is_none());
        assert_eq!(
            *map.get("appid").expect("appid"),
            zvariant::Value::new(5678 as u64)
        );
    }

    #[tokio::test]
    async fn ftrace_init() {
        let h = testing::start();
        let path = h.test.path();

        let tracefs = PathBuf::from(Ftrace::base());

        fs::create_dir_all(tracefs.join("events/oom/mark_victim")).expect("create_dir_all");
        unistd::mkfifo(
            tracefs.join("trace_pipe").as_path(),
            Mode::S_IRUSR | Mode::S_IWUSR,
        )
        .expect("trace_pipe");
        let dbus = Connection::session().await.expect("dbus");
        let ftrace = Ftrace::init(dbus).await.expect("ftrace");

        assert_eq!(
            fs::read_to_string(tracefs.join("events/oom/mark_victim/enable")).unwrap(),
            "1"
        );
    }

    #[tokio::test]
    async fn ftrace_relay() {
        let h = testing::start();
        let path = h.test.path();

        let tracefs = PathBuf::from(Ftrace::base());

        fs::create_dir_all(tracefs.join("events/oom/mark_victim")).expect("create_dir_all");
        unistd::mkfifo(
            tracefs.join("trace_pipe").as_path(),
            Mode::S_IRUSR | Mode::S_IWUSR,
        )
        .expect("trace_pipe");

        fs::create_dir_all(path.join("proc/14351")).expect("create_dir_all");
        fs::write(path.join("proc/14351/comm"), "ftrace\n").expect("write comm");
        fs::write(path.join("proc/14351/environ"), "SteamGameId=5678").expect("write environ");

        let (sender, mut receiver) = unbounded_channel();
        let trace = MockTrace { traces: sender };
        let dbus = zbus::connection::Builder::session()
            .unwrap()
            .name("com.steampowered.SteamOSLogSubmitter")
            .unwrap()
            .serve_at("/com/steampowered/SteamOSLogSubmitter/helpers/Trace", trace)
            .unwrap()
            .build()
            .await
            .expect("dbus");
        let mut ftrace = Ftrace::init(dbus).await.expect("ftrace");

        assert!(match receiver.try_recv() {
            Empty => true,
            _ => false,
        });
        ftrace
            .handle_event(
                " GamepadUI Input-4886    [003] .N.1. 23828.572941: mark_victim: pid=14351",
            )
            .await
            .expect("event");
        let (line, data) = match receiver.try_recv() {
            Ok((line, data)) => (line, data),
            _ => panic!("Test failed"),
        };
        assert_eq!(
            line,
            " GamepadUI Input-4886    [003] .N.1. 23828.572941: mark_victim: pid=14351"
        );
        assert_eq!(data.len(), 2);
        assert_eq!(
            data.get("appid").map(|v| v.downcast_ref()),
            Some(Ok(5678 as u64))
        );
        assert_eq!(
            data.get("comm").map(|v| v.downcast_ref()),
            Some(Ok("ftrace"))
        );

        ftrace
            .handle_event(" GamepadUI Input-4886    [003] .N.1. 23828.572941: split_lock_warn <-")
            .await
            .expect("event");
        let (line, data) = match receiver.try_recv() {
            Ok((line, data)) => (line, data),
            _ => panic!("Test failed"),
        };
        assert_eq!(
            line,
            " GamepadUI Input-4886    [003] .N.1. 23828.572941: split_lock_warn <-"
        );
        assert_eq!(data.len(), 0);
    }
}