push.rocks/smartvpn
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jürgen Kunz
2026-03-15 18:10:25 +00:00
parent 97bb148063
commit 9ee41348e0
15 changed files with 2152 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

235
rust/src/server.rs
View File

@@ -5,6 +5,7 @@ use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::net::Ipv4Addr;
use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpListener;
use tokio::sync::{mpsc, Mutex, RwLock};
use tokio_tungstenite::tungstenite::Message;
@@ -12,9 +13,14 @@ use tracing::{info, error, warn};
use crate::codec::{Frame, FrameCodec, PacketType};
use crate::crypto;
use crate::mtu::{MtuConfig, TunnelOverhead};
use crate::network::IpPool;
use crate::ratelimit::TokenBucket;
use crate::transport;
/// Dead-peer timeout: 3x max keepalive interval (Healthy=60s).
const DEAD_PEER_TIMEOUT: Duration = Duration::from_secs(180);
/// Server configuration (matches TS IVpnServerConfig).
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
@@ -29,6 +35,10 @@ pub struct ServerConfig {
pub mtu: Option<u16>,
pub keepalive_interval_secs: Option<u64>,
pub enable_nat: Option<bool>,
/// Default rate limit for new clients (bytes/sec). None = unlimited.
pub default_rate_limit_bytes_per_sec: Option<u64>,
/// Default burst size for new clients (bytes). None = unlimited.
pub default_burst_bytes: Option<u64>,
}
/// Information about a connected client.
@@ -40,6 +50,12 @@ pub struct ClientInfo {
pub connected_since: String,
pub bytes_sent: u64,
pub bytes_received: u64,
pub packets_dropped: u64,
pub bytes_dropped: u64,
pub last_keepalive_at: Option<String>,
pub keepalives_received: u64,
pub rate_limit_bytes_per_sec: Option<u64>,
pub burst_bytes: Option<u64>,
}
/// Server statistics.
@@ -63,6 +79,8 @@ pub struct ServerState {
pub ip_pool: Mutex<IpPool>,
pub clients: RwLock<HashMap<String, ClientInfo>>,
pub stats: RwLock<ServerStatistics>,
pub rate_limiters: Mutex<HashMap<String, TokenBucket>>,
pub mtu_config: MtuConfig,
pub started_at: std::time::Instant,
}
@@ -98,11 +116,18 @@ impl VpnServer {
}
}
let link_mtu = config.mtu.unwrap_or(1420);
// Compute effective MTU from overhead
let overhead = TunnelOverhead::default_overhead();
let mtu_config = MtuConfig::new(overhead.effective_tun_mtu(1500).max(link_mtu));
let state = Arc::new(ServerState {
config: config.clone(),
ip_pool: Mutex::new(ip_pool),
clients: RwLock::new(HashMap::new()),
stats: RwLock::new(ServerStatistics::default()),
rate_limiters: Mutex::new(HashMap::new()),
mtu_config,
started_at: std::time::Instant::now(),
});
@@ -166,11 +191,52 @@ impl VpnServer {
if let Some(client) = clients.remove(client_id) {
let ip: Ipv4Addr = client.assigned_ip.parse()?;
state.ip_pool.lock().await.release(&ip);
state.rate_limiters.lock().await.remove(client_id);
info!("Client {} disconnected", client_id);
}
}
Ok(())
}
/// Set a rate limit for a specific client.
pub async fn set_client_rate_limit(
&self,
client_id: &str,
rate_bytes_per_sec: u64,
burst_bytes: u64,
) -> Result<()> {
if let Some(ref state) = self.state {
let mut limiters = state.rate_limiters.lock().await;
if let Some(limiter) = limiters.get_mut(client_id) {
limiter.update_limits(rate_bytes_per_sec, burst_bytes);
} else {
limiters.insert(
client_id.to_string(),
TokenBucket::new(rate_bytes_per_sec, burst_bytes),
);
}
// Update client info
let mut clients = state.clients.write().await;
if let Some(info) = clients.get_mut(client_id) {
info.rate_limit_bytes_per_sec = Some(rate_bytes_per_sec);
info.burst_bytes = Some(burst_bytes);
}
}
Ok(())
}
/// Remove rate limit for a specific client (unlimited).
pub async fn remove_client_rate_limit(&self, client_id: &str) -> Result<()> {
if let Some(ref state) = self.state {
state.rate_limiters.lock().await.remove(client_id);
let mut clients = state.clients.write().await;
if let Some(info) = clients.get_mut(client_id) {
info.rate_limit_bytes_per_sec = None;
info.burst_bytes = None;
}
}
Ok(())
}
}
async fn run_listener(
@@ -257,25 +323,43 @@ async fn handle_client_connection(
let mut noise_transport = responder.into_transport_mode()?;
// Register client
let default_rate = state.config.default_rate_limit_bytes_per_sec;
let default_burst = state.config.default_burst_bytes;
let client_info = ClientInfo {
client_id: client_id.clone(),
assigned_ip: assigned_ip.to_string(),
connected_since: timestamp_now(),
bytes_sent: 0,
bytes_received: 0,
packets_dropped: 0,
bytes_dropped: 0,
last_keepalive_at: None,
keepalives_received: 0,
rate_limit_bytes_per_sec: default_rate,
burst_bytes: default_burst,
};
state.clients.write().await.insert(client_id.clone(), client_info);
// Set up rate limiter if defaults are configured
if let (Some(rate), Some(burst)) = (default_rate, default_burst) {
state
.rate_limiters
.lock()
.await
.insert(client_id.clone(), TokenBucket::new(rate, burst));
}
{
let mut stats = state.stats.write().await;
stats.total_connections += 1;
}
// Send assigned IP info (encrypted)
// Send assigned IP info (encrypted), include effective MTU
let ip_info = serde_json::json!({
"assignedIp": assigned_ip.to_string(),
"gateway": state.ip_pool.lock().await.gateway_addr().to_string(),
"mtu": state.config.mtu.unwrap_or(1420),
"effectiveMtu": state.mtu_config.effective_mtu,
});
let ip_info_bytes = serde_json::to_vec(&ip_info)?;
let len = noise_transport.write_message(&ip_info_bytes, &mut buf)?;
@@ -289,66 +373,116 @@ async fn handle_client_connection(
info!("Client {} connected with IP {}", client_id, assigned_ip);
// Main packet loop
// Main packet loop with dead-peer detection
let mut last_activity = tokio::time::Instant::now();
loop {
match ws_stream.next().await {
Some(Ok(Message::Binary(data))) => {
let mut frame_buf = BytesMut::from(&data[..][..]);
match <FrameCodec as tokio_util::codec::Decoder>::decode(&mut FrameCodec, &mut frame_buf) {
Ok(Some(frame)) => match frame.packet_type {
PacketType::IpPacket => {
match noise_transport.read_message(&frame.payload, &mut buf) {
Ok(len) => {
let mut stats = state.stats.write().await;
stats.bytes_received += len as u64;
stats.packets_received += 1;
tokio::select! {
msg = ws_stream.next() => {
match msg {
Some(Ok(Message::Binary(data))) => {
last_activity = tokio::time::Instant::now();
let mut frame_buf = BytesMut::from(&data[..][..]);
match <FrameCodec as tokio_util::codec::Decoder>::decode(&mut FrameCodec, &mut frame_buf) {
Ok(Some(frame)) => match frame.packet_type {
PacketType::IpPacket => {
match noise_transport.read_message(&frame.payload, &mut buf) {
Ok(len) => {
// Rate limiting check
let allowed = {
let mut limiters = state.rate_limiters.lock().await;
if let Some(limiter) = limiters.get_mut(&client_id) {
limiter.try_consume(len)
} else {
true
}
};
if !allowed {
let mut clients = state.clients.write().await;
if let Some(info) = clients.get_mut(&client_id) {
info.packets_dropped += 1;
info.bytes_dropped += len as u64;
}
continue;
}
let mut stats = state.stats.write().await;
stats.bytes_received += len as u64;
stats.packets_received += 1;
// Update per-client stats
drop(stats);
let mut clients = state.clients.write().await;
if let Some(info) = clients.get_mut(&client_id) {
info.bytes_received += len as u64;
}
}
Err(e) => {
warn!("Decrypt error from {}: {}", client_id, e);
break;
}
}
}
Err(e) => {
warn!("Decrypt error from {}: {}", client_id, e);
PacketType::Keepalive => {
// Echo the keepalive payload back in the ACK
let ack_frame = Frame {
packet_type: PacketType::KeepaliveAck,
payload: frame.payload.clone(),
};
let mut frame_bytes = BytesMut::new();
<FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, ack_frame, &mut frame_bytes)?;
ws_sink.send(Message::Binary(frame_bytes.to_vec().into())).await?;
let mut stats = state.stats.write().await;
stats.keepalives_received += 1;
stats.keepalives_sent += 1;
// Update per-client keepalive tracking
drop(stats);
let mut clients = state.clients.write().await;
if let Some(info) = clients.get_mut(&client_id) {
info.last_keepalive_at = Some(timestamp_now());
info.keepalives_received += 1;
}
}
PacketType::Disconnect => {
info!("Client {} sent disconnect", client_id);
break;
}
_ => {
warn!("Unexpected packet type from {}: {:?}", client_id, frame.packet_type);
}
},
Ok(None) => {
warn!("Incomplete frame from {}", client_id);
}
Err(e) => {
warn!("Frame decode error from {}: {}", client_id, e);
break;
}
}
PacketType::Keepalive => {
let ack_frame = Frame {
packet_type: PacketType::KeepaliveAck,
payload: vec![],
};
let mut frame_bytes = BytesMut::new();
<FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, ack_frame, &mut frame_bytes)?;
ws_sink.send(Message::Binary(frame_bytes.to_vec().into())).await?;
let mut stats = state.stats.write().await;
stats.keepalives_received += 1;
stats.keepalives_sent += 1;
}
PacketType::Disconnect => {
info!("Client {} sent disconnect", client_id);
break;
}
_ => {
warn!("Unexpected packet type from {}: {:?}", client_id, frame.packet_type);
}
},
Ok(None) => {
warn!("Incomplete frame from {}", client_id);
}
Err(e) => {
warn!("Frame decode error from {}: {}", client_id, e);
Some(Ok(Message::Close(_))) | None => {
info!("Client {} connection closed", client_id);
break;
}
Some(Ok(Message::Ping(data))) => {
last_activity = tokio::time::Instant::now();
ws_sink.send(Message::Pong(data)).await?;
}
Some(Ok(_)) => {
last_activity = tokio::time::Instant::now();
continue;
}
Some(Err(e)) => {
warn!("WebSocket error from {}: {}", client_id, e);
break;
}
}
}
Some(Ok(Message::Close(_))) | None => {
info!("Client {} connection closed", client_id);
break;
}
Some(Ok(Message::Ping(data))) => {
ws_sink.send(Message::Pong(data)).await?;
}
Some(Ok(_)) => continue,
Some(Err(e)) => {
warn!("WebSocket error from {}: {}", client_id, e);
_ = tokio::time::sleep_until(last_activity + DEAD_PEER_TIMEOUT) => {
warn!("Client {} dead-peer timeout ({}s inactivity)", client_id, DEAD_PEER_TIMEOUT.as_secs());
break;
}
}
@@ -357,6 +491,7 @@ async fn handle_client_connection(
// Cleanup
state.clients.write().await.remove(&client_id);
state.ip_pool.lock().await.release(&assigned_ip);
state.rate_limiters.lock().await.remove(&client_id);
info!("Client {} disconnected, released IP {}", client_id, assigned_ip);
Ok(())