You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
87 lines
2.5 KiB
Rust
87 lines
2.5 KiB
Rust
use crate::torrent_manager::TorrentManager;
|
|
use std::collections::{HashMap, VecDeque};
|
|
use std::net::SocketAddr;
|
|
use std::option::Option::Some;
|
|
use torment_core::infohash::v1::U160;
|
|
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 {
|
|
torrents: HashMap<U160, TorrentManager>,
|
|
peer_socket: HashMap<(U160, U160), SocketAddr>,
|
|
message_queue: VecDeque<QueuedMessage>,
|
|
}
|
|
|
|
impl SessionManager {
|
|
pub fn add_torrent_manager(&mut self, torrent: TorrentManager) {
|
|
self.torrents.insert(torrent.info_hash(), torrent);
|
|
}
|
|
|
|
fn remove_torrent_manager(&mut self, info_hash: U160) -> Option<TorrentManager> {
|
|
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) {
|
|
while let Some((peer_id, message)) = torrent_manager.next() {
|
|
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() -> SessionManager {
|
|
SessionManager {
|
|
torrents: 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 (_, torrent) in &mut self.torrents {
|
|
torrent.house_keeping();
|
|
}
|
|
}
|
|
}
|