diff --git a/src/client.rs b/src/client.rs index 3336475ca..8f535cc1d 100644 --- a/src/client.rs +++ b/src/client.rs @@ -8,6 +8,7 @@ mod device_registry; pub(crate) mod device_topology; #[cfg(feature = "client-lifecycle")] mod extension_lifecycle; +pub mod interceptor; mod iq_ops; mod lid_pn; mod lifecycle; @@ -342,6 +343,12 @@ pub struct MemoryReport { // -- Misc -- pub chatstate_handlers: usize, pub custom_enc_handlers: usize, + /// Interceptors currently registered. + /// + /// A handle that outlives its interest leaves one registered, and a leak + /// here costs a walk on every stanza — which the count is what makes + /// visible. + pub stanza_interceptors: usize, } impl MemoryReport { @@ -505,6 +512,7 @@ impl std::fmt::Display for MemoryReport { writeln!(f, "--- Misc ---")?; writeln!(f, " chatstate_handlers: {}", self.chatstate_handlers)?; writeln!(f, " custom_enc_handlers: {}", self.custom_enc_handlers)?; + writeln!(f, " stanza_interceptors: {}", self.stanza_interceptors)?; writeln!( f, " total estimated: {} B", @@ -1413,6 +1421,15 @@ pub struct Client { /// Number of consumers currently requesting `Event::RawNode` forwarding. raw_node_forwarding: AtomicUsize, + /// 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. + stanza_interceptors: std::sync::RwLock>>, + /// Kept alongside so the read loop can skip the lock entirely while none + /// are registered. + stanza_interceptor_count: AtomicUsize, + next_interceptor_id: AtomicU64, + /// Active VoIP calls and their media-task abort handles. `abort_all` runs from the /// connection-cleanup path so a disconnect/reconnect tears down every in-flight call. Behind the /// `voip` feature: it is populated only by the `voip` media facade. diff --git a/src/client/accessors.rs b/src/client/accessors.rs index 0c28211d5..617541e26 100644 --- a/src/client/accessors.rs +++ b/src/client/accessors.rs @@ -1,6 +1,7 @@ //! Small accessors, config setters, node waiters and sync-error helpers. use super::*; +use crate::client::interceptor::{InterceptorHandle, Registration, StanzaInterceptor}; /// Identity for span/error tagging. Named fields, not a tuple — LID/PN transposition would /// otherwise be a silent, unchecked bug at call sites. @@ -60,6 +61,77 @@ impl Client { self.raw_node_forwarding.load(Ordering::Relaxed) != 0 } + /// Register an interceptor that sees each decoded stanza before the + /// built-in pipeline, and may take it. + /// + /// Interceptors run in registration order; the first to return + /// [`Interception::Handled`] wins and the rest are skipped. + /// + /// See [`crate::client::interceptor`] for what this is for and what it + /// costs. + /// + /// [`Interception::Handled`]: crate::client::interceptor::Interception::Handled + pub fn add_stanza_interceptor( + self: &Arc, + interceptor: Arc, + ) -> InterceptorHandle { + let id = self.next_interceptor_id.fetch_add(1, Ordering::Relaxed); + self.update_stanza_interceptors(|registered| { + registered.push(Registration { id, interceptor }); + }); + InterceptorHandle { + client: Arc::downgrade(self), + id, + } + } + + pub(crate) fn remove_stanza_interceptor(&self, id: u64) { + self.update_stanza_interceptors(|registered| { + registered.retain(|entry| entry.id != id); + }); + } + + /// Whether any interceptor is registered. + /// + /// One relaxed load, so the read loop pays nothing while none are. Relaxed + /// is enough because the lock behind it does the synchronising: a reader + /// racing a registration either sees the count in time or does not, and a + /// stanza that arrived before the registration finished was never that + /// interceptor's to see. + pub(crate) fn has_stanza_interceptors(&self) -> bool { + self.stanza_interceptor_count.load(Ordering::Relaxed) != 0 + } + + /// The current interceptors. + /// + /// A refcount bump, not a copy — and the snapshot is released before any + /// interceptor runs, so one that registers another cannot deadlock. + pub(crate) fn stanza_interceptors(&self) -> Arc> { + Arc::clone( + &self + .stanza_interceptors + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner), + ) + } + + /// Copy-on-write, as the event bus does it: a reader keeps whatever + /// snapshot it took, so registering never blocks the read loop for longer + /// than the swap. + fn update_stanza_interceptors(&self, edit: impl FnOnce(&mut Vec)) { + let mut guard = self + .stanza_interceptors + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut next = guard.as_ref().clone(); + edit(&mut next); + // Stored while the write lock is held, so a reader never sees a count + // promising more than the snapshot holds. + self.stanza_interceptor_count + .store(next.len(), Ordering::Relaxed); + *guard = Arc::new(next); + } + /// Enable or disable skipping of history sync notifications at runtime. /// /// When enabled, the client will acknowledge incoming history sync @@ -311,6 +383,7 @@ impl Client { ), chatstate_handlers, custom_enc_handlers: self.custom_enc_handlers.get().map_or(0, |m| m.len()), + stanza_interceptors: self.stanza_interceptors().len(), } } diff --git a/src/client/interceptor.rs b/src/client/interceptor.rs new file mode 100644 index 000000000..8fc79f18a --- /dev/null +++ b/src/client/interceptor.rs @@ -0,0 +1,219 @@ +//! Taking over a stanza before the built-in pipeline sees it. +//! +//! The client models the stanzas it knows about and nacks the rest, which is +//! the right default: a `` tells the server this client cannot act on +//! something, and silence would leave the stanza in the offline queue forever. +//! +//! But it leaves no room for a consumer that *can* act on it. A stanza this +//! version does not model gets nacked whether or not the application would have +//! known what to do with it, and there is no way to say otherwise: +//! [`StanzaRouter::register`] panics on a duplicate tag, so even a handler for +//! an existing tag cannot be replaced. +//! +//! An interceptor is that room. It runs where dispatch would have, and either +//! steps aside or claims the stanza — in which case the built-in handler is +//! skipped, and whatever answer the client owed the server becomes the +//! claimant's to send. For most tags that answer is a transport ack and the +//! client still sends it; see [Acknowledgement](#acknowledgement) for the ones +//! where it is not. +//! +//! # What an interceptor sees +//! +//! Stanzas that would have reached dispatch. A response the client already +//! correlated to a pending request, a `` that ends the stream, +//! and the connection-critical tags below never get there, so an interceptor +//! does not see them either. +//! +//! To observe *everything* decoded, including those, use +//! [`Event::RawNode`] — it is emitted before any of the early returns. +//! The two are different tools: one watches, this one takes over. +//! +//! [`Event::RawNode`]: wacore::types::events::Event::RawNode +//! +//! # What cannot be claimed +//! +//! `success`, `failure`, `stream:error` and `ack` settle connection state: +//! authentication, shutdown and reconnection, and the waiters a send blocks on. +//! An interceptor that took one would not extend the client — it would leave it +//! authenticated-but-unaware, or never reconnecting, or waiting forever on a +//! send that already completed. They are never offered. +//! +//! Housekeeping is likewise untouched. Offline-sync tracking and +//! response-waiter resolution run before dispatch and keep running whether or +//! not a stanza is claimed. +//! +//! A server-initiated `` ping is not offered either, for the same reason as +//! the four above: a claimed ping is a pong never sent, and the server drops +//! the connection over it. +//! +//! Every other `` *is* offered, including ones the client answers on its +//! own — a pairing step, a query it models. Most `` traffic is exactly what +//! a consumer would want to extend. Claiming one leaves the server without the +//! reply it expects, so match narrowly. +//! +//! # Acknowledgement +//! +//! A claim does not change what the server is owed. Where the client would have +//! acked, it still acks; where it would have nacked a tag it does not model, +//! the claim turns that into an ack, because someone did handle it — answering +//! nothing would leave the stanza in the offline queue and keep the stream +//! recycling. Both need `id` and `from`: without them there is nothing to +//! address, and the client would not have answered either. +//! +//! A tag the client *does* model but answers some other way is answered by +//! nobody once claimed. A direct `` draws a delivery `` and +//! an `` draws an ``; a generic `` is neither, so +//! the client does not send one. Claiming those means owing the reply. +//! +//! ``, `` and `` are not in that group: the client +//! answers those with a transport `` already, so a claim leaves the ack +//! exactly where it was. +//! +//! [`StanzaRouter::register`]: crate::handlers::router::StanzaRouter::register +//! +//! # Cost +//! +//! Nothing runs while no interceptor is registered: the read loop checks one +//! relaxed atomic and carries on. Registering is what turns the check into a +//! walk. +//! +//! # Example +//! +//! Handling a stanza the client does not model, instead of nacking it: +//! +//! ```no_run +//! use std::sync::Arc; +//! use whatsapp_rust::client::interceptor::{Interception, StanzaInterceptor}; +//! use wacore_binary::node::OwnedNodeRef; +//! +//! struct Vendor; +//! +//! impl StanzaInterceptor for Vendor { +//! fn intercept(&self, node: &OwnedNodeRef) -> Interception { +//! if node.tag() == "vendor:thing" { +//! // … act on it … +//! Interception::Handled +//! } else { +//! Interception::Pass +//! } +//! } +//! } +//! +//! # fn example(client: &Arc) { +//! let handle = client.add_stanza_interceptor(Arc::new(Vendor)); +//! // Dropping `handle` removes it. +//! # let _ = handle; +//! # } +//! ``` + +use std::sync::{Arc, Weak}; + +use wacore::sync_marker::MaybeSendSync; +use wacore_binary::node::OwnedNodeRef; + +use crate::Client; + +/// What an interceptor decided about a stanza. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +pub enum Interception { + /// Leave the stanza to the client. + /// + /// The default, so an interceptor that only cares about one tag needs no + /// branch for the rest. + #[default] + Pass, + /// The interceptor took the stanza. + /// + /// The built-in pipeline is skipped. Where the client's answer was a + /// transport ack it still sends one; where the answer was something else — + /// a delivery `` for a direct ``, an `` + /// — nothing is sent, and the claimant owes that reply. See the + /// [acknowledgement](index.html#acknowledgement) rules. + Handled, +} + +impl Interception { + /// Whether the built-in pipeline should be skipped. + #[must_use] + pub const fn is_handled(self) -> bool { + matches!(self, Self::Handled) + } +} + +/// Sees a stanza on its way to the built-in pipeline. +/// +/// Not every decoded stanza: see the module documentation for what never +/// reaches this point. +/// +/// Runs on the read loop, so it must return quickly: time spent here is time +/// the next stanza waits. Work that can take a while belongs on a task. +/// +/// Must not panic. Like [`EventHandler`], this is called directly by the read +/// loop, and an unwind there takes the connection with it. The plugin host +/// catches panics from plugins; a directly registered interceptor is trusted. +/// +/// [`EventHandler`]: wacore::types::events::EventHandler +pub trait StanzaInterceptor: MaybeSendSync + 'static { + /// Decide what happens to `node`. + fn intercept(&self, node: &OwnedNodeRef) -> Interception; +} + +impl StanzaInterceptor for F +where + F: Fn(&OwnedNodeRef) -> Interception + MaybeSendSync + 'static, +{ + fn intercept(&self, node: &OwnedNodeRef) -> Interception { + self(node) + } +} + +/// Keeps an interceptor registered. Dropping it removes the interceptor. +/// +/// Holds a weak client reference, so a forgotten handle cannot keep a client +/// alive. +#[must_use = "dropping the handle immediately removes the interceptor"] +#[derive(Debug)] +pub struct InterceptorHandle { + pub(crate) client: Weak, + pub(crate) id: u64, +} + +impl Drop for InterceptorHandle { + fn drop(&mut self) { + if let Some(client) = self.client.upgrade() { + client.remove_stanza_interceptor(self.id); + } + } +} + +/// One registered interceptor. +#[derive(Clone)] +pub(crate) struct Registration { + pub(crate) id: u64, + pub(crate) interceptor: Arc, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn pass_is_the_default_so_narrow_interceptors_stay_short() { + assert_eq!(Interception::default(), Interception::Pass); + assert!(!Interception::Pass.is_handled()); + assert!(Interception::Handled.is_handled()); + } + + #[test] + fn a_closure_is_an_interceptor() { + fn takes(_: impl StanzaInterceptor) {} + takes(|_node: &OwnedNodeRef| Interception::Pass); + } + + #[test] + fn interception_is_comparable_and_debuggable() { + assert_eq!(Interception::Pass, Interception::Pass); + assert_ne!(Interception::Pass, Interception::Handled); + assert!(!format!("{:?}", Interception::Handled).is_empty()); + } +} diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 66c5f9b0a..35f1a4f29 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -506,6 +506,9 @@ impl Client { saver_handle: std::sync::OnceLock::new(), alloc_meter: std::sync::OnceLock::new(), raw_node_forwarding: AtomicUsize::new(0), + stanza_interceptors: std::sync::RwLock::new(Arc::new(Vec::new())), + stanza_interceptor_count: AtomicUsize::new(0), + next_interceptor_id: AtomicU64::new(0), #[cfg(feature = "voip-runtime")] call_registry: Arc::new(wacore::voip::CallRegistry::new()), #[cfg(feature = "voip-runtime")] diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 7c3004f5b..3a4afb34a 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -55,6 +55,36 @@ fn from_jid_matches( /// The wire shape the server uses for E2EE status updates, carrying the same /// payload as ``. +/// Stanzas that carry connection state, which an interceptor may not claim. +/// +/// `success` and `failure` settle authentication, `stream:error` drives +/// shutdown and reconnection, and `ack` resolves the waiters a send is blocked +/// on. Letting a consumer take one would not extend the client — it would leave +/// it authenticated-but-unaware, or never reconnecting, or waiting forever on a +/// send that already completed. +/// +/// `zapo` protects the same two auth tags from its stanza filters, for the same +/// reason. +fn is_connection_critical(node: &wacore_binary::NodeRef<'_>) -> bool { + matches!( + node.tag.as_ref(), + "success" | "failure" | "stream:error" | "ack" + ) || (node.tag.as_ref() == "iq" && is_ping_request(node)) +} + +/// A server-initiated ping, which this client owes a pong. +/// +/// Type-agnostic on an absent type, like WA Web's `handleIq`, but never a +/// `type="result"`/`"error"` ping — that is a response to our own ping, and +/// ponging it back is wrong. +fn is_ping_request(node: &wacore_binary::NodeRef<'_>) -> bool { + node.get_attr("type").is_none_or(|s| s.as_str() == "get") + && (node.get_optional_child("ping").is_some() + || node + .get_attr("xmlns") + .is_some_and(|s| s.as_str() == "urn:xmpp:ping")) +} + fn is_status_broadcast_stanza(node: &wacore_binary::NodeRef<'_>) -> bool { from_jid_matches(node, |jid| jid.is_status_broadcast()) } @@ -508,6 +538,40 @@ impl Client { let should_ack = self.should_ack(nr); let deferred_ack_node = should_ack.then(|| Arc::clone(&node)); + // An interceptor runs before the built-in pipeline so a consumer can + // act on a stanza this version does not model, instead of watching it + // get nacked. + if self.has_stanza_interceptors() + && !is_connection_critical(nr) + && self.intercept_stanza(&node) + { + // A claim does not change what the server is owed. Where this + // client would have acked it still acks; where it would have nacked + // a tag it does not model, the claim turns that into an ack, + // because someone did handle it — and answering nothing would leave + // the stanza in the offline queue with the stream recycling. + // + // A tag the client models but answers some other way gets nothing + // here: a direct draws a delivery , an + // draws an , and a generic + // is neither. Inventing one is worse than silence — whoever claimed + // the stanza took on the reply. The tags `should_ack` covers are + // unaffected; they were already answered above. + // + // Same identity requirement as the nack path: without `id` and + // `from` there is nothing to address. + let ack = deferred_ack_node.or_else(|| { + (!self.stanza_router.models(nr.tag.as_ref()) + && nr.get_attr("id").is_some() + && nr.get_attr("from").is_some()) + .then(|| Arc::clone(&node)) + }); + if let Some(node) = ack { + self.maybe_deferred_ack(node).await; + } + return; + } + // Bypass async_trait's boxed future for the hot built-in handlers while // retaining router registration for direct router callers. match nr.tag.as_ref() { @@ -556,6 +620,25 @@ impl Client { } } + /// Offer a stanza to the registered interceptors. + /// + /// Returns whether one took it. The first to claim the stanza wins, so an + /// interceptor registered earlier can shadow a later one — registration + /// order is the priority order. + fn intercept_stanza(self: &Arc, node: &Arc) -> bool { + for registration in self.stanza_interceptors().iter() { + if registration.interceptor.intercept(node).is_handled() { + debug!( + target: "Client/Recv", + "Stanza <{}> taken by an interceptor", + node.tag() + ); + return true; + } + } + false + } + /// Whether a decrypted node must stay on the read loop instead of moving to /// a spawned task. success/failure/stream:error carry connection state the /// rest depends on, and `ib` sets up offline-sync tracking before the batch @@ -1982,17 +2065,11 @@ impl Client { tracing::instrument(name = "wa.conn.iq_in", level = "debug", skip_all) )] pub(crate) async fn handle_iq(self: &Arc, node: &wacore_binary::NodeRef<'_>) -> bool { - // Pong a server-initiated ping (a request: type="get" or, like WA Web's - // type-agnostic handleIq, an absent type), but not a type="result"/"error" - // ping — that's a response to our own ping, and ponging it back is wrong. - // The previous gate required type=="get" exactly, dropping an absent-type - // ping and risking a keepalive timeout/disconnect. - let is_ping_request = node.get_attr("type").is_none_or(|s| s.as_str() == "get") - && (node.get_optional_child("ping").is_some() - || node - .get_attr("xmlns") - .is_some_and(|s| s.as_str() == "urn:xmpp:ping")); - if is_ping_request { + // Pong a server-initiated ping. The gate is shared with + // `is_connection_critical`, which never offers one to an interceptor: + // a claimed ping is a pong never sent, and the server drops the + // connection for it. + if is_ping_request(node) { debug!("Received ping, sending pong."); let mut parser = node.attrs(); let from_jid = parser.jid("from"); diff --git a/src/client/tests.rs b/src/client/tests.rs index 337ffe454..89660e418 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -226,13 +226,13 @@ async fn test_ack_without_matching_waiter() { /// Round-trip a built `Node` into the raw-bytes shape `unpack()` produces /// from the network (marshal_ref prepends a 0x00 format byte that /// `OwnedNodeRef::new` does not expect). -fn to_owned_node(node: &Node) -> wacore_binary::OwnedNodeRef { +fn to_owned_node(node: &Node) -> OwnedNodeRef { wacore_binary::marshal::marshal_ref(&node.as_node_ref()) - .and_then(|buf| wacore_binary::OwnedNodeRef::new(bytes::Bytes::from(buf).slice(1..))) + .and_then(|buf| OwnedNodeRef::new(bytes::Bytes::from(buf).slice(1..))) .expect("valid node") } -fn owned_ack_node(id: &str) -> wacore_binary::OwnedNodeRef { +fn owned_ack_node(id: &str) -> OwnedNodeRef { to_owned_node( &NodeBuilder::new("ack") .attr("id", id) @@ -1446,7 +1446,7 @@ fn test_unified_session_protocol_node() { info!("✅ test_unified_session_protocol_node passed"); } -fn node_to_owned_ref(node: Node) -> Arc { +fn node_to_owned_ref(node: Node) -> Arc { crate::test_utils::node_to_owned_ref(&node) } @@ -4907,3 +4907,420 @@ async fn chatstate_dispatch_reaches_every_registered_handler() { "the event is built once and cloned per handler" ); } + +// --- stanza interceptors --------------------------------------------------- + +use crate::client::interceptor::{Interception, StanzaInterceptor}; +use wacore_binary::OwnedNodeRef; + +/// Builds a client with no transport, which is enough: interception happens +/// before anything is sent. +async fn create_interceptor_test_client() -> Arc { + create_offline_sync_test_client().await +} + +/// Records which stanzas it saw and claims the ones whose tag matches. +struct Recorder { + claim: &'static str, + seen: Arc>>, +} + +impl StanzaInterceptor for Recorder { + fn intercept(&self, node: &OwnedNodeRef) -> Interception { + self.seen + .lock() + .expect("recorder lock") + .push(node.tag().to_string()); + if node.tag() == self.claim { + Interception::Handled + } else { + Interception::Pass + } + } +} + +fn recorder(claim: &'static str) -> (Arc, Arc>>) { + let seen = Arc::new(std::sync::Mutex::new(Vec::new())); + ( + Arc::new(Recorder { + claim, + seen: Arc::clone(&seen), + }), + seen, + ) +} + +#[tokio::test] +async fn interceptors_cost_nothing_until_one_is_registered() { + let client = create_interceptor_test_client().await; + assert!(!client.has_stanza_interceptors()); + + let (interceptor, _seen) = recorder("nothing"); + let handle = client.add_stanza_interceptor(interceptor); + assert!(client.has_stanza_interceptors()); + + drop(handle); + assert!( + !client.has_stanza_interceptors(), + "dropping the handle unregisters it" + ); +} + +#[tokio::test] +async fn an_interceptor_sees_every_decoded_stanza() { + let client = create_interceptor_test_client().await; + let (interceptor, seen) = recorder("nothing-matches"); + let _handle = client.add_stanza_interceptor(interceptor); + + for tag in ["ib", "receipt", "notification"] { + client + .process_node(node_to_owned_ref(NodeBuilder::new(tag).build())) + .await; + } + + assert_eq!( + *seen.lock().expect("recorder lock"), + ["ib", "receipt", "notification"] + ); +} + +/// A `` the client dispatches, and the event that proves it did. +fn receipt_stanza() -> Node { + NodeBuilder::new("receipt") + .attr("from", "5511999998888@s.whatsapp.net") + .attr("id", "RCPT-INTERCEPT") + .build() +} + +#[tokio::test] +async fn passing_leaves_the_stanza_to_the_client() { + use wacore::types::events::ChannelEventHandler; + let client = create_interceptor_test_client().await; + let (handler, events) = ChannelEventHandler::new(); + client.subscribe_handler(handler).detach(); + + let (interceptor, seen) = recorder("nothing-matches"); + let _handle = client.add_stanza_interceptor(interceptor); + + client + .process_node(node_to_owned_ref(receipt_stanza())) + .await; + + assert_eq!(seen.lock().expect("recorder lock").len(), 1, "it ran"); + assert!( + events.try_recv().is_ok(), + "and the built-in handler dispatched its event" + ); +} + +#[tokio::test] +async fn claiming_a_stanza_skips_the_built_in_pipeline() { + use wacore::types::events::ChannelEventHandler; + let client = create_interceptor_test_client().await; + let (handler, events) = ChannelEventHandler::new(); + client.subscribe_handler(handler).detach(); + + let (interceptor, seen) = recorder("receipt"); + let _handle = client.add_stanza_interceptor(interceptor); + + client + .process_node(node_to_owned_ref(receipt_stanza())) + .await; + + assert_eq!(seen.lock().expect("recorder lock").len(), 1); + assert!( + events.try_recv().is_err(), + "the built-in receipt handler must not have dispatched" + ); +} + +#[tokio::test] +async fn interception_does_not_touch_connection_bookkeeping() { + // Offline-sync tracking runs before dispatch, and must keep running: it is + // what tells the client the drain finished. An interceptor exists to take + // over *handling* a stanza, not to opt out of staying connected. + let client = create_interceptor_test_client().await; + client + .offline_sync_metrics + .active + .store(true, Ordering::Release); + + let (interceptor, seen) = recorder("ib"); + let _handle = client.add_stanza_interceptor(interceptor); + + let node = NodeBuilder::new("ib") + .children([NodeBuilder::new("offline").attr("count", "0").build()]) + .build(); + client.process_node(node_to_owned_ref(node)).await; + + assert_eq!(seen.lock().expect("recorder lock").len(), 1, "it ran"); + assert!( + !client.offline_sync_metrics.active.load(Ordering::Acquire), + "offline-sync tracking still ran, claimed or not" + ); +} + +#[tokio::test] +async fn the_first_interceptor_to_claim_a_stanza_wins() { + let client = create_interceptor_test_client().await; + let (first, first_seen) = recorder("receipt"); + let (second, second_seen) = recorder("receipt"); + + let _a = client.add_stanza_interceptor(first); + let _b = client.add_stanza_interceptor(second); + + client + .process_node(node_to_owned_ref(NodeBuilder::new("receipt").build())) + .await; + + assert_eq!(first_seen.lock().expect("lock").len(), 1); + assert!( + second_seen.lock().expect("lock").is_empty(), + "registration order is priority order" + ); +} + +#[tokio::test] +async fn a_passing_interceptor_does_not_stop_the_next_one() { + let client = create_interceptor_test_client().await; + let (first, first_seen) = recorder("nothing"); + let (second, second_seen) = recorder("receipt"); + + let _a = client.add_stanza_interceptor(first); + let _b = client.add_stanza_interceptor(second); + + client + .process_node(node_to_owned_ref(NodeBuilder::new("receipt").build())) + .await; + + assert_eq!(first_seen.lock().expect("lock").len(), 1); + assert_eq!(second_seen.lock().expect("lock").len(), 1); +} + +#[tokio::test] +async fn an_unregistered_interceptor_stops_seeing_stanzas() { + let client = create_interceptor_test_client().await; + let (interceptor, seen) = recorder("nothing"); + let handle = client.add_stanza_interceptor(interceptor); + + client + .process_node(node_to_owned_ref(NodeBuilder::new("receipt").build())) + .await; + assert_eq!(seen.lock().expect("lock").len(), 1); + + drop(handle); + client + .process_node(node_to_owned_ref(NodeBuilder::new("receipt").build())) + .await; + assert_eq!( + seen.lock().expect("lock").len(), + 1, + "no further stanzas after unregistering" + ); +} + +#[tokio::test] +async fn a_claimed_unknown_stanza_is_acked_rather_than_nacked() { + // The reason this exists. A tag the client does not model is nacked, which + // tells the server this client cannot act on it. An interceptor that *can* + // act on it says the opposite — but it must still say something: answering + // nothing leaves the stanza in the offline queue and keeps the stream + // recycling. + let (client, transport) = crate::test_utils::create_iq_test_client().await; + let (interceptor, seen) = recorder("vendor:thing"); + let _handle = client.add_stanza_interceptor(interceptor); + + let node = NodeBuilder::new("vendor:thing") + .attr("id", "V-1") + .attr("from", "s.whatsapp.net") + .build(); + // `should_ack` covers only the tags the client models, so this stanza takes + // the claimed-stanza ack path rather than the deferred one. + assert!( + !client.should_ack(&node.as_node_ref()), + "fixture must exercise the path should_ack does not cover" + ); + client.process_node(node_to_owned_ref(node)).await; + assert_eq!(seen.lock().expect("lock").len(), 1, "the interceptor ran"); + + // What the server actually receives, which is the whole point: an ack, and + // not the nack this tag would otherwise have drawn. + let sent = crate::test_utils::decode_sent_iq(&transport, 0).await; + let sent = sent.get(); + assert_eq!(sent.tag.as_ref(), "ack"); + assert!( + sent.get_attr("class") + .is_some_and(|class| *class == "vendor:thing") + ); + assert!(sent.get_attr("id").is_some_and(|id| *id == "V-1")); + assert_eq!( + transport.sent_count(), + 1, + "one answer, so no nack followed the ack" + ); +} + +#[tokio::test] +async fn claiming_a_stanza_the_client_answers_differently_sends_no_generic_ack() { + // A direct is answered with a delivery , an with an + // . Neither is an , so the claimed-stanza + // path stays quiet rather than sending the server something it did not ask + // for. The interceptor that took the stanza owes the reply. + let (client, transport) = crate::test_utils::create_iq_test_client().await; + let (interceptor, seen) = recorder("message"); + let _handle = client.add_stanza_interceptor(interceptor); + + let node = NodeBuilder::new("message") + .attr("id", "M-1") + .attr("from", "5511999998888@s.whatsapp.net") + .build(); + assert!( + !client.should_ack(&node.as_node_ref()), + "a direct message is not ack-answered" + ); + client.process_node(node_to_owned_ref(node)).await; + + assert_eq!(seen.lock().expect("lock").len(), 1, "it was claimed"); + crate::test_utils::wait_for_outbound_tasks(&client).await; + assert_eq!( + transport.sent_count(), + 0, + "no invented answer for a tag the client models" + ); +} + +#[tokio::test] +async fn a_claimed_stanza_without_identity_is_left_alone() { + // An ack has nothing to address without `id` and `from`, which is the same + // condition the nack path checks. + let client = create_interceptor_test_client().await; + let (interceptor, seen) = recorder("vendor:thing"); + let _handle = client.add_stanza_interceptor(interceptor); + + client + .process_node(node_to_owned_ref(NodeBuilder::new("vendor:thing").build())) + .await; + + assert_eq!(seen.lock().expect("lock").len(), 1); +} + +#[tokio::test] +async fn connection_critical_stanzas_are_never_offered_to_an_interceptor() { + // An interceptor exists to extend a client, not to leave it + // authenticated-but-unaware, never reconnecting, or waiting forever on a + // send that already completed. These four settle connection state, so they + // do not reach an interceptor at all — an interceptor that tried to claim + // one never gets the chance. + let client = create_interceptor_test_client().await; + let (interceptor, seen) = recorder("claims-everything-it-sees"); + let _handle = client.add_stanza_interceptor(interceptor); + + for tag in ["success", "failure", "stream:error", "ack"] { + client + .process_node(node_to_owned_ref(NodeBuilder::new(tag).build())) + .await; + } + + assert!( + seen.lock().expect("lock").is_empty(), + "a connection-critical stanza must not reach an interceptor" + ); + + // A neighbouring tag is still offered, so the guard is a list and not a + // switch that turned interception off. + client + .process_node(node_to_owned_ref(NodeBuilder::new("notification").build())) + .await; + assert_eq!(seen.lock().expect("lock").len(), 1); +} + +#[tokio::test] +async fn a_server_ping_is_never_offered_to_an_interceptor() { + // A claimed ping is a pong never sent, and the server closes the connection + // over it — the same class of harm as claiming `success` or `ack`, so the + // same protection. Every other stays offered: that is the traffic an + // interceptor exists to extend. + let client = create_interceptor_test_client().await; + let (interceptor, seen) = recorder("claims-everything-it-sees"); + let _handle = client.add_stanza_interceptor(interceptor); + + for node in [ + NodeBuilder::new("iq") + .attr("from", "s.whatsapp.net") + .attr("id", "PING-1") + .attr("type", "get") + .children([NodeBuilder::new("ping").build()]) + .build(), + // WA Web's handleIq is type-agnostic, so an absent type is a ping too. + NodeBuilder::new("iq") + .attr("from", "s.whatsapp.net") + .attr("id", "PING-2") + .attr("xmlns", "urn:xmpp:ping") + .build(), + ] { + client.process_node(node_to_owned_ref(node)).await; + } + assert!( + seen.lock().expect("lock").is_empty(), + "a server ping must not reach an interceptor" + ); + + // A ping *response* is ours, not the server's, so it carries no pong + // obligation and stays offered. + let response = NodeBuilder::new("iq") + .attr("from", "s.whatsapp.net") + .attr("id", "PING-3") + .attr("type", "result") + .children([NodeBuilder::new("ping").build()]) + .build(); + client.process_node(node_to_owned_ref(response)).await; + + let other = NodeBuilder::new("iq") + .attr("from", "s.whatsapp.net") + .attr("id", "IQ-1") + .attr("type", "get") + .children([NodeBuilder::new("query").build()]) + .build(); + client.process_node(node_to_owned_ref(other)).await; + assert_eq!( + seen.lock().expect("lock").len(), + 2, + "a ping result and an ordinary are both still offered" + ); +} + +#[tokio::test] +async fn a_closure_can_be_an_interceptor() { + let client = create_interceptor_test_client().await; + let calls = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&calls); + + let _handle = client.add_stanza_interceptor(Arc::new(move |_node: &OwnedNodeRef| { + counter.fetch_add(1, Ordering::Relaxed); + Interception::Pass + })); + + client + .process_node(node_to_owned_ref(NodeBuilder::new("receipt").build())) + .await; + assert_eq!(calls.load(Ordering::Relaxed), 1); +} + +#[tokio::test] +async fn a_handle_outliving_its_client_does_not_keep_it_alive() { + // The handle holds a weak reference, so a forgotten one cannot pin a + // client — and dropping it afterwards must not panic either. + let handle = { + let client = create_interceptor_test_client().await; + let (interceptor, _seen) = recorder("nothing"); + let owners = Arc::strong_count(&client); + let handle = client.add_stanza_interceptor(interceptor); + assert_eq!( + Arc::strong_count(&client), + owners, + "registering must not make the handle an owner" + ); + handle + }; + drop(handle); +} diff --git a/src/handlers/router.rs b/src/handlers/router.rs index b76471bec..a9ab054c0 100644 --- a/src/handlers/router.rs +++ b/src/handlers/router.rs @@ -57,6 +57,14 @@ impl StanzaRouter { } } + /// Whether this client models `tag` at all. + /// + /// A tag with no handler is one [`dispatch`](Self::dispatch) would report + /// unhandled, which is what makes the caller nack it. + pub fn models(&self, tag: &str) -> bool { + self.handlers.contains_key(tag) + } + /// Get the number of registered handlers (useful for testing). pub fn handler_count(&self) -> usize { self.handlers.len() diff --git a/src/lib.rs b/src/lib.rs index b351ed6a0..393b04433 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -120,6 +120,7 @@ pub(crate) mod flush_scope; /// types embed it. pub use client::ClientError; pub use client::NodeFilter; +pub use client::interceptor::{Interception, InterceptorHandle, StanzaInterceptor}; pub use client::{ AllocSnapshot, CollectionStats, HttpResourceReport, MemoryReport, ResourceReport, StatsSnapshot, StorageResourceReport, TransportResourceReport,