Skip to content

feat: add automatic env var reload to prompt hooks #426

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ethanpailes wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat: add automatic env var reload to prompt hooks #426

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 76 additions & 5 deletions libshpool/src/daemon/server.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -623,11 +623,14 @@ impl Server {

let session_env_file = self.session_env_file(session_name);
info!("populating {:?}", session_env_file);
fs::write(
session_env_file,
header.local_env.iter().map(|(k, v)| format!("{k}={v}")).collect::<Vec<_>>().join("\n"),
)
.context("writing session env")?;
let content = format_forward_env(&header.local_env);
fs::write(&session_env_file, content).context("writing session env")?;

// Remove the stamp file if present so shells with 1-second timestamp
// resolution (e.g. bash 3.2 on macOS) reload unconditionally on
// reattach without needing a full second to elapse.
let stamp_file = format!("{}.stamp", session_env_file.display());
let _ = fs::remove_file(stamp_file);

Ok(())
}
Expand Down Expand Up @@ -1361,6 +1364,31 @@ impl Server {
}
}

fn format_forward_env<'a, I>(env_vars: I) -> String
where
I: IntoIterator<Item = &'a (String, String)>,
{
let mut content = String::new();
for (k, v) in env_vars {
if is_valid_env_key(k) {
content.push_str(&format!("export {k}='{}'\n", v.replace('\'', "'\\''")));
} else {
warn!("skipping invalid environment variable key: {k}");
}
}
content
}

fn is_valid_env_key(key: &str) -> bool {
let mut chars = key.chars();
match chars.next() {
Some(c) if c.is_ascii_alphabetic() || c == '_' => {
chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
}
_ => false,
}
}

// HACK: this is not a good way to detect shells that don't support our
// sentinel injection approach, but it is better than just hanging when a
// user tries to start one.
Expand Down Expand Up @@ -1497,3 +1525,46 @@ impl std::fmt::Display for ShellSelectionError {
}

impl std::error::Error for ShellSelectionError {}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn test_format_forward_env_basic() {
let vars = vec![
(String::from("FOO"), String::from("bar")),
(String::from("SPACES"), String::from("hello world")),
(String::from("QUOTES"), String::from("don't stop")),
(String::from("SPECIAL"), String::from("$(evil) `evil` $VAR")),
(String::from("MULTILINE"), String::from("line1\nline2")),
];
let res = format_forward_env(&vars);
assert_eq!(
res,
"export FOO='bar'\n\
export SPACES='hello world'\n\
export QUOTES='don'\\''t stop'\n\
export SPECIAL='$(evil) `evil` $VAR'\n\
export MULTILINE='line1\nline2'\n"
);
}

#[test]
fn test_format_forward_env_invalid_keys() {
let vars = vec![
(String::from("GOOD_KEY_1"), String::from("val")),
(String::from("_ALSO_GOOD"), String::from("val")),
(String::from("1BAD_KEY"), String::from("val")),
(String::from("BAD-DASH"), String::from("val")),
(String::from("BAD KEY"), String::from("val")),
(String::from(""), String::from("val")),
];
let res = format_forward_env(&vars);
assert_eq!(
res,
"export GOOD_KEY_1='val'\n\
export _ALSO_GOOD='val'\n"
);
}
}
130 changes: 106 additions & 24 deletions libshpool/src/daemon/shell_inject.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,45 +72,127 @@ pub fn maybe_setup(
let prompt_prefix = prompt_prefix.replace("$SHPOOL_SESSION_NAME", session_name);

let mut script = match (prompt_prefix.as_str(), shell_type) {
(_, Ok(KnownShell::Bash)) => format!(
r#"
if [[ -z "${{PROMPT_COMMAND+x}}" ]]; then
PS1="{prompt_prefix}${{PS1}}"
else
SHPOOL__OLD_PROMPT_COMMAND=("${{PROMPT_COMMAND[@]}}")
SHPOOL__OLD_PS1="${{PS1}}"
function __shpool__prompt_command() {{
PS1="${{SHPOOL__OLD_PS1}}"
for prompt_hook in "${{SHPOOL__OLD_PROMPT_COMMAND[@]}}"
do
eval "${{prompt_hook}}"
done
PS1="{prompt_prefix}${{PS1}}"
}}
PROMPT_COMMAND=__shpool__prompt_command
fi
(_, Ok(KnownShell::Bash)) => {
// In Bash 5.1+, PROMPT_COMMAND supports arrays. However, older
// versions of Bash (such as Bash 3.2, the default system shell on
// macOS) only execute PROMPT_COMMAND if it is a scalar string;
// assigning an array causes Bash 3.2 to silently ignore it.
// We capture any existing hooks (array or scalar) into
// SHPOOL__OLD_PROMPT_COMMAND, unset PROMPT_COMMAND, and assign
// PROMPT_COMMAND as a scalar string to ensure universal
// compatibility.
format!(
r#"
SHPOOL__OLD_PROMPT_COMMAND=("${{PROMPT_COMMAND[@]}}")
SHPOOL__OLD_PS1="${{PS1}}"
function __shpool__prompt_command() {{
local ret=$?
local env_file="${{SHPOOL_SESSION_DIR}}/forward.env"
local stamp_file="${{env_file}}.stamp"
if [ -n "${{SHPOOL_SESSION_DIR}}" ] && [ -f "${{env_file}}" ]; then
if [ ! -f "${{stamp_file}}" ] || [ "${{env_file}}" -nt "${{stamp_file}}" ]; then
touch -r "${{env_file}}" "${{stamp_file}}" 2>/dev/null

local allexport_was_set=0
case "$-" in
*a*) allexport_was_set=1 ;;
esac
set -a
. "${{env_file}}"
if [ "$allexport_was_set" -eq 0 ] ; then
set +a
fi
fi
fi

PS1="${{SHPOOL__OLD_PS1}}"
(exit $ret)
for prompt_hook in "${{SHPOOL__OLD_PROMPT_COMMAND[@]}}"
do
eval "${{prompt_hook}}"
ret=$?
done
PS1="{prompt_prefix}${{PS1}}"
return $ret
}}
unset PROMPT_COMMAND
PROMPT_COMMAND=__shpool__prompt_command
"#
),
)
}
(_, Ok(KnownShell::Zsh)) => format!(
r#"
typeset -a precmd_functions
SHPOOL__OLD_PROMPT="${{PROMPT}}"
function __shpool__reset_rprompt() {{
local ret=$?
local env_file="${{SHPOOL_SESSION_DIR:-}}/forward.env"
local stamp_file="${{env_file}}.stamp"
if [ -n "${{SHPOOL_SESSION_DIR:-}}" ] && [ -f "${{env_file}}" ]; then
if [ ! -f "${{stamp_file}}" ] || [ "${{env_file}}" -nt "${{stamp_file}}" ]; then
touch -r "${{env_file}}" "${{stamp_file}}" 2>/dev/null

local allexport_was_set=0
case "$-" in
*a*) allexport_was_set=1 ;;
esac
set -a
. "${{env_file}}"
if [ "$allexport_was_set" -eq 0 ] ; then
set +a
fi
fi
fi

PROMPT="${{SHPOOL__OLD_PROMPT}}"
return $ret
}}
precmd_functions[1,0]=(__shpool__reset_rprompt)
function __shpool__prompt_command() {{
local ret=$?
PROMPT="{prompt_prefix}${{PROMPT}}"
return $ret
}}
precmd_functions+=(__shpool__prompt_command)
"#
),
(_, Ok(KnownShell::Fish)) => format!(
r#"
functions --copy fish_prompt shpool__old_prompt
function fish_prompt; echo -n "{prompt_prefix}"; shpool__old_prompt; end
"#
),
(_, Ok(KnownShell::Fish)) => {
// Fish only added the `-nt` (newer-than) binary operator to its
// builtin `test` in fish 4.0b1. In older fish versions (such as
// fish 3.x), calling builtin `test -nt` errors with "unexpected
// argument". To maintain zero-fork prompt evaluation on
// fish 4+ while remaining compatible with fish 3, we
// probe for `-nt` support once at injection
// time and define `__shpool_is_newer` to use the builtin if
// available, falling back to `command test` (coreutils)
// otherwise.
format!(
r#"
functions --copy fish_prompt shpool__old_prompt
function __shpool_set_status; return $argv[1]; end
set -l __shpool_nt_err (test /dev/null -nt /dev/null 2>&1)
if test -z "$__shpool_nt_err"
function __shpool_is_newer; test $argv[1] -nt $argv[2]; end
else
function __shpool_is_newer; command test $argv[1] -nt $argv[2]; end
end
function fish_prompt
set -l last_status $status
set -l env_file "$SHPOOL_SESSION_DIR/forward.env"
set -l stamp_file "$env_file.stamp"
if test -n "$SHPOOL_SESSION_DIR"; and test -f "$env_file"
if test ! -f "$stamp_file"; or __shpool_is_newer "$env_file" "$stamp_file"
touch -r "$env_file" "$stamp_file" 2>/dev/null
source "$env_file"
end
end
echo -n "{prompt_prefix}"
__shpool_set_status $last_status
shpool__old_prompt
end
"#
)
}
(_, Err(e)) => {
warn!("could not sniff shell: {}", e);

Expand Down
Loading
Loading