feat(rust-core): add adaptive keepalive telemetry, MTU handling, and per-client rate limiting APIs

This commit is contained in:
2026-03-15 18:10:25 +00:00
parent 97bb148063
commit 9ee41348e0
15 changed files with 2152 additions and 101 deletions

View File

@@ -3,13 +3,14 @@ use bytes::BytesMut;
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{mpsc, RwLock};
use tokio::sync::{mpsc, watch, RwLock};
use tokio_tungstenite::tungstenite::Message;
use tracing::{info, error, warn};
use tracing::{info, error, warn, debug};
use crate::codec::{Frame, FrameCodec, PacketType};
use crate::crypto;
use crate::keepalive::{self, KeepaliveSignal, LinkHealth};
use crate::telemetry::ConnectionQuality;
use crate::transport;
/// Client configuration (matches TS IVpnClientConfig).
@@ -65,6 +66,8 @@ pub struct VpnClient {
assigned_ip: Arc<RwLock<Option<String>>>,
shutdown_tx: Option<mpsc::Sender<()>>,
connected_since: Arc<RwLock<Option<std::time::Instant>>>,
quality_rx: Option<watch::Receiver<ConnectionQuality>>,
link_health: Arc<RwLock<LinkHealth>>,
}
impl VpnClient {
@@ -75,6 +78,8 @@ impl VpnClient {
assigned_ip: Arc::new(RwLock::new(None)),
shutdown_tx: None,
connected_since: Arc::new(RwLock::new(None)),
quality_rx: None,
link_health: Arc::new(RwLock::new(LinkHealth::Degraded)),
}
}
@@ -93,6 +98,7 @@ impl VpnClient {
let stats = self.stats.clone();
let assigned_ip_ref = self.assigned_ip.clone();
let connected_since = self.connected_since.clone();
let link_health = self.link_health.clone();
// Decode server public key
let server_pub_key = base64::Engine::decode(
@@ -161,6 +167,13 @@ impl VpnClient {
info!("Connected to VPN, assigned IP: {}", assigned_ip);
// Create adaptive keepalive monitor
let (monitor, handle) = keepalive::create_keepalive(None);
self.quality_rx = Some(handle.quality_rx);
// Spawn the keepalive monitor
tokio::spawn(monitor.run());
// Spawn packet forwarding loop
let assigned_ip_clone = assigned_ip.clone();
tokio::spawn(client_loop(
@@ -170,7 +183,9 @@ impl VpnClient {
state,
stats,
shutdown_rx,
config.keepalive_interval_secs.unwrap_or(30),
handle.signal_rx,
handle.ack_tx,
link_health,
));
Ok(assigned_ip_clone)
@@ -184,6 +199,7 @@ impl VpnClient {
*self.assigned_ip.write().await = None;
*self.connected_since.write().await = None;
*self.state.write().await = ClientState::Disconnected;
self.quality_rx = None;
info!("Disconnected from VPN");
Ok(())
}
@@ -208,13 +224,14 @@ impl VpnClient {
status
}
/// Get traffic statistics.
/// Get traffic statistics (includes connection quality).
pub async fn get_statistics(&self) -> serde_json::Value {
let stats = self.stats.read().await;
let since = self.connected_since.read().await;
let uptime = since.map(|s| s.elapsed().as_secs()).unwrap_or(0);
let health = self.link_health.read().await;
serde_json::json!({
let mut result = serde_json::json!({
"bytesSent": stats.bytes_sent,
"bytesReceived": stats.bytes_received,
"packetsSent": stats.packets_sent,
@@ -222,7 +239,35 @@ impl VpnClient {
"keepalivesSent": stats.keepalives_sent,
"keepalivesReceived": stats.keepalives_received,
"uptimeSeconds": uptime,
})
});
// Include connection quality if available
if let Some(ref rx) = self.quality_rx {
let quality = rx.borrow().clone();
result["quality"] = serde_json::json!({
"srttMs": quality.srtt_ms,
"jitterMs": quality.jitter_ms,
"minRttMs": quality.min_rtt_ms,
"maxRttMs": quality.max_rtt_ms,
"lossRatio": quality.loss_ratio,
"consecutiveTimeouts": quality.consecutive_timeouts,
"linkHealth": format!("{}", *health),
"keepalivesSent": quality.keepalives_sent,
"keepalivesAcked": quality.keepalives_acked,
});
}
result
}
/// Get connection quality snapshot.
pub fn get_connection_quality(&self) -> Option<ConnectionQuality> {
self.quality_rx.as_ref().map(|rx| rx.borrow().clone())
}
/// Get current link health.
pub async fn get_link_health(&self) -> LinkHealth {
*self.link_health.read().await
}
}
@@ -234,11 +279,11 @@ async fn client_loop(
state: Arc<RwLock<ClientState>>,
stats: Arc<RwLock<ClientStatistics>>,
mut shutdown_rx: mpsc::Receiver<()>,
keepalive_secs: u64,
mut signal_rx: mpsc::Receiver<KeepaliveSignal>,
ack_tx: mpsc::Sender<()>,
link_health: Arc<RwLock<LinkHealth>>,
) {
let mut buf = vec![0u8; 65535];
let mut keepalive_ticker = tokio::time::interval(Duration::from_secs(keepalive_secs));
keepalive_ticker.tick().await; // skip first immediate tick
loop {
tokio::select! {
@@ -264,6 +309,8 @@ async fn client_loop(
}
PacketType::KeepaliveAck => {
stats.write().await.keepalives_received += 1;
// Signal the keepalive monitor that ACK was received
let _ = ack_tx.send(()).await;
}
PacketType::Disconnect => {
info!("Server sent disconnect");
@@ -290,19 +337,37 @@ async fn client_loop(
}
}
}
_ = keepalive_ticker.tick() => {
let ka_frame = Frame {
packet_type: PacketType::Keepalive,
payload: vec![],
};
let mut frame_bytes = BytesMut::new();
if <FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, ka_frame, &mut frame_bytes).is_ok() {
if ws_sink.send(Message::Binary(frame_bytes.to_vec().into())).await.is_err() {
warn!("Failed to send keepalive");
signal = signal_rx.recv() => {
match signal {
Some(KeepaliveSignal::SendPing(timestamp_ms)) => {
// Embed the timestamp in the keepalive payload (8 bytes, big-endian)
let ka_frame = Frame {
packet_type: PacketType::Keepalive,
payload: timestamp_ms.to_be_bytes().to_vec(),
};
let mut frame_bytes = BytesMut::new();
if <FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, ka_frame, &mut frame_bytes).is_ok() {
if ws_sink.send(Message::Binary(frame_bytes.to_vec().into())).await.is_err() {
warn!("Failed to send keepalive");
*state.write().await = ClientState::Disconnected;
break;
}
stats.write().await.keepalives_sent += 1;
}
}
Some(KeepaliveSignal::PeerDead) => {
warn!("Peer declared dead by keepalive monitor");
*state.write().await = ClientState::Disconnected;
break;
}
stats.write().await.keepalives_sent += 1;
Some(KeepaliveSignal::LinkHealthChanged(health)) => {
debug!("Link health changed to: {}", health);
*link_health.write().await = health;
}
None => {
// Keepalive monitor channel closed
break;
}
}
}
_ = shutdown_rx.recv() => {