aquatic_udp: start work on announce handler in glommio version

aquatic_udp/src/lib/common/announce.rs

@@ -0,0 +1,178 @@
+use rand::rngs::SmallRng;
+
+use aquatic_common::extract_response_peers;
+
+use crate::common::*;
+
+pub fn handle_announce_request<I: Ip>(
+    config: &Config,
+    rng: &mut SmallRng,
+    torrents: &mut TorrentMap<I>,
+    request: AnnounceRequest,
+    peer_ip: I,
+    peer_valid_until: ValidUntil,
+) -> AnnounceResponse {
+    let peer_key = PeerMapKey {
+        ip: peer_ip,
+        peer_id: request.peer_id,
+    };
+
+    let peer_status = PeerStatus::from_event_and_bytes_left(request.event, request.bytes_left);
+
+    let peer = Peer {
+        ip_address: peer_ip,
+        port: request.port,
+        status: peer_status,
+        valid_until: peer_valid_until,
+    };
+
+    let torrent_data = torrents.entry(request.info_hash).or_default();
+
+    let opt_removed_peer = match peer_status {
+        PeerStatus::Leeching => {
+            torrent_data.num_leechers += 1;
+
+            torrent_data.peers.insert(peer_key, peer)
+        }
+        PeerStatus::Seeding => {
+            torrent_data.num_seeders += 1;
+
+            torrent_data.peers.insert(peer_key, peer)
+        }
+        PeerStatus::Stopped => torrent_data.peers.remove(&peer_key),
+    };
+
+    match opt_removed_peer.map(|peer| peer.status) {
+        Some(PeerStatus::Leeching) => {
+            torrent_data.num_leechers -= 1;
+        }
+        Some(PeerStatus::Seeding) => {
+            torrent_data.num_seeders -= 1;
+        }
+        _ => {}
+    }
+
+    let max_num_peers_to_take = calc_max_num_peers_to_take(config, request.peers_wanted.0);
+
+    let response_peers = extract_response_peers(
+        rng,
+        &torrent_data.peers,
+        max_num_peers_to_take,
+        peer_key,
+        Peer::to_response_peer,
+    );
+
+    AnnounceResponse {
+        transaction_id: request.transaction_id,
+        announce_interval: AnnounceInterval(config.protocol.peer_announce_interval),
+        leechers: NumberOfPeers(torrent_data.num_leechers as i32),
+        seeders: NumberOfPeers(torrent_data.num_seeders as i32),
+        peers: response_peers,
+    }
+}
+
+#[inline]
+fn calc_max_num_peers_to_take(config: &Config, peers_wanted: i32) -> usize {
+    if peers_wanted <= 0 {
+        config.protocol.max_response_peers as usize
+    } else {
+        ::std::cmp::min(
+            config.protocol.max_response_peers as usize,
+            peers_wanted as usize,
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashSet;
+    use std::net::Ipv4Addr;
+
+    use indexmap::IndexMap;
+    use quickcheck::{quickcheck, TestResult};
+    use rand::thread_rng;
+
+    use super::*;
+
+    fn gen_peer_map_key_and_value(i: u32) -> (PeerMapKey<Ipv4Addr>, Peer<Ipv4Addr>) {
+        let ip_address = Ipv4Addr::from(i.to_be_bytes());
+        let peer_id = PeerId([0; 20]);
+
+        let key = PeerMapKey {
+            ip: ip_address,
+            peer_id,
+        };
+        let value = Peer {
+            ip_address,
+            port: Port(1),
+            status: PeerStatus::Leeching,
+            valid_until: ValidUntil::new(0),
+        };
+
+        (key, value)
+    }
+
+    #[test]
+    fn test_extract_response_peers() {
+        fn prop(data: (u16, u16)) -> TestResult {
+            let gen_num_peers = data.0 as u32;
+            let req_num_peers = data.1 as usize;
+
+            let mut peer_map: PeerMap<Ipv4Addr> = IndexMap::with_capacity(gen_num_peers as usize);
+
+            let mut opt_sender_key = None;
+            let mut opt_sender_peer = None;
+
+            for i in 0..gen_num_peers {
+                let (key, value) = gen_peer_map_key_and_value((i << 16) + i);
+
+                if i == 0 {
+                    opt_sender_key = Some(key);
+                    opt_sender_peer = Some(value.to_response_peer());
+                }
+
+                peer_map.insert(key, value);
+            }
+
+            let mut rng = thread_rng();
+
+            let peers = extract_response_peers(
+                &mut rng,
+                &peer_map,
+                req_num_peers,
+                opt_sender_key.unwrap_or_else(|| gen_peer_map_key_and_value(1).0),
+                Peer::to_response_peer,
+            );
+
+            // Check that number of returned peers is correct
+
+            let mut success = peers.len() <= req_num_peers;
+
+            if req_num_peers >= gen_num_peers as usize {
+                success &= peers.len() == gen_num_peers as usize
+                    || peers.len() + 1 == gen_num_peers as usize;
+            }
+
+            // Check that returned peers are unique (no overlap) and that sender
+            // isn't returned
+
+            let mut ip_addresses = HashSet::with_capacity(peers.len());
+
+            for peer in peers {
+                if peer == opt_sender_peer.clone().unwrap()
+                    || ip_addresses.contains(&peer.ip_address)
+                {
+                    success = false;
+
+                    break;
+                }
+
+                ip_addresses.insert(peer.ip_address);
+            }
+
+            TestResult::from_bool(success)
+        }
+
+        quickcheck(prop as fn((u16, u16)) -> TestResult);
+    }
+}

aquatic_udp/src/lib/common/mod.rs

@@ -0,0 +1,248 @@
+use std::hash::Hash;
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
+use std::sync::{atomic::AtomicUsize, Arc};
+use std::time::Instant;
+
+use hashbrown::HashMap;
+use indexmap::IndexMap;
+use parking_lot::Mutex;
+
+pub use aquatic_common::{access_list::AccessList, ValidUntil};
+pub use aquatic_udp_protocol::*;
+
+use crate::config::Config;
+
+pub mod announce;
+pub mod network;
+
+pub const MAX_PACKET_SIZE: usize = 4096;
+
+pub trait Ip: Hash + PartialEq + Eq + Clone + Copy {
+    fn ip_addr(self) -> IpAddr;
+}
+
+impl Ip for Ipv4Addr {
+    fn ip_addr(self) -> IpAddr {
+        IpAddr::V4(self)
+    }
+}
+
+impl Ip for Ipv6Addr {
+    fn ip_addr(self) -> IpAddr {
+        IpAddr::V6(self)
+    }
+}
+
+pub enum ConnectedRequest {
+    Announce(AnnounceRequest),
+    Scrape(ScrapeRequest),
+}
+
+pub enum ConnectedResponse {
+    Announce(AnnounceResponse),
+    Scrape(ScrapeResponse),
+}
+
+impl Into<Response> for ConnectedResponse {
+    fn into(self) -> Response {
+        match self {
+            Self::Announce(response) => Response::Announce(response),
+            Self::Scrape(response) => Response::Scrape(response),
+        }
+    }
+}
+
+#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
+pub enum PeerStatus {
+    Seeding,
+    Leeching,
+    Stopped,
+}
+
+impl PeerStatus {
+    /// Determine peer status from announce event and number of bytes left.
+    ///
+    /// Likely, the last branch will be taken most of the time.
+    #[inline]
+    pub fn from_event_and_bytes_left(event: AnnounceEvent, bytes_left: NumberOfBytes) -> Self {
+        if event == AnnounceEvent::Stopped {
+            Self::Stopped
+        } else if bytes_left.0 == 0 {
+            Self::Seeding
+        } else {
+            Self::Leeching
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct Peer<I: Ip> {
+    pub ip_address: I,
+    pub port: Port,
+    pub status: PeerStatus,
+    pub valid_until: ValidUntil,
+}
+
+impl<I: Ip> Peer<I> {
+    #[inline(always)]
+    pub fn to_response_peer(&self) -> ResponsePeer {
+        ResponsePeer {
+            ip_address: self.ip_address.ip_addr(),
+            port: self.port,
+        }
+    }
+}
+
+#[derive(PartialEq, Eq, Hash, Clone, Copy)]
+pub struct PeerMapKey<I: Ip> {
+    pub ip: I,
+    pub peer_id: PeerId,
+}
+
+pub type PeerMap<I> = IndexMap<PeerMapKey<I>, Peer<I>>;
+
+pub struct TorrentData<I: Ip> {
+    pub peers: PeerMap<I>,
+    pub num_seeders: usize,
+    pub num_leechers: usize,
+}
+
+impl<I: Ip> Default for TorrentData<I> {
+    fn default() -> Self {
+        Self {
+            peers: IndexMap::new(),
+            num_seeders: 0,
+            num_leechers: 0,
+        }
+    }
+}
+
+pub type TorrentMap<I> = HashMap<InfoHash, TorrentData<I>>;
+
+#[derive(Default)]
+pub struct TorrentMaps {
+    pub ipv4: TorrentMap<Ipv4Addr>,
+    pub ipv6: TorrentMap<Ipv6Addr>,
+}
+
+impl TorrentMaps {
+    /// Remove disallowed and inactive torrents
+    pub fn clean(&mut self, config: &Config, access_list: &Arc<AccessList>) {
+        let now = Instant::now();
+
+        let access_list_mode = config.access_list.mode;
+
+        self.ipv4.retain(|info_hash, torrent| {
+            access_list.allows(access_list_mode, &info_hash.0)
+                && Self::clean_torrent_and_peers(now, torrent)
+        });
+        self.ipv4.shrink_to_fit();
+
+        self.ipv6.retain(|info_hash, torrent| {
+            access_list.allows(access_list_mode, &info_hash.0)
+                && Self::clean_torrent_and_peers(now, torrent)
+        });
+        self.ipv6.shrink_to_fit();
+    }
+
+    /// Returns true if torrent is to be kept
+    #[inline]
+    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;
+
+            if !keep {
+                match peer.status {
+                    PeerStatus::Seeding => {
+                        *num_seeders -= 1;
+                    }
+                    PeerStatus::Leeching => {
+                        *num_leechers -= 1;
+                    }
+                    _ => (),
+                };
+            }
+
+            keep
+        });
+
+        !torrent.peers.is_empty()
+    }
+}
+
+#[derive(Default)]
+pub struct Statistics {
+    pub requests_received: AtomicUsize,
+    pub responses_sent: AtomicUsize,
+    pub bytes_received: AtomicUsize,
+    pub bytes_sent: AtomicUsize,
+}
+
+#[derive(Clone)]
+pub struct State {
+    pub access_list: Arc<AccessList>,
+    pub torrents: Arc<Mutex<TorrentMaps>>,
+    pub statistics: Arc<Statistics>,
+}
+
+impl Default for State {
+    fn default() -> Self {
+        Self {
+            access_list: Arc::new(AccessList::default()),
+            torrents: Arc::new(Mutex::new(TorrentMaps::default())),
+            statistics: Arc::new(Statistics::default()),
+        }
+    }
+}
+
+#[derive(Default)]
+pub struct ConnectionMap(HashMap<(ConnectionId, SocketAddr), ValidUntil>);
+
+impl ConnectionMap {
+    pub fn insert(
+        &mut self,
+        connection_id: ConnectionId,
+        socket_addr: SocketAddr,
+        valid_until: ValidUntil,
+    ) {
+        self.0.insert((connection_id, socket_addr), valid_until);
+    }
+
+    pub fn contains(&mut self, connection_id: ConnectionId, socket_addr: SocketAddr) -> bool {
+        self.0.contains_key(&(connection_id, socket_addr))
+    }
+
+    pub fn clean(&mut self) {
+        let now = Instant::now();
+
+        self.0.retain(|_, v| v.0 > now);
+        self.0.shrink_to_fit();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_peer_status_from_event_and_bytes_left() {
+        use crate::common::*;
+
+        use PeerStatus::*;
+
+        let f = PeerStatus::from_event_and_bytes_left;
+
+        assert_eq!(Stopped, f(AnnounceEvent::Stopped, NumberOfBytes(0)));
+        assert_eq!(Stopped, f(AnnounceEvent::Stopped, NumberOfBytes(1)));
+
+        assert_eq!(Seeding, f(AnnounceEvent::Started, NumberOfBytes(0)));
+        assert_eq!(Leeching, f(AnnounceEvent::Started, NumberOfBytes(1)));
+
+        assert_eq!(Seeding, f(AnnounceEvent::Completed, NumberOfBytes(0)));
+        assert_eq!(Leeching, f(AnnounceEvent::Completed, NumberOfBytes(1)));
+
+        assert_eq!(Seeding, f(AnnounceEvent::None, NumberOfBytes(0)));
+        assert_eq!(Leeching, f(AnnounceEvent::None, NumberOfBytes(1)));
+    }
+}