Skip to content
Merged
Show file tree
Hide file tree
Changes from all 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