wip
parent
b806c63df4
commit
ef40af30f4
GPG Key ID:
AD2560A0F84F0759
@ -1 +1,2 @@
|
|||||||
/target
|
/target
|
||||||
|
callgrind*
|
||||||
|
|||||||
@ -0,0 +1,20 @@
|
|||||||
|
name: torment-cli
|
||||||
|
version: "0.1.0"
|
||||||
|
about: "A torrent client"
|
||||||
|
subcommands:
|
||||||
|
- torrent:
|
||||||
|
about: "Show details about a torrent"
|
||||||
|
args:
|
||||||
|
- file:
|
||||||
|
required: true
|
||||||
|
- download:
|
||||||
|
about: "Show details about a torrent"
|
||||||
|
args:
|
||||||
|
- file:
|
||||||
|
help: Set the torrent to download
|
||||||
|
required: true
|
||||||
|
- peer:
|
||||||
|
short: p
|
||||||
|
long: peer
|
||||||
|
help: Set a custom peer
|
||||||
|
takes_value: true
|
||||||
@ -0,0 +1,144 @@
|
|||||||
|
use bytes::{Bytes, BytesMut};
|
||||||
|
use rand::random;
|
||||||
|
use std::cmp::min;
|
||||||
|
use std::convert::TryInto;
|
||||||
|
use std::io::{Read, Write};
|
||||||
|
use std::net::{SocketAddr, TcpStream};
|
||||||
|
use std::process::exit;
|
||||||
|
use torment_core::infohash::v1::U160;
|
||||||
|
use torment_core::infohash::InfoHashCapable;
|
||||||
|
use torment_core::metainfo::Torrent;
|
||||||
|
use torment_core::peer_id;
|
||||||
|
use torment_peer::message::Message::Request;
|
||||||
|
use torment_peer::message::{Handshake, Message, SelectionMessage};
|
||||||
|
use torment_peer::{Peer, PeerProtocol};
|
||||||
|
use torment_storage::ToStorageMap;
|
||||||
|
|
||||||
|
pub fn download(torrent: Torrent, peers: Vec<SocketAddr>) {
|
||||||
|
println!("Starting download of {}", torrent.name());
|
||||||
|
if peers.len() == 0 {
|
||||||
|
println!("No peers given, stopping");
|
||||||
|
exit(1);
|
||||||
|
}
|
||||||
|
let peer = peers.first().unwrap();
|
||||||
|
let mut stream = TcpStream::connect(peer).expect("Failed to open tcp connection to peer");
|
||||||
|
stream.set_nodelay(true).expect("Failed setting nodelay");
|
||||||
|
let id = U160::from(peer_id(random()));
|
||||||
|
|
||||||
|
let our_header = Handshake::new(id, torrent.info_hash());
|
||||||
|
stream.write(&our_header.to_bytes());
|
||||||
|
let mut buffer = [0u8; 4096];
|
||||||
|
let mut done_header = false;
|
||||||
|
let mut header = Handshake::new(U160::random(), U160::random());
|
||||||
|
let mut peer = None;
|
||||||
|
let mut message_buffer = BytesMut::with_capacity(4096 * 10);
|
||||||
|
let mut storage = torrent.to_storage_map("/tmp", false);
|
||||||
|
let mut asked = 0;
|
||||||
|
let request_length = 2u32.pow(14);
|
||||||
|
let piece_length = torrent.meta_info().piece_length() as u32;
|
||||||
|
let mut progress = 0;
|
||||||
|
let mut header_buffer = BytesMut::new();
|
||||||
|
loop {
|
||||||
|
let size = stream.read(&mut buffer).expect("Failed to receive data");
|
||||||
|
if done_header == false {
|
||||||
|
header_buffer.extend_from_slice(&buffer[..size]);
|
||||||
|
if header_buffer.len() < 68 {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
header = Handshake::from_bytes(Bytes::from(buffer[..size].to_vec()))
|
||||||
|
.expect("Failed to parse header");
|
||||||
|
if header.info_hash() != torrent.info_hash() {
|
||||||
|
panic!(
|
||||||
|
"Peer wants to serve different torrent ({}) :(",
|
||||||
|
header.info_hash()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
println!(
|
||||||
|
"Peer[{}] connected",
|
||||||
|
String::from_utf8_lossy(&header.peer_id().to_bytes())
|
||||||
|
);
|
||||||
|
peer = Some(Peer::new(
|
||||||
|
stream.peer_addr().unwrap(),
|
||||||
|
PeerProtocol::TCP,
|
||||||
|
header,
|
||||||
|
torrent.meta_info(),
|
||||||
|
));
|
||||||
|
done_header = true;
|
||||||
|
// stream.write_all(&*Message::Request(SelectionMessage::new(0, 0, 1024)).to_bytes());
|
||||||
|
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut peer = peer.as_mut().unwrap();
|
||||||
|
message_buffer.extend_from_slice(&buffer[..size]);
|
||||||
|
while message_buffer.len() >= 4 {
|
||||||
|
let length = u32::from_be_bytes(message_buffer[..4].try_into().unwrap());
|
||||||
|
if length == 0 {
|
||||||
|
message_buffer.split_to(4);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if message_buffer.len() < (length + 4) as usize {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
let message = message_buffer.split_to((4 + length) as usize).freeze();
|
||||||
|
|
||||||
|
let msg = Message::from_bytes(message.slice(4..)).expect("Failed parsing message");
|
||||||
|
// println!("=> {:?}", msg);
|
||||||
|
if msg == Message::Unchoke {
|
||||||
|
stream.write_all(&Message::Interested.to_length_prefixed_bytes());
|
||||||
|
}
|
||||||
|
peer.process(msg);
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut messages = vec![];
|
||||||
|
let was_above_90 = peer.count_open_requests() > 90;
|
||||||
|
while !peer.is_choked()
|
||||||
|
&& peer.count_open_requests() < 500
|
||||||
|
&& !was_above_90
|
||||||
|
&& asked < torrent.meta_info().pieces()
|
||||||
|
{
|
||||||
|
let curr_piece_length = if asked + 1 == torrent.meta_info().pieces() {
|
||||||
|
piece_length - (storage.size() as u32 % piece_length)
|
||||||
|
} else {
|
||||||
|
piece_length
|
||||||
|
};
|
||||||
|
let size = min(request_length, curr_piece_length - progress);
|
||||||
|
let sel = SelectionMessage::new(asked as u32, progress, size);
|
||||||
|
peer.requested(sel);
|
||||||
|
let msg = Message::Request(sel);
|
||||||
|
// println!("<= {:?}", msg);
|
||||||
|
messages.extend(msg.to_length_prefixed_bytes());
|
||||||
|
progress += size;
|
||||||
|
if progress >= curr_piece_length {
|
||||||
|
asked += 1;
|
||||||
|
progress = 0;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if messages.len() > 0 {
|
||||||
|
stream
|
||||||
|
.write_all(&messages)
|
||||||
|
.expect("Failed writing to network socket");
|
||||||
|
}
|
||||||
|
|
||||||
|
while let Some(piece) = peer.next_piece() {
|
||||||
|
// println!("> {} {} {}", piece.index(), piece.offset(), piece.length());
|
||||||
|
storage
|
||||||
|
.write(
|
||||||
|
piece.index() as usize,
|
||||||
|
piece.offset() as usize,
|
||||||
|
piece.piece(),
|
||||||
|
)
|
||||||
|
.expect("Failed writing piece to storage");
|
||||||
|
}
|
||||||
|
|
||||||
|
if peer.count_open_requests() == 0 && asked > 1 {
|
||||||
|
println!("done");
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,15 @@
|
|||||||
|
[package]
|
||||||
|
name = "torment-daemon"
|
||||||
|
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-peer = { path = "../torment-peer" }
|
||||||
|
torment-manager = { path = "../torment-manager" }
|
||||||
|
bytes = "0.5.6"
|
||||||
|
polling = "0.1.6"
|
||||||
|
rand = "0.7.3"
|
||||||
@ -0,0 +1,151 @@
|
|||||||
|
use bytes::BytesMut;
|
||||||
|
use polling::{Event, Poller};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::convert::TryInto;
|
||||||
|
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 torment_core::infohash::v1::U160;
|
||||||
|
use torment_core::metainfo::Torrent;
|
||||||
|
use torment_core::peer_id;
|
||||||
|
use torment_manager::session_manager::SessionManager;
|
||||||
|
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,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn main() {
|
||||||
|
let mut session_manager = SessionManager::new();
|
||||||
|
|
||||||
|
let mut buffer = vec![];
|
||||||
|
File::open("/home/eater/Downloads/[Commie] Senyuu. - 23 [150B93D5].mkv.torrent")
|
||||||
|
.unwrap()
|
||||||
|
.read_to_end(&mut buffer)
|
||||||
|
.unwrap();
|
||||||
|
let torrent_manager = TorrentManager::from_torrent(
|
||||||
|
Torrent::from_bytes(
|
||||||
|
buffer,
|
||||||
|
Some("[Commie] Senyuu. - 23 [150B93D5].mkv".to_string()),
|
||||||
|
)
|
||||||
|
.unwrap(),
|
||||||
|
TorrentTarget {
|
||||||
|
path: PathBuf::from("/tmp"),
|
||||||
|
is_base_path: true,
|
||||||
|
},
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
|
||||||
|
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();
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,16 @@
|
|||||||
|
[package]
|
||||||
|
name = "torment-manager"
|
||||||
|
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-peer = { path = "../torment-peer" }
|
||||||
|
torment-storage = { path = "../torment-storage" }
|
||||||
|
bytes = "0.5.6"
|
||||||
|
url = "2.1.1"
|
||||||
|
lazy_static = "1.4.0"
|
||||||
|
ring = "0.16.15"
|
||||||
@ -0,0 +1,3 @@
|
|||||||
|
pub mod session_manager;
|
||||||
|
pub mod torrent_manager;
|
||||||
|
pub mod tracker_manager;
|
||||||
@ -0,0 +1,86 @@
|
|||||||
|
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();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,275 @@
|
|||||||
|
use crate::tracker_manager::{TrackerId, TrackerManager};
|
||||||
|
use bytes::Bytes;
|
||||||
|
use ring::digest::{digest, SHA1_FOR_LEGACY_USE_ONLY};
|
||||||
|
use std::collections::{HashMap, HashSet, VecDeque};
|
||||||
|
use std::net::SocketAddr;
|
||||||
|
use std::ops::Index;
|
||||||
|
use std::option::Option::Some;
|
||||||
|
use std::path::PathBuf;
|
||||||
|
use torment_core::infohash::v1::U160;
|
||||||
|
use torment_core::metainfo::{MetaInfo, Torrent};
|
||||||
|
use torment_core::Bitfield;
|
||||||
|
use torment_peer::message::{Handshake, Message, PieceMessage, SelectionMessage};
|
||||||
|
use torment_peer::{Peer, PeerProtocol};
|
||||||
|
use torment_storage::{StorageMap, ToStorageMap};
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub enum Source {
|
||||||
|
Torrent(TorrentSource),
|
||||||
|
MetaInfo(MetaInfoSource),
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct MetaInfoSource {
|
||||||
|
info_hash: U160,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct TorrentSource {
|
||||||
|
file: Option<PathBuf>,
|
||||||
|
torrent: Option<Torrent>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct TorrentTarget {
|
||||||
|
pub path: PathBuf,
|
||||||
|
pub is_base_path: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct TorrentManager {
|
||||||
|
info_hash: U160,
|
||||||
|
bitfield: Bitfield,
|
||||||
|
trackers: Vec<Vec<TrackerId>>,
|
||||||
|
source: Source,
|
||||||
|
target: TorrentTarget,
|
||||||
|
peers: HashMap<U160, Peer>,
|
||||||
|
queue: VecDeque<(U160, Message)>,
|
||||||
|
storage_map: StorageMap,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub const REQUEST_SIZE: usize = 16384;
|
||||||
|
|
||||||
|
impl TorrentManager {
|
||||||
|
pub fn info_hash(&self) -> U160 {
|
||||||
|
self.info_hash
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn from_torrent(
|
||||||
|
torrent: Torrent,
|
||||||
|
target: TorrentTarget,
|
||||||
|
tracker_manager: Option<&mut TrackerManager>,
|
||||||
|
) -> TorrentManager {
|
||||||
|
let trackers = tracker_manager
|
||||||
|
.map(|manager| {
|
||||||
|
torrent
|
||||||
|
.announce_list()
|
||||||
|
.iter()
|
||||||
|
.map(|tier_list| {
|
||||||
|
tier_list
|
||||||
|
.iter()
|
||||||
|
.filter_map(|tracker| manager.get_tracker_id(tracker))
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
})
|
||||||
|
.unwrap_or(vec![]);
|
||||||
|
|
||||||
|
TorrentManager {
|
||||||
|
info_hash: torrent.info_hash(),
|
||||||
|
bitfield: Bitfield::with_size(torrent.meta_info().pieces()),
|
||||||
|
trackers,
|
||||||
|
source: Source::Torrent(TorrentSource {
|
||||||
|
file: None,
|
||||||
|
torrent: Some(torrent.clone()),
|
||||||
|
}),
|
||||||
|
storage_map: torrent.to_storage_map(&target.path, target.is_base_path),
|
||||||
|
target,
|
||||||
|
peers: HashMap::new(),
|
||||||
|
queue: Default::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn handshake(
|
||||||
|
&mut self,
|
||||||
|
handshake: Handshake,
|
||||||
|
addr: SocketAddr,
|
||||||
|
protocol: PeerProtocol,
|
||||||
|
) -> bool {
|
||||||
|
if handshake.info_hash() != self.info_hash {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(peer) = self.peers.get(&handshake.peer_id()) {
|
||||||
|
if peer.addr() != addr {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let meta_info = self.meta_info();
|
||||||
|
let peer = Peer::new(addr, protocol, handshake, meta_info);
|
||||||
|
self.peers.insert(peer.id(), peer);
|
||||||
|
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
fn meta_info(&mut self) -> &MetaInfo {
|
||||||
|
if let Source::Torrent(torrent) = &self.source {
|
||||||
|
return torrent
|
||||||
|
.torrent
|
||||||
|
.as_ref()
|
||||||
|
.expect("No torrent loaded uh oh missing functionality")
|
||||||
|
.meta_info();
|
||||||
|
}
|
||||||
|
|
||||||
|
panic!("Can't resolve MetaInfo for torrent")
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn process(&mut self, peer_id: U160, message: Message) -> bool {
|
||||||
|
let mut queue = vec![];
|
||||||
|
let ok = if let Some(peer) = self.peers.get_mut(&peer_id) {
|
||||||
|
if peer.process(message) {
|
||||||
|
while let Some(message) = peer.next() {
|
||||||
|
self.queue.push_back((peer_id, message));
|
||||||
|
}
|
||||||
|
|
||||||
|
while let Some(piece) = peer.next_piece() {
|
||||||
|
if self.storage_map.has_piece(piece.index() as usize) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if self
|
||||||
|
.storage_map
|
||||||
|
.write(
|
||||||
|
piece.index() as usize,
|
||||||
|
piece.offset() as usize,
|
||||||
|
piece.piece(),
|
||||||
|
)
|
||||||
|
.unwrap()
|
||||||
|
{
|
||||||
|
queue.push(piece.index() as usize);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
while let Some(piece_request) = peer.next_request() {
|
||||||
|
if !self.bitfield.get(piece_request.index()) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut buffer = vec![0u8; piece_request.length() as usize];
|
||||||
|
self.storage_map
|
||||||
|
.read(
|
||||||
|
piece_request.index() as usize,
|
||||||
|
piece_request.offset() as usize,
|
||||||
|
&mut buffer,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
self.queue.push_back((
|
||||||
|
peer_id,
|
||||||
|
Message::Piece(PieceMessage {
|
||||||
|
index: piece_request.index(),
|
||||||
|
offset: piece_request.offset(),
|
||||||
|
piece: Bytes::from(buffer),
|
||||||
|
}),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
true
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
};
|
||||||
|
|
||||||
|
for have in queue {
|
||||||
|
let piece_hash = self.meta_info().hash(have);
|
||||||
|
let piece_data = self.storage_map.read_piece(have).unwrap();
|
||||||
|
let res = digest(&SHA1_FOR_LEGACY_USE_ONLY, &piece_data);
|
||||||
|
if piece_hash != res.as_ref() {
|
||||||
|
println!("=> Piece#{} failed verification", have);
|
||||||
|
self.storage_map.wipe_piece(have);
|
||||||
|
continue;
|
||||||
|
} else {
|
||||||
|
println!("=> Piece#{} verified", have);
|
||||||
|
}
|
||||||
|
|
||||||
|
self.bitfield.set(have as u32);
|
||||||
|
|
||||||
|
let keys = self.peers.keys().copied().collect::<Vec<_>>();
|
||||||
|
for key in keys {
|
||||||
|
self.queue.push_back((key, Message::Have(have as u32)));
|
||||||
|
|
||||||
|
if !self.peers[&key].has_piece(have as u32) && self.peers[&key].we_choked() {
|
||||||
|
self.peers.get_mut(&key).unwrap().set_we_choked(false);
|
||||||
|
self.queue.push_back((key, Message::Unchoke));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
ok
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn next(&mut self) -> Option<(U160, Message)> {
|
||||||
|
self.queue.pop_front()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn house_keeping(&mut self) {
|
||||||
|
let mut peers = self
|
||||||
|
.peers
|
||||||
|
.iter()
|
||||||
|
.filter_map(|(_, peer)| {
|
||||||
|
if peer.has()1
|
||||||
|
|
||||||
|
if peer.is_choked() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(peer.id())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect::<HashSet<_>>();
|
||||||
|
|
||||||
|
let pieces = self.meta_info().pieces() as u32;
|
||||||
|
for i in 0u32..pieces {
|
||||||
|
if self.bitfield.get(i) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let length = self.storage_map.get_piece_length(i as usize);
|
||||||
|
let mut offset = 0;
|
||||||
|
|
||||||
|
for peer in &peers.clone() {
|
||||||
|
if !self.peers[peer].has_piece(i) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
while offset < length && self.peers[peer].count_open_requests() < 25 {
|
||||||
|
if !self.storage_map.has_piece_bit(i as usize, offset) {
|
||||||
|
let request_length = if (offset + REQUEST_SIZE) > length {
|
||||||
|
length - offset
|
||||||
|
} else {
|
||||||
|
REQUEST_SIZE
|
||||||
|
};
|
||||||
|
|
||||||
|
let msg = SelectionMessage::new(i, offset as u32, request_length as u32);
|
||||||
|
|
||||||
|
if self.peers.get_mut(peer).unwrap().requested(msg) {
|
||||||
|
self.queue.push_back((*peer, Message::Request(msg)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
offset += REQUEST_SIZE;
|
||||||
|
}
|
||||||
|
|
||||||
|
if self.peers[peer].count_open_requests() >= 25 {
|
||||||
|
peers.remove(peer);
|
||||||
|
}
|
||||||
|
|
||||||
|
if offset >= length {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,34 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
|
use url::Url;
|
||||||
|
|
||||||
|
pub struct TrackerManager {
|
||||||
|
counter: AtomicUsize,
|
||||||
|
trackers: HashMap<usize, Tracker>,
|
||||||
|
url_index: HashMap<String, usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct Tracker {
|
||||||
|
url: Url,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub type TrackerId = usize;
|
||||||
|
|
||||||
|
impl TrackerManager {
|
||||||
|
pub fn get_tracker_id(&mut self, url: &Url) -> Option<TrackerId> {
|
||||||
|
// dht is not a tracker bye
|
||||||
|
if url.scheme() == "dht" {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
let url_str = url.as_str();
|
||||||
|
if let Some(id) = self.url_index.get(url_str) {
|
||||||
|
return Some(*id);
|
||||||
|
}
|
||||||
|
|
||||||
|
let new_id = self.counter.fetch_add(1, Ordering::AcqRel);
|
||||||
|
self.url_index.insert(url_str.to_string(), new_id);
|
||||||
|
|
||||||
|
Some(new_id)
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,292 @@
|
|||||||
|
use bytes::{Bytes, BytesMut};
|
||||||
|
use std::convert::TryInto;
|
||||||
|
use std::error::Error;
|
||||||
|
use std::fmt::{Debug, Display, Formatter};
|
||||||
|
use torment_core::infohash::v1::U160;
|
||||||
|
use torment_core::infohash::InfoHashCapable;
|
||||||
|
use torment_core::Bitfield;
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, Ord, PartialOrd, Eq, PartialEq)]
|
||||||
|
pub enum Message {
|
||||||
|
Choke,
|
||||||
|
Unchoke,
|
||||||
|
Interested,
|
||||||
|
NotInterested,
|
||||||
|
Have(u32),
|
||||||
|
Bitfield(Bitfield),
|
||||||
|
Request(SelectionMessage),
|
||||||
|
Cancel(SelectionMessage),
|
||||||
|
Piece(PieceMessage),
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug)]
|
||||||
|
pub enum MessageParsingError {
|
||||||
|
TooShort,
|
||||||
|
UnknownType,
|
||||||
|
InvalidPrefix,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Error for MessageParsingError {}
|
||||||
|
|
||||||
|
impl Display for MessageParsingError {
|
||||||
|
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||||
|
write!(f, "{:?}", self)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Message {
|
||||||
|
pub fn from_bytes(bytes: Bytes) -> Result<Message, MessageParsingError> {
|
||||||
|
Ok(match bytes[0] {
|
||||||
|
0 => Message::Choke,
|
||||||
|
1 => Message::Unchoke,
|
||||||
|
2 => Message::Interested,
|
||||||
|
3 => Message::NotInterested,
|
||||||
|
4 => {
|
||||||
|
if bytes.len() < 5 {
|
||||||
|
return Err(MessageParsingError::TooShort);
|
||||||
|
}
|
||||||
|
|
||||||
|
Message::Have(u32::from_be_bytes(bytes[1..5].try_into().unwrap()))
|
||||||
|
}
|
||||||
|
|
||||||
|
5 => Message::Bitfield(Bitfield::new(bytes.slice(1..))),
|
||||||
|
|
||||||
|
6 | 8 => {
|
||||||
|
if bytes.len() < 13 {
|
||||||
|
return Err(MessageParsingError::TooShort);
|
||||||
|
}
|
||||||
|
|
||||||
|
let selection = SelectionMessage {
|
||||||
|
index: u32::from_be_bytes(bytes[1..5].try_into().unwrap()),
|
||||||
|
offset: u32::from_be_bytes(bytes[5..9].try_into().unwrap()),
|
||||||
|
length: u32::from_be_bytes(bytes[9..13].try_into().unwrap()),
|
||||||
|
};
|
||||||
|
|
||||||
|
if bytes[0] == 6 {
|
||||||
|
Message::Request(selection)
|
||||||
|
} else {
|
||||||
|
Message::Cancel(selection)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
7 => {
|
||||||
|
if bytes.len() < 9 {
|
||||||
|
return Err(MessageParsingError::TooShort);
|
||||||
|
}
|
||||||
|
|
||||||
|
Message::Piece(PieceMessage {
|
||||||
|
index: u32::from_be_bytes(bytes[1..5].try_into().unwrap()),
|
||||||
|
offset: u32::from_be_bytes(bytes[5..9].try_into().unwrap()),
|
||||||
|
piece: bytes.slice(9..),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
_ => return Err(MessageParsingError::UnknownType),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn to_length_prefixed_bytes(&self) -> Vec<u8> {
|
||||||
|
let buffer = self.to_bytes();
|
||||||
|
let size = (buffer.len() as u32).to_be_bytes();
|
||||||
|
let mut map = Vec::with_capacity(buffer.len() + 4);
|
||||||
|
map.extend_from_slice(&size);
|
||||||
|
map.extend(buffer);
|
||||||
|
map
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn to_bytes(&self) -> Bytes {
|
||||||
|
match self {
|
||||||
|
Message::Choke => Bytes::from_static(&[0u8]),
|
||||||
|
Message::Unchoke => Bytes::from_static(&[1u8]),
|
||||||
|
Message::Interested => Bytes::from_static(&[2u8]),
|
||||||
|
Message::NotInterested => Bytes::from_static(&[3u8]),
|
||||||
|
Message::Have(piece) => {
|
||||||
|
let mut buffer = BytesMut::with_capacity(5);
|
||||||
|
buffer.extend(&[4]);
|
||||||
|
buffer.extend_from_slice(&piece.to_be_bytes());
|
||||||
|
buffer.freeze()
|
||||||
|
}
|
||||||
|
Message::Bitfield(bitfield) => {
|
||||||
|
let mut buffer = BytesMut::with_capacity(1 + bitfield.as_ref().len());
|
||||||
|
buffer.extend(&[5]);
|
||||||
|
buffer.extend_from_slice(bitfield.as_ref());
|
||||||
|
buffer.freeze()
|
||||||
|
}
|
||||||
|
Message::Request(request) => {
|
||||||
|
let mut buffer = BytesMut::with_capacity(13);
|
||||||
|
buffer.extend(&[6]);
|
||||||
|
buffer.extend_from_slice(&request.index.to_be_bytes());
|
||||||
|
buffer.extend_from_slice(&request.offset.to_be_bytes());
|
||||||
|
buffer.extend_from_slice(&request.length.to_be_bytes());
|
||||||
|
buffer.freeze()
|
||||||
|
}
|
||||||
|
Message::Piece(piece) => {
|
||||||
|
let mut buffer = BytesMut::with_capacity(9 + piece.piece.len());
|
||||||
|
buffer.extend(&[7]);
|
||||||
|
buffer.extend_from_slice(&piece.index.to_be_bytes());
|
||||||
|
buffer.extend_from_slice(&piece.offset.to_be_bytes());
|
||||||
|
buffer.extend_from_slice(piece.piece.as_ref());
|
||||||
|
buffer.freeze()
|
||||||
|
}
|
||||||
|
Message::Cancel(cancel) => {
|
||||||
|
let mut buffer = BytesMut::with_capacity(13);
|
||||||
|
buffer.extend(&[8]);
|
||||||
|
buffer.extend_from_slice(&cancel.index.to_be_bytes());
|
||||||
|
buffer.extend_from_slice(&cancel.offset.to_be_bytes());
|
||||||
|
buffer.extend_from_slice(&cancel.length.to_be_bytes());
|
||||||
|
buffer.freeze()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, Copy)]
|
||||||
|
pub struct Handshake {
|
||||||
|
reserved: [u8; 8],
|
||||||
|
info_hash: U160,
|
||||||
|
peer_id: U160,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Handshake {
|
||||||
|
pub fn from_bytes(bytes: Bytes) -> Result<Handshake, MessageParsingError> {
|
||||||
|
if bytes.len() < 68 {
|
||||||
|
return Err(MessageParsingError::TooShort);
|
||||||
|
}
|
||||||
|
|
||||||
|
if &bytes[0..20] != b"\x13BitTorrent protocol" {
|
||||||
|
return Err(MessageParsingError::InvalidPrefix);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(Handshake {
|
||||||
|
reserved: bytes[20..28].try_into().unwrap(),
|
||||||
|
info_hash: U160::from_bytes(&bytes[28..48]).unwrap(),
|
||||||
|
peer_id: U160::from_bytes(&bytes[48..68]).unwrap(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn to_bytes(&self) -> [u8; 68] {
|
||||||
|
let mut header =
|
||||||
|
*b"\x13BitTorrent protocol\0\0\0\0\0\0\0\0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
|
||||||
|
header[28..48].copy_from_slice(self.info_hash.to_bytes().as_ref());
|
||||||
|
header[48..68].copy_from_slice(self.peer_id.to_bytes().as_ref());
|
||||||
|
header
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn new(peer_id: U160, info_hash: U160) -> Handshake {
|
||||||
|
Handshake {
|
||||||
|
reserved: [0; 8],
|
||||||
|
info_hash,
|
||||||
|
peer_id,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn peer_id(&self) -> U160 {
|
||||||
|
self.peer_id
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn info_hash(&self) -> U160 {
|
||||||
|
self.info_hash
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, Hash, Copy, Ord, PartialOrd, Eq, PartialEq)]
|
||||||
|
pub struct SelectionMessage {
|
||||||
|
index: u32,
|
||||||
|
offset: u32,
|
||||||
|
length: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SelectionMessage {
|
||||||
|
pub fn new(index: u32, offset: u32, length: u32) -> SelectionMessage {
|
||||||
|
SelectionMessage {
|
||||||
|
index,
|
||||||
|
offset,
|
||||||
|
length,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn index(&self) -> u32 {
|
||||||
|
self.index
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn offset(&self) -> u32 {
|
||||||
|
self.offset
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn length(&self) -> u32 {
|
||||||
|
self.length
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Ord, PartialOrd, Eq, PartialEq)]
|
||||||
|
pub struct PieceMessage {
|
||||||
|
pub index: u32,
|
||||||
|
pub offset: u32,
|
||||||
|
pub piece: Bytes,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Debug for PieceMessage {
|
||||||
|
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||||
|
f.debug_struct("PieceMessage")
|
||||||
|
.field("index", &self.index)
|
||||||
|
.field("offset", &self.offset)
|
||||||
|
.field("piece", &self.piece.len())
|
||||||
|
.finish()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PieceMessage {
|
||||||
|
pub fn index(&self) -> u32 {
|
||||||
|
self.index
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn offset(&self) -> u32 {
|
||||||
|
self.offset
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn length(&self) -> u32 {
|
||||||
|
self.piece.len() as u32
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn piece(&self) -> Bytes {
|
||||||
|
self.piece.clone()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use crate::message::{Message, PieceMessage, SelectionMessage};
|
||||||
|
use bytes::Bytes;
|
||||||
|
use torment_core::Bitfield;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn round_trip() {
|
||||||
|
let msgs = [
|
||||||
|
Message::Choke,
|
||||||
|
Message::Unchoke,
|
||||||
|
Message::Interested,
|
||||||
|
Message::NotInterested,
|
||||||
|
Message::Have(42),
|
||||||
|
Message::Bitfield(Bitfield::with_size(4)),
|
||||||
|
Message::Request(SelectionMessage {
|
||||||
|
index: 69,
|
||||||
|
offset: 1337,
|
||||||
|
length: 42,
|
||||||
|
}),
|
||||||
|
Message::Piece(PieceMessage {
|
||||||
|
index: 69,
|
||||||
|
offset: 1337,
|
||||||
|
piece: Bytes::from_static(b"hewwo"),
|
||||||
|
}),
|
||||||
|
Message::Cancel(SelectionMessage {
|
||||||
|
index: 69,
|
||||||
|
offset: 1337,
|
||||||
|
length: 42,
|
||||||
|
}),
|
||||||
|
];
|
||||||
|
|
||||||
|
for msg in &msgs {
|
||||||
|
assert_eq!(msg, &Message::from_bytes(msg.to_bytes()).unwrap())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,13 @@
|
|||||||
|
[package]
|
||||||
|
name = "torment-storage"
|
||||||
|
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" }
|
||||||
|
remem = "0.1.0"
|
||||||
|
bytes = "0.5.6"
|
||||||
|
lazy_static = "1.4.0"
|
||||||
@ -0,0 +1,348 @@
|
|||||||
|
use bytes::{Bytes, BytesMut};
|
||||||
|
use lazy_static::lazy_static;
|
||||||
|
use remem::{ItemGuard, Pool};
|
||||||
|
use std::borrow::Borrow;
|
||||||
|
use std::cmp::min;
|
||||||
|
use std::collections::{BTreeMap, HashMap, HashSet};
|
||||||
|
use std::fmt::{Debug, Formatter};
|
||||||
|
use std::fs::{DirBuilder, File, OpenOptions};
|
||||||
|
use std::io;
|
||||||
|
use std::io::{Read, Seek, SeekFrom, Write};
|
||||||
|
use std::ops::{Add, Deref, Range};
|
||||||
|
use std::path::{Path, PathBuf, MAIN_SEPARATOR};
|
||||||
|
use std::rc::Rc;
|
||||||
|
use std::sync::Mutex;
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
use torment_core::metainfo::{MetaInfo, MetaInfoObject, Torrent};
|
||||||
|
use torment_core::utils::EphemeralMap;
|
||||||
|
|
||||||
|
lazy_static! {
|
||||||
|
static ref MEMORY_POOL: Pool<Vec<u8>> = Pool::new(|| vec![0u8; 4194304]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct StorageMap {
|
||||||
|
base_path: PathBuf,
|
||||||
|
piece_length: usize,
|
||||||
|
pieces: usize,
|
||||||
|
open_files: HashMap<usize, Rc<Mutex<File>>>,
|
||||||
|
buffer: HashMap<usize, StoragePiece>,
|
||||||
|
mapping: BTreeMap<usize, StorageMapping>,
|
||||||
|
size: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct StoragePiece {
|
||||||
|
index: usize,
|
||||||
|
pieces: usize,
|
||||||
|
ranges: HashSet<usize>,
|
||||||
|
buffer: ItemGuard<'static, Vec<u8>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl StoragePiece {
|
||||||
|
fn is_complete(&self) -> bool {
|
||||||
|
self.pieces == self.ranges.len()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Debug for StoragePiece {
|
||||||
|
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||||
|
f.debug_struct("StoragePiece")
|
||||||
|
.field("index", &self.index)
|
||||||
|
.field("ranges", &self.ranges)
|
||||||
|
.finish()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl StorageMap {
|
||||||
|
fn get_offsets(&self, offset: usize, length: usize) -> Vec<usize> {
|
||||||
|
let x = 0..1;
|
||||||
|
let end = offset + length;
|
||||||
|
let mut range = self.mapping.range(offset..);
|
||||||
|
let mut items = vec![];
|
||||||
|
while let Some((item_end, _)) = range.next() {
|
||||||
|
items.push(*item_end);
|
||||||
|
|
||||||
|
if *item_end > end {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
items
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_file(&mut self, item: &StorageMapping) -> io::Result<Rc<Mutex<File>>> {
|
||||||
|
if let Some(file) = self.open_files.get(&item.offset) {
|
||||||
|
return Ok(Rc::clone(file));
|
||||||
|
}
|
||||||
|
|
||||||
|
let path = item.path.join(&format!("{}", MAIN_SEPARATOR));
|
||||||
|
let file_path = self.base_path.join(path);
|
||||||
|
println!("Opening file: {:?}", file_path);
|
||||||
|
|
||||||
|
if let Some(parent) = file_path.parent() {
|
||||||
|
DirBuilder::new().recursive(true).create(parent)?;
|
||||||
|
}
|
||||||
|
|
||||||
|
let file = OpenOptions::new()
|
||||||
|
.write(true)
|
||||||
|
.read(true)
|
||||||
|
.create(true)
|
||||||
|
.open(file_path.as_path())?;
|
||||||
|
|
||||||
|
let rc = Rc::new(Mutex::new(file));
|
||||||
|
self.open_files.insert(item.offset, Rc::clone(&rc));
|
||||||
|
Ok(rc)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn internal_write<T: AsRef<[u8]>>(&mut self, offset: usize, data: T) -> io::Result<()> {
|
||||||
|
let data = data.as_ref();
|
||||||
|
let mut written = 0;
|
||||||
|
for item_end in self.get_offsets(offset, data.len()) {
|
||||||
|
let item = self.mapping.get(&item_end).unwrap().clone();
|
||||||
|
let file = self.get_file(&item)?;
|
||||||
|
let mut file = file.lock().unwrap();
|
||||||
|
let curr_offset = (offset + written) - item.offset;
|
||||||
|
file.seek(SeekFrom::Start(curr_offset as u64))?;
|
||||||
|
let to_write = min(item.length - curr_offset, data.len() - written);
|
||||||
|
file.write_all(&data[written..written + to_write])?;
|
||||||
|
written += to_write;
|
||||||
|
}
|
||||||
|
|
||||||
|
assert_eq!(written, data.len(), "Failed to write all data");
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn internal_read<T: AsMut<[u8]>>(&mut self, offset: usize, mut data: T) -> io::Result<()> {
|
||||||
|
let data = data.as_mut();
|
||||||
|
let mut readed = 0;
|
||||||
|
for item_end in self.get_offsets(offset, data.len()) {
|
||||||
|
let item = self.mapping.get(&item_end).unwrap().clone();
|
||||||
|
let file = self.get_file(&item)?;
|
||||||
|
let mut file = file.lock().unwrap();
|
||||||
|
let curr_offset = (offset + readed) - item.offset;
|
||||||
|
file.seek(SeekFrom::Start(curr_offset as u64))?;
|
||||||
|
let to_read = min(item.length - curr_offset, data.len() - readed);
|
||||||
|
file.read_exact(&mut data[readed..readed + to_read])?;
|
||||||
|
readed += to_read;
|
||||||
|
}
|
||||||
|
|
||||||
|
assert_eq!(readed, data.len(), "Failed to read all data");
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn has_piece(&self, index: usize) -> bool {
|
||||||
|
if let Some(piece) = self.buffer.get(&index) {
|
||||||
|
piece.is_complete()
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn read_piece(&mut self, index: usize) -> io::Result<Vec<u8>> {
|
||||||
|
let mut bytes = vec![0; self.get_piece_length(index)];
|
||||||
|
self.read(index, 0, &mut bytes);
|
||||||
|
Ok(bytes)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn read<T: AsMut<[u8]>>(
|
||||||
|
&mut self,
|
||||||
|
index: usize,
|
||||||
|
offset: usize,
|
||||||
|
mut data: T,
|
||||||
|
) -> io::Result<()> {
|
||||||
|
let data = data.as_mut();
|
||||||
|
self.internal_read((index * self.piece_length) + offset, data)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn write<T: AsRef<[u8]>>(
|
||||||
|
&mut self,
|
||||||
|
index: usize,
|
||||||
|
offset: usize,
|
||||||
|
data: T,
|
||||||
|
) -> io::Result<bool> {
|
||||||
|
let data = data.as_ref();
|
||||||
|
let item = if let Some(item) = self.buffer.get_mut(&index) {
|
||||||
|
item
|
||||||
|
} else {
|
||||||
|
let request_size = 2usize.pow(14);
|
||||||
|
let piece_length = self.get_piece_length(index);
|
||||||
|
let pieces = (piece_length / request_size)
|
||||||
|
+ if piece_length % request_size > 0 {
|
||||||
|
1
|
||||||
|
} else {
|
||||||
|
0
|
||||||
|
};
|
||||||
|
|
||||||
|
self.buffer.insert(index, {
|
||||||
|
StoragePiece {
|
||||||
|
index,
|
||||||
|
ranges: Default::default(),
|
||||||
|
buffer: MEMORY_POOL.get(),
|
||||||
|
pieces,
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
self.buffer.get_mut(&index).unwrap()
|
||||||
|
};
|
||||||
|
|
||||||
|
item.buffer[offset..offset + data.len()].copy_from_slice(data);
|
||||||
|
item.ranges.insert(offset);
|
||||||
|
|
||||||
|
if item.is_complete() {
|
||||||
|
if let Some(item) = self.buffer.remove(&index) {
|
||||||
|
self.internal_write(
|
||||||
|
item.index * self.piece_length,
|
||||||
|
&item.buffer[0..self.get_piece_length(item.index)],
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(true)
|
||||||
|
} else {
|
||||||
|
Ok(false)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn wipe_piece(&mut self, index: usize) {
|
||||||
|
self.buffer.remove(&index);
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn get_piece_length(&self, index: usize) -> usize {
|
||||||
|
if index + 1 == self.pieces {
|
||||||
|
let len = (self.size % self.piece_length);
|
||||||
|
if len == 0 {
|
||||||
|
self.piece_length
|
||||||
|
} else {
|
||||||
|
len
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
self.piece_length
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn has_piece_bit(&self, index: usize, offset: usize) -> bool {
|
||||||
|
self.buffer
|
||||||
|
.get(&index)
|
||||||
|
.map(|item| item.ranges.contains(&offset))
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn house_keeping(&mut self) {
|
||||||
|
// self.open_files.clean()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn size(&self) -> usize {
|
||||||
|
self.size
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub trait ToStorageMap {
|
||||||
|
fn to_storage_map<P: AsRef<Path>>(&self, path: P, is_base_path: bool) -> StorageMap;
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ToStorageMap for Torrent {
|
||||||
|
fn to_storage_map<P: AsRef<Path>>(&self, path: P, is_base_path: bool) -> StorageMap {
|
||||||
|
let name = self.name().clone();
|
||||||
|
match self.meta_info().object() {
|
||||||
|
MetaInfoObject::File(size) => {
|
||||||
|
if is_base_path {
|
||||||
|
StorageMapBuilder::create(path, self.meta_info().piece_length())
|
||||||
|
.insert(vec![name.to_string()], *size)
|
||||||
|
.build()
|
||||||
|
} else {
|
||||||
|
let path = path.as_ref().to_path_buf();
|
||||||
|
let parent = path.parent().unwrap_or(if cfg!(target_os = "win") {
|
||||||
|
Path::new("C:\\")
|
||||||
|
} else {
|
||||||
|
Path::new("/")
|
||||||
|
});
|
||||||
|
|
||||||
|
StorageMapBuilder::create(parent, self.meta_info().piece_length())
|
||||||
|
.insert(
|
||||||
|
vec![path
|
||||||
|
.file_name()
|
||||||
|
.map(|os_str| os_str.to_string_lossy().to_string())
|
||||||
|
.unwrap_or_else(|| self.name().to_string())],
|
||||||
|
*size,
|
||||||
|
)
|
||||||
|
.build()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
MetaInfoObject::Files(files) => {
|
||||||
|
let mut builder = if is_base_path {
|
||||||
|
StorageMapBuilder::create(path, self.meta_info().piece_length())
|
||||||
|
} else {
|
||||||
|
StorageMapBuilder::create(
|
||||||
|
path.as_ref().join(self.name()),
|
||||||
|
self.meta_info().piece_length(),
|
||||||
|
)
|
||||||
|
};
|
||||||
|
|
||||||
|
for file in files {
|
||||||
|
builder = builder.insert(file.path().to_vec(), file.length());
|
||||||
|
}
|
||||||
|
|
||||||
|
builder.build()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct StorageMapping {
|
||||||
|
path: Vec<String>,
|
||||||
|
length: usize,
|
||||||
|
offset: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct StorageMapBuilder {
|
||||||
|
base_path: PathBuf,
|
||||||
|
offset: usize,
|
||||||
|
piece_length: usize,
|
||||||
|
items: BTreeMap<usize, StorageMapping>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl StorageMapBuilder {
|
||||||
|
pub fn create<P: AsRef<Path>>(path: P, piece_length: usize) -> StorageMapBuilder {
|
||||||
|
StorageMapBuilder {
|
||||||
|
base_path: path.as_ref().to_path_buf(),
|
||||||
|
offset: 0,
|
||||||
|
piece_length,
|
||||||
|
items: BTreeMap::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn insert(mut self, path: Vec<String>, length: usize) -> Self {
|
||||||
|
let offset = self.offset;
|
||||||
|
self.offset += length;
|
||||||
|
self.items.insert(
|
||||||
|
self.offset - 1,
|
||||||
|
StorageMapping {
|
||||||
|
offset,
|
||||||
|
length,
|
||||||
|
path,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn build(self) -> StorageMap {
|
||||||
|
StorageMap {
|
||||||
|
base_path: self.base_path,
|
||||||
|
piece_length: self.piece_length,
|
||||||
|
pieces: (self.offset / self.piece_length)
|
||||||
|
+ if self.offset % self.piece_length > 0 {
|
||||||
|
1
|
||||||
|
} else {
|
||||||
|
0
|
||||||
|
},
|
||||||
|
open_files: Default::default(),
|
||||||
|
buffer: Default::default(),
|
||||||
|
mapping: self.items,
|
||||||
|
size: self.offset,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue