aquatic/aquatic_udp/src/lib/handlers.rs

use std::net::SocketAddr;
use std::time::Duration;
use std::vec::Drain;

use crossbeam_channel::{Sender, Receiver};
use indexmap::IndexMap;
use parking_lot::MutexGuard;
use rand::{SeedableRng, Rng, rngs::{SmallRng, StdRng}};

use bittorrent_udp::types::*;

use crate::common::*;
use crate::config::Config;


pub fn run_request_worker(
    state: State,
    config: Config,
    request_receiver: Receiver<(Request, SocketAddr)>,
    response_sender: Sender<(Response, SocketAddr)>,
){
    let mut connect_requests: Vec<(ConnectRequest, SocketAddr)> = Vec::new();
    let mut announce_requests: Vec<(AnnounceRequest, SocketAddr)> = Vec::new();
    let mut scrape_requests: Vec<(ScrapeRequest, SocketAddr)> = Vec::new();

    let mut responses: Vec<(Response, SocketAddr)> = Vec::new();

    let mut std_rng = StdRng::from_entropy();
    let mut small_rng = SmallRng::from_rng(&mut std_rng).unwrap();

    let timeout = Duration::from_micros(
        config.handlers.channel_recv_timeout_microseconds
    );

    loop {
        let mut opt_connections = 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 ConnectionMap mutex isn't locked.
        for i in 0..config.handlers.max_requests_per_iter {
            let (request, src): (Request, 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.connections.try_lock(){
                            opt_connections = Some(guard);

                            break
                        } else {
                            continue
                        }
                    },
                }
            };

            match request {
                Request::Connect(r) => {
                    connect_requests.push((r, src))
                },
                Request::Announce(r) => {
                    announce_requests.push((r, src))
                },
                Request::Scrape(r) => {
                    scrape_requests.push((r, src))
                },
            }
        }

        let mut connections: MutexGuard<ConnectionMap> = opt_connections.unwrap_or_else(||
            state.connections.lock()
        );

        handle_connect_requests(
            &config,
            &mut connections,
            &mut std_rng,
            connect_requests.drain(..),
            &mut responses
        );

        announce_requests.retain(|(request, src)| {
            let connection_key = ConnectionKey {
                connection_id: request.connection_id,
                socket_addr: *src,
            };
    
            if connections.contains_key(&connection_key){
                true
            } else {
                let response = ErrorResponse {
                    transaction_id: request.transaction_id,
                    message: "Connection invalid or expired".to_string()
                };
    
                responses.push((response.into(), *src));

                false
            }
        });

        scrape_requests.retain(|(request, src)| {
            let connection_key = ConnectionKey {
                connection_id: request.connection_id,
                socket_addr: *src,
            };
    
            if connections.contains_key(&connection_key){
                true
            } else {
                let response = ErrorResponse {
                    transaction_id: request.transaction_id,
                    message: "Connection invalid or expired".to_string()
                };
    
                responses.push((response.into(), *src));

                false
            }
        });

        ::std::mem::drop(connections);

        if !(announce_requests.is_empty() && scrape_requests.is_empty()){
            let mut torrents = 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){
                eprintln!("error sending response to channel: {}", err);
            }
        }
    }
}


#[inline]
pub fn handle_connect_requests(
    config: &Config,
    connections: &mut MutexGuard<ConnectionMap>,
    rng: &mut StdRng,
    requests: Drain<(ConnectRequest, SocketAddr)>,
    responses: &mut Vec<(Response, SocketAddr)>,
){
    let valid_until = ValidUntil::new(config.cleaning.max_connection_age);

    responses.extend(requests.map(|(request, src)| {
        let connection_id = ConnectionId(rng.gen());

        let key = ConnectionKey {
            connection_id,
            socket_addr: src,
        };

        connections.insert(key, valid_until);

        let response = Response::Connect(
            ConnectResponse {
                connection_id,
                transaction_id: request.transaction_id,
            }
        );
        
        (response, src)
    }));
}


#[inline]
pub fn handle_announce_requests(
    config: &Config,
    torrents: &mut MutexGuard<TorrentMap>,
    rng: &mut SmallRng,
    requests: Drain<(AnnounceRequest, SocketAddr)>,
    responses: &mut Vec<(Response, SocketAddr)>,
){
    let peer_valid_until = ValidUntil::new(config.cleaning.max_peer_age);

    responses.extend(requests.map(|(request, src)| {
        let peer_ip = src.ip();

        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::to_response_peer
        );

        let response = Response::Announce(AnnounceResponse {
            transaction_id: request.transaction_id,
            announce_interval: AnnounceInterval(config.network.peer_announce_interval),
            leechers: NumberOfPeers(torrent_data.num_leechers as i32),
            seeders: NumberOfPeers(torrent_data.num_seeders as i32),
            peers: response_peers
        });

        (response, src)
    }));
}


#[inline]
pub fn handle_scrape_requests(
    torrents: &mut MutexGuard<TorrentMap>,
    requests: Drain<(ScrapeRequest, SocketAddr)>,
    responses: &mut Vec<(Response, 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()
        );

        for info_hash in request.info_hashes.iter() {
            if let Some(torrent_data) = torrents.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 = Response::Scrape(ScrapeResponse {
            transaction_id: request.transaction_id,
            torrent_stats: stats,
        });

        (response, src)
    }));
}


/// Extract response peers
/// 
/// If there are more peers in map than `max_num_peers_to_take`, do a
/// half-random selection of peers from first and second halves of map,
/// in order to avoid returning too homogeneous peers.
/// 
/// Don't care if we send back announcing peer.
#[inline]
pub fn extract_response_peers<K, V, R, F>(
    rng: &mut impl Rng,
    peer_map: &IndexMap<K, V>,
    max_num_peers_to_take: usize,
    peer_conversion_function: F
) -> Vec<R>
    where
        K: Eq + ::std::hash::Hash,
        F: Fn(&V) -> R

{
    let peer_map_len = peer_map.len();

    if peer_map_len <= max_num_peers_to_take {
        peer_map.values()
            .map(peer_conversion_function)
            .collect()
    } else {
        let half_num_to_take = max_num_peers_to_take / 2;
        let half_peer_map_len = peer_map_len / 2;

        let offset_first_half = rng.gen_range(
            0,
            (half_peer_map_len + (peer_map_len % 2)) - half_num_to_take
        );
        let offset_second_half = rng.gen_range(
            half_peer_map_len,
            peer_map_len - half_num_to_take
        );

        let end_first_half = offset_first_half + half_num_to_take;
        let end_second_half = offset_second_half + half_num_to_take + (max_num_peers_to_take % 2);

        let mut peers: Vec<R> = Vec::with_capacity(max_num_peers_to_take);

        for i in offset_first_half..end_first_half {
            if let Some((_, peer)) = peer_map.get_index(i){
                peers.push(peer_conversion_function(peer))
            }
        }
        for i in offset_second_half..end_second_half {
            if let Some((_, peer)) = peer_map.get_index(i){
                peers.push(peer_conversion_function(peer))
            }
        }
        
        debug_assert_eq!(peers.len(), max_num_peers_to_take);

        peers
    }
}


#[inline]
fn calc_max_num_peers_to_take(
    config: &Config,
    peers_wanted: i32,
) -> usize {
    if peers_wanted <= 0 {
        config.network.max_response_peers as usize
    } else {
        ::std::cmp::min(
            config.network.max_response_peers as usize,
            peers_wanted as usize
        )
    }
}


#[inline(always)]
pub fn create_torrent_scrape_statistics(
    seeders: i32,
    leechers: i32
) -> TorrentScrapeStatistics {
    TorrentScrapeStatistics {
        seeders: NumberOfPeers(seeders),
        completed: NumberOfDownloads(0), // No implementation planned
        leechers: NumberOfPeers(leechers)
    }
}


#[cfg(test)]
mod tests {
    use std::net::IpAddr;
    use std::collections::HashSet;

    use indexmap::IndexMap;
    use rand::thread_rng;
    use quickcheck::{TestResult, quickcheck};

    use super::*;

    fn gen_peer_map_key_and_value(i: u32) -> (PeerMapKey, Peer) {
        let ip_address = IpAddr::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: (u32, u16)) -> TestResult {
            let gen_num_peers = data.0;
            let req_num_peers = data.1 as usize;

            let mut peer_map: PeerMap = IndexMap::new();

            for i in 0..gen_num_peers {
                let (key, value) = gen_peer_map_key_and_value(i);

                peer_map.insert(key, value);
            }

            let mut rng = thread_rng();

            let peers = extract_response_peers(
                &mut rng,
                &peer_map,
                req_num_peers,
                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;
            }

            // Check that returned peers are unique (no overlap)

            let mut ip_addresses = HashSet::new();

            for peer in peers {
                if ip_addresses.contains(&peer.ip_address){
                    success = false;

                    break;
                }

                ip_addresses.insert(peer.ip_address);
            }

            TestResult::from_bool(success)
        }   

        quickcheck(prop as fn((u32, u16)) -> TestResult);
    }
}