pray_core/
ssh_client.rs

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

pray_core/
ssh_client.rs
