whatsapp_rust/client/
lid_pn.rs

1//! LID-PN (Linked ID to Phone Number) mapping methods for Client.
2//!
3//! This module contains methods for managing the bidirectional mapping
4//! between LIDs (Linked IDs) and phone numbers.
5//!
6//! Key features:
7//! - Cache warm-up from persistent storage
8//! - Adding new LID-PN mappings with automatic migration
9//! - Resolving JIDs to their LID equivalents
10//! - Bidirectional lookup (LID to PN and PN to LID)
11
12use anyhow::Result;
13use log::debug;
14use wacore_binary::jid::Jid;
15
16use super::Client;
17use crate::lid_pn_cache::{LearningSource, LidPnEntry};
18
19impl Client {
20    /// Warm up the LID-PN cache from persistent storage.
21    /// This is called during client initialization to populate the in-memory cache
22    /// with previously learned LID-PN mappings.
23    pub(crate) async fn warm_up_lid_pn_cache(&self) -> Result<(), anyhow::Error> {
24        let backend = self.persistence_manager.backend();
25        let entries = backend.get_all_lid_mappings().await?;
26
27        if entries.is_empty() {
28            debug!("LID-PN cache warm-up: no entries found in storage");
29            return Ok(());
30        }
31
32        let cache_entries: Vec<LidPnEntry> = entries
33            .into_iter()
34            .map(|e| {
35                LidPnEntry::with_timestamp(
36                    e.lid,
37                    e.phone_number,
38                    e.created_at,
39                    LearningSource::parse(&e.learning_source),
40                )
41            })
42            .collect();
43
44        self.lid_pn_cache.warm_up(cache_entries).await;
45        Ok(())
46    }
47
48    /// Add a LID-PN mapping to both the in-memory cache and persistent storage.
49    /// This is called when we learn about a mapping from messages, usync, etc.
50    /// Also migrates any existing PN-keyed device registry entries to LID.
51    pub(crate) async fn add_lid_pn_mapping(
52        &self,
53        lid: &str,
54        phone_number: &str,
55        source: LearningSource,
56    ) -> Result<()> {
57        use anyhow::anyhow;
58        use wacore::store::traits::LidPnMappingEntry;
59
60        // Check if this is a new mapping (not just an update)
61        let is_new_mapping = self
62            .lid_pn_cache
63            .get_current_lid(phone_number)
64            .await
65            .is_none();
66
67        // Add to in-memory cache
68        let entry = LidPnEntry::new(lid.to_string(), phone_number.to_string(), source);
69        self.lid_pn_cache.add(entry.clone()).await;
70
71        // Persist to storage
72        let backend = self.persistence_manager.backend();
73        let storage_entry = LidPnMappingEntry {
74            lid: entry.lid,
75            phone_number: entry.phone_number,
76            created_at: entry.created_at,
77            updated_at: entry.created_at,
78            learning_source: entry.learning_source.as_str().to_string(),
79        };
80
81        backend
82            .put_lid_mapping(&storage_entry)
83            .await
84            .map_err(|e| anyhow!("persisting LID-PN mapping: {e}"))?;
85
86        // If this is a new LID mapping, migrate any existing PN-keyed device registry entries
87        if is_new_mapping {
88            self.migrate_device_registry_on_lid_discovery(phone_number, lid)
89                .await;
90        }
91
92        Ok(())
93    }
94
95    /// Ensure phone-to-LID mappings are resolved for the given JIDs.
96    /// Matches WhatsApp Web's WAWebManagePhoneNumberMappingJob.ensurePhoneNumberToLidMapping().
97    /// Should be called before establishing new E2E sessions to avoid duplicate sessions.
98    ///
99    /// This checks the local cache for existing mappings. For JIDs without cached mappings,
100    /// the caller should consider fetching them via usync query if establishing sessions.
101    pub(crate) async fn resolve_lid_mappings(&self, jids: &[Jid]) -> Vec<Jid> {
102        let mut resolved = Vec::with_capacity(jids.len());
103
104        for jid in jids {
105            // Only resolve for user JIDs (not groups, status, etc.)
106            if !jid.is_pn() && !jid.is_lid() {
107                resolved.push(jid.clone());
108                continue;
109            }
110
111            // If it's already a LID, use as-is
112            if jid.is_lid() {
113                resolved.push(jid.clone());
114                continue;
115            }
116
117            // Try to resolve PN to LID from cache
118            if let Some(lid_user) = self.lid_pn_cache.get_current_lid(&jid.user).await {
119                resolved.push(Jid::lid_device(lid_user, jid.device));
120            } else {
121                // No cached mapping — use original JID. Mapping will be learned
122                // organically from incoming messages or usync responses.
123                resolved.push(jid.clone());
124            }
125        }
126
127        resolved
128    }
129
130    /// Resolve the encryption JID for a given target JID.
131    /// This uses the same logic as the receiving path to ensure consistent
132    /// lock keys between sending and receiving.
133    ///
134    /// For PN JIDs, this checks if a LID mapping exists and returns the LID.
135    /// This ensures that sending and receiving use the same session lock.
136    pub(crate) async fn resolve_encryption_jid(&self, target: &Jid) -> Jid {
137        let pn_server = wacore_binary::jid::DEFAULT_USER_SERVER;
138        let lid_server = wacore_binary::jid::HIDDEN_USER_SERVER;
139
140        if target.server == lid_server {
141            // Already a LID - use it directly
142            target.clone()
143        } else if target.server == pn_server {
144            // PN JID - check if we have a LID mapping
145            if let Some(lid_user) = self.lid_pn_cache.get_current_lid(&target.user).await {
146                let lid_jid = Jid {
147                    user: lid_user,
148                    server: wacore_binary::jid::cow_server_from_str(lid_server),
149                    device: target.device,
150                    agent: target.agent,
151                    integrator: target.integrator,
152                };
153                debug!(
154                    "[SEND-LOCK] Resolved {} to LID {} for session lock",
155                    target, lid_jid
156                );
157                lid_jid
158            } else {
159                // No LID mapping - use PN as-is
160                debug!("[SEND-LOCK] No LID mapping for {}, using PN", target);
161                target.clone()
162            }
163        } else {
164            // Other server type - use as-is
165            target.clone()
166        }
167    }
168
169    /// Get the phone number (user part) for a given LID.
170    /// Looks up the LID-PN mapping from the in-memory cache.
171    ///
172    /// # Arguments
173    ///
174    /// * `lid` - The LID user part (e.g., "100000012345678") or full JID (e.g., "100000012345678@lid")
175    ///
176    /// # Returns
177    ///
178    /// The phone number user part if a mapping exists, None otherwise.
179    pub async fn get_phone_number_from_lid(&self, lid: &str) -> Option<String> {
180        // Handle both full JID (e.g., "100000012345678@lid") and user part only
181        let lid_user = if lid.contains('@') {
182            lid.split('@').next().unwrap_or(lid)
183        } else {
184            lid
185        };
186        self.lid_pn_cache.get_phone_number(lid_user).await
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193    use crate::lid_pn_cache::LearningSource;
194    use crate::test_utils::create_test_client;
195    use std::sync::Arc;
196    use wacore_binary::jid::HIDDEN_USER_SERVER;
197
198    #[tokio::test]
199    async fn test_resolve_encryption_jid_pn_to_lid() {
200        let client: Arc<Client> = create_test_client().await;
201        let pn = "55999999999";
202        let lid = "100000012345678";
203
204        // Add mapping to cache
205        client
206            .add_lid_pn_mapping(lid, pn, LearningSource::PeerPnMessage)
207            .await
208            .unwrap();
209
210        let pn_jid = Jid::pn(pn);
211        let resolved = client.resolve_encryption_jid(&pn_jid).await;
212
213        assert_eq!(resolved.user, lid);
214        assert_eq!(resolved.server, HIDDEN_USER_SERVER);
215    }
216
217    #[tokio::test]
218    async fn test_resolve_encryption_jid_preserves_lid() {
219        let client: Arc<Client> = create_test_client().await;
220        let lid = "100000012345678";
221        let lid_jid = Jid::lid(lid);
222
223        let resolved = client.resolve_encryption_jid(&lid_jid).await;
224
225        assert_eq!(resolved, lid_jid);
226    }
227
228    #[tokio::test]
229    async fn test_resolve_encryption_jid_no_mapping_returns_pn() {
230        let client: Arc<Client> = create_test_client().await;
231        let pn = "55999999999";
232        let pn_jid = Jid::pn(pn);
233
234        let resolved = client.resolve_encryption_jid(&pn_jid).await;
235
236        assert_eq!(resolved, pn_jid);
237    }
238}
whatsapp_rust/client/lid_pn.rs

whatsapp_rust/client/
lid_pn.rs
