GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/openvmm

Public

mirrored from https://github.com/microsoft/openvmmAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
release/2411-fork

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

openvmm/openhcl/diag_client/src

openhcl/diag_client/src/lib.rs

928lines · modecode

RawDownload
Latest commit unavailable.
1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! The client for connecting to the Underhill diagnostics server.
5
6#![warn(missing_docs)]
7
8pub mod kmsg_stream;
9
10use anyhow::Context;
11use diag_proto::network_packet_capture_request::OpData;
12use diag_proto::network_packet_capture_request::Operation;
13use diag_proto::ExecRequest;
14use diag_proto::WaitRequest;
15use diag_proto::WaitResponse;
16use futures::AsyncReadExt;
17use futures::AsyncWrite;
18use futures::AsyncWriteExt;
19use inspect::Node;
20use inspect::ValueKind;
21use inspect_proto::InspectResponse2;
22use inspect_proto::UpdateResponse2;
23use kmsg_stream::KmsgStream;
24use mesh_rpc::service::Status;
25use pal_async::driver::Driver;
26use pal_async::socket::PolledSocket;
27use pal_async::task::Spawn;
28use std::io::ErrorKind;
29use std::path::Path;
30use std::path::PathBuf;
31use std::time::Duration;
32use thiserror::Error;
33
34#[cfg(windows)]
35/// Functions for Hyper-V
36pub mod hyperv {
37 use super::ConnectError;
38 use anyhow::Context;
39 use guid::Guid;
40 use pal_async::driver::Driver;
41 use pal_async::socket::PolledSocket;
42 use pal_async::windows::pipe::NamedPipeServer;
43 use std::fs::File;
44 use std::io::Write;
45 use std::process::Command;
46 use std::time::Duration;
47 use vmsocket::VmAddress;
48 use vmsocket::VmSocket;
49 use vmsocket::VmStream;
50
51 /// Defines how to access the serial port
52 pub enum ComPortAccessInfo {
53 /// Access by number
54 PortNumber(u32),
55 /// Access through a named pipe
56 PortPipePath(String),
57 }
58
59 /// Get ID from name
60 pub fn vm_id_from_name(name: &str) -> anyhow::Result<Guid> {
61 let output = Command::new("hvc.exe")
62 .arg("id")
63 .arg(name)
64 .output()
65 .context("failed to launch hvc")?;
66
67 if output.status.success() {
68 let stdout = std::str::from_utf8(&output.stdout)
69 .context("failed to parse hvc output")?
70 .trim();
71 Ok(stdout
72 .parse()
73 .with_context(|| format!("failed to parse VM ID '{}'", &stdout))?)
74 } else {
75 anyhow::bail!(
76 "{}",
77 std::str::from_utf8(&output.stderr).context("failed to parse hvc error output")?
78 )
79 }
80 }
81
82 /// Connect to Hyper-V socket
83 pub async fn connect_vsock(
84 driver: &(impl Driver + ?Sized),
85 vm_id: Guid,
86 port: u32,
87 ) -> Result<VmStream, ConnectError> {
88 let socket = VmSocket::new()
89 .context("failed to create AF_HYPERV socket")
90 .map_err(ConnectError::other)?;
91
92 socket
93 .set_connect_timeout(Duration::from_secs(1))
94 .context("failed to set connect timeout")
95 .map_err(ConnectError::other)?;
96
97 socket
98 .set_high_vtl(true)
99 .context("failed to set socket for VTL2")
100 .map_err(ConnectError::other)?;
101
102 let mut socket: PolledSocket<socket2::Socket> = PolledSocket::new(driver, socket.into())
103 .context("failed to create polled socket")
104 .map_err(ConnectError::other)?;
105
106 socket
107 .connect(&VmAddress::hyperv_vsock(vm_id, port).into())
108 .await
109 .map_err(ConnectError::connect)?;
110
111 Ok(socket.convert().into_inner())
112 }
113
114 /// Opens a serial port on a Hyper-V VM.
115 ///
116 /// If the VM is not running, it will act as a server for the pipe. Hyper-V
117 /// will connect to the server once the VM starts running.
118 pub async fn open_serial_port(
119 driver: &(impl Driver + ?Sized),
120 vm: &str,
121 port: ComPortAccessInfo,
122 ) -> anyhow::Result<File> {
123 let path = match port {
124 ComPortAccessInfo::PortNumber(num) => {
125 let output = Command::new("powershell.exe")
126 .arg(format!(
127 r#"$x = Get-VMComPort "{vm}" -Number {num} -ErrorAction Stop; $x.Path"#,
128 ))
129 .output()
130 .context("failed to query VM com port")?;
131
132 if !output.status.success() {
133 let _ = std::io::stderr().write_all(&output.stderr);
134 anyhow::bail!(
135 "failed to query VM com port: exit status {}",
136 output.status.code().unwrap()
137 );
138 }
139 String::from_utf8(output.stdout)?
140 }
141 ComPortAccessInfo::PortPipePath(path) => path,
142 };
143
144 let path = path.trim();
145 if path.is_empty() {
146 anyhow::bail!("Requested VM COM port is not configured");
147 }
148
149 let pipe = match fs_err::OpenOptions::new().read(true).write(true).open(path) {
150 Ok(pipe) => pipe.into(),
151 Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
152 let server =
153 NamedPipeServer::create(path).context("failed to create pipe server")?;
154
155 server.accept(driver)?.await?
156 }
157 Err(err) => Err(err)?,
158 };
159
160 Ok(pipe)
161 }
162}
163
164/// Connect to a vsock with port and path
165pub async fn connect_hybrid_vsock(
166 driver: &(impl Driver + ?Sized),
167 path: &Path,
168 port: u32,
169) -> Result<PolledSocket<socket2::Socket>, ConnectError> {
170 let socket = unix_socket::UnixStream::connect(path).map_err(ConnectError::connect)?;
171 let mut socket = PolledSocket::new(driver, socket).map_err(ConnectError::other)?;
172 socket
173 .write_all(format!("CONNECT {port}\n").as_bytes())
174 .await
175 .map_err(ConnectError::other)?;
176
177 let mut ok = [0; 3];
178 socket
179 .read_exact(&mut ok)
180 .await
181 .map_err(ConnectError::other)?;
182 if &ok != b"OK " {
183 // FUTURE: consider returning an error that can be retried. This may
184 // require some changes to the hybrid vsock protocol, unclear.
185 return Err(ConnectError::other(anyhow::anyhow!(
186 "missing hybrid vsock response"
187 )));
188 }
189
190 for _ in 0.."4294967295\n".len() {
191 let mut b = [0];
192 socket
193 .read_exact(&mut b)
194 .await
195 .map_err(ConnectError::other)?;
196 if b[0] == b'\n' {
197 // Don't need to parse the host port number.
198 return Ok(socket.convert());
199 }
200 }
201 Err(ConnectError::other(anyhow::anyhow!(
202 "invalid hybrid vsock response"
203 )))
204}
205
206enum SocketType<'a> {
207 #[cfg(windows)]
208 VmId {
209 vm_id: guid::Guid,
210 port: u32,
211 },
212 HybridVsock {
213 path: &'a Path,
214 port: u32,
215 },
216}
217
218async fn new_data_connection(
219 driver: &(impl Driver + ?Sized),
220 typ: SocketType<'_>,
221) -> anyhow::Result<(u64, PolledSocket<socket2::Socket>)> {
222 let mut socket = match typ {
223 #[cfg(windows)]
224 SocketType::VmId { vm_id, port } => {
225 let socket = hyperv::connect_vsock(driver, vm_id, port).await?;
226 PolledSocket::new(driver, socket2::Socket::from(socket))?
227 }
228 SocketType::HybridVsock { path, port } => connect_hybrid_vsock(driver, path, port).await?,
229 };
230
231 // Read the 8 byte connection id which is always sent first on the connection.
232 let mut id = [0; 8];
233 socket
234 .read_exact(&mut id)
235 .await
236 .context("reading connection id")?;
237 let id = u64::from_ne_bytes(id);
238 Ok((id, socket))
239}
240
241/// Represents different VM types.
242enum VmType {
243 /// A Hyper-V VM represented by a VM ID GUID, which uses a VmSocket to connect.
244 #[cfg(windows)]
245 HyperV(guid::Guid),
246 /// A VM which uses hybrid vsock over Unix sockets.
247 HybridVsock(PathBuf),
248 /// A VM that cannot be used for data connections.
249 None,
250}
251
252/// The diagnostics client.
253pub struct DiagClient {
254 vm: VmType,
255 ttrpc: mesh_rpc::Client,
256 driver: Box<dyn Driver>,
257}
258
259/// Defines packet capture operations.
260#[derive(PartialEq)]
261pub enum PacketCaptureOperation {
262 /// Query details.
263 Query,
264 /// Start packet capture.
265 Start,
266 /// Stop packet capture.
267 Stop,
268}
269
270/// An error connecting to the diagnostics server.
271#[derive(Debug, Error)]
272#[error("failed to connect")]
273pub struct ConnectError {
274 #[source]
275 err: anyhow::Error,
276 kind: ConnectErrorKind,
277}
278
279#[derive(Debug)]
280enum ConnectErrorKind {
281 Other,
282 VmNotStarted,
283 ServerTimedOut,
284}
285
286impl ConnectError {
287 /// Returns the time to wait before retrying the connection. If `None`, the
288 /// connection should not be retried.
289 pub fn retry_timeout(&self) -> Option<Duration> {
290 match self.kind {
291 ConnectErrorKind::VmNotStarted => Some(Duration::from_secs(1)),
292 ConnectErrorKind::ServerTimedOut => {
293 // The socket infrastructure has an internal timeout.
294 Some(Duration::ZERO)
295 }
296 _ => None,
297 }
298 }
299
300 fn other(err: impl Into<anyhow::Error>) -> Self {
301 Self {
302 err: err.into(),
303 kind: ConnectErrorKind::Other,
304 }
305 }
306
307 fn connect(err: std::io::Error) -> Self {
308 let kind = match err.kind() {
309 ErrorKind::AddrNotAvailable => ConnectErrorKind::VmNotStarted,
310 ErrorKind::TimedOut => ConnectErrorKind::ServerTimedOut,
311 _ => match err.raw_os_error() {
312 #[cfg(windows)]
313 Some(windows_sys::Win32::Networking::WinSock::WSAENETUNREACH) => {
314 ConnectErrorKind::VmNotStarted
315 }
316 _ => ConnectErrorKind::Other,
317 },
318 };
319 Self {
320 err: anyhow::Error::from(err).context("failed to connect"),
321 kind,
322 }
323 }
324}
325
326impl DiagClient {
327 /// Creates a client from Hyper-V VM name.
328 #[cfg(windows)]
329 pub async fn from_hyperv_name(
330 driver: impl Driver + Spawn,
331 name: &str,
332 ) -> Result<Self, ConnectError> {
333 Self::from_hyperv_id(
334 driver,
335 hyperv::vm_id_from_name(name).map_err(ConnectError::other)?,
336 )
337 .await
338 }
339
340 /// Creates a client from a Hyper-V or HCS VM ID.
341 #[cfg(windows)]
342 pub async fn from_hyperv_id(
343 driver: impl Driver + Spawn,
344 vm_id: guid::Guid,
345 ) -> Result<Self, ConnectError> {
346 let socket = hyperv::connect_vsock(&driver, vm_id, diag_proto::VSOCK_CONTROL_PORT).await?;
347 Ok(Self {
348 vm: VmType::HyperV(vm_id),
349 ttrpc: mesh_rpc::Client::new(&driver, socket),
350 driver: Box::new(driver),
351 })
352 }
353
354 /// Creates a client from a Unix socket path.
355 pub async fn from_socket(
356 driver: impl Driver + Spawn,
357 path: &Path,
358 ) -> Result<Self, ConnectError> {
359 let socket = unix_socket::UnixStream::connect(path).map_err(ConnectError::other)?;
360 Ok(Self {
361 vm: VmType::None,
362 ttrpc: mesh_rpc::Client::new(&driver, socket),
363 driver: Box::new(driver),
364 })
365 }
366
367 /// Creates a client from a hybrid vsock Unix socket path.
368 pub async fn from_hybrid_vsock(
369 driver: impl Driver + Spawn,
370 path: &Path,
371 ) -> Result<Self, ConnectError> {
372 let socket = connect_hybrid_vsock(&driver, path, diag_proto::VSOCK_CONTROL_PORT)
373 .await?
374 .into_inner();
375 Ok(Self {
376 vm: VmType::HybridVsock(path.into()),
377 ttrpc: mesh_rpc::Client::new(&driver, socket),
378 driver: Box::new(driver),
379 })
380 }
381
382 /// Creates a client from an existing connection.
383 ///
384 /// This client won't be usable with operations that require additional connections.
385 pub fn from_conn<T>(driver: impl Driver + Spawn, conn: T) -> Self
386 where
387 T: 'static + Send + Sync + pal_async::socket::AsSockRef + std::io::Read + std::io::Write,
388 {
389 Self {
390 vm: VmType::None,
391 ttrpc: mesh_rpc::Client::new(&driver, conn),
392 driver: Box::new(driver),
393 }
394 }
395
396 /// Creates a builder for execing a command.
397 pub fn exec(&self, command: impl AsRef<str>) -> ExecBuilder<'_> {
398 ExecBuilder {
399 client: self,
400 with_stdin: false,
401 with_stdout: false,
402 with_stderr: false,
403 request: ExecRequest {
404 command: command.as_ref().to_owned(),
405 ..Default::default()
406 },
407 }
408 }
409
410 /// Creates a new data connection socket.
411 ///
412 /// This can be used with [`DiagClient::custom_call`].
413 pub async fn connect_data(&self) -> anyhow::Result<(u64, PolledSocket<socket2::Socket>)> {
414 let socket_type = match &self.vm {
415 #[cfg(windows)]
416 VmType::HyperV(guid) => SocketType::VmId {
417 vm_id: *guid,
418 port: diag_proto::VSOCK_DATA_PORT,
419 },
420 VmType::HybridVsock(path) => SocketType::HybridVsock {
421 path,
422 port: diag_proto::VSOCK_DATA_PORT,
423 },
424 VmType::None => {
425 anyhow::bail!("cannot make additional connections with this client")
426 }
427 };
428 new_data_connection(self.driver.as_ref(), socket_type).await
429 }
430
431 /// Sends an inspection request to the server.
432 pub async fn inspect(
433 &self,
434 path: impl Into<String>,
435 depth: Option<usize>,
436 timeout: Option<Duration>,
437 ) -> anyhow::Result<Node> {
438 let response = self.ttrpc.start_call(
439 inspect_proto::InspectService::Inspect,
440 inspect_proto::InspectRequest {
441 path: path.into(),
442 // It would be better to pass an Option<u32> in the proto, but that would break backcompat.
443 depth: depth.unwrap_or(u32::MAX as usize) as u32,
444 },
445 timeout,
446 );
447
448 let response = response.upcast::<Result<InspectResponse2, Status>>();
449 let response = response.await?.map_err(grpc_status)?;
450
451 Ok(response.result)
452 }
453
454 /// Updates an inspectable value.
455 pub async fn update(
456 &self,
457 path: impl Into<String>,
458 value: impl Into<String>,
459 ) -> anyhow::Result<inspect::Value> {
460 let response = self.ttrpc.start_call(
461 inspect_proto::InspectService::Update,
462 inspect_proto::UpdateRequest {
463 path: path.into(),
464 value: value.into(),
465 },
466 None,
467 );
468
469 let response = response.upcast::<Result<UpdateResponse2, Status>>();
470 let response = response.await?.map_err(grpc_status)?;
471
472 Ok(response.new_value)
473 }
474
475 /// Get PID of a given process
476 pub async fn get_pid(&self, name: &str) -> anyhow::Result<i32> {
477 let hosts = self.inspect("mesh/hosts", Some(1), None).await?;
478 let mut plist = Vec::new();
479
480 let Node::Dir(processes) = hosts else {
481 anyhow::bail!("Hosts node is not a dir");
482 };
483 for process in processes {
484 let Node::Dir(pnode) = process.node else {
485 anyhow::bail!("Process node is not a dir");
486 };
487 for entry in pnode {
488 if entry.name == "name" {
489 let Node::Value(value) = entry.node else {
490 anyhow::bail!("Name node is not a value");
491 };
492 let ValueKind::String(strval) = value.kind else {
493 anyhow::bail!("Name node is not a string");
494 };
495 if strval == name {
496 return Ok(process.name.parse()?);
497 }
498 plist.push(strval);
499 }
500 }
501 }
502
503 anyhow::bail!("PID of {name} not found. Processes: {:?}", plist)
504 }
505
506 /// Starts the VM.
507 pub async fn start(
508 &self,
509 env: impl IntoIterator<Item = (String, Option<String>)>,
510 args: impl IntoIterator<Item = String>,
511 ) -> anyhow::Result<()> {
512 let request = diag_proto::StartRequest {
513 env: env
514 .into_iter()
515 .map(|(name, value)| diag_proto::EnvPair { name, value })
516 .collect(),
517 args: args.into_iter().collect(),
518 };
519 self.ttrpc
520 .call(diag_proto::UnderhillDiag::Start, request)
521 .await?
522 .map_err(grpc_status)?;
523
524 Ok(())
525 }
526
527 /// Gets the contents of /dev/kmsg
528 pub async fn kmsg(&self, follow: bool) -> anyhow::Result<KmsgStream> {
529 let (conn, socket) = self.connect_data().await?;
530
531 self.ttrpc
532 .call(
533 diag_proto::UnderhillDiag::Kmsg,
534 diag_proto::KmsgRequest { follow, conn },
535 )
536 .await?
537 .map_err(grpc_status)?;
538
539 Ok(KmsgStream::new(socket))
540 }
541
542 /// Gets the contents of the file
543 pub async fn read_file(
544 &self,
545 follow: bool,
546 file_path: String,
547 ) -> anyhow::Result<PolledSocket<socket2::Socket>> {
548 let (conn, socket) = self.connect_data().await?;
549
550 self.ttrpc
551 .call(
552 diag_proto::UnderhillDiag::ReadFile,
553 diag_proto::FileRequest {
554 follow,
555 conn,
556 file_path,
557 },
558 )
559 .await?
560 .map_err(grpc_status)?;
561
562 Ok(socket)
563 }
564
565 /// Issues a call to the server using a custom RPC.
566 ///
567 /// This can be used to support extension RPCs that are not part of the main
568 /// diagnostics service.
569 pub fn custom_call<F, R, T, U>(
570 &self,
571 rpc: F,
572 input: T,
573 timeout: Option<Duration>,
574 ) -> mesh::OneshotReceiver<Result<U, Status>>
575 where
576 F: FnOnce(T, mesh::OneshotSender<Result<U, Status>>) -> R,
577 R: mesh_rpc::service::ServiceRpc,
578 U: mesh::MeshPayload,
579 {
580 self.ttrpc.start_call(rpc, input, timeout)
581 }
582
583 /// Crashes the VM.
584 pub async fn crash(&self, pid: i32) -> anyhow::Result<()> {
585 self.ttrpc
586 .call(
587 diag_proto::UnderhillDiag::Crash,
588 diag_proto::CrashRequest { pid },
589 )
590 .await?
591 .map_err(grpc_status)?;
592
593 Ok(())
594 }
595
596 /// Sets up network packet capture trace.
597 pub async fn packet_capture(
598 &self,
599 op: PacketCaptureOperation,
600 num_streams: u32,
601 snaplen: u16,
602 ) -> anyhow::Result<(Vec<PolledSocket<socket2::Socket>>, u32)> {
603 let mut sockets = Vec::new();
604 let op_data = match op {
605 PacketCaptureOperation::Start => {
606 let mut conns = Vec::new();
607 for _ in 0..num_streams {
608 let (conn, socket) = self.connect_data().await?;
609 conns.push(conn);
610 sockets.push(socket);
611 }
612 Some(OpData::StartData(diag_proto::StartPacketCaptureData {
613 snaplen: snaplen.into(),
614 conns,
615 }))
616 }
617 _ => None,
618 };
619
620 let operation = match op {
621 PacketCaptureOperation::Query => Operation::Query,
622 PacketCaptureOperation::Start => Operation::Start,
623 PacketCaptureOperation::Stop => Operation::Stop,
624 };
625
626 let response = self
627 .ttrpc
628 .call(
629 diag_proto::UnderhillDiag::PacketCapture,
630 diag_proto::NetworkPacketCaptureRequest {
631 operation: operation.into(),
632 op_data,
633 },
634 )
635 .await?
636 .map_err(grpc_status)?;
637
638 Ok((sockets, response.num_streams))
639 }
640
641 /// Saves a core dump file being streamed from Underhill
642 pub async fn core_dump(
643 &self,
644 pid: i32,
645 mut writer: impl AsyncWrite + Unpin,
646 mut stderr: impl AsyncWrite + Unpin,
647 verbose: bool,
648 ) -> anyhow::Result<()> {
649 // Launch hcl-dump to dump the target process. Use raw_socket_io so that
650 // the diagnostics process does not have to be running during the core
651 // dump process; this ensures that we can dump the diagnostics process,
652 // too.
653 let mut process = self.exec("/bin/underhill-dump");
654 if verbose {
655 process.args(["-v"]);
656 }
657 let mut process = process
658 .args([pid.to_string()])
659 .stdin(false)
660 .stdout(true)
661 .stderr(true)
662 .raw_socket_io(true)
663 .spawn()
664 .await
665 .context("failed to launch underhill-dump")?;
666
667 let process_stdout = PolledSocket::new(&self.driver, process.stdout.take().unwrap())?;
668 let process_stderr = PolledSocket::new(&self.driver, process.stderr.take().unwrap())?;
669
670 let out = futures::io::copy(process_stdout, &mut writer);
671 let err = futures::io::copy(process_stderr, &mut stderr);
672
673 futures::try_join!(out, err)?;
674
675 let status = process
676 .wait()
677 .await
678 .context("failed to wait for underhill-dump")?;
679
680 if !status.success() {
681 anyhow::bail!(
682 "underhill-dump failed with exit code {}",
683 status.exit_code()
684 );
685 }
686 Ok(())
687 }
688
689 /// Restarts the Underhill worker.
690 pub async fn restart(&self) -> anyhow::Result<()> {
691 self.ttrpc
692 .call(diag_proto::UnderhillDiag::Restart, ())
693 .await?
694 .map_err(grpc_status)?;
695
696 Ok(())
697 }
698
699 /// Pause the VM (including all devices).
700 pub async fn pause(&self) -> anyhow::Result<()> {
701 self.ttrpc
702 .call(diag_proto::UnderhillDiag::Pause, ())
703 .await?
704 .map_err(grpc_status)?;
705
706 Ok(())
707 }
708
709 /// Resume the VM.
710 pub async fn resume(&self) -> anyhow::Result<()> {
711 self.ttrpc
712 .call(diag_proto::UnderhillDiag::Resume, ())
713 .await?
714 .map_err(grpc_status)?;
715
716 Ok(())
717 }
718
719 /// Dumps the VM's VTL2 saved state.
720 pub async fn dump_saved_state(&self) -> anyhow::Result<Vec<u8>> {
721 let state = self
722 .ttrpc
723 .call(diag_proto::UnderhillDiag::DumpSavedState, ())
724 .await?
725 .map_err(grpc_status)?;
726
727 Ok(state.data)
728 }
729}
730
731fn grpc_status(status: Status) -> anyhow::Error {
732 anyhow::anyhow!(status.message)
733}
734
735/// A builder for launching a command in VTL2.
736pub struct ExecBuilder<'a> {
737 client: &'a DiagClient,
738 with_stdin: bool,
739 with_stdout: bool,
740 with_stderr: bool,
741 request: ExecRequest,
742}
743
744impl ExecBuilder<'_> {
745 /// Adds `args` to the argument list.
746 pub fn args<T: AsRef<str>>(&mut self, args: impl IntoIterator<Item = T>) -> &mut Self {
747 self.request
748 .args
749 .extend(args.into_iter().map(|s| s.as_ref().to_owned()));
750 self
751 }
752
753 /// Sets whether the process is spawned with a TTY.
754 pub fn tty(&mut self, tty: bool) -> &mut Self {
755 self.request.tty = tty;
756 self
757 }
758
759 /// Specifies whether a stdin socket should be opened.
760 pub fn stdin(&mut self, stdin: bool) -> &mut Self {
761 self.with_stdin = stdin;
762 self
763 }
764
765 /// Specifies whether a stdout socket should be opened.
766 pub fn stdout(&mut self, stdout: bool) -> &mut Self {
767 self.with_stdout = stdout;
768 self
769 }
770
771 /// Specifies whether a stderr socket should be opened.
772 pub fn stderr(&mut self, stderr: bool) -> &mut Self {
773 self.with_stderr = stderr;
774 self
775 }
776
777 /// Specifies whether the processes's stdout and stderr should be combined
778 /// into a single stream (the stdout socket).
779 pub fn combine_stderr(&mut self, combine_stderr: bool) -> &mut Self {
780 self.request.combine_stderr = combine_stderr;
781 self
782 }
783
784 /// Specifies whether the vsock sockets used for stdio should be passed
785 /// directly to the launched process instead of going through relays.
786 pub fn raw_socket_io(&mut self, raw_socket_io: bool) -> &mut Self {
787 self.request.raw_socket_io = raw_socket_io;
788 self
789 }
790
791 /// Clears the default environment.
792 pub fn env_clear(&mut self) -> &mut Self {
793 self.request.clear_env = true;
794 self
795 }
796
797 /// Removes an environment variable.
798 pub fn env_remove(&mut self, name: impl AsRef<str>) -> &mut Self {
799 self.request.env.push(diag_proto::EnvPair {
800 name: name.as_ref().to_owned(),
801 value: None,
802 });
803 self
804 }
805
806 /// Sets an environment variable.
807 pub fn env(&mut self, name: impl AsRef<str>, value: impl AsRef<str>) -> &mut Self {
808 self.request.env.push(diag_proto::EnvPair {
809 name: name.as_ref().to_owned(),
810 value: Some(value.as_ref().to_owned()),
811 });
812 self
813 }
814
815 /// Spawns the process.
816 pub async fn spawn(&self) -> anyhow::Result<Process> {
817 let mut request = self.request.clone();
818
819 let stdin = if self.with_stdin {
820 let (id, stdin) = self
821 .client
822 .connect_data()
823 .await
824 .context("failed to connect stdin")?;
825 request.stdin = id;
826
827 Some(stdin.into_inner())
828 } else {
829 None
830 };
831
832 let stdout = if self.with_stdout {
833 let (id, stdout) = self
834 .client
835 .connect_data()
836 .await
837 .context("failed to connect stdout")?;
838 request.stdout = id;
839
840 Some(stdout.into_inner())
841 } else {
842 None
843 };
844
845 let stderr = if self.with_stdout {
846 let (id, stderr) = self
847 .client
848 .connect_data()
849 .await
850 .context("failed to connect stderr")?;
851 request.stderr = id;
852
853 Some(stderr.into_inner())
854 } else {
855 None
856 };
857
858 let response = self
859 .client
860 .ttrpc
861 .call(diag_proto::UnderhillDiag::Exec, request)
862 .await?
863 .map_err(grpc_status)?;
864
865 let wait = self.client.ttrpc.start_call(
866 diag_proto::UnderhillDiag::Wait,
867 WaitRequest { pid: response.pid },
868 None,
869 );
870
871 Ok(Process {
872 stdin,
873 stdout,
874 stderr,
875 wait,
876 pid: response.pid,
877 })
878 }
879}
880
881/// A process running in VTL2.
882#[derive(Debug)]
883pub struct Process {
884 /// The standard input stream.
885 pub stdin: Option<socket2::Socket>,
886 /// The standard output stream.
887 pub stdout: Option<socket2::Socket>,
888 /// The standard error stream.
889 pub stderr: Option<socket2::Socket>,
890 pid: i32,
891 wait: mesh::OneshotReceiver<Result<WaitResponse, Status>>,
892}
893
894impl Process {
895 /// Returns the process ID.
896 pub fn id(&self) -> i32 {
897 self.pid
898 }
899
900 /// Waits for the process to exit.
901 pub async fn wait(self) -> anyhow::Result<ExitStatus> {
902 let response = self
903 .wait
904 .await
905 .context("disconnected")?
906 .map_err(|err| anyhow::anyhow!("{}", err.message))?;
907
908 Ok(ExitStatus { response })
909 }
910}
911
912/// Process exit status.
913#[derive(Debug)]
914pub struct ExitStatus {
915 response: WaitResponse,
916}
917
918impl ExitStatus {
919 /// The exit code.
920 pub fn exit_code(&self) -> i32 {
921 self.response.exit_code
922 }
923
924 /// Whether the process successfully terminated.
925 pub fn success(&self) -> bool {
926 self.response.exit_code == 0
927 }
928}
929