use std::env; use std::fmt::{self, Debug}; use std::io::Write; use std::path::Path; use std::process::{Command, Stdio}; use std::time::Duration; use process_control::{ChildExt, Control, Output}; use super::{Context, Module, ModuleConfig}; use crate::{ config::Either, configs::custom::CustomConfig, formatter::StringFormatter, utils::create_command, }; /// Creates a custom module with some configuration /// /// The relevant TOML config will set the files, extensions, and directories needed /// for the module to be displayed. If none of them match, and optional "when" /// command can be run -- if its result is 0, the module will be shown. /// /// Finally, the content of the module itself is also set by a command. pub fn module<'a>(name: &str, context: &'a Context) -> Option> { let toml_config = get_config(name, context)?; let config = CustomConfig::load(toml_config); if config.disabled { return None; } if let Some(os) = config.os { if os != env::consts::OS && !(os == "unix" && cfg!(unix)) { return None; } } if config.require_repo && context.get_repo().is_err() { return None; } // Note: Forward config if `Module` ends up needing `config` let mut module = Module::new(&format!("custom.{name}"), config.description, None); let mut is_match = context .try_begin_scan()? .set_extensions(&config.detect_extensions) .set_files(&config.detect_files) .set_folders(&config.detect_folders) .is_match(); if !is_match { is_match = match config.when { Either::First(b) => b, Either::Second(s) => exec_when(s, &config, context), }; if !is_match { return None; } } let parsed = StringFormatter::new(config.format).and_then(|formatter| { formatter .map_meta(|var, _| match var { "symbol" => Some(config.symbol), _ => None, }) .map_style(|variable| match variable { "style" => Some(Ok(config.style)), _ => None, }) .map_no_escaping(|variable| match variable { "output" => { let output = exec_command(config.command, context, &config)?; let trimmed = output.trim(); if trimmed.is_empty() { None } else { Some(Ok(trimmed.to_string())) } } _ => None, }) .parse(None, Some(context)) }); match parsed { Ok(segments) => module.set_segments(segments), Err(error) => { log::warn!("Error in module `custom.{}`:\n{}", name, error); } }; Some(module) } /// Gets the TOML config for the custom module, handling the case where the module is not defined fn get_config<'a>(module_name: &str, context: &'a Context<'a>) -> Option<&'a toml::Value> { struct DebugCustomModules<'tmp>(&'tmp toml::value::Table); impl Debug for DebugCustomModules<'_> { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { f.debug_list().entries(self.0.keys()).finish() } } let config = context.config.get_custom_module_config(module_name); if config.is_some() { return config; } else if let Some(modules) = context.config.get_custom_modules() { log::debug!( "top level format contains custom module {module_name:?}, but no configuration was provided. Configuration for the following modules were provided: {:?}", DebugCustomModules(modules), ); } else { log::debug!( "top level format contains custom module {module_name:?}, but no configuration was provided.", ); }; None } /// Return the invoking shell, using `shell` and fallbacking in order to `STARSHIP_SHELL` and "sh"/"cmd" fn get_shell<'a, 'b>( shell_args: &'b [&'a str], context: &Context, ) -> (std::borrow::Cow<'a, str>, &'b [&'a str]) { if !shell_args.is_empty() { (shell_args[0].into(), &shell_args[1..]) } else if let Some(env_shell) = context.get_env("STARSHIP_SHELL") { (env_shell.into(), &[] as &[&str]) } else if cfg!(windows) { // `/C` is added by `handle_shell` ("cmd".into(), &[] as &[&str]) } else { ("sh".into(), &[] as &[&str]) } } /// Attempt to run the given command in a shell by passing it as either `stdin` or an argument to `get_shell()`, /// depending on the configuration or by invoking a platform-specific fallback shell if `shell` is empty. fn shell_command(cmd: &str, config: &CustomConfig, context: &Context) -> Option { let (shell, shell_args) = get_shell(config.shell.0.as_ref(), context); let mut use_stdin = config.use_stdin; let mut command = match create_command(shell.as_ref()) { Ok(command) => command, // Don't attempt to use fallback shell if the user specified a shell Err(error) if !shell_args.is_empty() => { log::debug!( "Error creating command with STARSHIP_SHELL, falling back to fallback shell: {}", error ); // Skip `handle_shell` and just set the shell and command use_stdin = Some(!cfg!(windows)); if cfg!(windows) { let mut c = create_command("cmd").ok()?; c.arg("/C"); c } else { let mut c = create_command("/usr/bin/env").ok()?; c.arg("sh"); c } } _ => return None, }; command .current_dir(&context.current_dir) .args(shell_args) .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::piped()); let use_stdin = use_stdin.unwrap_or_else(|| handle_shell(&mut command, &shell, shell_args)); if !use_stdin { command.arg(cmd); } let mut child = match command.spawn() { Ok(child) => child, Err(error) => { log::debug!( "Failed to run command with given shell or STARSHIP_SHELL env variable:: {}", error ); return None; } }; if use_stdin { child.stdin.as_mut()?.write_all(cmd.as_bytes()).ok()?; } let mut output = child.controlled_with_output(); if !config.ignore_timeout { output = output .time_limit(Duration::from_millis(context.root_config.command_timeout)) .terminate_for_timeout() } match output.wait().ok()? { None => { log::warn!("Executing custom command {cmd:?} timed out."); log::warn!("You can set command_timeout in your config to a higher value or set ignore_timeout to true for this module to allow longer-running commands to keep executing."); None } Some(status) => Some(status), } } /// Execute the given command capturing all output, and return whether it return 0 fn exec_when(cmd: &str, config: &CustomConfig, context: &Context) -> bool { log::trace!("Running '{}'", cmd); if let Some(output) = shell_command(cmd, config, context) { if !output.status.success() { log::trace!("non-zero exit code '{:?}'", output.status.code()); log::trace!( "stdout: {}", std::str::from_utf8(&output.stdout).unwrap_or("") ); log::trace!( "stderr: {}", std::str::from_utf8(&output.stderr).unwrap_or("") ); } output.status.success() } else { log::debug!("Cannot start command"); false } } /// Execute the given command, returning its output on success fn exec_command(cmd: &str, context: &Context, config: &CustomConfig) -> Option { log::trace!("Running '{cmd}'"); if let Some(output) = shell_command(cmd, config, context) { if !output.status.success() { log::trace!("Non-zero exit code '{:?}'", output.status.code()); log::trace!( "stdout: {}", std::str::from_utf8(&output.stdout).unwrap_or("") ); log::trace!( "stderr: {}", std::str::from_utf8(&output.stderr).unwrap_or("") ); return None; } Some(String::from_utf8_lossy(&output.stdout).into()) } else { None } } /// If the specified shell refers to `PowerShell`, adds the arguments "-Command -" to the /// given command. /// Returns `false` if the shell shell expects scripts as arguments, `true` if as `stdin`. fn handle_shell(command: &mut Command, shell: &str, shell_args: &[&str]) -> bool { let shell_exe = Path::new(shell).file_stem(); let no_args = shell_args.is_empty(); match shell_exe.and_then(std::ffi::OsStr::to_str) { Some("pwsh" | "powershell") => { if no_args { command.arg("-NoProfile").arg("-Command").arg("-"); } true } Some("cmd") => { if no_args { command.arg("/C"); } false } Some("nu") => { if no_args { command.arg("-c"); } false } _ => true, } } #[cfg(test)] mod tests { use super::*; use crate::test::{fixture_repo, FixtureProvider, ModuleRenderer}; use nu_ansi_term::Color; use std::fs::File; use std::io; #[cfg(not(windows))] const SHELL: &[&str] = &["/bin/sh"]; #[cfg(windows)] const SHELL: &[&str] = &["cmd"]; #[cfg(not(windows))] const FAILING_COMMAND: &str = "false"; #[cfg(windows)] const FAILING_COMMAND: &str = "color 00"; const UNKNOWN_COMMAND: &str = "ydelsyiedsieudleylse dyesdesl"; fn render_cmd(cmd: &str) -> io::Result> { let dir = tempfile::tempdir()?; let cmd = cmd.to_owned(); let shell = SHELL .iter() .map(std::borrow::ToOwned::to_owned) .collect::>(); let out = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { [custom.test] format = "$output" command = cmd shell = shell when = true ignore_timeout = true }) .collect(); dir.close()?; Ok(out) } fn render_when(cmd: &str) -> io::Result { let dir = tempfile::tempdir()?; let cmd = cmd.to_owned(); let shell = SHELL .iter() .map(std::borrow::ToOwned::to_owned) .collect::>(); let out = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { [custom.test] format = "test" when = cmd shell = shell ignore_timeout = true }) .collect() .is_some(); dir.close()?; Ok(out) } #[test] fn when_returns_right_value() -> io::Result<()> { assert!(render_cmd("echo hello")?.is_some()); assert!(render_cmd(FAILING_COMMAND)?.is_none()); Ok(()) } #[test] fn when_returns_false_if_invalid_command() -> io::Result<()> { assert!(!render_when(UNKNOWN_COMMAND)?); Ok(()) } #[test] #[cfg(not(windows))] fn command_returns_right_string() -> io::Result<()> { assert_eq!(render_cmd("echo hello")?, Some("hello".into())); assert_eq!(render_cmd("echo 강남스타일")?, Some("강남스타일".into())); Ok(()) } #[test] #[cfg(windows)] fn command_returns_right_string() -> io::Result<()> { assert_eq!(render_cmd("echo hello")?, Some("hello".into())); assert_eq!(render_cmd("echo 강남스타일")?, Some("강남스타일".into())); Ok(()) } #[test] #[cfg(not(windows))] fn command_ignores_stderr() -> io::Result<()> { assert_eq!(render_cmd("echo foo 1>&2; echo bar")?, Some("bar".into())); assert_eq!(render_cmd("echo foo; echo bar 1>&2")?, Some("foo".into())); Ok(()) } #[test] #[cfg(windows)] fn command_ignores_stderr() -> io::Result<()> { assert_eq!(render_cmd("echo foo 1>&2 & echo bar")?, Some("bar".into())); assert_eq!(render_cmd("echo foo& echo bar 1>&2")?, Some("foo".into())); Ok(()) } #[test] fn command_can_fail() -> io::Result<()> { assert_eq!(render_cmd(FAILING_COMMAND)?, None); assert_eq!(render_cmd(UNKNOWN_COMMAND)?, None); Ok(()) } #[test] fn cwd_command() -> io::Result<()> { let dir = tempfile::tempdir()?; let mut f = File::create(dir.path().join("a.txt"))?; write!(f, "hello")?; f.sync_all()?; let cat = if cfg!(windows) { "type" } else { "cat" }; let cmd = format!("{cat} a.txt"); let actual = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { [custom.test] command = cmd when = true ignore_timeout = true }) .collect(); let expected = Some(format!("{}", Color::Green.bold().paint("hello "))); assert_eq!(expected, actual); dir.close() } #[test] fn cwd_when() -> io::Result<()> { let dir = tempfile::tempdir()?; File::create(dir.path().join("a.txt"))?.sync_all()?; let cat = if cfg!(windows) { "type" } else { "cat" }; let cmd = format!("{cat} a.txt"); let actual = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { [custom.test] format = "test" when = cmd ignore_timeout = true }) .collect(); let expected = Some("test".to_owned()); assert_eq!(expected, actual); dir.close() } #[test] fn use_stdin_false() -> io::Result<()> { let dir = tempfile::tempdir()?; let shell = if cfg!(windows) { vec![ "powershell".to_owned(), "-NoProfile".to_owned(), "-Command".to_owned(), ] } else { vec!["sh".to_owned(), "-c".to_owned()] }; // `use_stdin = false` doesn't like Korean on Windows let actual = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { [custom.test] command = "echo test" when = true use_stdin = false shell = shell ignore_timeout = true }) .collect(); let expected = Some(format!("{}", Color::Green.bold().paint("test "))); assert_eq!(expected, actual); dir.close() } #[test] fn use_stdin_true() -> io::Result<()> { let dir = tempfile::tempdir()?; let shell = if cfg!(windows) { vec![ "powershell".to_owned(), "-NoProfile".to_owned(), "-Command".to_owned(), "-".to_owned(), ] } else { vec!["sh".to_owned()] }; let actual = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { [custom.test] command = "echo 강남스타일" when = true use_stdin = true ignore_timeout = true shell = shell }) .collect(); let expected = Some(format!("{}", Color::Green.bold().paint("강남스타일 "))); assert_eq!(expected, actual); dir.close() } #[test] #[cfg(not(windows))] fn when_true_with_string() -> std::io::Result<()> { let dir = tempfile::tempdir()?; let actual = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { [custom.test] format = "test" shell = ["sh"] when = "true" ignore_timeout = true }) .collect(); let expected = Some("test".to_string()); assert_eq!(expected, actual); dir.close() } #[test] #[cfg(not(windows))] fn when_false_with_string() -> std::io::Result<()> { let dir = tempfile::tempdir()?; let actual = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { [custom.test] format = "test" shell = ["sh"] when = "false" ignore_timeout = true }) .collect(); let expected = None; assert_eq!(expected, actual); dir.close() } #[test] fn when_true_with_bool() -> std::io::Result<()> { let dir = tempfile::tempdir()?; let actual = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { [custom.test] format = "test" when = true }) .collect(); let expected = Some("test".to_string()); assert_eq!(expected, actual); dir.close() } #[test] #[cfg(not(windows))] fn when_false_with_bool() -> std::io::Result<()> { let dir = tempfile::tempdir()?; let actual = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { [custom.test] format = "test" when = false }) .collect(); let expected = None; assert_eq!(expected, actual); dir.close() } #[test] fn timeout_short_cmd() -> std::io::Result<()> { let dir = tempfile::tempdir()?; let shell = if cfg!(windows) { "powershell".to_owned() } else { "sh".to_owned() }; let when = if cfg!(windows) { "$true".to_owned() } else { "true".to_owned() }; // Use a long timeout to ensure that the test doesn't fail let actual = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { command_timeout = 100_000 [custom.test] format = "test" when = when shell = shell ignore_timeout = false }) .collect(); let expected = Some("test".to_owned()); assert_eq!(expected, actual); dir.close() } #[test] fn timeout_cmd() -> std::io::Result<()> { let dir = tempfile::tempdir()?; let shell = if cfg!(windows) { "powershell".to_owned() } else { "sh".to_owned() }; // Use a long timeout to ensure that the test doesn't fail let actual = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { [custom.test] format = "test" when = "sleep 3" shell = shell ignore_timeout = false }) .collect(); let expected = None; assert_eq!(expected, actual); dir.close() } #[test] fn config_aliases_work() -> std::io::Result<()> { let dir = tempfile::tempdir()?; File::create(dir.path().join("a.txt"))?; std::fs::create_dir(dir.path().join("dir"))?; let actual = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { [custom.test] format = "test" files = ["a.txt"] }) .collect(); let expected = Some("test".to_string()); assert_eq!(expected, actual); let actual = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { [custom.test] format = "test" extensions = ["txt"] }) .collect(); let expected = Some("test".to_string()); assert_eq!(expected, actual); let actual = ModuleRenderer::new("custom.test") .path(dir.path()) .config(toml::toml! { [custom.test] format = "test" directories = ["dir"] }) .collect(); let expected = Some("test".to_string()); assert_eq!(expected, actual); dir.close() } #[test] fn disabled() { let actual = ModuleRenderer::new("custom.test") .config(toml::toml! { [custom.test] disabled = true when = true format = "test" }) .collect(); let expected = None; assert_eq!(expected, actual); } #[test] fn test_render_require_repo_not_in() -> io::Result<()> { let repo_dir = tempfile::tempdir()?; let actual = ModuleRenderer::new("custom.test") .path(repo_dir.path()) .config(toml::toml! { [custom.test] when = true require_repo = true format = "test" }) .collect(); let expected = None; assert_eq!(expected, actual); repo_dir.close() } #[test] fn test_render_require_repo_in() -> io::Result<()> { let repo_dir = fixture_repo(FixtureProvider::Git)?; let actual = ModuleRenderer::new("custom.test") .path(repo_dir.path()) .config(toml::toml! { [custom.test] when = true require_repo = true format = "test" }) .collect(); let expected = Some("test".to_string()); assert_eq!(expected, actual); repo_dir.close() } }