From b5a2b81bd72c909773cde9c7f5258492a360b2e8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= Date: Fri, 15 Oct 2021 02:30:49 +0200 Subject: [PATCH] aquatic_udp: simplify access list logic --- aquatic_common/src/lib.rs | 16 +++++-- aquatic_udp/src/lib/handlers.rs | 40 ++++++----------- aquatic_udp/src/lib/tasks.rs | 78 ++++++++++++++------------------- 3 files changed, 59 insertions(+), 75 deletions(-) diff --git a/aquatic_common/src/lib.rs b/aquatic_common/src/lib.rs index 1b97684..ae39748 100644 --- a/aquatic_common/src/lib.rs +++ b/aquatic_common/src/lib.rs @@ -9,7 +9,7 @@ use indexmap::IndexMap; use rand::Rng; use serde::{Deserialize, Serialize}; -#[derive(Clone, Debug, Serialize, Deserialize)] +#[derive(Clone, Copy, Debug, Serialize, Deserialize)] pub enum AccessListType { Allow, Deny, @@ -65,8 +65,18 @@ impl AccessList { Ok(()) } - pub fn contains(&self, info_hash_bytes: &[u8; 20]) -> bool { - self.0.contains(info_hash_bytes) + pub fn allows(&self, list_type: AccessListType, info_hash_bytes: &[u8; 20]) -> bool { + match list_type { + AccessListType::Allow => { + self.0.contains(info_hash_bytes) + } + AccessListType::Deny => { + !self.0.contains(info_hash_bytes) + } + AccessListType::Ignore => { + true + } + } } } diff --git a/aquatic_udp/src/lib/handlers.rs b/aquatic_udp/src/lib/handlers.rs index d2e93aa..d6cf58a 100644 --- a/aquatic_udp/src/lib/handlers.rs +++ b/aquatic_udp/src/lib/handlers.rs @@ -9,7 +9,7 @@ use rand::{ Rng, SeedableRng, }; -use aquatic_common::{convert_ipv4_mapped_ipv6, extract_response_peers}; +use aquatic_common::{AccessListType, convert_ipv4_mapped_ipv6, extract_response_peers}; use aquatic_udp_protocol::*; use crate::common::*; @@ -123,43 +123,29 @@ pub fn run_request_worker( ::std::mem::drop(connections); - // Check announce requests for allowed info hash + // Check announce requests for allowed info hashes - let access_list: MutexGuard = state.access_list.lock(); + match config.access_list.list_type { + access_list_type@(AccessListType::Allow | AccessListType::Deny) => { + let access_list: MutexGuard = state.access_list.lock(); - announce_requests.retain(|(request, src)| { - match config.access_list.list_type { - aquatic_common::AccessListType::Allow => { - if !access_list.contains(&request.info_hash.0) { + announce_requests.retain(|(request, src)| { + if !access_list.allows(access_list_type, &request.info_hash.0) { let response = ErrorResponse { transaction_id: request.transaction_id, - message: "Forbidden info hash".to_string(), + message: "Info hash not allowed".to_string(), }; responses.push((response.into(), *src)); return false; } - }, - aquatic_common::AccessListType::Deny => { - if access_list.contains(&request.info_hash.0) { - let response = ErrorResponse { - transaction_id: request.transaction_id, - message: "Forbidden info hash".to_string(), - }; - responses.push((response.into(), *src)); - - return false; - } - }, - aquatic_common::AccessListType::Ignore => {}, - } - - true - }); - - ::std::mem::drop(access_list); + true + }); + }, + AccessListType::Ignore => {}, + }; // Handle announce and scrape requests diff --git a/aquatic_udp/src/lib/tasks.rs b/aquatic_udp/src/lib/tasks.rs index 795d033..43e1711 100644 --- a/aquatic_udp/src/lib/tasks.rs +++ b/aquatic_udp/src/lib/tasks.rs @@ -19,72 +19,60 @@ pub fn clean_connections_and_torrents(config: &Config, state: &State) { } match config.access_list.list_type { - AccessListType::Allow => { + access_list_type@(AccessListType::Allow | AccessListType::Deny) => { let mut access_list = state.access_list.lock(); access_list.update_from_path(&config.access_list.path); let mut torrents = state.torrents.lock(); - torrents.ipv4.retain(|info_hash, _| access_list.contains(&info_hash.0)); - clean_torrent_map(&mut torrents.ipv4, now); + torrents.ipv4.retain(|info_hash, torrent| { + access_list.allows(access_list_type, &info_hash.0) && clean_torrent_and_peers(now, torrent) + }); + torrents.ipv4.shrink_to_fit(); - torrents.ipv6.retain(|info_hash, _| access_list.contains(&info_hash.0)); - clean_torrent_map(&mut torrents.ipv6, now); - }, - AccessListType::Deny => { - let mut access_list = state.access_list.lock(); - - access_list.update_from_path(&config.access_list.path); - - let mut torrents = state.torrents.lock(); - - torrents.ipv4.retain(|info_hash, _| !access_list.contains(&info_hash.0)); - clean_torrent_map(&mut torrents.ipv4, now); - - torrents.ipv6.retain(|info_hash, _| !access_list.contains(&info_hash.0)); - clean_torrent_map(&mut torrents.ipv6, now); + torrents.ipv6.retain(|info_hash, torrent| { + access_list.allows(access_list_type, &info_hash.0) && clean_torrent_and_peers(now, torrent) + }); + torrents.ipv6.shrink_to_fit(); }, AccessListType::Ignore => { let mut torrents = state.torrents.lock(); - clean_torrent_map(&mut torrents.ipv4, now); - clean_torrent_map(&mut torrents.ipv6, now); + torrents.ipv4.retain(|_, torrent| clean_torrent_and_peers(now, torrent)); + torrents.ipv4.shrink_to_fit(); + + torrents.ipv6.retain(|_, torrent| clean_torrent_and_peers(now, torrent)); + torrents.ipv6.shrink_to_fit(); } } } +/// Returns true if torrent is to be kept #[inline] -fn clean_torrent_map( - torrents: &mut TorrentMap, - now: Instant -) { - torrents.retain(|_, torrent| { - let num_seeders = &mut torrent.num_seeders; - let num_leechers = &mut torrent.num_leechers; +fn clean_torrent_and_peers(now: Instant, torrent: &mut TorrentData) -> bool { + let num_seeders = &mut torrent.num_seeders; + let num_leechers = &mut torrent.num_leechers; - torrent.peers.retain(|_, peer| { - let keep = peer.valid_until.0 > now; + torrent.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; - } - _ => (), - }; - } + if !keep { + match peer.status { + PeerStatus::Seeding => { + *num_seeders -= 1; + } + PeerStatus::Leeching => { + *num_leechers -= 1; + } + _ => (), + }; + } - keep - }); - - !torrent.peers.is_empty() + keep }); - torrents.shrink_to_fit(); + !torrent.peers.is_empty() } pub fn gather_and_print_statistics(state: &State, config: &Config) {