diff --git a/src/bot.rs b/src/bot.rs index 0e3da0ef5..3a5800f54 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -487,6 +487,7 @@ pub struct BotBuilder< // Optional fields event_handlers: Vec, raw_handlers: Vec>, + pre_ack_message_hook: Option, custom_enc_handlers: HashMap>, override_version: Option<(u32, u32, u32)>, device_props_override: Option, @@ -508,6 +509,7 @@ impl BotBuilder BotBuilder { runtime: self.runtime, event_handlers: self.event_handlers, raw_handlers: self.raw_handlers, + pre_ack_message_hook: self.pre_ack_message_hook, custom_enc_handlers: self.custom_enc_handlers, override_version: self.override_version, device_props_override: self.device_props_override, @@ -639,6 +642,27 @@ impl BotBuilder { }) } + /// Run `hook` inline before acknowledging each parsed inbound message. + /// + /// This is the durability hook for consumers that need strict + /// "ACK-after-commit" semantics. The hook receives the same normalized + /// parsed message and [`MessageInfo`] that will be emitted as + /// `Event::Message` after ACK. Return `Ok(())` only after your durable write + /// has committed. Returning `Err` suppresses both the ACK/receipt and the + /// normal message event for this delivery attempt so WhatsApp can retry. + /// + /// Keep this hook narrow and fast: it runs on the receive lane before ACK. + pub fn on_pre_ack_message(mut self, hook: F) -> Self + where + F: Fn(crate::message::ParsedMessagePreAckContext) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + self.pre_ack_message_hook = Some(Arc::new( + move |ctx| -> crate::message::ParsedMessagePreAckHookFuture { Box::pin(hook(ctx)) }, + )); + self + } + /// Run `handler` with the QR payload (and validity window) each time a /// pairing QR code is issued. Render `code` as a QR image for scanning. pub fn on_qr_code(self, handler: F) -> Self @@ -969,6 +993,10 @@ impl BotBuilder { // map once; the receive hot path then reads it lock-free. let _ = client.custom_enc_handlers.set(self.custom_enc_handlers); + if let Some(hook) = self.pre_ack_message_hook { + let _ = client.set_parsed_message_pre_ack_hook_arc(hook); + } + if self.skip_history_sync { client.set_skip_history_sync(true); } diff --git a/src/client.rs b/src/client.rs index a5b4f5414..7e646f440 100644 --- a/src/client.rs +++ b/src/client.rs @@ -501,6 +501,16 @@ pub struct Client { /// in `WAWebMessageProcessPlaceholder`. pub(crate) undecryptable_dispatched: Cache, + /// Parsed messages whose pre-ACK hook failed after Signal state advanced. + /// A same-process server redelivery may then arrive as a Signal duplicate; + /// this cache lets that duplicate retry the durable hook with the original + /// normalized message instead of ACKing an uncommitted delivery. + pub(crate) pending_parsed_message_pre_ack: + Cache, + /// Reliable barrier for Signal flush deferral. Kept separate from the cache + /// because `PortableCache::entry_count()` is intentionally best-effort. + pub(crate) pending_parsed_message_pre_ack_count: AtomicUsize, + pub enable_auto_reconnect: Arc, pub auto_reconnect_errors: Arc, @@ -563,6 +573,11 @@ pub struct Client { /// `OnceLock::get` (no lock) and no per-node guard acquisition. pub custom_enc_handlers: std::sync::OnceLock>>, + /// Optional awaited hook for consumers that need to durably commit a parsed + /// inbound message before the SDK sends the delivery receipt / transport ACK. + pub(crate) parsed_message_pre_ack_hook: + std::sync::OnceLock, + /// Chat state (typing indicator) handlers registered by external consumers. /// Each handler receives a `ChatStateEvent` describing the chat, optional participant and state. pub(crate) chatstate_handlers: Arc>>, diff --git a/src/client/accessors.rs b/src/client/accessors.rs index 3a979498f..551fef516 100644 --- a/src/client/accessors.rs +++ b/src/client/accessors.rs @@ -23,6 +23,50 @@ impl Client { self.core.event_bus.add_handler(handler); } + /// Register an awaited hook that runs after an inbound message has been + /// parsed and normalized, but before the SDK sends its delivery receipt or + /// transport ACK. + /// + /// The hook is intended for durability-sensitive consumers: persist the + /// message, commit the transaction, then return `Ok(())`. If the hook + /// returns an error, `whatsapp-rust` suppresses the ACK and the normal + /// `Event::Message` dispatch for this attempt, leaving the server free to + /// redeliver/retry instead of marking the message delivered. + /// + /// Only one hook can be registered for a `Client`. Applications using + /// [`BotBuilder`](crate::bot::BotBuilder) can prefer + /// [`BotBuilder::on_pre_ack_message`](crate::bot::BotBuilder::on_pre_ack_message) + /// to wire the hook during construction. + /// + /// Registering this hook also keeps inbound message processing serial. That + /// preserves the durability contract when a hook fails after Signal decrypt: + /// the SDK can defer Signal cache flushes while the uncommitted parsed + /// message is pending, without allowing another inbound message to race the + /// pending retry path. + pub fn set_parsed_message_pre_ack_hook( + &self, + hook: F, + ) -> Result<(), crate::message::ParsedMessagePreAckHookAlreadySet> + where + F: Fn(crate::message::ParsedMessagePreAckContext) -> Fut + Send + Sync + 'static, + Fut: std::future::Future> + Send + 'static, + { + self.set_parsed_message_pre_ack_hook_arc(Arc::new( + move |ctx| -> crate::message::ParsedMessagePreAckHookFuture { Box::pin(hook(ctx)) }, + )) + } + + pub(crate) fn set_parsed_message_pre_ack_hook_arc( + &self, + hook: crate::message::ParsedMessagePreAckHook, + ) -> Result<(), crate::message::ParsedMessagePreAckHookAlreadySet> { + self.parsed_message_pre_ack_hook + .set(hook) + .map_err(|_| crate::message::ParsedMessagePreAckHookAlreadySet)?; + self.swap_message_semaphore(1); + Ok(()) + } + /// Enable or disable raw node forwarding. /// When enabled, `Event::RawNode` is emitted for every decoded stanza before /// the stanza router dispatches it. Only enable when external consumers need diff --git a/src/client/adapters.rs b/src/client/adapters.rs index fa0d73cdb..367633598 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -59,6 +59,17 @@ impl Client { /// Flush the in-memory signal cache to the database backend. /// Called after each message is decrypted or after encryption operations. pub(crate) async fn flush_signal_cache(&self) -> Result<(), anyhow::Error> { + if self.parsed_message_pre_ack_hook.get().is_some() + && self + .pending_parsed_message_pre_ack_count + .load(Ordering::Acquire) + > 0 + { + return Err(anyhow::anyhow!( + "Signal cache flush deferred while parsed-message pre-ACK commit is pending" + )); + } + // Hold no device guard across the flush: this per-message batched SQLite // write would otherwise block every concurrent Device write for its duration. let backend = self diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index d713560cf..74ada06b3 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -192,6 +192,13 @@ impl Client { ), undecryptable_dispatched: cache_config.undecryptable_dispatched.build_with_ttl(), + // Intentionally unbounded/no-TTL: an entry means Signal decrypt may + // have advanced volatile state for a message the application has + // not durably committed. Expiring it would allow a later Signal + // flush to make that advance durable and turn server redelivery + // into an ACKed duplicate without rerunning the hook. + pending_parsed_message_pre_ack: Cache::builder().build(), + pending_parsed_message_pre_ack_count: AtomicUsize::new(0), offline_sync_metrics: Arc::new(OfflineSyncMetrics { active: AtomicBool::new(false), @@ -226,6 +233,7 @@ impl Client { pairing_cancellation_tx: Arc::new(Mutex::new(None)), pair_code_state: Arc::new(Mutex::new(wacore::pair_code::PairCodeState::default())), custom_enc_handlers: std::sync::OnceLock::new(), + parsed_message_pre_ack_hook: std::sync::OnceLock::new(), chatstate_handlers: Arc::new(RwLock::new(Vec::new())), pdo_pending_requests: cache_config.pdo_pending_requests.build_with_ttl(), pdo_requested: cache_config.pdo_requested.build_with_ttl(), diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 527d26025..b676996f6 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -21,6 +21,15 @@ impl Client { /// a consistent (generation, Arc) pair. Must be called from a non-async /// context or inside a scoped block (MutexGuard is !Send). pub(crate) fn swap_message_semaphore(&self, permits: usize) { + let permits = if self.parsed_message_pre_ack_hook.get().is_some() { + // A pre-ACK hook can intentionally fail after Signal decrypt has + // advanced volatile state. Keep message processing serial while the + // hook is installed so pending retry/flush deferral has one inbound + // message to reason about at a time. + 1 + } else { + permits + }; let mut guard = match self.message_processing_semaphore.lock() { Ok(g) => g, Err(poisoned) => poisoned.into_inner(), @@ -415,9 +424,13 @@ impl Client { }; match tag { "receipt" | "notification" | "call" => true, - "message" => from - .to_jid() - .is_some_and(|j| j.is_newsletter() || j.is_status_broadcast()), + "message" => { + if self.parsed_message_pre_ack_hook.get().is_some() { + return false; + } + from.to_jid() + .is_some_and(|j| j.is_newsletter() || j.is_status_broadcast()) + } _ => false, } } diff --git a/src/lib.rs b/src/lib.rs index 14291c917..b37e5e004 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -121,6 +121,7 @@ pub mod version; pub mod prelude { pub use crate::bot::{Bot, BotBuilder, BotHandle, MessageContext}; pub use crate::client::{Client, ClientError}; + pub use crate::message::{ParsedMessagePreAckContext, ParsedMessagePreAckHookAlreadySet}; pub use crate::request::IqError; #[cfg(feature = "tokio-runtime")] pub use crate::runtime_impl::TokioRuntime; diff --git a/src/message.rs b/src/message.rs index be81b7d22..0ff5d315c 100644 --- a/src/message.rs +++ b/src/message.rs @@ -4,6 +4,8 @@ use crate::types::message::MessageInfo; use log::{debug, warn}; use prost::Message as ProtoMessage; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use wacore::libsignal::crypto::DecryptionError; use wacore::libsignal::protocol::SenderKeyDistributionMessage; @@ -28,6 +30,61 @@ use waproto::whatsapp::{self as wa}; /// After this many retries, we stop sending retry receipts and rely solely on PDO. const MAX_DECRYPT_RETRIES: u8 = 5; +/// Future returned by a parsed-message pre-ACK hook. +/// +/// The hook is awaited inline on the receive path after the inbound payload has +/// been parsed and normalized, but before the SDK sends the delivery receipt or +/// transport ACK. Returning an error intentionally suppresses the ACK so the +/// server can redeliver the message later. This is useful for consumers that +/// must durably commit an inbound message before WhatsApp considers it +/// delivered. +pub(crate) type ParsedMessagePreAckHookFuture = + Pin> + Send + 'static>>; + +/// Callback registered with [`Client::set_parsed_message_pre_ack_hook`]. +/// +/// [`Client::set_parsed_message_pre_ack_hook`]: crate::Client::set_parsed_message_pre_ack_hook +pub(crate) type ParsedMessagePreAckHook = + Arc ParsedMessagePreAckHookFuture + Send + Sync>; + +/// Message context passed to the awaited pre-ACK hook. +/// +/// Both `message` and `info` are `Arc`-wrapped so the hook can cheaply share the +/// same normalized values that the event bus will receive after ACK. The +/// `client` handle is included for applications that need account or storage +/// context while performing a durable commit. +#[derive(Clone)] +pub struct ParsedMessagePreAckContext { + pub message: Arc, + pub info: Arc, + pub client: Arc, +} + +impl ParsedMessagePreAckContext { + pub(crate) fn new( + message: Arc, + info: Arc, + client: Arc, + ) -> Self { + Self { + message, + info, + client, + } + } +} + +/// Returned when a parsed-message pre-ACK hook has already been registered. +#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)] +#[error("parsed message pre-ACK hook is already registered")] +pub struct ParsedMessagePreAckHookAlreadySet; + +#[derive(Clone)] +pub(crate) struct PendingParsedMessagePreAck { + pub message: Arc, + pub info: Arc, +} + /// Pre-extracted enc node payload. Holds owned copies of the fields needed for /// decryption so the async decrypt phase doesn't borrow the original NodeRef tree. pub(crate) struct EncPayload { @@ -77,16 +134,23 @@ pub(crate) struct SessionBatchOutcome { duplicate: bool, undecryptable: bool, dispatched: bool, + pre_ack_hook_failed: bool, skdm_only: bool, plaintext_failed: bool, had_failure: bool, } +#[derive(Clone, Copy, Debug, Default)] +struct GroupBatchOutcome { + pre_ack_hook_failed: bool, +} + #[derive(Clone, Copy, Debug, Default)] struct MigrationDecryptOutcome { decrypted: bool, duplicate: bool, dispatched: bool, + pre_ack_hook_failed: bool, skdm_only: bool, plaintext_failed: bool, } @@ -94,6 +158,7 @@ struct MigrationDecryptOutcome { #[derive(Clone, Copy, Debug, Default)] pub(crate) struct PlaintextHandleOutcome { dispatched: bool, + pre_ack_hook_failed: bool, skdm_only: bool, } diff --git a/src/message/dispatch.rs b/src/message/dispatch.rs index 70b867ab2..d5eff5c1f 100644 --- a/src/message/dispatch.rs +++ b/src/message/dispatch.rs @@ -9,7 +9,7 @@ impl Client { self: &Arc, msg: wa::Message, info: &Arc, - ) { + ) -> bool { use wacore::proto_helpers::MessageExt; wacore::telemetry::recv("decrypted"); @@ -36,12 +36,153 @@ impl Client { { Arc::make_mut(&mut info).comment_target = Some(target); } - let dispatch_msg = decrypted.unwrap_or(msg); - self.ack_received_message(&info); + let dispatch_msg = Arc::new(decrypted.unwrap_or(msg)); + + if !self + .await_parsed_message_pre_ack(Arc::clone(&dispatch_msg), Arc::clone(&info), true) + .await + { + return false; + } + + self.remove_pending_pre_ack_message(&Self::pre_ack_message_key(&info)) + .await; + self.ack_parsed_message_after_pre_ack(&info); self.core .event_bus - .dispatch(Event::Message(Arc::new(dispatch_msg), info)); + .dispatch(Event::Message(dispatch_msg, info)); + true + } + + async fn await_parsed_message_pre_ack( + self: &Arc, + message: Arc, + info: Arc, + store_pending_on_failure: bool, + ) -> bool { + let Some(hook) = self.parsed_message_pre_ack_hook.get().cloned() else { + return true; + }; + + if let Err(err) = hook(crate::message::ParsedMessagePreAckContext::new( + Arc::clone(&message), + Arc::clone(&info), + Arc::clone(self), + )) + .await + { + warn!( + "Parsed-message pre-ACK hook failed for message {}; suppressing ACK for retry/redelivery: {err:#}", + info.id + ); + if store_pending_on_failure { + self.store_pending_pre_ack_message(message, info).await; + } + return false; + } + + true + } + + pub(crate) async fn dispatch_pending_pre_ack_message( + self: &Arc, + info: &Arc, + ) -> Option { + let key = Self::pre_ack_message_key(info); + let pending = self.pending_parsed_message_pre_ack.get(&key).await?; + + if !self + .await_parsed_message_pre_ack( + Arc::clone(&pending.message), + Arc::clone(&pending.info), + false, + ) + .await + { + return Some(false); + } + + self.remove_pending_pre_ack_message(&key).await; + self.ack_parsed_message_after_pre_ack(&pending.info); + self.core + .event_bus + .dispatch(Event::Message(pending.message, pending.info)); + Some(true) + } + + fn pre_ack_message_key(info: &MessageInfo) -> wacore::types::message::ChatMessageId { + wacore::types::message::ChatMessageId::new(info.source.chat.clone(), info.id.clone()) + } + + async fn store_pending_pre_ack_message( + self: &Arc, + message: Arc, + info: Arc, + ) { + let key = Self::pre_ack_message_key(&info); + if self + .pending_parsed_message_pre_ack + .get(&key) + .await + .is_none() + { + self.pending_parsed_message_pre_ack_count + .fetch_add(1, std::sync::atomic::Ordering::Release); + } + self.pending_parsed_message_pre_ack + .insert( + key, + crate::message::PendingParsedMessagePreAck { message, info }, + ) + .await; + } + + async fn remove_pending_pre_ack_message( + self: &Arc, + key: &wacore::types::message::ChatMessageId, + ) { + if self + .pending_parsed_message_pre_ack + .remove(key) + .await + .is_some() + { + self.pending_parsed_message_pre_ack_count + .fetch_sub(1, std::sync::atomic::Ordering::AcqRel); + } + } + + fn ack_parsed_message_after_pre_ack(self: &Arc, info: &Arc) { + if self.parsed_message_pre_ack_hook.get().is_some() + && info.source.chat.is_status_broadcast() + { + self.ack_received_message(info); + self.spawn_message_ack(info); + } else if self.parsed_message_pre_ack_hook.get().is_some() + && info.source.chat.is_newsletter() + { + self.spawn_message_ack(info); + } else { + self.ack_received_message(info); + } + } + + pub(crate) fn ack_status_drop_after_pre_ack_hook(self: &Arc, info: &Arc) { + if self.parsed_message_pre_ack_hook.get().is_some() + && info.source.chat.is_status_broadcast() + { + self.spawn_message_ack(info); + } + } + + pub(crate) fn ack_newsletter_drop_after_pre_ack_hook( + self: &Arc, + info: &Arc, + ) { + if self.parsed_message_pre_ack_hook.get().is_some() && info.source.chat.is_newsletter() { + self.spawn_message_ack(info); + } } /// Acknowledge a received message so the server drops it from the offline diff --git a/src/message/receive.rs b/src/message/receive.rs index 407641097..d6b2f3b1b 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -111,7 +111,8 @@ impl Client { .await; let client = Arc::clone(self); let info2 = Arc::clone(&info); - let skip_ack = info.source.chat.is_status_broadcast(); + let skip_ack = info.source.chat.is_status_broadcast() + && self.parsed_message_pre_ack_hook.get().is_none(); self.outbound_flush.spawn(&*self.runtime, async move { // Only ack once the PDO request is out (or skipped as ancient); // a transient send failure leaves it queued for redelivery. @@ -250,7 +251,9 @@ impl Client { "[msg:{}] All enc payloads unrecognized; transport-acking to drop from offline queue", info.id ); - if !info.source.chat.is_status_broadcast() { + if !info.source.chat.is_status_broadcast() + || self.parsed_message_pre_ack_hook.get().is_some() + { self.spawn_node_transport_ack(nr).await; } return None; @@ -386,7 +389,7 @@ impl Client { should_process_skmsg_after_session(session_payloads.is_empty(), session_outcome); if should_process_skmsg { - match self + let group_outcome = match self .clone() .process_group_enc_batch( &group_payloads, @@ -396,8 +399,9 @@ impl Client { ) .await { - Ok(()) => { + Ok(outcome) => { // Processed successfully or handled errors (e.g. sent retry receipt) + outcome } Err(e) => { log::warn!( @@ -406,7 +410,11 @@ impl Client { info.source.sender.observe(), info.source.chat.observe() ); + GroupBatchOutcome::default() } + }; + if group_outcome.pre_ack_hook_failed { + return; } } else { // Only show warning if session messages actually FAILED (not duplicates) @@ -474,14 +482,19 @@ impl Client { } else if session_had_duplicates && !session_decrypted_successfully && !session_dispatched_undecryptable - && !info.source.chat.is_status_broadcast() + && (!info.source.chat.is_status_broadcast() + || self.parsed_message_pre_ack_hook.get().is_some()) { // Duplicate (already-processed) with no group content: ack it so the // server drops it from the offline queue (whatsmeow/WA Web treat // old-counter like success). status is acked by the should_ack gate // (a status SKDM pkmsg can reach here), so skip it to avoid a // redundant receipt. - self.ack_received_message(&info); + match self.dispatch_pending_pre_ack_message(&info).await { + Some(true) => {} + Some(false) => return, + None => self.ack_received_message(&info), + } } else if should_ack_skdm_only_session_fallback(session_outcome, bot_payloads.is_empty()) { // SKDM-only session decrypts skip dispatch, so this stanza would // otherwise stay queued. WA Web and whatsmeow ack every decrypted @@ -491,6 +504,10 @@ impl Client { self.ack_received_message(&info); } + if session_outcome.pre_ack_hook_failed { + return; + } + // Bot-secret (msmsg) payloads run inline here so they're serialised // with the session/group decrypt batches under the same global // permit + per-chat enqueue lock acquired upstream. @@ -673,6 +690,8 @@ impl Client { Ok(plaintext_outcome) => { outcome.decrypted = true; outcome.dispatched |= plaintext_outcome.dispatched; + outcome.pre_ack_hook_failed |= plaintext_outcome.pre_ack_hook_failed; + outcome.had_failure |= plaintext_outcome.pre_ack_hook_failed; outcome.skdm_only |= plaintext_outcome.skdm_only; } Err(e) => { @@ -791,6 +810,10 @@ impl Client { Ok(plaintext_outcome) => { outcome.decrypted = true; outcome.dispatched |= plaintext_outcome.dispatched; + outcome.pre_ack_hook_failed |= + plaintext_outcome.pre_ack_hook_failed; + outcome.had_failure |= + plaintext_outcome.pre_ack_hook_failed; outcome.skdm_only |= plaintext_outcome.skdm_only; } Err(e) => { @@ -842,6 +865,7 @@ impl Client { if migration_outcome.decrypted || migration_outcome.duplicate || migration_outcome.plaintext_failed + || migration_outcome.pre_ack_hook_failed { outcome.decrypted |= migration_outcome.decrypted; outcome.duplicate |= migration_outcome.duplicate; @@ -850,6 +874,10 @@ impl Client { outcome.plaintext_failed |= migration_outcome.plaintext_failed; outcome.had_failure |= migration_outcome.plaintext_failed; + outcome.pre_ack_hook_failed |= + migration_outcome.pre_ack_hook_failed; + outcome.had_failure |= + migration_outcome.pre_ack_hook_failed; if migration_outcome.plaintext_failed { outcome.undecryptable |= self .handle_plaintext_failure(info, decrypt_fail_mode) @@ -926,6 +954,7 @@ impl Client { if migration_outcome.decrypted || migration_outcome.duplicate || migration_outcome.plaintext_failed + || migration_outcome.pre_ack_hook_failed { outcome.decrypted |= migration_outcome.decrypted; outcome.duplicate |= migration_outcome.duplicate; @@ -933,6 +962,8 @@ impl Client { outcome.skdm_only |= migration_outcome.skdm_only; outcome.plaintext_failed |= migration_outcome.plaintext_failed; outcome.had_failure |= migration_outcome.plaintext_failed; + outcome.pre_ack_hook_failed |= migration_outcome.pre_ack_hook_failed; + outcome.had_failure |= migration_outcome.pre_ack_hook_failed; if migration_outcome.plaintext_failed { outcome.undecryptable |= self.handle_plaintext_failure(info, decrypt_fail_mode).await; @@ -974,6 +1005,7 @@ impl Client { if migration_outcome.decrypted || migration_outcome.duplicate || migration_outcome.plaintext_failed + || migration_outcome.pre_ack_hook_failed { outcome.decrypted |= migration_outcome.decrypted; outcome.duplicate |= migration_outcome.duplicate; @@ -981,6 +1013,8 @@ impl Client { outcome.skdm_only |= migration_outcome.skdm_only; outcome.plaintext_failed |= migration_outcome.plaintext_failed; outcome.had_failure |= migration_outcome.plaintext_failed; + outcome.pre_ack_hook_failed |= migration_outcome.pre_ack_hook_failed; + outcome.had_failure |= migration_outcome.pre_ack_hook_failed; if migration_outcome.plaintext_failed { outcome.undecryptable |= self.handle_plaintext_failure(info, decrypt_fail_mode).await; @@ -1031,6 +1065,7 @@ impl Client { if migration_outcome.decrypted || migration_outcome.duplicate || migration_outcome.plaintext_failed + || migration_outcome.pre_ack_hook_failed { outcome.decrypted |= migration_outcome.decrypted; outcome.duplicate |= migration_outcome.duplicate; @@ -1038,6 +1073,8 @@ impl Client { outcome.skdm_only |= migration_outcome.skdm_only; outcome.plaintext_failed |= migration_outcome.plaintext_failed; outcome.had_failure |= migration_outcome.plaintext_failed; + outcome.pre_ack_hook_failed |= migration_outcome.pre_ack_hook_failed; + outcome.had_failure |= migration_outcome.pre_ack_hook_failed; if migration_outcome.plaintext_failed { outcome.undecryptable |= self.handle_plaintext_failure(info, decrypt_fail_mode).await; @@ -1098,11 +1135,12 @@ impl Client { info: &Arc, _sender_encryption_jid: &Jid, decrypt_fail_mode: crate::types::events::DecryptFailMode, - ) -> Result<(), DecryptionError> { + ) -> Result { if payloads.is_empty() { - return Ok(()); + return Ok(GroupBatchOutcome::default()); } let mut adapter = self.signal_adapter().await; + let mut outcome = GroupBatchOutcome::default(); // Always use bare sender for sender key operations. Real WA delivers // skmsg with bare participant but pkmsg (SKDM) with device-qualified @@ -1141,7 +1179,7 @@ impl Client { self.handle_unknown_device_sync(info).await; } - if let Err(e) = self + match self .clone() .handle_decrypted_plaintext( "skmsg", @@ -1151,7 +1189,12 @@ impl Client { ) .await { - log::warn!("Failed processing group plaintext (batch): {e:?}"); + Ok(plaintext_outcome) => { + outcome.pre_ack_hook_failed |= plaintext_outcome.pre_ack_hook_failed; + } + Err(e) => { + log::warn!("Failed processing group plaintext (batch): {e:?}"); + } } } Err(SignalProtocolError::DuplicatedMessage(iteration, counter)) => { @@ -1165,8 +1208,15 @@ impl Client { // Redelivered duplicate: ack it so the server drops it from the // offline queue. status is already acked by the should_ack gate, // so skip it to avoid a redundant receipt. - if !info.source.chat.is_status_broadcast() { - self.ack_received_message(info); + match self.dispatch_pending_pre_ack_message(info).await { + Some(true) => {} + Some(false) => outcome.pre_ack_hook_failed = true, + None if !info.source.chat.is_status_broadcast() + || self.parsed_message_pre_ack_hook.get().is_some() => + { + self.ack_received_message(info); + } + None => {} } } Err(SignalProtocolError::NoSenderKeyState(msg)) => { @@ -1176,6 +1226,7 @@ impl Client { info.id, info.source.sender.observe() ); + self.ack_status_drop_after_pre_ack_hook(info); continue; } @@ -1208,6 +1259,7 @@ impl Client { info.source.sender.observe(), e ); + self.ack_status_drop_after_pre_ack_hook(info); continue; } @@ -1231,11 +1283,13 @@ impl Client { .await; if !info.source.chat.is_status_broadcast() { self.spawn_nack(info, NackReason::UnhandledError, None); + } else { + self.ack_status_drop_after_pre_ack_hook(info); } } } } - Ok(()) + Ok(outcome) } /// WA Web: online → `syncDeviceListJob`, offline → `OfflinePendingDeviceCache`. @@ -1406,9 +1460,10 @@ impl Client { ..Default::default() }) } else { - self.dispatch_parsed_message(msg, info).await; + let dispatched = self.dispatch_parsed_message(msg, info).await; Ok(PlaintextHandleOutcome { - dispatched: true, + dispatched, + pre_ack_hook_failed: !dispatched, ..Default::default() }) } @@ -1506,6 +1561,7 @@ impl Client { Ok(plaintext_outcome) => MigrationDecryptOutcome { decrypted: true, dispatched: plaintext_outcome.dispatched, + pre_ack_hook_failed: plaintext_outcome.pre_ack_hook_failed, skdm_only: plaintext_outcome.skdm_only, ..Default::default() }, diff --git a/src/message/special.rs b/src/message/special.rs index 618b21d81..6c3e18825 100644 --- a/src/message/special.rs +++ b/src/message/special.rs @@ -19,6 +19,7 @@ impl Client { info.id, node.tag ); + self.ack_newsletter_drop_after_pre_ack_hook(info); return; }; @@ -37,6 +38,7 @@ impl Client { "[msg:{}] Failed to decode newsletter plaintext: {e}", info.id ); + self.ack_newsletter_drop_after_pre_ack_hook(info); } } } else { @@ -45,6 +47,7 @@ impl Client { info.id, info.source.chat.observe() ); + self.ack_newsletter_drop_after_pre_ack_hook(info); } } diff --git a/src/message/tests.rs b/src/message/tests.rs index 2c190fa58..3b6a1531e 100644 --- a/src/message/tests.rs +++ b/src/message/tests.rs @@ -5339,6 +5339,301 @@ async fn capturing_client( (client, transport) } +#[derive(Clone)] +struct PreAckOrderRecorder { + order: Arc>>, +} + +impl wacore::types::events::EventHandler for PreAckOrderRecorder { + fn handle_event(&self, event: Arc) { + if matches!(event.as_ref(), Event::Message(_, _)) { + self.order + .lock() + .expect("order recorder mutex") + .push("event"); + } + } +} + +#[tokio::test] +async fn parsed_message_pre_ack_hook_runs_before_ack_and_event() { + let (client, transport) = capturing_client("pre_ack_hook_success").await; + let order = Arc::new(std::sync::Mutex::new(Vec::new())); + client.register_handler(Arc::new(PreAckOrderRecorder { + order: Arc::clone(&order), + })); + + let hook_order = Arc::clone(&order); + let hook_transport = Arc::clone(&transport); + client + .set_parsed_message_pre_ack_hook(move |ctx| { + let hook_order = Arc::clone(&hook_order); + let hook_transport = Arc::clone(&hook_transport); + async move { + assert_eq!(ctx.info.id, "PRE_ACK_SUCCESS"); + assert_eq!(ctx.message.conversation.as_deref(), Some("commit me")); + assert!( + hook_transport.sent().is_empty(), + "ACK/receipt must not be sent before the pre-ACK hook completes" + ); + hook_order + .lock() + .expect("order recorder mutex") + .push("hook"); + Ok::<(), anyhow::Error>(()) + } + }) + .expect("hook should register once"); + + let info = Arc::new(create_test_message_info( + "5511999998888@s.whatsapp.net", + "PRE_ACK_SUCCESS", + "5511777776666@s.whatsapp.net", + )); + let msg = wa::Message { + conversation: Some("commit me".to_string()), + ..Default::default() + }; + + assert!(client.dispatch_parsed_message(msg, &info).await); + + assert_eq!( + order.lock().expect("order recorder mutex").as_slice(), + ["hook", "event"], + "message event must be dispatched only after the awaited pre-ACK hook succeeds" + ); + + for _ in 0..40 { + if !transport.sent().is_empty() { + break; + } + client + .runtime + .sleep(std::time::Duration::from_millis(10)) + .await; + } + assert!( + !transport.sent().is_empty(), + "successful hook should allow the normal ACK/receipt to be sent" + ); +} + +#[tokio::test] +async fn parsed_message_pre_ack_hook_failure_suppresses_ack_and_event() { + let (client, transport) = capturing_client("pre_ack_hook_failure").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + client + .set_parsed_message_pre_ack_hook(|ctx| async move { + assert_eq!(ctx.info.id, "PRE_ACK_FAILURE"); + Err::<(), anyhow::Error>(anyhow::anyhow!("durable commit failed")) + }) + .expect("hook should register once"); + + let info = Arc::new(create_test_message_info( + "5511999998888@s.whatsapp.net", + "PRE_ACK_FAILURE", + "5511777776666@s.whatsapp.net", + )); + let msg = wa::Message { + conversation: Some("retry me".to_string()), + ..Default::default() + }; + + assert!(!client.dispatch_parsed_message(msg, &info).await); + + client + .runtime + .sleep(std::time::Duration::from_millis(50)) + .await; + assert!( + transport.sent().is_empty(), + "failing hook must suppress ACK/receipt so the server can redeliver" + ); + assert!( + collector.events().is_empty(), + "failing hook must not emit Event::Message for an uncommitted delivery attempt" + ); +} + +#[tokio::test] +async fn pending_pre_ack_message_retries_hook_on_duplicate_redelivery() { + let (client, transport) = capturing_client("pre_ack_hook_retry_pending").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let attempts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let hook_attempts = Arc::clone(&attempts); + client + .set_parsed_message_pre_ack_hook(move |_ctx| { + let hook_attempts = Arc::clone(&hook_attempts); + async move { + let attempt = hook_attempts.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if attempt == 0 { + Err::<(), anyhow::Error>(anyhow::anyhow!("first commit attempt failed")) + } else { + Ok::<(), anyhow::Error>(()) + } + } + }) + .expect("hook should register once"); + + let info = Arc::new(create_test_message_info( + "5511999998888@s.whatsapp.net", + "PRE_ACK_RETRY", + "5511777776666@s.whatsapp.net", + )); + let msg = wa::Message { + conversation: Some("retry pending".to_string()), + ..Default::default() + }; + + assert!(!client.dispatch_parsed_message(msg, &info).await); + assert!( + transport.sent().is_empty(), + "first hook failure must leave the message unacknowledged" + ); + + assert_eq!( + client.dispatch_pending_pre_ack_message(&info).await, + Some(true), + "duplicate redelivery should retry the stored parsed message hook" + ); + + let got = collect_event( + &client, + collector, + |e| matches!(e, Event::Message(_, info) if info.id == "PRE_ACK_RETRY"), + 500, + ) + .await; + assert!( + got.is_some(), + "successful retry must dispatch the message event" + ); + + for _ in 0..40 { + if !transport.sent().is_empty() { + break; + } + client + .runtime + .sleep(std::time::Duration::from_millis(10)) + .await; + } + assert!( + !transport.sent().is_empty(), + "successful retry must allow the ACK/receipt" + ); + assert_eq!( + client.dispatch_pending_pre_ack_message(&info).await, + None, + "successful pending retry must remove the stored parsed message" + ); +} + +#[tokio::test] +async fn parsed_message_pre_ack_hook_disables_status_deferred_ack_gate() { + let client = crate::test_utils::create_test_client_with_name("pre_ack_status_gate").await; + let status_node = NodeBuilder::new("message") + .attr("from", "status@broadcast") + .attr("id", "PRE_ACK_STATUS") + .build(); + + assert!( + client.should_ack(&status_node.as_node_ref()), + "without the pre-ACK hook, status messages keep the existing deferred ACK fallback" + ); + + client + .set_parsed_message_pre_ack_hook(|_ctx| async { Ok::<(), anyhow::Error>(()) }) + .expect("hook should register once"); + + assert!( + !client.should_ack(&status_node.as_node_ref()), + "with the pre-ACK hook, status ACK must wait for parsed-message hook success" + ); +} + +#[tokio::test] +async fn successful_normal_redelivery_clears_stale_pre_ack_pending_entry() { + let (client, _transport) = capturing_client("pre_ack_normal_redelivery").await; + + let attempts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let hook_attempts = Arc::clone(&attempts); + client + .set_parsed_message_pre_ack_hook(move |_ctx| { + let hook_attempts = Arc::clone(&hook_attempts); + async move { + let attempt = hook_attempts.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if attempt == 0 { + Err::<(), anyhow::Error>(anyhow::anyhow!("first commit attempt failed")) + } else { + Ok::<(), anyhow::Error>(()) + } + } + }) + .expect("hook should register once"); + + let info = Arc::new(create_test_message_info( + "5511999998888@s.whatsapp.net", + "PRE_ACK_NORMAL_REDELIVERY", + "5511777776666@s.whatsapp.net", + )); + let msg = wa::Message { + conversation: Some("first attempt".to_string()), + ..Default::default() + }; + assert!(!client.dispatch_parsed_message(msg, &info).await); + + let redelivery_msg = wa::Message { + conversation: Some("normal redelivery".to_string()), + ..Default::default() + }; + assert!(client.dispatch_parsed_message(redelivery_msg, &info).await); + assert_eq!( + client.dispatch_pending_pre_ack_message(&info).await, + None, + "normal successful redelivery must clear the old pending parsed message" + ); +} + +#[tokio::test] +async fn own_status_success_keeps_transport_ack_with_pre_ack_hook() { + let (client, transport) = capturing_client("pre_ack_own_status_ack").await; + client + .set_parsed_message_pre_ack_hook(|_ctx| async { Ok::<(), anyhow::Error>(()) }) + .expect("hook should register once"); + + let mut info = create_test_message_info( + "status@broadcast", + "PRE_ACK_OWN_STATUS", + "5511777776666@s.whatsapp.net", + ); + info.source.is_from_me = true; + let info = Arc::new(info); + let msg = wa::Message { + conversation: Some("own status".to_string()), + ..Default::default() + }; + + assert!(client.dispatch_parsed_message(msg, &info).await); + for _ in 0..40 { + if !transport.sent().is_empty() { + break; + } + client + .runtime + .sleep(std::time::Duration::from_millis(10)) + .await; + } + assert!( + !transport.sent().is_empty(), + "own status success still needs a transport ACK when the deferred gate is suppressed" + ); +} + /// Regression: a malformed pkmsg used to fall through silently. Now /// it dispatches the consumer event AND emits a nack on the wire so /// the server stops retransmitting.