From ec361a321de320aa106d9f0fe260db960a05317d Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 8 Aug 2026 04:08:57 +0000 Subject: [PATCH 1/7] feat(client): observe what the client sends, not only what it receives The receive side hands every decoded stanza to consumers through Event::RawNode, lease-gated so it costs nothing when unused. The send side had no equivalent: wait_for_sent_node is a filtered one-shot for a single expected stanza, and it only sees stanzas that go through marshal_node_for_send, so acks, delivery receipts and direct-encoded IQs were invisible even to that. A test could not assert what reached the wire without wrapping the transport, and a stanza the server rejected could not be read back without a rebuild. Event::SentFrame carries the marshaled plaintext of each frame the transport accepted. It is emitted from the noise sender task, which is the one point all five send paths cross (send_node, send_raw_bytes, send_raw_bytes_burst, and the ack and receipt workers through the burst) and the same chokepoint SessionStats already records wire bytes at. Observing bytes rather than nodes is what makes that possible: two of those paths never build a Node at all. Cost follows the house pattern. The gate is a relaxed atomic load per frame through the tap the sender already holds, so an unwatched send allocates exactly what it did before. A watched one costs two allocations per frame, neither of which scales with the stanza: the payload is handed over as a refcount bump on the caller's own buffer, pinned by a pointer assertion. The dispatch is caught, since it runs on the task every send depends on. wait_for_sent_node stays as it is. It observes the Node before marshalling, resolves before the write, and is filtered, so it answers a different question than a stream of sent bytes; reimplementing it on top of this one would have to unmarshal every frame to test its filter, and would widen what its callers see. send_raw_bytes keeps its bypass semantics, which are now only about node logging and sent-node waiters. --- agent_docs/observability.md | 11 ++ src/client.rs | 100 ++++++++++++ src/client/accessors.rs | 27 ++++ src/client/lifecycle.rs | 5 +- src/client/messaging.rs | 4 +- src/client/tests.rs | 197 ++++++++++++++++++++++- src/handshake.rs | 6 +- src/lib.rs | 4 +- src/plugins/mod.rs | 23 ++- src/socket/noise_socket.rs | 276 +++++++++++++++++++++++++++++++-- src/test_utils.rs | 74 ++++++++- tests/handshake_integration.rs | 19 +-- wacore/src/types/events.rs | 49 +++++- 13 files changed, 760 insertions(+), 35 deletions(-) diff --git a/agent_docs/observability.md b/agent_docs/observability.md index 00e58cbff..a5ee01eb8 100644 --- a/agent_docs/observability.md +++ b/agent_docs/observability.md @@ -32,6 +32,17 @@ chokepoints: transport write — post-noise wire bytes (frame header + AEAD tag included). - **Received**: the read loop (`node_io.rs`) per `DataReceived` batch. +That sent chokepoint is the *only* place every outbound frame crosses — five +distinct send paths reach it (`send_node`, `send_raw_bytes`, +`send_raw_bytes_burst`, and the ack/receipt workers through the burst), and +`send_raw_bytes` deliberately bypasses node logging and sent-node waiters — so +anything that has to see *everything* the client sends belongs there and nowhere +else. `Event::SentFrame` is the other thing wired into it +(`Client::acquire_sent_frame_forwarding()`, lease-gated like `RawNode`): it hands +over the marshaled plaintext of each frame the transport accepted. Both halves +travel to the socket as `SendObservers`, so the next observer plugs in there +instead of widening `do_handshake` again. + It also owns the activity timestamps the keepalive dead-socket watchdog reads: `last_data_received_ms` (one clock read per received transport event, plus one more when that event carries several frames, so a slow drain is not read as diff --git a/src/client.rs b/src/client.rs index 7d15045c8..0ed1c4e4e 100644 --- a/src/client.rs +++ b/src/client.rs @@ -103,6 +103,102 @@ impl Drop for RawNodeLease { } } +/// Lease that keeps sent-frame events enabled for one consumer. +/// +/// Dropping the final lease disables forwarding. The lease holds only a weak +/// client reference, so it cannot keep the client alive. +#[must_use = "dropping the lease immediately releases sent-frame forwarding"] +pub struct SentFrameLease { + client: std::sync::Weak, +} + +impl Drop for SentFrameLease { + fn drop(&mut self) { + let Some(client) = self.client.upgrade() else { + return; + }; + client.sent_frame_tap.release(); + } +} + +/// Publishes the plaintext frames that reached the transport as +/// [`Event::SentFrame`](wacore::types::events::Event::SentFrame). +/// +/// The client owns it and hands the noise sender a clone of the `Arc`, the same +/// way it hands over [`SessionStats`](wacore::stats::SessionStats): the gate has +/// to be readable from the one point every send crosses, and that task cannot +/// hold the client without keeping it alive. +pub(crate) struct SentFrameTap { + /// Number of consumers currently requesting the event. + forwarding: AtomicUsize, + bus: wacore::types::events::CoreEventBus, + /// Proves the no-lease path builds nothing, rather than only that it + /// dispatches nothing. + #[cfg(test)] + published: AtomicUsize, +} + +impl SentFrameTap { + pub(crate) fn new(bus: wacore::types::events::CoreEventBus) -> Self { + Self { + forwarding: AtomicUsize::new(0), + bus, + #[cfg(test)] + published: AtomicUsize::new(0), + } + } + + /// Enable forwarding for one consumer. The public door is + /// [`Client::acquire_sent_frame_forwarding`], which pairs this with a lease + /// that releases it on drop; a caller here owns that pairing itself. + pub(crate) fn acquire(&self) { + let incremented = self + .forwarding + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |count| { + count.checked_add(1) + }) + .is_ok(); + assert!(incremented, "sent-frame forwarding lease counter overflow"); + } + + pub(crate) fn release(&self) { + let previous = self.forwarding.fetch_sub(1, Ordering::Relaxed); + debug_assert!(previous > 0, "sent-frame forwarding lease underflow"); + } + + #[inline] + pub(crate) fn enabled(&self) -> bool { + self.forwarding.load(Ordering::Relaxed) != 0 + } + + /// Hand one frame to the observers. + /// + /// The dispatch is caught: a consumer that only watches must not be able to + /// take the send pipeline down with it, and this runs on the noise sender + /// task, whose death would end every send on the connection. A handler that + /// *blocks* still stalls sends, which is the same contract every event + /// handler already has on the read loop. + pub(crate) fn publish(&self, plaintext: bytes::Bytes) { + #[cfg(test)] + self.published.fetch_add(1, Ordering::Relaxed); + let dispatch = std::panic::AssertUnwindSafe(|| { + self.bus.dispatch(Event::SentFrame( + wacore::types::events::SentFrame::builder() + .plaintext(plaintext) + .build(), + )); + }); + if std::panic::catch_unwind(dispatch).is_err() { + warn!("A sent-frame observer panicked; the send pipeline is unaffected."); + } + } + + #[cfg(test)] + pub(crate) fn published(&self) -> usize { + self.published.load(Ordering::Relaxed) + } +} + /// Filter for matching incoming stanzas (nodes) by tag and attributes. /// /// Used with [`Client::wait_for_node`] to wait for specific stanzas. @@ -1453,6 +1549,10 @@ pub struct Client { /// forwarding. decrypted_payload_forwarding: AtomicUsize, + /// Gate and publisher for `Event::SentFrame`. Behind an `Arc` because the + /// noise sender task reads it; see [`SentFrameTap`]. + pub(crate) sent_frame_tap: Arc, + /// Stanza interceptors, behind the same copy-on-write snapshot the event /// bus uses: reading one costs a refcount bump, so the read loop allocates /// nothing per stanza. Registering is the rare side, and pays the copy. diff --git a/src/client/accessors.rs b/src/client/accessors.rs index 696f80cf1..6830f6332 100644 --- a/src/client/accessors.rs +++ b/src/client/accessors.rs @@ -88,6 +88,33 @@ impl Client { self.decrypted_payload_forwarding.load(Ordering::Relaxed) != 0 } + /// Acquire sent-frame forwarding for one consumer. + /// + /// [`Event::SentFrame`] stays enabled until every acquired lease is dropped. + /// While none is held nothing is emitted and nothing is cloned: the send + /// path costs one relaxed atomic load. + /// + /// This is the outbound counterpart of + /// [`acquire_raw_node_forwarding`](Self::acquire_raw_node_forwarding), and + /// unlike [`wait_for_sent_node`](Self::wait_for_sent_node) it is neither + /// filtered nor one-shot and covers every send path, including the ones that + /// never build a `Node`. + /// + /// [`Event::SentFrame`]: wacore::types::events::Event::SentFrame + pub fn acquire_sent_frame_forwarding(self: &Arc) -> SentFrameLease { + self.sent_frame_tap.acquire(); + SentFrameLease { + client: Arc::downgrade(self), + } + } + + /// Only tests ask this: the send path reads the gate through the tap the + /// noise sender already holds, not through the client. + #[cfg(test)] + pub(crate) fn sent_frame_forwarding_enabled(&self) -> bool { + self.sent_frame_tap.enabled() + } + /// Register an interceptor that sees each decoded stanza before the /// built-in pipeline, and may take it. /// diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index c556dc825..45a4234a1 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -315,6 +315,7 @@ impl Client { let (tx, rx) = async_channel::bounded(32); let device_topology = device_topology::DeviceTopology::new(); + let sent_frame_tap = Arc::new(SentFrameTap::new(core.event_bus.clone())); let this = Self { runtime: runtime.clone(), core, @@ -507,6 +508,7 @@ impl Client { alloc_meter: std::sync::OnceLock::new(), raw_node_forwarding: AtomicUsize::new(0), decrypted_payload_forwarding: AtomicUsize::new(0), + sent_frame_tap, stanza_interceptors: std::sync::RwLock::new(Arc::new(Vec::new())), stanza_interceptor_count: AtomicUsize::new(0), next_interceptor_id: AtomicU64::new(0), @@ -823,7 +825,8 @@ impl Client { &self.ik_handshake_failures, transport.clone(), &mut transport_events, - Some(self.stats.clone()), + crate::socket::noise_socket::SendObservers::with_stats(self.stats.clone()) + .with_sent_frames(self.sent_frame_tap.clone()), ) .await { diff --git a/src/client/messaging.rs b/src/client/messaging.rs index 67d826a48..702f9a5ba 100644 --- a/src/client/messaging.rs +++ b/src/client/messaging.rs @@ -10,7 +10,9 @@ impl Client { /// cause the server to close the connection. /// /// This bypasses node logging and `sent_node_waiter` resolution — use - /// [`send_node`](Client::send_node) for normal stanza sending. + /// [`send_node`](Client::send_node) for normal stanza sending. It is still + /// observed: `Event::SentFrame` is emitted from the noise sender, past every + /// bypass here. pub async fn send_raw_bytes(&self, plaintext: Vec) -> Result<(), ClientError> { let noise_socket = self.get_noise_socket()?; // Wire bytes and the last-sent timestamp are recorded by the noise diff --git a/src/client/tests.rs b/src/client/tests.rs index 89660e418..f05098d03 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -3506,17 +3506,210 @@ async fn install_test_noise_socket( use crate::socket::NoiseSocket; use wacore::handshake::NoiseCipher; - let key = [0u8; 32]; - let noise_socket = NoiseSocket::new( + let key = TEST_NOISE_KEY; + // Wired to the client's own observers, like the real socket: a test that + // watches its sends needs the same plumbing production has. + let noise_socket = NoiseSocket::with_observers( runtime, transport, NoiseCipher::new(&key).expect("valid key"), NoiseCipher::new(&key).expect("valid key"), + crate::socket::noise_socket::SendObservers::default() + .with_sent_frames(client.sent_frame_tap.clone()), ); *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); client.set_connected_for_test(true); } +/// The key `install_test_noise_socket` builds its socket with, so a test can +/// decrypt what the client wrote. +const TEST_NOISE_KEY: [u8; 32] = [0u8; 32]; + +/// Every distinct way a stanza leaves the client, driven end to end, with the +/// observer's view compared against the transport's. +/// +/// `send_node` marshals and resolves sent-node waiters; the ack, receipt and +/// direct-encoded IQ paths hand pre-marshaled bytes straight to the socket and +/// were invisible to those waiters; the ack and receipt workers reach the wire +/// only through the burst. All of them cross the noise sender, which is why one +/// observation point covers the lot. +#[tokio::test] +async fn every_send_path_is_observed_exactly_as_it_reached_the_wire() { + use crate::transport::mock::CapturingMockTransport; + + let client = crate::test_utils::create_test_client().await; + let transport = Arc::new(CapturingMockTransport::new()); + install_test_noise_socket( + &client, + transport.clone(), + Arc::new(crate::runtime_impl::TokioRuntime), + ) + .await; + + let recorder = Arc::new(crate::test_utils::SentFrameRecorder::default()); + let _subscription = client.subscribe_handler(recorder.clone()); + let _lease = client.acquire_sent_frame_forwarding(); + + // 1. send_node: the only path that builds a Node the caller could inspect. + let presence = NodeBuilder::new("presence") + .attr("type", "available") + .attr("name", "observer") + .build(); + client + .send_node(presence.clone()) + .await + .expect("presence must send"); + + // 2. send_raw_bytes, through a real caller of it: the ack path documents + // that it bypasses node logging and sent-node waiters. + let incoming = NodeBuilder::new("notification") + .attr("id", "OBSERVED-1") + .attr("type", "w:gp2") + .attr("from", "5550000@g.us") + .build(); + client + .send_ack_for(&incoming.as_node_ref()) + .await + .expect("ack must send"); + + // 3. send_raw_bytes_burst, the shape the ack and receipt workers use: two + // frames coalesced into a single transport write. + let mut frames = vec![ + wacore_binary::marshal::marshal_exact( + &NodeBuilder::new("iq").attr("id", "BURST-1").build(), + ) + .expect("marshal"), + wacore_binary::marshal::marshal_exact( + &NodeBuilder::new("iq").attr("id", "BURST-2").build(), + ) + .expect("marshal"), + ]; + let mut results = Vec::new(); + client + .send_raw_bytes_burst(&mut frames, &mut results) + .await + .expect("burst must send"); + assert!(results.iter().all(|result| result.is_ok())); + + let wire = crate::test_utils::decrypt_wire_frames(&transport.sent(), &TEST_NOISE_KEY); + assert_eq!(wire.len(), 4, "four frames must have reached the transport"); + let observed: Vec> = recorder + .frames() + .iter() + .map(|frame| frame.to_vec()) + .collect(); + assert_eq!( + observed, wire, + "every send path must be observed, byte for byte and in wire order" + ); + + // The bytes are the stanza, not a rendering of it: the first frame decodes + // back to the node that was sent. + let decoded = wacore_binary::marshal::unmarshal_ref(&observed[0][1..]) + .expect("an observed frame must decode as the stanza it carried"); + assert_eq!(decoded.tag.as_ref(), "presence"); + assert_eq!( + decoded.attrs().optional_string("name").as_deref(), + Some("observer") + ); +} + +/// While nobody holds a lease the send path publishes nothing and builds +/// nothing, and releasing the last lease puts it back to that state. +#[tokio::test] +async fn sends_are_unobserved_until_a_consumer_asks() { + use crate::transport::mock::CapturingMockTransport; + + let client = crate::test_utils::create_test_client().await; + let transport = Arc::new(CapturingMockTransport::new()); + install_test_noise_socket( + &client, + transport.clone(), + Arc::new(crate::runtime_impl::TokioRuntime), + ) + .await; + + let recorder = Arc::new(crate::test_utils::SentFrameRecorder::default()); + let _subscription = client.subscribe_handler(recorder.clone()); + + assert!( + !client.sent_frame_forwarding_enabled(), + "forwarding must be off until a consumer acquires it" + ); + client + .send_node(NodeBuilder::new("presence").build()) + .await + .expect("presence must send"); + assert_eq!(client.sent_frame_tap.published(), 0); + assert!(recorder.frames().is_empty()); + + let lease = client.acquire_sent_frame_forwarding(); + assert!(client.sent_frame_forwarding_enabled()); + client + .send_node(NodeBuilder::new("presence").build()) + .await + .expect("presence must send"); + assert_eq!(recorder.frames().len(), 1); + + drop(lease); + assert!( + !client.sent_frame_forwarding_enabled(), + "the last lease dropping must disable forwarding again" + ); + client + .send_node(NodeBuilder::new("presence").build()) + .await + .expect("presence must send"); + assert_eq!( + recorder.frames().len(), + 1, + "no frame may be observed after the lease is gone" + ); + assert_eq!(client.sent_frame_tap.published(), 1); +} + +/// An observer that panics must not cost the client its send path: the dispatch +/// runs on the noise sender task, and an unwinding panic there would end every +/// send on the connection. +#[tokio::test] +async fn a_panicking_observer_leaves_the_client_sending() { + use crate::transport::mock::CapturingMockTransport; + + struct PanickingObserver; + impl wacore::types::events::EventHandler for PanickingObserver { + fn handle_event(&self, _event: Arc) { + panic!("observer panics on every frame"); + } + fn interest(&self) -> wacore::types::events::EventInterest { + wacore::types::events::EventInterest::of(&[wacore::types::events::EventKind::SentFrame]) + } + } + + let client = crate::test_utils::create_test_client().await; + let transport = Arc::new(CapturingMockTransport::new()); + install_test_noise_socket( + &client, + transport.clone(), + Arc::new(crate::runtime_impl::TokioRuntime), + ) + .await; + + let _subscription = client.subscribe_handler(Arc::new(PanickingObserver)); + let _lease = client.acquire_sent_frame_forwarding(); + + for attempt in 0..3 { + client + .send_node(NodeBuilder::new("presence").attr("t", attempt).build()) + .await + .unwrap_or_else(|e| panic!("send {attempt} must survive the observer: {e:?}")); + } + assert_eq!( + transport.sent_count(), + 3, + "every stanza must still reach the wire" + ); +} + fn receipt_test_info(id: &str) -> Arc { Arc::new(crate::types::message::MessageInfo { id: id.to_string(), diff --git a/src/handshake.rs b/src/handshake.rs index 6fbcb63c5..68394c20f 100644 --- a/src/handshake.rs +++ b/src/handshake.rs @@ -178,7 +178,7 @@ pub async fn do_handshake( ik_handshake_failures: &AtomicU32, transport: Arc, transport_events: &mut async_channel::Receiver, - stats: Option>, + observers: crate::socket::noise_socket::SendObservers, ) -> Result> { let device_snapshot = persistence_manager.get_device_snapshot(); let now_secs = wacore::time::now_secs(); @@ -225,12 +225,12 @@ pub async fn do_handshake( .await; } ik_handshake_failures.store(0, Ordering::Release); - Ok(Arc::new(NoiseSocket::with_stats( + Ok(Arc::new(NoiseSocket::with_observers( runtime, transport, success.write_cipher, success.read_cipher, - stats, + observers, ))) } Err(e) => { diff --git a/src/lib.rs b/src/lib.rs index 0e12a2d7f..e7737ab61 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -128,6 +128,7 @@ pub use client::{ pub use client::{CallError, Voip}; pub use client::{ Client, ClientBuild, ClientBuilder, ClientBuilderError, DecryptedPayloadLease, RawNodeLease, + SentFrameLease, }; #[cfg(feature = "client-lifecycle")] #[cfg_attr(docsrs, doc(cfg(feature = "client-lifecycle")))] @@ -238,7 +239,8 @@ pub mod version; pub mod prelude { pub use crate::bot::{Bot, BotBuilder, BotHandle, EventDelivery, MessageContext}; pub use crate::client::{ - Client, ClientBuilder, ClientBuilderError, ClientError, DecryptedPayloadLease, RawNodeLease, + Client, ClientBuilder, ClientBuilderError, ClientError, DecryptedPayloadLease, + RawNodeLease, SentFrameLease, }; #[cfg(feature = "client-lifecycle")] #[cfg_attr(docsrs, doc(cfg(feature = "client-lifecycle")))] diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index 01e35b550..52aae866f 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -35,6 +35,7 @@ use crate::Client; use crate::client::interceptor::{Interception, InterceptorHandle, StanzaInterceptor}; use crate::client::{ ClientLifecycle, ConnectionScope, ConnectionScopeState, DecryptedPayloadLease, RawNodeLease, + SentFrameLease, }; use crate::request::IqError; use crate::send::{SendError, SendResult}; @@ -492,6 +493,7 @@ impl Drop for TaskLease { struct GatedForwarding { raw_node: Option, decrypted_payload: Option, + sent_frame: Option, } impl GatedForwarding { @@ -502,6 +504,7 @@ impl GatedForwarding { fn is_short_for(&self, interest: EventInterest) -> bool { (interest.wants(EventKind::RawNode) && self.raw_node.is_none()) || (interest.wants(EventKind::DecryptedPayload) && self.decrypted_payload.is_none()) + || (interest.wants(EventKind::SentFrame) && self.sent_frame.is_none()) } /// Acquire what `interest` needs and this does not hold yet. @@ -515,6 +518,8 @@ impl GatedForwarding { decrypted_payload: (interest.wants(EventKind::DecryptedPayload) && self.decrypted_payload.is_none()) .then(|| client.acquire_decrypted_payload_forwarding()), + sent_frame: (interest.wants(EventKind::SentFrame) && self.sent_frame.is_none()) + .then(|| client.acquire_sent_frame_forwarding()), } } @@ -522,6 +527,7 @@ impl GatedForwarding { fn commit(&mut self, acquired: Self) { self.raw_node = self.raw_node.take().or(acquired.raw_node); self.decrypted_payload = self.decrypted_payload.take().or(acquired.decrypted_payload); + self.sent_frame = self.sent_frame.take().or(acquired.sent_frame); } /// Give up what `interest` no longer asks for. @@ -537,6 +543,9 @@ impl GatedForwarding { decrypted_payload: (!interest.wants(EventKind::DecryptedPayload)) .then(|| self.decrypted_payload.take()) .flatten(), + sent_frame: (!interest.wants(EventKind::SentFrame)) + .then(|| self.sent_frame.take()) + .flatten(), } } } @@ -5996,17 +6005,19 @@ mod tests { "the kind that is no longer wanted releases its own lease" ); - // Both at once, then each removed on its own. + // All at once, then each removed on its own. assert!( subscription .update_interest(EventInterest::of(&[ EventKind::RawNode, EventKind::DecryptedPayload, + EventKind::SentFrame, ])) .expect("interest update") ); assert!(client.raw_node_forwarding_enabled()); assert!(client.decrypted_payload_forwarding_enabled()); + assert!(client.sent_frame_forwarding_enabled()); assert!( subscription @@ -6015,10 +6026,20 @@ mod tests { ); assert!(client.raw_node_forwarding_enabled(), "kept"); assert!(!client.decrypted_payload_forwarding_enabled(), "released"); + assert!(!client.sent_frame_forwarding_enabled(), "released"); + + assert!( + subscription + .update_interest(EventInterest::of(&[EventKind::SentFrame])) + .expect("interest update") + ); + assert!(client.sent_frame_forwarding_enabled()); + assert!(!client.raw_node_forwarding_enabled(), "released"); assert!(subscription.unsubscribe()); assert!(!client.raw_node_forwarding_enabled()); assert!(!client.decrypted_payload_forwarding_enabled()); + assert!(!client.sent_frame_forwarding_enabled()); } #[tokio::test] diff --git a/src/socket/noise_socket.rs b/src/socket/noise_socket.rs index 68fb74a9f..7b296bd4b 100644 --- a/src/socket/noise_socket.rs +++ b/src/socket/noise_socket.rs @@ -92,6 +92,36 @@ struct SendJob { response_tx: oneshot::Sender, } +/// What a socket reports its sends to. Both halves belong to the `Client`; a +/// VoIP relay socket and most tests pass [`Default`], reporting to neither. +/// +/// One struct rather than one parameter each: the observation point is shared, +/// so the next thing that wants to watch sends plugs in here instead of widening +/// every constructor between here and `connect()` again. +#[derive(Default, Clone)] +pub struct SendObservers { + /// Wire-byte accounting, recorded after the transport write. + stats: Option>, + /// Publisher for the plaintext of each frame that reached the transport. + sent_frames: Option>, +} + +impl SendObservers { + /// Report wire bytes into `stats` and nothing else. + pub fn with_stats(stats: Arc) -> Self { + Self { + stats: Some(stats), + sent_frames: None, + } + } + + /// Also publish each sent frame's plaintext through `tap`. + pub(crate) fn with_sent_frames(mut self, tap: Arc) -> Self { + self.sent_frames = Some(tap); + self + } +} + pub struct NoiseSocket { read_key: Arc, read_counter: Arc, @@ -112,18 +142,24 @@ impl NoiseSocket { write_key: NoiseCipher, read_key: NoiseCipher, ) -> Self { - Self::with_stats(runtime, transport, write_key, read_key, None) + Self::with_observers( + runtime, + transport, + write_key, + read_key, + SendObservers::default(), + ) } - /// Like [`Self::new`], recording sent frames into `stats` (the main WA - /// session socket passes the client's [`SessionStats`](wacore::stats::SessionStats); VoIP relay - /// sockets and tests pass `None`). - pub fn with_stats( + /// Like [`Self::new`], reporting each send to `observers` (the main WA + /// session socket passes the client's; VoIP relay sockets and most tests + /// report to nothing). + pub fn with_observers( runtime: Arc, transport: Arc, write_key: NoiseCipher, read_key: NoiseCipher, - stats: Option>, + observers: SendObservers, ) -> Self { let write_key = Arc::new(write_key); let read_key = Arc::new(read_key); @@ -142,7 +178,7 @@ impl NoiseSocket { transport_clone, write_key_clone, send_job_rx, - stats, + observers, ))); Self { @@ -161,8 +197,9 @@ impl NoiseSocket { transport: Arc, write_key: Arc, send_job_rx: async_channel::Receiver, - stats: Option>, + observers: SendObservers, ) { + let SendObservers { stats, sent_frames } = observers; let mut write_counter: u32 = 0; // BytesMut: split().freeze() yields a zero-copy Bytes while retaining // the underlying allocation for the next frame. @@ -182,6 +219,12 @@ impl NoiseSocket { // drops its response channel, which the caller sees as a closed sender: // a held-over job can be lost, but it can never hang its caller. let mut carry_over: Option = None; + // Plaintexts of this batch's frames, held only while a consumer is + // watching: each entry is a refcount bump on the buffer the caller + // marshalled, and the whole `Vec` stays empty (unallocated) otherwise. + // They are kept until after the write so what is published is what the + // transport actually accepted. + let mut observed: Vec = Vec::new(); loop { let job = match carry_over.take() { @@ -207,6 +250,12 @@ impl NoiseSocket { let mut job = job; loop { let response_tx = job.response_tx; + // Cloned before the plaintext is consumed, dropped again if the + // frame never makes it into the buffer. + let to_observe = match sent_frames.as_deref() { + Some(tap) if tap.enabled() => Some(job.plaintext.clone()), + _ => None, + }; match Self::encrypt_frame_into( &runtime, &write_key, @@ -216,7 +265,12 @@ impl NoiseSocket { ) .await { - Ok(wire_bytes) => waiters.push((response_tx, wire_bytes)), + Ok(wire_bytes) => { + waiters.push((response_tx, wire_bytes)); + if let Some(plaintext) = to_observe { + observed.push(plaintext); + } + } Err(e) => { // The counter is untouched on this frame, and every // frame already in the buffer must still go out so the @@ -268,11 +322,19 @@ impl NoiseSocket { stats.record_frame_sent(*wire_bytes); } } + if let Some(tap) = sent_frames.as_deref() { + for plaintext in observed.drain(..) { + tap.publish(plaintext); + } + } Ok(()) } Err(e) => Err(EncryptSendError::transport(e)), } }; + // A write that failed says nothing about what the peer received, so + // its frames are not reported as sent. + observed.clear(); { // Crypto and framing failures are rejected before any byte @@ -1436,12 +1498,12 @@ mod tests { let key = [0u8; 32]; let stats = Arc::new(wacore::stats::SessionStats::new()); - let socket = NoiseSocket::with_stats( + let socket = NoiseSocket::with_observers( Arc::new(crate::runtime_impl::TokioRuntime), transport.clone(), NoiseCipher::new(&key).expect("32-byte key"), NoiseCipher::new(&key).expect("32-byte key"), - Some(stats.clone()), + SendObservers::with_stats(stats.clone()), ); for size in [0usize, 100, 5000] { @@ -1508,4 +1570,196 @@ mod tests { "Payload above inline threshold should encrypt successfully" ); } + + use crate::test_utils::SentFrameRecorder; + + /// A tap wired to a fresh bus, forwarding enabled, plus the observer behind + /// it and the subscription that has to outlive the test. + fn watched_tap() -> ( + Arc, + Arc, + wacore::types::events::Subscription, + ) { + let bus = wacore::types::events::CoreEventBus::new(); + let observer = Arc::new(SentFrameRecorder::default()); + let subscription = bus.subscribe_handler(observer.clone()); + let tap = Arc::new(crate::client::SentFrameTap::new(bus)); + tap.acquire(); + (tap, observer, subscription) + } + + fn socket_watched_by( + transport: Arc, + tap: Arc, + ) -> NoiseSocket { + let key = [0x9Cu8; 32]; + NoiseSocket::with_observers( + Arc::new(crate::runtime_impl::TokioRuntime), + transport, + NoiseCipher::new(&key).expect("32-byte key"), + NoiseCipher::new(&key).expect("32-byte key"), + SendObservers::default().with_sent_frames(tap), + ) + } + + /// The observer receives the caller's own buffer. Handing over a copy would + /// double the cost of every send the moment anyone watched, which is the + /// difference between a recorder a consumer can leave on and one it cannot. + #[tokio::test] + async fn an_observed_frame_is_the_buffer_the_caller_handed_over() { + let (tap, observer, _subscription) = watched_tap(); + let socket = + socket_watched_by(Arc::new(crate::transport::mock::MockTransport), tap.clone()); + + let payload = bytes::Bytes::from(vec![0x5Au8; 4096]); + let payload_ptr = payload.as_ptr(); + socket + .encrypt_and_send(payload.clone()) + .await + .expect("send must succeed"); + + let observed = observer.frames(); + assert_eq!(observed.len(), 1, "the frame must be observed exactly once"); + assert_eq!( + observed[0].as_ptr(), + payload_ptr, + "the observer must receive the sent buffer itself, not a copy of it" + ); + assert_eq!(&observed[0][..], &payload[..]); + } + + /// Nothing is observed while no consumer holds forwarding, and nothing is + /// built either: the tap counts its publications, so this also fails on a + /// build that is merely thrown away. + #[tokio::test] + async fn an_unwatched_send_publishes_nothing() { + let bus = wacore::types::events::CoreEventBus::new(); + let observer = Arc::new(SentFrameRecorder::default()); + let _subscription = bus.subscribe_handler(observer.clone()); + let tap = Arc::new(crate::client::SentFrameTap::new(bus)); + let socket = + socket_watched_by(Arc::new(crate::transport::mock::MockTransport), tap.clone()); + + assert!(!tap.enabled(), "no lease has been acquired"); + for _ in 0..4 { + socket + .encrypt_and_send(bytes::Bytes::from(vec![1u8; 64])) + .await + .expect("send must succeed"); + } + + assert_eq!( + tap.published(), + 0, + "nothing may be built without a consumer" + ); + assert!(observer.frames().is_empty()); + } + + /// A frame the transport refused is not reported as sent: observing after + /// the write is what makes what arrives equal what left. + #[tokio::test] + async fn a_refused_write_is_not_observed() { + let (tap, observer, _subscription) = watched_tap(); + let transport = Arc::new(crate::transport::mock::CapturingMockTransport::new()); + transport.fail_next_sends(1); + let socket = socket_watched_by(transport.clone(), tap.clone()); + + socket + .encrypt_and_send(bytes::Bytes::from(vec![2u8; 64])) + .await + .expect_err("injected transport failure"); + + assert_eq!(tap.published(), 0); + assert!(observer.frames().is_empty()); + } + + /// An observer that panics must not take the sender task with it, or one + /// consumer watching would end every send on the connection. + #[tokio::test] + async fn a_panicking_observer_does_not_break_the_send() { + struct PanickingObserver; + impl wacore::types::events::EventHandler for PanickingObserver { + fn handle_event(&self, _event: Arc) { + panic!("observer panics on every frame"); + } + fn interest(&self) -> wacore::types::events::EventInterest { + wacore::types::events::EventInterest::of(&[ + wacore::types::events::EventKind::SentFrame, + ]) + } + } + + let bus = wacore::types::events::CoreEventBus::new(); + let _subscription = bus.subscribe_handler(Arc::new(PanickingObserver)); + let tap = Arc::new(crate::client::SentFrameTap::new(bus)); + tap.acquire(); + let transport = Arc::new(crate::transport::mock::CapturingMockTransport::new()); + let socket = socket_watched_by(transport.clone(), tap); + + for attempt in 0..3u8 { + socket + .encrypt_and_send(bytes::Bytes::from(vec![attempt; 32])) + .await + .unwrap_or_else(|e| panic!("send {attempt} must survive the observer: {e:?}")); + } + assert_eq!( + transport.sent_count(), + 3, + "every frame must still reach the transport" + ); + } + + /// What watching costs per frame, measured rather than argued. The idle path + /// must not move at all, and a watched send must not copy the payload: the + /// delta is the `Bytes` promotion to a shared handle plus the `Arc` + /// the bus dispatches, neither of which scales with the stanza. + #[tokio::test] + async fn watching_costs_a_constant_per_frame_and_idling_costs_nothing() { + async fn min_allocs_per_send(socket: &NoiseSocket, payload_len: usize) -> u64 { + let mut min = u64::MAX; + // Enough windows that one lands without a sibling test thread + // allocating inside it; the buffers the sender reuses have long + // stopped growing by then. + for _ in 0..2_000 { + let before = crate::test_alloc::ALLOCS.load(Ordering::Relaxed); + socket + .encrypt_and_send(bytes::Bytes::from(vec![3u8; payload_len])) + .await + .expect("send must succeed"); + let after = crate::test_alloc::ALLOCS.load(Ordering::Relaxed); + min = min.min(after - before); + } + min + } + + let (tap, _observer, _subscription) = watched_tap(); + let idle_tap = Arc::new(crate::client::SentFrameTap::new( + wacore::types::events::CoreEventBus::new(), + )); + + let idle = min_allocs_per_send( + &socket_watched_by(Arc::new(crate::transport::mock::MockTransport), idle_tap), + 256, + ) + .await; + let watched = min_allocs_per_send( + &socket_watched_by(Arc::new(crate::transport::mock::MockTransport), tap), + 256, + ) + .await; + + // 4 = the payload each window allocates for itself, plus the three the + // send path already cost before any of this existed. A ceiling rather + // than an equality: this must fail on a regression, not on a saving. + assert!( + idle <= 4, + "an unwatched send must cost what it always did, got {idle}" + ); + assert_eq!( + watched - idle, + 2, + "watching must cost a constant per frame (idle {idle}, watched {watched})" + ); + } } diff --git a/src/test_utils.rs b/src/test_utils.rs index 81e082dcd..679b993a4 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -14,6 +14,69 @@ pub fn node_to_owned_ref(node: &Node) -> Arc { } } +/// Records every [`Event::SentFrame`] +/// it is dispatched, keeping the `Bytes` it arrived with so a test can check the +/// pointer as well as the contents. +#[cfg(test)] +#[derive(Default)] +pub(crate) struct SentFrameRecorder { + frames: Mutex>, +} + +#[cfg(test)] +impl SentFrameRecorder { + pub(crate) fn frames(&self) -> Vec { + self.frames.lock().expect("recorded frames mutex").clone() + } +} + +#[cfg(test)] +impl EventHandler for SentFrameRecorder { + fn handle_event(&self, event: Arc) { + if let Event::SentFrame(sent) = &*event { + self.frames + .lock() + .expect("recorded frames mutex") + .push(sent.plaintext.clone()); + } + } + + fn interest(&self) -> EventInterest { + EventInterest::of(&[EventKind::SentFrame]) + } +} + +/// Splits a run of length-prefixed noise frames and decrypts each under the +/// counter its position implies, yielding the plaintexts in wire order. +/// +/// The counter is the AES-GCM nonce, so this doubles as an order check: a frame +/// read out of position cannot authenticate. +#[cfg(test)] +pub(crate) fn decrypt_wire_frames(writes: &[bytes::Bytes], key: &[u8; 32]) -> Vec> { + use wacore::handshake::NoiseCipher; + + let cipher = NoiseCipher::new(key).expect("32-byte key"); + let mut plaintexts = Vec::new(); + for write in writes { + let mut wire = &write[..]; + while !wire.is_empty() { + let mut len = 0usize; + for byte in &wire[..wacore::framing::FRAME_LENGTH_SIZE] { + len = (len << 8) | *byte as usize; + } + let body = + &wire[wacore::framing::FRAME_LENGTH_SIZE..wacore::framing::FRAME_LENGTH_SIZE + len]; + let mut body = bytes::BytesMut::from(body); + cipher + .decrypt_in_place_with_counter(plaintexts.len() as u32, &mut body) + .expect("each captured frame must decrypt under its own counter"); + plaintexts.push(body.to_vec()); + wire = &wire[wacore::framing::FRAME_LENGTH_SIZE + len..]; + } + } + plaintexts +} + pub async fn wait_for_lock_waiter(lock: &Arc>, baseline: usize) { poll_until("a task to reach the contested lock", || { Arc::strong_count(lock) > baseline @@ -248,7 +311,7 @@ pub async fn seed_peer_session(client: &Arc, peer: &Jid) { } use std::sync::Mutex; -use wacore::types::events::{Event, EventHandler}; +use wacore::types::events::{Event, EventHandler, EventInterest, EventKind}; #[derive(Default)] pub struct TestEventCollector { @@ -320,14 +383,15 @@ pub(crate) async fn create_iq_test_client() -> ( ) .await; - // Wired to the client's stats like the real socket is, so per-frame - // bookkeeping is part of what tests observe. - let noise_socket = crate::socket::NoiseSocket::with_stats( + // Wired to the client's own observers like the real socket is, so per-frame + // bookkeeping and sent-frame forwarding are part of what tests observe. + let noise_socket = crate::socket::NoiseSocket::with_observers( Arc::new(TokioRuntime), transport.clone() as Arc, NoiseCipher::new(&[0u8; 32]).expect("32-byte key"), NoiseCipher::new(&[0u8; 32]).expect("32-byte key"), - Some(client.stats.clone()), + crate::socket::noise_socket::SendObservers::with_stats(client.stats.clone()) + .with_sent_frames(client.sent_frame_tap.clone()), ); *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); client.set_connected_for_test(true); diff --git a/tests/handshake_integration.rs b/tests/handshake_integration.rs index 7b1a62ba6..7d2723d99 100644 --- a/tests/handshake_integration.rs +++ b/tests/handshake_integration.rs @@ -39,6 +39,7 @@ use wacore_noise::test_util::build_cert_chain_bytes; use whatsapp_rust::waproto::whatsapp as wa; use whatsapp_rust::handshake::do_handshake; +use whatsapp_rust::socket::noise_socket::SendObservers; use whatsapp_rust::transport::{Transport, TransportEvent}; /// In-process responder driving Noise XX or IK from the server side. @@ -311,7 +312,7 @@ async fn cold_start_xx_then_cached_ik_reconnect() { counter2.as_ref(), transport.clone(), &mut events_rx, - None, + SendObservers::default(), ) .await; @@ -355,7 +356,7 @@ async fn cold_start_xx_then_cached_ik_reconnect() { counter3.as_ref(), transport2.clone(), &mut events_rx2, - None, + SendObservers::default(), ) .await; @@ -455,7 +456,7 @@ async fn post_xxfallback_failure_does_not_invalidate_ik_cache() { counter.as_ref(), transport.clone(), &mut events_rx, - None, + SendObservers::default(), ) .await; task.await.unwrap(); @@ -525,7 +526,7 @@ async fn ik_continue_does_not_overwrite_cached_chain() { counter.as_ref(), transport.clone(), &mut events_rx, - None, + SendObservers::default(), ) .await; task.await.unwrap(); @@ -573,7 +574,7 @@ async fn xx_after_pair_success_persists_cert_chain() { counter.as_ref(), transport1.clone(), &mut events_rx1, - None, + SendObservers::default(), ) .await .expect("unpaired XX must succeed"); @@ -603,7 +604,7 @@ async fn xx_after_pair_success_persists_cert_chain() { counter.as_ref(), transport2.clone(), &mut events_rx2, - None, + SendObservers::default(), ) .await .expect("paired XX must succeed"); @@ -646,7 +647,7 @@ async fn unpaired_xx_does_not_persist_cert_chain() { counter.as_ref(), transport.clone(), &mut events_rx, - None, + SendObservers::default(), ) .await; task.await.unwrap(); @@ -778,7 +779,7 @@ async fn ik_rejected_recovers_via_xxfallback_and_repopulates_cache() { counter.as_ref(), transport.clone(), &mut events_rx, - None, + SendObservers::default(), ) .await; task.await.unwrap(); @@ -872,7 +873,7 @@ async fn ik_with_stale_cache_invalidates_and_increments_counter() { counter.as_ref(), transport.clone(), &mut events_rx, - None, + SendObservers::default(), ) .await; diff --git a/wacore/src/types/events.rs b/wacore/src/types/events.rs index 798cbae1a..aa4fc7739 100755 --- a/wacore/src/types/events.rs +++ b/wacore/src/types/events.rs @@ -276,6 +276,7 @@ pub enum EventKind { PairingCodeError, AppStateSyncFailed, DecryptedPayload, + SentFrame, // When adding a variant, mind the 128-kind ceiling below (EventInterest packs // each discriminant as a bit in a u128) and keep the guard pointing at the // last variant. @@ -289,7 +290,7 @@ impl EventKind { // Build-time tripwire: a new variant that would overflow EventInterest's bitmask // fails compilation instead of silently corrupting the mask at runtime. -const _: () = assert!((EventKind::DecryptedPayload as u8) < EventKind::CAPACITY); +const _: () = assert!((EventKind::SentFrame as u8) < EventKind::CAPACITY); /// A set of [`EventKind`]s a handler wants delivered. Producers can query the /// aggregate interest before building expensive payloads, and dispatch avoids @@ -992,6 +993,13 @@ pub enum Event { /// Last, like every new variant: a binary `Serialize` format writes the /// variant index, so inserting in the middle renumbers everything after it. DecryptedPayload(DecryptedPayload), + + /// One marshaled stanza that reached the transport, emitted after the write. + /// + /// The outbound counterpart of [`Event::RawNode`]. Library extension — no WA + /// Web equivalent. Gated by `Client::acquire_sent_frame_forwarding()` so + /// nothing is cloned while unused. + SentFrame(SentFrame), } /// Payload for [`Event::PairPasskeyRequest`]. @@ -1095,6 +1103,7 @@ impl Event { Event::NewsletterLiveUpdate(_) => EventKind::NewsletterLiveUpdate, Event::RawNode(_) => EventKind::RawNode, Event::DecryptedPayload(_) => EventKind::DecryptedPayload, + Event::SentFrame(_) => EventKind::SentFrame, Event::MexNotification(_) => EventKind::MexNotification, Event::PairPasskeyRequest(_) => EventKind::PairPasskeyRequest, Event::PairPasskeyConfirmation(_) => EventKind::PairPasskeyConfirmation, @@ -1688,6 +1697,44 @@ pub struct DecryptedPayload { pub payload: Bytes, } +/// Payload of [`Event::SentFrame`]: one marshaled stanza, exactly as it was +/// encrypted and written. +/// +/// The send side had no observer at all. [`Event::RawNode`] hands over every +/// decoded stanza that arrives, but the only thing watching what leaves was a +/// filtered one-shot waiter for a single expected stanza, and the paths that +/// never build a `Node` at all (acks, delivery receipts, direct-encoded IQs) +/// were invisible even to that. So a test could not assert what went to the wire +/// without wrapping the transport, and a malformed stanza in production could not +/// be read back without a rebuild. +/// +/// Reasons to want it: recording a session for replay, asserting the wire form in +/// an integration test, and diagnosing a stanza the server rejected. +/// +/// Emitted from the noise sender once the transport accepted the write, which is +/// the single point every send crosses. A frame that failed to encrypt or to +/// write never appears here, and neither do the pre-noise handshake frames. It is +/// dispatched before the send it belongs to resolves, so a caller that awaited a +/// send can already see its frame. +/// +/// Gated by `Client::acquire_sent_frame_forwarding()`: nothing is emitted, and +/// nothing is cloned, while no consumer holds a lease. +#[derive(Debug, Clone, Serialize, bon::Builder)] +#[non_exhaustive] +pub struct SentFrame { + /// The marshaled stanza, keeping the leading format byte the binary + /// protocol writes, so decoding it is + /// `wacore_binary::marshal::unmarshal_ref(&plaintext[1..])`. + /// + /// A `Bytes`, so forwarding it costs a refcount bump rather than a copy. + /// + /// **Not serialized**, for the same reason as + /// [`DecryptedPayload::payload`]: no text format carries raw bytes without + /// an encoding choice this type has no business making. + #[serde(skip)] + pub plaintext: Bytes, +} + #[derive(Debug, Clone, Serialize, bon::Builder)] #[non_exhaustive] pub struct UndecryptableMessage { From 36246230fbd47b716551aeffaa1e0b56afbdc810 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 8 Aug 2026 04:18:05 +0000 Subject: [PATCH 2/7] fix(client): stop publishing a frame after its last observer let go The gate was read once, at capture time, and the frame was published after the write it belongs to. A lease released in between still saw that frame arrive, which contradicts the documented contract that nothing is emitted while no consumer holds one. Reading the gate again before the publish closes it, for one relaxed load per batch. The test parks the sender inside the transport write, so the frame is captured and encrypted before the lease goes away: without the second read it fails. Also record what the whole-dispatch panic containment costs, since it is per dispatch and not per handler: a panicking observer loses that frame for the observers registered behind it. The bus has no per-handler isolation for any kind, and the plugin host already wraps its own handlers, so isolating this one kind would be a change to the bus rather than to this path. --- src/client.rs | 9 ++++--- src/socket/noise_socket.rs | 53 +++++++++++++++++++++++++++++++++++++- 2 files changed, 58 insertions(+), 4 deletions(-) diff --git a/src/client.rs b/src/client.rs index 0ed1c4e4e..b60a8787c 100644 --- a/src/client.rs +++ b/src/client.rs @@ -175,9 +175,12 @@ impl SentFrameTap { /// /// The dispatch is caught: a consumer that only watches must not be able to /// take the send pipeline down with it, and this runs on the noise sender - /// task, whose death would end every send on the connection. A handler that - /// *blocks* still stalls sends, which is the same contract every event - /// handler already has on the read loop. + /// task, whose death would end every send on the connection. Containment is + /// per dispatch, not per handler, so a panicking observer costs this frame + /// for the observers behind it — the bus offers no per-handler isolation for + /// any kind, and plugins already wrap their own handlers. A handler that + /// *blocks* still stalls sends, the contract every handler has on the read + /// loop. pub(crate) fn publish(&self, plaintext: bytes::Bytes) { #[cfg(test)] self.published.fetch_add(1, Ordering::Relaxed); diff --git a/src/socket/noise_socket.rs b/src/socket/noise_socket.rs index 7b296bd4b..9f1dbe496 100644 --- a/src/socket/noise_socket.rs +++ b/src/socket/noise_socket.rs @@ -322,7 +322,12 @@ impl NoiseSocket { stats.record_frame_sent(*wire_bytes); } } - if let Some(tap) = sent_frames.as_deref() { + // Re-read the gate rather than trusting the one read at + // capture time: a lease released while this batch was in + // flight must not see a frame arrive after it let go. + if let Some(tap) = sent_frames.as_deref() + && tap.enabled() + { for plaintext in observed.drain(..) { tap.publish(plaintext); } @@ -1184,6 +1189,10 @@ mod tests { struct GatedTransport { writes: std::sync::Mutex>, gate: tokio::sync::Semaphore, + /// Writes that have reached the gate, counted before it is awaited: the + /// only way a test can tell "the sender is parked mid-write" from "the + /// sender has not started yet". + arrivals: std::sync::atomic::AtomicUsize, } impl GatedTransport { @@ -1191,18 +1200,24 @@ mod tests { Arc::new(Self { writes: std::sync::Mutex::new(Vec::new()), gate: tokio::sync::Semaphore::new(0), + arrivals: std::sync::atomic::AtomicUsize::new(0), }) } fn writes(&self) -> Vec { self.writes.lock().expect("writes mutex").clone() } + + fn arrivals(&self) -> usize { + self.arrivals.load(Ordering::SeqCst) + } } #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] impl Transport for GatedTransport { async fn send(&self, data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> { + self.arrivals.fetch_add(1, Ordering::SeqCst); let permit = self.gate.acquire().await.expect("gate open"); permit.forget(); self.writes.lock().expect("writes mutex").push(data); @@ -1674,6 +1689,42 @@ mod tests { assert!(observer.frames().is_empty()); } + /// A lease released while a frame is already encrypted and waiting on the + /// transport must still turn the frame away: the gate is read at capture + /// time, so without the second read at publish time a consumer that stopped + /// watching would get one more frame after it let go. The gated transport + /// holds the write open across the release, which is the whole window. + #[tokio::test] + async fn a_lease_released_mid_write_turns_its_frame_away() { + let (tap, observer, _subscription) = watched_tap(); + let transport = GatedTransport::closed(); + let socket = Arc::new(socket_watched_by(transport.clone(), tap.clone())); + + let mut send = queue_all(&socket, [bytes::Bytes::from(vec![4u8; 64])].into_iter()); + // Parked inside the write, which is past capture and past encryption: + // releasing before this point would prove nothing, since the frame + // would never have been captured at all. + crate::test_utils::poll_until("the write to reach the gate", || transport.arrivals() == 1) + .await; + tap.release(); + transport.gate.add_permits(1); + for result in (&mut send).await { + result.expect("the send itself must still succeed"); + } + + assert_eq!( + transport.writes().len(), + 1, + "the frame must still reach the wire" + ); + assert_eq!( + tap.published(), + 0, + "a frame must not be published after the last lease is released" + ); + assert!(observer.frames().is_empty()); + } + /// An observer that panics must not take the sender task with it, or one /// consumer watching would end every send on the connection. #[tokio::test] From 323ba7de5d6207fa9be161df09ea45e5264df448 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 8 Aug 2026 04:21:51 +0000 Subject: [PATCH 3/7] docs(events): say that a sent frame is plaintext, not a transport frame "Exactly as it was encrypted and written" reads as if the payload were the bytes that went on the socket. It is the input to that encryption: no length prefix, no AEAD tag. A consumer building a replay tool has to know which of the two it is holding. --- wacore/src/types/events.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/wacore/src/types/events.rs b/wacore/src/types/events.rs index aa4fc7739..d2042b666 100755 --- a/wacore/src/types/events.rs +++ b/wacore/src/types/events.rs @@ -1698,7 +1698,7 @@ pub struct DecryptedPayload { } /// Payload of [`Event::SentFrame`]: one marshaled stanza, exactly as it was -/// encrypted and written. +/// handed to the noise frame encryption. /// /// The send side had no observer at all. [`Event::RawNode`] hands over every /// decoded stanza that arrives, but the only thing watching what leaves was a @@ -1726,6 +1726,11 @@ pub struct SentFrame { /// protocol writes, so decoding it is /// `wacore_binary::marshal::unmarshal_ref(&plaintext[1..])`. /// + /// Plaintext, as the name says, not a transport frame: the length prefix and + /// the AEAD tag are added after this, and only + /// [`Client::stats`](crate::stats::SessionStats) accounts for those. Replay + /// it as a stanza, not as bytes to put on a socket. + /// /// A `Bytes`, so forwarding it costs a refcount bump rather than a copy. /// /// **Not serialized**, for the same reason as From c3035455e1b25094181884c09cd6fc910054ba1c Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 8 Aug 2026 04:25:08 +0000 Subject: [PATCH 4/7] docs(client): state what the sent-frame lease guarantees, and what it does not Reviewers read the recheck before publishing as a promise that no frame can follow the last drop, and asked for a lock to close the remainder. A lock there would deadlock an observer that drops its lease from inside its own handler, and the guarantee it would buy is not one the gate ever offered: the count is aggregate, and delivery is decided by subscription, exactly as for the other gated kinds. Say so where a consumer reads it. --- src/client.rs | 8 ++++++++ src/socket/noise_socket.rs | 7 ++++--- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/src/client.rs b/src/client.rs index b60a8787c..e822b10e8 100644 --- a/src/client.rs +++ b/src/client.rs @@ -107,6 +107,14 @@ impl Drop for RawNodeLease { /// /// Dropping the final lease disables forwarding. The lease holds only a weak /// client reference, so it cannot keep the client alive. +/// +/// It gates, it does not fence, and one aggregate count gates them all rather +/// than one per lease: a frame captured while any lease was alive can still +/// arrive just after this one drops, and which handlers receive it is a matter of +/// subscription, not of who holds a lease. Every gated kind works this way. +/// Making the drop wait for in-flight dispatches to drain would instead deadlock +/// an observer that drops its lease from inside its own handler, which is the +/// natural way to record one frame and stop. #[must_use = "dropping the lease immediately releases sent-frame forwarding"] pub struct SentFrameLease { client: std::sync::Weak, diff --git a/src/socket/noise_socket.rs b/src/socket/noise_socket.rs index 9f1dbe496..197d77098 100644 --- a/src/socket/noise_socket.rs +++ b/src/socket/noise_socket.rs @@ -322,9 +322,10 @@ impl NoiseSocket { stats.record_frame_sent(*wire_bytes); } } - // Re-read the gate rather than trusting the one read at - // capture time: a lease released while this batch was in - // flight must not see a frame arrive after it let go. + // Re-read the gate rather than trusting the read at + // capture time, so a batch that outlived its last lease + // stays quiet. A release racing this instant may still + // lose: the lease gates, it does not fence. if let Some(tap) = sent_frames.as_deref() && tap.enabled() { From ce7e57478944dd4dc9b1bc7a5f1593480fbd4197 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 8 Aug 2026 06:20:55 +0000 Subject: [PATCH 5/7] fix(test): decrypt captured frames without reparsing their length `CapturingMockTransport::sent()` already splits each write into its frames, prefix intact, so the helper was walking a length prefix that had been walked for it and adding a third copy of that parse to the crate. It now checks the declared length against the entry instead, which turns a caller that passes raw writes into a failed assertion rather than a decrypt of garbage. The noise socket the client-test helper installs now reports to the client's stats as well as its tap, matching the real socket and the other fixture: a test that later asserts on `frames_sent` would have read zero and blamed the code under test. The allocation delta is compared as a signed difference. Both sides are empirical minima off a process-global counter, and an inversion has to print the two numbers rather than panic inside the subtraction. Also fixes the observability doc: it claimed five distinct send paths and then listed callers that funnel through two, and still described VoIP sockets as passing `None`, which this branch replaced with `SendObservers::default()`. --- agent_docs/observability.md | 18 +++++++------- src/client/tests.rs | 2 +- src/socket/noise_socket.rs | 4 +++- src/test_utils.rs | 47 +++++++++++++++++++++---------------- 4 files changed, 41 insertions(+), 30 deletions(-) diff --git a/agent_docs/observability.md b/agent_docs/observability.md index 26332cb48..01670e073 100644 --- a/agent_docs/observability.md +++ b/agent_docs/observability.md @@ -32,12 +32,14 @@ chokepoints: transport write — post-noise wire bytes (frame header + AEAD tag included). - **Received**: the read loop (`node_io.rs`) per `DataReceived` batch. -That sent chokepoint is the *only* place every outbound frame crosses — five -distinct send paths reach it (`send_node`, `send_raw_bytes`, -`send_raw_bytes_burst`, and the ack/receipt workers through the burst), and -`send_raw_bytes` deliberately bypasses node logging and sent-node waiters — so -anything that has to see *everything* the client sends belongs there and nowhere -else. `Event::SentFrame` is the other thing wired into it +That sent chokepoint is the *only* place every outbound frame crosses. Two +functions reach it — `send_raw_bytes` and `send_raw_bytes_burst` — and everything +else funnels through one of those: `send_node` and every IQ through the first, +the ack and delivery-receipt workers through the second. `send_raw_bytes` +deliberately bypasses node logging and sent-node waiters, so anything that has to +see *everything* the client sends belongs at the chokepoint and nowhere else. + +`Event::SentFrame` is the other thing wired into it (`Client::acquire_sent_frame_forwarding()`, lease-gated like `RawNode`): it hands over the marshaled plaintext of each frame the transport accepted. Both halves travel to the socket as `SendObservers`, so the next observer plugs in there @@ -53,8 +55,8 @@ cost a clock read on every frame written, which is the client's hottest path and a call out of the module on wasm32/embedded. `frames_sent` answers "is it still sending?" for free. Message-level counters piggyback on the existing `telemetry::send`/`recv` chokepoints; reconnect attempts are counted in the -run loop. VoIP relay sockets pass `None` and are not counted — this is the -main WA session socket only. +run loop. VoIP relay sockets pass `SendObservers::default()` and are not counted +— this is the main WA session socket only. ### 2. `Client::memory_report()` — retained memory (on demand) diff --git a/src/client/tests.rs b/src/client/tests.rs index 0f9beba19..87803d4b1 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -3514,7 +3514,7 @@ async fn install_test_noise_socket( transport, NoiseCipher::new(&key).expect("valid key"), NoiseCipher::new(&key).expect("valid key"), - crate::socket::noise_socket::SendObservers::default() + crate::socket::noise_socket::SendObservers::with_stats(client.stats.clone()) .with_sent_frames(client.sent_frame_tap.clone()), ); *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); diff --git a/src/socket/noise_socket.rs b/src/socket/noise_socket.rs index 4c52b0017..9eb69dd2c 100644 --- a/src/socket/noise_socket.rs +++ b/src/socket/noise_socket.rs @@ -2061,8 +2061,10 @@ mod tests { idle <= 4, "an unwatched send must cost what it always did, got {idle}" ); + // Signed: both are empirical minima off a process-global counter, and an + // inversion has to report the numbers rather than panic on underflow. assert_eq!( - watched - idle, + watched as i64 - idle as i64, 2, "watching must cost a constant per frame (idle {idle}, watched {watched})" ); diff --git a/src/test_utils.rs b/src/test_utils.rs index 679b993a4..45c2c7577 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -46,33 +46,40 @@ impl EventHandler for SentFrameRecorder { } } -/// Splits a run of length-prefixed noise frames and decrypts each under the -/// counter its position implies, yielding the plaintexts in wire order. +/// Decrypts captured noise frames under the counter each one's position implies, +/// yielding the plaintexts in wire order. +/// +/// Takes what [`CapturingMockTransport::sent`] yields: one framed frame per +/// entry, prefix intact, already split out of whatever writes carried them. The +/// length is checked rather than reparsed, so a caller passing raw writes (where +/// one entry can hold several frames) fails here instead of decrypting garbage. /// /// The counter is the AES-GCM nonce, so this doubles as an order check: a frame /// read out of position cannot authenticate. +/// +/// [`CapturingMockTransport::sent`]: crate::transport::mock::CapturingMockTransport::sent #[cfg(test)] -pub(crate) fn decrypt_wire_frames(writes: &[bytes::Bytes], key: &[u8; 32]) -> Vec> { - use wacore::handshake::NoiseCipher; +pub(crate) fn decrypt_wire_frames(frames: &[bytes::Bytes], key: &[u8; 32]) -> Vec> { + use wacore::framing::FRAME_LENGTH_SIZE; + use wacore::noise::NoiseCipher; let cipher = NoiseCipher::new(key).expect("32-byte key"); - let mut plaintexts = Vec::new(); - for write in writes { - let mut wire = &write[..]; - while !wire.is_empty() { - let mut len = 0usize; - for byte in &wire[..wacore::framing::FRAME_LENGTH_SIZE] { - len = (len << 8) | *byte as usize; - } - let body = - &wire[wacore::framing::FRAME_LENGTH_SIZE..wacore::framing::FRAME_LENGTH_SIZE + len]; - let mut body = bytes::BytesMut::from(body); - cipher - .decrypt_in_place_with_counter(plaintexts.len() as u32, &mut body) - .expect("each captured frame must decrypt under its own counter"); - plaintexts.push(body.to_vec()); - wire = &wire[wacore::framing::FRAME_LENGTH_SIZE + len..]; + let mut plaintexts = Vec::with_capacity(frames.len()); + for (counter, frame) in frames.iter().enumerate() { + let mut declared = 0usize; + for byte in &frame[..FRAME_LENGTH_SIZE] { + declared = (declared << 8) | *byte as usize; } + assert_eq!( + declared, + frame.len() - FRAME_LENGTH_SIZE, + "expected one frame per entry, as `sent()` yields them" + ); + let mut body = bytes::BytesMut::from(&frame[FRAME_LENGTH_SIZE..]); + cipher + .decrypt_in_place_with_counter(counter as u32, &mut body) + .expect("each captured frame must decrypt under its own counter"); + plaintexts.push(body.to_vec()); } plaintexts } From 8219c0537605845743280fab15996e3d25d588f0 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 8 Aug 2026 07:52:33 +0000 Subject: [PATCH 6/7] Merge branch 'main', decoding a sent frame through the checked entry point main added `unmarshal_packed_ref`, which validates the leading format byte before decoding instead of assuming it. The `SentFrame` doc and the send-path test both hand-rolled `&plaintext[1..]`, which is the same thing right up until a payload arrives with the compressed bit set. Point them at the helper. No conflicts this time. Rechecked the allocation test against main's send-path work (#1256): the idle path and the two-per-frame watching cost are unchanged. --- src/client/tests.rs | 2 +- wacore/src/types/events.rs | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/client/tests.rs b/src/client/tests.rs index 10fbb600c..3bdcc2425 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -3770,7 +3770,7 @@ async fn every_send_path_is_observed_exactly_as_it_reached_the_wire() { // The bytes are the stanza, not a rendering of it: the first frame decodes // back to the node that was sent. - let decoded = wacore_binary::marshal::unmarshal_ref(&observed[0][1..]) + let decoded = wacore_binary::marshal::unmarshal_packed_ref(&observed[0]) .expect("an observed frame must decode as the stanza it carried"); assert_eq!(decoded.tag.as_ref(), "presence"); assert_eq!( diff --git a/wacore/src/types/events.rs b/wacore/src/types/events.rs index d2042b666..ca45f3cc0 100755 --- a/wacore/src/types/events.rs +++ b/wacore/src/types/events.rs @@ -1724,7 +1724,8 @@ pub struct DecryptedPayload { pub struct SentFrame { /// The marshaled stanza, keeping the leading format byte the binary /// protocol writes, so decoding it is - /// `wacore_binary::marshal::unmarshal_ref(&plaintext[1..])`. + /// `wacore_binary::marshal::unmarshal_packed_ref(&plaintext)`, which checks + /// that byte rather than assuming it. /// /// Plaintext, as the name says, not a transport frame: the length prefix and /// the AEAD tag are added after this, and only From f5e9abf2678cfdaa45438a20cea356fb846ae9d9 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 8 Aug 2026 12:34:13 +0000 Subject: [PATCH 7/7] docs(observability): bound the chokepoint claim to the session socket "Every outbound frame" overstated it: the XX/IK exchange writes to the transport before the noise socket exists, so handshake frames cross no chokepoint and are not forwarded. Say post-handshake, and scope "everything the client sends" to the session socket. --- agent_docs/observability.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/agent_docs/observability.md b/agent_docs/observability.md index 01670e073..5060660b7 100644 --- a/agent_docs/observability.md +++ b/agent_docs/observability.md @@ -32,12 +32,14 @@ chokepoints: transport write — post-noise wire bytes (frame header + AEAD tag included). - **Received**: the read loop (`node_io.rs`) per `DataReceived` batch. -That sent chokepoint is the *only* place every outbound frame crosses. Two +That sent chokepoint is the *only* place every post-handshake frame crosses (the +XX/IK exchange writes to the transport directly, before this socket exists). Two functions reach it — `send_raw_bytes` and `send_raw_bytes_burst` — and everything else funnels through one of those: `send_node` and every IQ through the first, the ack and delivery-receipt workers through the second. `send_raw_bytes` deliberately bypasses node logging and sent-node waiters, so anything that has to -see *everything* the client sends belongs at the chokepoint and nowhere else. +see everything the client sends on the session socket belongs at the chokepoint +and nowhere else. `Event::SentFrame` is the other thing wired into it (`Client::acquire_sent_frame_forwarding()`, lease-gated like `RawNode`): it hands