YGGverse/aquatic
1
0
Fork 0
mirror of https://github.com/YGGverse/aquatic.git synced 2026-04-01 02:05:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

aquatic_udp: simplify access list logic

Browse source
This commit is contained in:
Joakim Frostegård 2021-10-15 02:30:49 +02:00
parent 8639f380f4
commit b5a2b81bd7
3 changed files with 59 additions and 75 deletions
Download patch file Download diff file Expand all files Collapse all files

16
aquatic_common/src/lib.rs
View file

@ -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
}
}
}
}

40
aquatic_udp/src/lib/handlers.rs
View file

@ -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<AccessList> = state.access_list.lock();
match config.access_list.list_type {
access_list_type@(AccessListType::Allow | AccessListType::Deny) => {
let access_list: MutexGuard<AccessList> = 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

78
aquatic_udp/src/lib/tasks.rs
View file

@ -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<I: Ip>(
torrents: &mut TorrentMap<I>,
now: Instant
) {
torrents.retain(|_, torrent| {
let num_seeders = &mut torrent.num_seeders;
let num_leechers = &mut torrent.num_leechers;
fn clean_torrent_and_peers<I: Ip>(now: Instant, torrent: &mut TorrentData<I>) -> 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) {