Skip to main content

cli/ssh/
mod.rs

1//! `shine ssh`: wraps the system `ssh` binary to establish an interactive
2//! session that also carries a session-scoped file-transfer channel back to
3//! the local machine (see docs/ssh-local-transfer-prd.md).
4//!
5//! Architecture, confirmed against a real host via `scripts/spike-ssh-forward.sh`:
6//! - We prepend our own `-R <remote-sock>:<local-forward-target>` to the
7//!   user's ssh args (safe: ssh options may appear in any order before the
8//!   destination).
9//! - We replace the remote command with a wrapper that sets
10//!   `SHINE_SSH_SESSION`/`SHINE_SSH_TOKEN`/`SHINE_SSH_REMOTE_SOCK` via `env`
11//!   (not `SetEnv`/`SendEnv`, which most sshd configs don't accept), then
12//!   `exec`s either the user's original remote command or their login shell.
13//!   Explicit `--with`/`--with-secret` values join that process environment.
14//! - sshd does NOT clean up the forwarded remote socket file on disconnect
15//!   (confirmed by the spike), so the wrapper registers its own `trap ...
16//!   EXIT` to remove it.
17//!
18//! The default remote mode is POSIX: it uses a Unix socket and a POSIX shell
19//! wrapper regardless of the *local* platform. `--remote-shell windows` is
20//! an explicit environment-forwarding-only mode: it sends a Base64-encoded
21//! PowerShell bootstrap, preferring PowerShell 7 (`pwsh.exe`) and falling
22//! back to Windows PowerShell (`powershell.exe`), and deliberately creates no
23//! transfer listener or `-R` forward. Locally, the POSIX path's
24//! `bind_local_listener` uses a Unix socket on macOS/Linux, or loopback TCP on
25//! Windows.
26
27mod agent;
28// Drives the real agent over an in-process Unix socket pair, so it is
29// unix-only (Windows is the local side only and has no `UnixListener`).
30#[cfg(all(test, unix))]
31mod integration_tests;
32mod protocol;
33mod session_context;
34// `remote_client` dials the forwarded socket via a Unix stream: it only
35// ever runs on the *remote* end of a session, which is always assumed
36// Linux/macOS (see module docs), so it is unconditionally unix-only —
37// unlike `agent`, which must compile on Windows too since Windows is
38// supported as the *local* side.
39#[cfg(unix)]
40mod remote_client;
41
42use std::collections::{BTreeMap, BTreeSet};
43
44use anyhow::{Context, Result, bail};
45use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
46
47use crate::commands::RemoteShell;
48use crate::config::Config;
49use crate::env::{EnvConfig, parse_env_specs, secret_key};
50use crate::secret;
51use crate::theme;
52
53/// Grace period given to still-running per-connection transfer tasks to
54/// notice the (by now closed) `ssh` tunnel and finish their own cleanup
55/// before the session directory is removed. See
56/// `agent::drain_connection_tasks`.
57const CONNECTION_DRAIN_GRACE_PERIOD: std::time::Duration = std::time::Duration::from_secs(5);
58
59#[cfg(not(unix))]
60const WINDOWS_REMOTE_UNSUPPORTED: &str = "`shine local` commands require this machine to be the \
61    remote (Linux/macOS) side of a `shine ssh` session; Windows is currently supported as the \
62    local side only";
63
64#[cfg(unix)]
65pub async fn handle_local_download(
66    remote_source: &str,
67    local_destination: Option<&str>,
68    force: bool,
69    dry_run: bool,
70    use_scp: bool,
71) -> Result<()> {
72    remote_client::handle_download(remote_source, local_destination, force, dry_run, use_scp).await
73}
74
75#[cfg(not(unix))]
76pub async fn handle_local_download(
77    _remote_source: &str,
78    _local_destination: Option<&str>,
79    _force: bool,
80    _dry_run: bool,
81    _use_scp: bool,
82) -> Result<()> {
83    bail!(WINDOWS_REMOTE_UNSUPPORTED)
84}
85
86#[cfg(unix)]
87pub async fn handle_local_upload(
88    local_source: &str,
89    remote_destination: Option<&str>,
90    force: bool,
91    dry_run: bool,
92    use_scp: bool,
93) -> Result<()> {
94    remote_client::handle_upload(local_source, remote_destination, force, dry_run, use_scp).await
95}
96
97#[cfg(not(unix))]
98pub async fn handle_local_upload(
99    _local_source: &str,
100    _remote_destination: Option<&str>,
101    _force: bool,
102    _dry_run: bool,
103    _use_scp: bool,
104) -> Result<()> {
105    bail!(WINDOWS_REMOTE_UNSUPPORTED)
106}
107
108#[cfg(unix)]
109pub async fn handle_local_status() -> Result<()> {
110    remote_client::handle_status().await
111}
112
113#[cfg(not(unix))]
114pub async fn handle_local_status() -> Result<()> {
115    bail!(WINDOWS_REMOTE_UNSUPPORTED)
116}
117
118/// Single-letter ssh options that consume a separate value, per ssh(1).
119/// Used only to locate the destination/command boundary in the user's
120/// argument list — never to reinterpret what the options mean.
121const VALUE_OPTION_LETTERS: &[char] = &[
122    'B', 'b', 'c', 'D', 'E', 'e', 'F', 'I', 'i', 'J', 'L', 'l', 'm', 'O', 'o', 'p', 'Q', 'R', 'S',
123    'W', 'w',
124];
125
126pub async fn handle_ssh(
127    config: &Config,
128    remote_shell: RemoteShell,
129    with: &[String],
130    with_secret: &[String],
131    args: &[String],
132) -> Result<()> {
133    let (ssh_options, host, remote_command) = split_ssh_args(args)?;
134    let forwarded_env = resolve_forwarded_env(config, with, with_secret).await?;
135
136    // Windows OpenSSH executes remote commands through cmd.exe by default.
137    // Do not create the POSIX-only transfer channel there; the encoded
138    // PowerShell command has no user-controlled syntax in cmd.exe.
139    if remote_shell == RemoteShell::Windows {
140        let session_id = uuid::Uuid::new_v4().to_string();
141        let local_theme = theme::resolve_local_terminal_theme_for_injection();
142        let wrapped_command = build_windows_wrapped_remote_command(
143            &session_id,
144            local_theme.map(theme::Theme::as_str),
145            &forwarded_env,
146            &remote_command,
147        )?;
148        let mut cmd = tokio::process::Command::new("ssh");
149        cmd.args(build_windows_ssh_invocation_args(
150            &ssh_options,
151            &host,
152            &wrapped_command,
153        ));
154        return finish_ssh_status(run_ssh_with_ctrl_c(&mut cmd).await?);
155    }
156
157    let session_id = uuid::Uuid::new_v4().to_string();
158    let token = uuid::Uuid::new_v4().to_string();
159
160    let session_dir = config.shine_dir().join("run").join("ssh").join(&session_id);
161    tokio::fs::create_dir_all(&session_dir)
162        .await
163        .with_context(|| format!("creating {}", session_dir.display()))?;
164    // The remote host is always assumed Linux/macOS (see module docs), so
165    // its socket is always a Unix socket regardless of the local platform.
166    let remote_sock = format!("/tmp/.shine-ssh-{session_id}.sock");
167
168    let (listener, local_forward_target) = bind_local_listener(&session_dir).await?;
169    let session_local_dir = std::env::current_dir().context("reading current directory")?;
170
171    // Reuse the interactive connection as a control master so the rsync/scp
172    // child reconnects over it with no second authentication (ADR 0011). Skip
173    // if the user already configured their own multiplexing, so we don't fight
174    // their settings.
175    let control_options = if session_context::user_set_control_options(&ssh_options) {
176        None
177    } else {
178        Some(session_dir.join("ctl.sock"))
179    };
180
181    let context = std::sync::Arc::new(session_context::SessionContext {
182        host: host.clone(),
183        ssh_options: ssh_options.clone(),
184        local_dir: session_local_dir.clone(),
185        control_path: control_options.clone(),
186    });
187    context.save(&session_dir).await?;
188
189    let connection_tasks = agent::new_connection_tasks();
190    let agent_handle =
191        tokio::spawn(listener.serve(token.clone(), context.clone(), connection_tasks.clone()));
192
193    // Query the *local* terminal — same-host, sub-millisecond round trip,
194    // no fragmentation risk unlike a remote OSC query (PRD §2.2/§6.1) — so
195    // the remote login shell never has to guess at its own theme.
196    let local_theme = theme::resolve_local_terminal_theme_for_injection();
197    let wrapped_command = build_wrapped_remote_command(
198        &session_id,
199        &token,
200        &remote_sock,
201        local_theme.map(theme::Theme::as_str),
202        &forwarded_env,
203        &remote_command,
204    );
205
206    let mut cmd = tokio::process::Command::new("ssh");
207    cmd.args(build_ssh_invocation_args(
208        &ssh_options,
209        &remote_sock,
210        &local_forward_target,
211        control_options.as_deref(),
212        &host,
213        &wrapped_command,
214    ));
215
216    // Racing against ctrl_c() (rather than just awaiting cmd.status()) is
217    // what makes the cleanup below actually run on Ctrl-C: installing this
218    // listener overrides SIGINT's default disposition for the process, so a
219    // Ctrl-C no longer kills us before we get a chance to clean up. The ssh
220    // child is in the same foreground process group and receives SIGINT
221    // independently; we still await its exit so we don't race it.
222    let status = run_ssh_with_ctrl_c(&mut cmd).await?;
223
224    // Stop accepting new connections, then give any still-running transfer
225    // a bounded chance to notice the tunnel is gone and run its own
226    // cleanup before we remove the session directory out from under it.
227    agent_handle.abort();
228    agent::drain_connection_tasks(&connection_tasks, CONNECTION_DRAIN_GRACE_PERIOD).await;
229    let _ = tokio::fs::remove_dir_all(&session_dir).await;
230
231    finish_ssh_status(status)
232}
233
234async fn run_ssh_with_ctrl_c(
235    cmd: &mut tokio::process::Command,
236) -> Result<std::process::ExitStatus> {
237    let mut ssh_run = std::pin::pin!(cmd.status());
238    tokio::select! {
239        status = &mut ssh_run => status,
240        _ = tokio::signal::ctrl_c() => ssh_run.await,
241    }
242    .context("failed to run ssh")
243}
244
245fn finish_ssh_status(status: std::process::ExitStatus) -> Result<()> {
246    if status.success() {
247        return Ok(());
248    }
249    if let Some(code) = status.code() {
250        std::process::exit(code);
251    }
252    #[cfg(unix)]
253    {
254        use std::os::unix::process::ExitStatusExt;
255        std::process::exit(128 + status.signal().unwrap_or(1));
256    }
257    #[cfg(not(unix))]
258    std::process::exit(1);
259}
260
261const RESERVED_REMOTE_ENV: &[&str] = &[
262    "SHINE_SSH_SESSION",
263    "SHINE_SSH_TOKEN",
264    "SHINE_SSH_REMOTE_SOCK",
265    "SHINE_TERMINAL_THEME",
266];
267
268/// Resolves only explicitly selected config values. Plaintext selection is
269/// deliberately exact: unlike `shine env run --with`, it never falls through
270/// to `<KEY>_SECRET`. Sending decrypted material to another host requires the
271/// visibly distinct `--with-secret` opt-in.
272async fn resolve_forwarded_env(
273    config: &Config,
274    with: &[String],
275    with_secret: &[String],
276) -> Result<BTreeMap<String, String>> {
277    let plain_specs = parse_env_specs(with)?;
278    let secret_specs = parse_env_specs(with_secret)?;
279    let env = EnvConfig::load_or_init(config).await?;
280    let mut targets = BTreeSet::new();
281    let mut resolved = BTreeMap::new();
282
283    for spec in plain_specs {
284        validate_forward_target(&spec.target, &mut targets)?;
285        if spec.source.ends_with("_SECRET") {
286            bail!(
287                "--with does not inject secret storage key {}; use --with-secret with the base key instead",
288                spec.source
289            );
290        }
291        let value = env.get(&spec.source).with_context(|| {
292            let encrypted = secret_key(&spec.source);
293            if env.get(&encrypted).is_some() {
294                format!(
295                    "{} is stored as {encrypted}; use --with-secret {} to decrypt and inject it",
296                    spec.source, spec.source
297                )
298            } else {
299                format!("{} is not set in the active config [env]", spec.source)
300            }
301        })?;
302        resolved.insert(spec.target, value.to_string());
303    }
304
305    for spec in secret_specs {
306        validate_forward_target(&spec.target, &mut targets)?;
307        if spec.source.ends_with("_SECRET") {
308            bail!(
309                "--with-secret expects a base key without the _SECRET suffix: {}",
310                spec.source
311            );
312        }
313        let encrypted = secret_key(&spec.source);
314        let ciphertext = env
315            .get(&encrypted)
316            .with_context(|| format!("{encrypted} is not set in the active config [env]"))?;
317        let value = secret::decrypt_secret(ciphertext, &config.age_identities())
318            .await
319            .with_context(|| format!("decrypting {encrypted}"))?;
320        resolved.insert(spec.target, value);
321    }
322
323    Ok(resolved)
324}
325
326fn validate_forward_target(target: &str, targets: &mut BTreeSet<String>) -> Result<()> {
327    if RESERVED_REMOTE_ENV.contains(&target) {
328        bail!("cannot override shine-managed SSH variable {target}");
329    }
330    if !targets.insert(target.to_string()) {
331        bail!("duplicate target variable: {target}");
332    }
333    Ok(())
334}
335
336/// Binds the local end of the session's transfer channel and returns it
337/// together with the target to embed in `ssh`'s `-R <remote-sock>:<target>`
338/// argument.
339#[cfg(unix)]
340async fn bind_local_listener(
341    session_dir: &std::path::Path,
342) -> Result<(agent::LocalListener, String)> {
343    let local_sock = session_dir.join("local.sock");
344    let listener = tokio::net::UnixListener::bind(&local_sock)
345        .with_context(|| format!("binding local transfer socket {}", local_sock.display()))?;
346    Ok((
347        agent::LocalListener::Unix(listener),
348        local_sock.display().to_string(),
349    ))
350}
351
352/// Windows lacks the mature, well-tested Unix-domain-socket support that
353/// macOS/Linux have, so the local end uses a loopback TCP socket instead;
354/// `ssh -R` supports mixing this with the remote's Unix-socket endpoint
355/// (verified via `scripts/spike-ssh-forward-windows.ps1`).
356#[cfg(windows)]
357async fn bind_local_listener(
358    _session_dir: &std::path::Path,
359) -> Result<(agent::LocalListener, String)> {
360    let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
361        .await
362        .context("binding local transfer TCP listener")?;
363    let port = listener
364        .local_addr()
365        .context("reading local TCP listener port")?
366        .port();
367    Ok((
368        agent::LocalListener::Tcp(listener),
369        format!("127.0.0.1:{port}"),
370    ))
371}
372
373/// Splits a raw `shine ssh` argument list into ssh options, the destination,
374/// and an optional remote command — mirroring what `ssh` itself would infer,
375/// without reinterpreting the options' meaning (see module docs). An
376/// explicit `--` may be used to disambiguate; it is consumed here and not
377/// forwarded to the real `ssh` invocation.
378fn split_ssh_args(args: &[String]) -> Result<(Vec<String>, String, Vec<String>)> {
379    let mut ssh_options = Vec::new();
380    let mut i = 0;
381    while i < args.len() {
382        let token = &args[i];
383        if token == "--" {
384            i += 1;
385            break;
386        }
387        if token == "-" || !token.starts_with('-') {
388            let host = token.clone();
389            let remote_command = args[i + 1..].to_vec();
390            return Ok((ssh_options, host, remote_command));
391        }
392
393        ssh_options.push(token.clone());
394        let letters: Vec<char> = token.chars().skip(1).collect();
395        let mut consumes_next = false;
396        for (idx, letter) in letters.iter().enumerate() {
397            if VALUE_OPTION_LETTERS.contains(letter) {
398                consumes_next = idx == letters.len() - 1;
399                break;
400            }
401        }
402        i += 1;
403        if consumes_next {
404            let Some(value) = args.get(i) else {
405                bail!("ssh option {token} requires a value");
406            };
407            ssh_options.push(value.clone());
408            i += 1;
409        }
410    }
411
412    let Some(host) = args.get(i) else {
413        bail!("no SSH destination given; usage: shine ssh [SSH_ARGS]... <HOST> [COMMAND]");
414    };
415    let remote_command = args[i + 1..].to_vec();
416    Ok((ssh_options, host.clone(), remote_command))
417}
418
419/// Assembles the argument list passed to the `ssh` binary: the user's own
420/// options first (untouched, per module docs), then our `-t`/`-R` forward,
421/// the destination, and the wrapped remote command. Kept as a pure function
422/// so the composition can be unit-tested without spawning a real `ssh`.
423fn build_ssh_invocation_args(
424    ssh_options: &[String],
425    remote_sock: &str,
426    local_forward_target: &str,
427    control_path: Option<&std::path::Path>,
428    host: &str,
429    wrapped_command: &str,
430) -> Vec<String> {
431    let mut args = ssh_options.to_vec();
432    // Enable connection multiplexing so a later `rsync`/`scp` child can reuse
433    // this authenticated master connection (ADR 0011). Only injected when the
434    // user didn't set their own ControlMaster/ControlPath.
435    if let Some(control_path) = control_path {
436        args.push("-o".to_string());
437        args.push("ControlMaster=auto".to_string());
438        args.push("-o".to_string());
439        args.push(format!("ControlPath={}", control_path.display()));
440        args.push("-o".to_string());
441        args.push("ControlPersist=60".to_string());
442    }
443    args.push("-t".to_string());
444    args.push("-R".to_string());
445    args.push(format!("{remote_sock}:{local_forward_target}"));
446    args.push(host.to_string());
447    args.push(wrapped_command.to_string());
448    args
449}
450
451/// Windows does not receive a transfer listener, so its SSH invocation is a
452/// deliberately small normal TTY session followed by one opaque PowerShell
453/// command. Keeping the encoded payload as a single argv item prevents CMD
454/// from seeing secret values or PowerShell metacharacters.
455fn build_windows_ssh_invocation_args(
456    ssh_options: &[String],
457    host: &str,
458    wrapped_command: &str,
459) -> Vec<String> {
460    let mut args = ssh_options.to_vec();
461    args.push("-t".to_string());
462    args.push(host.to_string());
463    args.push(wrapped_command.to_string());
464    args
465}
466
467fn build_wrapped_remote_command(
468    session_id: &str,
469    token: &str,
470    remote_sock: &str,
471    local_theme: Option<&str>,
472    forwarded_env: &BTreeMap<String, String>,
473    remote_command: &[String],
474) -> String {
475    let inner_exec = if remote_command.is_empty() {
476        r#"exec "$SHELL" -l"#.to_string()
477    } else {
478        let quoted = remote_command
479            .iter()
480            .map(|token| single_quote(token))
481            .collect::<Vec<_>>()
482            .join(" ");
483        format!("exec {quoted}")
484    };
485    // Double quotes here are safe: this text is only ever embedded through
486    // `single_quote`, which POSIX-escapes it as one opaque literal for the
487    // outer shell, so nothing inside (single or double quotes, `$`, etc.)
488    // is interpreted until the inner `sh -c` re-parses it.
489    let inner_script = format!(r#"trap "rm -f $SHINE_SSH_REMOTE_SOCK" EXIT; {inner_exec}"#);
490
491    let mut env_prefix = format!(
492        "SHINE_SSH_SESSION={session_id} SHINE_SSH_TOKEN={token} SHINE_SSH_REMOTE_SOCK={remote_sock}"
493    );
494    // Unlike the three values above (internally generated UUIDs/hex/paths,
495    // never user input), this one is quoted defensively per
496    // docs/terminal-theme-sync-prd.md §6.1/§10 even though its source
497    // (`Theme::as_str`) only ever produces the literal `light` or `dark`.
498    if let Some(theme) = local_theme {
499        env_prefix.push_str(&format!(" SHINE_TERMINAL_THEME={}", single_quote(theme)));
500    }
501    for (key, value) in forwarded_env {
502        env_prefix.push_str(&format!(" {key}={}", single_quote(value)));
503    }
504
505    format!("env {env_prefix} sh -c {}", single_quote(&inner_script))
506}
507
508/// Builds an opaque Windows PowerShell remote command. OpenSSH passes this
509/// through cmd.exe on typical Windows servers, so only the Base64 alphabet is
510/// allowed to carry user-controlled values across that boundary. A small
511/// Windows PowerShell bootstrap probes for PowerShell 7 before launching the
512/// real encoded payload in the selected shell. Keeping the outer command free
513/// of CMD operators also works when the SSH server's default shell has already
514/// been changed from CMD to PowerShell.
515fn build_windows_wrapped_remote_command(
516    session_id: &str,
517    local_theme: Option<&str>,
518    forwarded_env: &BTreeMap<String, String>,
519    remote_command: &[String],
520) -> Result<String> {
521    let mut script = String::new();
522    push_powershell_env_assignment(&mut script, "SHINE_SSH_SESSION", session_id)?;
523    if let Some(theme) = local_theme {
524        push_powershell_env_assignment(&mut script, "SHINE_TERMINAL_THEME", theme)?;
525    }
526    for (key, value) in forwarded_env {
527        push_powershell_env_assignment(&mut script, key, value)?;
528    }
529
530    let interactive = remote_command.is_empty();
531    if !interactive {
532        // Reset the native-program status so a PowerShell command does not
533        // accidentally inherit a status from profile/startup execution.
534        script.push_str("$global:LASTEXITCODE = 0\n& ");
535        for (index, argument) in remote_command.iter().enumerate() {
536            if index > 0 {
537                script.push(' ');
538            }
539            script.push_str(&powershell_single_quoted_literal(argument)?);
540        }
541        script.push_str(
542            "\nif ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE }\nif (-not $?) { exit 1 }\n",
543        );
544    }
545
546    let payload_encoded = BASE64.encode(utf16le_bytes(&script)?);
547    let no_profile = if interactive { "" } else { " -NoProfile" };
548    let no_exit = if interactive { " -NoExit" } else { "" };
549    let bootstrap = format!(
550        "$pwsh = Get-Command pwsh.exe -CommandType Application -ErrorAction SilentlyContinue | Select-Object -First 1\n\
551         $shell = if ($null -ne $pwsh) {{ $pwsh.Source }} else {{ 'powershell.exe' }}\n\
552         & $shell{no_profile}{no_exit} -EncodedCommand '{payload_encoded}'\n\
553         $ok = $?\n\
554         $code = $LASTEXITCODE\n\
555         if ($null -ne $code -and $code -ne 0) {{ exit $code }}\n\
556         if (-not $ok) {{ exit 1 }}\n"
557    );
558    let bootstrap_encoded = BASE64.encode(utf16le_bytes(&bootstrap)?);
559    Ok(format!(
560        "powershell.exe -NoProfile -EncodedCommand {bootstrap_encoded}"
561    ))
562}
563
564fn push_powershell_env_assignment(script: &mut String, key: &str, value: &str) -> Result<()> {
565    // Targets have already been parsed as environment identifiers by
566    // `parse_env_specs`; keep this check local so this builder remains safe
567    // if it is reused independently later.
568    if !key
569        .bytes()
570        .all(|byte| byte.is_ascii_alphanumeric() || byte == b'_')
571    {
572        bail!("cannot safely represent Windows environment variable name {key}");
573    }
574    script.push_str("$env:");
575    script.push_str(key);
576    script.push_str(" = ");
577    script.push_str(&powershell_single_quoted_literal(value)?);
578    script.push('\n');
579    Ok(())
580}
581
582/// Produces a PowerShell single-quoted string. PowerShell represents a literal
583/// apostrophe by doubling it; NUL is rejected because Windows environment
584/// values and command-line APIs cannot represent it safely.
585fn powershell_single_quoted_literal(value: &str) -> Result<String> {
586    if value.contains('\0') {
587        bail!("cannot forward values containing NUL bytes to Windows PowerShell");
588    }
589    Ok(format!("'{}'", value.replace('\'', "''")))
590}
591
592fn utf16le_bytes(value: &str) -> Result<Vec<u8>> {
593    if value.contains('\0') {
594        bail!("cannot encode PowerShell commands containing NUL bytes");
595    }
596    Ok(value
597        .encode_utf16()
598        .flat_map(u16::to_le_bytes)
599        .collect::<Vec<_>>())
600}
601
602/// POSIX single-quotes `s` for safe embedding as one shell word, escaping
603/// any literal `'` via the standard `'\''` idiom (close quote, escaped
604/// quote, reopen quote). Applying this at each nesting level independently
605/// composes correctly regardless of how many quoting layers are involved.
606fn single_quote(s: &str) -> String {
607    format!("'{}'", s.replace('\'', r"'\''"))
608}
609
610#[cfg(test)]
611mod tests {
612    use super::*;
613
614    #[test]
615    fn plain_host_with_no_options_or_command() {
616        let (options, host, command) = split_ssh_args(&["dev".to_string()]).unwrap();
617        assert!(options.is_empty());
618        assert_eq!(host, "dev");
619        assert!(command.is_empty());
620    }
621
622    #[test]
623    fn host_followed_by_a_remote_command() {
624        let args = vec!["dev".to_string(), "ls".to_string(), "-la".to_string()];
625        let (options, host, command) = split_ssh_args(&args).unwrap();
626        assert!(options.is_empty());
627        assert_eq!(host, "dev");
628        assert_eq!(command, vec!["ls", "-la"]);
629    }
630
631    #[test]
632    fn value_option_with_separate_token() {
633        let args = vec!["-p".to_string(), "2222".to_string(), "dev".to_string()];
634        let (options, host, command) = split_ssh_args(&args).unwrap();
635        assert_eq!(options, vec!["-p", "2222"]);
636        assert_eq!(host, "dev");
637        assert!(command.is_empty());
638    }
639
640    #[test]
641    fn value_option_with_attached_value() {
642        let args = vec!["-p2222".to_string(), "dev".to_string()];
643        let (options, host, _command) = split_ssh_args(&args).unwrap();
644        assert_eq!(options, vec!["-p2222"]);
645        assert_eq!(host, "dev");
646    }
647
648    #[test]
649    fn repeated_o_option() {
650        let args = vec![
651            "-o".to_string(),
652            "ProxyJump=bastion".to_string(),
653            "dev".to_string(),
654        ];
655        let (options, host, _command) = split_ssh_args(&args).unwrap();
656        assert_eq!(options, vec!["-o", "ProxyJump=bastion"]);
657        assert_eq!(host, "dev");
658    }
659
660    #[test]
661    fn bundled_boolean_flags_consume_no_value() {
662        let args = vec!["-vvv".to_string(), "dev".to_string()];
663        let (options, host, _command) = split_ssh_args(&args).unwrap();
664        assert_eq!(options, vec!["-vvv"]);
665        assert_eq!(host, "dev");
666    }
667
668    #[test]
669    fn explicit_double_dash_separator() {
670        let args = vec!["--".to_string(), "dev".to_string(), "ls".to_string()];
671        let (options, host, command) = split_ssh_args(&args).unwrap();
672        assert!(options.is_empty());
673        assert_eq!(host, "dev");
674        assert_eq!(command, vec!["ls"]);
675    }
676
677    #[test]
678    fn no_destination_is_an_error() {
679        assert!(split_ssh_args(&[]).is_err());
680    }
681
682    #[test]
683    fn dangling_value_option_is_an_error() {
684        assert!(split_ssh_args(&["-p".to_string()]).is_err());
685    }
686
687    #[test]
688    fn wrapped_command_round_trips_through_a_real_shell() {
689        // Exercises the full nested-quoting composition end to end: run the
690        // wrapped command through `sh -c` and check the remote command's
691        // stdout, rather than reasoning about escaping by hand.
692        let wrapped = build_wrapped_remote_command(
693            "sid",
694            "tok",
695            "/tmp/shine-ssh-mod-test-sid.sock",
696            None,
697            &BTreeMap::new(),
698            &["echo".to_string(), "it's a test".to_string()],
699        );
700        let output = std::process::Command::new("sh")
701            .arg("-c")
702            .arg(&wrapped)
703            .output()
704            .expect("failed to run sh");
705        assert!(output.status.success(), "stderr: {:?}", output.stderr);
706        assert_eq!(
707            String::from_utf8_lossy(&output.stdout).trim_end(),
708            "it's a test"
709        );
710    }
711
712    #[test]
713    fn ssh_invocation_args_keep_user_options_verbatim_and_ahead_of_our_own() {
714        let args = vec!["-J".to_string(), "bastion".to_string()];
715        let (parsed_options, host, _command) =
716            split_ssh_args(&[args.clone(), vec!["dev".to_string()]].concat()).unwrap();
717
718        let invocation = build_ssh_invocation_args(
719            &parsed_options,
720            "/tmp/.shine-ssh-sid.sock",
721            "/tmp/shine-ssh-sid/local.sock",
722            None,
723            &host,
724            "wrapped-command",
725        );
726
727        assert_eq!(
728            invocation,
729            vec![
730                "-J",
731                "bastion",
732                "-t",
733                "-R",
734                "/tmp/.shine-ssh-sid.sock:/tmp/shine-ssh-sid/local.sock",
735                "dev",
736                "wrapped-command",
737            ]
738        );
739    }
740
741    #[test]
742    fn windows_ssh_invocation_has_no_transfer_or_posix_wrapper() {
743        let invocation = build_windows_ssh_invocation_args(
744            &["-p".to_string(), "2222".to_string()],
745            "windows-host",
746            "powershell.exe -NoProfile -EncodedCommand QQ==",
747        );
748
749        assert_eq!(
750            invocation,
751            vec![
752                "-p",
753                "2222",
754                "-t",
755                "windows-host",
756                "powershell.exe -NoProfile -EncodedCommand QQ==",
757            ]
758        );
759        assert!(!invocation.iter().any(|arg| arg == "-R"));
760        assert!(!invocation.iter().any(|arg| arg.contains("env ")));
761        assert!(!invocation.iter().any(|arg| arg.contains("sh -c")));
762    }
763
764    #[test]
765    fn ssh_invocation_args_inject_control_master_when_control_path_given() {
766        let (parsed_options, host, _command) = split_ssh_args(&["dev".to_string()]).unwrap();
767        let invocation = build_ssh_invocation_args(
768            &parsed_options,
769            "/tmp/.shine-ssh-sid.sock",
770            "/tmp/shine-ssh-sid/local.sock",
771            Some(std::path::Path::new("/tmp/shine-ssh-sid/ctl.sock")),
772            &host,
773            "wrapped-command",
774        );
775
776        assert_eq!(
777            invocation,
778            vec![
779                "-o",
780                "ControlMaster=auto",
781                "-o",
782                "ControlPath=/tmp/shine-ssh-sid/ctl.sock",
783                "-o",
784                "ControlPersist=60",
785                "-t",
786                "-R",
787                "/tmp/.shine-ssh-sid.sock:/tmp/shine-ssh-sid/local.sock",
788                "dev",
789                "wrapped-command",
790            ]
791        );
792    }
793
794    #[test]
795    fn ssh_invocation_args_preserve_repeated_o_options_in_order() {
796        let args = vec![
797            "-o".to_string(),
798            "ProxyJump=bastion".to_string(),
799            "-o".to_string(),
800            "ServerAliveInterval=30".to_string(),
801            "dev".to_string(),
802            "ls".to_string(),
803            "-la".to_string(),
804        ];
805        let (parsed_options, host, command) = split_ssh_args(&args).unwrap();
806        assert_eq!(command, vec!["ls", "-la"]);
807
808        let invocation = build_ssh_invocation_args(
809            &parsed_options,
810            "/tmp/.shine-ssh-sid.sock",
811            "/tmp/shine-ssh-sid/local.sock",
812            None,
813            &host,
814            "wrapped-command",
815        );
816
817        // The user's repeated -o options must appear verbatim, in order, and
818        // ahead of our own -t/-R/host/command — never reordered or merged.
819        assert_eq!(
820            invocation,
821            vec![
822                "-o",
823                "ProxyJump=bastion",
824                "-o",
825                "ServerAliveInterval=30",
826                "-t",
827                "-R",
828                "/tmp/.shine-ssh-sid.sock:/tmp/shine-ssh-sid/local.sock",
829                "dev",
830                "wrapped-command",
831            ]
832        );
833    }
834
835    #[test]
836    fn wrapped_command_defaults_to_login_shell() {
837        let wrapped = build_wrapped_remote_command(
838            "sid",
839            "tok",
840            "/tmp/.shine-ssh-sid.sock",
841            None,
842            &BTreeMap::new(),
843            &[],
844        );
845        assert!(wrapped.contains(r#"exec "$SHELL" -l"#));
846        assert!(wrapped.contains("trap \"rm -f $SHINE_SSH_REMOTE_SOCK\" EXIT"));
847    }
848
849    #[test]
850    fn wrapped_command_omits_theme_var_when_none() {
851        let wrapped = build_wrapped_remote_command(
852            "sid",
853            "tok",
854            "/tmp/.shine-ssh-sid.sock",
855            None,
856            &BTreeMap::new(),
857            &[],
858        );
859        assert!(!wrapped.contains("SHINE_TERMINAL_THEME"));
860    }
861
862    #[test]
863    fn wrapped_command_injects_quoted_theme_var_when_present() {
864        let wrapped = build_wrapped_remote_command(
865            "sid",
866            "tok",
867            "/tmp/.shine-ssh-sid.sock",
868            Some("dark"),
869            &BTreeMap::new(),
870            &[],
871        );
872        assert!(wrapped.contains("SHINE_TERMINAL_THEME='dark'"));
873        // Must appear inside the `env ...` prefix, before the `sh -c` handoff.
874        assert!(
875            wrapped.find("SHINE_TERMINAL_THEME").unwrap() < wrapped.find("sh -c").unwrap(),
876            "theme var must be part of the env prefix: {wrapped}"
877        );
878    }
879
880    #[test]
881    fn wrapped_command_theme_injection_round_trips_through_a_real_shell() {
882        // Same rationale as wrapped_command_round_trips_through_a_real_shell:
883        // verify the quoting composition by actually running it, rather than
884        // reasoning about escaping by hand.
885        let wrapped = build_wrapped_remote_command(
886            "sid",
887            "tok",
888            "/tmp/shine-ssh-mod-test-theme-sid.sock",
889            Some("dark"),
890            &BTreeMap::new(),
891            &["printenv".to_string(), "SHINE_TERMINAL_THEME".to_string()],
892        );
893        let output = std::process::Command::new("sh")
894            .arg("-c")
895            .arg(&wrapped)
896            .output()
897            .expect("failed to run sh");
898        assert!(output.status.success(), "stderr: {:?}", output.stderr);
899        assert_eq!(String::from_utf8_lossy(&output.stdout).trim_end(), "dark");
900    }
901
902    #[test]
903    fn wrapped_command_forwarded_env_round_trips_special_characters() {
904        let forwarded = BTreeMap::from([(
905            "REMOTE_VALUE".to_string(),
906            "space ' quote $dollar\nand newline".to_string(),
907        )]);
908        let wrapped = build_wrapped_remote_command(
909            "sid",
910            "tok",
911            "/tmp/shine-ssh-mod-test-env-sid.sock",
912            None,
913            &forwarded,
914            &["printenv".to_string(), "REMOTE_VALUE".to_string()],
915        );
916        let output = std::process::Command::new("sh")
917            .arg("-c")
918            .arg(&wrapped)
919            .output()
920            .expect("failed to run sh");
921        assert!(output.status.success(), "stderr: {:?}", output.stderr);
922        assert_eq!(
923            String::from_utf8_lossy(&output.stdout),
924            "space ' quote $dollar\nand newline\n"
925        );
926    }
927
928    #[test]
929    fn windows_wrapped_command_decodes_special_values_and_command_argv() {
930        let forwarded = BTreeMap::from([(
931            "REMOTE_VALUE".to_string(),
932            "space ' quote \" $dollar & amp % percent ! bang\nand newline".to_string(),
933        )]);
934        let wrapped = build_windows_wrapped_remote_command(
935            "sid",
936            Some("dark"),
937            &forwarded,
938            &[
939                "C:\\Program Files\\tool.exe".to_string(),
940                "one ' $ & % !".to_string(),
941            ],
942        )
943        .unwrap();
944
945        assert!(wrapped.starts_with("powershell.exe -NoProfile -EncodedCommand "));
946        assert!(!wrapped.contains("REMOTE_VALUE"));
947        assert!(!wrapped.contains("$dollar"));
948        let bootstrap = decode_powershell_script(wrapped.rsplit_once(' ').unwrap().1);
949        let encoded = bootstrap
950            .split(" -EncodedCommand '")
951            .nth(1)
952            .unwrap()
953            .split('\'')
954            .next()
955            .unwrap();
956        let script = decode_powershell_script(encoded);
957        assert!(!bootstrap.contains("REMOTE_VALUE"));
958        assert!(!bootstrap.contains("$dollar"));
959        assert!(bootstrap.contains("Get-Command pwsh.exe"));
960        assert!(bootstrap.contains("else { 'powershell.exe' }"));
961        assert!(bootstrap.contains("& $shell -NoProfile -EncodedCommand"));
962        assert!(bootstrap.contains("exit $code"));
963        assert!(script.contains("$env:SHINE_SSH_SESSION = 'sid'"));
964        assert!(script.contains("$env:SHINE_TERMINAL_THEME = 'dark'"));
965        assert!(script.contains(
966            "$env:REMOTE_VALUE = 'space '' quote \" $dollar & amp % percent ! bang\nand newline'"
967        ));
968        assert!(script.contains("& 'C:\\Program Files\\tool.exe' 'one '' $ & % !'"));
969        assert!(script.contains("exit $LASTEXITCODE"));
970    }
971
972    fn decode_powershell_script(encoded: &str) -> String {
973        let bytes = BASE64.decode(encoded).unwrap();
974        let units = bytes
975            .chunks_exact(2)
976            .map(|pair| u16::from_le_bytes([pair[0], pair[1]]))
977            .collect::<Vec<_>>();
978        String::from_utf16(&units).unwrap()
979    }
980
981    #[test]
982    fn windows_wrapped_command_uses_no_exit_for_interactive_session() {
983        let wrapped =
984            build_windows_wrapped_remote_command("sid", None, &BTreeMap::new(), &[]).unwrap();
985        let bootstrap = decode_powershell_script(wrapped.rsplit_once(' ').unwrap().1);
986        assert!(bootstrap.contains("& $shell -NoExit -EncodedCommand "));
987        assert!(!bootstrap.contains("& $shell -NoProfile"));
988    }
989
990    #[test]
991    fn windows_wrapped_command_selects_shell_before_running_payload() {
992        let wrapped = build_windows_wrapped_remote_command(
993            "sid",
994            None,
995            &BTreeMap::new(),
996            &["exit".to_string(), "7".to_string()],
997        )
998        .unwrap();
999
1000        assert!(wrapped.starts_with("powershell.exe -NoProfile -EncodedCommand "));
1001        let bootstrap = decode_powershell_script(wrapped.rsplit_once(' ').unwrap().1);
1002        assert!(bootstrap.contains("Get-Command pwsh.exe -CommandType Application"));
1003        assert!(bootstrap.contains("$pwsh.Source"));
1004        assert!(bootstrap.contains("else { 'powershell.exe' }"));
1005        assert_eq!(bootstrap.matches("& $shell").count(), 1);
1006        assert!(!wrapped.contains("&&"));
1007        assert!(!wrapped.contains("||"));
1008    }
1009
1010    #[test]
1011    fn windows_wrapped_command_rejects_nul_values() {
1012        let forwarded = BTreeMap::from([("REMOTE_VALUE".to_string(), "bad\0value".to_string())]);
1013        let error = build_windows_wrapped_remote_command("sid", None, &forwarded, &[]).unwrap_err();
1014        assert!(error.to_string().contains("NUL"));
1015    }
1016
1017    #[tokio::test]
1018    async fn forwarded_plain_env_uses_exact_key_and_alias() {
1019        let dir = std::env::temp_dir().join(format!("shine-ssh-env-{}", uuid::Uuid::new_v4()));
1020        let mut config = Config::new_for_test(&dir);
1021        config.env.insert("LOCAL_NAME".into(), "local value".into());
1022        config
1023            .env
1024            .insert("LOCAL_NAME_SECRET".into(), "encrypted value".into());
1025
1026        let resolved = resolve_forwarded_env(&config, &["LOCAL_NAME=REMOTE_NAME".to_string()], &[])
1027            .await
1028            .unwrap();
1029
1030        assert_eq!(
1031            resolved.get("REMOTE_NAME").map(String::as_str),
1032            Some("local value")
1033        );
1034    }
1035
1036    #[tokio::test]
1037    async fn forwarded_plain_env_requires_explicit_secret_opt_in() {
1038        let dir = std::env::temp_dir().join(format!("shine-ssh-env-{}", uuid::Uuid::new_v4()));
1039        let mut config = Config::new_for_test(&dir);
1040        config
1041            .env
1042            .insert("API_TOKEN_SECRET".into(), "ciphertext".into());
1043
1044        let error = resolve_forwarded_env(&config, &["API_TOKEN".to_string()], &[])
1045            .await
1046            .unwrap_err();
1047
1048        assert!(error.to_string().contains("use --with-secret API_TOKEN"));
1049    }
1050
1051    #[tokio::test]
1052    async fn forwarded_secret_requires_base_key_and_encrypted_storage() {
1053        let dir = std::env::temp_dir().join(format!("shine-ssh-env-{}", uuid::Uuid::new_v4()));
1054        let mut config = Config::new_for_test(&dir);
1055        config.env.insert("API_TOKEN".into(), "plaintext".into());
1056        config
1057            .env
1058            .insert("OTHER_SECRET".into(), "ciphertext".into());
1059
1060        let plaintext_only = resolve_forwarded_env(&config, &[], &["API_TOKEN".to_string()])
1061            .await
1062            .unwrap_err();
1063        assert!(
1064            plaintext_only
1065                .to_string()
1066                .contains("API_TOKEN_SECRET is not set")
1067        );
1068
1069        let suffixed = resolve_forwarded_env(&config, &[], &["OTHER_SECRET".to_string()])
1070            .await
1071            .unwrap_err();
1072        assert!(
1073            suffixed
1074                .to_string()
1075                .contains("expects a base key without the _SECRET suffix")
1076        );
1077    }
1078
1079    #[tokio::test]
1080    async fn forwarded_env_rejects_duplicate_and_reserved_targets() {
1081        let dir = std::env::temp_dir().join(format!("shine-ssh-env-{}", uuid::Uuid::new_v4()));
1082        let mut config = Config::new_for_test(&dir);
1083        config.env.insert("ONE".into(), "1".into());
1084        config.env.insert("TWO".into(), "2".into());
1085
1086        let duplicate = resolve_forwarded_env(
1087            &config,
1088            &["ONE=REMOTE".to_string(), "TWO=REMOTE".to_string()],
1089            &[],
1090        )
1091        .await
1092        .unwrap_err();
1093        assert!(
1094            duplicate
1095                .to_string()
1096                .contains("duplicate target variable: REMOTE")
1097        );
1098
1099        let reserved = resolve_forwarded_env(&config, &["ONE=SHINE_SSH_TOKEN".to_string()], &[])
1100            .await
1101            .unwrap_err();
1102        assert!(
1103            reserved
1104                .to_string()
1105                .contains("cannot override shine-managed SSH variable SHINE_SSH_TOKEN")
1106        );
1107    }
1108}