Skip to content

refactor: targeted single-device DM retry (matches WA Web) #549

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jlucaso1 merged 9 commits into from
Apr 15, 2026
Merged

refactor: targeted single-device DM retry (matches WA Web) #549

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
d3fb42b
Initial plan
Copilot Apr 15, 2026
a0abb3b
fix: clear DM retry sessions for all known devices
Copilot Apr 15, 2026
748ef5d
test: add DM retry e2e regression coverage
Copilot Apr 15, 2026
75b4d51
chore: polish retry regression helpers
Copilot Apr 15, 2026
820d39e
perf: avoid fallback allocation in DM retry cleanup
Copilot Apr 15, 2026
cb928e0
refactor: target DM retries to a single device
Copilot Apr 15, 2026
bd12256
test: verify targeted DM retry resend
Copilot Apr 15, 2026
74b4956
chore: polish targeted retry follow-up
Copilot Apr 15, 2026
c9a7585
fix: dedupe and lock targeted retry paths
Copilot Apr 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
269 changes: 181 additions & 88 deletions src/retry.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ use wacore::libsignal::store::PreKeyStore;
use wacore::protocol::ProtocolNode;
use wacore::types::jid::JidExt;
use wacore_binary::JidExt as _;
use wacore_binary::OwnedNodeRef;
use wacore_binary::builder::NodeBuilder;
use wacore_binary::{Jid, OwnedNodeRef};
#[cfg(test)]
use wacore_binary::{Node, NodeContent};
use wacore_binary::{NodeContentRef, NodeRef};
Expand Down Expand Up @@ -387,86 +387,8 @@ impl Client {
);
}
} else {
// For DMs, handle base key tracking for collision detection (matches WhatsApp Web).
// This detects when we haven't regenerated our session despite receiving retry receipts,
// which can cause infinite retry loops where both sides are stuck with stale keys.
let signal_address = resolved_jid.to_protocol_address();
let device_store = self.persistence_manager.get_device_arc().await;

// Check for base key collision before deleting the session.
// Read session through cache for consistent state.
{
let device_guard = device_store.read().await;
let session = self
.signal_cache
.peek_session(&signal_address, &*device_guard.backend)
.await
.ok()
.flatten();

if let Some(session) = session
&& let Ok(current_base_key) = session.alice_base_key()
{
let addr_str = signal_address.as_str();
if retry_count == MIN_RETRY_FOR_BASE_KEY_CHECK {
// On retry 2: Save the base key for later comparison
if let Err(e) = device_guard
.backend
.save_base_key(addr_str, &message_id, current_base_key)
.await
{
warn!("Failed to save base key for {}: {}", signal_address, e);
} else {
info!(
"Saved base key for {} at retry #{} for collision detection",
signal_address, retry_count
);
}
} else if retry_count > MIN_RETRY_FOR_BASE_KEY_CHECK {
// On retry > 2: Check if base key is the same (collision detection)
match device_guard
.backend
.has_same_base_key(addr_str, &message_id, current_base_key)
.await
{
Ok(true) => {
// Collision detected! We haven't regenerated our session.
warn!(
"Base key collision detected for {} at retry #{}. \
Session hasn't been regenerated. Forcing fresh session.",
signal_address, retry_count
);
// Clean up base key entry since we're deleting the session
let _ = device_guard
.backend
.delete_base_key(addr_str, &message_id)
.await;
}
Ok(false) => {
// Base key changed, session was regenerated - good!
info!(
"Base key changed for {} at retry #{} - session regenerated",
signal_address, retry_count
);
// Clean up old base key entry
let _ = device_guard
.backend
.delete_base_key(addr_str, &message_id)
.await;
}
Err(e) => {
warn!("Failed to check base key for {}: {}", signal_address, e);
}
}
}
}
}

// Delete the old session through the signal cache so encryption uses a fresh session.
// IMPORTANT: Must go through cache, not backend, to avoid stale cached sessions.
self.signal_cache.delete_session(&signal_address).await;
self.flush_signal_cache().await?;
info!("Deleted session for {signal_address} due to retry receipt");
self.delete_dm_retry_session_target(&resolved_jid, &message_id, retry_count)
.await?;
}

// Status broadcasts can't resend (requires explicit recipient list).
Expand Down Expand Up @@ -514,19 +436,115 @@ impl Client {
self.send_node(stanza).await?;
self.flush_signal_cache().await?;
} else {
// DM retry: re-encrypt via normal send path (already targets single recipient)
self.send_message_impl(
// DM retry: pairwise resend to the requesting device only.
self.ensure_e2e_sessions(std::slice::from_ref(&resolved_jid))
Comment on lines +438 to +439

@chatgpt-codex-connector chatgpt-codex-connector Bot Apr 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Include device in DM retry dedupe key

This branch now resends a DM retry only to the requesting device, but DM deduplication is still keyed only by (chat, msg_id) (src/retry.rs:131-145). In multi-device chats, once one device’s retry is handled, subsequent retry receipts from the user’s other devices are treated as duplicates and skipped, so those devices never get their targeted resend and remain undecryptable. Deduping needs to include the sender device when DM retries are device-targeted.

Useful? React with 👍 / 👎.

.await?;

let device_snapshot = self.persistence_manager.get_device_snapshot().await;
let mut store_adapter = self.signal_adapter().await;

let stanza = wacore::send::prepare_dm_retry_stanza(
&mut store_adapter.session_store,
&mut store_adapter.identity_store,
receipt.source.chat.clone(),
Comment on lines +448 to 451

@chatgpt-codex-connector chatgpt-codex-connector Bot Apr 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Acquire session lock before DM retry encryption

The new DM retry flow encrypts directly with prepare_dm_retry_stanza without taking the per-device session lock, unlike the normal DM send path (src/send.rs:1137-1142). If this retry runs concurrently with another send/decrypt for the same Signal address, both paths can mutate the same ratchet state at once and cause additional undecryptable messages or retry loops. Wrap retry encryption with the same session-locking discipline used by regular DM sends.

Useful? React with 👍 / 👎.

participant_jid,
resolved_jid.clone(),
&original_msg,
Some(message_id),
false,
true,
None,
vec![],
message_id,
retry_count,
device_snapshot.account.as_ref(),
)
.await?;

self.send_node(stanza).await?;
self.flush_signal_cache().await?;
}

Ok(())
}

async fn delete_dm_retry_session_target(
&self,
device_jid: &Jid,
message_id: &str,
retry_count: u8,
) -> Result<(), anyhow::Error> {
// Base key collision detection prevents stale-session retry loops.
let signal_address = device_jid.to_protocol_address();
let device_store = self.persistence_manager.get_device_arc().await;

// Check for base key collision before deleting the session.
// Read session through cache for consistent state.
{
let device_guard = device_store.read().await;
let session = self
.signal_cache
.peek_session(&signal_address, &*device_guard.backend)
.await
.ok()
.flatten();

if let Some(session) = session
&& let Ok(current_base_key) = session.alice_base_key()
{
let addr_str = signal_address.as_str();
if retry_count == MIN_RETRY_FOR_BASE_KEY_CHECK {
// Save retry #2's base key so later retries can prove regeneration happened.
if let Err(e) = device_guard
.backend
.save_base_key(addr_str, message_id, current_base_key)
.await
{
warn!("Failed to save base key for {}: {}", signal_address, e);
} else {
info!(
"Saved base key for {} at retry #{} for collision detection",
signal_address, retry_count
);
}
} else if retry_count > MIN_RETRY_FOR_BASE_KEY_CHECK {
// An unchanged base key means we never rebuilt the session.
match device_guard
.backend
.has_same_base_key(addr_str, message_id, current_base_key)
.await
{
Ok(true) => {
// Collision detected! We haven't regenerated our session.
warn!(
"Base key collision detected for {} at retry #{}. \
Session hasn't been regenerated. Forcing fresh session.",
signal_address, retry_count
);
// Clean up base key entry since we're deleting the session
let _ = device_guard
.backend
.delete_base_key(addr_str, message_id)
.await;
}
Ok(false) => {
// Base key changed, session was regenerated - good!
info!(
"Base key changed for {} at retry #{} - session regenerated",
signal_address, retry_count
);
// Clean up old base key entry
let _ = device_guard
.backend
.delete_base_key(addr_str, message_id)
.await;
}
Err(e) => {
warn!("Failed to check base key for {}: {}", signal_address, e);
}
}
}
}
}

// Delete through the cache so resend can't revive a stale in-memory session.
self.signal_cache.delete_session(&signal_address).await;
info!("Deleted session for {signal_address} due to retry receipt");
Ok(())
}
Comment on lines +468 to 553

@coderabbitai coderabbitai Bot Apr 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Missing session lock before session deletion.

The delete_dm_retry_session_target method reads and deletes the session without acquiring the per-sender session lock. This could race with concurrent decrypt operations that also read/modify session state.

The existing code in process_retry_key_bundle (lines 672-673) and the registration ID mismatch path (lines 265-268) both acquire the session lock before session deletion:

let lock = self.session_lock_for(signal_address.as_str()).await;
let _guard = lock.lock().await;
self.signal_cache.delete_session(&signal_address).await;

Consider adding the same locking pattern here:

Proposed fix 
     async fn delete_dm_retry_session_target(
         &self,
         device_jid: &Jid,
         message_id: &str,
         retry_count: u8,
     ) -> Result<(), anyhow::Error> {
         // Base key collision detection prevents stale-session retry loops.
         let signal_address = device_jid.to_protocol_address();
+        let session_mutex = self.session_lock_for(signal_address.as_str()).await;
+        let _session_guard = session_mutex.lock().await;
+
         let device_store = self.persistence_manager.get_device_arc().await;

As per coding guidelines: "Use session_locks to serialize per-sender Signal encrypt/decrypt operations".

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@src/retry.rs` around lines 466 - 549, The delete_dm_retry_session_target
function deletes a session without acquiring the per-sender session lock; before
calling self.signal_cache.delete_session(&signal_address).await (and before the
subsequent info!("Deleted session...")), acquire the per-sender lock via let
lock = self.session_lock_for(signal_address.as_str()).await; then await
lock.lock() and hold the guard while calling delete_session so session deletion
is serialized with concurrent encrypt/decrypt operations (mirrors the pattern
used in process_retry_key_bundle and the registration ID mismatch path).


Expand Down Expand Up @@ -913,6 +931,8 @@ mod tests {
use crate::store::persistence_manager::PersistenceManager;
use crate::test_utils::MockHttpClient;
use std::borrow::Cow;
use std::sync::Arc;
use wacore::types::jid::JidExt as _;
use wacore_binary::{Jid, JidExt};
use waproto::whatsapp as wa;

Expand Down Expand Up @@ -1427,6 +1447,79 @@ mod tests {
);
}

#[tokio::test]
async fn dm_retry_deletes_only_requested_session() {
let client =
crate::test_utils::create_test_client_with_failing_http("retry_dm_devices").await;
let user = "100000000000088".to_string();
let resolved_jid = Jid::lid_device(user.clone(), 33);

let backend = client.persistence_manager.backend();
let device_0 = Jid::lid_device(user.clone(), 0).to_protocol_address();
let device_33 = Jid::lid_device(user, 33).to_protocol_address();

backend
.put_session(device_0.as_str(), b"invalid-session")
.await
.unwrap();
backend
.put_session(device_33.as_str(), b"invalid-session")
.await
.unwrap();

client
.delete_dm_retry_session_target(&resolved_jid, "MSG-ONE-DEVICE", 1)
.await
.unwrap();
client.flush_signal_cache().await.unwrap();

assert!(
backend
.get_session(device_0.as_str())
.await
.unwrap()
.is_some(),
"non-requesting device session should be preserved"
);
assert!(
backend
.get_session(device_33.as_str())
.await
.unwrap()
.is_none(),
"requesting device session should be deleted"
);
}

#[tokio::test]
async fn dm_retry_deletes_resolved_session_without_registry_devices() {
let client =
crate::test_utils::create_test_client_with_failing_http("retry_dm_fallback").await;
let resolved_jid = Jid::lid("100000000000099");
let signal_address = resolved_jid.to_protocol_address();
let backend = client.persistence_manager.backend();

backend
.put_session(signal_address.as_str(), b"invalid-session")
.await
.unwrap();

client
.delete_dm_retry_session_target(&resolved_jid, "MSG-FALLBACK", 1)
.await
.unwrap();
client.flush_signal_cache().await.unwrap();

assert!(
backend
.get_session(signal_address.as_str())
.await
.unwrap()
.is_none(),
"resolved session should be deleted when registry has no device list"
);
}

#[test]
fn bot_jid_detection() {
// Test bot JID detection for bot message filtering
Expand Down
Loading
Loading