Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 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
2 changes: 1 addition & 1 deletion src/client/context_impl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ impl SendContextResolver for Client {
self.groups().query_info(jid).await
}

async fn get_lid_for_phone(&self, phone_user: &str) -> Option<String> {
async fn get_lid_for_phone(&self, phone_user: &str) -> Option<wacore_binary::CompactString> {
self.lid_pn_cache.get_current_lid(phone_user).await
}
}
2 changes: 1 addition & 1 deletion src/client/device_registry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ impl Client {
// Check if user is a PN (has a LID mapping)
if let Some(lid) = self.lid_pn_cache.get_current_lid(user).await {
return UserLookupKeys::PnWithLid {
lid: lid.into(),
lid,
pn: user.into(),
};
}
Expand Down
4 changes: 2 additions & 2 deletions src/client/lid_pn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -350,7 +350,7 @@ impl Client {
};
match self.lid_pn_cache.get_current_lid(&target.user).await {
Some(lid_user) => Jid {
user: lid_user.into(),
user: lid_user,
server: lid_server,
device: target.device,
agent: target.agent,
Expand All @@ -375,7 +375,7 @@ impl Client {
} else if jid.is_pn() {
let lid_user = self.lid_pn_cache.get_current_lid(&jid.user).await?;
Some(Jid {
user: lid_user.into(),
user: lid_user,
server: wacore_binary::Server::Lid,
device: jid.device,
agent: jid.agent,
Expand Down
10 changes: 7 additions & 3 deletions src/features/groups.rs
Original file line number Diff line number Diff line change
Expand Up @@ -876,15 +876,19 @@ impl<'a> Groups<'a> {

/// Resolve JID to tc_token store key. When `only_lid`, PN JIDs without a
/// LID mapping return `None` instead of falling back to the PN user.
async fn resolve_token_key(&self, jid: &Jid, only_lid: bool) -> Option<String> {
async fn resolve_token_key(
&self,
jid: &Jid,
only_lid: bool,
) -> Option<wacore_binary::CompactString> {
if jid.is_lid() {
Some(jid.user.to_string())
Some(jid.user.clone())
} else {
let lid = self.client.lid_pn_cache.get_current_lid(&jid.user).await;
if only_lid {
lid
} else {
Some(lid.unwrap_or_else(|| jid.user.to_string()))
Some(lid.unwrap_or_else(|| jid.user.clone()))
}
}
}
Expand Down
4 changes: 2 additions & 2 deletions src/handlers/notification.rs
Original file line number Diff line number Diff line change
Expand Up @@ -653,8 +653,8 @@ async fn handle_privacy_token_notification(client: &Arc<Client>, node: &NodeRef<
.filter(|j| !j.user.is_empty());

// Resolve to a LID key. We borrow from Jid.user (CompactString) or from
// get_current_lid (String), then pass as &str to the storage layer.
let resolved_lid: Option<String>;
// get_current_lid (CompactString), then pass as &str to the storage layer.
let resolved_lid: Option<wacore_binary::CompactString>;
let sender_lid: &str = if let Some(ref lid_jid) = sender_lid_jid {
&lid_jid.user
} else {
Expand Down
33 changes: 20 additions & 13 deletions src/lid_pn_cache.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@

use std::sync::Arc;

use wacore_binary::CompactString;

use crate::cache_config::{CacheConfig, CacheEntryConfig};
use crate::cache_store::TypedCache;
pub use wacore::types::{LearningSource, LidPnEntry};
Expand Down Expand Up @@ -101,8 +103,13 @@ impl LidPnCache {
/// Get the current LID for a phone number.
///
/// Returns the LID user part if a mapping exists, None otherwise.
pub async fn get_current_lid(&self, phone: &str) -> Option<String> {
self.pn_to_entry.get(phone).await.map(|e| e.lid.clone())
/// The cache holds `Arc<LidPnEntry>`; return the LID user as an (inline for
/// typical ~15-digit LIDs) `CompactString` instead of deep-cloning a `String`.
pub async fn get_current_lid(&self, phone: &str) -> Option<CompactString> {
self.pn_to_entry
.get(phone)
.await
.map(|e| e.lid.as_str().into())
}

/// Whether the learn fast path can skip re-recording `phone <-> lid`: the
Expand Down Expand Up @@ -257,8 +264,8 @@ mod tests {

// Should be retrievable both ways
assert_eq!(
cache.get_current_lid("559980000001").await,
Some("100000012345678".to_string())
cache.get_current_lid("559980000001").await.as_deref(),
Some("100000012345678")
);
assert_eq!(
cache.get_phone_number("100000012345678").await,
Expand All @@ -280,8 +287,8 @@ mod tests {
cache.add(&old_entry).await;

assert_eq!(
cache.get_current_lid("559980000001").await,
Some("100000012345678".to_string())
cache.get_current_lid("559980000001").await.as_deref(),
Some("100000012345678")
);

// Add newer mapping for same phone (different LID)
Expand All @@ -295,8 +302,8 @@ mod tests {

// Should return the newer LID for PN lookup
assert_eq!(
cache.get_current_lid("559980000001").await,
Some("100000087654321".to_string())
cache.get_current_lid("559980000001").await.as_deref(),
Some("100000087654321")
);

// Both LIDs should still be in the LID -> Entry map
Expand Down Expand Up @@ -334,8 +341,8 @@ mod tests {

// PN -> LID should still return the newer one
assert_eq!(
cache.get_current_lid("559980000001").await,
Some("100000087654321".to_string())
cache.get_current_lid("559980000001").await.as_deref(),
Some("100000087654321")
);
}

Expand Down Expand Up @@ -369,9 +376,9 @@ mod tests {
assert_eq!(cache.lid_count().await, 3);
assert_eq!(cache.pn_count().await, 3);

assert_eq!(cache.get_current_lid("pn1").await, Some("lid1".to_string()));
assert_eq!(cache.get_current_lid("pn2").await, Some("lid2".to_string()));
assert_eq!(cache.get_current_lid("pn3").await, Some("lid3".to_string()));
assert_eq!(cache.get_current_lid("pn1").await.as_deref(), Some("lid1"));
assert_eq!(cache.get_current_lid("pn2").await.as_deref(), Some("lid2"));
assert_eq!(cache.get_current_lid("pn3").await.as_deref(), Some("lid3"));
}

#[tokio::test]
Expand Down
26 changes: 17 additions & 9 deletions src/message.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3057,8 +3057,12 @@ mod tests {

// Cache learned the mapping in both directions.
assert_eq!(
client.lid_pn_cache.get_current_lid(pn_user).await,
Some(lid_user.to_string()),
client
.lid_pn_cache
.get_current_lid(pn_user)
.await
.as_deref(),
Some(lid_user),
"PN→LID lookup must hit"
);
assert_eq!(
Expand Down Expand Up @@ -3152,8 +3156,12 @@ mod tests {
// Hosted variant must reach the cache; without it, learn_lid_pn_mapping
// is skipped and the hosted-device fix is incomplete.
assert_eq!(
client.lid_pn_cache.get_current_lid(pn_user).await,
Some(lid_user.to_string()),
client
.lid_pn_cache
.get_current_lid(pn_user)
.await
.as_deref(),
Some(lid_user),
"PN→LID lookup must work for hosted family"
);
assert_eq!(
Expand Down Expand Up @@ -5901,8 +5909,8 @@ mod tests {
// Verify the cache has the mapping
let cached_lid = client.lid_pn_cache.get_current_lid(phone).await;
assert_eq!(
cached_lid,
Some(lid.to_string()),
cached_lid.as_deref(),
Some(lid),
"Cache should have the LID-PN mapping"
);

Expand Down Expand Up @@ -5963,7 +5971,7 @@ mod tests {
} else if let Some(lid_user) = client.lid_pn_cache.get_current_lid(&sender.user).await {
// Use the cached LID
Jid {
user: lid_user.into(),
user: lid_user,
server: wacore_binary::Server::Lid,
device: sender.device,
agent: sender.agent,
Expand Down Expand Up @@ -6078,7 +6086,7 @@ mod tests {
} else if let Some(lid_user) = client.lid_pn_cache.get_current_lid(&sender.user).await {
// This is the path we're testing - fallback to cached LID
Jid {
user: lid_user.into(),
user: lid_user,
server: wacore_binary::Server::Lid,
device: sender.device,
agent: sender.agent,
Expand Down Expand Up @@ -6179,7 +6187,7 @@ mod tests {
}
} else if let Some(lid_user) = client.lid_pn_cache.get_current_lid(&sender.user).await {
Jid {
user: lid_user.into(),
user: lid_user,
server: wacore_binary::Server::Lid,
device: sender.device,
agent: sender.agent,
Expand Down
2 changes: 1 addition & 1 deletion src/send.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1962,7 +1962,7 @@ impl Client {
}

if let Some(lid_user) = self.lid_pn_cache.get_current_lid(&jid.user).await {
Jid::new(&lid_user, Server::Lid)
Jid::new(lid_user, Server::Lid)
} else {
jid.to_non_ad()
}
Expand Down
2 changes: 1 addition & 1 deletion wacore/src/client/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,7 @@ pub trait SendContextResolver: crate::sync_marker::MaybeSendSync {
/// when sending to a phone number address.
///
/// Returns None if no LID mapping is known for this phone number.
async fn get_lid_for_phone(&self, phone_user: &str) -> Option<String> {
async fn get_lid_for_phone(&self, phone_user: &str) -> Option<wacore_binary::CompactString> {
// Default implementation returns None - subclasses can override
let _ = phone_user;
None
Expand Down
65 changes: 4 additions & 61 deletions wacore/src/history_sync.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,15 +66,15 @@ pub fn process_history_sync(

let buf = Bytes::from(decompressed);
let mut pos = 0;
// Pre-size the secret-record accumulator from a single allocation-free count
// pass, so it never grows by repeated doubling as conversations stream in.
let estimated_records = count_history_sync_messages(&buf);
let mut result = HistorySyncResult {
own_pushname: None,
nct_salt: None,
conversations_processed: 0,
tc_token_candidates: Vec::new(),
msg_secret_records: Vec::with_capacity(estimated_records),
// Grown on demand: a full pre-count pass scanned the whole blob just to
// size a Vec that only holds the secret-record subset (it over-allocated
// and cost ~2.5% of the decode); plain growth is cheaper here.
msg_secret_records: Vec::new(),
// Always retained on this path: the `!retain_blob` case returned above
// and ran the streaming variant, so control only reaches here when the
// caller wants the blob.
Expand Down Expand Up @@ -357,63 +357,6 @@ fn skip_field(wire_type: u32, buf: &[u8], pos: usize) -> Result<usize, HistorySy
}
}

/// Best-effort count of total `HistorySyncMsg` entries (field 2 inside each
/// field-2 conversation) in a decompressed HistorySync blob. Allocation-free;
/// used to pre-size the message-secret accumulator in one shot instead of letting
/// it grow by repeated doubling. An under-count (e.g. on a malformed tail) only
/// costs a few late re-grows, never correctness — the decode loop re-validates.
fn count_history_sync_messages(buf: &[u8]) -> usize {
let mut pos = 0;
let mut total = 0;
while pos < buf.len() {
let Ok((tag, br)) = read_varint(&buf[pos..]) else {
break;
};
pos += br;
let field = (tag >> 3) as u32;
let wt = (tag & 0x7) as u32;
if field == 2 && wt == wire_type::LENGTH_DELIMITED {
let Ok((len, vl)) = read_varint(&buf[pos..]) else {
break;
};
pos += vl;
let Ok(end) = checked_end(pos, len, buf.len(), "conv-count") else {
break;
};
total += count_conversation_messages(&buf[pos..end]);
pos = end;
} else {
match skip_field(wt, buf, pos) {
Ok(np) => pos = np,
Err(_) => break,
}
}
}
total
}

/// Count field-2 (message) entries within a single conversation's bytes.
fn count_conversation_messages(buf: &[u8]) -> usize {
let mut pos = 0;
let mut n = 0;
while pos < buf.len() {
let Ok((tag, br)) = read_varint(&buf[pos..]) else {
break;
};
pos += br;
let field = (tag >> 3) as u32;
let wt = (tag & 0x7) as u32;
if field == 2 && wt == wire_type::LENGTH_DELIMITED {
n += 1;
}
match skip_field(wt, buf, pos) {
Ok(np) => pos = np,
Err(_) => break,
}
}
n
}

/// Manual pushname parser — Pushname proto has fields: id (tag 1) and pushname (tag 2).
/// Checks id first and only allocates the pushname string if id matches `own_user`.
fn extract_own_pushname(data: &[u8], own_user: &str) -> Option<String> {
Expand Down
10 changes: 7 additions & 3 deletions wacore/src/iq/usync.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ pub enum UsyncContext {
pub struct IsOnWhatsAppUser {
pub jid: Jid,
/// Helps server optimize the lookup (WA Web pre-populates this from its LID cache).
pub known_lid: Option<String>,
pub known_lid: Option<wacore_binary::CompactString>,
}

fn build_user_nodes(users: &[IsOnWhatsAppUser]) -> Vec<Node> {
Expand All @@ -106,7 +106,11 @@ fn build_user_nodes(users: &[IsOnWhatsAppUser]) -> Vec<Node> {
};
let mut children = vec![NodeBuilder::new("contact").string_content(phone).build()];
if let Some(lid) = &user.known_lid {
children.push(NodeBuilder::new("lid").attr("jid", Jid::lid(lid)).build());
children.push(
NodeBuilder::new("lid")
.attr("jid", Jid::lid(lid.as_str()))
.build(),
);
}
NodeBuilder::new("user").children(children).build()
} else {
Expand Down Expand Up @@ -855,7 +859,7 @@ mod tests {
let spec = IsOnWhatsAppSpec::new(
vec![IsOnWhatsAppUser {
jid: Jid::pn("1234567890"),
known_lid: Some("100000001".to_string()),
known_lid: Some("100000001".into()),
}],
"sid",
IsOnWhatsAppQueryType::Pn,
Expand Down
Loading
Loading