diff --git a/Cargo.lock b/Cargo.lock index 7539687..3f2a36f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -436,7 +436,7 @@ dependencies = [ [[package]] name = "dosh" -version = "1.0.0-rc22" +version = "1.0.0-rc23" dependencies = [ "anyhow", "base64", diff --git a/Cargo.toml b/Cargo.toml index dd88e34..7f82cc0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "dosh" -version = "1.0.0-rc22" +version = "1.0.0-rc23" edition = "2024" license = "MIT" diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index 7042c7a..6852348 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -36,7 +36,10 @@ use dosh::transport::{ DoshTransport, SessionEvent, SessionRole, SessionTransportConfig, TransportConfig, TransportEvent, }; -use dosh::udp::{is_transient_udp_error, is_transient_udp_send_error}; +use dosh::udp::{ + is_transient_udp_error, is_transient_udp_send_error, recv_udp_retrying_transient, + send_udp_retrying_transient, +}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use std::collections::BTreeMap; @@ -3789,14 +3792,11 @@ async fn try_native_auth( hello: hello.clone(), })?, )?; - socket.send_to(&packet, addr).await?; + let auth_timeout = Duration::from_millis(config.native_auth_timeout_ms.max(1)); + send_udp_retrying_transient(socket, &packet, addr, auth_timeout).await?; let mut buf = vec![0u8; 65535]; - let (n, _) = tokio::time::timeout( - Duration::from_millis(config.native_auth_timeout_ms.max(1)), - socket.recv_from(&mut buf), - ) - .await??; + let (n, _) = recv_udp_retrying_transient(socket, &mut buf, auth_timeout).await?; let packet = protocol::decode(&buf[..n])?; if packet.header.kind != PacketKind::NativeServerHello { if packet.header.kind == PacketKind::AttachReject { @@ -3859,13 +3859,9 @@ async fn try_native_auth( CLIENT_TO_SERVER, &auth_body, )?; - socket.send_to(&packet, addr).await?; + send_udp_retrying_transient(socket, &packet, addr, auth_timeout).await?; - let (n, _) = tokio::time::timeout( - Duration::from_millis(config.native_auth_timeout_ms.max(1)), - socket.recv_from(&mut buf), - ) - .await??; + let (n, _) = recv_udp_retrying_transient(socket, &mut buf, auth_timeout).await?; let packet = protocol::decode(&buf[..n])?; if packet.header.kind != PacketKind::NativeAuthOk { if packet.header.kind == PacketKind::AttachReject { @@ -3947,14 +3943,11 @@ async fn try_native_auth_check( hello: hello.clone(), })?, )?; - socket.send_to(&packet, addr).await?; + let auth_timeout = Duration::from_millis(config.native_auth_timeout_ms.max(1)); + send_udp_retrying_transient(socket, &packet, addr, auth_timeout).await?; let mut buf = vec![0u8; 65535]; - let (n, _) = tokio::time::timeout( - Duration::from_millis(config.native_auth_timeout_ms.max(1)), - socket.recv_from(&mut buf), - ) - .await??; + let (n, _) = recv_udp_retrying_transient(socket, &mut buf, auth_timeout).await?; let packet = protocol::decode(&buf[..n])?; if packet.header.kind != PacketKind::NativeServerHello { if packet.header.kind == PacketKind::AttachReject { @@ -4016,13 +4009,9 @@ async fn try_native_auth_check( CLIENT_TO_SERVER, &protocol::to_body(&NativeUserAuthBody { auth })?, )?; - socket.send_to(&packet, addr).await?; + send_udp_retrying_transient(socket, &packet, addr, auth_timeout).await?; - let (n, _) = tokio::time::timeout( - Duration::from_millis(config.native_auth_timeout_ms.max(1)), - socket.recv_from(&mut buf), - ) - .await??; + let (n, _) = recv_udp_retrying_transient(socket, &mut buf, auth_timeout).await?; let packet = protocol::decode(&buf[..n])?; if packet.header.kind != PacketKind::NativeAuthCheckOk { if packet.header.kind == PacketKind::AttachReject { @@ -4182,7 +4171,7 @@ async fn bootstrap_attach( let body = protocol::to_body(&req)?; let packet = protocol::encode_plain(PacketKind::BootstrapAttachRequest, [0u8; 16], 1, 0, &body)?; - socket.send_to(&packet, addr).await?; + send_udp_retrying_transient(socket, &packet, addr, Duration::from_secs(5)).await?; let expected_key_id = protocol::session_key_id(&bootstrap.session_key); let ok = recv_response_until(socket, Duration::from_secs(5), |packet| { if packet.header.kind == PacketKind::AttachReject && packet.header.conn_id == [0u8; 16] { @@ -4262,7 +4251,7 @@ async fn try_ticket_attach( 0, &protocol::to_body(&envelope)?, )?; - socket.send_to(&packet, addr).await?; + send_udp_retrying_transient(socket, &packet, addr, Duration::from_millis(700)).await?; let ok = recv_response_until(socket, Duration::from_millis(700), |packet| { if packet.header.kind == PacketKind::AttachReject && packet.header.conn_id == [0u8; 16] { @@ -4338,7 +4327,7 @@ async fn try_resume_fresh_process( CLIENT_TO_SERVER, &body, )?; - socket.send_to(&packet, addr).await?; + send_udp_retrying_transient(socket, &packet, addr, Duration::from_millis(700)).await?; let frame = recv_response_until(socket, Duration::from_millis(700), |packet| { if !is_resume_response_for_client(&packet, cached.client_id) { return Ok(None); @@ -4393,7 +4382,7 @@ async fn try_live_resume( &body, )?; *send_seq += 1; - socket.send_to(&packet, addr).await?; + send_udp_retrying_transient(socket, &packet, addr, Duration::from_millis(700)).await?; let frame = recv_response_until(socket, Duration::from_millis(700), |packet| { if !is_resume_response_for_client(&packet, cached.client_id) { return Ok(None); @@ -4440,14 +4429,9 @@ where return Err(anyhow!("timed out waiting for server response")); } let remaining = wait - elapsed; - let (n, _) = match tokio::time::timeout(remaining, socket.recv_from(&mut buf)) + let (n, _) = recv_udp_retrying_transient(socket, &mut buf, remaining) .await - .context("timed out waiting for server response")? - { - Ok(value) => value, - Err(err) if is_transient_udp_error(&err) => continue, - Err(err) => return Err(err.into()), - }; + .context("timed out waiting for server response")?; let Ok(packet) = protocol::decode(&buf[..n]) else { continue; }; diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index 71f948a..6e9feb5 100644 --- a/src/bin/dosh-server.rs +++ b/src/bin/dosh-server.rs @@ -361,6 +361,7 @@ impl NativeAuthRateLimiter { struct Session { pty: Option, parser: vt100::Parser, + restored_screen: Option, clients: HashMap<[u8; 16], ClientState>, output_seq: u64, recent: VecDeque>, @@ -387,6 +388,14 @@ struct Session { last_screen_persist_at: Instant, } +#[derive(Clone)] +struct RestoredScreenSnapshot { + cols: u16, + rows: u16, + output_seq: u64, + snapshot: Vec, +} + #[derive(Clone)] struct ClientState { endpoint: SocketAddr, @@ -512,6 +521,7 @@ impl ServerState { Session { pty, parser: vt100::Parser::new(rows.max(1), cols.max(1), self.config.scrollback), + restored_screen: None, clients: HashMap::new(), output_seq: 0, recent: VecDeque::with_capacity(self.config.scrollback), @@ -641,15 +651,23 @@ impl ServerState { .unwrap_or((80, 24)); let mut parser = vt100::Parser::new(rows, cols, self.config.scrollback); let mut output_seq = 0; + let mut restored_screen = None; if let Some(saved) = saved { parser.process(&saved.snapshot); output_seq = saved.output_seq; + restored_screen = Some(RestoredScreenSnapshot { + cols, + rows, + output_seq, + snapshot: saved.snapshot, + }); } self.sessions.insert( name.clone(), Session { pty: Some(pty), parser, + restored_screen, clients: HashMap::new(), output_seq, recent: VecDeque::with_capacity(self.config.scrollback), @@ -762,6 +780,19 @@ fn apply_terminal_size( Ok(()) } +fn attach_snapshot(session: &Session, cols: u16, rows: u16) -> Vec { + let cols = cols.max(1); + let rows = rows.max(1); + if let Some(restored) = &session.restored_screen + && restored.output_seq == session.output_seq + && restored.cols == cols + && restored.rows == rows + { + return restored.snapshot.clone(); + } + screen_snapshot(session.parser.screen()) +} + fn mode_allows_terminal_updates(mode: &str) -> bool { mode != "view-only" && mode != "forward-only" } @@ -1189,7 +1220,7 @@ async fn handle_native_user_auth( apply_terminal_size(session.pty.as_ref(), &mut session.parser, cols, rows)?; } let client_id = crypto::random_16(); - let snapshot = screen_snapshot(session.parser.screen()); + let snapshot = attach_snapshot(session, cols, rows); let output_seq = session.output_seq; locked.insert_client( &session_name, @@ -1337,7 +1368,7 @@ async fn handle_bootstrap_attach( )?; } let client_id = crypto::random_16(); - let snapshot = screen_snapshot(session.parser.screen()); + let snapshot = attach_snapshot(session, req.cols, req.rows); let output_seq = session.output_seq; let bootstrap_session = req.bootstrap.session.clone(); locked.insert_client( @@ -1473,7 +1504,7 @@ async fn handle_ticket_attach( )?; } let client_id = crypto::random_16(); - let snapshot = screen_snapshot(session.parser.screen()); + let snapshot = attach_snapshot(session, req.cols, req.rows); let output_seq = session.output_seq; let ticket_session = req.session.clone(); locked.insert_client( @@ -1586,20 +1617,23 @@ async fn handle_resume( .sessions .get_mut(&req.session) .ok_or_else(|| anyhow!("unknown session"))?; - let client = session - .clients - .get_mut(&packet.header.conn_id) - .ok_or_else(|| anyhow!("unknown client"))?; - if !client.replay.accept(packet.header.seq) { - return Ok(()); - } - client.endpoint = peer; - client.last_acked = req.last_rendered_seq; - client.cols = req.cols; - client.rows = req.rows; - client.last_seen = Instant::now(); - client.send_seq += 1; - if mode_allows_terminal_updates(&client.mode) { + let (send_seq, allow_terminal_updates) = { + let client = session + .clients + .get_mut(&packet.header.conn_id) + .ok_or_else(|| anyhow!("unknown client"))?; + if !client.replay.accept(packet.header.seq) { + return Ok(()); + } + client.endpoint = peer; + client.last_acked = req.last_rendered_seq; + client.cols = req.cols; + client.rows = req.rows; + client.last_seen = Instant::now(); + client.send_seq += 1; + (client.send_seq, mode_allows_terminal_updates(&client.mode)) + }; + if allow_terminal_updates { apply_terminal_size( session.pty.as_ref(), &mut session.parser, @@ -1607,8 +1641,8 @@ async fn handle_resume( req.rows, )?; } - let snapshot = screen_snapshot(session.parser.screen()); - (client.send_seq, session.output_seq, snapshot) + let snapshot = attach_snapshot(session, req.cols, req.rows); + (send_seq, session.output_seq, snapshot) }; let frame = Frame { session: session_name, @@ -3454,6 +3488,7 @@ async fn broadcast_output( .get_mut(&output.session) .ok_or_else(|| anyhow!("unknown session"))?; session.parser.process(&output.bytes); + session.restored_screen = None; session.output_seq += 1; let output_seq = session.output_seq; session.recent.push_back(output.bytes.clone()); @@ -4049,6 +4084,34 @@ mod tests { assert_eq!(ServerConfig::default().client_timeout_secs, 2_592_000); } + #[test] + fn attach_snapshot_prefers_restored_screen_for_same_size_reattach() { + let restored = b"\x1b[?1049lRESTORED_SCREEN_MARKER".to_vec(); + let mut session = Session { + pty: None, + parser: vt100::Parser::new(24, 80, 100), + restored_screen: Some(RestoredScreenSnapshot { + cols: 80, + rows: 24, + output_seq: 7, + snapshot: restored.clone(), + }), + clients: HashMap::new(), + output_seq: 7, + recent: VecDeque::new(), + empty_since: Some(Instant::now()), + holder_control: None, + persistent: true, + bytes_since_persist: 0, + last_persisted_seq: 7, + last_screen_persist_at: Instant::now(), + }; + + session.parser.set_size(24, 80); + assert_eq!(attach_snapshot(&session, 80, 24), restored); + assert_ne!(attach_snapshot(&session, 100, 30), restored); + } + #[tokio::test] async fn unknown_resume_reject_keeps_client_id() { let (pty_tx, _rx) = mpsc::unbounded_channel(); @@ -4399,6 +4462,7 @@ mod tests { Session { pty: None, parser: vt100::Parser::new(24, 80, 100), + restored_screen: None, clients: HashMap::from([(client_id, client)]), output_seq: 0, recent: VecDeque::new(), @@ -4463,6 +4527,7 @@ mod tests { Session { pty: None, parser: vt100::Parser::new(24, 80, 100), + restored_screen: None, clients: HashMap::from([(client_id, client)]), output_seq: 0, recent: VecDeque::new(), @@ -4539,6 +4604,7 @@ mod tests { Session { pty: None, parser: vt100::Parser::new(24, 80, 100), + restored_screen: None, clients: HashMap::from([(client_id, client)]), output_seq: 0, recent: VecDeque::new(), @@ -4679,6 +4745,7 @@ mod tests { Session { pty: None, parser: vt100::Parser::new(24, 80, 100), + restored_screen: None, clients: HashMap::from([( client_id, ClientState { @@ -4776,6 +4843,7 @@ mod tests { Session { pty: None, parser: vt100::Parser::new(24, 80, 100), + restored_screen: None, clients: HashMap::from([( client_id, ClientState { diff --git a/src/client.rs b/src/client.rs index 7edce19..36be25c 100644 --- a/src/client.rs +++ b/src/client.rs @@ -13,6 +13,7 @@ use crate::protocol::{ }; use crate::ssh_agent; use crate::transport::{DoshTransport, SessionRole, SessionTransportConfig, TransportConfig}; +use crate::udp::{recv_udp_retrying_transient, send_udp_retrying_transient}; use anyhow::{Context, Result, anyhow, bail}; use std::net::{SocketAddr, ToSocketAddrs}; use std::path::PathBuf; @@ -214,10 +215,10 @@ impl DoshClientBuilder { hello: hello.clone(), })?, )?; - socket.send_to(&packet, peer_addr).await?; + send_udp_retrying_transient(&socket, &packet, peer_addr, timeout).await?; let mut buf = vec![0u8; 65535]; - let (n, _) = tokio::time::timeout(timeout, socket.recv_from(&mut buf)).await??; + let (n, _) = recv_udp_retrying_transient(&socket, &mut buf, timeout).await?; let packet = protocol::decode(&buf[..n])?; if packet.header.kind != PacketKind::NativeServerHello { if packet.header.kind == PacketKind::AttachReject { @@ -261,8 +262,8 @@ impl DoshClientBuilder { CLIENT_TO_SERVER, &protocol::to_body(&NativeUserAuthBody { auth })?, )?; - socket.send_to(&auth_packet, peer_addr).await?; - let (n, _) = tokio::time::timeout(timeout, socket.recv_from(&mut buf)).await??; + send_udp_retrying_transient(&socket, &auth_packet, peer_addr, timeout).await?; + let (n, _) = recv_udp_retrying_transient(&socket, &mut buf, timeout).await?; let packet = protocol::decode(&buf[..n])?; if packet.header.kind != PacketKind::NativeAuthOk { if packet.header.kind == PacketKind::AttachReject { diff --git a/src/server.rs b/src/server.rs index 9fd4816..3ae3fe7 100644 --- a/src/server.rs +++ b/src/server.rs @@ -13,7 +13,7 @@ use crate::transport::{ DoshTransport, SessionEvent, SessionRole, SessionTransportConfig, TransportConfig, service_name_from_target, }; -use crate::udp::is_transient_udp_send_error; +use crate::udp::{is_transient_udp_error, is_transient_udp_send_error}; use anyhow::{Context, Result, anyhow, bail}; use ed25519_dalek::SigningKey; use std::collections::{HashMap, HashSet}; @@ -169,7 +169,11 @@ impl DoshServer { let mut buf = vec![0u8; 65535]; loop { self.expire_pending(); - let (n, peer) = self.socket.recv_from(&mut buf).await?; + let (n, peer) = match self.socket.recv_from(&mut buf).await { + Ok(value) => value, + Err(err) if is_transient_udp_error(&err) => continue, + Err(err) => return Err(err.into()), + }; let packet = match protocol::decode(&buf[..n]) { Ok(packet) => packet, Err(_) => continue, diff --git a/src/udp.rs b/src/udp.rs index 5ad4a18..082df41 100644 --- a/src/udp.rs +++ b/src/udp.rs @@ -1,5 +1,9 @@ //! UDP error classification shared by Dosh clients, servers, and embedders. +use std::net::SocketAddr; +use std::time::{Duration, Instant}; +use tokio::net::UdpSocket; + pub fn is_transient_udp_error(err: &std::io::Error) -> bool { matches!( err.kind(), @@ -13,6 +17,54 @@ pub fn is_transient_udp_send_error(err: &std::io::Error) -> bool { is_transient_udp_error(err) } +pub async fn send_udp_retrying_transient( + socket: &UdpSocket, + packet: &[u8], + peer: SocketAddr, + wait: Duration, +) -> std::io::Result<()> { + let deadline = Instant::now() + wait.max(Duration::from_millis(1)); + loop { + match socket.send_to(packet, peer).await { + Ok(_) => return Ok(()), + Err(err) if is_transient_udp_send_error(&err) && Instant::now() < deadline => { + sleep_until_retry(deadline).await; + } + Err(err) => return Err(err), + } + } +} + +pub async fn recv_udp_retrying_transient( + socket: &UdpSocket, + buf: &mut [u8], + wait: Duration, +) -> std::io::Result<(usize, SocketAddr)> { + let deadline = Instant::now() + wait.max(Duration::from_millis(1)); + loop { + let now = Instant::now(); + if now >= deadline { + return Err(std::io::Error::from(std::io::ErrorKind::TimedOut)); + } + match tokio::time::timeout(deadline - now, socket.recv_from(buf)).await { + Ok(Ok(value)) => return Ok(value), + Ok(Err(err)) if is_transient_udp_error(&err) && Instant::now() < deadline => { + sleep_until_retry(deadline).await; + } + Ok(Err(err)) => return Err(err), + Err(_) => return Err(std::io::Error::from(std::io::ErrorKind::TimedOut)), + } + } +} + +async fn sleep_until_retry(deadline: Instant) { + let now = Instant::now(); + if now >= deadline { + return; + } + tokio::time::sleep((deadline - now).min(Duration::from_millis(20))).await; +} + #[cfg(unix)] fn is_transient_udp_os_error(code: i32) -> bool { matches!( @@ -52,7 +104,12 @@ fn is_transient_udp_os_error(_code: i32) -> bool { #[cfg(test)] mod tests { - use super::{is_transient_udp_error, is_transient_udp_send_error}; + use super::{ + is_transient_udp_error, is_transient_udp_send_error, recv_udp_retrying_transient, + send_udp_retrying_transient, + }; + use std::time::Duration; + use tokio::net::UdpSocket; #[test] fn classifies_portable_transient_udp_errors() { @@ -70,6 +127,35 @@ mod tests { assert!(!is_transient_udp_send_error(&fatal)); } + #[tokio::test] + async fn retrying_udp_send_and_receive_round_trip() { + let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let sender = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let addr = receiver.local_addr().unwrap(); + + send_udp_retrying_transient(&sender, b"hello", addr, Duration::from_millis(50)) + .await + .unwrap(); + let mut buf = [0u8; 16]; + let (n, peer) = recv_udp_retrying_transient(&receiver, &mut buf, Duration::from_millis(50)) + .await + .unwrap(); + + assert_eq!(&buf[..n], b"hello"); + assert_eq!(peer, sender.local_addr().unwrap()); + } + + #[tokio::test] + async fn retrying_udp_receive_times_out_cleanly() { + let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let mut buf = [0u8; 16]; + let err = recv_udp_retrying_transient(&receiver, &mut buf, Duration::from_millis(1)) + .await + .unwrap_err(); + + assert_eq!(err.kind(), std::io::ErrorKind::TimedOut); + } + #[cfg(unix)] #[test] fn classifies_unix_network_churn_as_transient() { diff --git a/tests/integration_smoke.rs b/tests/integration_smoke.rs index 4b99628..04a2812 100644 --- a/tests/integration_smoke.rs +++ b/tests/integration_smoke.rs @@ -6,6 +6,7 @@ use std::net::{TcpListener, TcpStream, UdpSocket}; use std::os::unix::net::{UnixListener, UnixStream}; use std::path::{Path, PathBuf}; use std::process::{Child, Command, Stdio}; +use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::mpsc; use std::thread; use std::time::Duration; @@ -23,14 +24,35 @@ use dosh::protocol::{ }; use ed25519_dalek::{Signer, SigningKey, VerifyingKey}; +const TEST_PORT_START: u16 = 30_000; +const TEST_PORT_END: u16 = 60_999; +static NEXT_UDP_PORT: AtomicU32 = AtomicU32::new(TEST_PORT_START as u32); +static NEXT_TCP_PORT: AtomicU32 = AtomicU32::new(TEST_PORT_START as u32); + +fn next_test_port(counter: &AtomicU32) -> u16 { + let span = u32::from(TEST_PORT_END - TEST_PORT_START) + 1; + let offset = counter.fetch_add(1, Ordering::SeqCst) % span; + TEST_PORT_START + offset as u16 +} + fn free_udp_port() -> u16 { - let socket = UdpSocket::bind("127.0.0.1:0").unwrap(); - socket.local_addr().unwrap().port() + for _ in TEST_PORT_START..=TEST_PORT_END { + let port = next_test_port(&NEXT_UDP_PORT); + if UdpSocket::bind(("127.0.0.1", port)).is_ok() { + return port; + } + } + panic!("no free UDP test port in {TEST_PORT_START}-{TEST_PORT_END}"); } fn free_tcp_port() -> u16 { - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - listener.local_addr().unwrap().port() + for _ in TEST_PORT_START..=TEST_PORT_END { + let port = next_test_port(&NEXT_TCP_PORT); + if TcpListener::bind(("127.0.0.1", port)).is_ok() { + return port; + } + } + panic!("no free TCP test port in {TEST_PORT_START}-{TEST_PORT_END}"); } fn write_server_config(dir: &tempfile::TempDir, port: u16) -> std::path::PathBuf {