aquatic_udp: first access list implementation

aquatic_udp/src/lib/common.rs

@@ -6,7 +6,7 @@ use hashbrown::HashMap;
 use indexmap::IndexMap;
 use parking_lot::Mutex;
 
-pub use aquatic_common::ValidUntil;
+pub use aquatic_common::{AccessList, ValidUntil};
 pub use aquatic_udp_protocol::*;
 
 pub const MAX_PACKET_SIZE: usize = 4096;
@@ -121,6 +121,7 @@ pub struct Statistics {
 pub struct State {
     pub connections: Arc<Mutex<ConnectionMap>>,
     pub torrents: Arc<Mutex<TorrentMaps>>,
+    pub access_list: Arc<Mutex<AccessList>>,
     pub statistics: Arc<Statistics>,
 }
 
@@ -129,6 +130,7 @@ impl Default for State {
         Self {
             connections: Arc::new(Mutex::new(HashMap::new())),
             torrents: Arc::new(Mutex::new(TorrentMaps::default())),
+            access_list: Arc::new(Mutex::new(AccessList::default())),
             statistics: Arc::new(Statistics::default()),
         }
     }

aquatic_udp/src/lib/config.rs

@@ -1,5 +1,6 @@
 use std::net::SocketAddr;
 
+use aquatic_common::AccessListConfig;
 use serde::{Deserialize, Serialize};
 
 use aquatic_cli_helpers::LogLevel;
@@ -21,6 +22,7 @@ pub struct Config {
     pub statistics: StatisticsConfig,
     pub cleaning: CleaningConfig,
     pub privileges: PrivilegeConfig,
+    pub access_list: AccessListConfig,
 }
 
 impl aquatic_cli_helpers::Config for Config {
@@ -113,6 +115,7 @@ impl Default for Config {
             statistics: StatisticsConfig::default(),
             cleaning: CleaningConfig::default(),
             privileges: PrivilegeConfig::default(),
+            access_list: AccessListConfig::default(),
         }
     }
 }

aquatic_udp/src/lib/handlers.rs

@@ -79,6 +79,8 @@ pub fn run_request_worker(
             &mut responses,
         );
 
+        // Check announce and scrape requests for valid connection
+
         announce_requests.retain(|(request, src)| {
             let connection_key = ConnectionKey {
                 connection_id: request.connection_id,
@@ -95,7 +97,7 @@ pub fn run_request_worker(
 
                 responses.push((response.into(), *src));
 
-                false
+                return false;
             }
         });
 
@@ -121,6 +123,46 @@ pub fn run_request_worker(
 
         ::std::mem::drop(connections);
 
+        // Check announce requests for allowed info hash
+
+        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) {
+                        let response = ErrorResponse {
+                            transaction_id: request.transaction_id,
+                            message: "Forbidden info hash".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);
+
+        // Handle announce and scrape requests
+
         if !(announce_requests.is_empty() && scrape_requests.is_empty()) {
             let mut torrents = state.torrents.lock();
 

aquatic_udp/src/lib/lib.rs

@@ -31,8 +31,8 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
 
             if sockets == config.socket_workers {
                 PrivDrop::default()
-                    .chroot(config.privileges.chroot_path)
-                    .user(config.privileges.user)
+                    .chroot(config.privileges.chroot_path.clone())
+                    .user(config.privileges.user.clone())
                     .apply()?;
 
                 break;
@@ -45,7 +45,7 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
     loop {
         ::std::thread::sleep(Duration::from_secs(config.cleaning.interval));
 
-        tasks::clean_connections_and_torrents(&state);
+        tasks::clean_connections_and_torrents(&config, &state);
     }
 }
 

aquatic_udp/src/lib/tasks.rs

@@ -3,10 +3,12 @@ use std::time::Instant;
 
 use histogram::Histogram;
 
+use aquatic_common::AccessListType;
+
 use crate::common::*;
 use crate::config::Config;
 
-pub fn clean_connections_and_torrents(state: &State) {
+pub fn clean_connections_and_torrents(config: &Config, state: &State) {
     let now = Instant::now();
 
     {
@@ -16,14 +18,47 @@ pub fn clean_connections_and_torrents(state: &State) {
         connections.shrink_to_fit();
     }
 
-    let mut torrents = state.torrents.lock();
+    match config.access_list.list_type {
+        AccessListType::Allow => {
+            let mut access_list = state.access_list.lock();
 
-    clean_torrent_map(&mut torrents.ipv4, now);
-    clean_torrent_map(&mut torrents.ipv6, now);
+            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);
+        },
+        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);
+        },
+        AccessListType::Ignore => {
+            let mut torrents = state.torrents.lock();
+
+            clean_torrent_map(&mut torrents.ipv4, now);
+            clean_torrent_map(&mut torrents.ipv6, now);
+        }
+    }
 }
 
 #[inline]
-fn clean_torrent_map<I: Ip>(torrents: &mut TorrentMap<I>, now: Instant) {
+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;