flowey_core/
shell.rs

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}