diff --git a/src/config.rs b/src/config.rs index 429997fcc..4efcb2825 100644 --- a/src/config.rs +++ b/src/config.rs @@ -17,6 +17,7 @@ pub enum TrackerServer { #[derive(Serialize, Deserialize)] pub struct UdpTrackerConfig { + pub enabled: bool, pub bind_address: String, pub announce_interval: u32, } @@ -28,6 +29,7 @@ pub struct HttpTrackerConfig { pub on_reverse_proxy: bool, pub announce_interval: u32, pub ssl_enabled: bool, + pub ssl_bind_address: String, #[serde(serialize_with = "none_as_empty_string")] pub ssl_cert_path: Option, #[serde(serialize_with = "none_as_empty_string")] @@ -52,10 +54,14 @@ pub struct Configuration { pub log_level: Option, pub mode: TrackerMode, pub db_path: String, + pub persistence: bool, pub cleanup_interval: Option, + pub cleanup_peerless: bool, pub external_ip: Option, pub udp_tracker: UdpTrackerConfig, + pub udp_tracker_ipv6: UdpTrackerConfig, pub http_tracker: HttpTrackerConfig, + pub http_tracker_ipv6: HttpTrackerConfig, pub http_api: HttpApiConfig, } @@ -128,18 +134,37 @@ impl Configuration { log_level: Option::from(String::from("info")), mode: TrackerMode::PublicMode, db_path: String::from("data.db"), + persistence: false, cleanup_interval: Some(600), + cleanup_peerless: true, external_ip: Some(String::from("0.0.0.0")), udp_tracker: UdpTrackerConfig { + enabled: true, bind_address: String::from("0.0.0.0:6969"), announce_interval: 120, }, + udp_tracker_ipv6: UdpTrackerConfig { + enabled: false, + bind_address: String::from("[::]:6969"), + announce_interval: 120, + }, http_tracker: HttpTrackerConfig { enabled: false, bind_address: String::from("0.0.0.0:6969"), on_reverse_proxy: false, announce_interval: 120, ssl_enabled: false, + ssl_bind_address: String::from("0.0.0.0:6868"), + ssl_cert_path: None, + ssl_key_path: None + }, + http_tracker_ipv6: HttpTrackerConfig { + enabled: false, + bind_address: String::from("[::]:6969"), + on_reverse_proxy: false, + announce_interval: 120, + ssl_enabled: false, + ssl_bind_address: String::from("[::]:6868"), ssl_cert_path: None, ssl_key_path: None }, diff --git a/src/database.rs b/src/database.rs index fbec824a0..827c7dbed 100644 --- a/src/database.rs +++ b/src/database.rs @@ -1,10 +1,11 @@ -use crate::{InfoHash, AUTH_KEY_LENGTH}; +use crate::{InfoHash, AUTH_KEY_LENGTH, TorrentTracker}; use log::debug; use r2d2_sqlite::{SqliteConnectionManager, rusqlite}; use r2d2::{Pool}; use r2d2_sqlite::rusqlite::NO_PARAMS; use crate::key_manager::AuthKey; use std::str::FromStr; +use std::sync::Arc; pub struct SqliteDatabase { pool: Pool @@ -32,6 +33,13 @@ impl SqliteDatabase { info_hash VARCHAR(20) NOT NULL UNIQUE );".to_string(); + let create_torrents_table = " + CREATE TABLE IF NOT EXISTS torrents ( + id integer PRIMARY KEY AUTOINCREMENT, + info_hash VARCHAR(20) NOT NULL UNIQUE, + completed INTEGER DEFAULT 0 NOT NULL + );".to_string(); + let create_keys_table = format!(" CREATE TABLE IF NOT EXISTS keys ( id integer PRIMARY KEY AUTOINCREMENT, @@ -43,7 +51,15 @@ impl SqliteDatabase { match conn.execute(&create_whitelist_table, NO_PARAMS) { Ok(updated) => { match conn.execute(&create_keys_table, NO_PARAMS) { - Ok(updated2) => Ok(updated + updated2), + Ok(updated2) => { + match conn.execute(&create_torrents_table, NO_PARAMS) { + Ok(updated3) => Ok(updated + updated2 + updated3), + Err(e) => { + debug!("{:?}", e); + Err(e) + } + } + } Err(e) => { debug!("{:?}", e); Err(e) @@ -57,6 +73,42 @@ impl SqliteDatabase { } } + pub async fn load_persistent_torrent_data(&self, tracker: Arc) -> Result { + let tracker_copy = tracker.clone(); + let conn = self.pool.get().unwrap(); + let mut stmt = conn.prepare("SELECT info_hash, completed FROM torrents")?; + + let info_hash_iter = stmt.query_map(NO_PARAMS, |row| { + let info_hash: String = row.get(0)?; + let info_hash_converted = InfoHash::from_str(&info_hash).unwrap(); + let completed: u32 = row.get(1)?; + Ok((info_hash_converted, completed)) + })?; + + for info_hash_item in info_hash_iter { + let (info_hash, completed): (InfoHash, u32) = info_hash_item.unwrap(); + tracker_copy.add_torrent(&info_hash, 0u32, completed, 0u32).await; + } + + Ok(true) + } + + pub async fn save_persistent_torrent_data(&self, tracker: Arc) -> Result { + let tracker_copy = tracker.clone(); + let mut conn = self.pool.get().unwrap(); + let db = tracker_copy.get_torrents().await; + let db_transaction = conn.transaction()?; + let _: Vec<_> = db + .iter() + .map(|(info_hash, torrent_entry)| { + let (_seeders, completed, _leechers) = torrent_entry.get_stats(); + let _ = db_transaction.execute("INSERT OR REPLACE INTO torrents (info_hash, completed) VALUES (?, ?)", &[info_hash.to_string(), completed.to_string()]); + }) + .collect(); + let _ = db_transaction.commit(); + Ok(true) + } + pub async fn get_info_hash_from_whitelist(&self, info_hash: &str) -> Result { let conn = self.pool.get().unwrap(); let mut stmt = conn.prepare("SELECT info_hash FROM whitelist WHERE info_hash = ?")?; diff --git a/src/main.rs b/src/main.rs index d8a73854e..741f05061 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,6 @@ use std::net::SocketAddr; use log::{info}; -use torrust_tracker::{http_api_server, Configuration, TorrentTracker, UdpServer, HttpTrackerConfig, UdpTrackerConfig, HttpApiConfig, logging, TrackerServer}; +use torrust_tracker::{http_api_server, Configuration, TorrentTracker, UdpServer, HttpTrackerConfig, UdpTrackerConfig, HttpApiConfig, logging}; use std::sync::Arc; use tokio::task::JoinHandle; use torrust_tracker::torrust_http_tracker::server::HttpServer; @@ -19,6 +19,11 @@ async fn main() { // the singleton torrent tracker that gets passed to the HTTP and UDP server let tracker = Arc::new(TorrentTracker::new(config.clone())); + // Load torrents if enabled + if config.persistence { + load_torrents_into_memory(tracker.clone()).await; + } + // start torrent cleanup job (periodically removes old peers) let _torrent_cleanup_job = start_torrent_cleanup_job(config.clone(), tracker.clone()).unwrap(); @@ -27,21 +32,58 @@ async fn main() { let _api_server = start_api_server(&config.http_api, tracker.clone()); } - // check which tracker to run, UDP (Default) or HTTP - let _tracker_server = match config.get_tracker_server() { - TrackerServer::UDP => { - start_udp_tracker_server(&config.udp_tracker, tracker.clone()).await - } - TrackerServer::HTTP => { - start_http_tracker_server(&config.http_tracker, tracker.clone()) - } - }; + // start UDP tracker if enabled + if config.udp_tracker.enabled { + let _udp_server = start_udp_tracker_server(&config.udp_tracker, tracker.clone()).await; + } + + // start UDP tracker for IPv6 if enabled + if config.udp_tracker_ipv6.enabled { + let _udp_server_ipv6 = start_udp_ipv6_tracker_server(&config.udp_tracker_ipv6, tracker.clone()).await; + } + + // start HTTP tracker if enabled + if config.http_tracker.enabled { + let _http_server = start_http_tracker_server(&config.http_tracker, tracker.clone()); + } + // start HTTPS tracker if enabled + if config.http_tracker.ssl_enabled { + let _http_ssl_server = start_http_ssl_tracker_server(&config.http_tracker, tracker.clone()); + } + + //start HTTP tracker for IPv6 if enabled + if config.http_tracker_ipv6.enabled { + let _http_server_ipv6 = start_http_tracker_server(&config.http_tracker_ipv6, tracker.clone()); + } + + // start HTTPS tracker for IPv6 if enabled + if config.http_tracker_ipv6.ssl_enabled { + let _http_ssl_server_ipv6 = start_http_ssl_tracker_server(&config.http_tracker_ipv6, tracker.clone()); + } + + // handle the signals here let ctrl_c = tokio::signal::ctrl_c(); tokio::select! { - _ = _tracker_server => { panic!("Tracker server exited.") }, _ = ctrl_c => { info!("Torrust shutting down..") } } + + // Save torrents if enabled + if config.persistence { + save_torrents_into_memory(tracker.clone()).await; + } +} + +async fn load_torrents_into_memory(tracker: Arc) { + info!("Loading torrents from SQL into memory..."); + let _ = tracker.load_torrents(tracker.clone()).await; + info!("Torrents loaded"); +} + +async fn save_torrents_into_memory(tracker: Arc) { + info!("Saving torrents into SQL from memory..."); + let _ = tracker.save_torrents(tracker.clone()).await; + info!("Torrents saved"); } fn start_torrent_cleanup_job(config: Arc, tracker: Arc) -> Option> { @@ -77,19 +119,26 @@ fn start_api_server(config: &HttpApiConfig, tracker: Arc) -> Joi fn start_http_tracker_server(config: &HttpTrackerConfig, tracker: Arc) -> JoinHandle<()> { let http_tracker = HttpServer::new(tracker); let bind_addr = config.bind_address.parse::().unwrap(); - let ssl_enabled = config.ssl_enabled; + + tokio::spawn(async move { + // run with tls if ssl_enabled and cert and key path are set + info!("Starting HTTP server on: {}", bind_addr); + http_tracker.start(bind_addr).await; + }) +} + +fn start_http_ssl_tracker_server(config: &HttpTrackerConfig, tracker: Arc) -> JoinHandle<()> { + let http_tracker = HttpServer::new(tracker); + let ssl_bind_addr = config.ssl_bind_address.parse::().unwrap(); let ssl_cert_path = config.ssl_cert_path.clone(); let ssl_key_path = config.ssl_key_path.clone(); tokio::spawn(async move { // run with tls if ssl_enabled and cert and key path are set - if ssl_enabled && ssl_cert_path.is_some() && ssl_key_path.is_some() { - info!("Starting HTTPS server on: {} (TLS)", bind_addr); - http_tracker.start_tls(bind_addr, ssl_cert_path.as_ref().unwrap(), ssl_key_path.as_ref().unwrap()).await; - } else { - info!("Starting HTTP server on: {}", bind_addr); - http_tracker.start(bind_addr).await; + if ssl_cert_path.is_some() && ssl_key_path.is_some() { + info!("Starting HTTPS server on: {} (TLS)", ssl_bind_addr); + http_tracker.start_tls(ssl_bind_addr, ssl_cert_path.as_ref().unwrap(), ssl_key_path.as_ref().unwrap()).await; } }) } @@ -104,3 +153,14 @@ async fn start_udp_tracker_server(config: &UdpTrackerConfig, tracker: Arc) -> JoinHandle<()> { + let udp_server = UdpServer::new_ipv6(tracker).await.unwrap_or_else(|e| { + panic!("Could not start UDP server (IPv6): {}", e); + }); + + info!("Starting UDP server on: {}", config.bind_address); + tokio::spawn(async move { + udp_server.start().await; + }) +} diff --git a/src/torrust_udp_tracker/server.rs b/src/torrust_udp_tracker/server.rs index 0da4ce140..86cc9727f 100644 --- a/src/torrust_udp_tracker/server.rs +++ b/src/torrust_udp_tracker/server.rs @@ -22,6 +22,15 @@ impl UdpServer { }) } + pub async fn new_ipv6(tracker: Arc) -> Result { + let srv = UdpSocket::bind(&tracker.config.udp_tracker_ipv6.bind_address).await?; + + Ok(UdpServer { + socket: srv, + tracker, + }) + } + pub async fn start(&self) { loop { let mut data = [0; MAX_PACKET_SIZE]; diff --git a/src/tracker.rs b/src/tracker.rs index 67acd7583..c53367f95 100644 --- a/src/tracker.rs +++ b/src/tracker.rs @@ -321,6 +321,16 @@ impl TorrentTracker { Ok(()) } + // Loading the torrents into memory + pub async fn load_torrents(&self, tracker: Arc) -> Result { + self.database.load_persistent_torrent_data(tracker).await + } + + // Saving the torrents from memory + pub async fn save_torrents(&self, tracker: Arc) -> Result { + self.database.save_persistent_torrent_data(tracker).await + } + // Adding torrents is not relevant to public trackers. pub async fn add_torrent_to_whitelist(&self, info_hash: &InfoHash) -> Result { self.database.add_info_hash_to_whitelist(info_hash.clone()).await @@ -378,6 +388,22 @@ impl TorrentTracker { } } + pub async fn add_torrent(&self, info_hash: &InfoHash, seeders: u32, completed: u32, leechers: u32) -> TorrentStats { + let mut torrents = self.torrents.write().await; + + if !torrents.contains_key(&info_hash) { + let mut torrent_entry = TorrentEntry::new(); + torrent_entry.completed = completed; + torrents.insert(info_hash.clone(), torrent_entry); + } + + TorrentStats { + seeders, + completed, + leechers, + } + } + pub async fn get_torrents(&self) -> tokio::sync::RwLockReadGuard<'_, BTreeMap> { self.torrents.read().await } @@ -413,7 +439,7 @@ impl TorrentTracker { } } - if self.config.mode.clone() == TrackerMode::PublicMode { + if self.config.mode.clone() == TrackerMode::PublicMode && self.config.cleanup_peerless && !self.config.persistence { // peer-less torrents.. if torrent_entry.peers.len() == 0 { torrents_to_remove.push(k.clone());