petri/
openhcl_diag.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4use anyhow::Context;
5use diag_client::DiagClient;
6use diag_client::ExitStatus;
7use diag_client::kmsg_stream::KmsgStream;
8use futures::io::AllowStdIo;
9use std::io::Read;
10
11pub struct OpenHclDiagHandler(DiagClient);
12
13/// The result of running a VTL2 command.
14#[derive(Debug)]
15pub struct Vtl2CommandResult {
16    /// The stdout of the command.
17    pub stdout: String,
18    /// The stderr of the command.
19    pub stderr: String,
20    /// The raw stdout of the command.
21    pub stdout_raw: Vec<u8>,
22    /// The raw stderr of the command.
23    pub stderr_raw: Vec<u8>,
24    /// The exit status of the command.
25    pub exit_status: ExitStatus,
26}
27
28impl OpenHclDiagHandler {
29    pub(crate) fn new(client: DiagClient) -> Self {
30        Self(client)
31    }
32
33    pub(crate) async fn wait_for_vtl2(&self) -> anyhow::Result<()> {
34        self.0.wait_for_server().await
35    }
36
37    pub(crate) async fn run_vtl2_command(
38        &self,
39        command: impl AsRef<str>,
40        args: impl IntoIterator<Item = impl AsRef<str>>,
41    ) -> anyhow::Result<Vtl2CommandResult> {
42        let client = self.diag_client().await?;
43        let mut proc = client
44            .exec(command.as_ref())
45            .args(args)
46            .stdout(true)
47            .stderr(true)
48            .raw_socket_io(true)
49            .spawn()
50            .await?;
51
52        let (mut stdout, mut stderr) = (proc.stdout.take().unwrap(), proc.stderr.take().unwrap());
53        let exit_status = proc.wait().await?;
54
55        let mut stdout_buf = Vec::new();
56        stdout
57            .read_to_end(&mut stdout_buf)
58            .context("error reading stdout socket")?;
59        let stdout_str = String::from_utf8_lossy(&stdout_buf);
60
61        let mut stderr_buf = Vec::new();
62        stderr
63            .read_to_end(&mut stderr_buf)
64            .context("error reading stderr socket")?;
65        let stderr_str = String::from_utf8_lossy(&stderr_buf);
66
67        Ok(Vtl2CommandResult {
68            stdout: stdout_str.to_string(),
69            stderr: stderr_str.to_string(),
70            stdout_raw: stdout_buf,
71            stderr_raw: stderr_buf,
72            exit_status,
73        })
74    }
75
76    pub async fn core_dump(&self, name: &str, path: &std::path::Path) -> anyhow::Result<()> {
77        let client = self.diag_client().await?;
78        let pid = client.get_pid(name).await?;
79        client
80            .core_dump(
81                pid,
82                AllowStdIo::new(fs_err::File::create(path)?),
83                AllowStdIo::new(std::io::stderr()),
84                true,
85            )
86            .await
87    }
88
89    pub async fn crash(&self, name: &str) -> anyhow::Result<()> {
90        let client = self.diag_client().await?;
91        let pid = client.get_pid(name).await?;
92        client.crash(pid).await
93    }
94
95    pub async fn inspect(
96        &self,
97        path: impl Into<String>,
98        depth: Option<usize>,
99        timeout: Option<std::time::Duration>,
100    ) -> anyhow::Result<inspect::Node> {
101        self.diag_client()
102            .await?
103            .inspect(path, depth, timeout)
104            .await
105    }
106
107    pub async fn kmsg(&self) -> anyhow::Result<KmsgStream> {
108        self.diag_client().await?.kmsg(false).await
109    }
110
111    async fn diag_client(&self) -> anyhow::Result<&DiagClient> {
112        self.wait_for_vtl2().await?;
113        Ok(&self.0)
114    }
115}