feat(rust-core): add adaptive keepalive telemetry, MTU handling, and per-client rate limiting APIs
This commit is contained in:
@@ -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() => {
|
||||
|
||||
Reference in New Issue
Block a user