aquatic_udp: split some code into mio and glommio versions

aquatic_udp/src/lib/mio/handlers/announce.rs

@@ -0,0 +1,220 @@
+use std::net::{IpAddr, SocketAddr};
+use std::vec::Drain;
+
+use parking_lot::MutexGuard;
+use rand::rngs::SmallRng;
+
+use aquatic_common::{convert_ipv4_mapped_ipv6, extract_response_peers};
+use aquatic_udp_protocol::*;
+
+use crate::common::*;
+use crate::config::Config;
+
+#[inline]
+pub fn handle_announce_requests(
+    config: &Config,
+    torrents: &mut MutexGuard<TorrentMaps>,
+    rng: &mut SmallRng,
+    requests: Drain<(AnnounceRequest, SocketAddr)>,
+    responses: &mut Vec<(ConnectedResponse, SocketAddr)>,
+) {
+    let peer_valid_until = ValidUntil::new(config.cleaning.max_peer_age);
+
+    responses.extend(requests.map(|(request, src)| {
+        let peer_ip = convert_ipv4_mapped_ipv6(src.ip());
+
+        let response = match peer_ip {
+            IpAddr::V4(ip) => handle_announce_request(
+                config,
+                rng,
+                &mut torrents.ipv4,
+                request,
+                ip,
+                peer_valid_until,
+            ),
+            IpAddr::V6(ip) => handle_announce_request(
+                config,
+                rng,
+                &mut torrents.ipv6,
+                request,
+                ip,
+                peer_valid_until,
+            ),
+        };
+
+        (ConnectedResponse::Announce(response), src)
+    }));
+}
+
+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/mio/handlers/mod.rs

@@ -0,0 +1,88 @@
+use std::net::SocketAddr;
+use std::time::Duration;
+
+use crossbeam_channel::{Receiver, Sender};
+use rand::{rngs::SmallRng, SeedableRng};
+
+use aquatic_udp_protocol::*;
+
+use crate::common::*;
+use crate::config::Config;
+
+mod announce;
+mod scrape;
+
+use announce::handle_announce_requests;
+use scrape::handle_scrape_requests;
+
+pub fn run_request_worker(
+    state: State,
+    config: Config,
+    request_receiver: Receiver<(ConnectedRequest, SocketAddr)>,
+    response_sender: Sender<(ConnectedResponse, SocketAddr)>,
+) {
+    let mut announce_requests: Vec<(AnnounceRequest, SocketAddr)> = Vec::new();
+    let mut scrape_requests: Vec<(ScrapeRequest, SocketAddr)> = Vec::new();
+    let mut responses: Vec<(ConnectedResponse, SocketAddr)> = Vec::new();
+
+    let mut small_rng = SmallRng::from_entropy();
+
+    let timeout = Duration::from_micros(config.handlers.channel_recv_timeout_microseconds);
+
+    loop {
+        let mut opt_torrents = None;
+
+        // Collect requests from channel, divide them by type
+        //
+        // Collect a maximum number of request. Stop collecting before that
+        // number is reached if having waited for too long for a request, but
+        // only if TorrentMaps mutex isn't locked.
+        for i in 0..config.handlers.max_requests_per_iter {
+            let (request, src): (ConnectedRequest, SocketAddr) = if i == 0 {
+                match request_receiver.recv() {
+                    Ok(r) => r,
+                    Err(_) => break, // Really shouldn't happen
+                }
+            } else {
+                match request_receiver.recv_timeout(timeout) {
+                    Ok(r) => r,
+                    Err(_) => {
+                        if let Some(guard) = state.torrents.try_lock() {
+                            opt_torrents = Some(guard);
+
+                            break;
+                        } else {
+                            continue;
+                        }
+                    }
+                }
+            };
+
+            match request {
+                ConnectedRequest::Announce(r) => announce_requests.push((r, src)),
+                ConnectedRequest::Scrape(r) => scrape_requests.push((r, src)),
+            }
+        }
+
+        // Generate responses for announce and scrape requests, then drop MutexGuard.
+        {
+            let mut torrents = opt_torrents.unwrap_or_else(|| state.torrents.lock());
+
+            handle_announce_requests(
+                &config,
+                &mut torrents,
+                &mut small_rng,
+                announce_requests.drain(..),
+                &mut responses,
+            );
+
+            handle_scrape_requests(&mut torrents, scrape_requests.drain(..), &mut responses);
+        }
+
+        for r in responses.drain(..) {
+            if let Err(err) = response_sender.send(r) {
+                ::log::error!("error sending response to channel: {}", err);
+            }
+        }
+    }
+}

aquatic_udp/src/lib/mio/handlers/scrape.rs

@@ -0,0 +1,64 @@
+use std::net::SocketAddr;
+use std::vec::Drain;
+
+use parking_lot::MutexGuard;
+
+use aquatic_common::convert_ipv4_mapped_ipv6;
+use aquatic_udp_protocol::*;
+
+use crate::common::*;
+
+#[inline]
+pub fn handle_scrape_requests(
+    torrents: &mut MutexGuard<TorrentMaps>,
+    requests: Drain<(ScrapeRequest, SocketAddr)>,
+    responses: &mut Vec<(ConnectedResponse, SocketAddr)>,
+) {
+    let empty_stats = create_torrent_scrape_statistics(0, 0);
+
+    responses.extend(requests.map(|(request, src)| {
+        let mut stats: Vec<TorrentScrapeStatistics> = Vec::with_capacity(request.info_hashes.len());
+
+        let peer_ip = convert_ipv4_mapped_ipv6(src.ip());
+
+        if peer_ip.is_ipv4() {
+            for info_hash in request.info_hashes.iter() {
+                if let Some(torrent_data) = torrents.ipv4.get(info_hash) {
+                    stats.push(create_torrent_scrape_statistics(
+                        torrent_data.num_seeders as i32,
+                        torrent_data.num_leechers as i32,
+                    ));
+                } else {
+                    stats.push(empty_stats);
+                }
+            }
+        } else {
+            for info_hash in request.info_hashes.iter() {
+                if let Some(torrent_data) = torrents.ipv6.get(info_hash) {
+                    stats.push(create_torrent_scrape_statistics(
+                        torrent_data.num_seeders as i32,
+                        torrent_data.num_leechers as i32,
+                    ));
+                } else {
+                    stats.push(empty_stats);
+                }
+            }
+        }
+
+        let response = ConnectedResponse::Scrape(ScrapeResponse {
+            transaction_id: request.transaction_id,
+            torrent_stats: stats,
+        });
+
+        (response, src)
+    }));
+}
+
+#[inline(always)]
+fn create_torrent_scrape_statistics(seeders: i32, leechers: i32) -> TorrentScrapeStatistics {
+    TorrentScrapeStatistics {
+        seeders: NumberOfPeers(seeders),
+        completed: NumberOfDownloads(0), // No implementation planned
+        leechers: NumberOfPeers(leechers),
+    }
+}

aquatic_udp/src/lib/mio/mod.rs

@@ -0,0 +1,2 @@
+pub mod handlers;
+pub mod network;

aquatic_udp/src/lib/mio/network.rs

@@ -0,0 +1,312 @@
+use std::io::{Cursor, ErrorKind};
+use std::net::{IpAddr, SocketAddr};
+use std::sync::{
+    atomic::{AtomicUsize, Ordering},
+    Arc,
+};
+use std::time::{Duration, Instant};
+use std::vec::Drain;
+
+use crossbeam_channel::{Receiver, Sender};
+use mio::net::UdpSocket;
+use mio::{Events, Interest, Poll, Token};
+use rand::prelude::{Rng, SeedableRng, StdRng};
+use socket2::{Domain, Protocol, Socket, Type};
+
+use aquatic_udp_protocol::{IpVersion, Request, Response};
+
+use crate::common::*;
+use crate::config::Config;
+
+pub fn run_socket_worker(
+    state: State,
+    config: Config,
+    token_num: usize,
+    request_sender: Sender<(ConnectedRequest, SocketAddr)>,
+    response_receiver: Receiver<(ConnectedResponse, SocketAddr)>,
+    num_bound_sockets: Arc<AtomicUsize>,
+) {
+    let mut rng = StdRng::from_entropy();
+    let mut buffer = [0u8; MAX_PACKET_SIZE];
+
+    let mut socket = UdpSocket::from_std(create_socket(&config));
+    let mut poll = Poll::new().expect("create poll");
+
+    let interests = Interest::READABLE;
+
+    poll.registry()
+        .register(&mut socket, Token(token_num), interests)
+        .unwrap();
+
+    num_bound_sockets.fetch_add(1, Ordering::SeqCst);
+
+    let mut events = Events::with_capacity(config.network.poll_event_capacity);
+    let mut connections = ConnectionMap::default();
+
+    let mut local_responses: Vec<(Response, SocketAddr)> = Vec::new();
+
+    let timeout = Duration::from_millis(50);
+
+    let mut iter_counter = 0usize;
+
+    loop {
+        poll.poll(&mut events, Some(timeout))
+            .expect("failed polling");
+
+        for event in events.iter() {
+            let token = event.token();
+
+            if (token.0 == token_num) & event.is_readable() {
+                read_requests(
+                    &config,
+                    &state,
+                    &mut connections,
+                    &mut rng,
+                    &mut socket,
+                    &mut buffer,
+                    &request_sender,
+                    &mut local_responses,
+                );
+            }
+        }
+
+        send_responses(
+            &state,
+            &config,
+            &mut socket,
+            &mut buffer,
+            &response_receiver,
+            local_responses.drain(..),
+        );
+
+        iter_counter += 1;
+
+        if iter_counter == 1000 {
+            connections.clean();
+
+            iter_counter = 0;
+        }
+    }
+}
+
+fn create_socket(config: &Config) -> ::std::net::UdpSocket {
+    let socket = if config.network.address.is_ipv4() {
+        Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::UDP))
+    } else {
+        Socket::new(Domain::IPV6, Type::DGRAM, Some(Protocol::UDP))
+    }
+    .expect("create socket");
+
+    socket.set_reuse_port(true).expect("socket: set reuse port");
+
+    socket
+        .set_nonblocking(true)
+        .expect("socket: set nonblocking");
+
+    socket
+        .bind(&config.network.address.into())
+        .unwrap_or_else(|err| panic!("socket: bind to {}: {:?}", config.network.address, err));
+
+    let recv_buffer_size = config.network.socket_recv_buffer_size;
+
+    if recv_buffer_size != 0 {
+        if let Err(err) = socket.set_recv_buffer_size(recv_buffer_size) {
+            ::log::error!(
+                "socket: failed setting recv buffer to {}: {:?}",
+                recv_buffer_size,
+                err
+            );
+        }
+    }
+
+    socket.into()
+}
+
+#[inline]
+fn read_requests(
+    config: &Config,
+    state: &State,
+    connections: &mut ConnectionMap,
+    rng: &mut StdRng,
+    socket: &mut UdpSocket,
+    buffer: &mut [u8],
+    request_sender: &Sender<(ConnectedRequest, SocketAddr)>,
+    local_responses: &mut Vec<(Response, SocketAddr)>,
+) {
+    let mut requests_received: usize = 0;
+    let mut bytes_received: usize = 0;
+
+    let valid_until = ValidUntil::new(config.cleaning.max_connection_age);
+    let access_list_mode = config.access_list.mode;
+
+    loop {
+        match socket.recv_from(&mut buffer[..]) {
+            Ok((amt, src)) => {
+                let request =
+                    Request::from_bytes(&buffer[..amt], config.protocol.max_scrape_torrents);
+
+                bytes_received += amt;
+
+                if request.is_ok() {
+                    requests_received += 1;
+                }
+
+                match request {
+                    Ok(Request::Connect(request)) => {
+                        let connection_id = ConnectionId(rng.gen());
+
+                        connections.insert(connection_id, src, valid_until);
+
+                        let response = Response::Connect(ConnectResponse {
+                            connection_id,
+                            transaction_id: request.transaction_id,
+                        });
+
+                        local_responses.push((response, src))
+                    }
+                    Ok(Request::Announce(request)) => {
+                        if connections.contains(request.connection_id, src) {
+                            if state
+                                .access_list
+                                .allows(access_list_mode, &request.info_hash.0)
+                            {
+                                if let Err(err) = request_sender
+                                    .try_send((ConnectedRequest::Announce(request), src))
+                                {
+                                    ::log::warn!("request_sender.try_send failed: {:?}", err)
+                                }
+                            } else {
+                                let response = Response::Error(ErrorResponse {
+                                    transaction_id: request.transaction_id,
+                                    message: "Info hash not allowed".into(),
+                                });
+
+                                local_responses.push((response, src))
+                            }
+                        }
+                    }
+                    Ok(Request::Scrape(request)) => {
+                        if connections.contains(request.connection_id, src) {
+                            if let Err(err) =
+                                request_sender.try_send((ConnectedRequest::Scrape(request), src))
+                            {
+                                ::log::warn!("request_sender.try_send failed: {:?}", err)
+                            }
+                        }
+                    }
+                    Err(err) => {
+                        ::log::debug!("Request::from_bytes error: {:?}", err);
+
+                        if let RequestParseError::Sendable {
+                            connection_id,
+                            transaction_id,
+                            err,
+                        } = err
+                        {
+                            if connections.contains(connection_id, src) {
+                                let response = ErrorResponse {
+                                    transaction_id,
+                                    message: err.right_or("Parse error").into(),
+                                };
+
+                                local_responses.push((response.into(), src));
+                            }
+                        }
+                    }
+                }
+            }
+            Err(err) => {
+                if err.kind() == ErrorKind::WouldBlock {
+                    break;
+                }
+
+                ::log::info!("recv_from error: {}", err);
+            }
+        }
+    }
+
+    if config.statistics.interval != 0 {
+        state
+            .statistics
+            .requests_received
+            .fetch_add(requests_received, Ordering::SeqCst);
+        state
+            .statistics
+            .bytes_received
+            .fetch_add(bytes_received, Ordering::SeqCst);
+    }
+}
+
+#[inline]
+fn send_responses(
+    state: &State,
+    config: &Config,
+    socket: &mut UdpSocket,
+    buffer: &mut [u8],
+    response_receiver: &Receiver<(ConnectedResponse, SocketAddr)>,
+    local_responses: Drain<(Response, SocketAddr)>,
+) {
+    let mut responses_sent: usize = 0;
+    let mut bytes_sent: usize = 0;
+
+    let mut cursor = Cursor::new(buffer);
+
+    let response_iterator = local_responses.into_iter().chain(
+        response_receiver
+            .try_iter()
+            .map(|(response, addr)| (response.into(), addr)),
+    );
+
+    for (response, src) in response_iterator {
+        cursor.set_position(0);
+
+        let ip_version = ip_version_from_ip(src.ip());
+
+        match response.write(&mut cursor, ip_version) {
+            Ok(()) => {
+                let amt = cursor.position() as usize;
+
+                match socket.send_to(&cursor.get_ref()[..amt], src) {
+                    Ok(amt) => {
+                        responses_sent += 1;
+                        bytes_sent += amt;
+                    }
+                    Err(err) => {
+                        if err.kind() == ErrorKind::WouldBlock {
+                            break;
+                        }
+
+                        ::log::info!("send_to error: {}", err);
+                    }
+                }
+            }
+            Err(err) => {
+                ::log::error!("Response::write error: {:?}", err);
+            }
+        }
+    }
+
+    if config.statistics.interval != 0 {
+        state
+            .statistics
+            .responses_sent
+            .fetch_add(responses_sent, Ordering::SeqCst);
+        state
+            .statistics
+            .bytes_sent
+            .fetch_add(bytes_sent, Ordering::SeqCst);
+    }
+}
+
+fn ip_version_from_ip(ip: IpAddr) -> IpVersion {
+    match ip {
+        IpAddr::V4(_) => IpVersion::IPv4,
+        IpAddr::V6(ip) => {
+            if let [0, 0, 0, 0, 0, 0xffff, ..] = ip.segments() {
+                IpVersion::IPv4
+            } else {
+                IpVersion::IPv6
+            }
+        }
+    }
+}