Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
1 change: 1 addition & 0 deletions crates/atuin-desktop-runtime/src/blocks/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ pub(crate) mod script;
pub(crate) mod sql_block;
pub(crate) mod sqlite;
pub(crate) mod ssh_connect;

pub(crate) mod sub_runbook;
pub(crate) mod terminal;
pub(crate) mod var;
Expand Down
11 changes: 9 additions & 2 deletions crates/atuin-desktop-runtime/src/blocks/script.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use crate::events::GCEvent;
use crate::execution::{
CancellationToken, ExecutionContext, ExecutionHandle, ExecutionStatus, StreamingBlockOutput,
};
use crate::ssh::{OutputLine as SessionOutputLine, SshWarning};
use crate::ssh::{build_env_exports, OutputLine as SessionOutputLine, SshWarning};

use super::FromDocument;

Expand Down Expand Up @@ -697,8 +697,15 @@ impl Script {
None
};

let env_exports = build_env_exports(context.context_resolver.env_vars());

let code_to_run = if let Some(ref path) = remote_temp_path {
format!("export ATUIN_OUTPUT_VARS='{}'\n{}", path, code)
format!(
"{}export ATUIN_OUTPUT_VARS='{}'\n{}",
env_exports, path, code
)
} else if !env_exports.is_empty() {
format!("{}{}", env_exports, code)
} else {
code.to_string()
};
Expand Down
11 changes: 10 additions & 1 deletion crates/atuin-desktop-runtime/src/blocks/terminal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ use crate::execution::{
CancellationToken, ExecutionContext, ExecutionHandle, ExecutionStatus, StreamingBlockOutput,
};
use crate::pty::{Pty, PtyLike};
use crate::ssh::{SshPty, SshWarning};
use crate::ssh::{build_env_exports, SshPty, SshWarning};

/// Output structure for Terminal blocks that implements BlockExecutionOutput
/// for template access to terminal output.
Expand Down Expand Up @@ -491,6 +491,15 @@ impl Terminal {
.emit_gc_event(GCEvent::PtyOpened(metadata.clone()))
.await;

if ssh_host.is_some() {
let env_exports = build_env_exports(context.context_resolver.env_vars());
if !env_exports.is_empty() {
if let Err(e) = pty_store.write_pty(self.id, env_exports.into()).await {
tracing::warn!("Failed to write env exports to SSH PTY: {}", e);
}
}
}

if let Some(ref remote_path) = remote_var_path {
let export_cmd = format!("export ATUIN_OUTPUT_VARS='{}'\n", remote_path);
if let Err(e) = pty_store.write_pty(self.id, export_cmd.into()).await {
Expand Down
2 changes: 2 additions & 0 deletions crates/atuin-desktop-runtime/src/ssh/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,11 +22,13 @@

mod pool;
mod session;
mod ssh_env;
mod ssh_pool;

#[cfg(test)]
mod integration_tests;

pub use pool::Pool;
pub use session::{Authentication, CommandResult, OutputLine, Session, SshConfig, SshWarning};
pub use ssh_env::build_env_exports;
pub use ssh_pool::{SshPoolHandle, SshPty};
102 changes: 102 additions & 0 deletions crates/atuin-desktop-runtime/src/ssh/ssh_env.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
use std::collections::HashMap;

/// Escape a value for safe inclusion in a POSIX single-quoted string.
///
/// Replaces each `'` with `'\''` (end quote, escaped literal quote, start new quote)
/// and wraps the result in single quotes. POSIX single-quoted strings treat all
/// characters as literal, so this handles `$`, backticks, newlines, etc.
fn shell_escape_value(value: &str) -> String {
let mut escaped = String::with_capacity(value.len() + 2);
escaped.push('\'');
for ch in value.chars() {
if ch == '\'' {
escaped.push_str("'\\''");
} else {
escaped.push(ch);
}
}
escaped.push('\'');
escaped
}

/// Build `export K='V'\n` lines for all env vars.
///
/// Returns an empty string if the map is empty.
pub fn build_env_exports(env_vars: &HashMap<String, String>) -> String {
if env_vars.is_empty() {
return String::new();
}

let mut out = String::new();
for (key, value) in env_vars {
out.push_str("export ");
out.push_str(key);
out.push('=');
out.push_str(&shell_escape_value(value));
out.push('\n');
}
out
}

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

#[test]
fn test_shell_escape_simple_value() {
assert_eq!(shell_escape_value("hello"), "'hello'");
}

#[test]
fn test_shell_escape_empty_value() {
assert_eq!(shell_escape_value(""), "''");
}

#[test]
fn test_shell_escape_value_with_single_quotes() {
assert_eq!(shell_escape_value("it's"), "'it'\\''s'");
}

#[test]
fn test_shell_escape_value_with_special_chars() {
assert_eq!(shell_escape_value("$HOME"), "'$HOME'");
assert_eq!(shell_escape_value("`cmd`"), "'`cmd`'");
assert_eq!(shell_escape_value("a\nb"), "'a\nb'");
assert_eq!(shell_escape_value("a b"), "'a b'");
assert_eq!(shell_escape_value("a\"b"), "'a\"b'");
}

#[test]
fn test_build_env_exports_empty_map() {
let map = HashMap::new();
assert_eq!(build_env_exports(&map), "");
}

#[test]
fn test_build_env_exports_single_var() {
let mut map = HashMap::new();
map.insert("FOO".to_string(), "bar".to_string());
assert_eq!(build_env_exports(&map), "export FOO='bar'\n");
}

#[test]
fn test_build_env_exports_special_value() {
let mut map = HashMap::new();
map.insert("VAR".to_string(), "it's $complex".to_string());
assert_eq!(build_env_exports(&map), "export VAR='it'\\''s $complex'\n");
}

#[test]
fn test_build_env_exports_multiple_vars() {
let mut map = HashMap::new();
map.insert("A".to_string(), "1".to_string());
map.insert("B".to_string(), "2".to_string());

let result = build_env_exports(&map);
// HashMap order is not guaranteed, so check both lines are present
assert!(result.contains("export A='1'\n"));
assert!(result.contains("export B='2'\n"));
// Should have exactly 2 lines
assert_eq!(result.lines().count(), 2);
}
}
Loading