xvpn/src/client.rs

use anyhow::Result;
use chrono::Utc;
use clap::Args;
use ipnet::Ipv4Net;
use serde::{Deserialize, Serialize};

use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::tcp::{OwnedReadHalf, OwnedWriteHalf},
    time::Instant,
};

use crate::{
    network::ip_match_network,
    router::{CLIENT_REGISTER_TIMEOUT, CliRegMessages, RouterMessages, SERVER_PACKET_SIZE},
    tun::inti_tun_interface,
};

pub struct ClientStaTistic {
    pub last_keep_alive: Option<chrono::DateTime<Utc>>,
    pub keep_alive_count: usize,
    pub total_data_received: usize,
    pub total_data_sent: usize,
    pub last_data_received: Option<chrono::DateTime<Utc>>,
    pub last_data_sent: Option<chrono::DateTime<Utc>>,
    pub latency_ms: Option<f64>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Args)]
pub struct ClientCfg {
    /// The server endpoint to connect to in host:port or ip:port format.
    #[arg(long, short)]
    pub server: String,

    /// The local interface IP address (example: 10.8.0.2/32).
    #[arg(long = "interface-ip")]
    pub interface_ip: Ipv4Net,

    /// The local interface name.
    #[arg(long = "interface-name", default_value = "xvpn0")]
    pub interface_name: String,

    /// Local routes in CIDR format.
    /// Example: --local-route 1.1.1.1/32,10.0.0.0/24
    #[arg(long = "local-route", visible_alias = "lr", value_delimiter = ',')]
    pub local_routes: Vec<Ipv4Net>,
    /// MTU for the TUN interface.
    /// If not specified, the default MTU of the system will be used.
    #[arg(long = "mtu", default_value = "1400")]
    pub mtu: u16,
}

pub async fn start(config: ClientCfg) -> Result<()> {
    println!("Starting client with config: {:?}", config);

    let stream = tokio::net::TcpStream::connect(&config.server).await?;
    //stream.set_nodelay(true)?;
    let (mut rx, mut tx) = stream.into_split();
    // let client_stream = ClientStream::new(tx);

    let mut vpn_buf = vec![0u8; SERVER_PACKET_SIZE];
    let mut tun_buf = vec![0u8; config.mtu as usize];
    register_client(&mut rx, &mut tx, &config, &mut vpn_buf).await?;
    let tun_device = inti_tun_interface(&config).await?;

    println!("Client registration successful. Entering main loop to receive messages from router...");
    loop {
        tokio::select! {
                    msg = rx.read(&mut vpn_buf) => {
                        match msg {
                            Ok(0) => {
                                println!("Connection to router closed by peer.");
                                return Ok(());
                            }
                            Ok(n) => {
                                match RouterMessages::from_slice(&vpn_buf[..n]){
                                    RouterMessages::KeepAlive(timestamp) => {
                                        println!("Received keep-alive message from router with timestamp: {}, delta {} ms", timestamp, (Utc::now().timestamp_micros() - timestamp).abs() as f64 / 1000.0);
                                    }

                                    _ => println!("Received message from router: {:?}", RouterMessages::from_slice(&vpn_buf[..n]))
                                };
                            }
                            Err(e) => {
                                eprintln!("Error reading from router: {}", e);
                                return Err(anyhow::anyhow!(format!("Error reading from router: {}", e)));
                            }
                        }
                    }
                data = tun_device.recv(&mut tun_buf) => {
                    match data {
                        Ok(n) => {
                            let packet = etherparse::Ipv4HeaderSlice::from_slice(&tun_buf[..n])?;
                            let src = packet.source_addr();
                            match ip_match_network(src, &config.local_routes).await {
                                Some(net) => println!("Source IP {} matches local route {}", src, net),
                                None => {},
                            }
                        }
                        Err(e) => {
                            eprintln!("Error reading from TUN interface: {}", e);
                            return Err(anyhow::anyhow!(format!("Error reading from TUN interface: {}", e)));
                        }
                    }
                }
        }
    }

    // Ok(())
}

pub async fn register_client(
    rx: &mut OwnedReadHalf,
    tx: &mut OwnedWriteHalf,
    config: &ClientCfg,
    buf: &mut [u8],
) -> Result<()> {
    let register_msg = RouterMessages::CliReg(CliRegMessages::Reg(config.clone()));
    let mut client_registration_timeout =
        tokio::time::interval_at(Instant::now() + CLIENT_REGISTER_TIMEOUT, CLIENT_REGISTER_TIMEOUT);
    tx.write_all(&register_msg.to_bytes()).await?;

    loop {
        tokio::select! {
                msg = rx.read(buf) => {
                    match msg {
                        Ok(0) => {
                            let msg = "Connection closed by router while waiting for registration confirmation.";
                            eprintln!("{}", msg);
                            return Err(anyhow::anyhow!(msg));
                        }
                        Ok(n) => {
                            let response = RouterMessages::from_slice(&buf[..n]);
                            println!("Received registration response from router: {:?}", response);
                            match response {
                                RouterMessages::CliReg(CliRegMessages::RegOk(uuid)) => {
                                    println!("Client registration successful with UUID: {}", uuid);
                                    return Ok(());
                                }
                                RouterMessages::CliReg(CliRegMessages::RegFailed(err_msg)) => {
                                    eprintln!("Client registration failed: {}", err_msg);
                                    return Err(anyhow::anyhow!(format!("Client registration failed: {}", err_msg)));
                                }
                                _ => {
                                    let msg = "Unexpected message type received during client registration.";
                                    eprintln!("{}", msg);
                                    return Err(anyhow::anyhow!(msg));
                                }
                            }
                        }
                        Err(e) => {
                            eprintln!("Error reading from router during client registration: {}", e);
                            return Err(anyhow::anyhow!(format!("Error reading from router: {}", e)));
                        }
                    }


                }
                _= client_registration_timeout.tick() => {
                    let msg = "Client registration timed out waiting for confirmation from router.";
                    eprintln!("{}", msg);
                    eprintln!("Closing connection with Server");
                    tx.shutdown().await?;
                    return Err(anyhow::anyhow!(msg));
                }



        }
    }
    // Ok(())
}