diff --git a/crates/rds-desktop/src/client.rs b/crates/rds-desktop/src/client.rs index 6288b25..6678fcc 100644 --- a/crates/rds-desktop/src/client.rs +++ b/crates/rds-desktop/src/client.rs @@ -1028,11 +1028,15 @@ async fn read_one_with_receipt( // discards its partial buffer on cancellation and hid whether even a // header or any media bytes arrived during observed freezes. let mut body = Vec::new(); + let mut hasher = receipts.as_ref().map(|_| blake3::Hasher::new()); let mut chunk = [0u8; 16 * 1024]; while let Some(count) = stream.read(&mut chunk).await.ok()? { if count > MAX_FRAME_BYTES.saturating_sub(body.len()) { return None; } + if let Some(hasher) = hasher.as_mut() { + hasher.update(&chunk[..count]); + } body.extend_from_slice(&chunk[..count]); body_bytes = body.len(); } @@ -1040,7 +1044,8 @@ async fn read_one_with_receipt( if body.is_empty() { return None; } - let digest = *blake3::hash(&body).as_bytes(); + let hasher = hasher?; + let digest = *hasher.finalize().as_bytes(); let obsolete = header.seq < next_seq.load(Ordering::Relaxed); // The proof belongs to the bounded wire reader, before ordered // decode or local IPC can backpressure completed payloads. diff --git a/crates/rds-desktop/src/receipts.rs b/crates/rds-desktop/src/receipts.rs index 1a701de..22b718a 100644 --- a/crates/rds-desktop/src/receipts.rs +++ b/crates/rds-desktop/src/receipts.rs @@ -30,7 +30,7 @@ pub(crate) async fn wait( } } -type Pending = BTreeMap>)>; +type Pending = BTreeMap, Option>)>; #[derive(Clone)] pub(crate) struct Receipts(Arc>); @@ -39,19 +39,35 @@ impl Receipts { pub(crate) fn new() -> Self { Self(Arc::new(Mutex::new(BTreeMap::new()))) } - pub(crate) fn register(&self, seq: u64, digest: [u8; 32]) -> Option { + /// Reserve one receipt slot before the payload is sent. The digest is + /// filled incrementally by the sender as bytes actually leave the + /// stream, avoiding a second full-payload hash pass. + pub(crate) fn register(&self, seq: u64) -> Option { let mut pending = self.0.lock().ok()?; if pending.len() >= crate::session::MAX_PENDING_FRAME_ACKS || pending.contains_key(&seq) { return None; } let (send, receive) = oneshot::channel(); - pending.insert(seq, (digest, Some(send))); + pending.insert(seq, (None, Some(send))); Some(Ticket { owner: self.clone(), seq, receive, }) } + pub(crate) fn set_digest(&self, seq: u64, digest: [u8; 32]) -> bool { + let Ok(mut pending) = self.0.lock() else { + return false; + }; + let Some((expected, _)) = pending.get_mut(&seq) else { + return false; + }; + if expected.is_some() { + return false; + } + *expected = Some(digest); + true + } pub(crate) fn confirm(&self, seq: u64, digest: &[u8; 32], obsolete: bool) -> bool { let Ok(mut pending) = self.0.lock() else { return false; @@ -59,7 +75,7 @@ impl Receipts { let Some((expected, send)) = pending.get_mut(&seq) else { return false; }; - if expected != digest { + if expected.as_ref() != Some(digest) { return false; } send.take().is_some_and(|send| { @@ -94,11 +110,18 @@ impl Drop for Ticket { #[cfg(test)] mod tests { use super::*; + + fn registered(receipts: &Receipts, seq: u64, digest: [u8; 32]) -> Ticket { + let ticket = receipts.register(seq).unwrap(); + assert!(receipts.set_digest(seq, digest)); + ticket + } + #[tokio::test] async fn only_an_exact_live_payload_proof_completes_once() { let receipts = Receipts::new(); let digest = *blake3::hash(b"controlled payload").as_bytes(); - let mut ticket = receipts.register(7, digest).unwrap(); + let mut ticket = registered(&receipts, 7, digest); assert!(!receipts.confirm(6, &digest, false)); assert!(!receipts.confirm(7, &[0; 32], false)); assert!( @@ -111,21 +134,32 @@ mod tests { ticket.received().await.unwrap(); drop(ticket); assert!(!receipts.confirm(7, &digest, false)); - assert!(receipts.register(7, digest).is_some()); + assert!(receipts.register(7).is_some()); + } + + #[test] + fn payload_proof_stays_unconfirmable_until_sender_sets_digest() { + let receipts = Receipts::new(); + let digest = *blake3::hash(b"payload").as_bytes(); + let _ticket = receipts.register(7).unwrap(); + assert!(!receipts.confirm(7, &digest, false)); + assert!(receipts.set_digest(7, digest)); + assert!(receipts.confirm(7, &digest, false)); + assert!(!receipts.set_digest(7, digest)); } #[test] fn pending_proofs_are_bounded_and_cancellation_frees_the_exact_slot() { let receipts = Receipts::new(); let mut tickets: Vec<_> = (0..crate::session::MAX_PENDING_FRAME_ACKS) - .map(|seq| receipts.register(seq as u64, [seq as u8; 32]).unwrap()) + .map(|seq| receipts.register(seq as u64).unwrap()) .collect(); - assert!(receipts.register(0, [9; 32]).is_none()); - assert!(receipts.register(100, [9; 32]).is_none()); + assert!(receipts.register(0).is_none()); + assert!(receipts.register(100).is_none()); let retired = tickets.pop().unwrap(); let seq = retired.seq; drop(retired); assert!(!receipts.confirm(seq, &[seq as u8; 32], false)); - assert!(receipts.register(100, [9; 32]).is_some()); + assert!(receipts.register(100).is_some()); drop(tickets); assert!(receipts.0.lock().unwrap().is_empty()); } @@ -133,7 +167,7 @@ mod tests { async fn a_delayed_transport_fin_cannot_hold_a_validated_payload_ticket() { let receipts = Receipts::new(); let digest = *blake3::hash(b"already received frame").as_bytes(); - let mut ticket = Some(receipts.register(1796, digest).unwrap()); + let mut ticket = Some(registered(&receipts, 1796, digest)); let transport = std::future::pending::>(); assert!(receipts.confirm(1796, &digest, false)); assert_eq!( @@ -152,7 +186,7 @@ mod tests { #[tokio::test] async fn transport_fin_alone_is_not_a_negotiated_payload_proof() { let receipts = Receipts::new(); - let mut ticket = Some(receipts.register(1, [1; 32]).unwrap()); + let mut ticket = Some(registered(&receipts, 1, [1; 32])); assert!( tokio::time::timeout( std::time::Duration::ZERO, @@ -180,7 +214,7 @@ mod tests { async fn obsolete_payload_proofs_and_session_retirement_do_not_confirm_fresh_work() { let old = Receipts::new(); let digest = [9; 32]; - let mut ticket = Some(old.register(0, digest).unwrap()); + let mut ticket = Some(registered(&old, 0, digest)); assert!(old.confirm(0, &digest, true)); assert_eq!( wait(&mut ticket, std::future::pending()).await.unwrap(), @@ -188,7 +222,7 @@ mod tests { ); drop(ticket); let fresh = Receipts::new(); - let mut ticket = fresh.register(0, digest).unwrap(); + let mut ticket = registered(&fresh, 0, digest); assert!(!old.confirm(0, &digest, false)); assert!( tokio::time::timeout(std::time::Duration::ZERO, ticket.received()) diff --git a/crates/rds-desktop/src/session.rs b/crates/rds-desktop/src/session.rs index 43cfca7..a8a2bb3 100644 --- a/crates/rds-desktop/src/session.rs +++ b/crates/rds-desktop/src/session.rs @@ -8,8 +8,10 @@ //! on the bi-directional control stream, which outranks every frame //! stream. +use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::task::{Context, Poll}; use std::time::{Duration, Instant}; use bytes::Bytes; @@ -1959,11 +1961,9 @@ async fn send_frame_inner( generation, mut key, } = produced; - let mut payload_ticket = match &delivery.payload_receipts { - Some(receipts) => match receipts.register( - produced.header.seq, - *blake3::hash(&produced.payload).as_bytes(), - ) { + let payload_receipts = delivery.payload_receipts.clone(); + let mut payload_ticket = match &payload_receipts { + Some(receipts) => match receipts.register(produced.header.seq) { Some(ticket) => Some(ticket), None => { tracing::warn!( @@ -2008,7 +2008,8 @@ async fn send_frame_inner( tracing::debug!("frame header write failed: {e}"); return SendOutcome::Failed; } - let outcome = match send_payload(stream, &produced, rx).await { + let mut payload_hasher = payload_receipts.as_ref().map(|_| blake3::Hasher::new()); + let outcome = match send_payload(stream, &produced, rx, payload_hasher.as_mut()).await { Ok(PayloadOutcome::Abandoned(next)) => { return SendOutcome::Superseded(next); } @@ -2023,6 +2024,17 @@ async fn send_frame_inner( return SendOutcome::Failed; } }; + if matches!(outcome, SendOutcome::Sent | SendOutcome::Done) + && let (Some(receipts), Some(hasher)) = (&payload_receipts, payload_hasher) + { + let digest = *hasher.finalize().as_bytes(); + if !receipts.set_digest(produced.header.seq, digest) { + tracing::debug!( + frame_seq = produced.header.seq, + "payload receipt digest was retired before send completed" + ); + } + } if let Err(e) = stream.finish() { // finish() reports an erased ClosedStream. Only a ready, exact peer // disposition can classify it as obsolete; never wait on other errors. @@ -2138,11 +2150,51 @@ enum PayloadOutcome { Abandoned(T), } +/// Hash payload bytes as the async writer actually accepts them. Receipt +/// hashing used to scan the complete payload before the write and therefore +/// paid a second Full HD memory pass on every frame. +struct DigestingWriter<'a, W> { + inner: &'a mut W, + hasher: Option<&'a mut blake3::Hasher>, +} + +impl AsyncWrite for DigestingWriter<'_, W> { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + match Pin::new(&mut *this.inner).poll_write(cx, buf) { + Poll::Ready(Ok(written)) => { + if let Some(hasher) = this.hasher.as_deref_mut() { + hasher.update(&buf[..written]); + } + Poll::Ready(Ok(written)) + } + other => other, + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut *self.get_mut().inner).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut *self.get_mut().inner).poll_shutdown(cx) + } +} + async fn send_payload>( stream: &mut W, produced: &Produced, rx: &mut mpsc::Receiver, + hasher: Option<&mut blake3::Hasher>, ) -> std::io::Result> { + let mut stream = DigestingWriter { + inner: stream, + hasher, + }; // write_all is not cancellation-safe: keep its progress alive across a // producer event. Starting another write_all would duplicate the prefix. let writing = stream.write_all(&produced.payload); @@ -2449,7 +2501,7 @@ mod tests { let frame = produced(0, keyframe); let mut wire = vec![0; 17]; { - let mut sending = Box::pin(send_payload(&mut writer, &frame, &mut rx)); + let mut sending = Box::pin(send_payload(&mut writer, &frame, &mut rx, None)); // Poll until the tiny stream buffer is full, then consume only // a prefix. The producer event must interrupt a partial write. poll_fn(|cx| { @@ -2513,7 +2565,7 @@ mod tests { let (mut writer, reader) = tokio::io::duplex(64); let (tx, mut rx) = mpsc::channel(1); let frame = produced(0, true); - let mut sending = Box::pin(send_payload(&mut writer, &frame, &mut rx)); + let mut sending = Box::pin(send_payload(&mut writer, &frame, &mut rx, None)); poll_fn(|cx| { assert!(sending.as_mut().poll(cx).is_pending()); Poll::Ready(()) @@ -2548,7 +2600,7 @@ mod tests { let (mut writer, _reader) = tokio::io::duplex(64); let (tx, mut rx) = mpsc::channel(1); let frame = produced(0, false); - let mut sending = Box::pin(send_payload(&mut writer, &frame, &mut rx)); + let mut sending = Box::pin(send_payload(&mut writer, &frame, &mut rx, None)); poll_fn(|cx| { assert!(sending.as_mut().poll(cx).is_pending()); Poll::Ready(())