pray_core/
ssh_client.rs

1use crate::client_trust::{effective_trust_home, gate_pray_ssh_host, gate_pray_ssh_publisher};
2use crate::ssh_identity::active_ssh_user_fingerprint;
3use crate::ssh_rpc::{call_stdio, RpcRequest, RpcResponse, SSH_RPC_SPEC};
4use crate::{PrayError, PrayResult};
5use serde_json::Value;
6use std::io::{BufReader, Write};
7use std::path::{Path, PathBuf};
8use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
9use std::sync::atomic::{AtomicU64, Ordering};
10
11static REQUEST_COUNTER: AtomicU64 = AtomicU64::new(1);
12
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct PraySshTarget {
15    pub user: Option<String>,
16    pub host: String,
17    pub port: u16,
18    pub root: Option<PathBuf>,
19}
20
21pub fn is_pray_ssh_url(url: &str) -> bool {
22    url.starts_with("pray+ssh://") || url.starts_with("ssh+pray://")
23}
24
25pub fn parse_pray_ssh_url(url: &str) -> PrayResult<PraySshTarget> {
26    let remainder = url
27        .strip_prefix("pray+ssh://")
28        .or_else(|| url.strip_prefix("ssh+pray://"))
29        .ok_or_else(|| PrayError::Parse {
30            kind: "pray ssh url",
31            message: format!("expected pray+ssh:// url, got {url}"),
32        })?;
33
34    let (authority, root) = match remainder.split_once('/') {
35        Some((authority, path)) if !path.is_empty() => {
36            (authority, Some(PathBuf::from(format!("/{path}"))))
37        }
38        _ => (remainder, None),
39    };
40
41    let (credentials, host_port) = match authority.rsplit_once('@') {
42        Some((user, host_port)) => (Some(user.to_string()), host_port),
43        None => (None, authority),
44    };
45
46    let (host, port) = match host_port.rsplit_once(':') {
47        Some((host, port_text)) if !host.contains(']') => {
48            let port = port_text.parse::<u16>().map_err(|error| PrayError::Parse {
49                kind: "pray ssh url",
50                message: format!("invalid port in {url}: {error}"),
51            })?;
52            (host.to_string(), port)
53        }
54        _ => (host_port.to_string(), 22),
55    };
56
57    if host.is_empty() {
58        return Err(PrayError::Parse {
59            kind: "pray ssh url",
60            message: format!("missing host in {url}"),
61        });
62    }
63
64    Ok(PraySshTarget {
65        user: credentials,
66        host,
67        port,
68        root,
69    })
70}
71
72pub struct SshRpcSession {
73    child: Child,
74    stdin: Option<ChildStdin>,
75    reader: BufReader<ChildStdout>,
76}
77
78impl SshRpcSession {
79    pub fn connect(target: &PraySshTarget) -> PrayResult<Self> {
80        if target.host == "stdio-host" {
81            if let Ok(root) = std::env::var("PRAY_TEST_SSH_STDIO_ROOT") {
82                let mut command = Command::new(pray_program());
83                command.arg("serve").arg("--stdio").arg("--root").arg(root);
84                return Self::connect_stdio(command);
85            }
86        }
87
88        let mut command = Command::new(ssh_program());
89        command
90            .arg("-p")
91            .arg(target.port.to_string())
92            .arg("-o")
93            .arg("BatchMode=yes")
94            .arg("-o")
95            .arg("StrictHostKeyChecking=accept-new");
96        if let Some(user) = &target.user {
97            command.arg(format!("{user}@{}", target.host));
98        } else {
99            command.arg(&target.host);
100        }
101        let mut remote_command = String::from("pray serve --stdio");
102        if let Some(root) = &target.root {
103            remote_command.push_str(" --root ");
104            remote_command.push_str(&shell_escape(root.to_string_lossy().as_ref()));
105        }
106        command.arg(remote_command);
107        command.stdin(Stdio::piped());
108        command.stdout(Stdio::piped());
109        command.stderr(Stdio::piped());
110        let mut child = command.spawn().map_err(|error| {
111            PrayError::Unsupported(format!("failed to start ssh for pray rpc: {error}"))
112        })?;
113        let stdin = child
114            .stdin
115            .take()
116            .ok_or_else(|| PrayError::Unsupported("rpc stdin unavailable".to_string()))?;
117        let stdout = child
118            .stdout
119            .take()
120            .ok_or_else(|| PrayError::Unsupported("rpc stdout unavailable".to_string()))?;
121        Ok(Self {
122            child,
123            stdin: Some(stdin),
124            reader: BufReader::new(stdout),
125        })
126    }
127
128    pub fn connect_stdio(mut command: Command) -> PrayResult<Self> {
129        command.stdin(Stdio::piped());
130        command.stdout(Stdio::piped());
131        command.stderr(Stdio::piped());
132        let mut child = command.spawn().map_err(|error| {
133            PrayError::Unsupported(format!("failed to start stdio rpc: {error}"))
134        })?;
135        let stdin = child
136            .stdin
137            .take()
138            .ok_or_else(|| PrayError::Unsupported("rpc stdin unavailable".to_string()))?;
139        let stdout = child
140            .stdout
141            .take()
142            .ok_or_else(|| PrayError::Unsupported("rpc stdout unavailable".to_string()))?;
143        Ok(Self {
144            child,
145            stdin: Some(stdin),
146            reader: BufReader::new(stdout),
147        })
148    }
149
150    pub fn call(&mut self, method: &str, params: Value) -> PrayResult<RpcResponse> {
151        let request_id = REQUEST_COUNTER.fetch_add(1, Ordering::Relaxed).to_string();
152        let request = RpcRequest::new(request_id, method, params);
153        let stdin = self
154            .stdin
155            .as_mut()
156            .ok_or_else(|| PrayError::Unsupported("rpc stdin unavailable".to_string()))?;
157        let response = call_stdio(&mut self.reader, stdin, &request)?;
158        if response.spec != SSH_RPC_SPEC {
159            return Err(PrayError::Resolution(format!(
160                "unexpected rpc spec in response: {}",
161                response.spec
162            )));
163        }
164        Ok(response)
165    }
166
167    pub fn call_json(&mut self, method: &str, params: Value) -> PrayResult<Value> {
168        let response = self.call(method, params)?;
169        if response.status / 100 != 2 {
170            let message = response
171                .body
172                .get("error")
173                .and_then(Value::as_str)
174                .unwrap_or("rpc request failed");
175            return Err(PrayError::Resolution(format!(
176                "rpc {method} failed with status {}: {message}",
177                response.status
178            )));
179        }
180        Ok(response.body)
181    }
182
183    pub fn call_bytes(&mut self, method: &str, params: Value) -> PrayResult<Vec<u8>> {
184        let response = self.call(method, params)?;
185        if response.status / 100 != 2 {
186            let message = response
187                .body
188                .get("error")
189                .and_then(Value::as_str)
190                .unwrap_or("rpc request failed");
191            return Err(PrayError::Resolution(format!(
192                "rpc {method} failed with status {}: {message}",
193                response.status
194            )));
195        }
196        response.decode_body_bytes()
197    }
198}
199
200impl Drop for SshRpcSession {
201    fn drop(&mut self) {
202        if let Some(mut stdin) = self.stdin.take() {
203            let _ = stdin.flush();
204        }
205        let _ = self.child.wait();
206    }
207}
208
209pub fn with_pray_ssh_session<T>(
210    source_url: &str,
211    operation: impl FnOnce(&mut SshRpcSession) -> PrayResult<T>,
212) -> PrayResult<T> {
213    let target = parse_pray_ssh_url(source_url)?;
214    let home = effective_trust_home()?;
215    let _host_key = gate_pray_ssh_host(&home, source_url, &target.host, target.port)?;
216    if let Some(publisher) = active_ssh_user_fingerprint() {
217        gate_pray_ssh_publisher(&home, source_url, &publisher)?;
218    }
219    let mut session = SshRpcSession::connect(&target)?;
220    operation(&mut session)
221}
222
223pub fn ssh_program() -> String {
224    [
225        "/usr/bin/ssh",
226        "/opt/homebrew/bin/ssh",
227        "/usr/local/bin/ssh",
228        "ssh",
229    ]
230    .into_iter()
231    .find(|candidate| *candidate == "ssh" || Path::new(candidate).exists())
232    .unwrap_or("ssh")
233    .to_string()
234}
235
236fn pray_program() -> String {
237    if let Ok(path) = std::env::var("PRAY_TEST_BINARY") {
238        return path;
239    }
240    std::env::current_exe()
241        .ok()
242        .and_then(|path| path.to_str().map(str::to_string))
243        .unwrap_or_else(|| "pray".to_string())
244}
245
246fn shell_escape(value: &str) -> String {
247    if value.chars().all(|character| {
248        character.is_ascii_alphanumeric() || matches!(character, '-' | '_' | '.' | '/')
249    }) {
250        value.to_string()
251    } else {
252        format!("'{}'", value.replace('\'', "'\"'\"'"))
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259
260    #[test]
261    fn parse_pray_ssh_url_reads_user_host_port_and_root() {
262        let target = parse_pray_ssh_url("pray+ssh://pray@prayers.internal:2222/var/lib/pray")
263            .expect("parse url");
264        assert_eq!(target.user.as_deref(), Some("pray"));
265        assert_eq!(target.host, "prayers.internal");
266        assert_eq!(target.port, 2222);
267        assert_eq!(target.root, Some(PathBuf::from("/var/lib/pray")));
268    }
269}
pray_core/ssh_client.rs

pray_core/
ssh_client.rs
