From bc3cffb3ee88ce3f975c4b48c183b248b4fd4923 Mon Sep 17 00:00:00 2001 From: Zack Weinberg Date: Tue, 24 Nov 2020 13:45:40 -0500 Subject: [PATCH] More robust shell identification logic (#67) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All silver subcommands look up the parent process’s name, and use that, verbatim, to identify the shell in use. This has several problems: - On Windows, the parent process’s name may or may not end with “.exe”; there was special case logic to handle “powershell.exe” versus “powershell” but not “bash.exe” versus “bash”. Relatedly, the parent process’s name ought to be interpreted case- insensitively on Windows; BASH.EXE is the same as bash.exe. - On Unix, there is a notion of a _login_ shell (invoked directly by login(1) for console logins, or sshd(1) for remote logins); these are identified by a leading dash in their process name, *not* by a special option. There was no code to handle this at all (bug #67). - sysinfo is slow, because it reads a whole bunch of files in /proc whether it’s going to need the data within or not. Some but not all of this overhead can be eliminated by using `System::new_with_specifics(RefreshKind::new())` instead of `System::new_all()`, but even then the wall-clock time for `silver init` drops from 0.021s to 0.002s (release build) if sysinfo is not invoked at all. This patch addresses all these problems, as follows: - There is a new command line option `--shell` which may be used to explicitly set the shell in use. Also, support for the `SILVER_SHELL` environment variable (used for this purpose in 1.1 and previous) is restored. sysinfo is only consulted if neither the command line option, nor the environment variable, is available. - No matter where we get the process name from, it is lowercased, a trailing “.exe” is removed, a leading dash is removed, and any leading directory names are also removed (bug #16). - The initialization scripts set `SILVER_SHELL`, so sysinfo will be consulted at most once per session, not once per command. I also reorganized the code in the Command::Init branch of main for readability’s sake, and made the “unknown shell” error message not be a panic. (Panics are for _bugs_, not for errors with a cause outside the program.) --- src/cli.rs | 2 + src/init.bash | 1 + src/init.ion | 3 +- src/init.ps1 | 3 +- src/main.rs | 132 ++++++++++++++++++++++++++++++++++++++------------ 5 files changed, 107 insertions(+), 34 deletions(-) diff --git a/src/cli.rs b/src/cli.rs index e7eed99..6448184 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -9,6 +9,8 @@ pub use structopt::StructOpt; pub struct Silver { #[structopt(short, long)] pub config: Option, + #[structopt(short, long)] + pub shell: Option, #[structopt(subcommand)] pub cmd: Command, } diff --git a/src/init.bash b/src/init.bash index f356215..2cd580f 100644 --- a/src/init.bash +++ b/src/init.bash @@ -2,4 +2,5 @@ PROMPT_COMMAND=silver_prompt silver_prompt() { PS1="$(code=$? jobs=$(jobs -p | wc -l) silver lprint)" } +export SILVER_SHELL="@SILVER_SHELL@" export VIRTUAL_ENV_DISABLE_PROMPT=1 diff --git a/src/init.ion b/src/init.ion index 1b2ce24..c157c09 100644 --- a/src/init.ion +++ b/src/init.ion @@ -1,4 +1,5 @@ fn PROMPT env code=$? jobs=$(jobs -p | wc -l) silver lprint end -export VIRTUAL_ENV_DISABLE_PROMPT = 1 # Doesn't make any sense yet +export SILVER_SHELL = "@SILVER_SHELL@" +export VIRTUAL_ENV_DISABLE_PROMPT = 1 diff --git a/src/init.ps1 b/src/init.ps1 index 59ed790..16595ae 100644 --- a/src/init.ps1 +++ b/src/init.ps1 @@ -6,4 +6,5 @@ function prompt { Start-Process -Wait -NoNewWindow silver lprint "$([char]0x1b)[0m" } -$Env:VIRTUAL_ENV_DISABLE_PROMPT = 1 +$env:SILVER_SHELL = "@SILVER_SHELL@" +$env:VIRTUAL_ENV_DISABLE_PROMPT = 1 diff --git a/src/main.rs b/src/main.rs index bd535df..a74592b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,8 +7,9 @@ mod sh; use cli::*; use once_cell::sync::{Lazy, OnceCell}; +use std::env; use std::path::{Path, PathBuf}; -use sysinfo::{get_current_pid, ProcessExt, System, SystemExt}; +use sysinfo::{get_current_pid, ProcessExt, RefreshKind, System, SystemExt}; static CONFIG_PATH: OnceCell = OnceCell::new(); @@ -21,6 +22,10 @@ static CONFIG: Lazy = Lazy::new(|| { .expect("Failed to read config") }); +const INIT_BASH: &str = include_str!("init.bash"); +const INIT_PS1: &str = include_str!("init.ps1"); +const INIT_ION: &str = include_str!("init.ion"); + #[derive(Clone, Debug)] pub struct Segment { background: String, @@ -38,47 +43,110 @@ impl Default for Segment { } } -fn main() { - let sys = System::new_all(); - let process = sys.get_process(get_current_pid().unwrap()).unwrap(); - let parent = sys.get_process(process.parent().unwrap()).unwrap(); - let shell = parent.name().trim(); +/// Helper function for trimming a String in place. Derived from +/// https://users.rust-lang.org/t/trim-string-in-place/15809/9 +fn replace_with_subslice(this: &mut String, f: F) + where F: FnOnce(&str) -> &str +{ + let original_len = this.len(); + let new_slice: &str = f(this); + let start = (new_slice.as_ptr() as usize) + .wrapping_sub(this.as_ptr() as usize); + if original_len < start { + this.clear(); + return; + } + + let len = new_slice.len(); + if start != 0 { + this.drain(..start); + } + this.truncate(len); +} + +/// Identify the type of shell in use. +fn get_shell(opt: &cli::Silver) -> String { + let mut shell: String = + if let Some(ref s) = opt.shell { + // If --shell was given on the command line, use that. + s.clone() + } else if let Ok(s) = env::var("SILVER_SHELL") { + // For backward compatibility with 1.1 and earlier, + // use the value of the SILVER_SHELL environment variable. + s + } else { + // Use the name of the parent process, if we can. + // Minimize the amount of information loaded by sysinfo. + // FIXME: Proper error handling, not all these unwraps. + let mut sys = System::new_with_specifics(RefreshKind::new()); + // It'd be nice if either the stdlib or sysinfo exposed + // getppid() directly, but they don't. + let mypid = get_current_pid().unwrap(); + sys.refresh_process(mypid); + let parentpid = sys.get_process(mypid).unwrap().parent().unwrap(); + sys.refresh_process(parentpid); + sys.get_process(parentpid).unwrap().name().to_string() + }; + + // For Windows compatibility, lowercase the shell's name + // ("BASH.EXE" is the same as "bash.exe"). + shell.make_ascii_lowercase(); + + // Remove any leading or trailing whitespace from `shell`. + // Remove anything up to and including the last '/' or '\\', + // in case the shell was identified by absolute path. + // Remove a trailing ".exe" if present, for Windows compatibility. + // Remove a leading "-" if present; on Unix this is a flag indicating + // a login shell (see login(1) and sh(1)), not part of the name. + // These are done unconditionally, not just when we are using + // the name of the parent process to identify the shell, because + // the installation instructions for older versions of silver + // said to set SILVER_SHELL to "$0" without trimming anything + // from it. + replace_with_subslice(&mut shell, |s| { + let s = s.trim(); + let s = s.rsplit(&['/', '\\'][..]).next().unwrap_or(s); + let s = s.strip_prefix('-').unwrap_or(s); + let s = s.strip_suffix(".exe").unwrap_or(s); + s + }); + + shell +} + +fn main() { let opt = cli::Silver::from_args(); - if let Some(path) = opt.config { + if let Some(ref path) = opt.config { let path = Path::new(path.as_str()).canonicalize().unwrap(); CONFIG_PATH.set(path).unwrap() } + let _shell = get_shell(&opt); + let shell = _shell.as_str(); + match opt.cmd { Command::Init => { - print!( - "{}", - match shell { - "bash" => include_str!("init.bash"), - "powershell" | "pwsh" | "powershell.exe" | "pwsh.exe" => - include_str!("init.ps1"), - "ion" => include_str!("init.ion"), - _ => - panic!( - "unknown shell: \"{}\". Supported shells: bash, ion, powershell", - shell - ), + let script = match shell { + "bash" => INIT_BASH, + "powershell" | "pwsh" => INIT_PS1, + "ion" => INIT_ION, + _ => { + use std::process::exit; + eprintln!("silver: unknown shell: \"{}\".", shell); + eprintln!("silver: supported shells: bash, ion, powershell"); + exit(1); } - .replace( - "silver", - format!( - "silver{}", - if let Some(path) = CONFIG_PATH.get() { - format!(" --config {}", path.display()) - } else { - String::new() - } - ) - .as_str() - ) - ) + }; + let script = script.replace("@SILVER_SHELL@", shell); + let script = if let Some(path) = CONFIG_PATH.get() { + script.replace("silver", format!("silver --config {}", + path.display()).as_str()) + } else { + script + }; + print!("{}", script); } Command::Lprint => { print::prompt(&shell, &CONFIG.left, |_, (_, c, n)| {