torrust_tracker/core/databases/
mysql.rs

1//! The `MySQL` database driver.
2use std::str::FromStr;
3use std::time::Duration;
4
5use r2d2::Pool;
6use r2d2_mysql::mysql::prelude::Queryable;
7use r2d2_mysql::mysql::{params, Opts, OptsBuilder};
8use r2d2_mysql::MySqlConnectionManager;
9use torrust_tracker_primitives::info_hash::InfoHash;
10use torrust_tracker_primitives::PersistentTorrents;
11
12use super::driver::Driver;
13use super::{Database, Error};
14use crate::core::auth::{self, Key};
15use crate::shared::bit_torrent::common::AUTH_KEY_LENGTH;
16
17const DRIVER: Driver = Driver::MySQL;
18
19pub struct Mysql {
20    pool: Pool<MySqlConnectionManager>,
21}
22
23impl Database for Mysql {
24    /// It instantiates a new `MySQL` database driver.
25    ///
26    /// Refer to [`databases::Database::new`](crate::core::databases::Database::new).
27    ///
28    /// # Errors
29    ///
30    /// Will return `r2d2::Error` if `db_path` is not able to create `MySQL` database.
31    fn new(db_path: &str) -> Result<Self, Error> {
32        let opts = Opts::from_url(db_path)?;
33        let builder = OptsBuilder::from_opts(opts);
34        let manager = MySqlConnectionManager::new(builder);
35        let pool = r2d2::Pool::builder().build(manager).map_err(|e| (e, DRIVER))?;
36
37        Ok(Self { pool })
38    }
39
40    /// Refer to [`databases::Database::create_database_tables`](crate::core::databases::Database::create_database_tables).
41    fn create_database_tables(&self) -> Result<(), Error> {
42        let create_whitelist_table = "
43        CREATE TABLE IF NOT EXISTS whitelist (
44            id integer PRIMARY KEY AUTO_INCREMENT,
45            info_hash VARCHAR(40) NOT NULL UNIQUE
46        );"
47        .to_string();
48
49        let create_torrents_table = "
50        CREATE TABLE IF NOT EXISTS torrents (
51            id integer PRIMARY KEY AUTO_INCREMENT,
52            info_hash VARCHAR(40) NOT NULL UNIQUE,
53            completed INTEGER DEFAULT 0 NOT NULL
54        );"
55        .to_string();
56
57        let create_keys_table = format!(
58            "
59        CREATE TABLE IF NOT EXISTS `keys` (
60          `id` INT NOT NULL AUTO_INCREMENT,
61          `key` VARCHAR({}) NOT NULL,
62          `valid_until` INT(10),
63          PRIMARY KEY (`id`),
64          UNIQUE (`key`)
65        );",
66            i8::try_from(AUTH_KEY_LENGTH).expect("auth::Auth Key Length Should fit within a i8!")
67        );
68
69        let mut conn = self.pool.get().map_err(|e| (e, DRIVER))?;
70
71        conn.query_drop(&create_torrents_table)
72            .expect("Could not create torrents table.");
73        conn.query_drop(&create_keys_table).expect("Could not create keys table.");
74        conn.query_drop(&create_whitelist_table)
75            .expect("Could not create whitelist table.");
76
77        Ok(())
78    }
79
80    /// Refer to [`databases::Database::drop_database_tables`](crate::core::databases::Database::drop_database_tables).
81    fn drop_database_tables(&self) -> Result<(), Error> {
82        let drop_whitelist_table = "
83        DROP TABLE `whitelist`;"
84            .to_string();
85
86        let drop_torrents_table = "
87        DROP TABLE `torrents`;"
88            .to_string();
89
90        let drop_keys_table = "
91            DROP TABLE `keys`;"
92            .to_string();
93
94        let mut conn = self.pool.get().map_err(|e| (e, DRIVER))?;
95
96        conn.query_drop(&drop_whitelist_table)
97            .expect("Could not drop `whitelist` table.");
98        conn.query_drop(&drop_torrents_table)
99            .expect("Could not drop `torrents` table.");
100        conn.query_drop(&drop_keys_table).expect("Could not drop `keys` table.");
101
102        Ok(())
103    }
104
105    /// Refer to [`databases::Database::load_persistent_torrents`](crate::core::databases::Database::load_persistent_torrents).
106    fn load_persistent_torrents(&self) -> Result<PersistentTorrents, Error> {
107        let mut conn = self.pool.get().map_err(|e| (e, DRIVER))?;
108
109        let torrents = conn.query_map(
110            "SELECT info_hash, completed FROM torrents",
111            |(info_hash_string, completed): (String, u32)| {
112                let info_hash = InfoHash::from_str(&info_hash_string).unwrap();
113                (info_hash, completed)
114            },
115        )?;
116
117        Ok(torrents.iter().copied().collect())
118    }
119
120    /// Refer to [`databases::Database::load_keys`](crate::core::databases::Database::load_keys).
121    fn load_keys(&self) -> Result<Vec<auth::PeerKey>, Error> {
122        let mut conn = self.pool.get().map_err(|e| (e, DRIVER))?;
123
124        let keys = conn.query_map(
125            "SELECT `key`, valid_until FROM `keys`",
126            |(key, valid_until): (String, Option<i64>)| match valid_until {
127                Some(valid_until) => auth::PeerKey {
128                    key: key.parse::<Key>().unwrap(),
129                    valid_until: Some(Duration::from_secs(valid_until.unsigned_abs())),
130                },
131                None => auth::PeerKey {
132                    key: key.parse::<Key>().unwrap(),
133                    valid_until: None,
134                },
135            },
136        )?;
137
138        Ok(keys)
139    }
140
141    /// Refer to [`databases::Database::load_whitelist`](crate::core::databases::Database::load_whitelist).
142    fn load_whitelist(&self) -> Result<Vec<InfoHash>, Error> {
143        let mut conn = self.pool.get().map_err(|e| (e, DRIVER))?;
144
145        let info_hashes = conn.query_map("SELECT info_hash FROM whitelist", |info_hash: String| {
146            InfoHash::from_str(&info_hash).unwrap()
147        })?;
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        const COMMAND : &str = "INSERT INTO torrents (info_hash, completed) VALUES (:info_hash_str, :completed) ON DUPLICATE KEY UPDATE completed = VALUES(completed)";
155
156        let mut conn = self.pool.get().map_err(|e| (e, DRIVER))?;
157
158        let info_hash_str = info_hash.to_string();
159
160        tracing::debug!("{}", info_hash_str);
161
162        Ok(conn.exec_drop(COMMAND, params! { info_hash_str, completed })?)
163    }
164
165    /// Refer to [`databases::Database::get_info_hash_from_whitelist`](crate::core::databases::Database::get_info_hash_from_whitelist).
166    fn get_info_hash_from_whitelist(&self, info_hash: InfoHash) -> Result<Option<InfoHash>, Error> {
167        let mut conn = self.pool.get().map_err(|e| (e, DRIVER))?;
168
169        let select = conn.exec_first::<String, _, _>(
170            "SELECT info_hash FROM whitelist WHERE info_hash = :info_hash",
171            params! { "info_hash" => info_hash.to_hex_string() },
172        )?;
173
174        let info_hash = select.map(|f| InfoHash::from_str(&f).expect("Failed to decode InfoHash String from DB!"));
175
176        Ok(info_hash)
177    }
178
179    /// Refer to [`databases::Database::add_info_hash_to_whitelist`](crate::core::databases::Database::add_info_hash_to_whitelist).
180    fn add_info_hash_to_whitelist(&self, info_hash: InfoHash) -> Result<usize, Error> {
181        let mut conn = self.pool.get().map_err(|e| (e, DRIVER))?;
182
183        let info_hash_str = info_hash.to_string();
184
185        conn.exec_drop(
186            "INSERT INTO whitelist (info_hash) VALUES (:info_hash_str)",
187            params! { info_hash_str },
188        )?;
189
190        Ok(1)
191    }
192
193    /// Refer to [`databases::Database::remove_info_hash_from_whitelist`](crate::core::databases::Database::remove_info_hash_from_whitelist).
194    fn remove_info_hash_from_whitelist(&self, info_hash: InfoHash) -> Result<usize, Error> {
195        let mut conn = self.pool.get().map_err(|e| (e, DRIVER))?;
196
197        let info_hash = info_hash.to_string();
198
199        conn.exec_drop("DELETE FROM whitelist WHERE info_hash = :info_hash", params! { info_hash })?;
200
201        Ok(1)
202    }
203
204    /// Refer to [`databases::Database::get_key_from_keys`](crate::core::databases::Database::get_key_from_keys).
205    fn get_key_from_keys(&self, key: &Key) -> Result<Option<auth::PeerKey>, Error> {
206        let mut conn = self.pool.get().map_err(|e| (e, DRIVER))?;
207
208        let query = conn.exec_first::<(String, Option<i64>), _, _>(
209            "SELECT `key`, valid_until FROM `keys` WHERE `key` = :key",
210            params! { "key" => key.to_string() },
211        );
212
213        let key = query?;
214
215        Ok(key.map(|(key, opt_valid_until)| match opt_valid_until {
216            Some(valid_until) => auth::PeerKey {
217                key: key.parse::<Key>().unwrap(),
218                valid_until: Some(Duration::from_secs(valid_until.unsigned_abs())),
219            },
220            None => auth::PeerKey {
221                key: key.parse::<Key>().unwrap(),
222                valid_until: None,
223            },
224        }))
225    }
226
227    /// Refer to [`databases::Database::add_key_to_keys`](crate::core::databases::Database::add_key_to_keys).
228    fn add_key_to_keys(&self, auth_key: &auth::PeerKey) -> Result<usize, Error> {
229        let mut conn = self.pool.get().map_err(|e| (e, DRIVER))?;
230
231        let key = auth_key.key.to_string();
232        let valid_until = match auth_key.valid_until {
233            Some(valid_until) => valid_until.as_secs().to_string(),
234            None => todo!(),
235        };
236
237        conn.exec_drop(
238            "INSERT INTO `keys` (`key`, valid_until) VALUES (:key, :valid_until)",
239            params! { key, valid_until },
240        )?;
241
242        Ok(1)
243    }
244
245    /// Refer to [`databases::Database::remove_key_from_keys`](crate::core::databases::Database::remove_key_from_keys).
246    fn remove_key_from_keys(&self, key: &Key) -> Result<usize, Error> {
247        let mut conn = self.pool.get().map_err(|e| (e, DRIVER))?;
248
249        conn.exec_drop("DELETE FROM `keys` WHERE key = :key", params! { "key" => key.to_string() })?;
250
251        Ok(1)
252    }
253}
torrust_tracker/core/databases/mysql.rs

torrust_tracker/core/databases/
mysql.rs
