Skip to main content

imessage_database/tables/
table.rs

1/*!
2 Table traits, database connection helpers, and shared table constants.
3
4 # Row iteration
5
6 Prepare a statement with the appropriate table type, then use [`Table::rows`]
7 to decode one row at a time. Use any statement that selects the table's columns.
8
9 ```no_run
10 use imessage_database::{
11    tables::{handle::Handle, table::{Table, get_connection}},
12    util::dirs::default_db_path,
13 };
14
15 let db_path = default_db_path().unwrap();
16 let db = get_connection(&db_path).unwrap();
17
18 let mut statement = Handle::get(&db).unwrap();
19 for handle_result in Handle::rows(&mut statement, []).unwrap() {
20     match handle_result {
21         Ok(handle) => println!("Handle: {}", handle.id),
22         Err(e) => eprintln!("Error: {:?}", e),
23     }
24 }
25 ```
26*/
27
28use std::{
29    fs::{File, metadata},
30    io::{ErrorKind, Read},
31    path::Path,
32};
33
34use rusqlite::{Connection, Error, OpenFlags, Params, Result, Row, Statement, blob::Blob};
35
36use crate::{
37    error::table::{TableConnectError, TableError},
38    util::path_hint::{Presence, diagnose, presence},
39};
40
41// MARK: Traits
42/// Deserialize a database record from a stepped SQLite row.
43///
44/// Implement [`from_row`](Self::from_row), then use [`rows`](Self::rows) and
45/// [`row`](Self::row) with any statement that selects the table's columns.
46pub trait Table: Sized {
47    /// Deserialize a single row into `Self`. Returns [`rusqlite::Result`]
48    /// for direct use inside `rusqlite::query_map` / `query_row`
49    /// callbacks. For high-level iteration, prefer [`Table::rows`] or
50    /// [`Table::row`].
51    fn from_row(row: &Row) -> Result<Self>;
52
53    /// Iterate over the results of `stmt`, deserializing each row through
54    /// [`from_row`](Self::from_row). Return [`TableError`] for both fetch and
55    /// deserialization failures. Pass either a [`rusqlite::Statement`] or a
56    /// [`rusqlite::CachedStatement`] (via deref coercion).
57    fn rows<'stmt, P: Params>(
58        stmt: &'stmt mut Statement<'_>,
59        params: P,
60    ) -> Result<impl Iterator<Item = Result<Self, TableError>> + 'stmt, TableError>
61    where
62        Self: 'stmt,
63    {
64        let mapped = stmt.query_map(params, |row| Ok(Self::from_row(row)))?;
65        Ok(mapped.map(flatten_row))
66    }
67
68    /// Fetch exactly one row from `stmt`. Returns
69    /// [`TableError::QueryError`] if the row is missing or fails to
70    /// deserialize. Accepts both [`rusqlite::Statement`] and
71    /// [`rusqlite::CachedStatement`] (the latter via deref coercion).
72    fn row<P: Params>(stmt: &mut Statement<'_>, params: P) -> Result<Self, TableError> {
73        flatten_row(stmt.query_row(params, |row| Ok(Self::from_row(row))))
74    }
75}
76
77/// Open a `BLOB` column for `rowid`, read-only.
78pub fn open_blob<'a>(
79    db: &'a Connection,
80    table: &str,
81    column: &str,
82    rowid: i64,
83) -> Option<Blob<'a>> {
84    db.blob_open(rusqlite::MAIN_DB, table, column, rowid, true)
85        .ok()
86}
87
88/// Flatten the doubly-nested result produced by `rusqlite::query_map` /
89/// `query_row` callbacks into a single [`TableError`]. The outer layer
90/// represents row-fetch failures, the inner layer represents row-deserialize
91/// failures from [`Table::from_row`].
92pub(crate) fn flatten_row<T>(item: Result<Result<T, Error>, Error>) -> Result<T, TableError> {
93    match item {
94        Ok(Ok(row)) => Ok(row),
95        Err(why) | Ok(Err(why)) => Err(TableError::QueryError(why)),
96    }
97}
98
99// MARK: Database
100/// Open the Messages `SQLite` database read-only.
101/// # Example:
102///
103/// ```
104/// use imessage_database::{
105///     util::dirs::default_db_path,
106///     tables::table::get_connection
107/// };
108///
109/// let db_path = default_db_path().unwrap();
110/// let connection = get_connection(&db_path);
111/// ```
112pub fn get_connection(path: &Path) -> Result<Connection, TableError> {
113    let connect_error = match presence(path) {
114        Err(why) => TableConnectError::Inaccessible(path.to_path_buf(), why),
115        Ok(Presence::Missing) => {
116            TableConnectError::DoesNotExist(path.to_path_buf(), diagnose(path))
117        }
118        Ok(Presence::Directory) => {
119            TableConnectError::IsDirectory(path.to_path_buf(), diagnose(path))
120        }
121        // Inspect only regular files to avoid blocking on a FIFO or device.
122        Ok(Presence::Other) => TableConnectError::NotADatabase(path.to_path_buf()),
123        Ok(Presence::File) => {
124            if !has_sqlite_header(path)? {
125                return Err(TableError::CannotConnect(TableConnectError::NotADatabase(
126                    path.to_path_buf(),
127                )));
128            }
129            return match Connection::open_with_flags(
130                path,
131                OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
132            ) {
133                Ok(connection) => {
134                    // Read pages from the mapped region where SQLite supports it.
135                    let _ = connection.pragma_update(None, "mmap_size", 8_589_934_592_i64); // up to 8 GiB
136                    let _ = connection.pragma_update(None, "cache_size", -65_536_i64); // ~64 MiB
137                    // Since SQLite 3.54, clearing the bindings of a cached statement with a
138                    // `col = ?` term expires the cache
139                    let _ = connection.set_db_config(
140                        rusqlite::config::DbConfig::SQLITE_DBCONFIG_ENABLE_QPSG,
141                        true,
142                    );
143                    Ok(connection)
144                }
145                Err(why) => Err(TableError::CannotConnect(TableConnectError::Permissions(
146                    why,
147                ))),
148            };
149        }
150    };
151    Err(TableError::CannotConnect(connect_error))
152}
153
154/// 16-byte magic string at the start of every `SQLite` file.
155const SQLITE_HEADER: &[u8; 16] = b"SQLite format 3\0";
156
157/// Check the magic bytes before opening a connection: with lazy file validation in `SQLite`,
158/// you would otherwise encounter a query failure on first use for a foreign file.
159fn has_sqlite_header(path: &Path) -> Result<bool, TableError> {
160    let mut file = File::open(path).map_err(|why| match why.kind() {
161        ErrorKind::PermissionDenied => {
162            TableError::CannotConnect(TableConnectError::Inaccessible(path.to_path_buf(), why))
163        }
164        _ => TableError::CannotRead(why),
165    })?;
166    let mut header = [0u8; SQLITE_HEADER.len()];
167    match file.read_exact(&mut header) {
168        Ok(()) => Ok(&header == SQLITE_HEADER),
169        Err(why) if why.kind() == ErrorKind::UnexpectedEof => Ok(false),
170        Err(why) => Err(TableError::CannotRead(why)),
171    }
172}
173
174/// Return the database file size on disk.
175/// # Example:
176///
177/// ```
178/// use imessage_database::{
179///     util::dirs::default_db_path,
180///     tables::table::get_db_size
181/// };
182///
183/// let db_path = default_db_path().unwrap();
184/// let database_size_in_bytes = get_db_size(&db_path);
185/// ```
186pub fn get_db_size(path: &Path) -> Result<u64, TableError> {
187    Ok(metadata(path)?.len())
188}
189
190// MARK: Constants
191// Table Names
192/// Handle table name.
193pub const HANDLE: &str = "handle";
194/// Message table name.
195pub const MESSAGE: &str = "message";
196/// Chat table name.
197pub const CHAT: &str = "chat";
198/// Attachment table name.
199pub const ATTACHMENT: &str = "attachment";
200/// Chat-to-message join table name.
201pub const CHAT_MESSAGE_JOIN: &str = "chat_message_join";
202/// Message-to-attachment join table name.
203pub const MESSAGE_ATTACHMENT_JOIN: &str = "message_attachment_join";
204/// Chat-to-handle join table name.
205pub const CHAT_HANDLE_JOIN: &str = "chat_handle_join";
206/// Recently deleted messages table.
207pub const RECENTLY_DELETED: &str = "chat_recoverable_message_join";
208
209// Column names
210/// [`plist`](crate::util::plist)-encoded app-message payload column.
211pub const MESSAGE_PAYLOAD: &str = "payload_data";
212/// [`plist`](crate::util::plist)-encoded message summary column.
213pub const MESSAGE_SUMMARY_INFO: &str = "message_summary_info";
214/// [`plist`](crate::util::plist)-encoded sticker metadata column.
215pub const STICKER_USER_INFO: &str = "sticker_user_info";
216/// [`plist`](crate::util::plist)-encoded attachment attribution column.
217pub const ATTRIBUTION_INFO: &str = "attribution_info";
218/// [`plist`](crate::util::plist)-encoded chat properties column.
219pub const PROPERTIES: &str = "properties";
220
221// Default information
222/// First-person display name for the database owner.
223pub const ME: &str = "Me";
224/// Default macOS Messages database path.
225pub const DEFAULT_PATH_MACOS: &str = "Library/Messages/chat.db";
226/// Default Messages database path inside an iOS backup.
227pub const DEFAULT_PATH_IOS: &str = "3d/3d0d7e5fb2ce288813306e4d4636395e047a3d28";
228/// Replacement token found in Fitness.app messages.
229pub const FITNESS_RECEIVER: &str = "$(kIMTranscriptPluginBreadcrumbTextReceiverIdentifier)";
230
231#[cfg(test)]
232mod tests {
233    use std::{
234        env::temp_dir,
235        fs::{create_dir_all, remove_dir_all, write},
236        path::PathBuf,
237        process,
238    };
239
240    use rusqlite::{Connection, Result, Row};
241
242    use crate::error::table::{TableConnectError, TableError};
243
244    use super::{Table, get_connection};
245
246    fn scratch(label: &str) -> PathBuf {
247        let path = temp_dir().join(format!("imessage-database-table-{}-{label}", process::id()));
248        let _ = remove_dir_all(&path);
249        create_dir_all(&path).unwrap();
250        path
251    }
252
253    #[test]
254    fn connection_rejects_directory() {
255        let dir = scratch("dir");
256
257        assert!(matches!(
258            get_connection(&dir),
259            Err(TableError::CannotConnect(TableConnectError::IsDirectory(
260                _,
261                _
262            )))
263        ));
264
265        remove_dir_all(&dir).unwrap();
266    }
267
268    #[test]
269    fn connection_rejects_missing_file() {
270        let dir = scratch("missing");
271
272        assert!(matches!(
273            get_connection(&dir.join("chat.db")),
274            Err(TableError::CannotConnect(TableConnectError::DoesNotExist(
275                _,
276                _
277            )))
278        ));
279
280        remove_dir_all(&dir).unwrap();
281    }
282
283    #[test]
284    fn connection_rejects_file_without_sqlite_header() {
285        let dir = scratch("header");
286        let text = dir.join("chat.db");
287        write(&text, b"<?xml version=\"1.0\"?>").unwrap();
288        let empty = dir.join("empty.db");
289        write(&empty, b"").unwrap();
290
291        assert!(matches!(
292            get_connection(&text),
293            Err(TableError::CannotConnect(TableConnectError::NotADatabase(
294                _
295            )))
296        ));
297        assert!(matches!(
298            get_connection(&empty),
299            Err(TableError::CannotConnect(TableConnectError::NotADatabase(
300                _
301            )))
302        ));
303
304        remove_dir_all(&dir).unwrap();
305    }
306
307    #[test]
308    fn connection_opens_sqlite_file() {
309        let dir = scratch("sqlite");
310        let db = dir.join("chat.db");
311        Connection::open(&db)
312            .unwrap()
313            .execute_batch("CREATE TABLE t (x)")
314            .unwrap();
315
316        assert!(get_connection(&db).is_ok());
317
318        remove_dir_all(&dir).unwrap();
319    }
320
321    struct TestRow(i64);
322
323    impl Table for TestRow {
324        fn from_row(row: &Row) -> Result<Self> {
325            Ok(Self(row.get(0)?))
326        }
327    }
328
329    #[test]
330    fn rows_decodes_every_row_in_order() {
331        let db = Connection::open_in_memory().unwrap();
332        let mut statement = db
333            .prepare("SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3")
334            .unwrap();
335
336        let seen: Vec<i64> = TestRow::rows(&mut statement, [])
337            .unwrap()
338            .map(|row| row.unwrap().0)
339            .collect();
340
341        assert_eq!(seen, vec![1, 2, 3]);
342    }
343
344    #[test]
345    fn row_without_a_result_is_a_query_error() {
346        let db = Connection::open_in_memory().unwrap();
347        let mut statement = db.prepare("SELECT 1 WHERE 0").unwrap();
348
349        let result = TestRow::row(&mut statement, []);
350
351        assert!(matches!(result, Err(TableError::QueryError(_))));
352    }
353}