torment/torment-dht-node/src/main.rs

#![recursion_limit = "512"]
use async_std::net::{IpAddr, UdpSocket};
use futures::select;
use futures::FutureExt;
use std::net::SocketAddr;
use std::pin::Pin;
use std::str::FromStr;
use std::time::{Duration, Instant};
use torment_core::peer_storage::{PeerStorage, PeerStorageHolder};
use torment_dht::host_node::HostNode;
use torment_dht::krpc::{FromBencode, Message, ToBencode};

async fn recv_from(socket: &UdpSocket) -> std::io::Result<(Vec<u8>, SocketAddr)> {
    let mut buffer = [0u8; 2048];
    socket
        .recv_from(&mut buffer)
        .await
        .map(|(len, addr)| (buffer[..len].to_vec(), addr))
}

#[async_std::main]
async fn main() {
    let socket = UdpSocket::bind(SocketAddr::from_str("[::]:50002").unwrap())
        .await
        .unwrap();
    let peer_storage = PeerStorage::new();
    let mut node = HostNode::new(PeerStorageHolder::new(peer_storage), Some(50002));
    node.add_bootstrap(SocketAddr::from_str("67.215.246.10:6881").unwrap(), None);
    node.add_bootstrap(
        SocketAddr::from_str("[2001:41d0:c:5ac:5::1]:6881").unwrap(),
        None,
    );

    let mut fut_holder = recv_from(&socket).fuse();
    let mut socket_fut = unsafe { Pin::new_unchecked(&mut fut_holder) };
    let mut fut_sleep_holder = async_std::task::sleep(Duration::from_secs(10)).fuse();
    let mut event_fut = unsafe { Pin::new_unchecked(&mut fut_sleep_holder) };

    loop {
        while let Some((message, to)) = node.next() {
            println!("| {} <= {:?}", to, message);
            socket
                .send_to(&message.to_bencode().unwrap(), to)
                .await
                .unwrap();
        }

        select! {
            res = socket_fut => {
                if let Ok((msg, from)) = res {
                    let from = if let IpAddr::V6(ipv6) = from.ip() {
                        if let Some(ipv4) = ipv6.to_ipv4() {
                           SocketAddr::new(IpAddr::V4(ipv4), from.port())
                        } else {
                           from
                        }
                    } else {
                        from
                    };

                    match Message::from_bencode(&msg) {
                        Ok(msg) => {
                            println!("| {} => {:?}", from, msg);
                            node.process(msg, from);
                        },
                        Err(err) => eprintln!("{} => Failed parsing UDP message: {}", from, err),
                    }
                }

                fut_holder = recv_from(&socket).fuse();
                socket_fut = unsafe { Pin::new_unchecked(&mut fut_holder) };
            },
            _ = event_fut => {
                println!("| Housekeeping");
                let instant = Instant::now();
                node.housekeeping();
                println!("| Housekeeping took {:?} (seen nodes: {}, table[ipv4={}, ipv6={}])", Instant::now() - instant, node.num_tracking_nodes(), node.num_ipv4_table_nodes(), node.num_ipv6_table_nodes());
                fut_sleep_holder = async_std::task::sleep(Duration::from_secs(10)).fuse();
                event_fut = unsafe { Pin::new_unchecked(&mut fut_sleep_holder) };
            },
        };
    }
}