torment/torment-cli/src/download.rs

use bytes::{Buf, 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::{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()).unwrap();
    let mut buffer = [0u8; 4096];
    let mut done_header = false;
    #[allow(unused_assignments)]
    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 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.advance(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())
                    .unwrap();
            }
            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;
        }
    }
}