Skip to content
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
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
26 changes: 26 additions & 0 deletions src/client/lid_pn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,32 @@ impl Client {
}
}

/// Swap a JID's namespace between PN and LID, preserving device/agent/integrator.
/// Returns `None` if no mapping exists or the JID is neither PN nor LID.
pub(crate) async fn swap_pn_lid_namespace(&self, jid: &Jid) -> Option<Jid> {
if jid.is_lid() {
let pn_user = self.lid_pn_cache.get_phone_number(&jid.user).await?;
Some(Jid {
user: pn_user.into(),
server: wacore_binary::Server::Pn,
device: jid.device,
agent: jid.agent,
integrator: jid.integrator,
})
} else if jid.is_pn() {
let lid_user = self.lid_pn_cache.get_current_lid(&jid.user).await?;
Some(Jid {
user: lid_user.into(),
server: wacore_binary::Server::Lid,
device: jid.device,
agent: jid.agent,
integrator: jid.integrator,
})
} else {
None
}
}

/// Migrate Signal sessions and identity keys from PN to LID address.
///
/// All reads/writes go through `signal_cache` to avoid reading stale data
Expand Down
79 changes: 62 additions & 17 deletions src/client/sender_keys.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
//! Sender key tracking and message cache methods for Client.

use anyhow::Result;
use wacore::types::message::ChatMessageId;
use wacore_binary::Jid;
use waproto::whatsapp as wa;

Expand Down Expand Up @@ -63,38 +64,77 @@ impl Client {
}

/// Take a sent message for retry handling. Checks L1 cache first (if enabled),
/// then falls back to DB. Matches WA Web's getMessageTable().get() pattern.
pub(crate) async fn take_recent_message(&self, to: &Jid, id: &str) -> Option<wa::Message> {
/// then falls back to DB. On miss, tries an alternate PN/LID key to handle
/// mapping changes between send time and retry time (WAWebLidMigrationUtils
/// `getAlternateMsgKey`).
/// Returns `(message, alternate_chat)`. When the message was found via the
/// alternate PN/LID key, `alternate_chat` contains the namespace that
/// matched -- the caller should use it for session operations instead of
/// `resolve_encryption_jid` (which would map back to the primary).
pub(crate) async fn take_recent_message(
&self,
to: &Jid,
id: &str,
) -> Option<(wa::Message, Option<Jid>)> {
let primary_key = self.make_chat_message_id(to, id).await;
if let Some(msg) = self.try_take_by_key(&primary_key).await {
return Some((msg, None));
}

// Primary miss -- try alternate PN<->LID key.
// If resolve_encryption_jid changed the namespace (PN→LID), the
// original `to` is already the alternate -- skip the cache lookup.
// Otherwise (LID input), swap via cache to try the PN form.
let alt_chat = if primary_key.chat.server != to.server {
Some(to.clone())
} else {
self.swap_pn_lid_namespace(&primary_key.chat).await
};

if let Some(alt_chat) = alt_chat {
log::debug!(
"Primary key miss for {}:{}, trying alternate {}",
primary_key.chat,
id,
alt_chat
);
let alt_key = ChatMessageId {
chat: alt_chat,
id: primary_key.id,
};
if let Some(msg) = self.try_take_by_key(&alt_key).await {
return Some((msg, Some(alt_key.chat)));
}
}

None
}

/// Look up and consume a message by exact `ChatMessageId` (L1 cache then DB).
async fn try_take_by_key(&self, key: &ChatMessageId) -> Option<wa::Message> {
use prost::Message;
let key = self.make_chat_message_id(to, id).await;
let chat_str = key.chat.to_string();
let has_l1_cache = self.cache_config.recent_messages.capacity > 0;

// L1 cache check (if capacity > 0)
if has_l1_cache && let Some(bytes) = self.recent_messages.remove(&key).await {
if has_l1_cache && let Some(bytes) = self.recent_messages.remove(key).await {
if let Ok(msg) = wa::Message::decode(bytes.as_slice()) {
// Cache hit — consume the DB row in the background to avoid orphans.
// Note: if the background DB write from add_recent_message hasn't completed
// yet, this delete may run first and the write creates an orphan. This is
// harmless — periodic cleanup (sent_message_ttl_secs) purges it. The race
// window is negligible since retry receipts arrive seconds after send.
let backend = self.persistence_manager.backend();
let cs = chat_str.clone();
let mid = key.id.clone();
self.runtime
.spawn(Box::pin(async move {
if let Err(e) = backend.take_sent_message(&cs, &mid).await {
log::warn!("Failed to clean up sent message {cs}:{mid}: {e}");
if let Err(e) = backend.take_sent_message(&chat_str, &mid).await {
log::warn!("Failed to clean up sent message {chat_str}:{mid}: {e}");
}
}))
.detach();
return Some(msg);
}
// Cache decode failed — fall through to DB
log::warn!(
"Failed to decode cached message for {}:{}, trying DB",
to,
id
key.chat,
key.id
);
}

Expand All @@ -108,16 +148,21 @@ impl Client {
Ok(Some(bytes)) => match wa::Message::decode(bytes.as_slice()) {
Ok(msg) => Some(msg),
Err(e) => {
log::warn!("Failed to decode DB message for {}:{}: {}", to, id, e);
log::warn!(
"Failed to decode DB message for {}:{}: {}",
key.chat,
key.id,
e
);
None
}
},
Ok(None) => None,
Err(e) => {
log::warn!(
"Failed to read sent message from DB for {}:{}: {}",
to,
id,
key.chat,
key.id,
e
);
None
Expand Down
Loading
Loading