From bca21e317c497da952ddd438c8866c4e102fd745 Mon Sep 17 00:00:00 2001 From: Kang Date: Tue, 7 Jul 2026 17:16:38 +0800 Subject: [PATCH] perf(zmodem): implement upload and download drain mechanisms (#234) - Introduced `ZmodemUploadDrain` and `ZmodemDownloadOoDrain` to manage residual byte suppression during Zmodem transfers. - Updated the SSH I/O loop and Telnet session task to utilize the new drain mechanisms for improved data handling. - Enhanced Zmodem transfer logic to track upload and download states, ensuring efficient data processing and user feedback. - Added tests for the new drain functionalities to validate their behavior and performance. --- src-tauri/src/core/ssh/io.rs | 26 ++- .../src/core/terminal_session/telnet/mod.rs | 4 +- .../core/terminal_session/telnet/session.rs | 45 ++++- src-tauri/src/core/zmodem/mod.rs | 4 +- src-tauri/src/core/zmodem/tests.rs | 73 ++++++- src-tauri/src/core/zmodem/transfer.rs | 180 +++++++++++++++++- src-tauri/vendor/zmodem2/src/crc.rs | 76 ++++---- src-tauri/vendor/zmodem2/src/transmission.rs | 66 +++++-- src-tauri/vendor/zmodem2/tests/protocol.rs | 101 +++++++++- 9 files changed, 505 insertions(+), 70 deletions(-) diff --git a/src-tauri/src/core/ssh/io.rs b/src-tauri/src/core/ssh/io.rs index 75d41d7d..67dc3a17 100644 --- a/src-tauri/src/core/ssh/io.rs +++ b/src-tauri/src/core/ssh/io.rs @@ -3,8 +3,8 @@ use crate::core::capture::OutputCaptureProcessor; use crate::core::input::remap_del_to_bs; use crate::core::ssh::osc::{self, OscStripper, ShellKind}; use crate::core::zmodem::{ - ZmodemAction, ZmodemDetectResult, ZmodemDetector, ZmodemDirection, ZmodemEvent, ZmodemTransfer, - start_zmodem_transfer, + ZmodemAction, ZmodemDetectResult, ZmodemDetector, ZmodemDirection, ZmodemDownloadOoDrain, + ZmodemEvent, ZmodemTransfer, ZmodemUploadDrain, start_zmodem_transfer, }; use crate::core::{ RecordingManager, SessionCommand, SessionManager, SessionOutputCoalescer, SharedCwd, @@ -423,6 +423,8 @@ pub(super) async fn ssh_io_loop( let mut zmodem_detector = ZmodemDetector::new(); let mut zmodem_transfer: Option = None; + let mut zmodem_upload_drain = ZmodemUploadDrain::new(); + let mut zmodem_download_oo_drain = ZmodemDownloadOoDrain::new(); let zmodem_event_name = format!("zmodem-event-{session_id}"); let initial_inject_deadline = @@ -442,7 +444,9 @@ pub(super) async fn ssh_io_loop( output.attach(); } Some(SessionCommand::Write(mut data)) => { - if zmodem_transfer.is_some() { + if zmodem_transfer.is_some() + || zmodem_upload_drain.should_suppress(std::time::Instant::now()) + { continue; } if backspace_as_bs { @@ -515,15 +519,31 @@ pub(super) async fn ssh_io_loop( Some(ChannelMsg::Data { ref data }) => { // ZMODEM: if a transfer is active, route raw bytes to it. if let Some(ref mut transfer) = zmodem_transfer { + let direction = transfer.direction(); let actions = transfer.feed_incoming(data); handle_zmodem_actions(&app, &zmodem_event_name, &mut channel, actions).await; if transfer.is_done() { zmodem_transfer = None; zmodem_detector.reset(); + if direction == ZmodemDirection::Upload { + zmodem_upload_drain.start(std::time::Instant::now()); + } else if direction == ZmodemDirection::Download { + zmodem_download_oo_drain.start(std::time::Instant::now()); + } } continue; } + let data = zmodem_upload_drain.filter(data, std::time::Instant::now()); + if data.is_empty() { + continue; + } + let data = zmodem_download_oo_drain + .filter(data, std::time::Instant::now()); + if data.is_empty() { + continue; + } + // ZMODEM: detect header in raw bytes before lossy UTF-8 conversion. // Detection must run in every phase so an early `rz` is not missed // while shell-integration output is still being suppressed. diff --git a/src-tauri/src/core/terminal_session/telnet/mod.rs b/src-tauri/src/core/terminal_session/telnet/mod.rs index 5ffa8858..5687569a 100644 --- a/src-tauri/src/core/terminal_session/telnet/mod.rs +++ b/src-tauri/src/core/terminal_session/telnet/mod.rs @@ -7,8 +7,8 @@ use crate::core::session::{ SessionCommand, SessionHandle, SessionInfo, SessionManager, SessionType, SharedCwd, }; use crate::core::zmodem::{ - ZmodemAction, ZmodemDetectResult, ZmodemDetector, ZmodemDirection, ZmodemEvent, ZmodemTransfer, - start_zmodem_transfer, + ZmodemAction, ZmodemDetectResult, ZmodemDetector, ZmodemDirection, ZmodemDownloadOoDrain, + ZmodemEvent, ZmodemTransfer, ZmodemUploadDrain, start_zmodem_transfer, }; use crate::core::{RecordingManager, SessionOutputCoalescer}; use crate::error::AppResult; diff --git a/src-tauri/src/core/terminal_session/telnet/session.rs b/src-tauri/src/core/terminal_session/telnet/session.rs index 72c991ed..218c2d90 100644 --- a/src-tauri/src/core/terminal_session/telnet/session.rs +++ b/src-tauri/src/core/terminal_session/telnet/session.rs @@ -53,6 +53,10 @@ async fn telnet_session_task( let zmodem_state: Arc>> = Arc::new(TokioMutex::new(None)); let zmodem_state_reader = zmodem_state.clone(); + let zmodem_upload_drain = Arc::new(TokioMutex::new(ZmodemUploadDrain::new())); + let zmodem_upload_drain_reader = zmodem_upload_drain.clone(); + let zmodem_download_oo_drain = Arc::new(TokioMutex::new(ZmodemDownloadOoDrain::new())); + let zmodem_download_oo_drain_reader = zmodem_download_oo_drain.clone(); let zmodem_event_name = format!("zmodem-event-{session_id}"); let zmodem_event_reader = zmodem_event_name.clone(); let (zmodem_out_tx, mut zmodem_out_rx) = mpsc::unbounded_channel::>(); @@ -101,10 +105,21 @@ async fn telnet_session_task( continue; } + let visible = if zmodem_upload_drain_reader.lock().await.is_active() { + let mut drain = zmodem_upload_drain_reader.lock().await; + drain.filter(&visible, std::time::Instant::now()).to_vec() + } else { + visible + }; + if visible.is_empty() { + continue; + } + // ZMODEM: if active, route to transfer. { let mut zm = zmodem_state_reader.lock().await; if let Some(ref mut transfer) = *zm { + let direction = transfer.direction(); let actions = transfer.feed_incoming(&visible); for action in actions { match action { @@ -119,11 +134,34 @@ async fn telnet_session_task( if transfer.is_done() { *zm = None; zmodem_detector.reset(); + if direction == ZmodemDirection::Upload { + zmodem_upload_drain_reader + .lock() + .await + .start(std::time::Instant::now()); + } else if direction == ZmodemDirection::Download { + zmodem_download_oo_drain_reader + .lock() + .await + .start(std::time::Instant::now()); + } } continue; } } + let visible = if zmodem_download_oo_drain_reader.lock().await.is_active() { + let mut drain = zmodem_download_oo_drain_reader.lock().await; + drain + .filter(&visible, std::time::Instant::now()) + .to_vec() + } else { + visible + }; + if visible.is_empty() { + continue; + } + // ZMODEM: detect header. let process_visible = match zmodem_detector.feed(&visible) { ZmodemDetectResult::Detected { @@ -245,7 +283,12 @@ async fn telnet_session_task( output.attach(); } Some(SessionCommand::Write(mut data)) => { - if zmodem_state.lock().await.is_some() { + if zmodem_state.lock().await.is_some() + || zmodem_upload_drain + .lock() + .await + .should_suppress(std::time::Instant::now()) + { continue; } diff --git a/src-tauri/src/core/zmodem/mod.rs b/src-tauri/src/core/zmodem/mod.rs index 6c51842c..e510aafd 100644 --- a/src-tauri/src/core/zmodem/mod.rs +++ b/src-tauri/src/core/zmodem/mod.rs @@ -7,7 +7,7 @@ //! transitions to an active [`ZmodemTransfer`] that owns the protocol state. use serde::Serialize; -use std::io::{Read, Seek, SeekFrom, Write}; +use std::io::{BufWriter, Read, Seek, SeekFrom, Write}; use std::path::PathBuf; use std::time::{Duration, Instant}; @@ -26,6 +26,8 @@ const CANCEL_SEQ_LEN: usize = 5; const ZMODEM_PROGRESS_INTERVAL: Duration = Duration::from_millis(100); const ZMODEM_PROGRESS_BYTES: u64 = 256 * 1024; +const ZMODEM_FINISH_DRAIN_IDLE: Duration = Duration::from_millis(250); +const ZMODEM_FILE_WRITE_BUFFER_SIZE: usize = 1024 * 1024; include!("types.rs"); include!("detector.rs"); diff --git a/src-tauri/src/core/zmodem/tests.rs b/src-tauri/src/core/zmodem/tests.rs index c65e9a95..66e6e448 100644 --- a/src-tauri/src/core/zmodem/tests.rs +++ b/src-tauri/src/core/zmodem/tests.rs @@ -1,8 +1,9 @@ #[cfg(test)] mod tests { use super::{ - ProgressThrottle, ZMODEM_PROGRESS_BYTES, ZMODEM_PROGRESS_INTERVAL, ZmodemDetectResult, - ZmodemDetector, ZmodemDirection, ZmodemEvent, + ProgressThrottle, ZMODEM_FINISH_DRAIN_IDLE, ZMODEM_PROGRESS_BYTES, + ZMODEM_PROGRESS_INTERVAL, ZmodemDetectResult, ZmodemDetector, ZmodemDirection, + ZmodemDownloadOoDrain, ZmodemEvent, ZmodemUploadDrain, }; use serde_json::json; use std::time::{Duration, Instant}; @@ -219,4 +220,72 @@ mod tests { assert!(throttle.should_emit_at(128, false, start)); assert!(throttle.should_emit_at(129, true, start + Duration::from_millis(1))); } + + #[test] + fn upload_drain_suppresses_residual_bytes_until_idle() { + let mut drain = ZmodemUploadDrain::new(); + let start = Instant::now(); + + assert!(!drain.should_suppress(start)); + + drain.start(start); + assert!(drain.should_suppress(start + ZMODEM_FINISH_DRAIN_IDLE / 2)); + assert!(drain.should_suppress(start + ZMODEM_FINISH_DRAIN_IDLE)); + assert!(!drain.should_suppress(start + ZMODEM_FINISH_DRAIN_IDLE * 3)); + } + + #[test] + fn upload_drain_allows_prompt_text() { + let mut drain = ZmodemUploadDrain::new(); + let start = Instant::now(); + + drain.start(start); + assert_eq!( + drain.filter( + b"\r\n[root@ubuntu ~]# ", + start + ZMODEM_FINISH_DRAIN_IDLE / 2 + ), + b"\r\n[root@ubuntu ~]# " + ); + assert_eq!(drain.filter(b"next", start + Duration::from_millis(1)), b"next"); + } + + #[test] + fn upload_drain_suppresses_binary_residue() { + let mut drain = ZmodemUploadDrain::new(); + let start = Instant::now(); + + drain.start(start); + assert_eq!( + drain.filter(b"\x18\x18\x00\xff", start + Duration::from_millis(1)), + b"" + ); + assert!(drain.should_suppress(start + Duration::from_millis(2))); + } + + #[test] + fn download_oo_drain_strips_only_zmodem_trailer() { + let mut drain = ZmodemDownloadOoDrain::new(); + let start = Instant::now(); + + drain.start(start); + assert_eq!(drain.filter(b"O", start + Duration::from_millis(1)), b""); + assert_eq!( + drain.filter(b"Oprompt", start + Duration::from_millis(2)), + b"prompt" + ); + assert_eq!(drain.filter(b"next", start + Duration::from_millis(3)), b"next"); + } + + #[test] + fn download_oo_drain_expires_without_stripping_prompt() { + let mut drain = ZmodemDownloadOoDrain::new(); + let start = Instant::now(); + + drain.start(start); + assert_eq!( + drain.filter(b"Output", start + ZMODEM_FINISH_DRAIN_IDLE * 2), + b"Output" + ); + } } diff --git a/src-tauri/src/core/zmodem/transfer.rs b/src-tauri/src/core/zmodem/transfer.rs index 103b8ca7..7f1c08ed 100644 --- a/src-tauri/src/core/zmodem/transfer.rs +++ b/src-tauri/src/core/zmodem/transfer.rs @@ -20,6 +20,7 @@ enum TransferState { receiver: zmodem2::Receiver, save_dir: PathBuf, current_file: Option, + stats: TransferStats, }, /// Actively sending files (upload / remote `rz`). Sending { @@ -27,6 +28,7 @@ enum TransferState { files: Vec, file_index: usize, current_file: Option, + stats: TransferStats, }, /// Transfer finished or aborted. Done, @@ -35,7 +37,7 @@ enum TransferState { struct ReceiveFile { name: String, size: u64, - file: std::fs::File, + file: BufWriter, written: u64, } @@ -86,6 +88,140 @@ impl ProgressThrottle { } } +struct TransferStats { + started_at: Instant, + bytes: u64, + ack_count: u64, + file_write_count: u64, +} + +impl TransferStats { + fn new() -> Self { + Self { + started_at: Instant::now(), + bytes: 0, + ack_count: 0, + file_write_count: 0, + } + } + + fn elapsed_secs(&self) -> f64 { + self.started_at.elapsed().as_secs_f64() + } + + fn mbps(&self) -> f64 { + let elapsed = self.elapsed_secs(); + if elapsed <= f64::EPSILON { + 0.0 + } else { + self.bytes as f64 / 1024.0 / 1024.0 / elapsed + } + } +} + +#[derive(Default)] +pub struct ZmodemUploadDrain { + suppress_until: Option, +} + +impl ZmodemUploadDrain { + pub fn new() -> Self { + Self { + suppress_until: None, + } + } + + pub fn start(&mut self, now: Instant) { + self.suppress_until = now.checked_add(ZMODEM_FINISH_DRAIN_IDLE); + } + + pub fn is_active(&self) -> bool { + self.suppress_until.is_some() + } + + pub fn should_suppress(&mut self, now: Instant) -> bool { + let Some(until) = self.suppress_until else { + return false; + }; + + if now <= until { + self.start(now); + return true; + } + + self.suppress_until = None; + false + } + + pub fn filter<'a>(&mut self, data: &'a [u8], now: Instant) -> &'a [u8] { + let Some(until) = self.suppress_until else { + return data; + }; + + if now > until { + self.suppress_until = None; + return data; + } + + if looks_like_terminal_text(data) { + self.suppress_until = None; + return data; + } + + self.start(now); + &data[data.len()..] + } +} + +#[derive(Default)] +pub struct ZmodemDownloadOoDrain { + suppress_until: Option, + remaining_o: usize, +} + +impl ZmodemDownloadOoDrain { + pub fn new() -> Self { + Self { + suppress_until: None, + remaining_o: 0, + } + } + + pub fn start(&mut self, now: Instant) { + self.suppress_until = now.checked_add(ZMODEM_FINISH_DRAIN_IDLE); + self.remaining_o = 2; + } + + pub fn is_active(&self) -> bool { + self.suppress_until.is_some() + } + + pub fn filter<'a>(&mut self, data: &'a [u8], now: Instant) -> &'a [u8] { + let Some(until) = self.suppress_until else { + return data; + }; + + if now > until { + self.suppress_until = None; + self.remaining_o = 0; + return data; + } + + let mut offset = 0; + while self.remaining_o > 0 && data.get(offset) == Some(&b'O') { + offset += 1; + self.remaining_o -= 1; + } + + if self.remaining_o == 0 || offset < data.len() { + self.suppress_until = None; + self.remaining_o = 0; + } + + &data[offset..] + } +} + impl ZmodemTransfer { pub fn new(direction: ZmodemDirection, initial_bytes: &[u8]) -> Self { Self { @@ -141,6 +277,7 @@ impl ZmodemTransfer { receiver, save_dir, current_file: None, + stats: TransferStats::new(), }; let mut actions = self.drain_outgoing(); @@ -171,6 +308,7 @@ impl ZmodemTransfer { files, file_index: 0, current_file: None, + stats: TransferStats::new(), }; let mut actions = Vec::new(); @@ -221,6 +359,7 @@ impl ZmodemTransfer { receiver, save_dir, current_file, + stats, } = &mut self.state else { return actions; @@ -258,6 +397,7 @@ impl ZmodemTransfer { let out = out.to_vec(); let n = out.len(); actions.push(ZmodemAction::SendToRemote(out)); + stats.ack_count += 1; receiver.advance_outgoing(n); } @@ -282,7 +422,10 @@ impl ZmodemTransfer { *current_file = Some(ReceiveFile { name: name.clone(), size, - file, + file: BufWriter::with_capacity( + ZMODEM_FILE_WRITE_BUFFER_SIZE, + file, + ), written: 0, }); if self.progress_throttle.should_emit(0, true) { @@ -324,6 +467,7 @@ impl ZmodemTransfer { *current_file = None; } zmodem2::ReceiverEvent::SessionComplete => { + log_transfer_complete(ZmodemDirection::Download, self.file_count, stats); self.state = TransferState::Done; actions.push(ZmodemAction::EmitEvent(ZmodemEvent::Complete { direction: ZmodemDirection::Download, @@ -350,6 +494,8 @@ impl ZmodemTransfer { return actions; } rf.written += len as u64; + stats.bytes += len as u64; + stats.file_write_count += 1; if self.progress_throttle.should_emit(rf.written, false) { actions.push(ZmodemAction::EmitEvent(ZmodemEvent::Progress { file_name: rf.name.clone(), @@ -377,6 +523,7 @@ impl ZmodemTransfer { files, file_index, current_file, + stats, } = &mut self.state else { return actions; @@ -434,6 +581,7 @@ impl ZmodemTransfer { tracing::warn!("feed_file error: {e}"); break; } + stats.bytes += n as u64; sf.sent = sf.sent.max(requested_offset + n as u64); sf.position += n as u64; if self.progress_throttle.should_emit(sf.sent, false) { @@ -494,6 +642,7 @@ impl ZmodemTransfer { } } zmodem2::SenderEvent::SessionComplete => { + log_transfer_complete(ZmodemDirection::Upload, self.file_count, stats); self.state = TransferState::Done; actions.push(ZmodemAction::EmitEvent(ZmodemEvent::Complete { direction: ZmodemDirection::Upload, @@ -516,6 +665,7 @@ impl ZmodemTransfer { files, file_index, current_file, + .. } = &mut self.state else { return vec![]; @@ -630,3 +780,29 @@ fn drain_sender_outgoing(sender: &mut zmodem2::Sender, actions: &mut Vec bool { + if data.is_empty() || std::str::from_utf8(data).is_err() { + return false; + } + + data.iter().all(|&byte| { + matches!( + byte, + b'\t' | b'\n' | b'\r' | 0x07 | 0x08 | 0x1b | 0x20..=0x7e | 0x80..=0xff + ) + }) +} + diff --git a/src-tauri/vendor/zmodem2/src/crc.rs b/src-tauri/vendor/zmodem2/src/crc.rs index c6e8ec15..13fe2491 100644 --- a/src-tauri/vendor/zmodem2/src/crc.rs +++ b/src-tauri/vendor/zmodem2/src/crc.rs @@ -3,7 +3,7 @@ /// Computes the CRC-16-XMODEM checksum. #[must_use] -pub const fn crc16_xmodem(data: &[u8]) -> u16 { +pub fn crc16_xmodem(data: &[u8]) -> u16 { let mut crc: u16 = 0x0000; let mut i = 0; while i < data.len() { @@ -15,14 +15,8 @@ pub const fn crc16_xmodem(data: &[u8]) -> u16 { /// Computes the CRC-32-ISO-HDLC checksum. #[must_use] -pub const fn crc32_iso_hdlc(data: &[u8]) -> u32 { - let mut crc: u32 = 0xFFFF_FFFF; - let mut i = 0; - while i < data.len() { - crc = crc32_update(crc, data[i]); - i += 1; - } - !crc +pub fn crc32_iso_hdlc(data: &[u8]) -> u32 { + crc32fast::hash(data) } /// A stateful, iterative CRC-16-XMODEM calculator. @@ -40,7 +34,7 @@ impl Crc16 { /// Updates the CRC state with a slice of bytes. pub fn update(&mut self, data: &[u8]) { for &byte in data { - self.update_byte(byte); + self.crc = crc16_update(self.crc, byte); } } @@ -58,61 +52,59 @@ impl Crc16 { /// A stateful, iterative CRC-32-ISO-HDLC calculator. pub struct Crc32 { - crc: u32, + hasher: crc32fast::Hasher, } impl Crc32 { /// Creates a new CRC-32 calculator. #[must_use] - pub const fn new() -> Self { - Self { crc: 0xFFFF_FFFF } + pub fn new() -> Self { + Self { + hasher: crc32fast::Hasher::new(), + } } /// Updates the CRC state with a slice of bytes. pub fn update(&mut self, data: &[u8]) { - for &byte in data { - self.update_byte(byte); - } + self.hasher.update(data); } /// Updates the CRC state with a single byte. pub fn update_byte(&mut self, byte: u8) { - self.crc = crc32_update(self.crc, byte); + self.hasher.update(&[byte]); } /// Finalizes the CRC calculation and returns the checksum. #[must_use] - pub const fn finalize(&self) -> u32 { - !self.crc + pub fn finalize(&self) -> u32 { + self.hasher.clone().finalize() } } /// Performs a single byte update for CRC-16-XMODEM. -const fn crc16_update(mut crc: u16, byte: u8) -> u16 { - crc ^= (byte as u16) << 8; - let mut i = 0; - while i < 8 { - if (crc & 0x8000) != 0 { - crc = (crc << 1) ^ 0x1021; - } else { - crc <<= 1; - } - i += 1; - } - crc +const fn crc16_update(crc: u16, byte: u8) -> u16 { + let index = ((crc >> 8) as u8 ^ byte) as usize; + (crc << 8) ^ CRC16_XMODEM_TABLE[index] } -/// Performs a single byte update for CRC-32-ISO-HDLC. -const fn crc32_update(mut crc: u32, byte: u8) -> u32 { - crc ^= byte as u32; - let mut i = 0; - while i < 8 { - if (crc & 1) != 0 { - crc = (crc >> 1) ^ 0xEDB8_8320; - } else { - crc >>= 1; +const fn make_crc16_xmodem_table() -> [u16; 256] { + let mut table = [0u16; 256]; + let mut byte = 0usize; + while byte < 256 { + let mut crc = (byte as u16) << 8; + let mut bit = 0; + while bit < 8 { + if (crc & 0x8000) != 0 { + crc = (crc << 1) ^ 0x1021; + } else { + crc <<= 1; + } + bit += 1; } - i += 1; + table[byte] = crc; + byte += 1; } - crc + table } + +const CRC16_XMODEM_TABLE: [u16; 256] = make_crc16_xmodem_table(); diff --git a/src-tauri/vendor/zmodem2/src/transmission.rs b/src-tauri/vendor/zmodem2/src/transmission.rs index 1cee80f2..f28c0999 100644 --- a/src-tauri/vendor/zmodem2/src/transmission.rs +++ b/src-tauri/vendor/zmodem2/src/transmission.rs @@ -18,10 +18,13 @@ use crate::{ZDLE, ZPAD}; use core::cmp::min; use core::fmt::Write as _; -/// Size of the unescaped subpacket payload. The size is picked from the -/// original ZMODEM specification. -const SUBPACKET_MAX_SIZE: usize = 1024; -const SUBPACKET_PER_ACK: usize = 10; +/// Size of the unescaped subpacket payload. +/// +/// 8 KiB keeps the fixed-capacity buffers modest while avoiding the severe +/// throughput penalty of 1 KiB frames over SSH and other high-latency links. +const SUBPACKET_MAX_SIZE: usize = 8 * 1024; +const SUBPACKET_PER_ACK: usize = 8; +const RECEIVER_BUFFER_SIZE: u16 = u16::MAX; const MAX_HEADER_ESCAPED: usize = 128; const MAX_SUBPACKET_ESCAPED: usize = SUBPACKET_MAX_SIZE * 2 + 2 + 8; const WIRE_BUF_SIZE: usize = MAX_HEADER_ESCAPED + MAX_SUBPACKET_ESCAPED; @@ -243,6 +246,14 @@ impl<'a> SliceReader<'a> { fn consumed(&self) -> usize { self.pos } + + fn remaining(&self) -> &[u8] { + &self.buf[self.pos..] + } + + fn advance(&mut self, n: usize) { + self.pos = min(self.pos.saturating_add(n), self.buf.len()); + } } impl Read for SliceReader<'_> { @@ -330,6 +341,14 @@ impl RxCrc { } } + fn update_slice(&mut self, data: &[u8], encoding: Encoding) { + if encoding == Encoding::ZBIN32 { + self.calc32.update(data); + } else { + self.calc16.update(data); + } + } + fn process( &mut self, port: &mut P, @@ -1167,10 +1186,7 @@ impl Receiver { } /// Handles reading a single byte for the `SubpacketState::Reading` state. - fn receive_subpacket_data_byte

(&mut self, port: &mut P) -> Result, Error> - where - P: Read + ?Sized, - { + fn receive_subpacket_data_byte(&mut self, port: &mut SliceReader<'_>) -> Result, Error> { let handle_followup = |this: &mut Self, byte: u8| -> Result, Error> { if let Ok(packet) = SubpacketType::try_from(byte) { this.crc.update(packet as u8, this.data_encoding); @@ -1191,6 +1207,29 @@ impl Receiver { return handle_followup(self, byte); } + let available = port.remaining(); + if !available.is_empty() { + let spare = SUBPACKET_MAX_SIZE.saturating_sub(self.buf.len()); + if spare == 0 { + return Err(Error::OutOfMemory); + } + let run = available + .iter() + .position(|&byte| byte == ZDLE) + .unwrap_or(available.len()) + .min(spare); + + if run > 0 { + let data = &available[..run]; + self.buf + .extend_from_slice(data) + .map_err(|_| Error::OutOfMemory)?; + self.crc.update_slice(data, self.data_encoding); + port.advance(run); + return Ok(Some(())); + } + } + let Some(byte) = port.read_byte()? else { return Ok(None); }; @@ -1207,10 +1246,7 @@ impl Receiver { Ok(Some(())) } - fn process_subpacket

(&mut self, port: &mut P) -> Result, Error> - where - P: Read + ?Sized, - { + fn process_subpacket(&mut self, port: &mut SliceReader<'_>) -> Result, Error> { match self.subpacket_state { SubpacketState::Reading => self.receive_subpacket_data_byte(port), SubpacketState::Crc(packet) => { @@ -1347,10 +1383,8 @@ fn write_zrinit

(port: &mut P) -> Result, Error> where P: Write + ?Sized, { - let zrinit = Zrinit::CANFDX | Zrinit::CANFC32; - let buffer_size = u16::try_from(SUBPACKET_MAX_SIZE) - .map_err(|_| Error::Unsupported)? - .to_le_bytes(); + let zrinit = Zrinit::CANFDX | Zrinit::CANOVIO | Zrinit::CANFC32; + let buffer_size = RECEIVER_BUFFER_SIZE.to_le_bytes(); Header::new( Encoding::ZHEX, Frame::ZRINIT, diff --git a/src-tauri/vendor/zmodem2/tests/protocol.rs b/src-tauri/vendor/zmodem2/tests/protocol.rs index fdb6b84f..12bd4b25 100644 --- a/src-tauri/vendor/zmodem2/tests/protocol.rs +++ b/src-tauri/vendor/zmodem2/tests/protocol.rs @@ -3,7 +3,10 @@ // Copyright (c) 2023-2026 Jarkko Sakkinen use rstest::rstest; -use zmodem2::{Encoding, Frame, Header, Receiver, ReceiverEvent, Sender, XON, ZDLE, ZPAD}; +use zmodem2::{ + Encoding, Frame, Header, Receiver, ReceiverEvent, Sender, SubpacketType, XON, ZDLE, ZPAD, + Zrinit, +}; #[rstest] #[case(Encoding::ZBIN, Frame::ZRQINIT, &[0; 4], &[ZPAD, ZDLE, Encoding::ZBIN as u8, 0, 0, 0, 0, 0, 0, 0])] @@ -38,6 +41,102 @@ pub fn test_header_read( assert!(Header::read(port) == Ok(Some(Header::new(encoding, frame, flags)))); } +fn zhex_header_from_wire(wire: &[u8]) -> Header { + assert!( + wire.starts_with(&[ZPAD, ZPAD, ZDLE, Encoding::ZHEX as u8]), + "expected ZHEX header, got {wire:?}" + ); + let mut port = &wire[3..]; + Header::read(&mut port) + .expect("read header") + .expect("header") +} + +fn write_zrinit(flags: &[u8; 4]) -> Vec { + let mut wire = Vec::new(); + Header::new(Encoding::ZHEX, Frame::ZRINIT, flags) + .write(&mut wire) + .expect("write zrinit") + .expect("complete zrinit"); + wire +} + +fn write_zrpos(offset: u32) -> Vec { + let mut wire = Vec::new(); + Header::new(Encoding::ZHEX, Frame::ZRPOS, &offset.to_le_bytes()) + .write(&mut wire) + .expect("write zrpos") + .expect("complete zrpos"); + wire +} + +fn has_subpacket_kind(wire: &[u8], kind: SubpacketType) -> bool { + wire.windows(2).any(|window| window == [ZDLE, kind as u8]) +} + +#[test] +fn test_receiver_zrinit_advertises_large_overlap_window() { + let receiver = Receiver::new().unwrap(); + let header = zhex_header_from_wire(receiver.drain_outgoing()); + let flags = header.count().to_le_bytes(); + let advertised_buffer_size = u16::from_le_bytes([flags[0], flags[1]]); + let caps = Zrinit::from_bits_truncate(flags[2] | flags[3]); + + assert!(header == Header::new(Encoding::ZHEX, Frame::ZRINIT, &flags)); + assert_eq!(advertised_buffer_size, u16::MAX); + assert!(caps.contains(Zrinit::CANFDX)); + assert!(caps.contains(Zrinit::CANOVIO)); + assert!(caps.contains(Zrinit::CANFC32)); +} + +#[test] +fn test_sender_requests_8k_file_chunks_after_large_zrinit() { + let mut sender = Sender::new().unwrap(); + sender.start_file(b"large.bin", 20_000).unwrap(); + sender.advance_outgoing(sender.drain_outgoing().len()); + + let flags = [0xff, 0xff, 0, (Zrinit::CANFDX | Zrinit::CANOVIO | Zrinit::CANFC32).bits()]; + sender.feed_incoming(&write_zrinit(&flags)).unwrap(); + sender.advance_outgoing(sender.drain_outgoing().len()); + + sender.feed_incoming(&write_zrpos(0)).unwrap(); + + let request = sender.poll_file().expect("file data request"); + assert_eq!(request.offset, 0); + assert_eq!(request.len, 8 * 1024); +} + +#[test] +fn test_sender_uses_streaming_subpackets_until_ack_boundary() { + let mut sender = Sender::new().unwrap(); + sender.start_file(b"window.bin", 70_000).unwrap(); + sender.advance_outgoing(sender.drain_outgoing().len()); + + let flags = [0, 0, 0, (Zrinit::CANFDX | Zrinit::CANOVIO | Zrinit::CANFC32).bits()]; + sender.feed_incoming(&write_zrinit(&flags)).unwrap(); + sender.advance_outgoing(sender.drain_outgoing().len()); + + sender.feed_incoming(&write_zrpos(0)).unwrap(); + + for packet_index in 0..8 { + let request = sender.poll_file().expect("file data request"); + assert_eq!(request.len, 8 * 1024); + sender.feed_file(&vec![0x55; request.len]).unwrap(); + let wire = sender.drain_outgoing().to_vec(); + + if packet_index == 0 { + assert!(has_subpacket_kind(&wire, SubpacketType::ZCRCG)); + } + if packet_index == 7 { + assert!(has_subpacket_kind(&wire, SubpacketType::ZCRCW)); + } + + sender.advance_outgoing(wire.len()); + } + + assert!(sender.poll_file().is_none()); +} + #[test] fn test_receive_malformed_header() { let mut receiver = Receiver::new().unwrap();