torment/torment-core/src/peer_storage.rs

use crate::infohash::v1::U160;
use crate::LookupFilter;
use std::collections::{HashMap, HashSet};
use std::net::SocketAddr;
use std::ops::Add;
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};
use url::Url;

#[derive(Debug, Default)]
pub struct PeerStorage {
    torrents: HashMap<U160, PeerCollection>,
}

#[derive(Debug, Default)]
pub struct PeerCollection {
    sets: Vec<PeerSet>,
    peers: HashSet<SocketAddr>,
    peers6: HashSet<SocketAddr>,
}

#[derive(Debug)]
pub struct PeerSet {
    source: PeerStorageSource,
    expires: Instant,
    peers: Vec<SocketAddr>,
}

impl PeerCollection {
    fn add_peers(&mut self, peers: Vec<SocketAddr>, source: PeerStorageSource) {
        let set = PeerSet {
            source,
            peers,
            expires: Instant::now().add(Duration::from_secs(60 * 30)),
        };

        for peer in &set.peers {
            if peer.is_ipv4() {
                self.peers.insert(*peer);
            } else {
                self.peers6.insert(*peer);
            }
        }

        self.sets.push(set);
    }
}

#[derive(Debug, Clone)]
pub enum PeerStorageSource {
    DHT,
    Tracker(Url),
}

impl PeerStorage {
    pub fn new() -> PeerStorage {
        PeerStorage {
            torrents: Default::default(),
        }
    }

    pub fn add_peers(
        &mut self,
        info_hash: U160,
        peers: Vec<SocketAddr>,
        source: PeerStorageSource,
    ) {
        self.torrents
            .entry(info_hash)
            .or_default()
            .add_peers(peers, source)
    }

    pub fn get_peers(&self, info_hash: U160, filter: LookupFilter) -> Vec<SocketAddr> {
        let collection = if let Some(collection) = self.torrents.get(&info_hash) {
            collection
        } else {
            return vec![];
        };

        let mut buffer = vec![];

        if filter != LookupFilter::IPv6 {
            buffer.extend(&collection.peers);
        }

        if filter != LookupFilter::IPv4 {
            buffer.extend(&collection.peers6);
        }

        buffer
    }
}

#[derive(Debug, Clone)]
pub struct PeerStorageHolder {
    inner: Arc<RwLock<PeerStorage>>,
}

impl PeerStorageHolder {
    pub fn new(peer_storage: PeerStorage) -> PeerStorageHolder {
        PeerStorageHolder {
            inner: Arc::new(RwLock::new(peer_storage)),
        }
    }

    pub fn get_peers(&self, info_hash: U160, filter: LookupFilter) -> Vec<SocketAddr> {
        self.inner.read().unwrap().get_peers(info_hash, filter)
    }

    pub fn add_peers(
        &mut self,
        info_hash: U160,
        peers: Vec<SocketAddr>,
        source: PeerStorageSource,
    ) {
        self.inner
            .write()
            .unwrap()
            .add_peers(info_hash, peers, source)
    }
}