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                    Ok(connection)
138                }
139                Err(why) => Err(TableError::CannotConnect(TableConnectError::Permissions(
140                    why,
141                ))),
142            };
143        }
144    };
145    Err(TableError::CannotConnect(connect_error))
146}
147
148/// 16-byte magic string at the start of every `SQLite` file.
149const SQLITE_HEADER: &[u8; 16] = b"SQLite format 3\0";
150
151/// Check the magic bytes before opening a connection: with lazy file validation in `SQLite`,
152/// you would otherwise encounter a query failure on first use for a foreign file.
153fn has_sqlite_header(path: &Path) -> Result<bool, TableError> {
154    let mut file = File::open(path).map_err(|why| match why.kind() {
155        ErrorKind::PermissionDenied => {
156            TableError::CannotConnect(TableConnectError::Inaccessible(path.to_path_buf(), why))
157        }
158        _ => TableError::CannotRead(why),
159    })?;
160    let mut header = [0u8; SQLITE_HEADER.len()];
161    match file.read_exact(&mut header) {
162        Ok(()) => Ok(&header == SQLITE_HEADER),
163        Err(why) if why.kind() == ErrorKind::UnexpectedEof => Ok(false),
164        Err(why) => Err(TableError::CannotRead(why)),
165    }
166}
167
168/// Return the database file size on disk.
169/// # Example:
170///
171/// ```
172/// use imessage_database::{
173///     util::dirs::default_db_path,
174///     tables::table::get_db_size
175/// };
176///
177/// let db_path = default_db_path().unwrap();
178/// let database_size_in_bytes = get_db_size(&db_path);
179/// ```
180pub fn get_db_size(path: &Path) -> Result<u64, TableError> {
181    Ok(metadata(path)?.len())
182}
183
184// MARK: Constants
185// Table Names
186/// Handle table name.
187pub const HANDLE: &str = "handle";
188/// Message table name.
189pub const MESSAGE: &str = "message";
190/// Chat table name.
191pub const CHAT: &str = "chat";
192/// Attachment table name.
193pub const ATTACHMENT: &str = "attachment";
194/// Chat-to-message join table name.
195pub const CHAT_MESSAGE_JOIN: &str = "chat_message_join";
196/// Message-to-attachment join table name.
197pub const MESSAGE_ATTACHMENT_JOIN: &str = "message_attachment_join";
198/// Chat-to-handle join table name.
199pub const CHAT_HANDLE_JOIN: &str = "chat_handle_join";
200/// Recently deleted messages table.
201pub const RECENTLY_DELETED: &str = "chat_recoverable_message_join";
202
203// Column names
204/// [`plist`](crate::util::plist)-encoded app-message payload column.
205pub const MESSAGE_PAYLOAD: &str = "payload_data";
206/// [`plist`](crate::util::plist)-encoded message summary column.
207pub const MESSAGE_SUMMARY_INFO: &str = "message_summary_info";
208/// [`plist`](crate::util::plist)-encoded sticker metadata column.
209pub const STICKER_USER_INFO: &str = "sticker_user_info";
210/// [`plist`](crate::util::plist)-encoded attachment attribution column.
211pub const ATTRIBUTION_INFO: &str = "attribution_info";
212/// [`plist`](crate::util::plist)-encoded chat properties column.
213pub const PROPERTIES: &str = "properties";
214
215// Default information
216/// First-person display name for the database owner.
217pub const ME: &str = "Me";
218/// Default macOS Messages database path.
219pub const DEFAULT_PATH_MACOS: &str = "Library/Messages/chat.db";
220/// Default Messages database path inside an iOS backup.
221pub const DEFAULT_PATH_IOS: &str = "3d/3d0d7e5fb2ce288813306e4d4636395e047a3d28";
222/// Replacement token found in Fitness.app messages.
223pub const FITNESS_RECEIVER: &str = "$(kIMTranscriptPluginBreadcrumbTextReceiverIdentifier)";
224
225#[cfg(test)]
226mod tests {
227    use std::{
228        env::temp_dir,
229        fs::{create_dir_all, remove_dir_all, write},
230        path::PathBuf,
231        process,
232    };
233
234    use rusqlite::{Connection, Result, Row};
235
236    use crate::error::table::{TableConnectError, TableError};
237
238    use super::{Table, get_connection};
239
240    fn scratch(label: &str) -> PathBuf {
241        let path = temp_dir().join(format!("imessage-database-table-{}-{label}", process::id()));
242        let _ = remove_dir_all(&path);
243        create_dir_all(&path).unwrap();
244        path
245    }
246
247    #[test]
248    fn connection_rejects_directory() {
249        let dir = scratch("dir");
250
251        assert!(matches!(
252            get_connection(&dir),
253            Err(TableError::CannotConnect(TableConnectError::IsDirectory(
254                _,
255                _
256            )))
257        ));
258
259        remove_dir_all(&dir).unwrap();
260    }
261
262    #[test]
263    fn connection_rejects_missing_file() {
264        let dir = scratch("missing");
265
266        assert!(matches!(
267            get_connection(&dir.join("chat.db")),
268            Err(TableError::CannotConnect(TableConnectError::DoesNotExist(
269                _,
270                _
271            )))
272        ));
273
274        remove_dir_all(&dir).unwrap();
275    }
276
277    #[test]
278    fn connection_rejects_file_without_sqlite_header() {
279        let dir = scratch("header");
280        let text = dir.join("chat.db");
281        write(&text, b"<?xml version=\"1.0\"?>").unwrap();
282        let empty = dir.join("empty.db");
283        write(&empty, b"").unwrap();
284
285        assert!(matches!(
286            get_connection(&text),
287            Err(TableError::CannotConnect(TableConnectError::NotADatabase(
288                _
289            )))
290        ));
291        assert!(matches!(
292            get_connection(&empty),
293            Err(TableError::CannotConnect(TableConnectError::NotADatabase(
294                _
295            )))
296        ));
297
298        remove_dir_all(&dir).unwrap();
299    }
300
301    #[test]
302    fn connection_opens_sqlite_file() {
303        let dir = scratch("sqlite");
304        let db = dir.join("chat.db");
305        Connection::open(&db)
306            .unwrap()
307            .execute_batch("CREATE TABLE t (x)")
308            .unwrap();
309
310        assert!(get_connection(&db).is_ok());
311
312        remove_dir_all(&dir).unwrap();
313    }
314
315    struct TestRow(i64);
316
317    impl Table for TestRow {
318        fn from_row(row: &Row) -> Result<Self> {
319            Ok(Self(row.get(0)?))
320        }
321    }
322
323    #[test]
324    fn rows_decodes_every_row_in_order() {
325        let db = Connection::open_in_memory().unwrap();
326        let mut statement = db
327            .prepare("SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3")
328            .unwrap();
329
330        let seen: Vec<i64> = TestRow::rows(&mut statement, [])
331            .unwrap()
332            .map(|row| row.unwrap().0)
333            .collect();
334
335        assert_eq!(seen, vec![1, 2, 3]);
336    }
337
338    #[test]
339    fn row_without_a_result_is_a_query_error() {
340        let db = Connection::open_in_memory().unwrap();
341        let mut statement = db.prepare("SELECT 1 WHERE 0").unwrap();
342
343        let result = TestRow::row(&mut statement, []);
344
345        assert!(matches!(result, Err(TableError::QueryError(_))));
346    }
347}