diff --git a/crates/tty7-core/src/daemon/install/ssh_ops.rs b/crates/tty7-core/src/daemon/install/ssh_ops.rs index 72be02e7..bf252a2d 100644 --- a/crates/tty7-core/src/daemon/install/ssh_ops.rs +++ b/crates/tty7-core/src/daemon/install/ssh_ops.rs @@ -155,8 +155,10 @@ fn is_not_found(msg: &str) -> bool { } async fn exec(conn: &Arc, cmd: &str) -> Result { + // The timeouts in `run` and `spawn_detached` drop this future mid-drain + // when the remote does not finish; the channel closes on that drop too. let mut channel = conn - .open_session_channel() + .open_command_channel() .await .map_err(|e| format!("could not open a command channel: {e}"))?; channel @@ -187,6 +189,7 @@ async fn exec(conn: &Arc, cmd: &str) -> Result log::debug!("ssh {key:?}: no remote shell integration"), } - self.probes.lock().unwrap().insert(key, probed.clone()); + // A probe the link would not carry says nothing about the + // shell. Remembered, its nothing would keep integration off + // this host for the rest of the run, fresh links included. + if probed.is_some() || conn.is_alive() { + self.probes.lock().unwrap().insert(key, probed.clone()); + } probed } }; @@ -603,7 +611,7 @@ const PROBE_TIMEOUT: Duration = Duration::from_secs(5); const PROBE_OUTPUT_LIMIT: usize = 8 * 1024; async fn probe_remote_shell(conn: &SshConnection) -> Option<(remote::RemoteShell, String)> { - let mut channel = conn.open_session_channel().await.ok()?; + let mut channel = conn.open_command_channel().await.ok()?; channel.exec(true, remote::PROBE_COMMAND).await.ok()?; let mut out: Vec = Vec::new(); @@ -627,7 +635,7 @@ async fn probe_remote_shell(conn: &SshConnection) -> Option<(remote::RemoteShell } async fn probe_remote_env(conn: &SshConnection) -> Option { - let mut channel = conn.open_session_channel().await.ok()?; + let mut channel = conn.open_command_channel().await.ok()?; channel .exec(true, remote_link::REMOTE_ENV_PROBE) .await @@ -671,37 +679,51 @@ fn sane_terminal_modes() -> Vec<(Pty, u32)> { #[cfg(test)] mod tests { + use super::test_support::{Exec, FakeSshd, base_spec}; use super::*; use crate::daemon::protocol::{SshAuthMode, SshProxy}; - fn base_spec() -> NativeSshSpec { - NativeSshSpec { - host: "h".into(), - port: 22, - user: "u".into(), - auth_mode: SshAuthMode::Auto, - identity_files: vec![], - agent_forward: false, - password: None, - key_passphrases: None, - proxy: SshProxy::None, - jump: None, - forwards: vec![], - keepalive_interval_s: None, - keepalive_count_max: None, - connect_timeout_s: None, - algorithms: Default::default(), - x11: false, - term: "xterm-256color".into(), - verify_host_keys: true, - skip_banner: false, - shell_integration: true, - login_script: vec![], - display_name: None, - profile_id: None, + /// A manager of its own, so nothing here touches the global cache. Its + /// runtime hosts the fake server and the connection under test. + fn manager() -> SshManager { + SshManager { + runtime: tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("build test runtime"), + conns: Mutex::new(HashMap::new()), + forwards: SshForwardRegistry::default(), + probes: Mutex::new(HashMap::new()), } } + #[test] + fn a_probe_the_link_would_not_carry_is_not_remembered() { + let mgr = manager(); + mgr.runtime.block_on(async { + let sshd = FakeSshd::connect(Exec::Hangs, Some(0)).await; + assert!(mgr.remote_bootstrap(&sshd.conn).await.is_none()); + assert!(!sshd.conn.is_alive(), "a refused session retires the link"); + assert!( + !mgr.probes.lock().unwrap().contains_key(sshd.conn.key()), + "the next session, on a fresh link, must ask again" + ); + }); + } + + #[test] + fn a_probe_answered_with_no_integration_is_remembered() { + let mgr = manager(); + mgr.runtime.block_on(async { + let sshd = FakeSshd::connect(Exec::Exits, None).await; + assert!(mgr.remote_bootstrap(&sshd.conn).await.is_none()); + assert!(sshd.conn.is_alive()); + assert!(mgr.probes.lock().unwrap().contains_key(sshd.conn.key())); + sshd.wait_for_closed(1).await; + assert_eq!(sshd.opened(), 1); + }); + } + #[test] fn connection_key_distinguishes_user_host_port_and_proxy() { let a = ConnectionKey::from_spec(&base_spec()); diff --git a/crates/tty7-core/src/daemon/ssh/session.rs b/crates/tty7-core/src/daemon/ssh/session.rs index 82483c57..dd25437d 100644 --- a/crates/tty7-core/src/daemon/ssh/session.rs +++ b/crates/tty7-core/src/daemon/ssh/session.rs @@ -148,14 +148,79 @@ pub async fn drive_channel( .window_change(u32::from(size.cols), u32::from(size.rows), pw, ph) .await; } - Some(ChannelCmd::Close) | None => { - let _ = channel.eof().await; - let _ = channel.close().await; - break; - } + Some(ChannelCmd::Close) | None => break, } } } + // Every way out closes, not only the pane's own Close. A pane whose + // reader is gone leaves a shell still running on the far side, and only + // a CHANNEL_CLOSE from here gives sshd that session back. After a close + // the server sent first, both of these put nothing on the wire. + let _ = channel.eof().await; + let _ = channel.close().await; +} + +/// A session channel for one command, closed whichever way its holder leaves. +/// +/// russh sends nothing for a `Channel` that is dropped. Its one close-on-drop +/// sits behind `into_stream`, and a command's output is read with `wait`, not +/// through a stream. Left to the holder, a `?` after the open or the timeout +/// wrapped around the whole exchange leaves the channel open on the server +/// for as long as the connection lives — and sshd counts every one of those +/// against `MaxSessions`, refusing the eleventh. +pub struct CommandChannel { + channel: Option>, + runtime: tokio::runtime::Handle, +} + +impl CommandChannel { + fn new(channel: Channel) -> Self { + Self { + channel: Some(channel), + // Taken here, on the runtime by construction, rather than looked + // up in `drop`, which runs wherever the holder is let go. + runtime: tokio::runtime::Handle::current(), + } + } +} + +impl std::ops::Deref for CommandChannel { + type Target = Channel; + + fn deref(&self) -> &Channel { + self.channel.as_ref().expect("held until drop") + } +} + +impl std::ops::DerefMut for CommandChannel { + fn deref_mut(&mut self) -> &mut Channel { + self.channel.as_mut().expect("held until drop") + } +} + +impl Drop for CommandChannel { + fn drop(&mut self) { + // `close` only queues the CHANNEL_CLOSE for the session task, which + // outlives this channel and writes it — the same best effort russh's + // own close-on-drop makes. A channel the server closed first is + // already out of the session's table, so this puts nothing more on + // the wire for it. + if let Some(channel) = self.channel.take() { + self.runtime.spawn(async move { + let _ = channel.close().await; + }); + } + } +} + +/// sshd's answer to a session open once the connection's `MaxSessions` are +/// all taken, by channels still running or by ones never closed. It says +/// nothing about this channel and everything about the connection. +fn refuses_every_session(e: &russh::Error) -> bool { + matches!( + e, + russh::Error::ChannelOpenFailure(russh::ChannelOpenFailure::ConnectFailed) + ) } pub struct SshConnection { @@ -202,7 +267,20 @@ impl SshConnection { } pub async fn open_session_channel(&self) -> Result, russh::Error> { - self.handle.lock().await.channel_open_session().await + let opened = self.handle.lock().await.channel_open_session().await; + // `is_alive` is what decides whether the cache hands this connection + // out again, and one that refuses sessions keeps refusing them until + // it is dropped: every retry on it would fail exactly this way. + if opened.as_ref().is_err_and(refuses_every_session) { + self.mark_dead(); + } + opened + } + + /// A session channel for one command, closed whichever way the caller + /// leaves it — see [`CommandChannel`]. + pub async fn open_command_channel(&self) -> Result { + self.open_session_channel().await.map(CommandChannel::new) } pub async fn open_direct_tcpip( @@ -309,6 +387,7 @@ impl SshConnection { #[cfg(test)] mod tests { + use super::super::test_support::{Exec, FakeSshd}; use super::*; use std::io::Read; @@ -355,6 +434,74 @@ mod tests { assert_eq!(n, 1); assert!(bridge.data_tx.try_send(vec![0xff]).is_ok()); } + + /// The pane's reader is gone before the shell's first byte, and the + /// shell keeps running: the one exit from `drive_channel` that used to + /// drop the channel with the far side none the wiser. + #[tokio::test] + async fn a_pane_that_is_gone_closes_the_shell_channel_behind_it() { + let sshd = FakeSshd::connect(Exec::Hangs, None).await; + let channel = sshd + .conn + .open_session_channel() + .await + .expect("open the shell channel"); + channel.request_shell(true).await.expect("shell request"); + + let BridgeEnds { + reader, + writer: _writer, + handle: _handle, + data_tx, + cmd_rx, + } = make_bridge(); + drop(reader); + + drive_channel(channel, data_tx, cmd_rx, sshd.conn.clone()).await; + sshd.wait_for_closed(1).await; + assert_eq!(sshd.opened(), 1); + } + + #[tokio::test] + async fn a_link_that_refuses_a_session_is_not_handed_out_again() { + let sshd = FakeSshd::connect(Exec::Hangs, Some(1)).await; + let _held = sshd + .conn + .open_session_channel() + .await + .expect("the one session the server allows"); + assert!(sshd.conn.is_alive()); + + let refused = sshd.conn.open_session_channel().await; + assert!( + matches!( + refused, + Err(russh::Error::ChannelOpenFailure( + russh::ChannelOpenFailure::ConnectFailed + )) + ), + "the fake answers like sshd at MaxSessions: {refused:?}" + ); + assert!( + !sshd.conn.is_alive(), + "the cache must dial afresh rather than retry this link" + ); + assert_eq!(sshd.refused(), 1); + } + + #[test] + fn only_a_refused_connect_means_the_link_is_spent() { + use russh::ChannelOpenFailure; + assert!(refuses_every_session(&russh::Error::ChannelOpenFailure( + ChannelOpenFailure::ConnectFailed + ))); + // A user whose account may not open sessions gets this on a fresh + // link too; retiring it would only dial again to be told the same. + assert!(!refuses_every_session(&russh::Error::ChannelOpenFailure( + ChannelOpenFailure::AdministrativelyProhibited + ))); + assert!(!refuses_every_session(&russh::Error::Disconnect)); + } } impl Drop for SshConnection { diff --git a/crates/tty7-core/src/daemon/ssh/test_support.rs b/crates/tty7-core/src/daemon/ssh/test_support.rs new file mode 100644 index 00000000..70cc2ceb --- /dev/null +++ b/crates/tty7-core/src/daemon/ssh/test_support.rs @@ -0,0 +1,234 @@ +//! An SSH server in the test process, for what `FakeRemote` cannot see. +//! +//! The install layer's fake answers `RemoteOps` calls without a wire, so +//! whether the channels behind those calls are ever closed — what sshd's +//! `MaxSessions` counts — was outside every test. This server takes one +//! connection, accepts session channels up to a limit, answers `exec` the way +//! the test asks, and counts what the client opened and closed. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; + +use russh::keys::PrivateKey; +use russh::keys::ssh_key::private::Ed25519Keypair; +use russh::server::{self, Auth, ChannelOpenHandle, Session}; +use russh::{Channel, ChannelId, ChannelOpenFailure}; + +use crate::daemon::protocol::{NativeSshSpec, SshAuthMode, SshProxy}; + +use super::broker::PromptBroker; +use super::forward::RemoteForwardTable; +use super::handler::ClientHandler; +use super::{ConnectionKey, SshConnection}; + +/// How the server answers an `exec` request. +#[derive(Clone, Copy)] +pub(crate) enum Exec { + /// Success, a line of output, exit status 0, EOF and CLOSE: a command + /// that ran and finished, after which the server closes first. + Exits, + /// Success and nothing more: a command that never finishes. The server + /// never closes, so only the client can. + Hangs, +} + +#[derive(Default)] +struct Counts { + opened: AtomicUsize, + closed: AtomicUsize, + refused: AtomicUsize, +} + +struct Sshd { + exec: Exec, + max_sessions: Option, + counts: Arc, +} + +impl server::Handler for Sshd { + type Error = russh::Error; + + async fn auth_none(&mut self, _user: &str) -> Result { + Ok(Auth::Accept) + } + + async fn channel_open_session( + &mut self, + _channel: Channel, + reply: ChannelOpenHandle, + _session: &mut Session, + ) -> Result<(), Self::Error> { + let open = self + .counts + .opened + .load(Ordering::SeqCst) + .saturating_sub(self.counts.closed.load(Ordering::SeqCst)); + if self.max_sessions.is_some_and(|max| open >= max) { + self.counts.refused.fetch_add(1, Ordering::SeqCst); + // sshd's answer once MaxSessions are all taken. + reply.reject(ChannelOpenFailure::ConnectFailed).await; + return Ok(()); + } + self.counts.opened.fetch_add(1, Ordering::SeqCst); + reply.accept().await; + Ok(()) + } + + async fn exec_request( + &mut self, + channel: ChannelId, + _command: &[u8], + session: &mut Session, + ) -> Result<(), Self::Error> { + session.channel_success(channel)?; + if let Exec::Exits = self.exec { + session.data(channel, &b"ok\n"[..])?; + session.exit_status_request(channel, 0)?; + session.eof(channel)?; + session.close(channel)?; + } + Ok(()) + } + + async fn shell_request( + &mut self, + channel: ChannelId, + session: &mut Session, + ) -> Result<(), Self::Error> { + session.channel_success(channel)?; + session.data(channel, &b"$ "[..])?; + Ok(()) + } + + async fn channel_close( + &mut self, + _channel: ChannelId, + _session: &mut Session, + ) -> Result<(), Self::Error> { + self.counts.closed.fetch_add(1, Ordering::SeqCst); + Ok(()) + } +} + +/// One connection to a server in this process, and what the server counted. +pub(crate) struct FakeSshd { + pub(crate) conn: Arc, + counts: Arc, +} + +impl FakeSshd { + pub(crate) async fn connect(exec: Exec, max_sessions: Option) -> FakeSshd { + let counts = Arc::new(Counts::default()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind a loopback port"); + let addr = listener.local_addr().expect("bound address"); + + let mut config = server::Config::default(); + config.inactivity_timeout = None; + config + .keys + .push(PrivateKey::from(Ed25519Keypair::from_seed(&[7; 32]))); + let handler = Sshd { + exec, + max_sessions, + counts: Arc::clone(&counts), + }; + tokio::spawn(async move { + let (socket, _) = listener.accept().await.expect("accept the test client"); + let running = server::run_stream(Arc::new(config), socket, handler) + .await + .expect("server handshake"); + let _ = running.await; + }); + + let spec = spec_for(addr.port()); + let remote_forwards = RemoteForwardTable::default(); + let handler = ClientHandler { + host: spec.host.clone(), + port: spec.port, + verify_host_keys: false, + skip_banner: true, + broker: PromptBroker::new(Box::new(|_| true)), + remote_forwards: remote_forwards.clone(), + }; + let mut handle = + russh::client::connect(Arc::new(russh::client::Config::default()), addr, handler) + .await + .expect("client handshake"); + let auth = handle + .authenticate_none("tester") + .await + .expect("auth round trip"); + assert!(auth.success(), "the fake accepts everyone"); + let conn = SshConnection::new(handle, ConnectionKey::from_spec(&spec), remote_forwards); + FakeSshd { conn, counts } + } + + /// Session channels the server accepted. + pub(crate) fn opened(&self) -> usize { + self.counts.opened.load(Ordering::SeqCst) + } + + /// CHANNEL_CLOSEs the server received. + pub(crate) fn closed(&self) -> usize { + self.counts.closed.load(Ordering::SeqCst) + } + + /// Session opens the server refused at its limit. + pub(crate) fn refused(&self) -> usize { + self.counts.refused.load(Ordering::SeqCst) + } + + /// Waits for the server to have received `n` closes. A close sent on + /// drop is queued for the client's session task, so it reaches the server + /// a moment after the holder is gone, never before. + pub(crate) async fn wait_for_closed(&self, n: usize) { + let deadline = Instant::now() + Duration::from_secs(5); + while self.closed() < n { + assert!( + Instant::now() < deadline, + "the server received {} of {n} closes", + self.closed() + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + } +} + +pub(crate) fn base_spec() -> NativeSshSpec { + NativeSshSpec { + host: "h".into(), + port: 22, + user: "u".into(), + auth_mode: SshAuthMode::Auto, + identity_files: vec![], + agent_forward: false, + password: None, + key_passphrases: None, + proxy: SshProxy::None, + jump: None, + forwards: vec![], + keepalive_interval_s: None, + keepalive_count_max: None, + connect_timeout_s: None, + algorithms: Default::default(), + x11: false, + term: "xterm-256color".into(), + verify_host_keys: true, + skip_banner: false, + shell_integration: true, + login_script: vec![], + display_name: None, + profile_id: None, + } +} + +fn spec_for(port: u16) -> NativeSshSpec { + let mut spec = base_spec(); + spec.host = "127.0.0.1".into(); + spec.port = port; + spec.verify_host_keys = false; + spec +}