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 polling::{Event, Poller};
|
||||
use std::collections::HashMap;
|
||||
use std::convert::TryInto;
|
||||
use rand::random;
|
||||
use std::fs::File;
|
||||
use std::io::{Read, Write};
|
||||
use std::net::TcpStream;
|
||||
use std::option::Option::Some;
|
||||
use std::path::PathBuf;
|
||||
use std::process::exit;
|
||||
use std::result::Result::Ok;
|
||||
use std::time::Duration;
|
||||
use std::io::Read;
|
||||
use std::net::IpAddr;
|
||||
use std::str::FromStr;
|
||||
use torment_core::infohash::v1::U160;
|
||||
use torment_core::metainfo::Torrent;
|
||||
use torment_core::peer_id;
|
||||
use torment_core::peer_storage::{PeerStorage, PeerStorageHolder};
|
||||
use torment_manager::session_manager::SessionManager;
|
||||
use torment_manager::torment_instance::TormentInstance;
|
||||
use torment_manager::torrent_manager::{TorrentManager, TorrentTarget};
|
||||
use torment_peer::message::{Handshake, Message};
|
||||
use torment_peer::PeerProtocol::TCP;
|
||||
|
||||
struct TcpState {
|
||||
builder: BytesMut,
|
||||
handshake: Option<Handshake>,
|
||||
stream: TcpStream,
|
||||
}
|
||||
use torment_manager::tracker_manager::TrackerManager;
|
||||
|
||||
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![];
|
||||
File::open("/home/eater/Downloads/[Commie] Senyuu. - 23 [150B93D5].mkv.torrent")
|
||||
let mut file = vec![];
|
||||
File::open("test.torrent")
|
||||
.unwrap()
|
||||
.read_to_end(&mut buffer)
|
||||
.read_to_end(&mut file)
|
||||
.unwrap();
|
||||
let torrent = Torrent::from_bytes(file, Some("test".to_string())).unwrap();
|
||||
|
||||
println!("Downloading {:#?}", torrent);
|
||||
let torrent_manager = TorrentManager::from_torrent(
|
||||
Torrent::from_bytes(
|
||||
buffer,
|
||||
Some("[Commie] Senyuu. - 23 [150B93D5].mkv".to_string()),
|
||||
)
|
||||
.unwrap(),
|
||||
torrent,
|
||||
TorrentTarget {
|
||||
path: PathBuf::from("/tmp"),
|
||||
path: "/tmp/test".into(),
|
||||
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);
|
||||
|
||||
let mut tcp_stream_ktorrent = TcpStream::connect("127.0.0.1:6881").unwrap();
|
||||
tcp_stream_ktorrent.set_nodelay(true).unwrap();
|
||||
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();
|
||||
}
|
||||
let mut instance = TormentInstance::new(50002, session_manager);
|
||||
instance.logic_loop();
|
||||
}
|
||||
|
||||
@ -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 torment_instance;
|
||||
pub mod torrent_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