torment/torment-daemon/src/main.rs

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();
    }
}