agentmux_srv/backend/remote/
files.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Files on an SSH host, through AgentMux's helper (spec §6.3 of
5//! SPEC_REMOTE_TERMINALS_AND_DURABLE_SESSIONS_2026_10_02.md): one
6//! `ssh -T host -- agentmux-remote serve --stdio` per connection, kept open
7//! while it is used and closed after [`IDLE_CLOSE`] without a call, speaking
8//! `agentmux_remote::fsproto`. Requests carry ids, so any number can be in
9//! flight on one connection.
10//!
11//! A host without this version's helper gets it installed on first use
12//! (`helper_install::ensure`), as a durable pane does. ssh's prompts go to
13//! the user in the window of the pane that asked ([`connect`]'s `ask`).
14
15use std::collections::HashMap;
16use std::sync::{Arc, Mutex, OnceLock};
17use std::time::{Duration, Instant};
18
19use agentmux_remote::fsproto::{self, Entry, ErrKind, Reply, Request, Splitter};
20use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
21use tokio::sync::{mpsc, oneshot};
22
23use super::host::HostSsh;
24use super::sessions::AskIn;
25
26/// A connection nobody has used for this long is closed.
27pub const IDLE_CLOSE: Duration = Duration::from_secs(10 * 60);
28
29/// How long the handshake may take: ssh may first ask the user something
30/// (askpass, whose dialog waits up to two minutes).
31const HANDSHAKE: Duration = Duration::from_secs(150);
32
33/// How long one request may take once connected.
34const CALL: Duration = Duration::from_secs(120);
35
36/// Why a remote file operation failed.
37#[derive(Debug, Clone, PartialEq, Eq)]
38pub struct RemoteError {
39    /// What the host said went wrong; `None` when the connection itself
40    /// failed (it could not be made, it dropped, or a request timed out).
41    pub kind: Option<ErrKind>,
42    pub message: String,
43}
44
45impl RemoteError {
46    fn link(message: impl Into<String>) -> Self {
47        Self {
48            kind: None,
49            message: message.into(),
50        }
51    }
52
53    pub fn is_not_found(&self) -> bool {
54        self.kind == Some(ErrKind::NotFound)
55    }
56}
57
58impl std::fmt::Display for RemoteError {
59    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
60        f.write_str(&self.message)
61    }
62}
63
64struct Call {
65    req: Request,
66    reply: oneshot::Sender<Result<Reply, String>>,
67}
68
69/// One open connection to a host's helper.
70pub struct RemoteFiles {
71    /// The connection's canonical name.
72    pub conn: String,
73    /// The user's home directory there (`~`).
74    pub home: String,
75    /// `std::env::consts::OS` there.
76    pub os: String,
77    calls: mpsc::UnboundedSender<Call>,
78    last_used: Mutex<Instant>,
79}
80
81impl RemoteFiles {
82    /// Speak the protocol over `reader` and `writer` (an ssh's stdout and
83    /// stdin), handshake included. `keep` lives as long as the connection
84    /// (the ssh child and its askpass grant).
85    pub async fn over<R, W>(
86        conn: &str,
87        reader: R,
88        writer: W,
89        keep: Box<dyn std::any::Any + Send>,
90    ) -> Result<Arc<Self>, RemoteError>
91    where
92        R: AsyncRead + Unpin + Send + 'static,
93        W: AsyncWrite + Unpin + Send + 'static,
94    {
95        let (tx, rx) = mpsc::unbounded_channel();
96        tokio::spawn(pump(conn.to_string(), reader, writer, rx, keep));
97        let hello = call(&tx, Request::Hello, HANDSHAKE).await?;
98        let Reply::Hello { protocol, home, os } = hello else {
99            return Err(RemoteError::link(format!(
100                "the helper on {conn} did not greet"
101            )));
102        };
103        if protocol != fsproto::PROTOCOL {
104            return Err(RemoteError::link(format!(
105                "the helper on {conn} speaks file protocol {protocol}, this AgentMux {}",
106                fsproto::PROTOCOL
107            )));
108        }
109        Ok(Arc::new(Self {
110            conn: conn.to_string(),
111            home,
112            os,
113            calls: tx,
114            last_used: Mutex::new(Instant::now()),
115        }))
116    }
117
118    /// Whether the connection is still up.
119    pub fn is_open(&self) -> bool {
120        !self.calls.is_closed()
121    }
122
123    /// One request; a host's refusal is `Err` with its kind.
124    pub async fn call(&self, req: Request) -> Result<Reply, RemoteError> {
125        *self.last_used.lock().unwrap() = Instant::now();
126        match call(&self.calls, req, CALL).await? {
127            Reply::Err { kind, message } => Err(RemoteError {
128                kind: Some(kind),
129                message,
130            }),
131            reply => Ok(reply),
132        }
133    }
134
135    pub async fn stat(&self, path: &str) -> Result<Entry, RemoteError> {
136        match self.call(Request::Stat { path: path.into() }).await? {
137            Reply::Stat(e) => Ok(e),
138            r => Err(unexpected(&r)),
139        }
140    }
141
142    /// Every entry of a directory, a page at a time.
143    pub async fn list(&self, path: &str) -> Result<Vec<Entry>, RemoteError> {
144        let mut all = Vec::new();
145        loop {
146            let req = Request::List {
147                path: path.into(),
148                offset: all.len() as u32,
149                limit: 1000,
150            };
151            match self.call(req).await? {
152                Reply::List { entries, total } => {
153                    let done = entries.is_empty() || all.len() + entries.len() >= total as usize;
154                    all.extend(entries);
155                    if done {
156                        return Ok(all);
157                    }
158                }
159                r => return Err(unexpected(&r)),
160            }
161        }
162    }
163
164    /// A whole file of at most `max` bytes, read in ranges.
165    pub async fn read(&self, path: &str, max: u64) -> Result<Vec<u8>, RemoteError> {
166        let size = self.stat(path).await?.size;
167        if size > max {
168            return Err(RemoteError {
169                kind: Some(ErrKind::TooLarge),
170                message: format!("{path} is {size} bytes, over the {max} this can open"),
171            });
172        }
173        let mut data = Vec::with_capacity(size as usize);
174        loop {
175            let req = Request::Read {
176                path: path.into(),
177                offset: data.len() as u64,
178                len: fsproto::MAX_READ,
179            };
180            match self.call(req).await? {
181                Reply::Read { data: part, eof } => {
182                    data.extend_from_slice(&part);
183                    if data.len() as u64 > max {
184                        return Err(RemoteError {
185                            kind: Some(ErrKind::TooLarge),
186                            message: format!("{path} grew past {max} bytes while it was read"),
187                        });
188                    }
189                    if eof || part.is_empty() {
190                        return Ok(data);
191                    }
192                }
193                r => return Err(unexpected(&r)),
194            }
195        }
196    }
197
198    /// Replace a file with `data`, atomically on the host.
199    pub async fn write(&self, path: &str, data: Vec<u8>) -> Result<(), RemoteError> {
200        if data.len() > fsproto::MAX_WRITE {
201            return Err(RemoteError {
202                kind: Some(ErrKind::TooLarge),
203                message: format!(
204                    "over the {} MB a remote write may be",
205                    fsproto::MAX_WRITE >> 20
206                ),
207            });
208        }
209        self.done(Request::Write {
210            path: path.into(),
211            data,
212        })
213        .await
214    }
215
216    /// Add `data` to the end of an existing file (a big upload, in pieces).
217    pub async fn append(&self, path: &str, data: Vec<u8>) -> Result<(), RemoteError> {
218        self.done(Request::Append {
219            path: path.into(),
220            data,
221        })
222        .await
223    }
224
225    /// One range of a file: up to `len` bytes from `offset`, and whether the
226    /// file ends there.
227    pub async fn read_range(
228        &self,
229        path: &str,
230        offset: u64,
231        len: u32,
232    ) -> Result<(Vec<u8>, bool), RemoteError> {
233        match self
234            .call(Request::Read {
235                path: path.into(),
236                offset,
237                len: len.min(fsproto::MAX_READ),
238            })
239            .await?
240        {
241            Reply::Read { data, eof } => Ok((data, eof)),
242            r => Err(unexpected(&r)),
243        }
244    }
245
246    /// Rename the file `from` over the file `to` in one step.
247    pub async fn replace(&self, from: &str, to: &str) -> Result<(), RemoteError> {
248        self.done(Request::Replace {
249            from: from.into(),
250            to: to.into(),
251        })
252        .await
253    }
254
255    /// `path` with its links and `..` resolved.
256    pub async fn realpath(&self, path: &str) -> Result<String, RemoteError> {
257        match self.call(Request::Realpath { path: path.into() }).await? {
258            Reply::Path(p) => Ok(p),
259            r => Err(unexpected(&r)),
260        }
261    }
262
263    /// A file's permission bits (`mode`, 0 to leave) and modification time
264    /// (`mtime_ms`, 0 to leave).
265    pub async fn set_meta(&self, path: &str, mode: u32, mtime_ms: i64) -> Result<(), RemoteError> {
266        self.done(Request::SetMeta {
267            path: path.into(),
268            mode,
269            mtime_ms,
270        })
271        .await
272    }
273
274    /// Flush a file's contents to disk.
275    pub async fn sync(&self, path: &str) -> Result<(), RemoteError> {
276        self.done(Request::Sync { path: path.into() }).await
277    }
278
279    /// One pane's processes on the host (`PaneProcs`): durable `session`'s
280    /// shell, else those marked with `block_id`.
281    pub async fn pane_procs(&self, block_id: &str, session: &str) -> Result<fsproto::PaneFrame, RemoteError> {
282        match self
283            .call(Request::PaneProcs {
284                block_id: block_id.to_string(),
285                session: session.to_string(),
286            })
287            .await?
288        {
289            Reply::PaneProcs(f) => Ok(f),
290            r => Err(unexpected(&r)),
291        }
292    }
293
294    /// The host's processes for Tower: the `top` busiest and largest among
295    /// those matching `filter`, with totals (`agentmux_remote::procs`).
296    pub async fn procs(&self, top: u32, filter: &str) -> Result<fsproto::ProcFrame, RemoteError> {
297        match self
298            .call(Request::Procs {
299                top,
300                filter: filter.to_string(),
301            })
302            .await?
303        {
304            Reply::Procs(f) => Ok(f),
305            r => Err(unexpected(&r)),
306        }
307    }
308
309    pub async fn mkdir(&self, path: &str, parents: bool) -> Result<(), RemoteError> {
310        self.done(Request::Mkdir {
311            path: path.into(),
312            parents,
313        })
314        .await
315    }
316
317    pub async fn rename(&self, from: &str, to: &str) -> Result<(), RemoteError> {
318        self.done(Request::Rename {
319            from: from.into(),
320            to: to.into(),
321        })
322        .await
323    }
324
325    pub async fn delete(&self, path: &str, recursive: bool) -> Result<(), RemoteError> {
326        self.done(Request::Delete {
327            path: path.into(),
328            recursive,
329        })
330        .await
331    }
332
333    async fn done(&self, req: Request) -> Result<(), RemoteError> {
334        match self.call(req).await? {
335            Reply::Done => Ok(()),
336            r => Err(unexpected(&r)),
337        }
338    }
339
340    fn idle_for(&self) -> Duration {
341        self.last_used.lock().unwrap().elapsed()
342    }
343}
344
345fn unexpected(r: &Reply) -> RemoteError {
346    RemoteError::link(format!("the helper answered out of turn: {r:?}"))
347}
348
349async fn call(
350    calls: &mpsc::UnboundedSender<Call>,
351    req: Request,
352    limit: Duration,
353) -> Result<Reply, RemoteError> {
354    let (reply, rx) = oneshot::channel();
355    calls
356        .send(Call { req, reply })
357        .map_err(|_| RemoteError::link("the connection to the host is closed"))?;
358    match tokio::time::timeout(limit, rx).await {
359        Ok(Ok(Ok(reply))) => Ok(reply),
360        Ok(Ok(Err(e))) => Err(RemoteError::link(e)),
361        Ok(Err(_)) => Err(RemoteError::link("the connection to the host closed")),
362        Err(_) => Err(RemoteError::link("the host took too long to answer")),
363    }
364}
365
366/// The connection's one task: requests out, replies back to whoever asked,
367/// until every handle is gone or the stream ends; then every call still
368/// waiting hears why, and `keep` (the ssh) is dropped.
369async fn pump<R, W>(
370    conn: String,
371    mut reader: R,
372    writer: W,
373    mut calls: mpsc::UnboundedReceiver<Call>,
374    keep: Box<dyn std::any::Any + Send>,
375) where
376    R: AsyncRead + Unpin,
377    W: AsyncWrite + Unpin + Send + 'static,
378{
379    // Writes on their own task: replies are read while a big request (up to
380    // 32 MB) is still going out. The helper answers one request at a time and
381    // stops reading while it writes a reply, so a writer that waited here,
382    // not reading, would leave both pipes full and both ends stuck.
383    let (out_tx, out_rx) = mpsc::unbounded_channel::<Vec<u8>>();
384    let (failed_tx, mut failed_rx) = oneshot::channel::<String>();
385    tokio::spawn(write_out(conn.clone(), writer, out_rx, failed_tx));
386    let mut waiting: HashMap<u32, oneshot::Sender<Result<Reply, String>>> = HashMap::new();
387    let mut next_id = 0u32;
388    let mut splitter = Splitter::new();
389    let mut buf = vec![0u8; 64 * 1024];
390    let why = loop {
391        tokio::select! {
392            c = calls.recv() => {
393                let Some(c) = c else { break "closed".to_string() };
394                next_id = next_id.wrapping_add(1);
395                let id = next_id;
396                if out_tx.send(c.req.encode(id)).is_err() {
397                    let _ = c.reply.send(Err(format!("the connection to {conn} closed")));
398                    break format!("the connection to {conn} closed");
399                }
400                waiting.insert(id, c.reply);
401            }
402            failed = &mut failed_rx => {
403                break failed.unwrap_or_else(|_| format!("the connection to {conn} closed"));
404            }
405            n = reader.read(&mut buf) => {
406                let n = match n {
407                    Ok(0) => break format!("the connection to {conn} closed"),
408                    Ok(n) => n,
409                    Err(e) => break format!("the connection to {conn} dropped ({e})"),
410                };
411                let bodies = match splitter.push(&buf[..n]) {
412                    Ok(b) => b,
413                    Err(e) => break format!("the helper on {conn} sent something unreadable ({e})"),
414                };
415                for body in bodies {
416                    match Reply::decode(&body) {
417                        Ok((id, reply)) => {
418                            if let Some(w) = waiting.remove(&id) {
419                                let _ = w.send(Ok(reply));
420                            }
421                        }
422                        Err(e) => {
423                            tracing::warn!(conn = %conn, error = %e, "remote files: unreadable reply");
424                        }
425                    }
426                }
427            }
428        }
429    };
430    for (_, w) in waiting.drain() {
431        let _ = w.send(Err(why.clone()));
432    }
433    // Calls that arrived after the end hear it too.
434    calls.close();
435    while let Ok(c) = calls.try_recv() {
436        let _ = c.reply.send(Err(why.clone()));
437    }
438    // Ends the writer (its queue closes), then the ssh.
439    drop(out_tx);
440    drop(keep);
441}
442
443/// The connection's writer: each request's bytes, in order. A failed write
444/// is reported once, and ends the connection.
445async fn write_out<W: AsyncWrite + Unpin>(
446    conn: String,
447    mut writer: W,
448    mut out: mpsc::UnboundedReceiver<Vec<u8>>,
449    failed: oneshot::Sender<String>,
450) {
451    while let Some(bytes) = out.recv().await {
452        let result = match writer.write_all(&bytes).await {
453            Ok(()) => writer.flush().await,
454            Err(e) => Err(e),
455        };
456        if let Err(e) = result {
457            let _ = failed.send(format!("the connection to {conn} dropped ({e})"));
458            return;
459        }
460    }
461}
462
463/// One host's slot in the pool: a lock so two first calls make one
464/// connection, around the connection once made.
465type Slot = Arc<tokio::sync::Mutex<Option<Arc<RemoteFiles>>>>;
466
467fn pool() -> &'static Mutex<HashMap<String, Slot>> {
468    static POOL: OnceLock<Mutex<HashMap<String, Slot>>> = OnceLock::new();
469    POOL.get_or_init(|| {
470        // Close what nobody uses (once, for the whole pool).
471        tokio::spawn(async {
472            loop {
473                tokio::time::sleep(Duration::from_secs(60)).await;
474                close_idle(IDLE_CLOSE).await;
475            }
476        });
477        Mutex::new(HashMap::new())
478    })
479}
480
481/// Drop every connection idle for `after` or more (each ssh then ends).
482pub async fn close_idle(after: Duration) {
483    let slots: Vec<Slot> = pool().lock().unwrap().values().cloned().collect();
484    for slot in slots {
485        let mut s = slot.lock().await;
486        if s.as_ref()
487            .is_some_and(|f| !f.is_open() || f.idle_for() >= after)
488        {
489            *s = None;
490        }
491    }
492}
493
494/// Drop the connection to `conn`'s helper, if there is one (its ssh ends):
495/// the helper was removed from the host.
496pub async fn close(conn: &str) {
497    let Ok(name) = canonical(conn) else { return };
498    let slot = pool().lock().unwrap().get(&name).cloned();
499    if let Some(slot) = slot {
500        *slot.lock().await = None;
501    }
502}
503
504/// The open connection to `conn`'s helper, made (and the helper installed)
505/// if there is none. `ask`: where ssh's prompts go, if any are needed.
506pub async fn connect(conn: &str, ask: Option<AskIn<'_>>) -> Result<Arc<RemoteFiles>, RemoteError> {
507    let name = canonical(conn)?;
508    let slot = pool()
509        .lock()
510        .unwrap()
511        .entry(name.clone())
512        .or_default()
513        .clone();
514    let mut s = slot.lock().await;
515    if let Some(f) = s.as_ref().filter(|f| f.is_open()) {
516        return Ok(f.clone());
517    }
518    let mut host = HostSsh::for_connection(conn).map_err(RemoteError::link)?;
519    let files = match start(&name, host.clone(), ask).await {
520        Ok(f) => f,
521        // The helper is not there (or not this version): install it, once,
522        // with ssh's prompts going to the user as for the connection itself
523        // (the grant lives until the install is done).
524        Err((e, true)) => {
525            super::helper_consent::allow_install(&name, ask.map(|a| a.block_id), super::helper_consent::Need::Files)
526                .await
527                .map_err(|why| {
528                    RemoteError::link(format!(
529                        "File browsing on {name} needs AgentMux's helper, which is not installed there: {why}."
530                    ))
531                })?;
532            let _install_grant = ask.and_then(|a| host.ask_user_in(a.block_id, &name, a.auth_key));
533            super::helper_install::ensure(&host, &name, "files", |_| async {})
534                .await
535                .map_err(|i| {
536                    RemoteError::link(format!(
537                        "{} (and installing the helper failed: {i})",
538                        e.message
539                    ))
540                })?;
541            start(&name, host, ask).await.map_err(|(e, _)| e)?
542        }
543        Err((e, false)) => return Err(e),
544    };
545    // The helper answered: the host has it (spec §7.1; Remotes shows it), and
546    // the connection is up. A Files pane's overlay showed "Disconnected" over
547    // a listing that had loaded (#4513), since nothing said so.
548    super::helper_hosts::remember(&name);
549    super::link::used(&name);
550    *s = Some(files.clone());
551    Ok(files)
552}
553
554/// [`connect`] for something that asks every few seconds on its own (a
555/// pane's stats badge), not on the user's word: the open connection, or a new
556/// one only on a host the helper is known to be on (`helper_hosts`), never an
557/// install, a consent question or an ssh prompt (a login that needs one
558/// fails instead; a pane's shared connection usually makes it unneeded).
559pub async fn connect_quiet(conn: &str) -> Result<Arc<RemoteFiles>, RemoteError> {
560    let name = canonical(conn)?;
561    let slot = pool().lock().unwrap().entry(name.clone()).or_default().clone();
562    let mut s = slot.lock().await;
563    if let Some(f) = s.as_ref().filter(|f| f.is_open()) {
564        return Ok(f.clone());
565    }
566    // A host this failed on lately (a helper of another version, a login
567    // that needs a prompt) isn't tried again on every poll.
568    if let Some(why) = quiet_failures().lock().unwrap().get(&name).filter(|(at, _)| at.elapsed() < QUIET_RETRY_AFTER) {
569        return Err(RemoteError::link(why.1.clone()));
570    }
571    if !super::helper_hosts::known(conn) {
572        return Err(RemoteError::link(format!("AgentMux's helper isn't on {name}")));
573    }
574    let host = HostSsh::for_connection(conn).map_err(RemoteError::link)?;
575    // Its own deadline, inside: a caller's timeout would drop this before a
576    // failure is recorded, and an unreachable host would be tried every poll.
577    let started = tokio::time::timeout(QUIET_START_TIMEOUT, start(&name, host, None)).await;
578    let files = match started {
579        Ok(Ok(f)) => f,
580        Ok(Err((e, _))) => {
581            quiet_failures().lock().unwrap().insert(name, (Instant::now(), e.message.clone()));
582            return Err(e);
583        }
584        Err(_) => {
585            let why = format!("{name} didn't answer in time");
586            quiet_failures().lock().unwrap().insert(name, (Instant::now(), why.clone()));
587            return Err(RemoteError::link(why));
588        }
589    };
590    quiet_failures().lock().unwrap().remove(&name);
591    *s = Some(files.clone());
592    Ok(files)
593}
594
595/// How long [`connect_quiet`] leaves a host it failed on alone.
596const QUIET_RETRY_AFTER: Duration = Duration::from_secs(60);
597
598/// How long [`connect_quiet`] waits for the helper to start: under the
599/// stats badge's own 5 s, so a host that doesn't answer is recorded as a
600/// failure rather than cut off and tried again on the next poll.
601const QUIET_START_TIMEOUT: Duration = Duration::from_secs(4);
602
603/// Hosts [`connect_quiet`] failed on: when, and why (said again meanwhile).
604fn quiet_failures() -> &'static Mutex<HashMap<String, (Instant, String)>> {
605    static FAILURES: OnceLock<Mutex<HashMap<String, (Instant, String)>>> = OnceLock::new();
606    FAILURES.get_or_init(Default::default)
607}
608
609/// The pool's key for `conn`: its canonical name, one per host however it is
610/// spelled. `Err` unless it is an SSH connection.
611fn canonical(conn: &str) -> Result<String, RemoteError> {
612    match super::ConnTarget::parse(conn).map_err(RemoteError::link)? {
613        t @ super::ConnTarget::Ssh(_) => Ok(t.name()),
614        _ => Err(RemoteError::link(format!(
615            "{conn:?} is not an SSH connection"
616        ))),
617    }
618}
619
620/// Tests: a connection to `conn` that is the real helper's request handling
621/// in this process (the files under `home`), put in the pool so [`connect`]
622/// returns it as it would a connection over ssh.
623#[cfg(test)]
624pub(crate) async fn connect_in_process(conn: &str, home: std::path::PathBuf) -> Arc<RemoteFiles> {
625    let (r, w) = testing::helper(home);
626    let f = RemoteFiles::over(conn, r, w, Box::new(())).await.unwrap();
627    let name = canonical(conn).unwrap();
628    let slot = pool().lock().unwrap().entry(name).or_default().clone();
629    *slot.lock().await = Some(f.clone());
630    f
631}
632
633#[cfg(test)]
634mod quiet_tests {
635    use super::*;
636
637    /// A host `connect_quiet` failed on lately gets the same answer again,
638    /// without another ssh, until the back-off has passed.
639    #[tokio::test]
640    async fn a_recent_failure_is_answered_without_trying_again() {
641        let name = canonical("me@quiet-backoff-test").unwrap();
642        quiet_failures().lock().unwrap().insert(name.clone(), (Instant::now(), "the helper is another version".into()));
643        let err = connect_quiet("me@quiet-backoff-test").await.err().expect("refused");
644        assert_eq!(err.message, "the helper is another version");
645        let long_ago = Instant::now().checked_sub(QUIET_RETRY_AFTER + Duration::from_secs(1)).unwrap();
646        quiet_failures().lock().unwrap().insert(name, (long_ago, "old".into()));
647        let err = connect_quiet("me@quiet-backoff-test").await.err().expect("still no helper here");
648        assert_ne!(err.message, "old", "after the back-off it is tried again");
649    }
650}
651
652#[cfg(test)]
653pub(crate) mod testing {
654    use super::*;
655
656    /// A helper in this process: the real `serve::handle` behind an
657    /// in-memory pipe, as ssh would carry it.
658    pub(crate) fn helper(
659        home: std::path::PathBuf,
660    ) -> (tokio::io::DuplexStream, tokio::io::DuplexStream) {
661        helper_with_pipes(home, 1 << 20)
662    }
663
664    /// [`helper`] with pipes of `size` bytes each way (an ssh's are small).
665    pub(crate) fn helper_with_pipes(
666        home: std::path::PathBuf,
667        size: usize,
668    ) -> (tokio::io::DuplexStream, tokio::io::DuplexStream) {
669        helper_with(home, size, Duration::ZERO)
670    }
671
672    /// [`helper`] that takes `delay` over each request, as a slow link
673    /// would: long enough to cancel something halfway.
674    #[cfg(target_os = "linux")]
675    pub(crate) fn slow_helper(
676        home: std::path::PathBuf,
677        delay: Duration,
678    ) -> (tokio::io::DuplexStream, tokio::io::DuplexStream) {
679        helper_with(home, 1 << 20, delay)
680    }
681
682    fn helper_with(
683        home: std::path::PathBuf,
684        size: usize,
685        delay: Duration,
686    ) -> (tokio::io::DuplexStream, tokio::io::DuplexStream) {
687        let (client_read, mut helper_write) = tokio::io::duplex(size);
688        let (mut helper_read, client_write) = tokio::io::duplex(size);
689        tokio::spawn(async move {
690            let mut server = agentmux_remote::serve::Server::new(home);
691            let mut s = Splitter::new();
692            let mut buf = vec![0u8; 65536];
693            loop {
694                let n = match helper_read.read(&mut buf).await {
695                    Ok(0) | Err(_) => return,
696                    Ok(n) => n,
697                };
698                for body in s.push(&buf[..n]).unwrap() {
699                    if !delay.is_zero() {
700                        tokio::time::sleep(delay).await;
701                    }
702                    let (id, req) = Request::decode(&body).unwrap();
703                    let reply = server.handle(req);
704                    helper_write.write_all(&reply.encode(id)).await.unwrap();
705                }
706            }
707        });
708        (client_read, client_write)
709    }
710}
711
712/// Start `serve --stdio` on the host. `Err((why, true))` when the helper
713/// looks missing (the stream ended before the greeting), so installing it
714/// is worth a try.
715async fn start(
716    name: &str,
717    mut host: HostSsh,
718    ask: Option<AskIn<'_>>,
719) -> Result<Arc<RemoteFiles>, (RemoteError, bool)> {
720    let grant = ask.and_then(|a| host.ask_user_in(a.block_id, name, a.auth_key));
721    let remote = format!(
722        "{} serve --stdio",
723        crate::backend::blockcontroller::durable_ssh::helper_path()
724    );
725    let mut cmd = host.command(&remote);
726    cmd.stdin(std::process::Stdio::piped())
727        .stdout(std::process::Stdio::piped())
728        .stderr(std::process::Stdio::piped());
729    let mut child = cmd
730        .spawn()
731        .map_err(|e| (RemoteError::link(format!("could not run ssh: {e}")), false))?;
732    let (Some(stdin), Some(stdout), Some(mut stderr)) =
733        (child.stdin.take(), child.stdout.take(), child.stderr.take())
734    else {
735        return Err((RemoteError::link("ssh started without its pipes"), false));
736    };
737    // ssh's complaints, kept for the error if the greeting never comes.
738    let said = Arc::new(Mutex::new(String::new()));
739    {
740        let said = said.clone();
741        tokio::spawn(async move {
742            let mut buf = Vec::new();
743            let _ = stderr.read_to_end(&mut buf).await;
744            *said.lock().unwrap() = String::from_utf8_lossy(&buf).into_owned();
745        });
746    }
747    match RemoteFiles::over(name, stdout, stdin, Box::new((child, grant))).await {
748        Ok(f) => Ok(f),
749        Err(e) => {
750            // Give stderr a moment to arrive after the stream ended.
751            tokio::time::sleep(Duration::from_millis(200)).await;
752            let said = said.lock().unwrap().clone();
753            let last = said
754                .lines()
755                .rev()
756                .find(|l| !l.trim().is_empty())
757                .unwrap_or("")
758                .trim()
759                .to_string();
760            let missing = said.contains("not found") || said.contains("No such file");
761            let message = if last.is_empty() {
762                e.message
763            } else {
764                format!("{} ({last})", e.message)
765            };
766            Err((RemoteError::link(message), missing))
767        }
768    }
769}
770
771#[cfg(test)]
772mod tests {
773    use super::*;
774
775    use super::testing::helper;
776
777    #[tokio::test]
778    async fn files_round_trip_through_the_helper() {
779        let home = tempfile::tempdir().unwrap();
780        let (r, w) = helper(home.path().to_path_buf());
781        let f = RemoteFiles::over("testhost", r, w, Box::new(()))
782            .await
783            .unwrap();
784        assert_eq!(std::path::Path::new(&f.home), home.path());
785
786        f.mkdir("~/d", false).await.unwrap();
787        // Bigger than one ranged read, so `read` has to take several.
788        let big: Vec<u8> = (0..(fsproto::MAX_READ as usize + 1000))
789            .map(|i| i as u8)
790            .collect();
791        f.write("~/d/big.bin", big.clone()).await.unwrap();
792        assert_eq!(f.read("~/d/big.bin", 64 << 20).await.unwrap(), big);
793        let e = f.read("~/d/big.bin", 10).await.unwrap_err();
794        assert_eq!(e.kind, Some(ErrKind::TooLarge));
795
796        for i in 0..1205 {
797            f.write(&format!("~/d/f{i:04}"), vec![]).await.unwrap();
798        }
799        // Several pages, all of them.
800        assert_eq!(f.list("~/d").await.unwrap().len(), 1206);
801
802        f.rename("~/d/f0000", "~/d/renamed").await.unwrap();
803        assert!(f.stat("~/d/f0000").await.unwrap_err().is_not_found());
804        assert_eq!(f.stat("~/d/renamed").await.unwrap().size, 0);
805        f.delete("~/d", true).await.unwrap();
806        assert!(f.stat("~/d").await.unwrap_err().is_not_found());
807    }
808
809    /// A big read's reply and a big write's request at once, over pipes far
810    /// smaller than either, against a helper that stops reading while it
811    /// writes: both finish (a writer that blocked the reader would not).
812    #[tokio::test]
813    async fn a_big_write_and_a_big_read_at_once_do_not_jam_small_pipes() {
814        let home = tempfile::tempdir().unwrap();
815        let big: Vec<u8> = (0..(3 << 20)).map(|i| (i % 251) as u8).collect();
816        std::fs::write(home.path().join("in.bin"), &big).unwrap();
817        let (r, w) = super::testing::helper_with_pipes(home.path().to_path_buf(), 64 * 1024);
818        let f = RemoteFiles::over("testhost", r, w, Box::new(()))
819            .await
820            .unwrap();
821        let (read, wrote) = tokio::time::timeout(Duration::from_secs(30), async {
822            tokio::join!(
823                f.read("~/in.bin", 64 << 20),
824                f.write("~/out.bin", vec![7u8; 8 << 20])
825            )
826        })
827        .await
828        .expect("no deadlock");
829        assert_eq!(read.unwrap(), big);
830        wrote.unwrap();
831        assert_eq!(
832            std::fs::metadata(home.path().join("out.bin"))
833                .unwrap()
834                .len(),
835            8 << 20
836        );
837    }
838
839    #[tokio::test]
840    async fn many_calls_in_flight_each_get_their_own_answer() {
841        let home = tempfile::tempdir().unwrap();
842        for i in 0..50 {
843            std::fs::write(home.path().join(format!("n{i}")), i.to_string()).unwrap();
844        }
845        let (r, w) = helper(home.path().to_path_buf());
846        let f = RemoteFiles::over("testhost", r, w, Box::new(()))
847            .await
848            .unwrap();
849        let reads = (0..50).map(|i| {
850            let f = f.clone();
851            async move { (i, f.read(&format!("~/n{i}"), 100).await.unwrap()) }
852        });
853        for (i, data) in futures_util::future::join_all(reads).await {
854            assert_eq!(data, i.to_string().into_bytes());
855        }
856    }
857
858    #[tokio::test]
859    async fn a_dropped_link_fails_every_call_and_closes() {
860        let (client_read, helper_write) = tokio::io::duplex(1024);
861        let (mut helper_read, client_write) = tokio::io::duplex(1024);
862        // Greets, then drops at the next request.
863        let helper = tokio::spawn(async move {
864            let mut helper_write = helper_write;
865            let mut buf = vec![0u8; 1024];
866            let n = helper_read.read(&mut buf).await.unwrap();
867            let (id, _) = Request::decode(&Splitter::new().push(&buf[..n]).unwrap()[0]).unwrap();
868            let hello = Reply::Hello {
869                protocol: fsproto::PROTOCOL,
870                home: "/home/u".into(),
871                os: "linux".into(),
872            };
873            helper_write.write_all(&hello.encode(id)).await.unwrap();
874            let _ = helper_read.read(&mut buf).await;
875        });
876        let f = RemoteFiles::over("testhost", client_read, client_write, Box::new(()))
877            .await
878            .unwrap();
879        let pending = f.stat("~/x");
880        let (r, _) = tokio::join!(pending, helper);
881        let e = r.unwrap_err();
882        assert_eq!(e.kind, None, "{e:?}");
883        assert!(e.message.contains("closed"), "{e:?}");
884        assert!(!f.is_open());
885        assert!(f.stat("~/y").await.is_err());
886    }
887
888    #[tokio::test]
889    async fn a_helper_of_another_protocol_is_refused() {
890        let (client_read, mut helper_write) = tokio::io::duplex(1024);
891        let (mut helper_read, client_write) = tokio::io::duplex(1024);
892        tokio::spawn(async move {
893            let mut buf = vec![0u8; 1024];
894            let n = helper_read.read(&mut buf).await.unwrap();
895            let (id, _) = Request::decode(&Splitter::new().push(&buf[..n]).unwrap()[0]).unwrap();
896            let hello = Reply::Hello {
897                protocol: fsproto::PROTOCOL + 1,
898                home: "/".into(),
899                os: "linux".into(),
900            };
901            helper_write.write_all(&hello.encode(id)).await.unwrap();
902            let _ = helper_read.read(&mut buf).await;
903        });
904        let e = RemoteFiles::over("testhost", client_read, client_write, Box::new(()))
905            .await
906            .err()
907            .unwrap();
908        assert!(e.message.contains("file protocol"), "{e:?}");
909    }
910}
agentmux_srv/backend/remote/files.rs

agentmux_srv/backend/remote/
files.rs
