wip
parent
ef40af30f4
commit
c7cc7ff9ed
GPG Key ID:
AD2560A0F84F0759
@ -0,0 +1 @@
|
|||||||
|
pub const REQUEST_SIZE: usize = 16384;
|
||||||
@ -0,0 +1,122 @@
|
|||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -1,151 +1,42 @@
|
|||||||
use bytes::BytesMut;
|
use rand::random;
|
||||||
use polling::{Event, Poller};
|
|
||||||
use std::collections::HashMap;
|
|
||||||
use std::convert::TryInto;
|
|
||||||
use std::fs::File;
|
use std::fs::File;
|
||||||
use std::io::{Read, Write};
|
use std::io::Read;
|
||||||
use std::net::TcpStream;
|
use std::net::IpAddr;
|
||||||
use std::option::Option::Some;
|
use std::str::FromStr;
|
||||||
use std::path::PathBuf;
|
|
||||||
use std::process::exit;
|
|
||||||
use std::result::Result::Ok;
|
|
||||||
use std::time::Duration;
|
|
||||||
use torment_core::infohash::v1::U160;
|
use torment_core::infohash::v1::U160;
|
||||||
use torment_core::metainfo::Torrent;
|
use torment_core::metainfo::Torrent;
|
||||||
use torment_core::peer_id;
|
use torment_core::peer_id;
|
||||||
|
use torment_core::peer_storage::{PeerStorage, PeerStorageHolder};
|
||||||
use torment_manager::session_manager::SessionManager;
|
use torment_manager::session_manager::SessionManager;
|
||||||
|
use torment_manager::torment_instance::TormentInstance;
|
||||||
use torment_manager::torrent_manager::{TorrentManager, TorrentTarget};
|
use torment_manager::torrent_manager::{TorrentManager, TorrentTarget};
|
||||||
use torment_peer::message::{Handshake, Message};
|
use torment_manager::tracker_manager::TrackerManager;
|
||||||
use torment_peer::PeerProtocol::TCP;
|
|
||||||
|
|
||||||
struct TcpState {
|
|
||||||
builder: BytesMut,
|
|
||||||
handshake: Option<Handshake>,
|
|
||||||
stream: TcpStream,
|
|
||||||
}
|
|
||||||
|
|
||||||
fn main() {
|
fn main() {
|
||||||
let mut session_manager = SessionManager::new();
|
let ip = IpAddr::from_str("0.0.0.0").unwrap();
|
||||||
|
let peer_id = U160::from(peer_id(random()));
|
||||||
|
|
||||||
|
let mut session_manager = SessionManager::new(ip, 50002, peer_id);
|
||||||
|
|
||||||
let mut buffer = vec![];
|
let mut file = vec![];
|
||||||
File::open("/home/eater/Downloads/[Commie] Senyuu. - 23 [150B93D5].mkv.torrent")
|
File::open("test.torrent")
|
||||||
.unwrap()
|
.unwrap()
|
||||||
.read_to_end(&mut buffer)
|
.read_to_end(&mut file)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
let torrent = Torrent::from_bytes(file, Some("test".to_string())).unwrap();
|
||||||
|
|
||||||
|
println!("Downloading {:#?}", torrent);
|
||||||
let torrent_manager = TorrentManager::from_torrent(
|
let torrent_manager = TorrentManager::from_torrent(
|
||||||
Torrent::from_bytes(
|
torrent,
|
||||||
buffer,
|
|
||||||
Some("[Commie] Senyuu. - 23 [150B93D5].mkv".to_string()),
|
|
||||||
)
|
|
||||||
.unwrap(),
|
|
||||||
TorrentTarget {
|
TorrentTarget {
|
||||||
path: PathBuf::from("/tmp"),
|
path: "/tmp/test".into(),
|
||||||
is_base_path: true,
|
is_base_path: true,
|
||||||
},
|
},
|
||||||
None,
|
Some(session_manager.tracker_manager_mut()),
|
||||||
);
|
);
|
||||||
|
|
||||||
let peer_id = U160::from(peer_id(rand::random()));
|
|
||||||
let info_hash = torrent_manager.info_hash();
|
|
||||||
session_manager.add_torrent_manager(torrent_manager);
|
session_manager.add_torrent_manager(torrent_manager);
|
||||||
|
|
||||||
let mut tcp_stream_ktorrent = TcpStream::connect("127.0.0.1:6881").unwrap();
|
let mut instance = TormentInstance::new(50002, session_manager);
|
||||||
tcp_stream_ktorrent.set_nodelay(true).unwrap();
|
instance.logic_loop();
|
||||||
tcp_stream_ktorrent
|
|
||||||
.write_all(Handshake::new(peer_id, info_hash).to_bytes().as_ref())
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
let mut tcp_stream_transmission = TcpStream::connect("192.168.188.100:51413").unwrap();
|
|
||||||
tcp_stream_transmission.set_nodelay(true).unwrap();
|
|
||||||
tcp_stream_transmission
|
|
||||||
.write_all(Handshake::new(peer_id, info_hash).to_bytes().as_ref())
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
let mut buffer = vec![0u8; 4096 * 10];
|
|
||||||
let mut poller = Poller::new().unwrap();
|
|
||||||
poller.insert(&tcp_stream_ktorrent).unwrap();
|
|
||||||
poller.insert(&tcp_stream_transmission).unwrap();
|
|
||||||
poller.interest(&tcp_stream_ktorrent, Event::readable(0));
|
|
||||||
poller.interest(&tcp_stream_transmission, Event::readable(1));
|
|
||||||
|
|
||||||
let mut items = vec![
|
|
||||||
TcpState {
|
|
||||||
builder: Default::default(),
|
|
||||||
handshake: None,
|
|
||||||
stream: tcp_stream_ktorrent,
|
|
||||||
},
|
|
||||||
TcpState {
|
|
||||||
builder: Default::default(),
|
|
||||||
handshake: None,
|
|
||||||
stream: tcp_stream_transmission,
|
|
||||||
},
|
|
||||||
];
|
|
||||||
|
|
||||||
let mut peer_map: HashMap<(U160, U160), usize> = Default::default();
|
|
||||||
|
|
||||||
loop {
|
|
||||||
let mut events: Vec<Event> = vec![];
|
|
||||||
poller.wait(&mut events, Some(Duration::from_secs(10)));
|
|
||||||
|
|
||||||
for event in events {
|
|
||||||
println!("Event => {:?}", event);
|
|
||||||
if !event.readable {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
let item = &mut items[event.key];
|
|
||||||
let packet = item.stream.read(&mut buffer).unwrap();
|
|
||||||
item.builder.extend_from_slice(&buffer[..packet]);
|
|
||||||
let handshake = if let Some(handshake) = &item.handshake {
|
|
||||||
handshake
|
|
||||||
} else {
|
|
||||||
if item.builder.len() >= 68 {
|
|
||||||
item.handshake =
|
|
||||||
Some(Handshake::from_bytes(item.builder.split_to(68).freeze()).unwrap());
|
|
||||||
let handshake = item.handshake.as_ref().unwrap();
|
|
||||||
println!("{} => {:?}", item.stream.peer_addr().unwrap(), handshake);
|
|
||||||
peer_map.insert((handshake.info_hash(), handshake.peer_id()), event.key);
|
|
||||||
session_manager.handshake(*handshake, item.stream.peer_addr().unwrap(), TCP);
|
|
||||||
handshake
|
|
||||||
} else {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
while item.builder.len() >= 4 {
|
|
||||||
let len = u32::from_be_bytes(item.builder[..4].try_into().unwrap());
|
|
||||||
if len + 4 > item.builder.len() as u32 {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
|
|
||||||
if len == 0 {
|
|
||||||
item.builder.split_to(4);
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
let message_bytes = item.builder.split_to((4 + len) as usize).freeze();
|
|
||||||
let msg = Message::from_bytes(message_bytes.slice(4..)).unwrap();
|
|
||||||
println!("{} => {:?}", item.stream.peer_addr().unwrap(), msg);
|
|
||||||
if !session_manager.process(info_hash, handshake.peer_id(), msg) {
|
|
||||||
exit(1);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
poller.interest(&item.stream, Event::readable(event.key));
|
|
||||||
}
|
|
||||||
|
|
||||||
while let Some(queued) = session_manager.next() {
|
|
||||||
if let Some(key) = peer_map.get(&(queued.info_hash, queued.peer_id)) {
|
|
||||||
let item = &mut items[*key];
|
|
||||||
println!("{} <= {:?}", queued.addr, queued.message);
|
|
||||||
|
|
||||||
item.stream
|
|
||||||
.write_all(&queued.message.to_length_prefixed_bytes())
|
|
||||||
.unwrap();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
println!("=> Running house keeping");
|
|
||||||
session_manager.house_keeping();
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@ -0,0 +1,2 @@
|
|||||||
|
pub const MAX_OPEN_REQUESTS: usize = 250;
|
||||||
|
pub const MAX_CONNECTIONS_PER_THREAD: usize = 100;
|
||||||
@ -1,3 +1,6 @@
|
|||||||
|
mod consts;
|
||||||
pub mod session_manager;
|
pub mod session_manager;
|
||||||
|
pub mod torment_instance;
|
||||||
pub mod torrent_manager;
|
pub mod torrent_manager;
|
||||||
pub mod tracker_manager;
|
pub mod tracker_manager;
|
||||||
|
pub use consts::*;
|
||||||
|
|||||||
@ -0,0 +1,432 @@
|
|||||||
|
#![allow(dead_code)]
|
||||||
|
|
||||||
|
use crate::session_manager::SessionManager;
|
||||||
|
use bytes::{Buf, Bytes, BytesMut};
|
||||||
|
use crossbeam_channel::{bounded, unbounded, Receiver, Sender};
|
||||||
|
use polling::{Event, Poller};
|
||||||
|
use rand::seq::SliceRandom;
|
||||||
|
use std::cmp::max;
|
||||||
|
use std::collections::{HashMap, HashSet, VecDeque};
|
||||||
|
use std::convert::TryInto;
|
||||||
|
use std::io::{ErrorKind, Read, Write};
|
||||||
|
use std::net::{IpAddr, Ipv6Addr, SocketAddr, TcpListener, TcpStream};
|
||||||
|
use std::ops::Add;
|
||||||
|
use std::option::Option::Some;
|
||||||
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::thread::JoinHandle;
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
use torment_core::infohash::v1::U160;
|
||||||
|
use torment_core::LookupFilter;
|
||||||
|
use torment_peer::message::{Handshake, Message};
|
||||||
|
use torment_peer::PeerProtocol;
|
||||||
|
use torment_peer::PeerProtocol::TCP;
|
||||||
|
|
||||||
|
pub struct TormentInstance {
|
||||||
|
id: U160,
|
||||||
|
port: u16,
|
||||||
|
dht_enabled: bool,
|
||||||
|
connections: OffshoreConnections,
|
||||||
|
tracking: HashMap<SocketAddr, (U160, U160)>,
|
||||||
|
handshake_sent: HashSet<SocketAddr>,
|
||||||
|
session: SessionManager,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct OffshoreConnections {
|
||||||
|
control_receiver: Receiver<ControlMessage>,
|
||||||
|
thread: JoinHandle<()>,
|
||||||
|
thread_sender: Sender<ControlMessage>,
|
||||||
|
}
|
||||||
|
|
||||||
|
enum ControlMessage {
|
||||||
|
Data(SocketAddr, Bytes),
|
||||||
|
Handshake(Handshake, SocketAddr, PeerProtocol),
|
||||||
|
Own(TcpStream, SocketAddr),
|
||||||
|
Disown(SocketAddr),
|
||||||
|
}
|
||||||
|
|
||||||
|
struct Stream {
|
||||||
|
buffer: BytesMut,
|
||||||
|
addr: SocketAddr,
|
||||||
|
stream: TcpStream,
|
||||||
|
handshake_received: bool,
|
||||||
|
id: usize,
|
||||||
|
handshake_sent: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn connection_pool_thread(
|
||||||
|
port: u16,
|
||||||
|
control_sender: Sender<ControlMessage>,
|
||||||
|
thread_receiver: Receiver<ControlMessage>,
|
||||||
|
) {
|
||||||
|
let mut streams: HashMap<SocketAddr, Stream> = HashMap::new();
|
||||||
|
let counter = AtomicUsize::new(0);
|
||||||
|
let mut index: HashMap<usize, SocketAddr> = HashMap::new();
|
||||||
|
let listener =
|
||||||
|
TcpListener::bind(SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), port)).unwrap();
|
||||||
|
|
||||||
|
let (notify_sender, notify_receiver) = bounded(1);
|
||||||
|
let (notify_restart_sender, notify_restart_receiver) = bounded(1);
|
||||||
|
|
||||||
|
let poller = Arc::new(Poller::new().unwrap());
|
||||||
|
let thread_receiver_clone = thread_receiver.clone();
|
||||||
|
let poller_clone = Arc::clone(&poller);
|
||||||
|
|
||||||
|
std::thread::Builder::new()
|
||||||
|
.name(format!("io/notify"))
|
||||||
|
.spawn(move || {
|
||||||
|
for message in thread_receiver {
|
||||||
|
// println!("Received channel message, notifying wait");
|
||||||
|
poller_clone.notify().unwrap();
|
||||||
|
notify_sender.send(message).unwrap();
|
||||||
|
notify_restart_receiver.recv().unwrap();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
poller.insert(&listener).unwrap();
|
||||||
|
poller.interest(&listener, Event::readable(usize::MAX));
|
||||||
|
|
||||||
|
let mut buffer = [0u8; 1024 * 1024];
|
||||||
|
loop {
|
||||||
|
let mut events = vec![];
|
||||||
|
match poller.wait(&mut events, None) {
|
||||||
|
Ok(_) => {}
|
||||||
|
Err(err) if err.kind() == ErrorKind::Interrupted => {}
|
||||||
|
Err(err) => Err(err).unwrap(),
|
||||||
|
}
|
||||||
|
|
||||||
|
// println!("Events => {:?}", events);
|
||||||
|
|
||||||
|
for event in events {
|
||||||
|
if !event.readable {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if event.key == usize::MAX {
|
||||||
|
while let Ok((stream, addr)) = listener.accept() {
|
||||||
|
let id = counter.fetch_add(1, Ordering::AcqRel);
|
||||||
|
stream.set_nodelay(true).unwrap();
|
||||||
|
poller.insert(&stream).unwrap();
|
||||||
|
poller.interest(&stream, Event::readable(id)).unwrap();
|
||||||
|
streams.insert(
|
||||||
|
addr,
|
||||||
|
Stream {
|
||||||
|
addr,
|
||||||
|
handshake_received: false,
|
||||||
|
handshake_sent: false,
|
||||||
|
buffer: BytesMut::new(),
|
||||||
|
stream,
|
||||||
|
id,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
index.insert(id, addr);
|
||||||
|
}
|
||||||
|
|
||||||
|
poller
|
||||||
|
.interest(&listener, Event::readable(usize::MAX))
|
||||||
|
.unwrap();
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let addr = index[&event.key];
|
||||||
|
let size = streams
|
||||||
|
.get_mut(&addr)
|
||||||
|
.unwrap()
|
||||||
|
.stream
|
||||||
|
.read(&mut buffer)
|
||||||
|
.unwrap_or(0);
|
||||||
|
|
||||||
|
if size == 0 {
|
||||||
|
control_sender.send(ControlMessage::Disown(addr)).unwrap();
|
||||||
|
index.remove(&event.key);
|
||||||
|
if let Some(stream) = streams.remove(&addr) {
|
||||||
|
poller.remove(&stream.stream).unwrap();
|
||||||
|
}
|
||||||
|
println!("{} == Disconnected", addr);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
// println!("{} => {} bytes", addr, size);
|
||||||
|
|
||||||
|
let stream = streams.get_mut(&addr).unwrap();
|
||||||
|
stream.buffer.extend_from_slice(&buffer[0..size]);
|
||||||
|
|
||||||
|
if stream.buffer.len() >= 68 && !stream.handshake_received {
|
||||||
|
let handshake = Handshake::from_bytes(stream.buffer.split_to(68).freeze());
|
||||||
|
let handshake = match handshake {
|
||||||
|
Err(_) => {
|
||||||
|
control_sender.send(ControlMessage::Disown(addr)).unwrap();
|
||||||
|
index.remove(&event.key);
|
||||||
|
if let Some(stream) = streams.remove(&addr) {
|
||||||
|
poller.remove(&stream.stream).unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
Ok(hand) => hand,
|
||||||
|
};
|
||||||
|
|
||||||
|
println!("{} => {:?}", stream.addr, handshake);
|
||||||
|
control_sender
|
||||||
|
.send(ControlMessage::Handshake(handshake, stream.addr, TCP))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
stream.handshake_received = true;
|
||||||
|
}
|
||||||
|
|
||||||
|
poller
|
||||||
|
.interest(&stream.stream, Event::readable(event.key))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
if !stream.handshake_received {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
while stream.buffer.len() >= 4 {
|
||||||
|
let len = u32::from_be_bytes(stream.buffer[..4].try_into().unwrap());
|
||||||
|
if len == 0 {
|
||||||
|
stream.buffer.advance(4);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if len + 4 > stream.buffer.len() as u32 {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut message_bytes = stream.buffer.split_to((4 + len) as usize);
|
||||||
|
message_bytes.advance(4);
|
||||||
|
control_sender
|
||||||
|
.send(ControlMessage::Data(stream.addr, message_bytes.freeze()))
|
||||||
|
.unwrap();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut msgs = vec![];
|
||||||
|
if let Ok(message) = notify_receiver.try_recv() {
|
||||||
|
msgs.push(message);
|
||||||
|
}
|
||||||
|
|
||||||
|
while let Ok(message) = thread_receiver_clone.try_recv() {
|
||||||
|
msgs.push(message);
|
||||||
|
}
|
||||||
|
|
||||||
|
for message in msgs {
|
||||||
|
match message {
|
||||||
|
ControlMessage::Data(target, data) => {
|
||||||
|
let ok = if let Some(stream) = streams.get_mut(&target) {
|
||||||
|
if stream.stream.write_all(&data).is_err() {
|
||||||
|
control_sender
|
||||||
|
.send(ControlMessage::Disown(stream.addr))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
false
|
||||||
|
} else {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
true
|
||||||
|
};
|
||||||
|
|
||||||
|
if !ok {
|
||||||
|
if let Some(stream) = streams.remove(&target) {
|
||||||
|
poller.remove(&stream.stream).unwrap();
|
||||||
|
index.remove(&stream.id);
|
||||||
|
println!("{} == Disconnected", stream.addr);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
ControlMessage::Disown(addr) => {
|
||||||
|
if let Some(stream) = streams.remove(&addr) {
|
||||||
|
index.remove(&stream.id);
|
||||||
|
poller.remove(&stream.stream).unwrap();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
ControlMessage::Own(stream, addr) => {
|
||||||
|
let id = counter.fetch_add(1, Ordering::AcqRel);
|
||||||
|
poller.insert(&stream).unwrap();
|
||||||
|
poller.interest(&stream, Event::readable(id)).unwrap();
|
||||||
|
streams.insert(
|
||||||
|
addr,
|
||||||
|
Stream {
|
||||||
|
addr,
|
||||||
|
handshake_received: false,
|
||||||
|
handshake_sent: true,
|
||||||
|
buffer: BytesMut::new(),
|
||||||
|
stream,
|
||||||
|
id,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
index.insert(id, addr);
|
||||||
|
}
|
||||||
|
|
||||||
|
_ => {}
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
notify_restart_sender.send(()).unwrap();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TormentInstance {
|
||||||
|
pub fn new(port: u16, session: SessionManager) -> TormentInstance {
|
||||||
|
let (control_sender, control_receiver) = unbounded();
|
||||||
|
|
||||||
|
let (thread_sender, thread_receiver) = unbounded();
|
||||||
|
let thread = std::thread::Builder::new()
|
||||||
|
.name(format!("io/poll"))
|
||||||
|
.spawn(move || connection_pool_thread(port, control_sender, thread_receiver))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
TormentInstance {
|
||||||
|
id: session.id(),
|
||||||
|
port,
|
||||||
|
dht_enabled: false,
|
||||||
|
connections: OffshoreConnections {
|
||||||
|
control_receiver,
|
||||||
|
thread,
|
||||||
|
thread_sender,
|
||||||
|
},
|
||||||
|
tracking: Default::default(),
|
||||||
|
handshake_sent: Default::default(),
|
||||||
|
session,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn tracker_logic(&mut self) {
|
||||||
|
self.session.announce();
|
||||||
|
self.session.tracker_manager_mut().house_keeping();
|
||||||
|
|
||||||
|
for torrent in self.session.torrents() {
|
||||||
|
let peer_count = self.session.peer_count(torrent);
|
||||||
|
if peer_count > 25 {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut peers = self
|
||||||
|
.session
|
||||||
|
.peer_storage
|
||||||
|
.get_peers(torrent, LookupFilter::All);
|
||||||
|
|
||||||
|
println!("have {} peers for {}", peers.len(), torrent);
|
||||||
|
|
||||||
|
peers.shuffle(&mut rand::thread_rng());
|
||||||
|
|
||||||
|
let mut todo = max(0, 25 - peer_count);
|
||||||
|
while todo > 0 && peers.len() > 0 {
|
||||||
|
let peer = peers.pop().unwrap();
|
||||||
|
if self.tracking.contains_key(&peer) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let thread_sender = self.connections.thread_sender.clone();
|
||||||
|
let peer_id = self.id;
|
||||||
|
|
||||||
|
self.handshake_sent.insert(peer);
|
||||||
|
|
||||||
|
std::thread::spawn(move || {
|
||||||
|
// println!("Trying connection with {}", peer);
|
||||||
|
if let Ok(mut stream) = TcpStream::connect(peer) {
|
||||||
|
stream.set_nodelay(true).unwrap();
|
||||||
|
stream
|
||||||
|
.write_all(&Handshake::new(peer_id, torrent).to_bytes())
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
// println!("{} <= Connecting", peer);
|
||||||
|
|
||||||
|
thread_sender
|
||||||
|
.send(ControlMessage::Own(stream, peer))
|
||||||
|
.unwrap();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
todo -= 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn logic_loop(&mut self) {
|
||||||
|
let mut next_house_keeping = Instant::now();
|
||||||
|
|
||||||
|
loop {
|
||||||
|
if next_house_keeping < Instant::now() {
|
||||||
|
self.session.house_keeping();
|
||||||
|
self.tracker_logic();
|
||||||
|
next_house_keeping = Instant::now().add(Duration::from_secs(10));
|
||||||
|
}
|
||||||
|
|
||||||
|
while next_house_keeping > Instant::now() {
|
||||||
|
if let Ok(message) = self
|
||||||
|
.connections
|
||||||
|
.control_receiver
|
||||||
|
.recv_timeout(Duration::from_secs(1))
|
||||||
|
{
|
||||||
|
match message {
|
||||||
|
ControlMessage::Handshake(handshake, addr, protocol) => {
|
||||||
|
let handshake_sent = self.handshake_sent.remove(&addr);
|
||||||
|
|
||||||
|
self.tracking
|
||||||
|
.insert(addr, (handshake.peer_id(), handshake.info_hash()));
|
||||||
|
if self.session.handshake(handshake, addr, protocol) {
|
||||||
|
if !handshake_sent {
|
||||||
|
self.connections
|
||||||
|
.thread_sender
|
||||||
|
.send(ControlMessage::Data(
|
||||||
|
addr,
|
||||||
|
Bytes::from(
|
||||||
|
Handshake::new(self.id, handshake.info_hash())
|
||||||
|
.to_bytes()
|
||||||
|
.to_vec(),
|
||||||
|
),
|
||||||
|
))
|
||||||
|
.unwrap();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
ControlMessage::Data(sock_addr, data) => {
|
||||||
|
if let Some((peer_id, info_hash)) =
|
||||||
|
self.tracking.get(&sock_addr).copied()
|
||||||
|
{
|
||||||
|
let message = Message::from_bytes(data);
|
||||||
|
// println!("{} => {:?}", sock_addr, message);
|
||||||
|
if message.is_err() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
self.session.process(info_hash, peer_id, message.unwrap());
|
||||||
|
} else {
|
||||||
|
// println!("{} => ????", sock_addr);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let queue: HashMap<_, VecDeque<_>> =
|
||||||
|
self.session
|
||||||
|
.dump_queue()
|
||||||
|
.into_iter()
|
||||||
|
.fold(HashMap::new(), |mut map, item| {
|
||||||
|
map.entry(item.addr).or_default().push_back(item.message);
|
||||||
|
map
|
||||||
|
});
|
||||||
|
for (addr, messages) in queue {
|
||||||
|
let mut bytes = BytesMut::new();
|
||||||
|
|
||||||
|
for msg in messages {
|
||||||
|
// println!("{} <= {:?}", addr, msg);
|
||||||
|
bytes.extend_from_slice(&msg.to_length_prefixed_bytes());
|
||||||
|
}
|
||||||
|
|
||||||
|
self.connections
|
||||||
|
.thread_sender
|
||||||
|
.send(ControlMessage::Data(addr, bytes.freeze()))
|
||||||
|
.unwrap()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,13 @@
|
|||||||
|
[package]
|
||||||
|
name = "torment-tracker"
|
||||||
|
version = "0.1.0"
|
||||||
|
authors = ["eater <=@eater.me>"]
|
||||||
|
edition = "2018"
|
||||||
|
|
||||||
|
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
torment-core = { path = "../torment-core" }
|
||||||
|
torment-bencode = { path = "../torment-bencode" }
|
||||||
|
reqwest = { version = "0.10", features = ["blocking"] }
|
||||||
|
url = "2.1.1"
|
||||||
@ -0,0 +1,269 @@
|
|||||||
|
use reqwest::blocking::Client;
|
||||||
|
use reqwest::{Method, Url};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::convert::TryInto;
|
||||||
|
use std::net::{IpAddr, SocketAddr};
|
||||||
|
use std::ops::Add;
|
||||||
|
use std::str::FromStr;
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
use torment_bencode::BencodeValue;
|
||||||
|
use torment_core::infohash::v1::U160;
|
||||||
|
use torment_core::infohash::InfoHashCapable;
|
||||||
|
use torment_core::peer_storage::{PeerStorageHolder, PeerStorageSource};
|
||||||
|
use torment_core::CompactContact;
|
||||||
|
|
||||||
|
pub struct Tracker {
|
||||||
|
url: Url,
|
||||||
|
is_udp: bool,
|
||||||
|
// last_scrape: Option<Instant>,
|
||||||
|
last_announce: HashMap<U160, Instant>,
|
||||||
|
requested_interval: Duration,
|
||||||
|
failure: HashMap<U160, String>,
|
||||||
|
peer_storage: PeerStorageHolder,
|
||||||
|
// supports_scrape: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub enum TrackerAnnounceRequest {
|
||||||
|
HTTP(Url, U160),
|
||||||
|
UDP(Url, U160),
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct TrackerAnnounceResponse {
|
||||||
|
info_hash: U160,
|
||||||
|
// url: Url,
|
||||||
|
interval: Option<Duration>,
|
||||||
|
source: PeerStorageSource,
|
||||||
|
result: Result<Vec<SocketAddr>, String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Tracker {
|
||||||
|
pub fn url(&self) -> &Url {
|
||||||
|
&self.url
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn new(url: Url, peer_storage: PeerStorageHolder) -> Tracker {
|
||||||
|
Tracker {
|
||||||
|
is_udp: url.scheme() == "udp",
|
||||||
|
url,
|
||||||
|
last_announce: HashMap::new(),
|
||||||
|
requested_interval: Duration::from_secs(30 * 60),
|
||||||
|
failure: HashMap::new(),
|
||||||
|
peer_storage,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn needs_update(&self, info_hash: U160) -> bool {
|
||||||
|
self.last_announce
|
||||||
|
.get(&info_hash)
|
||||||
|
.map(|inst| inst.add(self.requested_interval) < Instant::now())
|
||||||
|
.unwrap_or(true)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn process(&mut self, response: TrackerAnnounceResponse) {
|
||||||
|
match response.result {
|
||||||
|
Ok(peers) => {
|
||||||
|
self.peer_storage
|
||||||
|
.add_peers(response.info_hash, peers, response.source);
|
||||||
|
}
|
||||||
|
|
||||||
|
Err(err) => {
|
||||||
|
self.failure.insert(response.info_hash, err);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(interval) = response.interval {
|
||||||
|
self.requested_interval = interval;
|
||||||
|
} else {
|
||||||
|
self.requested_interval = Duration::from_secs(60 * 30);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn create_update_request(
|
||||||
|
&mut self,
|
||||||
|
info_hash: U160,
|
||||||
|
peer_id: U160,
|
||||||
|
ip: IpAddr,
|
||||||
|
port: u16,
|
||||||
|
uploaded: usize,
|
||||||
|
downloaded: usize,
|
||||||
|
left: i64,
|
||||||
|
) -> TrackerAnnounceRequest {
|
||||||
|
if self.is_udp {
|
||||||
|
return TrackerAnnounceRequest::UDP(self.url.clone(), info_hash);
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut url = self.url.clone();
|
||||||
|
info_hash.to_byte_array();
|
||||||
|
{
|
||||||
|
let mut query = url.query_pairs_mut();
|
||||||
|
query.append_pair("ip", &format!("{}", ip));
|
||||||
|
query.append_pair("port", &format!("{}", port));
|
||||||
|
query.append_pair("uploaded", &format!("{}", uploaded));
|
||||||
|
query.append_pair("downloaded", &format!("{}", downloaded));
|
||||||
|
query.append_pair("left", &format!("{}", left));
|
||||||
|
}
|
||||||
|
|
||||||
|
let encoded_info_hash = url::form_urlencoded::byte_serialize(&info_hash.to_bytes()).fold(
|
||||||
|
String::new(),
|
||||||
|
|mut coll, curr| {
|
||||||
|
coll.push_str(curr);
|
||||||
|
coll
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
let encoded_peer_id = url::form_urlencoded::byte_serialize(&peer_id.to_bytes()).fold(
|
||||||
|
String::new(),
|
||||||
|
|mut coll, curr| {
|
||||||
|
coll.push_str(curr);
|
||||||
|
coll
|
||||||
|
},
|
||||||
|
);
|
||||||
|
let query = url.query().unwrap_or("");
|
||||||
|
url.set_query(Some(&format!(
|
||||||
|
"{}&info_hash={}&peer_id={}",
|
||||||
|
query, encoded_info_hash, encoded_peer_id
|
||||||
|
)));
|
||||||
|
|
||||||
|
self.last_announce.insert(info_hash, Instant::now());
|
||||||
|
|
||||||
|
TrackerAnnounceRequest::HTTP(url, info_hash)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn block_do_announce(
|
||||||
|
request: TrackerAnnounceRequest,
|
||||||
|
client: &Client,
|
||||||
|
) -> TrackerAnnounceResponse {
|
||||||
|
let (url, info_hash) = match request {
|
||||||
|
TrackerAnnounceRequest::HTTP(url, info_hash) => (url, info_hash),
|
||||||
|
TrackerAnnounceRequest::UDP(url, info_hash) => {
|
||||||
|
return TrackerAnnounceResponse {
|
||||||
|
info_hash,
|
||||||
|
interval: None,
|
||||||
|
source: PeerStorageSource::Tracker(url.clone()),
|
||||||
|
// url,
|
||||||
|
result: Err(format!("UDP trackers not supported")),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
println!("Announcing {} on {}", info_hash, url);
|
||||||
|
|
||||||
|
let res = block_do_http_response_parse(url.clone(), client);
|
||||||
|
|
||||||
|
TrackerAnnounceResponse {
|
||||||
|
info_hash,
|
||||||
|
source: PeerStorageSource::Tracker(url.clone()),
|
||||||
|
// url,
|
||||||
|
interval: res.as_ref().map(|x| x.0).ok(),
|
||||||
|
result: res.map(|x| x.1),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn block_do_http_response_parse(
|
||||||
|
url: Url,
|
||||||
|
http_client: &Client,
|
||||||
|
) -> Result<(Duration, Vec<SocketAddr>), String> {
|
||||||
|
let resp = match http_client.request(Method::GET, url).send() {
|
||||||
|
Ok(resp) => resp,
|
||||||
|
|
||||||
|
Err(err) => {
|
||||||
|
return Err(format!("{}", err));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let body = resp.bytes();
|
||||||
|
if let Err(err) = body {
|
||||||
|
return Err(format!("{}", err));
|
||||||
|
}
|
||||||
|
|
||||||
|
let body = body.unwrap();
|
||||||
|
let body = BencodeValue::decode(body);
|
||||||
|
if let Err(err) = body {
|
||||||
|
return Err(format!("Bencoding error: {}", err));
|
||||||
|
}
|
||||||
|
|
||||||
|
let body = body.unwrap();
|
||||||
|
let body = body.dict();
|
||||||
|
if body.is_none() {
|
||||||
|
return Err(format!("Tracker gave invalid response"));
|
||||||
|
}
|
||||||
|
|
||||||
|
let body = body.unwrap();
|
||||||
|
|
||||||
|
if let Some(failure) = body.get(b"failure reason") {
|
||||||
|
let failure = failure.string().and_then(|x| x.ok()).unwrap_or(format!(
|
||||||
|
"Tracker failed to respond with valid failure, but did respond with a failure"
|
||||||
|
));
|
||||||
|
|
||||||
|
return Err(failure);
|
||||||
|
}
|
||||||
|
|
||||||
|
let interval = body
|
||||||
|
.get(b"interval")
|
||||||
|
.and_then(|item| item.int())
|
||||||
|
.unwrap_or(30 * 60);
|
||||||
|
|
||||||
|
let interval = Duration::from_secs(interval as u64);
|
||||||
|
|
||||||
|
let peers = body.get(b"peers");
|
||||||
|
|
||||||
|
if peers.is_none() {
|
||||||
|
return Ok((interval, vec![]));
|
||||||
|
}
|
||||||
|
|
||||||
|
let peers = peers.unwrap();
|
||||||
|
|
||||||
|
let peers: Vec<SocketAddr> = match peers {
|
||||||
|
BencodeValue::List(old) => {
|
||||||
|
let mut buffer = vec![];
|
||||||
|
for item in old.iter() {
|
||||||
|
let dict = if let Some(dict) = item.dict() {
|
||||||
|
dict
|
||||||
|
} else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
|
||||||
|
let ip = if let Some(BencodeValue::Bytes(ip)) = dict.get(b"ip") {
|
||||||
|
ip.clone()
|
||||||
|
} else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
|
||||||
|
let ip = if let Some(ip) = IpAddr::from_str(&String::from_utf8_lossy(&ip)).ok() {
|
||||||
|
ip
|
||||||
|
} else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
|
||||||
|
let port: u16 = if let Some(BencodeValue::Int(nr)) = dict.get(b"port") {
|
||||||
|
if *nr < u16::MAX as i64 {
|
||||||
|
if let Ok(port) = (*nr).try_into() {
|
||||||
|
port
|
||||||
|
} else {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
|
||||||
|
buffer.push(SocketAddr::new(ip, port));
|
||||||
|
}
|
||||||
|
|
||||||
|
buffer
|
||||||
|
}
|
||||||
|
BencodeValue::Bytes(compact) => compact
|
||||||
|
.chunks(6)
|
||||||
|
.filter_map(|x| SocketAddr::from_compact_contact(x).ok())
|
||||||
|
.collect(),
|
||||||
|
_ => {
|
||||||
|
return Err(format!("Tracker failed to respond with valid response"));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok((interval, peers))
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue