Skip to main content

imessage_database/tables/
handle.rs

1/*!
2 Handle table rows and handle grouping helpers.
3*/
4
5use rusqlite::{CachedStatement, Connection, Result, Row};
6use std::collections::HashMap;
7
8use crate::{
9    error::table::TableError,
10    tables::{
11        diagnostic::{HandleDiagnostic, column_exists, count_query},
12        table::{HANDLE, ME, Table},
13    },
14    util::union_find::UnionFind,
15};
16
17// MARK: Handle
18/// Row from the `handle` table.
19#[derive(Debug)]
20pub struct Handle {
21    /// Handle row ID.
22    pub rowid: i32,
23    /// Phone number, email address, or service identifier.
24    pub id: String,
25    /// Legacy merge ID that the iMessage server used to identify handles for
26    /// one sender account.
27    pub person_centric_id: Option<String>,
28}
29
30// MARK: Table
31impl Table for Handle {
32    fn from_row(row: &Row) -> Result<Handle> {
33        Ok(Handle {
34            rowid: row.get("rowid")?,
35            id: row.get("id")?,
36            person_centric_id: row.get("person_centric_id").unwrap_or(None),
37        })
38    }
39}
40
41// MARK: Rows
42impl Handle {
43    /// Prepare the `SELECT *` statement for the `handle` table.
44    pub fn get(db: &'_ Connection) -> Result<CachedStatement<'_>, TableError> {
45        Ok(db.prepare_cached(&format!("SELECT * from {HANDLE}"))?)
46    }
47
48    /// Load every row of the `handle` table in query order.
49    ///
50    /// Load once and share the rows: borrow them with [`HandleGroups::from_rows`]
51    /// before consuming them with [`Handle::ids_by_rowid`].
52    ///
53    /// # Example:
54    ///
55    /// ```no_run
56    /// use imessage_database::util::dirs::default_db_path;
57    /// use imessage_database::tables::table::get_connection;
58    /// use imessage_database::tables::handle::{Handle, HandleGroups};
59    ///
60    /// let db_path = default_db_path().unwrap();
61    /// let conn = get_connection(&db_path).unwrap();
62    /// let handles = Handle::load(&conn).unwrap();
63    /// let groups = HandleGroups::from_rows(&handles);
64    /// let ids = Handle::ids_by_rowid(handles);
65    /// ```
66    pub fn load(db: &Connection) -> Result<Vec<Handle>, TableError> {
67        let mut statement = Handle::get(db)?;
68        Handle::rows(&mut statement, [])?.collect()
69    }
70
71    /// Map handle identifiers by row ID.
72    ///
73    /// In group chats, treat row ID `0` as the local user and map it to [`ME`].
74    /// Use [`HandleGroups`] to group rows that belong to one person.
75    ///
76    /// # Example:
77    ///
78    /// ```no_run
79    /// use imessage_database::util::dirs::default_db_path;
80    /// use imessage_database::tables::table::get_connection;
81    /// use imessage_database::tables::handle::Handle;
82    ///
83    /// let db_path = default_db_path().unwrap();
84    /// let conn = get_connection(&db_path).unwrap();
85    /// let handles = Handle::ids_by_rowid(Handle::load(&conn).unwrap());
86    /// ```
87    #[must_use]
88    pub fn ids_by_rowid(handles: Vec<Handle>) -> HashMap<i32, String> {
89        let mut map = HashMap::with_capacity(handles.len() + 1);
90        map.insert(0, ME.to_string());
91
92        for handle in handles {
93            map.insert(handle.rowid, handle.id);
94        }
95
96        map
97    }
98}
99
100// MARK: Groups
101/// Groups of `handle` rows that belong to one person.
102///
103/// We group rows when their `id` values match byte-for-byte (one address
104/// registered on several services) or their `person_centric_id` values match.
105/// Call [`HandleGroups::union`] to add relationships from outside the database,
106/// such as contact matches. Call [`HandleGroups::finalize`] to assign one
107/// sequential group ID to each class.
108///
109/// Compare stored `id` values without normalization. Messages lowercases email
110/// addresses and canonicalizes phone numbers before storing them. Do not match
111/// phone numbers on digits alone; you could combine a short code with a full
112/// number that shares its suffix.
113///
114/// The iMessage server attached `person_centric_id` to incoming messages as a
115/// merge ID, but stopped issuing it for new senders long ago. Current databases
116/// retain this field only for the oldest handles. Keep a present ID as a
117/// grouping edge; it remains reliable even though current databases rarely
118/// include it.
119///
120/// # Example:
121///
122/// ```no_run
123/// use imessage_database::util::dirs::default_db_path;
124/// use imessage_database::tables::table::get_connection;
125/// use imessage_database::tables::handle::HandleGroups;
126///
127/// let db_path = default_db_path().unwrap();
128/// let conn = get_connection(&db_path).unwrap();
129/// let mut groups = HandleGroups::from_db(&conn).unwrap();
130/// groups.union(3, 7);
131/// let group_by_handle = groups.finalize();
132/// ```
133#[derive(Debug)]
134pub struct HandleGroups {
135    sets: UnionFind,
136}
137
138impl HandleGroups {
139    /// Seed groups from every row of the `handle` table.
140    ///
141    /// Load the table through [`Handle::load`], then seed the groups with
142    /// [`HandleGroups::from_rows`].
143    pub fn from_db(db: &Connection) -> Result<Self, TableError> {
144        Ok(Self::from_rows(&Handle::load(db)?))
145    }
146
147    /// Seed groups from loaded `handle` rows.
148    ///
149    /// Include row ID `0` as its own group. Chat participant lists can use it
150    /// for the local user.
151    #[must_use]
152    pub fn from_rows(handles: &[Handle]) -> Self {
153        let mut sets = UnionFind::new();
154        sets.make_set(0);
155
156        let mut first_by_id: HashMap<&str, i32> = HashMap::new();
157        let mut first_by_person: HashMap<&str, i32> = HashMap::new();
158
159        for handle in handles {
160            sets.make_set(handle.rowid);
161
162            if !handle.id.is_empty() {
163                let first = *first_by_id
164                    .entry(handle.id.as_str())
165                    .or_insert(handle.rowid);
166                sets.union(handle.rowid, first);
167            }
168
169            if let Some(person) = handle
170                .person_centric_id
171                .as_deref()
172                .filter(|p| !p.is_empty())
173            {
174                let first = *first_by_person.entry(person).or_insert(handle.rowid);
175                sets.union(handle.rowid, first);
176            }
177        }
178
179        Self { sets }
180    }
181
182    /// Merge the groups containing two handle row IDs.
183    ///
184    /// Pass row IDs from the table. Unknown IDs have no effect, so an external
185    /// source cannot create a phantom group with a stale ID.
186    pub fn union(&mut self, a: i32, b: i32) {
187        if self.sets.contains(a) && self.sets.contains(b) {
188            self.sets.union(a, b);
189        }
190    }
191
192    /// Assign a sequential group ID to each class and return it keyed by handle row ID.
193    ///
194    /// Order classes by their smallest row ID before assigning sequential IDs.
195    /// This keeps results stable for an unchanged table, even when union-find
196    /// picks a different representative.
197    pub fn finalize(mut self) -> HashMap<i32, i32> {
198        let ids: Vec<i32> = self.sets.ids().collect();
199        self.sets.label(ids)
200    }
201}
202
203// MARK: Diagnostic
204impl Handle {
205    /// Compute diagnostic data for the `handle` table.
206    ///
207    /// Group rows with [`HandleGroups`], then report the extra rows absorbed
208    /// into groups and the number of groups with multiple rows. Count
209    /// `person_centric_id` values separately when the column exists.
210    ///
211    /// # Example:
212    ///
213    /// ```no_run
214    /// use imessage_database::util::dirs::default_db_path;
215    /// use imessage_database::tables::table::get_connection;
216    /// use imessage_database::tables::handle::Handle;
217    ///
218    /// let db_path = default_db_path().unwrap();
219    /// let conn = get_connection(&db_path).unwrap();
220    /// Handle::run_diagnostic(&conn);
221    /// ```
222    pub fn run_diagnostic(db: &Connection) -> Result<HandleDiagnostic, TableError> {
223        let query = concat!(
224            "SELECT COUNT(DISTINCT person_centric_id) ",
225            "FROM handle ",
226            "WHERE person_centric_id NOT NULL AND person_centric_id != ''"
227        );
228
229        let person_centric_ids = if column_exists(db, HANDLE, "person_centric_id")? {
230            Some(count_query(db, query)?)
231        } else {
232            None
233        };
234
235        // Row 0 is synthetic; it has no handle row, so exclude it from the
236        // total and from group sizes.
237        let group_by_handle = HandleGroups::from_db(db)?.finalize();
238        let total_handles = group_by_handle.len() - 1;
239
240        let mut rows_per_group: HashMap<i32, usize> = HashMap::new();
241        for (rowid, group) in group_by_handle {
242            if rowid != 0 {
243                *rows_per_group.entry(group).or_insert(0) += 1;
244            }
245        }
246
247        Ok(HandleDiagnostic {
248            total_handles,
249            person_centric_ids,
250            multi_handle_groups: rows_per_group.values().filter(|rows| **rows > 1).count(),
251            merged_handles: total_handles - rows_per_group.len(),
252        })
253    }
254}
255
256// MARK: Tests
257#[cfg(test)]
258mod tests {
259    use crate::tables::handle::{Handle, HandleGroups};
260    use rusqlite::Connection;
261    use std::collections::{HashMap, HashSet};
262
263    /// Build an in-memory `handle` table from `(rowid, id, service, person_centric_id)` rows.
264    fn db_with(rows: &[(i32, &str, &str, Option<&str>)]) -> Connection {
265        let db = Connection::open_in_memory().unwrap();
266        db.execute(
267            "CREATE TABLE handle (
268                ROWID INTEGER PRIMARY KEY,
269                id TEXT NOT NULL,
270                service TEXT NOT NULL,
271                person_centric_id TEXT
272            )",
273            [],
274        )
275        .unwrap();
276        for (rowid, id, service, person) in rows {
277            db.execute(
278                "INSERT INTO handle (ROWID, id, service, person_centric_id) VALUES (?1, ?2, ?3, ?4)",
279                (rowid, id, service, person),
280            )
281            .unwrap();
282        }
283        db
284    }
285
286    fn groups_of(map: &HashMap<i32, i32>, rowids: &[i32]) -> HashSet<i32> {
287        rowids.iter().map(|r| map[r]).collect()
288    }
289
290    #[test]
291    fn load_returns_every_row() {
292        let db = db_with(&[
293            (1, "+15551234567", "iMessage", Some("p1")),
294            (2, "a@example.com", "SMS", None),
295            (3, "b@example.com", "iMessage", Some("")),
296        ]);
297
298        let handles = Handle::load(&db).unwrap();
299
300        let rows: Vec<(i32, &str, Option<&str>)> = handles
301            .iter()
302            .map(|h| (h.rowid, h.id.as_str(), h.person_centric_id.as_deref()))
303            .collect();
304        assert_eq!(
305            rows,
306            vec![
307                (1, "+15551234567", Some("p1")),
308                (2, "a@example.com", None),
309                (3, "b@example.com", Some("")),
310            ]
311        );
312    }
313
314    #[test]
315    fn ids_by_rowid_includes_me_and_every_row() {
316        let db = db_with(&[
317            (1, "+15551234567", "iMessage", None),
318            (2, "a@example.com", "SMS", None),
319        ]);
320
321        let ids = Handle::ids_by_rowid(Handle::load(&db).unwrap());
322
323        assert_eq!(ids.len(), 3);
324        assert_eq!(ids[&0], "Me");
325        assert_eq!(ids[&1], "+15551234567");
326        assert_eq!(ids[&2], "a@example.com");
327    }
328
329    #[test]
330    fn identical_ids_across_services_share_a_group() {
331        let db = db_with(&[
332            (1, "+15551234567", "iMessage", None),
333            (2, "+15551234567", "SMS", None),
334            (3, "+15551234567", "RCS", None),
335            (4, "a@example.com", "iMessage", None),
336        ]);
337        let map = HandleGroups::from_db(&db).unwrap().finalize();
338
339        assert_eq!(groups_of(&map, &[1, 2, 3]).len(), 1);
340        assert_ne!(map[&1], map[&4]);
341        assert_eq!(map[&0], 0, "the local user is group 0");
342        assert_eq!(map.len(), 5);
343    }
344
345    #[test]
346    fn ids_are_not_normalized() {
347        let db = db_with(&[
348            (1, "+15551234567", "iMessage", None),
349            (2, "15551234567", "SMS", None),
350            (3, "A@example.com", "iMessage", None),
351            (4, "a@example.com", "iMessage", None),
352        ]);
353        let map = HandleGroups::from_db(&db).unwrap().finalize();
354
355        assert_ne!(map[&1], map[&2]);
356        assert_ne!(map[&3], map[&4]);
357    }
358
359    #[test]
360    fn person_centric_id_joins_different_ids() {
361        let db = db_with(&[
362            (1, "+15551234567", "iMessage", Some("p1")),
363            (2, "a@example.com", "iMessage", Some("p1")),
364            (3, "b@example.com", "iMessage", Some("p2")),
365            (4, "+15559876543", "iMessage", None),
366        ]);
367        let map = HandleGroups::from_db(&db).unwrap().finalize();
368
369        assert_eq!(map[&1], map[&2]);
370        assert_ne!(map[&1], map[&3]);
371        assert_ne!(map[&1], map[&4]);
372    }
373
374    #[test]
375    fn edges_are_transitive() {
376        // Group rows 1, 2, and 3: connect 1 to 2 by person_centric_id, then 2
377        // to 3 by id.
378        let db = db_with(&[
379            (1, "+15551234567", "iMessage", Some("p1")),
380            (2, "a@example.com", "iMessage", Some("p1")),
381            (3, "a@example.com", "SMS", None),
382        ]);
383        let map = HandleGroups::from_db(&db).unwrap().finalize();
384
385        assert_eq!(groups_of(&map, &[1, 2, 3]).len(), 1);
386    }
387
388    #[test]
389    fn empty_person_centric_id_is_not_an_edge() {
390        let db = db_with(&[
391            (1, "+15551234567", "iMessage", Some("")),
392            (2, "a@example.com", "iMessage", Some("")),
393        ]);
394        let map = HandleGroups::from_db(&db).unwrap().finalize();
395
396        assert_ne!(map[&1], map[&2]);
397    }
398
399    #[test]
400    fn external_union_merges_groups() {
401        let db = db_with(&[
402            (1, "+15551234567", "iMessage", None),
403            (2, "a@example.com", "iMessage", None),
404            (3, "b@example.com", "iMessage", None),
405        ]);
406        let mut groups = HandleGroups::from_db(&db).unwrap();
407        groups.union(1, 2);
408        let map = groups.finalize();
409
410        assert_eq!(map[&1], map[&2]);
411        assert_ne!(map[&1], map[&3]);
412    }
413
414    #[test]
415    fn external_union_ignores_unknown_rowids() {
416        let db = db_with(&[(1, "+15551234567", "iMessage", None)]);
417        let mut groups = HandleGroups::from_db(&db).unwrap();
418        groups.union(1, 99);
419        groups.union(98, 99);
420        let map = groups.finalize();
421
422        assert_eq!(map.len(), 2);
423        assert!(!map.contains_key(&99));
424    }
425
426    #[test]
427    fn group_ids_are_sequential_in_rowid_order() {
428        let db = db_with(&[
429            (5, "c", "iMessage", None),
430            (3, "b", "iMessage", None),
431            (9, "b", "SMS", None),
432            (1, "a", "iMessage", None),
433        ]);
434        let map = HandleGroups::from_db(&db).unwrap().finalize();
435
436        assert_eq!(map[&0], 0);
437        assert_eq!(map[&1], 1);
438        assert_eq!(map[&3], 2);
439        assert_eq!(map[&5], 3);
440        assert_eq!(map[&9], 2);
441    }
442
443    #[test]
444    fn diagnostic_omits_person_centric_count_when_column_is_missing() {
445        let db = Connection::open_in_memory().unwrap();
446        db.execute(
447            "CREATE TABLE handle (ROWID INTEGER PRIMARY KEY, id TEXT NOT NULL)",
448            [],
449        )
450        .unwrap();
451        db.execute(
452            "INSERT INTO handle (ROWID, id) VALUES (1, 'first'), (2, 'second')",
453            [],
454        )
455        .unwrap();
456
457        let diagnostic = Handle::run_diagnostic(&db).unwrap();
458
459        assert_eq!(diagnostic.total_handles, 2);
460        assert_eq!(diagnostic.person_centric_ids, None);
461        assert_eq!(diagnostic.multi_handle_groups, 0);
462        assert_eq!(diagnostic.merged_handles, 0);
463    }
464
465    #[test]
466    fn diagnostic_counts_groups_and_merged_rows() {
467        let db = db_with(&[
468            (1, "+15551234567", "iMessage", Some("p1")),
469            (2, "+15551234567", "SMS", Some("p1")),
470            (3, "a@example.com", "iMessage", Some("p1")),
471            (4, "b@example.com", "iMessage", None),
472            (5, "b@example.com", "SMS", None),
473            (6, "+15559876543", "iMessage", Some("p2")),
474        ]);
475
476        let diagnostic = Handle::run_diagnostic(&db).unwrap();
477
478        assert_eq!(diagnostic.total_handles, 6);
479        assert_eq!(diagnostic.person_centric_ids, Some(2));
480        // {1,2,3}, {4,5}, {6}
481        assert_eq!(diagnostic.multi_handle_groups, 2);
482        assert_eq!(diagnostic.merged_handles, 3);
483    }
484
485    #[test]
486    fn diagnostic_ignores_empty_person_centric_ids() {
487        let db = db_with(&[
488            (1, "a@example.com", "iMessage", Some("")),
489            (2, "b@example.com", "iMessage", Some("")),
490        ]);
491
492        let diagnostic = Handle::run_diagnostic(&db).unwrap();
493
494        assert_eq!(diagnostic.person_centric_ids, Some(0));
495        assert_eq!(diagnostic.multi_handle_groups, 0);
496        assert_eq!(diagnostic.merged_handles, 0);
497    }
498}