torrust_tracker/core/databases/
sqlite.rs

1//! The `SQLite3` database driver.
2use std::panic::Location;
3use std::str::FromStr;
4
5use r2d2::Pool;
6use r2d2_sqlite::rusqlite::params;
7use r2d2_sqlite::rusqlite::types::Null;
8use r2d2_sqlite::SqliteConnectionManager;
9use torrust_tracker_primitives::info_hash::InfoHash;
10use torrust_tracker_primitives::{DurationSinceUnixEpoch, PersistentTorrents};
11
12use super::driver::Driver;
13use super::{Database, Error};
14use crate::core::auth::{self, Key};
15
16const DRIVER: Driver = Driver::Sqlite3;
17
18pub struct Sqlite {
19    pool: Pool<SqliteConnectionManager>,
20}
21
22impl Database for Sqlite {
23    /// It instantiates a new `SQLite3` database driver.
24    ///
25    /// Refer to [`databases::Database::new`](crate::core::databases::Database::new).
26    ///
27    /// # Errors
28    ///
29    /// Will return `r2d2::Error` if `db_path` is not able to create `SqLite` database.
30    fn new(db_path: &str) -> Result<Sqlite, Error> {
31        let cm = SqliteConnectionManager::file(db_path);
32        Pool::new(cm).map_or_else(|err| Err((err, Driver::Sqlite3).into()), |pool| Ok(Sqlite { pool }))
33    }
34
35    /// Refer to [`databases::Database::create_database_tables`](crate::core::databases::Database::create_database_tables).
36    fn create_database_tables(&self) -> Result<(), Error> {
37        let create_whitelist_table = "
38        CREATE TABLE IF NOT EXISTS whitelist (
39            id INTEGER PRIMARY KEY AUTOINCREMENT,
40            info_hash TEXT NOT NULL UNIQUE
41        );"
42        .to_string();
43
44        let create_torrents_table = "
45        CREATE TABLE IF NOT EXISTS torrents (
46            id INTEGER PRIMARY KEY AUTOINCREMENT,
47            info_hash TEXT NOT NULL UNIQUE,
48            completed INTEGER DEFAULT 0 NOT NULL
49        );"
50        .to_string();
51
52        let create_keys_table = "
53        CREATE TABLE IF NOT EXISTS keys (
54            id INTEGER PRIMARY KEY AUTOINCREMENT,
55            key TEXT NOT NULL UNIQUE,
56            valid_until INTEGER
57         );"
58        .to_string();
59
60        let conn = self.pool.get().map_err(|e| (e, DRIVER))?;
61
62        conn.execute(&create_whitelist_table, [])?;
63        conn.execute(&create_keys_table, [])?;
64        conn.execute(&create_torrents_table, [])?;
65
66        Ok(())
67    }
68
69    /// Refer to [`databases::Database::drop_database_tables`](crate::core::databases::Database::drop_database_tables).
70    fn drop_database_tables(&self) -> Result<(), Error> {
71        let drop_whitelist_table = "
72        DROP TABLE whitelist;"
73            .to_string();
74
75        let drop_torrents_table = "
76        DROP TABLE torrents;"
77            .to_string();
78
79        let drop_keys_table = "
80        DROP TABLE keys;"
81            .to_string();
82
83        let conn = self.pool.get().map_err(|e| (e, DRIVER))?;
84
85        conn.execute(&drop_whitelist_table, [])
86            .and_then(|_| conn.execute(&drop_torrents_table, []))
87            .and_then(|_| conn.execute(&drop_keys_table, []))?;
88
89        Ok(())
90    }
91
92    /// Refer to [`databases::Database::load_persistent_torrents`](crate::core::databases::Database::load_persistent_torrents).
93    fn load_persistent_torrents(&self) -> Result<PersistentTorrents, Error> {
94        let conn = self.pool.get().map_err(|e| (e, DRIVER))?;
95
96        let mut stmt = conn.prepare("SELECT info_hash, completed FROM torrents")?;
97
98        let torrent_iter = stmt.query_map([], |row| {
99            let info_hash_string: String = row.get(0)?;
100            let info_hash = InfoHash::from_str(&info_hash_string).unwrap();
101            let completed: u32 = row.get(1)?;
102            Ok((info_hash, completed))
103        })?;
104
105        Ok(torrent_iter.filter_map(std::result::Result::ok).collect())
106    }
107
108    /// Refer to [`databases::Database::load_keys`](crate::core::databases::Database::load_keys).
109    fn load_keys(&self) -> Result<Vec<auth::PeerKey>, Error> {
110        let conn = self.pool.get().map_err(|e| (e, DRIVER))?;
111
112        let mut stmt = conn.prepare("SELECT key, valid_until FROM keys")?;
113
114        let keys_iter = stmt.query_map([], |row| {
115            let key: String = row.get(0)?;
116            let opt_valid_until: Option<i64> = row.get(1)?;
117
118            match opt_valid_until {
119                Some(valid_until) => Ok(auth::PeerKey {
120                    key: key.parse::<Key>().unwrap(),
121                    valid_until: Some(DurationSinceUnixEpoch::from_secs(valid_until.unsigned_abs())),
122                }),
123                None => Ok(auth::PeerKey {
124                    key: key.parse::<Key>().unwrap(),
125                    valid_until: None,
126                }),
127            }
128        })?;
129
130        let keys: Vec<auth::PeerKey> = keys_iter.filter_map(std::result::Result::ok).collect();
131
132        Ok(keys)
133    }
134
135    /// Refer to [`databases::Database::load_whitelist`](crate::core::databases::Database::load_whitelist).
136    fn load_whitelist(&self) -> Result<Vec<InfoHash>, Error> {
137        let conn = self.pool.get().map_err(|e| (e, DRIVER))?;
138
139        let mut stmt = conn.prepare("SELECT info_hash FROM whitelist")?;
140
141        let info_hash_iter = stmt.query_map([], |row| {
142            let info_hash: String = row.get(0)?;
143
144            Ok(InfoHash::from_str(&info_hash).unwrap())
145        })?;
146
147        let info_hashes: Vec<InfoHash> = info_hash_iter.filter_map(std::result::Result::ok).collect();
148
149        Ok(info_hashes)
150    }
151
152    /// Refer to [`databases::Database::save_persistent_torrent`](crate::core::databases::Database::save_persistent_torrent).
153    fn save_persistent_torrent(&self, info_hash: &InfoHash, completed: u32) -> Result<(), Error> {
154        let conn = self.pool.get().map_err(|e| (e, DRIVER))?;
155
156        let insert = conn.execute(
157            "INSERT INTO torrents (info_hash, completed) VALUES (?1, ?2) ON CONFLICT(info_hash) DO UPDATE SET completed = ?2",
158            [info_hash.to_string(), completed.to_string()],
159        )?;
160
161        if insert == 0 {
162            Err(Error::InsertFailed {
163                location: Location::caller(),
164                driver: DRIVER,
165            })
166        } else {
167            Ok(())
168        }
169    }
170
171    /// Refer to [`databases::Database::get_info_hash_from_whitelist`](crate::core::databases::Database::get_info_hash_from_whitelist).
172    fn get_info_hash_from_whitelist(&self, info_hash: InfoHash) -> Result<Option<InfoHash>, Error> {
173        let conn = self.pool.get().map_err(|e| (e, DRIVER))?;
174
175        let mut stmt = conn.prepare("SELECT info_hash FROM whitelist WHERE info_hash = ?")?;
176
177        let mut rows = stmt.query([info_hash.to_hex_string()])?;
178
179        let query = rows.next()?;
180
181        Ok(query.map(|f| InfoHash::from_str(&f.get_unwrap::<_, String>(0)).unwrap()))
182    }
183
184    /// Refer to [`databases::Database::add_info_hash_to_whitelist`](crate::core::databases::Database::add_info_hash_to_whitelist).
185    fn add_info_hash_to_whitelist(&self, info_hash: InfoHash) -> Result<usize, Error> {
186        let conn = self.pool.get().map_err(|e| (e, DRIVER))?;
187
188        let insert = conn.execute("INSERT INTO whitelist (info_hash) VALUES (?)", [info_hash.to_string()])?;
189
190        if insert == 0 {
191            Err(Error::InsertFailed {
192                location: Location::caller(),
193                driver: DRIVER,
194            })
195        } else {
196            Ok(insert)
197        }
198    }
199
200    /// Refer to [`databases::Database::remove_info_hash_from_whitelist`](crate::core::databases::Database::remove_info_hash_from_whitelist).
201    fn remove_info_hash_from_whitelist(&self, info_hash: InfoHash) -> Result<usize, Error> {
202        let conn = self.pool.get().map_err(|e| (e, DRIVER))?;
203
204        let deleted = conn.execute("DELETE FROM whitelist WHERE info_hash = ?", [info_hash.to_string()])?;
205
206        if deleted == 1 {
207            // should only remove a single record.
208            Ok(deleted)
209        } else {
210            Err(Error::DeleteFailed {
211                location: Location::caller(),
212                error_code: deleted,
213                driver: DRIVER,
214            })
215        }
216    }
217
218    /// Refer to [`databases::Database::get_key_from_keys`](crate::core::databases::Database::get_key_from_keys).
219    fn get_key_from_keys(&self, key: &Key) -> Result<Option<auth::PeerKey>, Error> {
220        let conn = self.pool.get().map_err(|e| (e, DRIVER))?;
221
222        let mut stmt = conn.prepare("SELECT key, valid_until FROM keys WHERE key = ?")?;
223
224        let mut rows = stmt.query([key.to_string()])?;
225
226        let key = rows.next()?;
227
228        Ok(key.map(|f| {
229            let valid_until: Option<i64> = f.get(1).unwrap();
230            let key: String = f.get(0).unwrap();
231
232            match valid_until {
233                Some(valid_until) => auth::PeerKey {
234                    key: key.parse::<Key>().unwrap(),
235                    valid_until: Some(DurationSinceUnixEpoch::from_secs(valid_until.unsigned_abs())),
236                },
237                None => auth::PeerKey {
238                    key: key.parse::<Key>().unwrap(),
239                    valid_until: None,
240                },
241            }
242        }))
243    }
244
245    /// Refer to [`databases::Database::add_key_to_keys`](crate::core::databases::Database::add_key_to_keys).
246    fn add_key_to_keys(&self, auth_key: &auth::PeerKey) -> Result<usize, Error> {
247        let conn = self.pool.get().map_err(|e| (e, DRIVER))?;
248
249        let insert = match auth_key.valid_until {
250            Some(valid_until) => conn.execute(
251                "INSERT INTO keys (key, valid_until) VALUES (?1, ?2)",
252                [auth_key.key.to_string(), valid_until.as_secs().to_string()],
253            )?,
254            None => conn.execute(
255                "INSERT INTO keys (key, valid_until) VALUES (?1, ?2)",
256                params![auth_key.key.to_string(), Null],
257            )?,
258        };
259
260        if insert == 0 {
261            Err(Error::InsertFailed {
262                location: Location::caller(),
263                driver: DRIVER,
264            })
265        } else {
266            Ok(insert)
267        }
268    }
269
270    /// Refer to [`databases::Database::remove_key_from_keys`](crate::core::databases::Database::remove_key_from_keys).
271    fn remove_key_from_keys(&self, key: &Key) -> Result<usize, Error> {
272        let conn = self.pool.get().map_err(|e| (e, DRIVER))?;
273
274        let deleted = conn.execute("DELETE FROM keys WHERE key = ?", [key.to_string()])?;
275
276        if deleted == 1 {
277            // should only remove a single record.
278            Ok(deleted)
279        } else {
280            Err(Error::DeleteFailed {
281                location: Location::caller(),
282                error_code: deleted,
283                driver: DRIVER,
284            })
285        }
286    }
287}
torrust_tracker/core/databases/sqlite.rs

torrust_tracker/core/databases/
sqlite.rs
