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>, pub keep_alive_count: usize, pub total_data_received: usize, pub total_data_sent: usize, pub last_data_received: Option>, pub last_data_sent: Option>, pub latency_ms: Option, } #[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, /// 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(®ister_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(()) }