GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/openvmm

Public

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

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
2fa7fcaafa8f12253f263a3868985da4e52a7c73

Branches

Tags

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

Clone

HTTPS

Download ZIP

openvmm/flowey/flowey_core/src

flowey/flowey_core/src/shell.rs

559lines · modecode

RawDownload
Latest commit unavailable.
1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Shell abstraction for flowey command execution.
5//!
6//! Provides [`FloweyShell`] and [`FloweyCmd`] as thin wrappers around
7//! [`xshell::Shell`] and [`xshell::Cmd`] that enable command
8//! wrapping (e.g., running commands inside `nix-shell --pure --run`).
9
10use std::ffi::OsStr;
11use std::ffi::OsString;
12use std::ops::Deref;
13use std::process::Output;
14
15use serde::Deserialize;
16use serde::Serialize;
17
18/// A wrapper around [`xshell::Shell`] that supports transparent command
19/// wrapping via an optional [`CommandWrapperKind`].
20///
21/// Implements [`Deref<Target = xshell::Shell>`] so that existing usages like
22/// `rt.sh.change_dir()` and `rt.sh.set_var()` continue to work unchanged.
23pub struct FloweyShell {
24 inner: xshell::Shell,
25 wrapper: Option<CommandWrapperKind>,
26}
27
28impl FloweyShell {
29 /// Create a new `FloweyShell` with no command wrapper.
30 #[allow(clippy::disallowed_methods)]
31 pub fn new() -> anyhow::Result<Self> {
32 Ok(Self {
33 inner: xshell::Shell::new()?,
34 wrapper: None,
35 })
36 }
37
38 /// Set (or clear) the command wrapper used for all commands created
39 /// through this shell.
40 pub fn set_wrapper(&mut self, wrapper: Option<CommandWrapperKind>) {
41 self.wrapper = wrapper;
42 }
43
44 /// Access the underlying [`xshell::Shell`].
45 ///
46 /// This is primarily used by the [`shell_cmd!`](crate::shell_cmd)
47 /// macro to pass the shell reference into [`xshell::cmd!`].
48 pub fn xshell(&self) -> &xshell::Shell {
49 &self.inner
50 }
51
52 /// Wrap an [`xshell::Cmd`] into a [`FloweyCmd`] that will apply
53 /// this shell's [`CommandWrapperKind`] (if any) at execution time.
54 pub fn wrap<'a>(&'a self, cmd: xshell::Cmd<'a>) -> FloweyCmd<'a> {
55 FloweyCmd {
56 inner: cmd,
57 env_changes: Vec::new(),
58 stdin_contents: None,
59 ignore_status: false,
60 quiet: false,
61 secret: false,
62 ignore_stdout: false,
63 ignore_stderr: false,
64 wrapper: self.wrapper.clone(),
65 sh: &self.inner,
66 }
67 }
68}
69
70impl Deref for FloweyShell {
71 type Target = xshell::Shell;
72
73 fn deref(&self) -> &xshell::Shell {
74 &self.inner
75 }
76}
77
78/// Environment variable changes tracked by [`FloweyCmd`].
79enum EnvChange {
80 Set(OsString, OsString),
81 Remove(OsString),
82 Clear,
83}
84
85/// A wrapper around [`xshell::Cmd`] that applies a [`CommandWrapperKind`]
86/// at execution time.
87///
88/// Builder methods (`.arg()`, `.env()`, etc.) are accumulated on the
89/// inner [`xshell::Cmd`] (for args) or in shadow fields (for env, stdin,
90/// and flags). Execution methods (`.run()`, `.read()`, etc.) consume
91/// `self`, apply the wrapper transformation, re-apply the shadowed state,
92/// and then execute.
93pub struct FloweyCmd<'a> {
94 /// The inner command accumulates program + arguments only.
95 inner: xshell::Cmd<'a>,
96 // Shadow fields for state that must survive wrapping.
97 env_changes: Vec<EnvChange>,
98 stdin_contents: Option<Vec<u8>>,
99 ignore_status: bool,
100 quiet: bool,
101 secret: bool,
102 ignore_stdout: bool,
103 ignore_stderr: bool,
104 wrapper: Option<CommandWrapperKind>,
105 sh: &'a xshell::Shell,
106}
107
108// Mirrors xshell::Cmd's builder methods, but xshell doesn't export a common trait to implement
109impl<'a> FloweyCmd<'a> {
110 /// Adds an argument to this command.
111 pub fn arg<P: AsRef<OsStr>>(mut self, arg: P) -> Self {
112 self.inner = self.inner.arg(arg);
113 self
114 }
115
116 /// Adds all of the arguments to this command.
117 pub fn args<I>(mut self, args: I) -> Self
118 where
119 I: IntoIterator,
120 I::Item: AsRef<OsStr>,
121 {
122 self.inner = self.inner.args(args);
123 self
124 }
125
126 /// Overrides the value of an environmental variable for this command.
127 pub fn env<K: AsRef<OsStr>, V: AsRef<OsStr>>(mut self, key: K, val: V) -> Self {
128 self.env_changes.push(EnvChange::Set(
129 key.as_ref().to_owned(),
130 val.as_ref().to_owned(),
131 ));
132 self
133 }
134
135 /// Overrides the values of specified environmental variables for this
136 /// command.
137 pub fn envs<I, K, V>(mut self, vars: I) -> Self
138 where
139 I: IntoIterator<Item = (K, V)>,
140 K: AsRef<OsStr>,
141 V: AsRef<OsStr>,
142 {
143 for (k, v) in vars {
144 self.env_changes
145 .push(EnvChange::Set(k.as_ref().to_owned(), v.as_ref().to_owned()));
146 }
147 self
148 }
149
150 /// Removes an environment variable from this command.
151 pub fn env_remove<K: AsRef<OsStr>>(mut self, key: K) -> Self {
152 self.env_changes
153 .push(EnvChange::Remove(key.as_ref().to_owned()));
154 self
155 }
156
157 /// Removes all environment variables from this command.
158 pub fn env_clear(mut self) -> Self {
159 self.env_changes.push(EnvChange::Clear);
160 self
161 }
162
163 /// If set, the command's status code will not be checked, and
164 /// non-zero exit codes will not produce an error.
165 pub fn ignore_status(mut self) -> Self {
166 self.ignore_status = true;
167 self
168 }
169
170 /// Mutating variant of [`ignore_status`](Self::ignore_status).
171 pub fn set_ignore_status(&mut self, yes: bool) {
172 self.ignore_status = yes;
173 }
174
175 /// If set, the command's output will not be echoed to stdout.
176 pub fn quiet(mut self) -> Self {
177 self.quiet = true;
178 self
179 }
180
181 /// Mutating variant of [`quiet`](Self::quiet).
182 pub fn set_quiet(&mut self, yes: bool) {
183 self.quiet = yes;
184 }
185
186 /// If set, the command is treated as containing a secret and its
187 /// display will be redacted.
188 pub fn secret(mut self) -> Self {
189 self.secret = true;
190 self
191 }
192
193 /// Mutating variant of [`secret`](Self::secret).
194 pub fn set_secret(&mut self, yes: bool) {
195 self.secret = yes;
196 }
197
198 /// Passes data to the command's stdin.
199 pub fn stdin(mut self, stdin: impl AsRef<[u8]>) -> Self {
200 self.stdin_contents = Some(stdin.as_ref().to_vec());
201 self
202 }
203
204 /// If set, stdout is not captured.
205 pub fn ignore_stdout(mut self) -> Self {
206 self.ignore_stdout = true;
207 self
208 }
209
210 /// Mutating variant of [`ignore_stdout`](Self::ignore_stdout).
211 pub fn set_ignore_stdout(&mut self, yes: bool) {
212 self.ignore_stdout = yes;
213 }
214
215 /// If set, stderr is not captured.
216 pub fn ignore_stderr(mut self) -> Self {
217 self.ignore_stderr = true;
218 self
219 }
220
221 /// Mutating variant of [`ignore_stderr`](Self::ignore_stderr).
222 pub fn set_ignore_stderr(&mut self, yes: bool) {
223 self.ignore_stderr = yes;
224 }
225
226 /// Consume this command, apply the wrapper (if any), re-apply
227 /// shadowed state (env, stdin, flags), and return the final
228 /// [`xshell::Cmd`] ready for execution.
229 fn into_resolved(self) -> xshell::Cmd<'a> {
230 let mut cmd = match self.wrapper {
231 Some(wrapper) => wrapper.wrap_cmd(self.sh, self.inner),
232 None => self.inner,
233 };
234
235 // Re-apply env changes after wrapping to survive the wrapper's transformation
236 for change in self.env_changes {
237 match change {
238 EnvChange::Set(k, v) => cmd = cmd.env(k, v),
239 EnvChange::Remove(k) => cmd = cmd.env_remove(k),
240 EnvChange::Clear => cmd = cmd.env_clear(),
241 }
242 }
243 if let Some(stdin) = self.stdin_contents {
244 cmd = cmd.stdin(stdin);
245 }
246 cmd.set_ignore_status(self.ignore_status);
247 cmd.set_quiet(self.quiet);
248 cmd.set_secret(self.secret);
249 cmd.set_ignore_stdout(self.ignore_stdout);
250 cmd.set_ignore_stderr(self.ignore_stderr);
251
252 cmd
253 }
254
255 /// Run the command.
256 pub fn run(self) -> xshell::Result<()> {
257 self.into_resolved().run()
258 }
259
260 /// Run the command and return its stdout as a string, with leading
261 /// and trailing whitespace trimmed.
262 pub fn read(self) -> xshell::Result<String> {
263 self.into_resolved().read()
264 }
265
266 /// Run the command and return its stderr as a string, with leading
267 /// and trailing whitespace trimmed.
268 pub fn read_stderr(self) -> xshell::Result<String> {
269 self.into_resolved().read_stderr()
270 }
271
272 /// Run the command and return its full output.
273 pub fn output(self) -> xshell::Result<Output> {
274 self.into_resolved().output()
275 }
276}
277
278impl std::fmt::Display for FloweyCmd<'_> {
279 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
280 if self.secret {
281 return f.write_str("<secret>");
282 }
283 // Show the unwrapped command for user-facing logging.
284 std::fmt::Display::fmt(&self.inner, f)
285 }
286}
287
288/// Serializable description of a command wrapper.
289///
290/// This enum can be stored in `pipeline.json` so that CI backends can
291/// reconstruct the appropriate wrapper at runtime. It is also used
292/// directly by [`FloweyShell`] and [`FloweyCmd`] to transform commands.
293#[derive(Clone, Debug, Serialize, Deserialize)]
294pub enum CommandWrapperKind {
295 /// Wrap commands with `nix-shell --pure --run "..."`.
296 NixShell {
297 /// Optional path to a `shell.nix` file. If `None`, nix-shell
298 /// uses its default discovery (looking for `shell.nix` /
299 /// `default.nix` in the current directory).
300 path: Option<std::path::PathBuf>,
301 },
302 /// Wrap commands with `sh -c "..."` (test-only).
303 #[cfg(test)]
304 ShCmd,
305 /// Replace the command with `echo WRAPPED: <cmd>` (test-only).
306 #[cfg(test)]
307 Prefix,
308}
309
310impl CommandWrapperKind {
311 /// Transform a command before execution.
312 ///
313 /// The `cmd` parameter contains only the program and arguments.
314 /// Environment variables, stdin, and flags are applied by
315 /// [`FloweyCmd`] after this method returns.
316 fn wrap_cmd<'a>(self, sh: &'a xshell::Shell, cmd: xshell::Cmd<'a>) -> xshell::Cmd<'a> {
317 let cmd_str = format!("{cmd}");
318 match self {
319 CommandWrapperKind::NixShell { path } => {
320 let mut wrapped = sh.cmd("nix-shell");
321 if let Some(path) = path {
322 wrapped = wrapped.arg(path);
323 }
324 wrapped.arg("--pure").arg("--run").arg(cmd_str)
325 }
326 #[cfg(test)]
327 CommandWrapperKind::ShCmd => sh.cmd("sh").arg("-c").arg(cmd_str),
328 #[cfg(test)]
329 CommandWrapperKind::Prefix => sh.cmd("echo").arg(format!("WRAPPED: {cmd_str}")),
330 }
331 }
332}
333
334#[cfg(test)]
335#[allow(
336 clippy::disallowed_macros,
337 clippy::disallowed_methods,
338 reason = "test module"
339)]
340mod tests {
341 use super::*;
342
343 #[test]
344 fn no_wrapper_runs_command_directly() {
345 let sh = FloweyShell::new().unwrap();
346 let cmd = sh.wrap(xshell::cmd!(sh.xshell(), "echo hello"));
347 let output = cmd.read().unwrap();
348 assert_eq!(output, "hello");
349 }
350
351 #[test]
352 fn wrapper_transforms_command() {
353 let mut sh = FloweyShell::new().unwrap();
354 sh.set_wrapper(Some(CommandWrapperKind::Prefix));
355 let cmd = sh.wrap(xshell::cmd!(sh.xshell(), "my-program --flag"));
356 let output = cmd.read().unwrap();
357 assert_eq!(output, "WRAPPED: my-program --flag");
358 }
359
360 #[test]
361 fn env_vars_survive_with_wrapper() {
362 let mut sh = FloweyShell::new().unwrap();
363 sh.set_wrapper(Some(CommandWrapperKind::ShCmd));
364
365 let cmd = sh
366 .wrap(xshell::cmd!(sh.xshell(), "printenv MY_FLOWEY_WRAP_TEST"))
367 .env("MY_FLOWEY_WRAP_TEST", "survived_wrapping");
368 let output = cmd.read().unwrap();
369 assert_eq!(output, "survived_wrapping");
370 }
371
372 #[test]
373 fn stdin_survives_wrapping() {
374 let sh = FloweyShell::new().unwrap();
375 let cmd = sh
376 .wrap(xshell::cmd!(sh.xshell(), "cat"))
377 .stdin("test input");
378 let output = cmd.read().unwrap();
379 assert_eq!(output, "test input");
380 }
381
382 #[test]
383 fn stdin_survives_with_wrapper() {
384 let mut sh = FloweyShell::new().unwrap();
385 sh.set_wrapper(Some(CommandWrapperKind::ShCmd));
386
387 let cmd = sh
388 .wrap(xshell::cmd!(sh.xshell(), "cat"))
389 .stdin("wrapped stdin test");
390 let output = cmd.read().unwrap();
391 assert_eq!(output, "wrapped stdin test");
392 }
393
394 #[test]
395 fn ignore_status_survives_wrapping() {
396 let mut sh = FloweyShell::new().unwrap();
397 sh.set_wrapper(Some(CommandWrapperKind::ShCmd));
398
399 // `false` exits with status 1 — without ignore_status this would error.
400 let cmd = sh.wrap(xshell::cmd!(sh.xshell(), "false")).ignore_status();
401 assert!(cmd.run().is_ok());
402 }
403
404 #[test]
405 fn display_shows_unwrapped_command() {
406 let mut sh = FloweyShell::new().unwrap();
407 sh.set_wrapper(Some(CommandWrapperKind::Prefix));
408 let cmd = sh.wrap(xshell::cmd!(sh.xshell(), "my-program --flag value"));
409 assert_eq!(format!("{cmd}"), "my-program --flag value");
410 }
411
412 #[test]
413 fn nix_wrapper_display_without_path() {
414 let sh = FloweyShell::new().unwrap();
415 let cmd = CommandWrapperKind::NixShell { path: None }.wrap_cmd(
416 sh.xshell(),
417 xshell::cmd!(sh.xshell(), "cargo build --release"),
418 );
419 assert_eq!(
420 format!("{cmd}"),
421 "nix-shell --pure --run \"cargo build --release\""
422 );
423 }
424
425 #[test]
426 fn nix_wrapper_display_with_path() {
427 let sh = FloweyShell::new().unwrap();
428 let cmd = CommandWrapperKind::NixShell {
429 path: Some("/my/shell.nix".into()),
430 }
431 .wrap_cmd(sh.xshell(), xshell::cmd!(sh.xshell(), "cargo build"));
432 assert_eq!(
433 format!("{cmd}"),
434 "nix-shell /my/shell.nix --pure --run \"cargo build\""
435 );
436 }
437
438 #[test]
439 fn deref_exposes_shell_methods() {
440 let sh = FloweyShell::new().unwrap();
441 let _ = sh.current_dir();
442 }
443
444 #[test]
445 fn set_wrapper_clears_wrapper() {
446 let mut sh = FloweyShell::new().unwrap();
447 sh.set_wrapper(Some(CommandWrapperKind::Prefix));
448 sh.set_wrapper(None);
449 // With wrapper cleared, command should run directly.
450 let cmd = sh.wrap(xshell::cmd!(sh.xshell(), "echo direct"));
451 let output = cmd.read().unwrap();
452 assert_eq!(output, "direct");
453 }
454
455 #[test]
456 fn quiet_flag_survives_wrapping() {
457 let mut sh = FloweyShell::new().unwrap();
458 sh.set_wrapper(Some(CommandWrapperKind::Prefix));
459 // quiet() should not cause errors — just suppress echo to stderr.
460 let cmd = sh.wrap(xshell::cmd!(sh.xshell(), "echo test")).quiet();
461 let output = cmd.read().unwrap();
462 assert_eq!(output, "WRAPPED: echo test");
463 }
464
465 #[test]
466 fn args_accumulate_before_wrapping() {
467 let mut sh = FloweyShell::new().unwrap();
468 sh.set_wrapper(Some(CommandWrapperKind::Prefix));
469 let cmd = sh
470 .wrap(xshell::cmd!(sh.xshell(), "echo"))
471 .arg("one")
472 .arg("two");
473 let output = cmd.read().unwrap();
474 assert_eq!(output, "WRAPPED: echo one two");
475 }
476
477 #[test]
478 fn secret_display_is_redacted() {
479 let sh = FloweyShell::new().unwrap();
480 let cmd = sh
481 .wrap(xshell::cmd!(sh.xshell(), "curl --header secret-token"))
482 .secret();
483 assert_eq!(format!("{cmd}"), "<secret>");
484 }
485
486 #[test]
487 fn secret_display_redacted_with_wrapper() {
488 let mut sh = FloweyShell::new().unwrap();
489 sh.set_wrapper(Some(CommandWrapperKind::Prefix));
490 let cmd = sh
491 .wrap(xshell::cmd!(sh.xshell(), "curl --header secret-token"))
492 .secret();
493 assert_eq!(format!("{cmd}"), "<secret>");
494 }
495
496 #[test]
497 fn env_remove_survives_wrapping() {
498 let mut sh = FloweyShell::new().unwrap();
499 sh.set_wrapper(Some(CommandWrapperKind::ShCmd));
500
501 // Set a var via the shell, then remove it on the command.
502 // printenv should fail (exit 1) because the var is removed.
503 sh.set_var("FLOWEY_REMOVE_TEST", "present");
504 let cmd = sh
505 .wrap(xshell::cmd!(sh.xshell(), "printenv FLOWEY_REMOVE_TEST"))
506 .env_remove("FLOWEY_REMOVE_TEST")
507 .ignore_status();
508 let output = cmd.output().unwrap();
509 assert!(!output.status.success());
510 }
511
512 #[test]
513 fn env_clear_survives_wrapping() {
514 let mut sh = FloweyShell::new().unwrap();
515 sh.set_wrapper(Some(CommandWrapperKind::ShCmd));
516
517 // After env_clear, even PATH is gone. The wrapped command
518 // should still run (sh is resolved before env_clear applies),
519 // but the inner command won't find the var.
520 sh.set_var("FLOWEY_CLEAR_TEST", "present");
521 let cmd = sh
522 .wrap(xshell::cmd!(sh.xshell(), "printenv FLOWEY_CLEAR_TEST"))
523 .env_clear()
524 .ignore_status();
525 let output = cmd.output().unwrap();
526 assert!(!output.status.success());
527 }
528
529 #[test]
530 fn env_ordering_preserved_through_wrapping() {
531 let mut sh = FloweyShell::new().unwrap();
532 sh.set_wrapper(Some(CommandWrapperKind::ShCmd));
533
534 // Set, clear, then set again — only the final value should survive.
535 let cmd = sh
536 .wrap(xshell::cmd!(sh.xshell(), "printenv FLOWEY_ORDER_TEST"))
537 .env("FLOWEY_ORDER_TEST", "first")
538 .env_clear()
539 .env("FLOWEY_ORDER_TEST", "second");
540 let output = cmd.read().unwrap();
541 assert_eq!(output, "second");
542 }
543
544 #[test]
545 fn envs_plural_survives_wrapping() {
546 let mut sh = FloweyShell::new().unwrap();
547 sh.set_wrapper(Some(CommandWrapperKind::ShCmd));
548
549 let vars = vec![("FLOWEY_MULTI_A", "alpha"), ("FLOWEY_MULTI_B", "beta")];
550 // Print both vars separated by a space.
551 let cmd = sh
552 .wrap(xshell::cmd!(sh.xshell(), "sh"))
553 .arg("-c")
554 .arg("echo $FLOWEY_MULTI_A $FLOWEY_MULTI_B")
555 .envs(vars);
556 let output = cmd.read().unwrap();
557 assert_eq!(output, "alpha beta");
558 }
559}
560