diff --git a/aquatic_ws/src/lib/common.rs b/aquatic_ws/src/lib/common.rs index b524390..074a415 100644 --- a/aquatic_ws/src/lib/common.rs +++ b/aquatic_ws/src/lib/common.rs @@ -1,6 +1,8 @@ use std::net::{IpAddr, SocketAddr}; use std::sync::Arc; +use std::time::Instant; +use aquatic_common::access_list::AccessList; use crossbeam_channel::{Receiver, Sender}; use hashbrown::HashMap; use indexmap::IndexMap; @@ -12,6 +14,8 @@ pub use aquatic_common::ValidUntil; use aquatic_ws_protocol::*; +use crate::config::Config; + pub const LISTENER_TOKEN: Token = Token(0); pub const CHANNEL_TOKEN: Token = Token(1); @@ -82,6 +86,54 @@ pub type TorrentMap = HashMap; pub struct TorrentMaps { pub ipv4: TorrentMap, pub ipv6: TorrentMap, + pub access_list: AccessList, +} + +impl TorrentMaps { + pub fn clean(&mut self, config: &Config) { + Self::clean_torrent_map(config, &self.access_list, &mut self.ipv4); + Self::clean_torrent_map(config, &self.access_list, &mut self.ipv6); + } + + fn clean_torrent_map( + config: &Config, + access_list: &AccessList, + torrent_map: &mut TorrentMap + ) { + let now = Instant::now(); + + torrent_map.retain(|info_hash, torrent_data| { + if !access_list.allows(config.access_list.mode, &info_hash.0) { + return false; + } + + let num_seeders = &mut torrent_data.num_seeders; + let num_leechers = &mut torrent_data.num_leechers; + + torrent_data.peers.retain(|_, peer| { + let keep = peer.valid_until.0 >= now; + + if !keep { + match peer.status { + PeerStatus::Seeding => { + *num_seeders -= 1; + } + PeerStatus::Leeching => { + *num_leechers -= 1; + } + _ => (), + }; + } + + keep + }); + + !torrent_data.peers.is_empty() + }); + + torrent_map.shrink_to_fit(); + } + } #[derive(Clone)] diff --git a/aquatic_ws/src/lib/config.rs b/aquatic_ws/src/lib/config.rs index 1be0c47..b7ef10c 100644 --- a/aquatic_ws/src/lib/config.rs +++ b/aquatic_ws/src/lib/config.rs @@ -1,5 +1,6 @@ use std::net::SocketAddr; +use aquatic_common::access_list::AccessListConfig; use serde::{Deserialize, Serialize}; use aquatic_cli_helpers::LogLevel; @@ -21,6 +22,7 @@ pub struct Config { pub cleaning: CleaningConfig, pub statistics: StatisticsConfig, pub privileges: PrivilegeConfig, + pub access_list: AccessListConfig, } impl aquatic_cli_helpers::Config for Config { @@ -105,6 +107,7 @@ impl Default for Config { cleaning: CleaningConfig::default(), statistics: StatisticsConfig::default(), privileges: PrivilegeConfig::default(), + access_list: AccessListConfig::default(), } } } diff --git a/aquatic_ws/src/lib/handler.rs b/aquatic_ws/src/lib/handler.rs index 1f61a17..e140a74 100644 --- a/aquatic_ws/src/lib/handler.rs +++ b/aquatic_ws/src/lib/handler.rs @@ -99,6 +99,19 @@ pub fn handle_announce_requests( let valid_until = ValidUntil::new(config.cleaning.max_peer_age); for (request_sender_meta, request) in requests { + let info_hash_allowed = torrent_maps + .access_list + .allows(config.access_list.mode, &request.info_hash.0); + + if !info_hash_allowed { + // let response = OutMessage::ErrorResponse(); + + // out_message_sender.send(request_sender_meta, response); + // wake_socket_workers[request_sender_meta.worker_index] = true; + + continue; + } + let torrent_data: &mut TorrentData = if request_sender_meta.converted_peer_ip.is_ipv4() { torrent_maps.ipv4.entry(request.info_hash).or_default() } else { diff --git a/aquatic_ws/src/lib/lib.rs b/aquatic_ws/src/lib/lib.rs index e27c21b..c02e2da 100644 --- a/aquatic_ws/src/lib/lib.rs +++ b/aquatic_ws/src/lib/lib.rs @@ -24,12 +24,18 @@ pub const APP_NAME: &str = "aquatic_ws: WebTorrent tracker"; pub fn run(config: Config) -> anyhow::Result<()> { let state = State::default(); + tasks::update_access_list(&config, &mut state.torrent_maps.lock()); + start_workers(config.clone(), state.clone())?; loop { ::std::thread::sleep(Duration::from_secs(config.cleaning.interval)); - tasks::clean_torrents(&state); + let mut torrent_maps = state.torrent_maps.lock(); + + tasks::update_access_list(&config, &mut torrent_maps); + + torrent_maps.clean(&config); } } diff --git a/aquatic_ws/src/lib/tasks.rs b/aquatic_ws/src/lib/tasks.rs index 64b5338..18664fe 100644 --- a/aquatic_ws/src/lib/tasks.rs +++ b/aquatic_ws/src/lib/tasks.rs @@ -1,45 +1,23 @@ use std::time::Instant; +use aquatic_common::access_list::AccessListMode; use histogram::Histogram; use crate::common::*; +use crate::config::Config; -pub fn clean_torrents(state: &State) { - fn clean_torrent_map(torrent_map: &mut TorrentMap) { - let now = Instant::now(); - - torrent_map.retain(|_, torrent_data| { - let num_seeders = &mut torrent_data.num_seeders; - let num_leechers = &mut torrent_data.num_leechers; - - torrent_data.peers.retain(|_, peer| { - let keep = peer.valid_until.0 >= now; - - if !keep { - match peer.status { - PeerStatus::Seeding => { - *num_seeders -= 1; - } - PeerStatus::Leeching => { - *num_leechers -= 1; - } - _ => (), - }; - } - - keep - }); - - !torrent_data.peers.is_empty() - }); - - torrent_map.shrink_to_fit(); +pub fn update_access_list(config: &Config, torrent_maps: &mut TorrentMaps) { + match config.access_list.mode { + AccessListMode::Require | AccessListMode::Forbid => { + if let Err(err) = torrent_maps + .access_list + .update_from_path(&config.access_list.path) + { + ::log::error!("Couldn't update access list: {:?}", err); + } + } + AccessListMode::Ignore => {} } - - let mut torrent_maps = state.torrent_maps.lock(); - - clean_torrent_map(&mut torrent_maps.ipv4); - clean_torrent_map(&mut torrent_maps.ipv6); } pub fn print_statistics(state: &State) {