Skip to content

Commit 0ea0d31

Browse files
authored
fix(supervisor): restore canonical stdin after connection loss (#3852)
* fix(supervisor): restore canonical stdin after connection loss Probe idle SSH peers and enforce a receive deadline during transport I/O, including writes blocked by a stalled relay. Release the dead attachment's stdin lease through existing handler cleanup. Retry denied write intent on later ordinary input without displacing a healthy owner. Preserve explicit read-only, EOF and detach behavior, and discard control bytes retained while input ownership was denied. Cover half-open forwarding, blocked writes, healthy idle peers and competing reconnects through the production supervisor frame bridge and real SSH. Fixes #3648 Signed-off-by: Shiju <shiju@nvidia.com> * docs(skills): describe read-only reconnect input retry Explain what an openshell-cli user sees when automatic recovery reattaches before the supervisor closes the dead connection: the attachment reports read-only, later ordinary input retries stdin acquisition and prints `input enabled`, input typed while read-only is discarded, exit keys still detach, and an explicitly read-only viewer or a healthy owner is never affected. Signed-off-by: Shiju <shiju@nvidia.com> --------- Signed-off-by: Shiju <shiju@nvidia.com>
1 parent 33a8eac commit 0ea0d31

7 files changed

Lines changed: 787 additions & 11 deletions

File tree

‎crates/openshell-supervisor-process/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ libc = "0.2"
3939
[dev-dependencies]
4040
openshell-ocsf = { path = "../openshell-ocsf", features = ["test-support"] }
4141
tempfile = "3"
42+
tokio = { workspace = true, features = ["test-util"] }
4243

4344
[lints]
4445
workspace = true

‎crates/openshell-supervisor-process/src/ssh.rs‎

Lines changed: 61 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,14 @@ const NO_LOGIN_SHELL_ENV: (&str, &str) = ("OPENSHELL_NO_LOGIN_SHELL", "1");
2626
const MAIN_DETACH_PREFIX: u8 = 0x10;
2727
const MAIN_DETACH_KEY: u8 = 0x11;
2828
const MAIN_DETACH_EOF: u8 = 0x04;
29+
const SSH_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);
30+
const SSH_PEER_TIMEOUT: Duration = Duration::from_mins(1);
31+
32+
mod peer_stream;
33+
34+
#[cfg(test)]
35+
#[path = "ssh/reconnect_tests.rs"]
36+
mod reconnect_tests;
2937

3038
fn filter_main_detach_sequence(prefix_pending: &mut bool, data: &[u8]) -> (Vec<u8>, bool) {
3139
let mut forward = Vec::with_capacity(data.len() + usize::from(*prefix_pending));
@@ -70,6 +78,9 @@ fn ssh_server_init(
7078
let mut config = russh::server::Config {
7179
server_id: russh::SshId::Standard(Cow::Owned(format!("SSH-2.0-OpenShell_{VERSION}"))),
7280
auth_rejection_time: Duration::from_secs(1),
81+
// Peer replies keep idle attachments alive without relying on shell I/O.
82+
keepalive_interval: Some(SSH_KEEPALIVE_INTERVAL),
83+
keepalive_max: 3,
7384
..Default::default()
7485
};
7586
config.keys.push(host_key);
@@ -159,9 +170,15 @@ pub async fn run_ssh_server(
159170
let main_session = main_session.clone();
160171

161172
tokio::spawn(async move {
162-
if let Err(err) =
163-
handle_connection(stream, config, port_forward, boundary_exec, main_session)
164-
.await
173+
if let Err(err) = handle_connection(
174+
stream,
175+
config,
176+
port_forward,
177+
boundary_exec,
178+
main_session,
179+
SSH_PEER_TIMEOUT,
180+
)
181+
.await
165182
{
166183
ocsf_emit!(
167184
SshActivityBuilder::new(openshell_ocsf::ctx::ctx())
@@ -300,6 +317,7 @@ async fn handle_connection(
300317
port_forward: Arc<dyn openshell_isolation_interface::contract::BoundaryLoopbackConnector>,
301318
boundary_exec: Arc<dyn openshell_isolation_interface::contract::BoundaryExec>,
302319
main_session: Option<Arc<MainSession>>,
320+
peer_timeout: Duration,
303321
) -> Result<()> {
304322
// Access is gated by the Unix-socket filesystem permissions (root-only),
305323
// not by an application-level preface. The supervisor bridges the
@@ -316,9 +334,12 @@ async fn handle_connection(
316334
);
317335

318336
let handler = SshHandler::new(port_forward, boundary_exec, main_session);
337+
let stream = peer_stream::PeerStream::new(stream, peer_timeout);
319338
russh::server::run_stream(config, stream, handler)
320339
.await
321-
.map_err(|err| miette::miette!("ssh stream error: {err}"))?;
340+
.map_err(|err| miette::miette!("ssh stream error: {err}"))?
341+
.await
342+
.map_err(|err| miette::miette!("ssh session error: {err}"))?;
322343
Ok(())
323344
}
324345

@@ -339,6 +360,9 @@ struct ChannelState {
339360
main_input_owner: Option<u64>,
340361
main_attached: bool,
341362
main_read_only: bool,
363+
/// A writable attachment denied stdin may retry on later input. Explicit
364+
/// viewers and channels that sent EOF must never acquire a new lease.
365+
main_input_pending: bool,
342366
main_detach_prefix_pending: bool,
343367
main_output_task: Option<tokio::task::AbortHandle>,
344368
}
@@ -704,6 +728,7 @@ impl russh::server::Handler for SshHandler {
704728
Err(error) => (None, Some(error)),
705729
}
706730
};
731+
state.main_input_pending = warning.is_some();
707732
state.input_sender = input;
708733
state.main_detach_prefix_pending = false;
709734
let mut output = main_session.subscribe();
@@ -715,7 +740,7 @@ impl russh::server::Handler for SshHandler {
715740
.extended_data(
716741
channel,
717742
1,
718-
format!("openshell: {error}; attached read-only; press Ctrl-C or Ctrl-D to exit{line_ending}").into_bytes(),
743+
format!("openshell: {error}; attached read-only; retry input after the owner disconnects; press Ctrl-C or Ctrl-D to exit{line_ending}").into_bytes(),
719744
)
720745
.await;
721746
}
@@ -831,11 +856,35 @@ impl russh::server::Handler for SshHandler {
831856
.await;
832857
return Ok(());
833858
}
834-
let (forward, detach) = if state.main_attached {
859+
let denied_prefix = state.main_input_pending && state.main_detach_prefix_pending;
860+
let (mut forward, detach) = if state.main_attached {
835861
filter_main_detach_sequence(&mut state.main_detach_prefix_pending, data)
836862
} else {
837863
(data.to_vec(), false)
838864
};
865+
// Remember a denied viewer's prefix only to recognize split detach
866+
// sequences. It must not become process input when a later frame wins
867+
// the lease: that earlier keystroke was sent without write ownership.
868+
if denied_prefix && forward.first() == Some(&MAIN_DETACH_PREFIX) {
869+
forward.remove(0);
870+
}
871+
// A reconnect can precede detection of the old connection's failure.
872+
// Retry only on new input, using the same exclusive acquisition as a
873+
// fresh attachment. Detach keys retain their read-only behavior.
874+
if state.main_attached
875+
&& state.main_input_pending
876+
&& !state.main_read_only
877+
&& !detach
878+
&& !forward.is_empty()
879+
&& let Some(main_session) = self.main_session.as_ref()
880+
&& !main_session.finished()
881+
&& let Ok((owner, input)) = main_session.acquire_input()
882+
{
883+
state.main_input_owner = Some(owner);
884+
state.input_sender = Some(InputSender::Main(input));
885+
state.main_input_pending = false;
886+
session.extended_data(channel, 1, b"openshell: input enabled\r\n".to_vec())?;
887+
}
839888
let error = (!forward.is_empty())
840889
.then(|| state.input_sender.as_ref()?.send(forward).err())
841890
.flatten();
@@ -866,6 +915,7 @@ impl russh::server::Handler for SshHandler {
866915
main_session.release_input(owner);
867916
}
868917
state.input_sender.take();
918+
state.main_input_pending = false;
869919
state.main_detach_prefix_pending = false;
870920
} else {
871921
warn!("channel_eof on unknown channel {channel:?}");
@@ -1086,6 +1136,7 @@ impl SshHandler {
10861136
main_session.release_input(owner);
10871137
}
10881138
state.input_sender.take();
1139+
state.main_input_pending = false;
10891140
state.main_detach_prefix_pending = false;
10901141
if let Some(task) = state.main_output_task.take() {
10911142
task.abort();
@@ -1221,7 +1272,7 @@ mod tests {
12211272
use std::io::Write as _;
12221273
use std::process::{Command, Stdio};
12231274

1224-
struct AcceptAnyServerKey;
1275+
pub(super) struct AcceptAnyServerKey;
12251276

12261277
impl russh::client::Handler for AcceptAnyServerKey {
12271278
type Error = russh::Error;
@@ -1234,7 +1285,7 @@ mod tests {
12341285
}
12351286
}
12361287

1237-
struct TestLoopbackConnector;
1288+
pub(super) struct TestLoopbackConnector;
12381289

12391290
#[async_trait::async_trait]
12401291
impl openshell_isolation_interface::contract::BoundaryLoopbackConnector for TestLoopbackConnector {
@@ -1256,7 +1307,7 @@ mod tests {
12561307
}
12571308
}
12581309

1259-
struct RejectingExec;
1310+
pub(super) struct RejectingExec;
12601311

12611312
#[async_trait::async_trait]
12621313
impl openshell_isolation_interface::contract::BoundaryExec for RejectingExec {
@@ -1546,7 +1597,7 @@ mod tests {
15461597
assert_eq!(
15471598
String::from_utf8_lossy(&data),
15481599
format!(
1549-
"openshell: canonical main process already has an input owner; attached read-only; press Ctrl-C or Ctrl-D to exit{line_ending}"
1600+
"openshell: canonical main process already has an input owner; attached read-only; retry input after the owner disconnects; press Ctrl-C or Ctrl-D to exit{line_ending}"
15501601
)
15511602
);
15521603
}
Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
//! Bound SSH peer silence even while the session is blocked writing to its relay.
5+
6+
use std::future::Future;
7+
use std::io;
8+
use std::pin::Pin;
9+
use std::task::{Context, Poll};
10+
use std::time::Duration;
11+
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
12+
use tokio::time::{Instant, Sleep};
13+
14+
/// Only inbound bytes extend the deadline. SSH keepalive replies allow healthy
15+
/// idle clients to retain their connections without sending application input.
16+
pub(super) struct PeerStream<S> {
17+
stream: S,
18+
timeout: Duration,
19+
deadline: Pin<Box<Sleep>>,
20+
}
21+
22+
impl<S> PeerStream<S> {
23+
pub(super) fn new(stream: S, timeout: Duration) -> Self {
24+
Self {
25+
stream,
26+
timeout,
27+
deadline: Box::pin(tokio::time::sleep(timeout)),
28+
}
29+
}
30+
31+
fn poll_deadline(&mut self, cx: &mut Context<'_>) -> io::Result<()> {
32+
if self.deadline.as_mut().poll(cx).is_ready() {
33+
return Err(io::Error::new(
34+
io::ErrorKind::TimedOut,
35+
"SSH peer did not respond before the receive deadline",
36+
));
37+
}
38+
Ok(())
39+
}
40+
}
41+
42+
impl<S: AsyncRead + Unpin> AsyncRead for PeerStream<S> {
43+
fn poll_read(
44+
mut self: Pin<&mut Self>,
45+
cx: &mut Context<'_>,
46+
buf: &mut ReadBuf<'_>,
47+
) -> Poll<io::Result<()>> {
48+
self.poll_deadline(cx)?;
49+
let before = buf.filled().len();
50+
match Pin::new(&mut self.stream).poll_read(cx, buf) {
51+
Poll::Ready(Ok(())) => {
52+
if buf.filled().len() > before {
53+
let next = Instant::now() + self.timeout;
54+
self.deadline.as_mut().reset(next);
55+
}
56+
Poll::Ready(Ok(()))
57+
}
58+
other => other,
59+
}
60+
}
61+
}
62+
63+
impl<S: AsyncWrite + Unpin> AsyncWrite for PeerStream<S> {
64+
fn poll_write(
65+
mut self: Pin<&mut Self>,
66+
cx: &mut Context<'_>,
67+
buf: &[u8],
68+
) -> Poll<io::Result<usize>> {
69+
// Russh awaits packet writes outside its keepalive select. Polling here
70+
// makes an expired peer close even when the relay stops draining output.
71+
self.poll_deadline(cx)?;
72+
Pin::new(&mut self.stream).poll_write(cx, buf)
73+
}
74+
75+
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
76+
self.poll_deadline(cx)?;
77+
Pin::new(&mut self.stream).poll_flush(cx)
78+
}
79+
80+
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
81+
self.poll_deadline(cx)?;
82+
Pin::new(&mut self.stream).poll_shutdown(cx)
83+
}
84+
}
85+
86+
#[cfg(test)]
87+
mod tests {
88+
use super::*;
89+
use tokio::io::{AsyncReadExt, AsyncWriteExt};
90+
91+
#[tokio::test(start_paused = true)]
92+
async fn blocked_write_expires_at_receive_deadline() {
93+
let (stream, _peer) = tokio::io::duplex(1);
94+
let mut stream = PeerStream::new(stream, Duration::from_mins(1));
95+
let started = Instant::now();
96+
let error = stream.write_all(b"ab").await.unwrap_err();
97+
assert_eq!(error.kind(), io::ErrorKind::TimedOut);
98+
assert_eq!(started.elapsed(), Duration::from_mins(1));
99+
}
100+
101+
#[tokio::test(start_paused = true)]
102+
async fn only_received_bytes_extend_the_deadline() {
103+
let (stream, mut peer) = tokio::io::duplex(1);
104+
let mut stream = PeerStream::new(stream, Duration::from_mins(1));
105+
let started = Instant::now();
106+
tokio::time::advance(Duration::from_secs(40)).await;
107+
peer.write_all(b"r").await.unwrap();
108+
let mut reply = [0];
109+
stream.read_exact(&mut reply).await.unwrap();
110+
assert_eq!(&reply, b"r");
111+
tokio::time::advance(Duration::from_secs(40)).await;
112+
// A successful write must not renew peer liveness either.
113+
stream.write_all(b"a").await.unwrap();
114+
let error = stream.write_all(b"b").await.unwrap_err();
115+
assert_eq!(error.kind(), io::ErrorKind::TimedOut);
116+
assert_eq!(started.elapsed(), Duration::from_secs(100));
117+
}
118+
}

0 commit comments

Comments
 (0)