diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index f9b191b..458bc87 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -2509,6 +2509,7 @@ async fn proxy_stdio_loop(mut transport: DoshTransport, stream_id: u64) -> Resul })?; let mut stdout = tokio::io::stdout(); let mut maintenance = tokio::time::interval(Duration::from_millis(50)); + maintenance.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); let mut stdin_closed = false; loop { tokio::select! { diff --git a/src/bin/dosh-server/unix.rs b/src/bin/dosh-server/unix.rs index f93443d..e8080a3 100644 --- a/src/bin/dosh-server/unix.rs +++ b/src/bin/dosh-server/unix.rs @@ -27,7 +27,9 @@ use dosh::protocol::{ StreamClose, StreamData, StreamEof, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust, TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope, }; -use dosh::pty::{PtyHandle, PtyOutput, adopt_pty_from_fd, spawn_pty_session}; +use dosh::pty::{ + PTY_OUTPUT_QUEUE_CAPACITY, PtyHandle, PtyOutput, adopt_pty_from_fd, spawn_pty_session, +}; use dosh::udp::{is_transient_udp_error, is_transient_udp_send_error}; use sha2::{Digest, Sha256}; use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; @@ -209,7 +211,7 @@ async fn serve(config_path: Option) -> Result<()> { ], ); - let (pty_tx, mut pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, mut pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let state = Arc::new(Mutex::new(ServerState::new( config.clone(), secret, @@ -242,6 +244,7 @@ async fn serve(config_path: Option) -> Result<()> { let retransmit_socket = Arc::clone(&socket); tokio::spawn(async move { let mut interval = tokio::time::interval(Duration::from_millis(100)); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { interval.tick().await; if let Err(err) = retransmit_pending(&retransmit_state, &retransmit_socket).await { @@ -256,6 +259,7 @@ async fn serve(config_path: Option) -> Result<()> { let cleanup_state = Arc::clone(&state); tokio::spawn(async move { let mut interval = tokio::time::interval(Duration::from_secs(5)); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { interval.tick().await; cleanup_disconnected_clients(&cleanup_state); @@ -270,6 +274,7 @@ async fn serve(config_path: Option) -> Result<()> { let flush_state = Arc::clone(&state); tokio::spawn(async move { let mut interval = tokio::time::interval(SCREEN_PERSIST_MAX_AGE); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { interval.tick().await; flush_persistent_screens(&flush_state); @@ -293,7 +298,7 @@ async fn serve(config_path: Option) -> Result<()> { struct ServerState { config: ServerConfig, secret: [u8; 32], - pty_tx: mpsc::UnboundedSender, + pty_tx: mpsc::Sender, sessions: HashMap, pending_native: HashMap<[u8; 16], PendingNativeAuth>, next_server_stream_id: u64, @@ -504,7 +509,7 @@ impl ServerState { fn new( config: ServerConfig, secret: [u8; 32], - pty_tx: mpsc::UnboundedSender, + pty_tx: mpsc::Sender, ) -> Self { let per_minute = config.native_auth_rate_limit_per_minute; Self { @@ -4517,7 +4522,7 @@ mod tests { #[test] fn persists_terminal_sessions_when_enabled() { - let (pty_tx, _rx) = mpsc::unbounded_channel(); + let (pty_tx, _rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let config = ServerConfig { persist_sessions: true, prewarm_sessions: vec!["default".to_string()], @@ -4531,7 +4536,7 @@ mod tests { #[test] fn persist_disabled_never_persists() { - let (pty_tx, _rx) = mpsc::unbounded_channel(); + let (pty_tx, _rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let config = ServerConfig { persist_sessions: false, prewarm_sessions: vec!["default".to_string()], @@ -4655,7 +4660,7 @@ mod tests { #[tokio::test] async fn no_client_pty_output_before_first_reattach_keeps_restored_screen() { let restored = b"\x1b[?1049lRESTORED_AFTER_RESTART".to_vec(); - let (pty_tx, _rx) = mpsc::unbounded_channel(); + let (pty_tx, _rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let state = Arc::new(Mutex::new(ServerState::new( ServerConfig { persist_sessions: true, @@ -4733,7 +4738,7 @@ mod tests { #[tokio::test] async fn unknown_resume_reject_keeps_client_id() { - let (pty_tx, _rx) = mpsc::unbounded_channel(); + let (pty_tx, _rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let state = Arc::new(Mutex::new(ServerState::new( ServerConfig::default(), [0u8; 32], @@ -4957,7 +4962,7 @@ mod tests { #[test] fn client_index_stays_in_sync_with_session_clients() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); state .ensure_session("work", 80, 24, "forward-only", &[]) @@ -5004,7 +5009,7 @@ mod tests { #[test] fn cleanup_purges_timed_out_clients_from_index() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let config = ServerConfig { client_timeout_secs: 1, ..ServerConfig::default() @@ -5052,7 +5057,7 @@ mod tests { #[tokio::test] async fn forged_plaintext_detach_does_not_remove_client() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); state .ensure_session("work", 80, 24, "forward-only", &[]) @@ -5078,7 +5083,7 @@ mod tests { #[tokio::test] async fn authenticated_detach_removes_client() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); state .ensure_session("work", 80, 24, "forward-only", &[]) @@ -5109,7 +5114,7 @@ mod tests { #[tokio::test] async fn duplicate_stream_open_for_open_stream_resends_ok() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); let client_id = [11u8; 16]; let session_key = [12u8; 32]; @@ -5173,7 +5178,7 @@ mod tests { #[tokio::test] async fn retired_stream_open_is_rejected_not_reopened() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); let client_id = [15u8; 16]; let session_key = [16u8; 32]; @@ -5243,7 +5248,7 @@ mod tests { #[tokio::test] async fn pending_server_stream_open_is_retransmitted() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); let client_id = [13u8; 16]; let session_key = [14u8; 32]; @@ -5298,7 +5303,7 @@ mod tests { #[tokio::test] async fn server_stream_retransmit_uses_observed_rtt() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); let client_id = [15u8; 16]; let session_key = [16u8; 32]; @@ -5357,7 +5362,7 @@ mod tests { #[test] fn forward_only_session_does_not_allocate_pty() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); state @@ -5369,7 +5374,7 @@ mod tests { #[test] fn cleanup_reaps_abandoned_sessions_but_keeps_prewarmed() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let config = ServerConfig { client_timeout_secs: 1, prewarm_sessions: vec!["default".to_string()], @@ -5454,7 +5459,7 @@ mod tests { #[tokio::test] async fn stream_data_waits_for_open_and_credit() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); let client_id = [7u8; 16]; let session_key = [9u8; 32]; @@ -5558,7 +5563,7 @@ mod tests { #[tokio::test] async fn stream_eof_to_client_does_not_retire_stream() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); let client_id = [13u8; 16]; let session_key = [14u8; 32]; @@ -5647,7 +5652,7 @@ mod tests { #[tokio::test] async fn pending_server_stream_eof_is_retransmitted() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); let client_id = [17u8; 16]; let session_key = [18u8; 32]; @@ -5707,7 +5712,7 @@ mod tests { #[tokio::test] async fn pending_server_stream_close_is_retransmitted() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); let client_id = [19u8; 16]; let session_key = [20u8; 32]; @@ -5763,7 +5768,7 @@ mod tests { #[tokio::test] async fn pending_server_stream_close_expires_after_attempt_cap() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); let client_id = [21u8; 16]; let session_key = [22u8; 32]; @@ -5808,7 +5813,7 @@ mod tests { #[tokio::test] async fn pending_server_stream_window_adjust_is_retransmitted() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); let client_id = [23u8; 16]; let session_key = [24u8; 32]; @@ -5871,7 +5876,7 @@ mod tests { #[tokio::test] async fn pending_server_stream_window_adjust_expires_after_attempt_cap() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); let client_id = [25u8; 16]; let session_key = [26u8; 32]; @@ -5922,7 +5927,7 @@ mod tests { #[tokio::test] async fn server_stream_send_splits_large_writes() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); let client_id = [11u8; 16]; let session_key = [12u8; 32]; @@ -6029,7 +6034,7 @@ mod tests { #[tokio::test] async fn blocked_stream_data_does_not_block_terminal_frames() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); let client_id = [8u8; 16]; let session_key = [10u8; 32]; @@ -6134,7 +6139,7 @@ mod tests { #[tokio::test] async fn live_terminal_output_is_never_replaced_by_paced_snapshots() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let (pty_tx, _pty_rx) = mpsc::channel(PTY_OUTPUT_QUEUE_CAPACITY); let config = ServerConfig { output_frame_interval_ms: 1000, ..ServerConfig::default() diff --git a/src/pty.rs b/src/pty.rs index 3354b41..d0798a7 100644 --- a/src/pty.rs +++ b/src/pty.rs @@ -12,6 +12,10 @@ use tokio::sync::mpsc; // TUIs often write several KiB on the first draw; sending that as one UDP // datagram can fragment and vanish, leaving only a blank alternate screen. const PTY_OUTPUT_CHUNK_BYTES: usize = 1024; +/// Shared server queue bound. At the current chunk size this limits queued PTY +/// output to roughly 1 MiB before the producing shell receives normal PTY +/// backpressure. +pub const PTY_OUTPUT_QUEUE_CAPACITY: usize = 1024; /// Backing for a PTY master held by the server. /// @@ -126,7 +130,7 @@ pub fn spawn_pty_session( cols: u16, rows: u16, env: &[(String, String)], - tx: mpsc::UnboundedSender, + tx: mpsc::Sender, ) -> Result { let pty_system = NativePtySystem::default(); let pair = pty_system @@ -232,7 +236,7 @@ pub fn build_shell_command(shell: &str, env: &[(String, String)]) -> CommandBuil pub fn adopt_pty_from_fd( session: String, master_fd: RawFd, - tx: mpsc::UnboundedSender, + tx: mpsc::Sender, ) -> Result { // Take ownership of the fd. A clone gives us an independent reader so the // reader thread and the writer/resize side hold separate `File`s and don't @@ -252,7 +256,7 @@ pub fn adopt_pty_from_fd( fn spawn_reader_thread( session: String, mut reader: Box, - tx: mpsc::UnboundedSender, + tx: mpsc::Sender, ) -> Result<()> { let reader_session = session.clone(); thread::Builder::new() @@ -262,7 +266,7 @@ fn spawn_reader_thread( loop { match reader.read(&mut buf) { Ok(0) => { - let _ = tx.send(PtyOutput { + let _ = tx.blocking_send(PtyOutput { session: reader_session.clone(), bytes: Vec::new(), exited: true, @@ -271,15 +275,20 @@ fn spawn_reader_thread( } Ok(n) => { for chunk in buf[..n].chunks(PTY_OUTPUT_CHUNK_BYTES) { - let _ = tx.send(PtyOutput { - session: reader_session.clone(), - bytes: chunk.to_vec(), - exited: false, - }); + if tx + .blocking_send(PtyOutput { + session: reader_session.clone(), + bytes: chunk.to_vec(), + exited: false, + }) + .is_err() + { + return; + } } } Err(_) => { - let _ = tx.send(PtyOutput { + let _ = tx.blocking_send(PtyOutput { session: reader_session.clone(), bytes: Vec::new(), exited: true, @@ -298,6 +307,28 @@ mod tests { use super::*; const _: () = assert!(PTY_OUTPUT_CHUNK_BYTES <= 1200); + const _: () = assert!(PTY_OUTPUT_QUEUE_CAPACITY * PTY_OUTPUT_CHUNK_BYTES <= 1024 * 1024); + + #[test] + fn pty_output_queue_capacity_is_memory_bounded() { + let (tx, _rx) = mpsc::channel::(PTY_OUTPUT_QUEUE_CAPACITY); + for index in 0..PTY_OUTPUT_QUEUE_CAPACITY { + tx.try_send(PtyOutput { + session: "load".to_string(), + bytes: vec![index as u8; PTY_OUTPUT_CHUNK_BYTES], + exited: false, + }) + .unwrap(); + } + assert!(matches!( + tx.try_send(PtyOutput { + session: "load".to_string(), + bytes: vec![0; PTY_OUTPUT_CHUNK_BYTES], + exited: false, + }), + Err(mpsc::error::TrySendError::Full(_)) + )); + } #[test] fn terminfo_available_detects_known_and_unknown() {