inferlab_runtime/
ssh.rs

1//! The one authority for the SSH client invocation shape, shared by every
2//! module that reaches a remote machine.
3
4use crate::shell::shell_quote;
5use crate::{
6    container::{BoundedError, BoundedWait, CommandCleanupEvidence, run_with_bound},
7    operation_bound::OperationBound,
8};
9use std::process::Output;
10use thiserror::Error;
11
12const SSH_OPTIONS: &[&str] = &["-o", "BatchMode=yes", "--"];
13
14/// Dynamic-linker overrides stripped from the SSH client environment: the
15/// client is a host tool, not a serving-stack member, and a workspace Pixi
16/// or stack activation must not leak its library search path into it
17/// ([[RFC-0003:C-RUNTIME-WORKFLOWS]]).
18pub const SSH_ENV_REMOVE: &[&str] = &["LD_LIBRARY_PATH", "LD_PRELOAD"];
19
20/// The full SSH argv for bounded execution through an owning operation or
21/// cleanup deadline.
22pub fn ssh_argv(target: &str, script: &str) -> Vec<String> {
23    // Keep the remote account's interactive login initialization, which owns
24    // tool discovery and declared pass-through values, but replace that shell
25    // before the InferLab script runs. The non-login replacement preserves the
26    // initialized environment without running `.bash_logout`, whose exit status
27    // must not replace the remote operation's result.
28    let remote_command = format!("exec \"$BASH\" -c {}", shell_quote(script));
29    let mut argv: Vec<String> = ["ssh"]
30        .into_iter()
31        .map(str::to_owned)
32        .chain(SSH_OPTIONS.iter().map(|option| (*option).to_owned()))
33        .collect();
34    argv.extend([
35        target.to_owned(),
36        "bash".to_owned(),
37        "-lic".to_owned(),
38        shell_quote(&remote_command),
39    ]);
40    argv
41}
42
43#[derive(Debug, Error)]
44pub enum SshError {
45    #[error("failed to launch SSH for {target:?}: {source}")]
46    Launch {
47        target: String,
48        #[source]
49        source: std::io::Error,
50    },
51    #[error("failed to {operation} for SSH target {target:?}: {source}")]
52    Io {
53        target: String,
54        operation: &'static str,
55        #[source]
56        source: std::io::Error,
57    },
58    #[error(
59        "failed while waiting for SSH target {target:?} after {operation_elapsed_ms} ms: {source}; child cleanup: {cleanup:?}"
60    )]
61    WaitCleanup {
62        target: String,
63        operation_elapsed_ms: u64,
64        #[source]
65        source: std::io::Error,
66        cleanup: Box<CommandCleanupEvidence>,
67    },
68    #[error(
69        "SSH for target {target:?} was interrupted after {operation_elapsed_ms} ms; child cleanup: {cleanup:?}"
70    )]
71    Interrupted {
72        target: String,
73        operation_elapsed_ms: u64,
74        cleanup: Box<CommandCleanupEvidence>,
75    },
76    #[error(
77        "failed to clean up interrupted SSH for target {target:?} after {operation_elapsed_ms} ms: {source}; child cleanup: {cleanup:?}"
78    )]
79    InterruptCleanup {
80        target: String,
81        operation_elapsed_ms: u64,
82        #[source]
83        source: std::io::Error,
84        cleanup: Box<CommandCleanupEvidence>,
85    },
86    #[error(
87        "SSH supervisor for target {target:?} unexpectedly exhausted an unbounded operation after {operation_elapsed_ms} ms; child cleanup: {cleanup:?}"
88    )]
89    UnexpectedDeadline {
90        target: String,
91        operation_elapsed_ms: u64,
92        cleanup: Option<Box<CommandCleanupEvidence>>,
93    },
94}
95
96pub fn ssh_output(target: &str, script: &str) -> Result<Output, SshError> {
97    run_ssh(target, script, None)
98}
99
100pub fn ssh_output_with_input(target: &str, script: &str, input: &[u8]) -> Result<Output, SshError> {
101    run_ssh(target, script, Some(input))
102}
103
104fn run_ssh(target: &str, script: &str, input: Option<&[u8]>) -> Result<Output, SshError> {
105    let argv = ssh_argv(target, script);
106    match run_with_bound(
107        &argv,
108        SSH_ENV_REMOVE,
109        None,
110        input,
111        &OperationBound::unbounded(),
112        None,
113    ) {
114        Ok(BoundedWait::Exited {
115            status,
116            stdout,
117            stderr,
118        }) => Ok(Output {
119            status,
120            stdout,
121            stderr,
122        }),
123        Ok(BoundedWait::Expired {
124            kill,
125            operation_elapsed_ms,
126            cleanup,
127        }) => {
128            kill.map_err(|source| SshError::Io {
129                target: target.to_owned(),
130                operation: "clean up SSH after unexpected deadline",
131                source,
132            })?;
133            Err(SshError::UnexpectedDeadline {
134                target: target.to_owned(),
135                operation_elapsed_ms,
136                cleanup: cleanup.map(Box::new),
137            })
138        }
139        Ok(BoundedWait::Interrupted {
140            kill,
141            operation_elapsed_ms,
142            cleanup,
143        }) => match kill {
144            Ok(()) => Err(SshError::Interrupted {
145                target: target.to_owned(),
146                operation_elapsed_ms,
147                cleanup: Box::new(cleanup),
148            }),
149            Err(source) => Err(SshError::InterruptCleanup {
150                target: target.to_owned(),
151                operation_elapsed_ms,
152                source,
153                cleanup: Box::new(cleanup),
154            }),
155        },
156        Err(BoundedError::Launch(source)) => Err(SshError::Launch {
157            target: target.to_owned(),
158            source,
159        }),
160        Err(BoundedError::Stdin(source)) => Err(SshError::Io {
161            target: target.to_owned(),
162            operation: "write SSH stdin",
163            source,
164        }),
165        Err(BoundedError::Wait(source)) => Err(SshError::Io {
166            target: target.to_owned(),
167            operation: "wait for SSH",
168            source,
169        }),
170        Err(BoundedError::WaitCleanup {
171            source,
172            operation_elapsed_ms,
173            cleanup,
174        }) => Err(SshError::WaitCleanup {
175            target: target.to_owned(),
176            operation_elapsed_ms,
177            source,
178            cleanup: Box::new(cleanup),
179        }),
180    }
181}
182
183#[cfg(test)]
184mod tests {
185    use super::ssh_argv;
186    use std::error::Error;
187    use std::fs;
188    use std::io;
189    use std::path::Path;
190    use std::process::{Command, Output};
191
192    #[test]
193    fn ssh_argv_leaves_connection_policy_to_openssh_configuration() {
194        let argv = ssh_argv("fixture-target", "exit 0");
195
196        assert_eq!(
197            &argv[..5],
198            ["ssh", "-o", "BatchMode=yes", "--", "fixture-target"]
199        );
200        assert!(!argv.iter().any(|argument| {
201            argument.starts_with("ConnectTimeout=")
202                || argument.starts_with("ServerAliveInterval=")
203                || argument.starts_with("ServerAliveCountMax=")
204        }));
205    }
206
207    fn run_remote_command(home: &Path, script: &str) -> Result<Output, Box<dyn Error>> {
208        let argv = ssh_argv("fixture-target", script);
209        let separator = argv
210            .iter()
211            .position(|argument| argument == "--")
212            .ok_or_else(|| io::Error::other("SSH argv has no option separator"))?;
213        let remote_command = argv
214            .get(separator + 2..)
215            .ok_or_else(|| io::Error::other("SSH argv has no remote command"))?
216            .join(" ");
217        Ok(Command::new("sh")
218            .args(["-c", &remote_command])
219            .env("HOME", home)
220            .output()?)
221    }
222
223    #[test]
224    fn login_shell_teardown_does_not_override_remote_script_result() -> Result<(), Box<dyn Error>> {
225        let home = tempfile::tempdir()?;
226        fs::write(
227            home.path().join(".bash_profile"),
228            "export INFERLAB_LOGIN_MARKER=loaded\n",
229        )?;
230        fs::write(
231            home.path().join(".bash_logout"),
232            "printf 'logout-ran\\n' >&2\nfalse\n",
233        )?;
234
235        let alive = run_remote_command(
236            home.path(),
237            "set -eu; printf '%s\\n' \"$INFERLAB_LOGIN_MARKER\"; printf 'probe-stderr\\n' >&2; exit 0",
238        )?;
239        assert!(alive.status.success());
240        assert_eq!(String::from_utf8(alive.stdout)?, "loaded\n");
241        let alive_stderr = String::from_utf8(alive.stderr)?;
242        assert!(alive_stderr.contains("probe-stderr"));
243        assert!(!alive_stderr.contains("logout-ran"));
244
245        let dead = run_remote_command(home.path(), "set -eu; printf 'dead\\n'; exit 3")?;
246        assert_eq!(dead.status.code(), Some(3));
247        assert_eq!(String::from_utf8(dead.stdout)?, "dead\n");
248        Ok(())
249    }
250}
inferlab_runtime/ssh.rs

inferlab_runtime/
ssh.rs
