rustdesk-server/src/rendezvous_server.rs

use super::message_proto::*;
use bytes::Bytes;
use futures::{FutureExt, SinkExt};
use protobuf::{parse_from_bytes, Message as _};
use std::{
    collections::HashMap,
    error::Error,
    net::{Ipv4Addr, SocketAddr, SocketAddrV4},
    time::{SystemTime, UNIX_EPOCH},
};
use tokio::net::UdpSocket;
use tokio::stream::StreamExt;
use tokio_util::{codec::BytesCodec, udp::UdpFramed};

/// Certain router and firewalls scan the packet and if they
/// find an IP address belonging to their pool that they use to do the NAT mapping/translation, so here we mangle the ip address

pub struct V4AddrMangle(Vec<u8>);

impl V4AddrMangle {
    pub fn encode(addr: &SocketAddrV4) -> Self {
        let tm = (SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_micros() as u32) as u128;
        let ip = u32::from_ne_bytes(addr.ip().octets()) as u128;
        let port = addr.port() as u128;
        let v = ((ip + tm) << 49) | (tm << 17) | (port + (tm & 0xFFFF));
        let bytes = v.to_ne_bytes();
        let mut n_padding = 0;
        for i in bytes.iter().rev() {
            if i == &0u8 {
                n_padding += 1;
            } else {
                break;
            }
        }
        Self(bytes[..(16 - n_padding)].to_vec())
    }

    pub fn decode(&self) -> SocketAddrV4 {
        let mut padded = [0u8; 16];
        padded[..self.0.len()].copy_from_slice(&self.0);
        let number = u128::from_ne_bytes(padded);
        let tm = (number >> 17) & (u32::max_value() as u128);
        let ip = (((number >> 49) - tm) as u32).to_ne_bytes();
        let port = (number & 0xFFFFFF) - (tm & 0xFFFF);
        SocketAddrV4::new(Ipv4Addr::new(ip[0], ip[1], ip[2], ip[3]), port as u16)
    }
}

pub struct Peer {
    socket_addr: SocketAddrV4,
}

type PeerMap = HashMap<String, Peer>;

pub struct RendezvousServer {
    peer_map: PeerMap,
}

type FramedSocket = UdpFramed<BytesCodec>;
type ResultType = Result<(), Box<dyn Error>>;

impl RendezvousServer {
    pub async fn start(addr: &str) -> ResultType {
        let socket = UdpSocket::bind(addr).await?;
        let mut socket = UdpFramed::new(socket, BytesCodec::new());

        let mut rs = Self {
            peer_map: PeerMap::new(),
        };
        while let Some(Ok((bytes, addr))) = socket.next().await {
            if let SocketAddr::V4(addr_v4) = addr {
                if let Ok(msg_in) = parse_from_bytes::<Message>(&bytes) {
                    match msg_in.union {
                        Some(Message_oneof_union::register_peer(rp)) => {
                            if rp.hbb_addr.len() > 0 {
                                rs.peer_map.insert(
                                    rp.hbb_addr,
                                    Peer {
                                        socket_addr: addr_v4,
                                    },
                                );
                            }
                            tokio_timer::sleep(std::time::Duration::from_secs(60));
                        }
                        Some(Message_oneof_union::peek_peer(pp)) => {
                            rs.handle_peek_peer(&pp, addr, &mut socket).await?;
                            tokio_timer::sleep(std::time::Duration::from_secs(60));
                        }
                        _ => {}
                    }
                }
            }
        }
        Ok(())
    }

    pub async fn handle_peek_peer(
        &self,
        pp: &PeekPeer,
        addr: SocketAddr,
        socket: &mut FramedSocket,
    ) -> ResultType {
        if let Some(peer) = self.peer_map.get(&pp.hbb_addr) {
            let mut msg_out = Message::new();
            msg_out.set_peek_peer_response(PeekPeerResponse {
                socket_addr: V4AddrMangle::encode(&peer.socket_addr).0.to_vec(),
                ..Default::default()
            });
            send_to(&msg_out, addr, socket).await?;
        }
        Ok(())
    }
}

#[inline]
pub async fn send_to(msg: &Message, addr: SocketAddr, socket: &mut FramedSocket) -> ResultType {
    socket
        .send((Bytes::from(msg.write_to_bytes().unwrap()), addr))
        .await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_mangle() {
        let addr = SocketAddrV4::new(Ipv4Addr::new(192, 168, 16, 32), 21116);
        assert_eq!(addr, V4AddrMangle::encode(&addr).decode());
    }

    #[allow(unused_must_use)]
    #[tokio::main]
    async fn test_rs_async() {
        let server_addr = "0.0.0.0:21116";
        let f1 = RendezvousServer::start(server_addr);
        let to_addr = server_addr.parse().unwrap();
        let f2 = async {
            let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
            let mut socket = UdpFramed::new(socket, BytesCodec::new());
            let mut msg_out = Message::new();
            msg_out.set_peek_peer(PeekPeer {
                hbb_addr: "123".to_string(),
                ..Default::default()
            });
            send_to(&msg_out, to_addr, &mut socket).await;
        };
        tokio::join!(f1, f2);
    }

    #[test]
    fn test_rs() {
        self::test_rs_async();
    }
}