diff --git a/core/src/connection.rs b/core/src/connection.rs index e650ddc8..9d0bf41d 100644 --- a/core/src/connection.rs +++ b/core/src/connection.rs @@ -1,6 +1,7 @@ use crate::host::{AuthMethod, Host}; use crate::identity::Identity; use crate::known_hosts::KnownHosts; +use crate::transport::Transport; use crate::vault::Vault; use crate::{CoreError, Result}; use russh::client::{self, Config, Handle}; @@ -252,6 +253,14 @@ fn rsa_hash_for(advertised: Option>) -> HashAlg { } } +/// An authenticated SSH connection. +pub struct Connection { + pub handle: Handle, + /// Ends the connection when the peer will not take a goodbye; see + /// [`crate::transport`]. + pub transport: Transport, +} + /// The account, method and vault key a host signs in with. pub struct Credentials<'a> { pub username: &'a str, @@ -289,7 +298,7 @@ pub fn credentials<'a>(host: &'a Host, identity: Option<&'a Identity>) -> Result }) } -/// Open an authenticated SSH session and return the handle. +/// Open an authenticated SSH session and return it. /// The caller is responsible for opening a channel and starting the shell. /// The credentials come from `credentials`. pub async fn connect( @@ -299,7 +308,7 @@ pub async fn connect( vault: Option<&Vault>, passphrase: Option<&str>, password: Option<&str>, -) -> Result> { +) -> Result { let Credentials { username, auth, @@ -328,22 +337,33 @@ pub async fn connect( // connect failure would hide the two fingerprints that make the change // legible. `CoreError::Ssh` only reaches here from `check_server_key`, which // already names the address, so it is passed through unwrapped too. - let mut session = tokio::time::timeout( - CONNECT_TIMEOUT, - client::connect(client_config(), &addr, handler), - ) - .await - .map_err(|_| timed_out("connect", &addr))? - .map_err(|e| match e { - CoreError::HostKeyMismatch { .. } | CoreError::Ssh(_) => e, - other => CoreError::Ssh(format!( - "connect {addr}: {other}{}", - negotiation_hint(&other) - )), - })?; + let opened = async { + let config = client_config(); + let socket = tokio::net::TcpStream::connect(&addr) + .await + .map_err(|e| CoreError::from(russh::Error::IO(e)))?; + if config.nodelay { + if let Err(e) = socket.set_nodelay(true) { + tracing::warn!("set_nodelay failed for {addr}: {e}"); + } + } + let (stream, transport) = Transport::wrap(socket); + let handle = client::connect_stream(config, stream, handler).await?; + Ok::<_, CoreError>((handle, transport)) + }; + let (mut session, transport) = tokio::time::timeout(CONNECT_TIMEOUT, opened) + .await + .map_err(|_| timed_out("connect", &addr))? + .map_err(|e| match e { + CoreError::HostKeyMismatch { .. } | CoreError::Ssh(_) => e, + other => CoreError::Ssh(format!( + "connect {addr}: {other}{}", + negotiation_hint(&other) + )), + })?; // Authentication is bounded separately, and the two bounds cover disjoint - // stages. `client::connect` does not return until russh has read the + // stages. `client::connect_stream` does not return until russh has read the // server's banner and the key exchange has completed, so a peer that // answers the TCP handshake and then goes silent trips the connect timeout // above and never reaches this point. What is left for this bound is a @@ -402,12 +422,25 @@ pub async fn connect( outcome.map_err(|e| CoreError::Ssh(format!("auth: {e}"))) }) .await - .map_err(|_| timed_out("auth", &addr))??; - - if !auth_ok.success() { - return Err(CoreError::Ssh("authentication rejected by server".into())); + .map_err(|_| timed_out("auth", &addr)) + .and_then(|outcome| outcome) + .and_then(|outcome| { + if outcome.success() { + Ok(()) + } else { + Err(CoreError::Ssh("authentication rejected by server".into())) + } + }); + if let Err(e) = auth_ok { + // The session task would otherwise run on until it noticed the + // dropped handle, holding an unauthenticated connection open. + transport.cut(); + return Err(e); } - Ok(session) + Ok(Connection { + handle: session, + transport, + }) } #[cfg(test)] diff --git a/core/src/lib.rs b/core/src/lib.rs index 44315adc..f42233c8 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -32,6 +32,7 @@ pub mod output; pub mod session; pub mod sftp; pub mod store; +pub mod transport; pub mod vault; pub mod workspace; diff --git a/core/src/session.rs b/core/src/session.rs index 3bf22469..41f2f795 100644 --- a/core/src/session.rs +++ b/core/src/session.rs @@ -2,6 +2,7 @@ use crate::connection; use crate::host::Host; use crate::known_hosts::KnownHosts; use crate::output::OutputBatcher; +use crate::transport::{self, CloseSignal, Transport}; use crate::vault::Vault; use crate::{CoreError, Result}; use russh::client::Handle; @@ -9,7 +10,9 @@ use russh::ChannelMsg; use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::Instant; -use tokio::sync::{mpsc, Mutex}; +use tokio::sync::{watch, Mutex}; + +type SessionMap = HashMap>; /// Callback invoked when a session receives data from the remote end. /// The Tauri layer uses this to forward a batch of output to the frontend. @@ -28,7 +31,11 @@ async fn sleep_until(deadline: Option) { /// Manages all active terminal sessions (SSH and local). pub struct SessionManager { - sessions: Mutex>, + /// Live sessions by id. Held only to look an entry up, insert or remove + /// it, never across network I/O: one peer that stops reading would + /// otherwise stall writes, resizes and closes on every other session, and + /// keep a vault reset from draining the map. + sessions: Arc>, /// Every id held by an SSH session, from the moment it is claimed — before /// the connection exists — until the session is closed. `sessions` only /// learns an id once the connection is up, seconds later, so it cannot tell @@ -106,23 +113,49 @@ fn lock_claims( } struct SshSession { - #[allow(dead_code)] - id: String, - handle: Arc>>, + /// Shared with the goodbye sent after the session is closed, which needs + /// nothing else from it, so that goodbye does not keep the id claimed. + handle: Arc>, + transport: Transport, + /// Raised once the session is closed. A write stuck behind a peer that + /// stopped reading watches it and gives up at once rather than when the + /// goodbye ends the connection. + closed: CloseSignal, channel_id: russh::ChannelId, - resize_tx: mpsc::Sender<(u32, u32)>, + /// Serializes writes to this session so keystrokes reach the channel in + /// the order they were sent. Per session, so a stalled peer only holds up + /// its own input. + writes: Mutex<()>, + /// The size the remote terminal should have. Only the latest one matters, + /// so a resize replaces it and never waits; the session's resizer task + /// sends it on whenever the connection takes it. + size: watch::Sender<(u32, u32)>, /// Held so the id stays reserved while this entry exists. #[allow(dead_code)] claim: Arc, } +impl SshSession { + /// Mark the session closed and say goodbye to the peer, cutting the + /// connection if the goodbye does not get through in time. The goodbye + /// holds only the handle and the transport, so the id is released as soon + /// as the map entry and the reader are gone, and the reader goes once the + /// connection does, so within [`transport::GOODBYE_GRACE`] at most. + fn retire(&self) -> impl std::future::Future + Send + 'static { + self.closed.raise(); + let handle = self.handle.clone(); + let transport = self.transport.clone(); + async move { transport::goodbye(&handle, &transport).await } + } +} + // Re-export the handler type so the session manager can use it. pub use connection::SshHandler; impl SessionManager { pub fn new(data_callback: DataCallback, close_callback: CloseCallback) -> Self { Self { - sessions: Mutex::new(HashMap::new()), + sessions: Arc::new(Mutex::new(HashMap::new())), claimed: Arc::new(std::sync::Mutex::new(HashSet::new())), data_callback, close_callback, @@ -170,68 +203,93 @@ impl SessionManager { rows: u32, ) -> Result<()> { let session_id = claim.id.clone(); - let handle = + let connection::Connection { handle, transport } = connection::connect(host, identity, known_hosts, vault, passphrase, password).await?; - let channel = handle.channel_open_session().await.map_err(|e| { - CoreError::Ssh(format!("channel open: {e}")) - })?; - let channel_id = channel.id(); - - // Request PTY - channel - .request_pty( - false, - "xterm-256color", - cols, - rows, - 0, - 0, - &[], // default terminal modes - ) - .await - .map_err(|e| CoreError::Ssh(format!("request pty: {e}")))?; - - // Request shell - channel - .request_shell(false) - .await - .map_err(|e| CoreError::Ssh(format!("request shell: {e}")))?; - - // Run startup command if set - if let Some(cmd) = &host.startup_command { - if !cmd.is_empty() { - channel - .exec(false, cmd.clone()) - .await - .map_err(|e| CoreError::Ssh(format!("startup exec: {e}")))?; + let opened = async { + let channel = handle + .channel_open_session() + .await + .map_err(|e| CoreError::Ssh(format!("channel open: {e}")))?; + + // Request PTY + channel + .request_pty( + false, + "xterm-256color", + cols, + rows, + 0, + 0, + &[], // default terminal modes + ) + .await + .map_err(|e| CoreError::Ssh(format!("request pty: {e}")))?; + + // Request shell + channel + .request_shell(false) + .await + .map_err(|e| CoreError::Ssh(format!("request shell: {e}")))?; + + // Run startup command if set + if let Some(cmd) = &host.startup_command { + if !cmd.is_empty() { + channel + .exec(false, cmd.clone()) + .await + .map_err(|e| CoreError::Ssh(format!("startup exec: {e}")))?; + } } - } - - let handle = Arc::new(Mutex::new(handle)); - let (resize_tx, mut resize_rx) = mpsc::channel::<(u32, u32)>(32); + Ok::<_, CoreError>(channel) + }; + let channel = match opened.await { + Ok(channel) => channel, + Err(e) => { + transport.cut(); + return Err(e); + } + }; + let channel_id = channel.id(); + // The reader keeps only the read half, so nothing it does waits on the + // connection taking outbound messages; it always drains what arrives. + let (mut reader, writer) = channel.split(); + let (size, mut size_rx) = watch::channel((cols, rows)); let held = claim.commit(); - let session = SshSession { - id: session_id.clone(), - handle: handle.clone(), + let session = Arc::new(SshSession { + handle: Arc::new(handle), + transport, + closed: CloseSignal::new(), channel_id, - resize_tx, + writes: Mutex::new(()), + size, claim: held.clone(), - }; + }); self.sessions.lock().await.insert(session_id.clone(), session); + let sessions = self.sessions.clone(); + + // Window changes wait whenever the connection does, so they go out from + // a task of their own. It ends once the entry, and the sender with it, + // is gone, or once the connection is. + tokio::spawn(async move { + while size_rx.changed().await.is_ok() { + let (cols, rows) = *size_rx.borrow_and_update(); + if writer.window_change(cols, rows, 0, 0).await.is_err() { + break; + } + } + }); // Spawn the reading task let sid = session_id.clone(); let data_cb = self.data_callback.clone(); let close_cb = self.close_callback.clone(); - let mut channel = channel; tokio::spawn(async move { // Dropped when this task ends, after the close has been announced, // which is the moment the id is free for another session. - let _held = held; - let mut pending_resize: Option<(u32, u32)> = None; + let held = held; let mut batcher = OutputBatcher::new(Instant::now()); let flush = |batcher: &mut OutputBatcher| { if !batcher.is_empty() { @@ -244,7 +302,7 @@ impl SessionManager { // batcher while another branch's handler mutates it. let flush_at = batcher.deadline(); tokio::select! { - msg = channel.wait() => { + msg = reader.wait() => { match msg { Some(ChannelMsg::Data { data }) | Some(ChannelMsg::ExtendedData { data, .. }) => { @@ -259,63 +317,84 @@ impl SessionManager { // order; see `AppState::init`. flush(&mut batcher); close_cb(&sid, "session closed"); + // A shell that exits leaves the connection up; + // end it, unless a close already took the entry. + // The goodbye runs on its own so the id is free + // for a reconnect as soon as this task ends. + if let Some(session) = take_if_held(&sessions, &sid, &held).await { + tokio::spawn(session.retire()); + } break; } _ => {} } } - resize = resize_rx.recv() => { - pending_resize = resize; - } () = sleep_until(flush_at) => { flush(&mut batcher); } } - if let Some((c, r)) = pending_resize.take() { - let _ = channel.window_change(c, r, 0, 0).await; - } } }); Ok(()) } + /// The live session behind `session_id`, taken out from under the map lock. + async fn get(&self, session_id: &str) -> Result> { + self.sessions + .lock() + .await + .get(session_id) + .cloned() + .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string())) + } + + /// Whether an SSH session is live under `session_id`. + pub async fn contains(&self, session_id: &str) -> bool { + self.sessions.lock().await.contains_key(session_id) + } + /// Write data to a session's terminal. + /// + /// The send waits for as long as the peer leaves the connection unread, + /// and only this session's writes wait behind it. A close ends the wait. pub async fn write(&self, session_id: &str, data: &[u8]) -> Result<()> { - let sessions = self.sessions.lock().await; - let session = sessions - .get(session_id) - .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string()))?; - let handle = session.handle.lock().await; - handle - .data(session.channel_id, data.to_vec()) + let session = self.get(session_id).await?; + let send = async { + let _order = session.writes.lock().await; + session + .handle + .data(session.channel_id, data.to_vec()) + .await + .map_err(|_| CoreError::Ssh("write failed".into())) + }; + session + .closed + .guard(send, || CoreError::SessionNotFound(session_id.to_string())) .await - .map_err(|_| CoreError::Ssh("write failed".into()))?; - Ok(()) } - /// Resize a session's terminal. + /// Resize a session's terminal. Never waits on the peer: the new size + /// replaces any that has not gone out yet. pub async fn resize(&self, session_id: &str, cols: u32, rows: u32) -> Result<()> { - let sessions = self.sessions.lock().await; - let session = sessions - .get(session_id) - .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string()))?; - session - .resize_tx - .send((cols, rows)) - .await - .map_err(|_| CoreError::Ssh("resize channel closed".into()))?; + self.get(session_id).await?.size.send_replace((cols, rows)); Ok(()) } /// Close a session. + /// + /// The entry leaves the map before the goodbye is sent, so the session + /// stops being addressable at once even if the peer never answers, and the + /// entry's share of the id goes with it rather than waiting on the + /// goodbye. The reader task still holds a share until it has flushed and + /// announced the close, which the goodbye guarantees within + /// [`transport::GOODBYE_GRACE`]. pub async fn close(&self, session_id: &str) -> Result<()> { - let mut sessions = self.sessions.lock().await; - if let Some(session) = sessions.remove(session_id) { - // The entry's share of the id goes with it; the reader task still - // holds one until it has flushed and announced the close. - let handle = session.handle.lock().await; - let _ = handle.disconnect(russh::Disconnect::ByApplication, "", "en").await; + let removed = self.sessions.lock().await.remove(session_id); + if let Some(session) = removed { + let goodbye = session.retire(); + drop(session); + goodbye.await; } Ok(()) } @@ -329,24 +408,42 @@ impl SessionManager { /// /// The map is drained under the lock and the guard is released before any /// network work, so a session stops being addressable the moment this is - /// called. A remote that never acknowledges the disconnect can therefore - /// only delay the polite goodbye, never keep a session reachable. + /// called. The goodbyes then run side by side, each bounded by + /// [`transport::GOODBYE_GRACE`] and ending its connection either way, so + /// one peer that never answers cannot hold up the others or outlive the + /// call. pub async fn close_all(&self) -> usize { - let drained: Vec = { + let drained: Vec> = { let mut sessions = self.sessions.lock().await; sessions.drain().map(|(_, session)| session).collect() }; let count = drained.len(); - for session in &drained { - let handle = session.handle.lock().await; - let _ = handle - .disconnect(russh::Disconnect::ByApplication, "", "en") - .await; + let goodbyes: Vec<_> = drained + .into_iter() + .map(|session| tokio::spawn(session.retire())) + .collect(); + for goodbye in goodbyes { + let _ = goodbye.await; } count } } +/// Remove `session_id` from the map if the entry there is still the session +/// `held` belongs to. An id is reused once its session has ended, and a +/// reader finishing late must not take a newer session's entry. +async fn take_if_held( + sessions: &Mutex, + session_id: &str, + held: &Arc, +) -> Option> { + let mut sessions = sessions.lock().await; + match sessions.get(session_id) { + Some(session) if Arc::ptr_eq(&session.claim, held) => sessions.remove(session_id), + _ => None, + } +} + #[cfg(test)] mod claim_tests { use super::SessionManager; diff --git a/core/src/sftp.rs b/core/src/sftp.rs index 6c71aed4..45e598c7 100644 --- a/core/src/sftp.rs +++ b/core/src/sftp.rs @@ -2,15 +2,39 @@ use crate::connection; use crate::host::Host; use crate::identity::Identity; use crate::known_hosts::KnownHosts; +use crate::transport::{self, CloseSignal, Transport}; use crate::vault::Vault; use crate::{CoreError, Result}; use russh::client::Handle; -use russh_sftp::client::SftpSession; +use russh_sftp::client::error::Error as SftpError; +use russh_sftp::client::{RawSftpSession, SftpSession}; +use russh_sftp::protocol::{FileAttributes, OpenFlags, StatusCode}; use std::collections::HashMap; -use std::path::Path; +use std::future::Future; +use std::hash::Hash; +use std::path::{Path, PathBuf}; use std::sync::Arc; +use std::time::Duration; use tokio::io::AsyncWriteExt; -use tokio::sync::Mutex; +use tokio::sync::{Mutex, OwnedMutexGuard}; + +/// The most entries one listing returns. A server that keeps answering a +/// listing with more entries would otherwise keep it, and its memory, growing +/// for good; a directory past this is refused rather than shown in part. +const MAX_LIST_ENTRIES: usize = 200_000; + +/// The most name bytes one listing collects, for the same reason: entry +/// names are not bounded by the protocol, so the entry cap alone does not +/// bound the memory. +const MAX_LIST_NAME_BYTES: usize = 32 * 1024 * 1024; + +/// The longest a listing may take in total on a server without a listing +/// channel of its own; see [`SftpConnection::lister`]. +const FALLBACK_LIST_DEADLINE: Duration = Duration::from_secs(30); + +/// How long opening the listing channel may take before the connection does +/// without it. +const LISTER_OPEN_TIMEOUT: Duration = Duration::from_secs(10); /// A directory entry returned by SFTP list operations. #[derive(Debug, Clone, serde::Serialize)] @@ -25,21 +49,246 @@ pub struct SftpEntry { pub permissions: Option, } +impl SftpEntry { + fn new(name: String, long_name: String, attrs: &FileAttributes) -> Self { + let modified = attrs + .modified() + .ok() + .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()) + .map(|d| d.as_secs() as i64); + Self { + name, + long_name, + is_dir: attrs.is_dir(), + is_file: attrs.is_regular(), + is_symlink: attrs.is_symlink(), + size: attrs.len(), + modified, + permissions: attrs.permissions, + } + } +} + /// Manages SFTP sessions for multiple hosts. +/// +/// The session map is locked only to look an entry up, insert or remove it, +/// never across network I/O. A transfer or a listing therefore holds up +/// nothing but itself: other sessions keep working, and `close` and +/// `close_all` can always take an entry out of reach at once. pub struct SftpManager { - sessions: Mutex>, + sessions: Mutex>>, + /// The local files a download is under way to, across every session: the + /// local disk is one for all of them. See [`download_destination`]. + downloads: WriterLocks, } struct SftpConnection { - #[allow(dead_code)] - handle: Arc>>, + handle: Handle, + transport: Transport, sftp: SftpSession, + /// A second SFTP channel that serves directory listings, request by + /// request. `SftpSession::read_dir` collects a whole directory before it + /// returns, so it can neither cap a listing's size nor close the + /// directory handle when a listing is abandoned. `None` when the server + /// would not open the channel; listings then fall back to `read_dir` + /// under [`FALLBACK_LIST_DEADLINE`]. + lister: Option, + /// Raised once the connection is closed. Every operation watches it and + /// stops, so one that is stuck behind a peer that stopped reading, or + /// simply still running, cannot outlive the close or report success + /// after it. + closed: CloseSignal, + /// The remote paths an upload or write is under way to, each with the + /// lock the next writer to that path waits on; see + /// [`SftpConnection::claim_destination`]. + writers: WriterLocks, +} + +/// One lock per destination a writer is under way to, the next writer to the +/// same destination waiting on it. +type WriterLocks = std::sync::Mutex>>>; + +fn lock_writers( + writers: &WriterLocks, +) -> std::sync::MutexGuard<'_, HashMap>>> { + // The map only ever holds valid locks, whatever a panicking holder was doing. + writers.lock().unwrap_or_else(|e| e.into_inner()) +} + +/// Sole right to write one destination, released on drop. +struct DestinationGuard<'a, K: Eq + Hash> { + writers: &'a WriterLocks, + destination: K, + lock: Option>>, + held: Option>, +} + +impl Drop for DestinationGuard<'_, K> { + fn drop(&mut self) { + drop(self.held.take()); + let mut writers = lock_writers(self.writers); + drop(self.lock.take()); + // Every writer holding or waiting on the destination keeps a clone, so + // a count of one means the map's own copy is the last and the entry + // can go. The count is read under the map lock, which every clone is + // taken under, so no writer can be joining at the same moment. + if writers + .get(&self.destination) + .is_some_and(|lock| Arc::strong_count(lock) == 1) + { + writers.remove(&self.destination); + } + } +} + +/// Wait until no other writer in `writers` is writing `destination`, then hold +/// it until the guard drops. +async fn claim( + writers: &WriterLocks, + destination: K, +) -> DestinationGuard<'_, K> { + let lock = lock_writers(writers) + .entry(destination.clone()) + .or_default() + .clone(); + // Built before the wait so that a writer cancelled while waiting still + // takes its entry out of the map when nothing else needs it. + let mut guard = DestinationGuard { + writers, + destination, + lock: Some(lock.clone()), + held: None, + }; + guard.held = Some(lock.lock_owned().await); + guard +} + +/// The key a download to `local_path` is serialized under: the path with its +/// directory resolved, so that two spellings of one directory, through a +/// symlink or a `..`, still name one file. Falls back to the path as given +/// when the directory cannot be resolved; the open that follows then fails +/// on its own. +async fn download_destination(local_path: &Path) -> PathBuf { + let resolved = match (local_path.parent(), local_path.file_name()) { + (Some(parent), Some(name)) if !parent.as_os_str().is_empty() => { + tokio::fs::canonicalize(parent) + .await + .ok() + .map(|parent| parent.join(name)) + } + _ => None, + }; + resolved.unwrap_or_else(|| local_path.to_path_buf()) +} + +impl SftpConnection { + /// Wait until no other upload or write on this connection is writing + /// `path`, then hold it until the guard drops. SFTP calls on one connection + /// run side by side, and two writers on one file would otherwise truncate + /// and write it at once, leaving a mix of both that each reports as a + /// success. Paths are compared as given, which is how the browser builds + /// them. + async fn claim_destination(&self, path: &str) -> DestinationGuard<'_, String> { + claim(&self.writers, path.to_owned()).await + } + + /// Run `operation` until it finishes or the connection is closed. + async fn until_closed(&self, operation: impl Future>) -> Result { + self.closed + .guard(operation, || { + CoreError::Ssh("the SFTP session was closed".into()) + }) + .await + } + + /// Stop every request, listing and transfer on this connection, end its + /// SFTP channels, then the connection. The goodbye is bounded and cuts the + /// connection if the peer does not take it, so an unresponsive peer can + /// neither hold the caller nor keep the connection open. + async fn close(&self) { + self.closed.raise(); + let _ = self.sftp.close().await; + if let Some(lister) = &self.lister { + let _ = lister.close_session(); + } + transport::goodbye(&self.handle, &self.transport).await; + } +} + +/// Open a channel on `handle` and start the SFTP subsystem on it. +async fn open_sftp_channel( + handle: &Handle, +) -> Result> { + let channel = handle + .channel_open_session() + .await + .map_err(|e| CoreError::Ssh(format!("sftp channel open: {e}")))?; + channel + .request_subsystem(true, "sftp") + .await + .map_err(|e| CoreError::Ssh(format!("request sftp subsystem: {e}")))?; + Ok(channel.into_stream()) +} + +/// Open the listing channel, or `None` when the server will not have it, for +/// instance one that allows a single session per connection. +async fn open_lister(handle: &Handle) -> Option { + let opened = async { + let lister = RawSftpSession::new(open_sftp_channel(handle).await.ok()?); + lister.init().await.ok()?; + Some(lister) + }; + tokio::time::timeout(LISTER_OPEN_TIMEOUT, opened) + .await + .ok() + .flatten() +} + +/// Read the directory behind `handle` page by page, refusing it once it +/// passes [`MAX_LIST_ENTRIES`] or [`MAX_LIST_NAME_BYTES`]. Each page is bounded +/// by the session's per-request timeout, so a server that goes quiet ends the +/// listing too. +async fn read_listing(lister: &RawSftpSession, handle: &str) -> Result> { + let too_large = || { + CoreError::Ssh(format!( + "read_dir: the directory has more than {MAX_LIST_ENTRIES} entries or \ + {} MiB of names, more than Clavyn lists", + MAX_LIST_NAME_BYTES / (1024 * 1024) + )) + }; + let mut entries = Vec::new(); + let mut name_bytes = 0usize; + loop { + match lister.readdir(handle).await { + Ok(page) => { + for file in page.files { + if file.filename == "." || file.filename == ".." { + continue; + } + name_bytes = name_bytes.saturating_add(file.filename.len()); + if entries.len() >= MAX_LIST_ENTRIES || name_bytes > MAX_LIST_NAME_BYTES { + return Err(too_large()); + } + entries.push(SftpEntry::new( + file.filename.clone(), + file.filename, + &file.attrs, + )); + } + } + Err(SftpError::Status(status)) if status.status_code == StatusCode::Eof => { + return Ok(entries); + } + Err(e) => return Err(CoreError::Ssh(format!("read_dir: {e}"))), + } + } } impl SftpManager { pub fn new() -> Self { Self { sessions: Mutex::new(HashMap::new()), + downloads: WriterLocks::default(), } } @@ -68,37 +317,44 @@ impl SftpManager { } } - let handle = connection::connect(host, identity, known_hosts, vault, passphrase, password) - .await?; - - // Open a channel and request the SFTP subsystem - let channel = handle - .channel_open_session() - .await - .map_err(|e| CoreError::Ssh(format!("sftp channel open: {e}")))?; + let connection::Connection { handle, transport } = + connection::connect(host, identity, known_hosts, vault, passphrase, password).await?; - channel - .request_subsystem(true, "sftp") - .await - .map_err(|e| CoreError::Ssh(format!("request sftp subsystem: {e}")))?; - - // Convert the channel into an async stream for SftpSession - let stream = channel.into_stream(); - let sftp = SftpSession::new(stream) - .await - .map_err(|e| CoreError::Ssh(format!("sftp init: {e}")))?; - - let handle = Arc::new(Mutex::new(handle)); + // The two channels open side by side so the listing channel adds no + // round trips to the connect. + let session = async { + SftpSession::new(open_sftp_channel(&handle).await?) + .await + .map_err(|e| CoreError::Ssh(format!("sftp init: {e}"))) + }; + let (sftp, lister) = tokio::join!(session, open_lister(&handle)); + let sftp = match sftp { + Ok(sftp) => sftp, + Err(e) => { + transport.cut(); + return Err(e); + } + }; - let conn = SftpConnection { handle, sftp }; + let conn = Arc::new(SftpConnection { + handle, + transport, + sftp, + lister, + closed: CloseSignal::new(), + writers: WriterLocks::default(), + }); // The check at the top ran before a connect that takes seconds, so a // second connect on the same id may have finished in the meantime. The - // session already published keeps the id; this one is dropped rather + // session already published keeps the id; this one is ended rather // than silently replacing it. match self.sessions.lock().await.entry(session_id) { - std::collections::hash_map::Entry::Occupied(_) => Err(CoreError::InvalidInput( - "session id is already in use by an SFTP session".into(), - )), + std::collections::hash_map::Entry::Occupied(_) => { + conn.transport.cut(); + Err(CoreError::InvalidInput( + "session id is already in use by an SFTP session".into(), + )) + } std::collections::hash_map::Entry::Vacant(slot) => { slot.insert(conn); Ok(()) @@ -106,40 +362,80 @@ impl SftpManager { } } + /// The live connection behind `session_id`, taken out from under the map + /// lock. + async fn get(&self, session_id: &str) -> Result> { + self.sessions + .lock() + .await + .get(session_id) + .cloned() + .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string())) + } + + /// Run one request on `session_id`'s connection until it answers or the + /// connection is closed, and name the request in its error. + async fn request( + &self, + session_id: &str, + what: &str, + request: impl FnOnce(Arc) -> F, + ) -> Result + where + F: Future>, + { + let conn = self.get(session_id).await?; + let pending = request(conn.clone()); + conn.until_closed(async move { + pending + .await + .map_err(|e| CoreError::Ssh(format!("{what}: {e}"))) + }) + .await + } + /// List directory entries at the given path. pub async fn list_dir(&self, session_id: &str, path: &str) -> Result> { - let mut sessions = self.sessions.lock().await; - let conn = sessions - .get_mut(session_id) - .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string()))?; - - let mut entries = Vec::new(); - let mut dir = conn - .sftp - .read_dir(path) - .await - .map_err(|e| CoreError::Ssh(format!("read_dir: {e}")))?; - - // ReadDir is a plain Iterator (not async) - for entry in dir.by_ref() { - let name = entry.file_name(); - let attrs = entry.metadata(); - let modified = attrs - .modified() - .ok() - .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()) - .map(|d| d.as_secs() as i64); - entries.push(SftpEntry { - name: name.clone(), - long_name: name, - is_dir: attrs.is_dir(), - is_file: attrs.is_regular(), - is_symlink: attrs.is_symlink(), - size: attrs.len(), - modified, - permissions: attrs.permissions, - }); - } + let conn = self.get(session_id).await?; + + let mut entries = conn + .until_closed(async { + match &conn.lister { + Some(lister) => { + let handle = lister + .opendir(path) + .await + .map_err(|e| CoreError::Ssh(format!("read_dir: {e}")))? + .handle; + let listed = read_listing(lister, &handle).await; + // Closed however the listing went, so one that failed + // or was refused leaves no handle open on the server. + let _ = lister.close(handle).await; + listed + } + None => { + let dir = tokio::time::timeout( + FALLBACK_LIST_DEADLINE, + conn.sftp.read_dir(path), + ) + .await + .map_err(|_| { + CoreError::Ssh(format!( + "read_dir: the listing took longer than {}s", + FALLBACK_LIST_DEADLINE.as_secs() + )) + })? + .map_err(|e| CoreError::Ssh(format!("read_dir: {e}")))?; + Ok(dir + .map(|entry| { + let name = entry.file_name(); + SftpEntry::new(name.clone(), name, &entry.metadata()) + }) + .collect()) + } + } + }) + .await?; // Sort: directories first, then files, alphabetically entries.sort_by(|a, b| { @@ -155,15 +451,10 @@ impl SftpManager { /// Get the canonical absolute path (e.g., resolve "~"). pub async fn canonicalize(&self, session_id: &str, path: &str) -> Result { - let mut sessions = self.sessions.lock().await; - let conn = sessions - .get_mut(session_id) - .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string()))?; - - conn.sftp - .canonicalize(path) - .await - .map_err(|e| CoreError::Ssh(format!("canonicalize: {e}"))) + self.request(session_id, "canonicalize", |conn| async move { + conn.sftp.canonicalize(path).await + }) + .await } /// Read a file's contents into a byte vector. @@ -171,15 +462,10 @@ impl SftpManager { /// Kept for API compatibility. Interactive file transfers should use /// `download_to_local` so large files never cross Tauri IPC as JSON arrays. pub async fn read_file(&self, session_id: &str, path: &str) -> Result> { - let mut sessions = self.sessions.lock().await; - let conn = sessions - .get_mut(session_id) - .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string()))?; - - conn.sftp - .read(path) - .await - .map_err(|e| CoreError::Ssh(format!("read file: {e}"))) + self.request(session_id, "read file", |conn| async move { + conn.sftp.read(path).await + }) + .await } /// Write data to a remote file (creates or overwrites). @@ -187,20 +473,17 @@ impl SftpManager { /// Kept for API compatibility. Interactive file transfers should use /// `upload_from_local` so large files never cross Tauri IPC as JSON arrays. pub async fn write_file(&self, session_id: &str, path: &str, data: &[u8]) -> Result<()> { - let mut sessions = self.sessions.lock().await; - let conn = sessions - .get_mut(session_id) - .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string()))?; - - conn.sftp - .write(path, data) - .await - .map_err(|e| CoreError::Ssh(format!("write file: {e}"))) + self.request(session_id, "write file", |conn| async move { + let _destination = conn.claim_destination(path).await; + conn.sftp.write(path, data).await + }) + .await } /// Stream a remote file directly to a local path. An existing local file is /// refused unless the caller explicitly opts in to overwrite, mirroring - /// `upload_from_local`. + /// `upload_from_local`. A download waits for any other download to the same + /// local file, from any session, to finish first. pub async fn download_to_local( &self, session_id: &str, @@ -208,25 +491,32 @@ impl SftpManager { local_path: &Path, overwrite: bool, ) -> Result<()> { - let mut sessions = self.sessions.lock().await; - let conn = sessions - .get_mut(session_id) - .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string()))?; - - let mut remote = conn - .sftp - .open(remote_path) - .await - .map_err(|e| CoreError::Ssh(format!("open remote file: {e}")))?; - let mut local = open_download_target(local_path, overwrite).await?; + let conn = self.get(session_id).await?; + + conn.until_closed(async { + // Two downloads to one local file would otherwise truncate and + // write it at once, from any two sessions, leaving a mix of both + // that each reports as a success. + let destination = download_destination(local_path).await; + let _destination = claim(&self.downloads, destination).await; + let mut remote = conn + .sftp + .open(remote_path) + .await + .map_err(|e| CoreError::Ssh(format!("open remote file: {e}")))?; + let mut local = open_download_target(local_path, overwrite).await?; - tokio::io::copy(&mut remote, &mut local).await?; - local.flush().await?; - Ok(()) + tokio::io::copy(&mut remote, &mut local).await?; + local.flush().await?; + Ok(()) + }) + .await } /// Stream a local file directly to the remote host. Existing remote files - /// are rejected unless the caller explicitly opts in to overwrite. + /// are rejected unless the caller explicitly opts in to overwrite. An + /// upload waits for any other upload to the same path on this connection + /// to finish first. pub async fn upload_from_local( &self, session_id: &str, @@ -234,73 +524,76 @@ impl SftpManager { remote_path: &str, overwrite: bool, ) -> Result<()> { - let mut sessions = self.sessions.lock().await; - let conn = sessions - .get_mut(session_id) - .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string()))?; - - if !overwrite - && conn - .sftp - .try_exists(remote_path) - .await - .map_err(|e| CoreError::Ssh(format!("check remote file: {e}")))? - { - return Err(CoreError::InvalidInput(format!( - "remote file already exists: {remote_path}" - ))); - } - - let mut local = tokio::fs::File::open(local_path).await?; - let mut remote = conn - .sftp - .create(remote_path) - .await - .map_err(|e| CoreError::Ssh(format!("create remote file: {e}")))?; + let conn = self.get(session_id).await?; + + conn.until_closed(async { + let _destination = conn.claim_destination(remote_path).await; + if !overwrite + && conn + .sftp + .try_exists(remote_path) + .await + .map_err(|e| CoreError::Ssh(format!("check remote file: {e}")))? + { + return Err(CoreError::InvalidInput(format!( + "remote file already exists: {remote_path}" + ))); + } - tokio::io::copy(&mut local, &mut remote).await?; - remote.flush().await?; - remote.shutdown().await?; - Ok(()) + let mut local = tokio::fs::File::open(local_path).await?; + // Without `overwrite` the create is exclusive, so of two uploads to + // one path the second is refused even when both passed the check + // above before either had created the file. + let created = if overwrite { + conn.sftp.create(remote_path).await + } else { + conn.sftp + .open_with_flags( + remote_path, + OpenFlags::CREATE | OpenFlags::EXCLUDE | OpenFlags::WRITE, + ) + .await + }; + let mut remote = match created { + Ok(remote) => remote, + Err(_) if !overwrite && conn.sftp.try_exists(remote_path).await.unwrap_or(false) => { + return Err(CoreError::InvalidInput(format!( + "remote file already exists: {remote_path}" + ))); + } + Err(e) => return Err(CoreError::Ssh(format!("create remote file: {e}"))), + }; + + tokio::io::copy(&mut local, &mut remote).await?; + remote.flush().await?; + remote.shutdown().await?; + Ok(()) + }) + .await } /// Create a directory. pub async fn create_dir(&self, session_id: &str, path: &str) -> Result<()> { - let mut sessions = self.sessions.lock().await; - let conn = sessions - .get_mut(session_id) - .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string()))?; - - conn.sftp - .create_dir(path) - .await - .map_err(|e| CoreError::Ssh(format!("mkdir: {e}"))) + self.request(session_id, "mkdir", |conn| async move { + conn.sftp.create_dir(path).await + }) + .await } /// Remove a file. pub async fn remove_file(&self, session_id: &str, path: &str) -> Result<()> { - let mut sessions = self.sessions.lock().await; - let conn = sessions - .get_mut(session_id) - .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string()))?; - - conn.sftp - .remove_file(path) - .await - .map_err(|e| CoreError::Ssh(format!("remove file: {e}"))) + self.request(session_id, "remove file", |conn| async move { + conn.sftp.remove_file(path).await + }) + .await } /// Remove a directory. pub async fn remove_dir(&self, session_id: &str, path: &str) -> Result<()> { - let mut sessions = self.sessions.lock().await; - let conn = sessions - .get_mut(session_id) - .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string()))?; - - conn.sftp - .remove_dir(path) - .await - .map_err(|e| CoreError::Ssh(format!("remove dir: {e}"))) + self.request(session_id, "remove dir", |conn| async move { + conn.sftp.remove_dir(path).await + }) + .await } /// Rename a file or directory. @@ -310,53 +603,29 @@ impl SftpManager { old_path: &str, new_path: &str, ) -> Result<()> { - let mut sessions = self.sessions.lock().await; - let conn = sessions - .get_mut(session_id) - .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string()))?; - - conn.sftp - .rename(old_path, new_path) - .await - .map_err(|e| CoreError::Ssh(format!("rename: {e}"))) + self.request(session_id, "rename", |conn| async move { + conn.sftp.rename(old_path, new_path).await + }) + .await } /// Get file metadata (size, permissions, modified time). pub async fn stat(&self, session_id: &str, path: &str) -> Result { - let mut sessions = self.sessions.lock().await; - let conn = sessions - .get_mut(session_id) - .ok_or_else(|| CoreError::SessionNotFound(session_id.to_string()))?; - - let attrs = conn - .sftp - .metadata(path) - .await - .map_err(|e| CoreError::Ssh(format!("stat: {e}")))?; - + let attrs = self + .request(session_id, "stat", |conn| async move { + conn.sftp.metadata(path).await + }) + .await?; let name = path.rsplit('/').next().unwrap_or(path).to_string(); - let modified = attrs - .modified() - .ok() - .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()) - .map(|d| d.as_secs() as i64); - Ok(SftpEntry { - name, - long_name: path.to_string(), - is_dir: attrs.is_dir(), - is_file: attrs.is_regular(), - is_symlink: attrs.is_symlink(), - size: attrs.len(), - modified, - permissions: attrs.permissions, - }) + Ok(SftpEntry::new(name, path.to_string(), &attrs)) } - /// Close an SFTP session. + /// Close an SFTP session. The entry leaves the map before anything is + /// sent, and closing the connection ends any transfer still running on it. pub async fn close(&self, session_id: &str) -> Result<()> { - let mut sessions = self.sessions.lock().await; - if let Some(conn) = sessions.remove(session_id) { - let _ = conn.sftp.close().await; + let removed = self.sessions.lock().await.remove(session_id); + if let Some(conn) = removed { + conn.close().await; } Ok(()) } @@ -370,15 +639,24 @@ impl SftpManager { /// /// The map is drained under the lock and the guard is released before any /// network work, so a channel stops being addressable the moment this is - /// called even if the remote never answers the close. + /// called even if the remote never answers the close. Transfers and + /// listings never hold the lock, so the drain does not wait on them, and + /// every connection is marked closed before any goodbye is sent. pub async fn close_all(&self) -> usize { - let drained: Vec = { + let drained: Vec> = { let mut sessions = self.sessions.lock().await; sessions.drain().map(|(_, conn)| conn).collect() }; let count = drained.len(); - for conn in &drained { - let _ = conn.sftp.close().await; + let goodbyes: Vec<_> = drained + .into_iter() + .map(|conn| { + conn.closed.raise(); + tokio::spawn(async move { conn.close().await }) + }) + .collect(); + for goodbye in goodbyes { + let _ = goodbye.await; } count } diff --git a/core/src/transport.rs b/core/src/transport.rs new file mode 100644 index 00000000..fe6c798d --- /dev/null +++ b/core/src/transport.rs @@ -0,0 +1,276 @@ +//! Ending a connection from outside russh. +//! +//! russh has no call that stops a session's I/O task. `Handle::disconnect` only +//! queues a message, and the task stops reading that queue while a channel is +//! waiting on the peer's window, so the goodbye can sit there for as long as +//! the peer stays stalled. Dropping every `Handle` does not end the task +//! either while a channel half is still alive. [`Transport`] closes the TCP +//! stream underneath russh instead: every read and write on it fails, and +//! russh ends the session, and every channel on it, on that error. + +use crate::{CoreError, Result}; +use russh::client::{Handle, Handler}; +use std::future::Future; +use std::io; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, MutexGuard}; +use std::task::{Context, Poll, Waker}; +use std::time::Duration; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; +use tokio::sync::watch; + +/// How long a polite goodbye may take before the transport is cut. A peer that +/// has stopped reading never lets the goodbye out. +pub(crate) const GOODBYE_GRACE: Duration = Duration::from_secs(5); + +const READ: usize = 0; +const WRITE: usize = 1; + +struct Shared { + cut: AtomicBool, + /// The wakers of a read and a write parked on the inner stream, so that a + /// cut ends them at once rather than when the peer next sends something. + wakers: Mutex<[Option; 2]>, + /// Set once russh has dropped the stream, which it does when its session + /// task ends, however that happens. + ended: watch::Sender, +} + +fn lock_wakers(shared: &Shared) -> MutexGuard<'_, [Option; 2]> { + // A waker slot is valid whatever a panicking holder was doing with it. + shared.wakers.lock().unwrap_or_else(|e| e.into_inner()) +} + +/// The means to end one connection's TCP stream, held by whoever owns the +/// connection. +#[derive(Clone)] +pub struct Transport { + shared: Arc, +} + +impl Transport { + /// Wrap `stream` so that the returned transport can sever it. + pub(crate) fn wrap(stream: S) -> (Severable, Transport) { + let shared = Arc::new(Shared { + cut: AtomicBool::new(false), + wakers: Mutex::new([None, None]), + ended: watch::Sender::new(false), + }); + let transport = Transport { + shared: shared.clone(), + }; + ( + Severable { + inner: stream, + shared, + }, + transport, + ) + } + + /// Fail every pending and later read and write on the stream. russh ends + /// the session on the error, which closes its channels and fails every + /// call still waiting on it. + pub fn cut(&self) { + self.shared.cut.store(true, Ordering::SeqCst); + let parked = std::mem::take(&mut *lock_wakers(&self.shared)); + for waker in parked.into_iter().flatten() { + waker.wake(); + } + } + + /// Resolves once russh has let go of the stream, which means the session + /// task has ended, whether by a goodbye, an error or a cut. + pub async fn ended(&self) { + let mut ended = self.shared.ended.subscribe(); + let _ = ended.wait_for(|ended| *ended).await; + } +} + +/// A stream that fails every operation once its [`Transport`] is cut. +pub(crate) struct Severable { + inner: S, + shared: Arc, +} + +impl Severable { + /// Park `cx`'s waker in `slot`, then fail if the stream has been cut. The + /// waker is parked first so that a cut landing between the two still + /// finds it and wakes the caller. + fn check(&self, slot: usize, cx: &Context<'_>) -> io::Result<()> { + { + let mut wakers = lock_wakers(&self.shared); + match &mut wakers[slot] { + Some(parked) if parked.will_wake(cx.waker()) => {} + parked => *parked = Some(cx.waker().clone()), + } + } + if self.shared.cut.load(Ordering::SeqCst) { + Err(io::Error::new( + io::ErrorKind::ConnectionAborted, + "the connection was closed", + )) + } else { + Ok(()) + } + } +} + +impl Drop for Severable { + fn drop(&mut self) { + self.shared.ended.send_replace(true); + } +} + +impl AsyncRead for Severable { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + if let Err(e) = this.check(READ, cx) { + return Poll::Ready(Err(e)); + } + Pin::new(&mut this.inner).poll_read(cx, buf) + } +} + +impl AsyncWrite for Severable { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + if let Err(e) = this.check(WRITE, cx) { + return Poll::Ready(Err(e)); + } + Pin::new(&mut this.inner).poll_write(cx, buf) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if let Err(e) = this.check(WRITE, cx) { + return Poll::Ready(Err(e)); + } + Pin::new(&mut this.inner).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if let Err(e) = this.check(WRITE, cx) { + return Poll::Ready(Err(e)); + } + Pin::new(&mut this.inner).poll_shutdown(cx) + } +} + +/// Say goodbye to the peer and wait for the session to end, for at most +/// [`GOODBYE_GRACE`], then cut the transport whether or not that happened. +/// Once this returns the connection is over, or ends on its next poll. +pub(crate) async fn goodbye(handle: &Handle, transport: &Transport) { + let polite = async { + let _ = handle + .disconnect(russh::Disconnect::ByApplication, "", "en") + .await; + transport.ended().await; + }; + let _ = tokio::time::timeout(GOODBYE_GRACE, polite).await; + transport.cut(); +} + +/// Raised once, when a connection is closed. Work already running on the +/// connection watches it and stops, so it cannot keep the connection busy, or +/// report success, after the close. +pub(crate) struct CloseSignal(watch::Sender); + +impl CloseSignal { + pub(crate) fn new() -> Self { + Self(watch::Sender::new(false)) + } + + pub(crate) fn raise(&self) { + self.0.send_replace(true); + } + + /// Resolves once the signal has been raised, at once if it already was. + pub(crate) async fn raised(&self) { + let mut raised = self.0.subscribe(); + let _ = raised.wait_for(|raised| *raised).await; + } + + /// Run `operation` until it finishes or the signal is raised, and answer + /// `closed()` in the second case. + pub(crate) async fn guard( + &self, + operation: impl Future>, + closed: impl FnOnce() -> CoreError, + ) -> Result { + tokio::select! { + result = operation => result, + () = self.raised() => Err(closed()), + } + } +} + +#[cfg(test)] +mod tests { + use super::Transport; + use std::time::Duration; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + #[tokio::test] + async fn a_cut_ends_a_read_that_is_already_waiting() { + let (near, _far) = tokio::io::duplex(64); + let (mut stream, transport) = Transport::wrap(near); + let read = tokio::spawn(async move { + let mut buf = [0u8; 8]; + stream.read(&mut buf).await + }); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(!read.is_finished()); + transport.cut(); + let result = tokio::time::timeout(Duration::from_secs(1), read) + .await + .expect("the read outlived the cut") + .expect("read task"); + assert_eq!( + result.expect_err("read after cut").kind(), + std::io::ErrorKind::ConnectionAborted + ); + } + + #[tokio::test] + async fn a_cut_ends_a_write_blocked_on_a_full_buffer() { + let (near, _far) = tokio::io::duplex(16); + let (mut stream, transport) = Transport::wrap(near); + let write = tokio::spawn(async move { stream.write_all(&[0u8; 1024]).await }); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(!write.is_finished(), "nothing reads the far end"); + transport.cut(); + let result = tokio::time::timeout(Duration::from_secs(1), write) + .await + .expect("the write outlived the cut") + .expect("write task"); + assert!(result.is_err()); + } + + #[tokio::test] + async fn ended_resolves_once_the_stream_is_dropped() { + let (near, _far) = tokio::io::duplex(16); + let (stream, transport) = Transport::wrap(near); + let ended = tokio::spawn({ + let transport = transport.clone(); + async move { transport.ended().await } + }); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(!ended.is_finished()); + drop(stream); + tokio::time::timeout(Duration::from_secs(1), ended) + .await + .expect("ended never resolved") + .expect("ended task"); + } +} diff --git a/core/tests/session_isolation.rs b/core/tests/session_isolation.rs new file mode 100644 index 00000000..8d427850 --- /dev/null +++ b/core/tests/session_isolation.rs @@ -0,0 +1,813 @@ +//! One misbehaving peer must hold up nothing but its own session. +//! +//! These run the SSH and SFTP managers against an in-process russh server on +//! loopback whose SFTP side never finishes a directory listing and whose shell +//! exits as soon as it starts, or stops reading its input, and check that other +//! sessions, a close and the drain a vault reset relies on all carry on +//! regardless, and that a close always ends the connection and frees its id. + +use clavyn_core::host::Host; +use clavyn_core::known_hosts::KnownHosts; +use clavyn_core::session::SessionManager; +use clavyn_core::sftp::SftpManager; +use russh::keys::{Algorithm, PrivateKey}; +use russh::server::{self, Auth, ChannelOpenHandle, Msg, Session}; +use russh::{Channel, ChannelId}; +use russh_sftp::protocol::{ + Attrs, Data, File, FileAttributes, Handle, Name, OpenFlags, Status, StatusCode, Version, +}; +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::Arc; +use std::time::Duration; +use tempfile::TempDir; +use tokio::net::TcpListener; +use tokio::sync::Mutex; + +const USER: &str = "deploy"; +const PASSWORD: &str = "correct password"; + +/// An SFTP server whose directories never end: every READDIR answers with more +/// entries and never with end-of-file. Its files live in memory, enough to tell +/// an exclusive create from one that replaces a file, and to see what two +/// writers left in one. +#[derive(Clone)] +struct EndlessDirectory { + /// Entries in each READDIR answer. + page: usize, + /// Pause before each READDIR answer. + page_delay: Duration, + /// Pause before each STAT answer. + stat_delay: Duration, + /// Pause before each WRITE answer. + write_delay: Duration, + /// Pause before each READ answer. + read_delay: Duration, + /// Handles closed so far. + closes: Arc, + /// Every file created so far, by path, with its contents. + files: Arc>>>, +} + +impl Default for EndlessDirectory { + fn default() -> Self { + Self { + page: 1, + page_delay: Duration::from_millis(5), + stat_delay: Duration::ZERO, + write_delay: Duration::ZERO, + read_delay: Duration::ZERO, + closes: Arc::default(), + files: Arc::default(), + } + } +} + +fn ok_status(id: u32) -> Status { + Status { + id, + status_code: StatusCode::Ok, + error_message: "Ok".into(), + language_tag: "en-US".into(), + } +} + +impl russh_sftp::server::Handler for EndlessDirectory { + type Error = StatusCode; + + fn unimplemented(&self) -> Self::Error { + StatusCode::OpUnsupported + } + + async fn init( + &mut self, + _version: u32, + _extensions: HashMap, + ) -> Result { + Ok(Version::new()) + } + + async fn opendir(&mut self, id: u32, _path: String) -> Result { + Ok(Handle { + id, + handle: "dir".into(), + }) + } + + async fn readdir(&mut self, id: u32, _handle: String) -> Result { + tokio::time::sleep(self.page_delay).await; + Ok(Name { + id, + files: vec![File::dummy("entry"); self.page], + }) + } + + async fn realpath(&mut self, id: u32, _path: String) -> Result { + Ok(Name { + id, + files: vec![File::dummy("/home/deploy")], + }) + } + + async fn close(&mut self, id: u32, _handle: String) -> Result { + self.closes.fetch_add(1, Ordering::SeqCst); + Ok(ok_status(id)) + } + + async fn open( + &mut self, + id: u32, + filename: String, + pflags: OpenFlags, + _attrs: FileAttributes, + ) -> Result { + let mut files = self.files.lock().expect("files"); + if pflags.contains(OpenFlags::EXCLUDE) && files.contains_key(&filename) { + return Err(StatusCode::Failure); + } + let contents = files.entry(filename.clone()).or_default(); + if pflags.contains(OpenFlags::TRUNCATE) { + contents.clear(); + } + Ok(Handle { + id, + handle: filename, + }) + } + + async fn write( + &mut self, + id: u32, + handle: String, + offset: u64, + data: Vec, + ) -> Result { + tokio::time::sleep(self.write_delay).await; + let mut files = self.files.lock().expect("files"); + let contents = files.entry(handle).or_default(); + let start = offset as usize; + if contents.len() < start + data.len() { + contents.resize(start + data.len(), 0); + } + contents[start..start + data.len()].copy_from_slice(&data); + Ok(ok_status(id)) + } + + async fn read( + &mut self, + id: u32, + handle: String, + offset: u64, + len: u32, + ) -> Result { + tokio::time::sleep(self.read_delay).await; + let files = self.files.lock().expect("files"); + let contents = files.get(&handle).ok_or(StatusCode::NoSuchFile)?; + let start = offset as usize; + if start >= contents.len() { + return Err(StatusCode::Eof); + } + let end = contents.len().min(start + len as usize); + Ok(Data { + id, + data: contents[start..end].to_vec(), + }) + } + + async fn stat(&mut self, id: u32, path: String) -> Result { + tokio::time::sleep(self.stat_delay).await; + if self.files.lock().expect("files").contains_key(&path) { + Ok(Attrs { + id, + attrs: FileAttributes::default(), + }) + } else { + Err(StatusCode::NoSuchFile) + } + } +} + +#[derive(Clone, Default)] +struct TestServer { + channels: Arc>>>, + /// Stop reading once the first input arrives, the way a hung peer does. + stall_on_input: bool, + /// Stop reading for this long when the first input arrives, then read + /// normally again, the way a peer that recovers does. + stall_once: Option<(Duration, Arc)>, + sftp: EndlessDirectory, + /// Connections the server has seen end. + ended: Arc, +} + +impl server::Handler for TestServer { + type Error = russh::Error; + + async fn auth_password(&mut self, user: &str, password: &str) -> Result { + Ok(if user == USER && password == PASSWORD { + Auth::Accept + } else { + Auth::reject() + }) + } + + async fn channel_open_session( + &mut self, + channel: Channel, + reply: ChannelOpenHandle, + _session: &mut Session, + ) -> Result<(), Self::Error> { + self.channels + .lock() + .expect("channels") + .insert(channel.id(), channel); + reply.accept().await; + Ok(()) + } + + /// The shell says one line and exits, unless this server stalls on input, + /// in which case it stays open. + async fn shell_request( + &mut self, + channel: ChannelId, + session: &mut Session, + ) -> Result<(), Self::Error> { + if self.stall_on_input || self.stall_once.is_some() { + return Ok(()); + } + session.data(channel, b"bye\r\n".to_vec())?; + session.eof(channel)?; + session.close(channel)?; + Ok(()) + } + + async fn data( + &mut self, + _channel: ChannelId, + _data: &[u8], + _session: &mut Session, + ) -> Result<(), Self::Error> { + if self.stall_on_input { + std::future::pending::<()>().await; + } + if let Some((stall, stalled)) = &self.stall_once { + if !stalled.swap(true, Ordering::SeqCst) { + tokio::time::sleep(*stall).await; + } + } + Ok(()) + } + + async fn subsystem_request( + &mut self, + channel_id: ChannelId, + name: &str, + session: &mut Session, + ) -> Result<(), Self::Error> { + let channel = self.channels.lock().expect("channels").remove(&channel_id); + match channel { + Some(channel) if name == "sftp" => { + session.channel_success(channel_id)?; + russh_sftp::server::run(channel.into_stream(), self.sftp.clone()).await; + } + _ => session.channel_failure(channel_id)?, + } + Ok(()) + } +} + +async fn start_server() -> u16 { + serve(TestServer::default()).await +} + +async fn serve(handler: TestServer) -> u16 { + let mut rng = getrandom::rand_core::UnwrapErr(getrandom::SysRng); + let host_key = PrivateKey::random(&mut rng, Algorithm::Ed25519).expect("host key"); + let config = Arc::new(server::Config { + keys: vec![host_key], + auth_rejection_time: Duration::from_millis(10), + auth_rejection_time_initial: Some(Duration::ZERO), + ..Default::default() + }); + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let port = listener.local_addr().expect("addr").port(); + tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + let config = config.clone(); + let handler = handler.clone(); + tokio::spawn(async move { + if let Ok(running) = server::run_stream(config, stream, handler.clone()).await { + let _ = running.await; + } + handler.ended.fetch_add(1, Ordering::SeqCst); + }); + } + }); + port +} + +fn host(port: u16) -> Host { + serde_json::from_value(serde_json::json!({ + "id": "00000000-0000-0000-0000-000000000001", + "label": "Loopback", + "hostname": "127.0.0.1", + "port": port, + "username": USER, + "auth": {"password": {"credential_key": "unused"}}, + "tags": [] + })) + .expect("host") +} + +fn known_hosts(dir: &TempDir) -> Arc> { + Arc::new(Mutex::new( + KnownHosts::load(dir.path().join("known_hosts.json")).expect("load known hosts"), + )) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_listing_that_never_ends_holds_up_no_other_sftp_session() { + let dir = tempfile::tempdir().expect("temp dir"); + let port = start_server().await; + let host = host(port); + let pins = known_hosts(&dir); + let manager = Arc::new(SftpManager::new()); + for id in ["stuck", "other"] { + manager + .connect(id.into(), &host, None, pins.clone(), None, None, Some(PASSWORD)) + .await + .expect("sftp connect"); + } + + let listing = tokio::spawn({ + let manager = manager.clone(); + async move { manager.list_dir("stuck", "/").await } + }); + tokio::time::sleep(Duration::from_millis(300)).await; + assert!(!listing.is_finished(), "the endless listing ended on its own"); + + let home = tokio::time::timeout(Duration::from_secs(5), manager.canonicalize("other", ".")) + .await + .expect("another session waited on the listing") + .expect("canonicalize"); + assert_eq!(home, "/home/deploy"); + + // The drain a vault reset runs must take both sessions out of reach + // without waiting for the listing. + let closed = tokio::time::timeout(Duration::from_secs(5), manager.close_all()) + .await + .expect("close_all waited on the listing"); + assert_eq!(closed, 2); + assert!(manager.list().await.is_empty()); + + // Closing the session ends the listing rather than leaving it running. + let result = tokio::time::timeout(Duration::from_secs(2), listing) + .await + .expect("the listing outlived its session") + .expect("listing task"); + assert!(result.is_err()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn closing_one_sftp_session_does_not_wait_on_its_listing() { + let dir = tempfile::tempdir().expect("temp dir"); + let port = start_server().await; + let manager = Arc::new(SftpManager::new()); + manager + .connect("stuck".into(), &host(port), None, known_hosts(&dir), None, None, Some(PASSWORD)) + .await + .expect("sftp connect"); + + let listing = tokio::spawn({ + let manager = manager.clone(); + async move { manager.list_dir("stuck", "/").await } + }); + tokio::time::sleep(Duration::from_millis(300)).await; + + tokio::time::timeout(Duration::from_secs(5), manager.close("stuck")) + .await + .expect("close waited on the listing") + .expect("close"); + assert!(manager.list().await.is_empty()); + let result = tokio::time::timeout(Duration::from_secs(2), listing) + .await + .expect("the listing outlived its session") + .expect("listing task"); + assert!(result.is_err()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_shell_that_exits_leaves_no_session_behind() { + let dir = tempfile::tempdir().expect("temp dir"); + let port = start_server().await; + let (closed_tx, mut closed_rx) = tokio::sync::mpsc::unbounded_channel::(); + let output = Arc::new(std::sync::Mutex::new(Vec::::new())); + let sessions = SessionManager::new( + { + let output = output.clone(); + Arc::new(move |_, data: &[u8]| output.lock().expect("output").extend_from_slice(data)) + }, + Arc::new(move |id: &str, _| { + let _ = closed_tx.send(id.to_string()); + }), + ); + + let claim = sessions.claim("shell").expect("claim"); + sessions + .create_ssh_session(&claim, &host(port), None, known_hosts(&dir), None, None, Some(PASSWORD), 80, 24) + .await + .expect("ssh session"); + drop(claim); + + let closed = tokio::time::timeout(Duration::from_secs(10), closed_rx.recv()) + .await + .expect("the close was never announced"); + assert_eq!(closed.as_deref(), Some("shell")); + assert_eq!(output.lock().expect("output").as_slice(), b"bye\r\n"); + + // The reader removes its entry after announcing the close; give it the + // moment that takes, then the id must be gone and free again. + let gone = async { + while sessions.contains("shell").await || sessions.is_claimed("shell") { + tokio::time::sleep(Duration::from_millis(20)).await; + } + }; + tokio::time::timeout(Duration::from_secs(10), gone) + .await + .expect("a session whose shell exited stayed registered"); + assert!(sessions.write("shell", b"ls\n").await.is_err()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_write_stuck_behind_a_peer_that_stopped_reading_ends_with_the_close() { + let dir = tempfile::tempdir().expect("temp dir"); + let port = serve(TestServer { + stall_on_input: true, + ..TestServer::default() + }) + .await; + let sessions = Arc::new(SessionManager::new(Arc::new(|_, _| {}), Arc::new(|_, _| {}))); + connect_shell(&sessions, port, &dir).await; + let writer = write_until_stuck(&sessions).await; + + tokio::time::timeout(Duration::from_secs(10), sessions.close("stuck")) + .await + .expect("close waited on the stuck write") + .expect("close"); + let result = tokio::time::timeout(Duration::from_secs(2), writer) + .await + .expect("the stuck write outlived the close") + .expect("writer task"); + assert!(result.is_err()); + assert!(!sessions.contains("stuck").await); +} + +/// Open a shell on `port` under the id "stuck". +async fn connect_shell(sessions: &SessionManager, port: u16, dir: &TempDir) { + let claim = sessions.claim("stuck").expect("claim"); + sessions + .create_ssh_session(&claim, &host(port), None, known_hosts(dir), None, None, Some(PASSWORD), 80, 24) + .await + .expect("ssh session"); +} + +/// Keep writing to "stuck" until the peer's window and the socket buffers are +/// full and a write stops returning, and hand back the task still writing. +async fn write_until_stuck( + sessions: &Arc, +) -> tokio::task::JoinHandle> { + let written = Arc::new(AtomicUsize::new(0)); + let writer = tokio::spawn({ + let sessions = sessions.clone(); + let written = written.clone(); + async move { + loop { + sessions.write("stuck", &[b'x'; 64 * 1024]).await?; + written.fetch_add(1, Ordering::Relaxed); + } + #[allow(unreachable_code)] + Ok::<(), clavyn_core::CoreError>(()) + } + }); + let stuck = async { + loop { + let before = written.load(Ordering::Relaxed); + tokio::time::sleep(Duration::from_millis(500)).await; + if before > 0 && written.load(Ordering::Relaxed) == before { + return; + } + } + }; + tokio::time::timeout(Duration::from_secs(20), stuck) + .await + .expect("writes never stalled"); + assert!(!writer.is_finished()); + writer +} + +/// A manager whose close announcements arrive on the returned receiver. +fn announcing_manager() -> (Arc, tokio::sync::mpsc::UnboundedReceiver) { + let (closed_tx, closed_rx) = tokio::sync::mpsc::unbounded_channel::(); + let sessions = SessionManager::new( + Arc::new(|_, _| {}), + Arc::new(move |id: &str, _| { + let _ = closed_tx.send(id.to_string()); + }), + ); + (Arc::new(sessions), closed_rx) +} + +/// Wait for the close of "stuck" to be announced and its id to come free. +async fn expect_released( + sessions: &SessionManager, + closed_rx: &mut tokio::sync::mpsc::UnboundedReceiver, +) { + let announced = tokio::time::timeout(Duration::from_secs(10), closed_rx.recv()) + .await + .expect("the close was never announced"); + assert_eq!(announced.as_deref(), Some("stuck")); + let free = async { + while sessions.is_claimed("stuck") { + tokio::time::sleep(Duration::from_millis(20)).await; + } + }; + tokio::time::timeout(Duration::from_secs(5), free) + .await + .expect("the id stayed claimed after the close"); + assert!(sessions.claim("stuck").is_ok(), "the id could not be reused"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn closing_a_session_whose_peer_stopped_reading_frees_its_id() { + let dir = tempfile::tempdir().expect("temp dir"); + let port = serve(TestServer { + stall_on_input: true, + ..TestServer::default() + }) + .await; + let (sessions, mut closed_rx) = announcing_manager(); + connect_shell(&sessions, port, &dir).await; + let writer = write_until_stuck(&sessions).await; + + // The goodbye cannot get past the unread input, so it is the cut after + // the grace that has to end the connection, and with it the reader. + tokio::time::timeout(Duration::from_secs(10), sessions.close("stuck")) + .await + .expect("close waited on the stuck write") + .expect("close"); + let _ = writer.await; + expect_released(&sessions, &mut closed_rx).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_peer_that_recovers_after_the_goodbye_timed_out_is_still_disconnected() { + let dir = tempfile::tempdir().expect("temp dir"); + let server = TestServer { + stall_once: Some((Duration::from_secs(8), Arc::new(AtomicBool::new(false)))), + ..TestServer::default() + }; + let ended = server.ended.clone(); + let port = serve(server).await; + let (sessions, mut closed_rx) = announcing_manager(); + connect_shell(&sessions, port, &dir).await; + let writer = write_until_stuck(&sessions).await; + + tokio::time::timeout(Duration::from_secs(10), sessions.close("stuck")) + .await + .expect("close waited on the stuck write") + .expect("close"); + let _ = writer.await; + expect_released(&sessions, &mut closed_rx).await; + + // Once the peer reads again it finds the connection gone, rather than a + // session nobody can address any more. + let gone = async { + while ended.load(Ordering::SeqCst) == 0 { + tokio::time::sleep(Duration::from_millis(50)).await; + } + }; + tokio::time::timeout(Duration::from_secs(15), gone) + .await + .expect("the connection outlived the close on the server's side"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_resize_never_waits_on_a_peer_that_stopped_reading() { + let dir = tempfile::tempdir().expect("temp dir"); + let port = serve(TestServer { + stall_on_input: true, + ..TestServer::default() + }) + .await; + let (sessions, mut closed_rx) = announcing_manager(); + connect_shell(&sessions, port, &dir).await; + let writer = write_until_stuck(&sessions).await; + + for step in 0..100u32 { + tokio::time::timeout(Duration::from_secs(1), sessions.resize("stuck", 80 + step, 24)) + .await + .expect("a resize waited on the stalled peer") + .expect("resize"); + } + + tokio::time::timeout(Duration::from_secs(10), sessions.close("stuck")) + .await + .expect("close waited on the stuck resize") + .expect("close"); + let _ = writer.await; + expect_released(&sessions, &mut closed_rx).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_listing_past_the_entry_cap_is_refused_and_its_handle_closed() { + let dir = tempfile::tempdir().expect("temp dir"); + let sftp = EndlessDirectory { + page: 2_000, + page_delay: Duration::ZERO, + ..EndlessDirectory::default() + }; + let closes = sftp.closes.clone(); + let port = serve(TestServer { + sftp, + ..TestServer::default() + }) + .await; + let manager = SftpManager::new(); + manager + .connect("big".into(), &host(port), None, known_hosts(&dir), None, None, Some(PASSWORD)) + .await + .expect("sftp connect"); + + let error = tokio::time::timeout(Duration::from_secs(60), manager.list_dir("big", "/")) + .await + .expect("the listing ran past the entry cap") + .expect_err("a directory past the entry cap was listed"); + assert!(error.to_string().contains("more than"), "{error}"); + assert_eq!(closes.load(Ordering::SeqCst), 1, "the directory handle was left open"); + manager.close_all().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn of_two_uploads_to_one_path_without_overwrite_only_one_creates_it() { + let dir = tempfile::tempdir().expect("temp dir"); + let local = dir.path().join("report.txt"); + std::fs::write(&local, b"report").expect("local file"); + // Slow existence checks let both uploads pass theirs before either one + // creates the file. + let port = serve(TestServer { + sftp: EndlessDirectory { + stat_delay: Duration::from_millis(200), + ..EndlessDirectory::default() + }, + ..TestServer::default() + }) + .await; + let manager = Arc::new(SftpManager::new()); + manager + .connect("up".into(), &host(port), None, known_hosts(&dir), None, None, Some(PASSWORD)) + .await + .expect("sftp connect"); + + let upload = || { + let manager = manager.clone(); + let local = local.clone(); + tokio::spawn(async move { + manager + .upload_from_local("up", &local, "/report.txt", false) + .await + }) + }; + let (first, second) = tokio::join!(upload(), upload()); + let results = [first.expect("upload task"), second.expect("upload task")]; + assert_eq!( + results.iter().filter(|result| result.is_ok()).count(), + 1, + "{results:?}" + ); + let refused = results + .iter() + .find_map(|result| result.as_ref().err()) + .expect("one upload refused"); + assert!(refused.to_string().contains("already exists"), "{refused}"); + manager.close_all().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn an_overwrite_upload_waits_for_one_already_writing_the_same_path() { + let dir = tempfile::tempdir().expect("temp dir"); + let large = dir.path().join("large.bin"); + let small = dir.path().join("small.bin"); + std::fs::write(&large, vec![b'a'; 1024 * 1024]).expect("large file"); + std::fs::write(&small, vec![b'b'; 100]).expect("small file"); + // Slow writes keep the large upload running while the small one starts. + let sftp = EndlessDirectory { + write_delay: Duration::from_millis(2), + ..EndlessDirectory::default() + }; + let files = sftp.files.clone(); + let port = serve(TestServer { + sftp, + ..TestServer::default() + }) + .await; + let manager = Arc::new(SftpManager::new()); + manager + .connect("up".into(), &host(port), None, known_hosts(&dir), None, None, Some(PASSWORD)) + .await + .expect("sftp connect"); + + let first = tokio::spawn({ + let manager = manager.clone(); + async move { manager.upload_from_local("up", &large, "/data.bin", true).await } + }); + let started = async { + while files + .lock() + .expect("files") + .get("/data.bin") + .is_none_or(|contents| contents.is_empty()) + { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }; + tokio::time::timeout(Duration::from_secs(10), started) + .await + .expect("the first upload never started writing"); + let second = manager.upload_from_local("up", &small, "/data.bin", true).await; + + first.await.expect("upload task").expect("first upload"); + second.expect("second upload"); + // The second upload replaced the first's file once it was complete, + // rather than truncating and writing it while the first still was. + let contents = files.lock().expect("files").get("/data.bin").cloned().unwrap_or_default(); + assert!( + contents == [b'b'; 100], + "the file holds {} bytes, {} of them from the first upload, not the second upload's 100", + contents.len(), + contents.iter().filter(|&&byte| byte == b'a').count() + ); + manager.close_all().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_download_waits_for_one_already_writing_the_same_local_file() { + let dir = tempfile::tempdir().expect("temp dir"); + let target = dir.path().join("data.bin"); + // Slow reads keep the large download writing while the small one starts. + let sftp = EndlessDirectory { + read_delay: Duration::from_millis(2), + ..EndlessDirectory::default() + }; + { + let mut files = sftp.files.lock().expect("files"); + files.insert("/large.bin".into(), vec![b'a'; 1024 * 1024]); + files.insert("/small.bin".into(), vec![b'b'; 100]); + } + let port = serve(TestServer { + sftp, + ..TestServer::default() + }) + .await; + // Two sessions, since the local disk is shared by all of them. + let manager = Arc::new(SftpManager::new()); + for id in ["first", "second"] { + manager + .connect(id.into(), &host(port), None, known_hosts(&dir), None, None, Some(PASSWORD)) + .await + .expect("sftp connect"); + } + + let first = tokio::spawn({ + let manager = manager.clone(); + let target = target.clone(); + async move { manager.download_to_local("first", "/large.bin", &target, true).await } + }); + let started = async { + while std::fs::metadata(&target).map_or(true, |meta| meta.len() == 0) { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }; + tokio::time::timeout(Duration::from_secs(10), started) + .await + .expect("the first download never started writing"); + // Reached through a second spelling of the same directory, which a + // comparison of the paths as given would take for another file. + std::fs::create_dir(dir.path().join("sub")).expect("sub dir"); + let same_file = dir.path().join("sub").join("..").join("data.bin"); + let second = manager + .download_to_local("second", "/small.bin", &same_file, true) + .await; + + first.await.expect("download task").expect("first download"); + second.expect("second download"); + // The second download replaced the first's file once it was complete, + // rather than truncating and writing it while the first still was. + let contents = std::fs::read(&target).expect("downloaded file"); + assert!( + contents == [b'b'; 100], + "the file holds {} bytes, {} of them from the first download, not the second download's 100", + contents.len(), + contents.iter().filter(|&&byte| byte == b'a').count() + ); + manager.close_all().await; +} diff --git a/core/tests/ssh_roundtrip.rs b/core/tests/ssh_roundtrip.rs index 2f5160a6..f506e755 100644 --- a/core/tests/ssh_roundtrip.rs +++ b/core/tests/ssh_roundtrip.rs @@ -170,7 +170,8 @@ async fn password_auth_connects_and_pins_the_host_key() { let session = connect(&password_host(port), None, pins.clone(), None, None, Some(PASSWORD)) .await - .expect("password auth"); + .expect("password auth") + .handle; // The session carries data both ways. let mut channel = session.channel_open_session().await.expect("open channel"); diff --git a/desktop/src-tauri/src/commands.rs b/desktop/src-tauri/src/commands.rs index 5bc46dcd..90963418 100644 --- a/desktop/src-tauri/src/commands.rs +++ b/desktop/src-tauri/src/commands.rs @@ -919,9 +919,11 @@ pub async fn session_write( session_id: String, data: Vec, ) -> ApiResult<()> { - // Try SSH session first - if state.sessions.list().await.contains(&session_id) { - return state.sessions.write(&session_id, &data).await.map_err(err); + // Try SSH session first. One lookup: an id no SSH session holds comes back + // as `SessionNotFound` and falls through to the local terminals. + match state.sessions.write(&session_id, &data).await { + Err(clavyn_core::CoreError::SessionNotFound(_)) => {} + result => return result.map_err(err), } // Try local terminal. The writer is taken out from under the map lock and // the write runs on the blocking pool: it blocks for as long as the @@ -978,13 +980,15 @@ pub async fn session_resize( ) -> ApiResult<()> { let cols = pty_dimension("cols", cols)?; let rows = pty_dimension("rows", rows)?; - // Try SSH session first - if state.sessions.list().await.contains(&session_id) { - return state - .sessions - .resize(&session_id, cols.into(), rows.into()) - .await - .map_err(err); + // Try SSH session first, falling through to the local terminals on an id + // no SSH session holds. + match state + .sessions + .resize(&session_id, cols.into(), rows.into()) + .await + { + Err(clavyn_core::CoreError::SessionNotFound(_)) => {} + result => return result.map_err(err), } // Try local terminal let locals = state.local_terminals.lock().await; @@ -1010,7 +1014,7 @@ pub async fn close_session(state: State<'_, Arc>, session_id: String) // stream. Releasing it here would drop the sink first and discard the tail, // which also happens on the reconnect path. // Try SSH session - if state.sessions.list().await.contains(&session_id) { + if state.sessions.contains(&session_id).await { return state.sessions.close(&session_id).await.map_err(err); } // Try local terminal diff --git a/desktop/src/stores/sftp.concurrency.test.ts b/desktop/src/stores/sftp.concurrency.test.ts index cc6e2a68..d51c7f68 100644 --- a/desktop/src/stores/sftp.concurrency.test.ts +++ b/desktop/src/stores/sftp.concurrency.test.ts @@ -2,6 +2,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { createPinia, setActivePinia } from "pinia"; import { getInvokeMock, setInvokeHandler } from "../test/setup"; import { useSftpStore } from "./sftp"; +import type { SftpEntry } from "../api"; import type { Host } from "../types"; const host: Host = { @@ -108,3 +109,90 @@ describe("SFTP connection ownership", () => { expect(calls("sftp_close").map(([, args]) => args.sessionId)).not.toContain(freshId); }); }); + +describe("SFTP directory listings", () => { + const entry = (name: string): SftpEntry => ({ + name, + long_name: name, + is_dir: true, + is_file: false, + is_symlink: false, + size: 0, + modified: null, + permissions: null, + }); + + // Each listing waits until the test answers it, so the order in which + // listings answer is the test's to choose. + async function connectedWithHeldListings() { + const pending = new Map void; reject: (e: Error) => void }>(); + setInvokeHandler("sftp_connect", () => undefined); + const sftp = useSftpStore(); + await sftp.connect(host, "SECRET", "deploy"); + setInvokeHandler("sftp_list_dir", ({ path }: { path: string }) => + new Promise((resolve, reject) => { pending.set(path, { resolve, reject }); })); + return { sftp, pending }; + } + + it("keeps the newer directory when an older listing answers last", async () => { + const { sftp, pending } = await connectedWithHeldListings(); + const older = sftp.listDir("/var/log"); + const newer = sftp.listDir("/etc"); + await vi.waitFor(() => expect(pending.size).toBe(2)); + + pending.get("/etc")!.resolve([entry("nginx")]); + await newer; + pending.get("/var/log")!.resolve([entry("journal")]); + await older; + + expect(sftp.currentPath).toBe("/etc"); + expect(sftp.entries.map(e => e.name)).toEqual(["nginx"]); + expect(sftp.loading).toBe(false); + }); + + it("stays loading until the latest listing answers", async () => { + const { sftp, pending } = await connectedWithHeldListings(); + const older = sftp.listDir("/var/log"); + const newer = sftp.listDir("/etc"); + await vi.waitFor(() => expect(pending.size).toBe(2)); + + pending.get("/var/log")!.resolve([entry("journal")]); + await older; + expect(sftp.loading).toBe(true); + expect(sftp.currentPath).toBe("/home/deploy"); + + pending.get("/etc")!.resolve([entry("nginx")]); + await newer; + expect(sftp.loading).toBe(false); + expect(sftp.currentPath).toBe("/etc"); + }); + + it("ignores the failure of a listing the user has moved on from", async () => { + const { sftp, pending } = await connectedWithHeldListings(); + const older = sftp.listDir("/root"); + const newer = sftp.listDir("/etc"); + await vi.waitFor(() => expect(pending.size).toBe(2)); + + pending.get("/etc")!.resolve([entry("nginx")]); + await newer; + pending.get("/root")!.reject(new Error("permission denied")); + await expect(older).resolves.toBeUndefined(); + + expect(sftp.error).toBeNull(); + expect(sftp.currentPath).toBe("/etc"); + }); + + it("publishes nothing from a listing that answers after a disconnect", async () => { + const { sftp, pending } = await connectedWithHeldListings(); + const listing = sftp.listDir("/etc"); + await vi.waitFor(() => expect(pending.size).toBe(1)); + + await sftp.disconnect(); + pending.get("/etc")!.resolve([entry("nginx")]); + await listing; + + expect(sftp.sessionId).toBeNull(); + expect(sftp.entries).toEqual([]); + expect(sftp.currentPath).toBe("/"); + }); +}); diff --git a/desktop/src/stores/sftp.ts b/desktop/src/stores/sftp.ts index ad775861..816428ae 100644 --- a/desktop/src/stores/sftp.ts +++ b/desktop/src/stores/sftp.ts @@ -21,6 +21,12 @@ export const useSftpStore = defineStore("sftp", () => { let connectPendingKey: string | null = null; let connectGeneration = 0; + // Listings of one session run side by side in the backend, so a slow one can + // answer after a newer one. Only the latest listing of the current session + // publishes its result, so an older answer never moves the browser back to + // a directory the user has already left. + let listGeneration = 0; + function clearPublishedConnection() { sessionId.value = null; connectedHost.value = null; @@ -130,18 +136,25 @@ export const useSftpStore = defineStore("sftp", () => { } async function listDir(path: string) { - if (!sessionId.value) return; + const id = sessionId.value; + if (!id) return; + const generation = ++listGeneration; + const current = () => generation === listGeneration && sessionId.value === id; error.value = null; loading.value = true; try { - entries.value = await api.sftpListDir(sessionId.value, path); + const listed = await api.sftpListDir(id, path); + if (!current()) return; + entries.value = listed; currentPath.value = path; selectedEntry.value = null; } catch (e) { + // A superseded listing's failure is about a directory the user has left. + if (!current()) return; error.value = String(e); throw e; } finally { - loading.value = false; + if (generation === listGeneration) loading.value = false; } } @@ -220,8 +233,8 @@ export const useSftpStore = defineStore("sftp", () => { fullPath, overwrite, ); - entries.value = await api.sftpListDir(sessionId.value, currentPath.value); - selectedEntry.value = null; + // Through `listDir`, so the refresh cannot overwrite a newer listing. + await listDir(currentPath.value); } catch (e) { error.value = String(e); throw e;