torment/torment-manager/src/session_manager.rs

use crate::torrent_manager::TorrentManager;
use crate::tracker_manager::TrackerManager;
use std::collections::{HashMap, HashSet, VecDeque};
use std::net::{IpAddr, SocketAddr};
use std::option::Option::Some;
use torment_core::infohash::v1::U160;
use torment_core::peer_storage::{PeerStorage, PeerStorageHolder};
use torment_peer::message::{Handshake, Message};
use torment_peer::PeerProtocol;

pub struct QueuedMessage {
    pub message: Message,
    pub info_hash: U160,
    pub peer_id: U160,
    pub addr: SocketAddr,
}

pub struct SessionManager {
    id: U160,
    torrents: HashMap<U160, TorrentManager>,
    torrents_ids: HashSet<U160>,
    peer_socket: HashMap<(U160, U160), SocketAddr>,
    message_queue: VecDeque<QueuedMessage>,
    tracker_manager: TrackerManager,
    pub peer_storage: PeerStorageHolder,
}

impl SessionManager {
    pub fn peer_count(&self, info_hash: U160) -> usize {
        self.torrents[&info_hash].peer_count()
    }

    pub fn torrents(&self) -> Vec<U160> {
        self.torrents_ids.iter().copied().collect()
    }

    pub fn torrent_manager(&self, info_hash: U160) -> &TorrentManager {
        &self.torrents[&info_hash]
    }

    pub fn id(&self) -> U160 {
        self.id
    }

    pub fn tracker_manager(&self) -> &TrackerManager {
        &self.tracker_manager
    }

    pub fn tracker_manager_mut(&mut self) -> &mut TrackerManager {
        &mut self.tracker_manager
    }

    pub fn is_done(&self, info_hash: U160) -> bool {
        if let Some(torrent) = self.torrents.get(&info_hash) {
            torrent.is_done()
        } else {
            false
        }
    }

    pub fn dump_queue(&mut self) -> VecDeque<QueuedMessage> {
        std::mem::replace(&mut self.message_queue, VecDeque::new())
    }

    pub fn add_torrent_manager(&mut self, torrent: TorrentManager) {
        self.torrents_ids.insert(torrent.info_hash());
        self.torrents.insert(torrent.info_hash(), torrent);
    }

    #[allow(dead_code)]
    fn remove_torrent_manager(&mut self, info_hash: U160) -> Option<TorrentManager> {
        self.torrents_ids.remove(&info_hash);
        self.torrents.remove(&info_hash)
    }

    pub fn handshake(
        &mut self,
        handshake: Handshake,
        addr: SocketAddr,
        protocol: PeerProtocol,
    ) -> bool {
        if let Some(torrent_manager) = self.torrents.get_mut(&handshake.info_hash()) {
            self.peer_socket
                .insert((handshake.peer_id(), handshake.info_hash()), addr);
            torrent_manager.handshake(handshake, addr, protocol)
        } else {
            false
        }
    }

    pub fn process(&mut self, torrent: U160, peer_id: U160, message: Message) -> bool {
        if let Some(torrent_manager) = self.torrents.get_mut(&torrent) {
            if torrent_manager.process(peer_id, message) {
                let queue = torrent_manager.dump_queue();
                for (peer_id, message) in queue {
                    if let Some(addr) = self.peer_socket.get(&(peer_id, torrent)).copied() {
                        self.message_queue.push_back(QueuedMessage {
                            message,
                            addr,
                            peer_id,
                            info_hash: torrent,
                        });
                    }
                }

                true
            } else {
                false
            }
        } else {
            false
        }
    }

    pub fn new(ip: IpAddr, port: u16, peer_id: U160) -> SessionManager {
        let peer_storage = PeerStorageHolder::new(PeerStorage::new());
        SessionManager {
            id: peer_id,
            torrents: Default::default(),
            tracker_manager: TrackerManager::new(ip, port, peer_id, peer_storage.clone()),
            peer_storage,
            torrents_ids: Default::default(),
            peer_socket: Default::default(),
            message_queue: Default::default(),
        }
    }

    pub fn next(&mut self) -> Option<QueuedMessage> {
        self.message_queue.pop_front()
    }

    pub fn house_keeping(&mut self) {
        for (info_hash, torrent) in &mut self.torrents {
            torrent.house_keeping();

            let queue = torrent.dump_queue();
            for (peer_id, message) in queue {
                if let Some(addr) = self.peer_socket.get(&(peer_id, *info_hash)).copied() {
                    self.message_queue.push_back(QueuedMessage {
                        message,
                        addr,
                        peer_id,
                        info_hash: *info_hash,
                    });
                }
            }
        }
    }

    pub fn announce(&mut self) {
        println!("Announcing torrents...");
        for info_hash in self.torrents_ids.clone() {
            let trackers = self.torrents[&info_hash].trackers().clone();
            for tracker_tier in trackers {
                for tracker in tracker_tier {
                    self.tracker_manager
                        .announce(tracker, &self.torrents[&info_hash]);
                }
            }
        }
    }
}