|
|
use async_std::net::{SocketAddr, ToSocketAddrs, UdpSocket, TcpStream};
use async_std::prelude::*;
use async_std::io;
use async_std::task::spawn;
// use std::io::prelude::*;
// use std::net::TcpStream;
use std::error::Error;
use std::net::Ipv4Addr;
use std::convert::TryInto;
use std::str;
use byteorder::{BigEndian, ByteOrder};
use rand::Rng;
use sha1::{Sha1, Digest};
extern crate serde;
extern crate serde_bencode;
#[macro_use]
extern crate serde_derive;
extern crate serde_bytes;
use serde_bencode::{de, ser};
use serde_bytes::ByteBuf;
// use rustc_hex::ToHex;
const CONNECT_MSG: u32 = 0;
const ANNOUNCE_MSG: u32 = 1;
const BLOCK_LEN: u32 = 16384; // 2**14
#[derive(Debug, Deserialize)]
struct Node(String, i64);
#[derive(Debug, Serialize, Deserialize)]
struct File {
path: Vec<String>,
length: i64,
#[serde(default)]
md5sum: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
struct Info {
name: String,
pieces: ByteBuf,
#[serde(rename = "piece length")]
piece_length: i64,
#[serde(default)]
md5sum: Option<String>,
#[serde(default)]
length: Option<i64>,
#[serde(default)]
files: Option<Vec<File>>,
#[serde(default)]
private: Option<u8>,
#[serde(default)]
path: Option<Vec<String>>,
#[serde(default)]
#[serde(rename = "root hash")]
root_hash: Option<String>,
}
#[derive(Debug, Deserialize)]
struct Torrent {
info: Info,
#[serde(default)]
announce: Option<String>,
#[serde(default)]
nodes: Option<Vec<Node>>,
#[serde(default)]
encoding: Option<String>,
#[serde(default)]
httpseeds: Option<Vec<String>>,
#[serde(default)]
#[serde(rename = "announce-list")]
announce_list: Option<Vec<Vec<String>>>,
#[serde(default)]
#[serde(rename = "creation date")]
creation_date: Option<i64>,
#[serde(rename = "comment")]
comment: Option<String>,
#[serde(default)]
#[serde(rename = "created by")]
created_by: Option<String>,
}
impl Torrent {
pub fn open(path: &str) -> Torrent {
let b = std::fs::read(path).unwrap();
de::from_bytes::<Torrent>(&b).unwrap()
}
pub fn size(&self) -> u64 {
let files = self.info.files.as_ref().unwrap();
let mut size: u64 = 0;
for i in 0..files.len() {
size = size + files[i].length as u64;
}
size
}
pub fn info_hash(&self) -> Vec<u8> {
let info_bytes: Vec<u8> = ser::to_bytes(&self.info).unwrap();
let mut hasher = Sha1::new();
hasher.update(info_bytes);
let result = hasher.finalize();
result.to_vec()
}
pub fn piece_len(&self, piece_index: u32) -> u32 {
let total_len = self.size() as u32;
let piece_len = self.info.piece_length as u32;
let last_piece_len = total_len % piece_len;
let last_piece_index = ((total_len / piece_len) as f32).floor() as u32;
if last_piece_index == piece_index {
return last_piece_len;
}
piece_len
}
pub fn blocks_per_piece(&self, piece_index: u32) -> u32 {
let piece_len = self.piece_len(piece_index);
((piece_len as u32 / BLOCK_LEN) as f32).floor() as u32
}
pub fn block_len(&self, piece_index: u32, block_index: u32) -> u32 {
let piece_len = self.piece_len(piece_index);
let last_piece_len = piece_len % BLOCK_LEN;
let last_piece_index = ((piece_len / BLOCK_LEN) as f32).floor() as u32;
if block_index==last_piece_index {
return last_piece_len;
}
BLOCK_LEN
}
pub fn announce(&self) -> Option<String> {
let mut a: String = "".to_string();
match &self.announce {
Some(s) => {
a = s.to_string();
}
_ => (),
};
match &self.announce_list { // TODO refactor approach
Some(list) => {
a = list[0][0].clone();
}
_ => (),
};
Some(a)
}
pub fn protocol(&self) -> Option<String> {
let announce = &self.announce().unwrap();
let aux: Vec<&str> = announce.split(":").collect();
let protocol = aux[0];
Some(protocol.to_string())
}
pub async fn get_peers(&self) -> Result<Vec<Peer>, Box<dyn Error>> {
// TODO timming system to resend request if no answer in X seconds
if self.protocol().unwrap() != "udp" {
panic!("not udp: {:?}", self.protocol().unwrap()); // TODO remove panic, return error
}
let peer = self.announce.clone().unwrap().replace("udp://", "");
println!("peer {:?}", peer);
let socket = UdpSocket::bind("0.0.0.0:6681").await?;
let conn_req = msg_conn_req().to_vec();
println!("SENDING conn_req {:?}", conn_req);
socket.send_to(&conn_req, &peer).await?;
let mut buf = vec![0u8; 1024];
let peers = loop {
let (n, src) = socket.recv_from(&mut buf).await?;
let typ = msg_resp_type(&buf);
if typ == CONNECT_MSG {
println!("TYPE: CONNECT: {:?}", CONNECT_MSG);
// println!("HEX {:?}", &buf[0..n].to_hex());
let conn_resp = msg_parse_connect_resp(&buf[0..n].to_vec());
let announce_req = msg_ann_req(conn_resp.connection_id, self, 6681);
socket.send_to(&announce_req[..], &src).await?;
} else if typ==ANNOUNCE_MSG {
println!("TYPE: ANNOUNCE: {:?}", ANNOUNCE_MSG);
// println!("HEX {:?}", &buf[0..n].to_hex());
let ann_resp = msg_parse_announce_resp(&buf[0..n].to_vec());
break ann_resp.peers;
}
};
Ok(peers)
}
}
struct Block {
length: u32,
bytes: Vec<u8>,
}
struct Payload {
index: u32,
begin: u32,
length: u32,
block: Block,
}
struct Msg {
size: u32,
id: u8,
payload: Payload,
}
fn msg_parse(b: Vec<u8>) -> Result<Msg, &'static str> {
if b.len()<4 {
return Err("No id in msg.");
}
let id = b[4];
if b.len()<5 {
return Err("No payload in msg.");
}
let mut payload: Payload = Payload{
index: 0,
begin: 0,
length: 0,
block: Block{
length: 0,
bytes: Vec::new()
}
};
match id {
6 | 7 | 8 => {
payload.index = BigEndian::read_u32(&b[5..9]);
payload.begin = BigEndian::read_u32(&b[9..13]);
if id==7 {
payload.length = BigEndian::read_u32(&b[13..]);
}
}
_ => ()
};
Ok(Msg{
size: BigEndian::read_u32(&b[0..4]),
id: id,
payload: payload,
})
}
async fn msg_handler<'a>(
b: Vec<u8>,
mut stream: TcpStream,
torrent: &Torrent,
pieces: &mut Pieces,
queue: &mut Queue<'a>,
file: String,
) -> Result<(), Box<dyn Error>> {
if is_handshake(&b) {
stream.write(&msg_interested()).await?;
println!("IS HANDSHAKE");
return Ok(())
}
let m = msg_parse(b).unwrap();
match m.id {
0 => handler_choke(stream),
1 => handler_unchoke(stream, pieces, queue),
4 => handler_have(stream, pieces, queue, m.payload),
5 => handler_bitfield(stream, pieces, queue, m.payload),
7 => handler_piece(stream, pieces, queue, torrent, file, m.payload),
_ => {},
}
Ok(())
}
fn is_handshake(b: &Vec<u8>) -> bool {
b.len()==(b[0] as usize) +49 && b.len()==68 && &b[1..20] == b"BitTorrent protocol"
}
fn handler_choke(mut stream: TcpStream) {
// stream.shutdown()
}
fn handler_unchoke(mut stream: TcpStream, pieces: &mut Pieces, queue: &mut Queue) {
queue.chocked = false;
request_piece(stream, pieces, queue);
}
fn handler_have(mut stream: TcpStream, pieces: &mut Pieces, queue: &mut Queue, payload: Payload) {
}
fn handler_bitfield(mut stream: TcpStream, pieces: &mut Pieces, queue: &mut Queue, payload: Payload) {
}
fn handler_piece(mut stream: TcpStream, pieces: &mut Pieces, queue: &mut Queue, torrent: &Torrent, file: String, payload: Payload) {
}
async fn request_piece<'a>(mut stream: TcpStream, pieces: &mut Pieces, queue: &mut Queue<'a>) -> Result<(), Box<dyn Error>> {
if queue.chocked {
// return
}
while queue.length()>0 {
let piece_block = queue.get_last();
if pieces.needed(&piece_block) {
stream.write(&msg_request(&piece_block)).await?;
pieces.add_requested(&piece_block);
break;
}
};
Ok(())
}
fn msg_conn_req() -> [u8; 16] {
let mut b: [u8; 16] = [0; 16];
BigEndian::write_u64(&mut b[0..8], 0x41727101980); // connection_id
let random_bytes = rand::thread_rng().gen::<[u8; 4]>();
b[12..16].clone_from_slice(&random_bytes[..]);
b
}
fn msg_ann_req(connection_id: u64, torrent: &Torrent, port: u16) -> Vec<u8> {
let mut b: [u8; 98] = [0; 98];
BigEndian::write_u64(&mut b[0..8], connection_id);
BigEndian::write_u32(&mut b[8..12], ANNOUNCE_MSG); // action
let random_bytes = rand::thread_rng().gen::<[u8; 4]>();
b[12..16].clone_from_slice(&random_bytes[..]);
b[16..36].clone_from_slice(&torrent.info_hash());
println!("info_hash: {:?}", &b[16..36]);
// TODO [36..56] peerId b[36..56].clone_from_slice("todo".as_bytes());
// TODO [56..64] downloaded
println!("torrent.size(): {:?}", torrent.size());
BigEndian::write_u64(&mut b[64..72], torrent.size()); // [64..72] left
// TODO [72..80] uploaded
// TODO [80..84] event (0:none, 1:completed, 2:started, 3:stopped)
// TODO [84..88] ip address
let random_bytes = rand::thread_rng().gen::<[u8; 4]>();
b[88..92].clone_from_slice(&random_bytes[..]); // [88..92] key (random)
BigEndian::write_i32(&mut b[92..96], -1); // [92..96] num_want (default: -1)
BigEndian::write_u16(&mut b[96..98], port); // [96..98] port
b.to_vec()
}
fn msg_resp_type(b: &Vec<u8>) -> u32 {
let action = BigEndian::read_u32(&b[0..4]);
if action == 0 {
return CONNECT_MSG;
} else {
return ANNOUNCE_MSG;
}
}
#[derive(Debug)]
struct ConnResp {
action: u32,
transaction_id: u32,
connection_id: u64,
}
fn msg_parse_connect_resp(b: &Vec<u8>) -> ConnResp {
ConnResp {
action: BigEndian::read_u32(&b[0..4]),
transaction_id: BigEndian::read_u32(&b[4..8]),
connection_id: BigEndian::read_u64(&b[8..]),
}
}
#[derive(Debug)]
struct Peer {
ip: String,
port: u16,
}
#[derive(Debug)]
struct AnnResp {
action: u32,
transaction_id: u32,
interval: u32,
leechers: u32,
seeders: u32,
peers: Vec<Peer>,
}
fn msg_parse_announce_resp(b: &Vec<u8>) -> AnnResp {
let mut peers: Vec<Peer> = Vec::new();
let n_peers = (b.len() - 20) / 6;
for i in 0..n_peers {
let peer: Peer = Peer {
// ip: BigEndian::read_u32(&b[20+(6*i)..24+(6*i)]),
ip: bytes_to_ip(
&b[20 + (6 * i)..24 + (6 * i)]
.try_into()
.expect("err parsing peer ip"),
),
port: BigEndian::read_u16(&b[24 + (6 * i)..26 + (6 * i)]),
};
peers.push(peer);
}
let ann_resp: AnnResp = AnnResp {
action: BigEndian::read_u32(&b[0..4]),
transaction_id: BigEndian::read_u32(&b[4..8]),
interval: BigEndian::read_u32(&b[8..12]),
leechers: BigEndian::read_u32(&b[12..16]),
seeders: BigEndian::read_u32(&b[16..20]),
peers: peers,
};
ann_resp
}
fn bytes_to_ip(b: &[u8; 4]) -> String {
Ipv4Addr::new(b[0], b[1], b[2], b[3]).to_string()
}
fn msg_handshake(info_hash: Vec<u8>) -> Vec<u8> {
let mut b: [u8; 68] = [0; 68];
b[0] = 19;
b[1..20].clone_from_slice(b"BitTorrent protocol"); // TODO use constant for the string
b[28..48].clone_from_slice(&info_hash);
b.to_vec()
}
fn msg_choke() -> Vec<u8> {
let mut b: [u8; 5] = [0; 5];
BigEndian::write_u32(&mut b[0..4], 1);
b.to_vec()
}
fn msg_unchoke() -> Vec<u8> {
let mut b: [u8; 5] = [0; 5];
BigEndian::write_u32(&mut b[0..4], 1);
b[4] = 1;
b.to_vec()
}
fn msg_interested() -> Vec<u8> {
let mut b: [u8; 5] = [0; 5];
BigEndian::write_u32(&mut b[0..4], 1);
b[4] = 2;
b.to_vec()
}
fn msg_notinterested() -> Vec<u8> {
let mut b: [u8; 5] = [0; 5];
BigEndian::write_u32(&mut b[0..4], 1);
b[4] = 3;
b.to_vec()
}
fn msg_have(payload: Payload) -> Vec<u8> {
let mut b: [u8; 9] = [0; 9];
BigEndian::write_u32(&mut b[0..4], 5);
b[4] = 4;
BigEndian::write_u32(&mut b[5..9], payload.index); // TODO REVIEW CHECK payload.index
b.to_vec()
}
fn msg_bitfield(bitfield: Vec<u8>) -> Vec<u8> {
let mut b: [u8; 14] = [0; 14];
BigEndian::write_u32(&mut b[0..4], bitfield.len() as u32);
b[4] = 5;
b[5..].clone_from_slice(&bitfield);
b.to_vec()
}
fn msg_request(payload: &Payload) -> Vec<u8> {
let mut b: [u8; 17] = [0; 17];
BigEndian::write_u32(&mut b[0..4], 13);
b[4] = 6;
BigEndian::write_u32(&mut b[5..9], payload.index);
BigEndian::write_u32(&mut b[9..13], payload.begin);
BigEndian::write_u32(&mut b[13..17], payload.length);
b.to_vec()
}
fn msg_piece(payload: Payload) -> Vec<u8> {
let mut b: [u8; 13] = [0; 13]; // TODO payload.block.length
BigEndian::write_u32(&mut b[0..4], payload.block.length+9);
b[4] = 7;
BigEndian::write_u32(&mut b[5..9], payload.index);
BigEndian::write_u32(&mut b[9..13], payload.begin);
b[13..].clone_from_slice(&payload.block.bytes);
b.to_vec()
}
fn msg_cancel(payload: Payload) -> Vec<u8> {
let mut b: [u8; 17] = [0; 17];
BigEndian::write_u32(&mut b[0..4], 13);
b[4] = 8;
BigEndian::write_u32(&mut b[5..9], payload.index);
BigEndian::write_u32(&mut b[9..13], payload.begin);
BigEndian::write_u32(&mut b[13..17], payload.length);
b.to_vec()
}
fn msg_port(port: u16) -> Vec<u8> {
let mut b: [u8; 7] = [0; 7];
BigEndian::write_u32(&mut b[0..4], 3);
b[4] = 9;
BigEndian::write_u16(&mut b[5..7], port);
b.to_vec()
}
#[derive(Debug)]
struct Pieces {
requested: Vec<Vec<bool>>,
received: Vec<Vec<bool>>,
}
impl Pieces {
fn new(torrent: &Torrent) -> Pieces {
let sub_requested: Vec<bool> = vec![false; torrent.info.pieces.len()/20];
let sub_received: Vec<bool> = vec![false; torrent.info.pieces.len()/20];
let requested: Vec<Vec<bool>> = vec![sub_requested; torrent.info.pieces.len()/20];
let received: Vec<Vec<bool>> = vec![sub_received; torrent.info.pieces.len()/20];
Pieces{
requested: requested,
received: received,
}
}
fn add_requested(&mut self, block: &Payload) {
let index = block.begin / BLOCK_LEN;
self.requested[block.index as usize][index as usize] = true;
}
fn add_received(&mut self, block: &Payload) {
let index = block.begin / BLOCK_LEN;
self.requested[block.index as usize][index as usize] = true;
}
fn needed(&mut self, block: &Payload) -> bool {
let mut n = 1;
for i in 0..self.requested.len() {
for j in 0..self.requested[i].len() {
if self.requested[i][j] == false {
n = n*0;
}
}
}
if n==1 {
self.requested = self.received.clone();
}
let index = block.begin / BLOCK_LEN;
!self.requested[block.index as usize][index as usize]
}
fn is_done(&self, ) -> bool {
let mut n = 1;
for i in 0..self.requested.len() {
for j in 0..self.requested[i].len() {
if self.requested[i][j] == false {
n = n*0;
}
}
}
if n==1 {
return true;
}
false
}
fn print_percent_done(&self, ) {
}
}
struct Queue<'a>{
torrent: &'a Torrent,
queue: Vec<Payload>,
chocked: bool,
}
impl<'a> Queue<'a> {
fn new(torrent: &Torrent) -> Queue {
Queue{
torrent: torrent,
queue: Vec::new(),
chocked: true,
}
}
fn queue(&mut self, index: u32) {
let n_blocks = self.torrent.blocks_per_piece(index);
for i in 0..n_blocks {
let piece_block: Payload = Payload{
index: index,
begin: i*BLOCK_LEN,
length: self.torrent.block_len(index, i),
block: Block{
length: 0,
bytes: Vec::new(),
}
};
self.queue.push(piece_block);
}
}
fn get_last(&mut self) -> Payload {
self.queue.remove(self.queue.len())
}
fn peek(&self) -> &Payload {
&self.queue[0]
}
fn length(&self) -> usize {
self.queue.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
use rustc_hex::ToHex;
#[test]
fn torrent_structs() {
let t = Torrent::open("test.torrent");
assert_eq!(t.info.name, "Big Buck Bunny");
assert_eq!(t.announce().unwrap(), "udp://tracker.leechers-paradise.org:6969");
assert_eq!(t.info_hash().to_hex(), "dd8255ecdc7ca55fb0bbf81323d87062db1f6d1c");
assert_eq!(t.size(), 276_445_467);
}
#[test]
fn msgs() {
let t = Torrent::open("test.torrent");
let conn_req = msg_conn_req().to_vec();
assert_eq!(conn_req[0..12].to_hex(), "000004172710198000000000"); // [12..16] is random
let conn_resp_buf = b"00000000d4c575c2c078f9a83418bebc";
let conn_resp = msg_parse_connect_resp(&conn_resp_buf.to_vec());
println!("conn_resp {:?}", conn_resp);
assert_eq!(conn_resp.action, 808464432);
assert_eq!(conn_resp.transaction_id, 808464432);
assert_eq!(conn_resp.connection_id, 7220505182792409906);
let announce_req = msg_ann_req(conn_resp.connection_id, &t, 6681);
assert_eq!(announce_req[0..12].to_hex(), "643463353735633200000001");
assert_eq!(announce_req[16..72].to_hex(), "dd8255ecdc7ca55fb0bbf81323d87062db1f6d1c0000000000000000000000000000000000000000000000000000000000000000107a391b");
assert_eq!(announce_req[92..98].to_hex(), "ffffffff1a19");
let announce_resp_buf = b"000000017a638a60000006e8000000030000002429d412a11ae12d53dcbd1a19d048c0e51ae1d8241960c8d5d5c36e71c8d5c5b974a31ae2c11ef348d4aabc5f370ec8d5bc4e16321ae1b9cee13b2327b94187b66f7cb94186b1c8d5b75308331ae1b068c0e0c263ac6f857911389d27ac821ae19088c2c01ae172fe2468c00a699aebe6720f68a2f99a1ae167fce2141ae163582ee41ae1634921a5d2725b72aaf5c8d559482e8fc9d257fb394cc8d55740e9ebd166524044e11aff51b6dcb6c8d54e2ebe61c3574b87860fc8d54b6e11026dff473afc7a1ec944cd47cd1ae131249b051ae124ff6b521ae105b75c13c8d50551153a1ae105024e09cd52";
let ann_resp = msg_parse_announce_resp(&announce_resp_buf.to_vec());
assert_eq!(ann_resp.peers.len(), 81);
assert_eq!(ann_resp.peers[0].ip, "48.54.101.56");
assert_eq!(ann_resp.peers[0].port, 12336);
assert_eq!(ann_resp.peers[80].ip, "52.101.48.57");
assert_eq!(ann_resp.peers[80].port, 25444);
}
#[test]
fn pieces() {
let t = Torrent::open("test.torrent");
let mut p = Pieces::new(&t);
// println!("{:?}", p);
let payload: Payload = Payload{
index: 0,
begin: 0,
length: 0,
block: Block{
length: 0,
bytes: Vec::new()
}
};
assert_eq!(true, p.needed(&payload));
p.add_requested(&payload);
p.add_received(&payload);
assert_eq!(false, p.needed(&payload));
assert_eq!(false, p.is_done());
}
#[async_std::test]
async fn get_peers() {
let t = Torrent::open("test.torrent");
// println!("{:?}", t);
println!("torrent: {:?}", t.info.name);
println!("protocol: {:?}", t.protocol().unwrap());
println!("msg_conn_req: {:?}", msg_conn_req());
let peers_w = t.get_peers().await;
println!("peers getted");
let peers = peers_w.unwrap();
// println!("get_peers peers: {:?}", peers);
}
}
|
