torment/torment-dht/src/host_node.rs

use crate::krpc::{Message, MessageBody, RequestBody, ResponseBody, Want};
use crate::Table;
use rand::random;
use serde_derive::{Deserialize, Serialize};
use std::collections::{HashSet, VecDeque};
use std::convert::TryInto;
use std::net::{IpAddr, SocketAddr};
use std::ops::Add;
use std::time::{Duration, Instant};
use torment_core::infohash::v1::U160;
use torment_core::ip::IpAddrExt;
use torment_core::peer_storage::{PeerStorageHolder, PeerStorageSource};
use torment_core::utils::{EphemeralMap, EphemeralSet};
use torment_core::{ContactInfo, LookupFilter};

#[derive(Debug)]
enum OpenRequest {
    FindNode { target: U160 },
    GetPeers { info_hash: U160 },
    Ping,
    Announce { info_hash: U160 },
}

const BOOTSTRAP_INTERVAL: Duration = Duration::from_secs(60);

#[derive(Debug)]
pub struct HostNode {
    id: U160,
    port: Option<u16>,
    bootstrap_nodes: HashSet<(SocketAddr, Option<Want>)>,
    last_bootstrap: Option<Instant>,
    ipv4_table: Table,
    ipv6_table: Table,
    nodes: EphemeralMap<(U160, IpAddr, u16), PeerNode>,
    queue: VecDeque<(Message, SocketAddr)>,
    peer_storage: PeerStorageHolder,
}

#[derive(Debug)]
struct PeerNode {
    id: U160,
    addr: SocketAddr,
    transactions: EphemeralMap<u16, OpenRequest>,
    pre_announced: EphemeralSet<U160>,
    last_activity: Instant,
    tokens: EphemeralSet<u64>,
    received_tokens: EphemeralSet<Vec<u8>>,
    transaction_id: u16,
}

#[derive(Serialize, Deserialize)]
pub struct PersistentState {
    id: U160,
    nodes: Vec<ContactInfo>,
    nodes6: Vec<ContactInfo>,
}

fn validate_ip<T: Into<IpAddr>>(ip: T) -> bool {
    IpAddrExt::ext_is_global(&ip.into())
}

impl PeerNode {
    fn clean(&mut self) {
        self.transactions.clean();
        self.tokens.clean();
        self.received_tokens.clean();
        self.pre_announced.clean();
    }

    fn contact_info(&self) -> ContactInfo {
        ContactInfo {
            id: self.id,
            contact: self.addr,
        }
    }
}

impl PeerNode {
    fn create<T: Into<SocketAddr>>(id: U160, addr: T) -> PeerNode {
        PeerNode {
            id,
            addr: addr.into(),
            transactions: EphemeralMap::new(),
            pre_announced: Default::default(),
            last_activity: Instant::now(),
            tokens: Default::default(),
            received_tokens: Default::default(),
            transaction_id: 0,
        }
    }
}

impl HostNode {
    pub fn new(peer_storage: PeerStorageHolder, port: Option<u16>) -> HostNode {
        let id = U160::random();
        HostNode {
            id,
            port,
            bootstrap_nodes: Default::default(),
            last_bootstrap: None,
            ipv4_table: Table::new_with_id(id),
            ipv6_table: Table::new_with_id(id),
            nodes: EphemeralMap::new(),
            queue: Default::default(),
            peer_storage,
        }
    }

    pub fn from_persistent_state(
        peer_storage: PeerStorageHolder,
        port: Option<u16>,
        state: PersistentState,
    ) -> HostNode {
        let mut host_node = HostNode {
            id: state.id,
            port,
            bootstrap_nodes: Default::default(),
            last_bootstrap: None,
            ipv4_table: Table::new_with_id(state.id),
            ipv6_table: Table::new_with_id(state.id),
            nodes: EphemeralMap::new(),
            queue: Default::default(),
            peer_storage,
        };

        for node in state.nodes {
            host_node
                .ipv4_table
                .add_node(node.id, node.contact.ip(), node.contact.port());
        }

        for node in state.nodes6 {
            host_node
                .ipv6_table
                .add_node(node.id, node.contact.ip(), node.contact.port());
        }

        host_node.update_tables();
        host_node
    }

    fn get_next_t(&mut self, contact_info: &ContactInfo) -> [u8; 2] {
        let peer_node = self.get_peer_node(contact_info.id, contact_info.contact);
        let old_id = peer_node.transaction_id;
        peer_node.transaction_id = peer_node.transaction_id.overflowing_add(1).0;
        old_id.to_ne_bytes()
    }

    pub fn process(&mut self, message: Message, from: SocketAddr) {
        match message.body {
            MessageBody::Request(ref req) => self.handle_request(req, message.transaction_id, from),
            MessageBody::Response(ref resp) => {
                self.handle_response(resp, message.transaction_id, from)
            }
            _ => {}
        }
    }

    pub fn add_bootstrap(&mut self, socket_addr: SocketAddr, want: Option<Want>) {
        self.bootstrap_nodes.insert((socket_addr, want));
    }

    fn get_k_closest_nodes(&self, id: U160, want: Option<Want>) -> Vec<ContactInfo> {
        let mut nodes = vec![];
        if want.unwrap_or(Want::N4) == Want::N4 {
            nodes.extend(self.ipv4_table.find_k_closest_nodes(id, None))
        }

        if want.unwrap_or(Want::N6) == Want::N6 {
            nodes.extend(self.ipv6_table.find_k_closest_nodes(id, None))
        }

        nodes
    }

    fn find_node(&mut self, target: U160, want: Option<Want>) {
        if want.unwrap_or(Want::N4) == Want::N4 && self.ipv4_table.find_node(target).is_some() {
            // noop
            return;
        }

        if want.unwrap_or(Want::N6) == Want::N6 && self.ipv6_table.find_node(target).is_some() {
            // noop
            return;
        }

        for node in self.get_k_closest_nodes(target, want) {
            let t = self.get_next_t(&node);
            self.find_node_at(t, node.contact, target, want);
            self.get_peer_node(node.id, node.contact)
                .transactions
                .insert(
                    u16::from_ne_bytes(t),
                    OpenRequest::FindNode { target },
                    Instant::now().add(Duration::from_secs(5 * 60)),
                );

            self.notify_table_of_outgoing_activity(&node);
        }
    }

    fn find_node_at(
        &mut self,
        t: [u8; 2],
        socket_addr: SocketAddr,
        target: U160,
        want: Option<Want>,
    ) {
        let message = Message::request_find_node(self.id, t, target, want);
        self.queue.push_back((message, socket_addr))
    }

    fn get_peers(&mut self, info_hash: U160, want: Option<Want>) {
        let nodes = self.get_k_closest_nodes(info_hash, want);
        for node in nodes {
            self.get_peers_from(&node, info_hash, want)
        }
    }

    fn get_peers_from(&mut self, node: &ContactInfo, info_hash: U160, want: Option<Want>) {
        let message = Message::request_get_peers(self.id, self.get_next_t(&node), info_hash, want);
        self.get_peer_node(node.id, node.contact)
            .transactions
            .insert(
                u16::from_ne_bytes(message.transaction_id),
                OpenRequest::GetPeers { info_hash },
                Instant::now().add(Duration::from_secs(5 * 60)),
            );

        self.notify_table_of_outgoing_activity(node);
        self.queue.push_back((message, node.contact))
    }

    fn notify_table_of_outgoing_activity(&mut self, node: &ContactInfo) {
        if node.contact.is_ipv4() {
            self.ipv4_table.handle_outgoing_activity(
                node.id,
                node.contact.ip(),
                node.contact.port(),
            );
        } else {
            self.ipv6_table.handle_outgoing_activity(
                node.id,
                node.contact.ip(),
                node.contact.port(),
            );
        }
    }

    fn announce_peer_to(&mut self, node: &ContactInfo, info_hash: U160, token: Vec<u8>) {
        let message = Message::request_announce_peer(
            self.id,
            self.get_next_t(&node),
            info_hash,
            token,
            self.port.unwrap_or(0),
            self.port.is_none(),
        );
        self.get_peer_node(node.id, node.contact)
            .transactions
            .insert(
                u16::from_ne_bytes(message.transaction_id),
                OpenRequest::Announce { info_hash },
                Instant::now().add(Duration::from_secs(5 * 60)),
            );

        self.notify_table_of_outgoing_activity(node);
        self.queue.push_back((message, node.contact))
    }

    pub fn ping(&mut self, node: &ContactInfo) {
        let msg = Message::request_ping(self.id, self.get_next_t(node));
        self.get_peer_node(node.id, node.contact)
            .transactions
            .insert(
                u16::from_ne_bytes(msg.transaction_id),
                OpenRequest::Ping,
                Instant::now().add(Duration::from_secs(5 * 60)),
            );

        self.notify_table_of_outgoing_activity(node);
        self.queue.push_back((msg, node.contact))
    }

    fn announce_peer(&mut self, info_hash: U160) {
        let mut announce_to = vec![];
        let mut peers_from = vec![];
        for (_, details) in &mut self.nodes.iter_mut() {
            if let Some(token) = details.received_tokens.remove_first() {
                announce_to.push((details.contact_info(), token))
            } else {
                details
                    .pre_announced
                    .insert(info_hash, Instant::now().add(Duration::from_secs(5 * 60)));
                peers_from.push(details.contact_info());
            }
        }

        for (node, token) in announce_to {
            self.announce_peer_to(&node, info_hash, token)
        }

        for node in peers_from {
            self.get_peers_from(&node, info_hash, None)
        }
    }

    fn has_node(&self, node: &ContactInfo) -> bool {
        self.ipv4_table.has_node(node) || self.ipv6_table.has_node(node)
    }

    fn knows_node(&self, node: &ContactInfo) -> bool {
        self.nodes
            .contains_key(&(node.id, node.contact.ip(), node.contact.port()))
    }

    fn handle_response(&mut self, response: &ResponseBody, t: [u8; 2], from: SocketAddr) {
        let from_id = response.node_id();
        if from_id == self.id {
            // In this house we don't talk to imposters
            return;
        }

        match response {
            ResponseBody::GetPeers(get_peers) => {
                let node = self.get_peer_node(from_id, from);
                let info_hash = if let Some(OpenRequest::GetPeers { info_hash }) =
                    node.transactions.remove_value(&u16::from_ne_bytes(t))
                {
                    info_hash
                } else {
                    return;
                };

                let is_pre_announce = if node.pre_announced.remove(&info_hash) {
                    true
                } else {
                    node.received_tokens.insert(
                        get_peers.token.clone(),
                        Instant::now().add(Duration::from_secs(5 * 60)),
                    );

                    false
                };

                if is_pre_announce {
                    let contact_info = node.contact_info();
                    self.announce_peer_to(&contact_info, info_hash, get_peers.token.clone());
                    return;
                }

                if get_peers.values.len() > 0 {
                    self.peer_storage.add_peers(
                        info_hash,
                        get_peers.values.clone(),
                        PeerStorageSource::DHT,
                    );
                }

                for node in &get_peers.nodes.nodes {
                    if !validate_ip(node.contact.ip()) {
                        continue;
                    }

                    if self.has_node(node) {
                        continue;
                    }

                    self.ipv4_table
                        .add_node(node.id, node.contact.ip(), node.contact.port());
                    self.get_peers_from(node, info_hash, None);
                }

                for node in &get_peers.nodes.nodes6 {
                    if !validate_ip(node.contact.ip()) {
                        continue;
                    }

                    if self.has_node(node) {
                        continue;
                    }

                    self.ipv6_table
                        .add_node(node.id, node.contact.ip(), node.contact.port());
                    self.get_peers_from(node, info_hash, None);
                }
            }

            ResponseBody::FindNode(find_node) => {
                let node = self.get_peer_node(response.node_id(), from);

                // We don't really care if we had a transaction or not, since at bootstrap we didn't know the nodes ID yet
                let request = node.transactions.remove_value(&u16::from_ne_bytes(t));

                let target = if let Some(OpenRequest::FindNode { target }) = request {
                    target
                } else {
                    self.id
                };

                let poke_if_closer = |host: &mut HostNode, node: &ContactInfo| {
                    if !host.knows_node(node) && (node.id ^ host.id) < (from_id ^ host.id) {
                        let t = host.get_next_t(node);
                        host.find_node_at(t, node.contact, target, None);
                    }
                };

                for node in &find_node.nodes {
                    if !validate_ip(node.contact.ip()) {
                        continue;
                    }

                    if self
                        .ipv4_table
                        .add_node(node.id, node.contact.ip(), node.contact.port())
                    {
                        poke_if_closer(self, node);
                    }
                }

                for node in &find_node.nodes6 {
                    if !validate_ip(node.contact.ip()) {
                        continue;
                    }

                    if self
                        .ipv6_table
                        .add_node(node.id, node.contact.ip(), node.contact.port())
                    {
                        poke_if_closer(self, node);
                    }
                }
            }

            ResponseBody::Empty(id) => {
                let node = self.get_peer_node(response.node_id(), from);
                node.transactions.remove_value(&u16::from_ne_bytes(t));

                if node.addr.is_ipv4() {
                    self.ipv4_table.handle_activity(*id, from.ip(), from.port());
                } else {
                    self.ipv6_table.handle_activity(*id, from.ip(), from.port());
                }
            }
        }
    }

    fn get_peer_node<T: Into<SocketAddr>>(&mut self, id: U160, addr: T) -> &mut PeerNode {
        let addr = addr.into();
        self.nodes.update_expiry(
            &(id, addr.ip(), addr.port()),
            Instant::now().add(Duration::from_secs(60 * 15)),
        );
        self.nodes
            .entry((id, addr.ip(), addr.port()))
            .or_insert_with(Instant::now().add(Duration::from_secs(60 * 15)), || {
                PeerNode::create(id, addr)
            })
    }

    fn create_token(&mut self, id: U160, addr: SocketAddr, _info_hash: U160) -> Vec<u8> {
        let token: u64 = random();

        self.get_peer_node(id, addr).tokens.insert(
            token.clone(),
            Instant::now().add(Duration::from_secs(60 * 5)),
        );

        token.to_ne_bytes().to_vec()
    }

    fn handle_request(&mut self, request: &RequestBody, t: [u8; 2], from: SocketAddr) {
        if request.node_id() == self.id {
            // In this house we don't talk to imposters
            return;
        }

        let contact_info = ContactInfo {
            id: request.node_id(),
            contact: from,
        };

        if !self.has_node(&contact_info) {
            if from.is_ipv4() {
                self.ipv4_table.add_incoming_node(
                    contact_info.id,
                    contact_info.contact.ip(),
                    contact_info.contact.port(),
                );
            } else {
                self.ipv6_table.add_incoming_node(
                    contact_info.id,
                    contact_info.contact.ip(),
                    contact_info.contact.port(),
                );
            }
        }

        match request {
            RequestBody::GetPeers(get_peers) => {
                let table =
                    get_peers
                        .want
                        .unwrap_or_else(|| if from.is_ipv4() { Want::N4 } else { Want::N6 });
                let peers = self.peer_storage.get_peers(
                    get_peers.info_hash,
                    if table == Want::N4 {
                        LookupFilter::IPv4
                    } else {
                        LookupFilter::IPv6
                    },
                );

                let token = self.create_token(get_peers.id, from, get_peers.info_hash);
                if peers.len() == 0 {
                    self.queue.push_back((
                        Message::response_get_peers_not_found(
                            self.id,
                            t,
                            token,
                            if table == Want::N4 {
                                self.ipv4_table.find_k_closest_nodes(
                                    get_peers.info_hash,
                                    Some((get_peers.id, from.ip(), from.port())),
                                )
                            } else {
                                vec![]
                            },
                            if table == Want::N6 {
                                self.ipv6_table.find_k_closest_nodes(
                                    get_peers.info_hash,
                                    Some((get_peers.id, from.ip(), from.port())),
                                )
                            } else {
                                vec![]
                            },
                        ),
                        from,
                    ))
                } else {
                    self.queue.push_back((
                        Message::response_get_peers_found(self.id, t, token, peers),
                        from,
                    ));
                }
            }

            RequestBody::FindNode(find_node) => {
                let table =
                    find_node
                        .want
                        .unwrap_or_else(|| if from.is_ipv4() { Want::N4 } else { Want::N6 });

                let nodes = if table == Want::N4 {
                    if let Some(node) = self.ipv4_table.find_node(find_node.target) {
                        vec![node]
                    } else {
                        self.ipv4_table.find_k_closest_nodes(
                            find_node.target,
                            Some((find_node.id, from.ip(), from.port())),
                        )
                    }
                } else {
                    vec![]
                };

                let nodes6 = if table == Want::N6 {
                    if let Some(node) = self.ipv6_table.find_node(find_node.target) {
                        vec![node]
                    } else {
                        self.ipv6_table.find_k_closest_nodes(
                            find_node.target,
                            Some((find_node.id, from.ip(), from.port())),
                        )
                    }
                } else {
                    vec![]
                };

                self.queue
                    .push_back((Message::response_find_node(self.id, t, nodes, nodes6), from))
            }

            RequestBody::Ping(id) => {
                if from.is_ipv4() {
                    self.ipv4_table.handle_activity(*id, from.ip(), from.port())
                } else {
                    self.ipv6_table.handle_activity(*id, from.ip(), from.port())
                }

                self.queue
                    .push_back((Message::response_empty(self.id, t), from))
            }

            RequestBody::AnnouncePeer(announce) => {
                let ok = if let Some(node) =
                    self.nodes.get_mut(&(announce.id, from.ip(), from.port()))
                {
                    if announce.token.len() != 8 {
                        false
                    } else {
                        node.tokens.remove(&u64::from_ne_bytes(
                            announce.token[0..8].try_into().unwrap(),
                        ))
                    }
                } else {
                    false
                };

                if !ok {
                    self.queue
                        .push_back((Message::error(t, 203, "Bad token".to_string()), from))
                } else {
                    self.peer_storage.add_peers(
                        announce.info_hash,
                        vec![SocketAddr::new(
                            from.ip(),
                            if announce.implied_port {
                                from.port()
                            } else {
                                announce.port
                            },
                        )],
                        PeerStorageSource::DHT,
                    );
                    self.queue
                        .push_back((Message::response_empty(self.id, t), from))
                }
            }
        }
    }

    pub fn next(&mut self) -> Option<(Message, SocketAddr)> {
        self.queue.pop_front()
    }

    fn clean(&mut self) {
        for (_, node) in self.nodes.iter_mut() {
            node.clean();
        }
    }

    pub fn housekeeping(&mut self) {
        self.clean();
        self.update_tables();
        if self
            .last_bootstrap
            .map(|last| last.add(BOOTSTRAP_INTERVAL) < Instant::now())
            .unwrap_or(true)
        {
            for (addr, want) in self.bootstrap_nodes.clone() {
                self.find_node_at(random(), addr, self.id, want);
            }

            self.last_bootstrap = Some(Instant::now());
        }
    }

    fn update_tables(&mut self) {
        self.ipv4_table.update_nodes();
        self.ipv6_table.update_nodes();

        for item in self.ipv6_table.get_silent_nodes() {
            self.ping(&item);
        }

        for item in self.ipv4_table.get_silent_nodes() {
            self.ping(&item);
        }
    }

    pub fn num_tracking_nodes(&self) -> usize {
        self.nodes.len()
    }

    pub fn num_ipv4_table_nodes(&self) -> usize {
        self.ipv4_table.count_nodes()
    }

    pub fn num_ipv6_table_nodes(&self) -> usize {
        self.ipv6_table.count_nodes()
    }
}

#[cfg(test)]
mod test {
    use crate::host_node::HostNode;
    use crate::krpc::Message;
    use std::net::{IpAddr, Ipv4Addr, SocketAddr};
    use torment_core::infohash::v1::U160;
    use torment_core::peer_storage::{PeerStorage, PeerStorageHolder};
    use torment_core::ContactInfo;

    fn get_host() -> HostNode {
        let peer_storage = PeerStorageHolder::new(PeerStorage::new());
        HostNode::new(peer_storage, None)
    }

    #[test]
    fn ping_pong() {
        let mut host = get_host();
        let t = [0u8; 2];
        let ip = Ipv4Addr::LOCALHOST;
        let id = U160::random();
        let from = SocketAddr::new(IpAddr::from(ip), 4);

        host.process(Message::request_ping(id, t), from);
        assert_eq!(1, host.ipv4_table.count_nodes());
        assert_eq!(
            host.next(),
            Some((Message::response_empty(host.id, t), from))
        )
    }

    #[test]
    fn get_peers() {
        let mut host = get_host();
        let t = [0u8; 2];
        let ip = Ipv4Addr::LOCALHOST;
        let from = SocketAddr::new(IpAddr::from(ip), 4);
        let id = U160::random();

        host.process(Message::request_ping(id, t), from);
        assert_eq!(1, host.ipv4_table.count_nodes());
        host.next();

        let other = U160::random();
        host.process(Message::request_ping(other, t), from);
        assert_eq!(2, host.ipv4_table.count_nodes());
        host.next();

        let info_hash = U160::random();
        host.process(Message::request_get_peers(id, t, info_hash, None), from);

        let nodes = host
            .next()
            .unwrap()
            .0
            .expect_response()
            .1
            .expect_get_peers()
            .nodes
            .nodes;

        assert_eq!(1, nodes.len());
        assert_eq!(
            &ContactInfo {
                id: other,
                contact: SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 4),
            },
            nodes.get(0).unwrap()
        );
    }

    #[test]
    fn announce_peer() {
        let mut host = get_host();
        let t = [0u8; 2];
        let ip = Ipv4Addr::LOCALHOST;
        let from = SocketAddr::new(IpAddr::from(ip), 4);
        let id = U160::random();

        host.process(Message::request_ping(id, t), from);
        assert_eq!(1, host.ipv4_table.count_nodes());
        host.next();

        let other = U160::random();
        host.process(Message::request_ping(other, t), from);
        assert_eq!(2, host.ipv4_table.count_nodes());
        host.next();

        let info_hash = U160::random();
        host.process(Message::request_get_peers(other, t, info_hash, None), from);

        let resp = host
            .next()
            .unwrap()
            .0
            .expect_response()
            .1
            .expect_get_peers();

        host.process(
            Message::request_announce_peer(other, t, info_hash, resp.token.clone(), 4, false),
            from,
        );

        host.next().unwrap().0.expect_response().1.expect_empty();

        host.process(
            Message::request_announce_peer(other, t, info_hash, resp.token, 4, false),
            from,
        );

        host.next().unwrap().0.expect_error();
    }
}