initial
This commit is contained in:
385
rust/src/server.rs
Normal file
385
rust/src/server.rs
Normal file
@@ -0,0 +1,385 @@
|
||||
use anyhow::Result;
|
||||
use bytes::BytesMut;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::net::Ipv4Addr;
|
||||
use std::sync::Arc;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::{mpsc, Mutex, RwLock};
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tracing::{info, error, warn};
|
||||
|
||||
use crate::codec::{Frame, FrameCodec, PacketType};
|
||||
use crate::crypto;
|
||||
use crate::network::IpPool;
|
||||
use crate::transport;
|
||||
|
||||
/// Server configuration (matches TS IVpnServerConfig).
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ServerConfig {
|
||||
pub listen_addr: String,
|
||||
pub tls_cert: Option<String>,
|
||||
pub tls_key: Option<String>,
|
||||
pub private_key: String,
|
||||
pub public_key: String,
|
||||
pub subnet: String,
|
||||
pub dns: Option<Vec<String>>,
|
||||
pub mtu: Option<u16>,
|
||||
pub keepalive_interval_secs: Option<u64>,
|
||||
pub enable_nat: Option<bool>,
|
||||
}
|
||||
|
||||
/// Information about a connected client.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ClientInfo {
|
||||
pub client_id: String,
|
||||
pub assigned_ip: String,
|
||||
pub connected_since: String,
|
||||
pub bytes_sent: u64,
|
||||
pub bytes_received: u64,
|
||||
}
|
||||
|
||||
/// Server statistics.
|
||||
#[derive(Debug, Clone, Serialize, Default)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ServerStatistics {
|
||||
pub bytes_sent: u64,
|
||||
pub bytes_received: u64,
|
||||
pub packets_sent: u64,
|
||||
pub packets_received: u64,
|
||||
pub keepalives_sent: u64,
|
||||
pub keepalives_received: u64,
|
||||
pub uptime_seconds: u64,
|
||||
pub active_clients: u64,
|
||||
pub total_connections: u64,
|
||||
}
|
||||
|
||||
/// Shared server state.
|
||||
pub struct ServerState {
|
||||
pub config: ServerConfig,
|
||||
pub ip_pool: Mutex<IpPool>,
|
||||
pub clients: RwLock<HashMap<String, ClientInfo>>,
|
||||
pub stats: RwLock<ServerStatistics>,
|
||||
pub started_at: std::time::Instant,
|
||||
}
|
||||
|
||||
/// The VPN server.
|
||||
pub struct VpnServer {
|
||||
state: Option<Arc<ServerState>>,
|
||||
shutdown_tx: Option<mpsc::Sender<()>>,
|
||||
}
|
||||
|
||||
impl VpnServer {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
state: None,
|
||||
shutdown_tx: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn start(&mut self, config: ServerConfig) -> Result<()> {
|
||||
if self.state.is_some() {
|
||||
anyhow::bail!("Server is already running");
|
||||
}
|
||||
|
||||
let ip_pool = IpPool::new(&config.subnet)?;
|
||||
|
||||
if config.enable_nat.unwrap_or(false) {
|
||||
if let Err(e) = crate::network::enable_ip_forwarding() {
|
||||
warn!("Failed to enable IP forwarding: {}", e);
|
||||
}
|
||||
if let Ok(iface) = crate::network::get_default_interface() {
|
||||
if let Err(e) = crate::network::setup_nat(&config.subnet, &iface).await {
|
||||
warn!("Failed to setup NAT: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let state = Arc::new(ServerState {
|
||||
config: config.clone(),
|
||||
ip_pool: Mutex::new(ip_pool),
|
||||
clients: RwLock::new(HashMap::new()),
|
||||
stats: RwLock::new(ServerStatistics::default()),
|
||||
started_at: std::time::Instant::now(),
|
||||
});
|
||||
|
||||
let (shutdown_tx, mut shutdown_rx) = mpsc::channel::<()>(1);
|
||||
self.state = Some(state.clone());
|
||||
self.shutdown_tx = Some(shutdown_tx);
|
||||
|
||||
let listen_addr = config.listen_addr.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = run_listener(state, listen_addr, &mut shutdown_rx).await {
|
||||
error!("Server listener error: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
info!("VPN server started");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn stop(&mut self) -> Result<()> {
|
||||
if let Some(tx) = self.shutdown_tx.take() {
|
||||
let _ = tx.send(()).await;
|
||||
}
|
||||
self.state = None;
|
||||
info!("VPN server stopped");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn get_status(&self) -> serde_json::Value {
|
||||
if let Some(ref state) = self.state {
|
||||
serde_json::json!({
|
||||
"state": "connected",
|
||||
"connectedSince": format!("{:?}", state.started_at.elapsed()),
|
||||
})
|
||||
} else {
|
||||
serde_json::json!({ "state": "disconnected" })
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get_statistics(&self) -> ServerStatistics {
|
||||
if let Some(ref state) = self.state {
|
||||
let mut stats = state.stats.read().await.clone();
|
||||
stats.uptime_seconds = state.started_at.elapsed().as_secs();
|
||||
stats.active_clients = state.clients.read().await.len() as u64;
|
||||
stats
|
||||
} else {
|
||||
ServerStatistics::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn list_clients(&self) -> Vec<ClientInfo> {
|
||||
if let Some(ref state) = self.state {
|
||||
state.clients.read().await.values().cloned().collect()
|
||||
} else {
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn disconnect_client(&self, client_id: &str) -> Result<()> {
|
||||
if let Some(ref state) = self.state {
|
||||
let mut clients = state.clients.write().await;
|
||||
if let Some(client) = clients.remove(client_id) {
|
||||
let ip: Ipv4Addr = client.assigned_ip.parse()?;
|
||||
state.ip_pool.lock().await.release(&ip);
|
||||
info!("Client {} disconnected", client_id);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_listener(
|
||||
state: Arc<ServerState>,
|
||||
listen_addr: String,
|
||||
shutdown_rx: &mut mpsc::Receiver<()>,
|
||||
) -> Result<()> {
|
||||
let listener = TcpListener::bind(&listen_addr).await?;
|
||||
info!("WebSocket server listening on {}", listen_addr);
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
accept = listener.accept() => {
|
||||
match accept {
|
||||
Ok((stream, addr)) => {
|
||||
info!("New connection from {}", addr);
|
||||
let state = state.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = handle_client_connection(state, stream).await {
|
||||
warn!("Client connection error: {}", e);
|
||||
}
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Accept error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = shutdown_rx.recv() => {
|
||||
info!("Shutdown signal received");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_client_connection(
|
||||
state: Arc<ServerState>,
|
||||
stream: tokio::net::TcpStream,
|
||||
) -> Result<()> {
|
||||
let ws = transport::accept_connection(stream).await?;
|
||||
let (mut ws_sink, mut ws_stream) = ws.split();
|
||||
|
||||
let client_id = uuid_v4();
|
||||
|
||||
let assigned_ip = state.ip_pool.lock().await.allocate(&client_id)?;
|
||||
|
||||
let server_private_key = base64::Engine::decode(
|
||||
&base64::engine::general_purpose::STANDARD,
|
||||
&state.config.private_key,
|
||||
)?;
|
||||
|
||||
let mut responder = crypto::create_responder(&server_private_key)?;
|
||||
let mut buf = vec![0u8; 65535];
|
||||
|
||||
// Receive handshake init
|
||||
let init_msg = match ws_stream.next().await {
|
||||
Some(Ok(Message::Binary(data))) => data.to_vec(),
|
||||
_ => anyhow::bail!("Expected handshake init message"),
|
||||
};
|
||||
|
||||
let mut frame_buf = BytesMut::from(&init_msg[..]);
|
||||
let frame = <FrameCodec as tokio_util::codec::Decoder>::decode(&mut FrameCodec, &mut frame_buf)?
|
||||
.ok_or_else(|| anyhow::anyhow!("Incomplete handshake frame"))?;
|
||||
|
||||
if frame.packet_type != PacketType::HandshakeInit {
|
||||
anyhow::bail!("Expected HandshakeInit, got {:?}", frame.packet_type);
|
||||
}
|
||||
|
||||
responder.read_message(&frame.payload, &mut buf)?;
|
||||
let len = responder.write_message(&[], &mut buf)?;
|
||||
let response_payload = buf[..len].to_vec();
|
||||
|
||||
let response_frame = Frame {
|
||||
packet_type: PacketType::HandshakeResp,
|
||||
payload: response_payload,
|
||||
};
|
||||
let mut frame_bytes = BytesMut::new();
|
||||
<FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, response_frame, &mut frame_bytes)?;
|
||||
ws_sink.send(Message::Binary(frame_bytes.to_vec().into())).await?;
|
||||
|
||||
let mut noise_transport = responder.into_transport_mode()?;
|
||||
|
||||
// Register client
|
||||
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,
|
||||
};
|
||||
state.clients.write().await.insert(client_id.clone(), client_info);
|
||||
|
||||
{
|
||||
let mut stats = state.stats.write().await;
|
||||
stats.total_connections += 1;
|
||||
}
|
||||
|
||||
// Send assigned IP info (encrypted)
|
||||
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),
|
||||
});
|
||||
let ip_info_bytes = serde_json::to_vec(&ip_info)?;
|
||||
let len = noise_transport.write_message(&ip_info_bytes, &mut buf)?;
|
||||
let encrypted_info = Frame {
|
||||
packet_type: PacketType::IpPacket,
|
||||
payload: buf[..len].to_vec(),
|
||||
};
|
||||
let mut frame_bytes = BytesMut::new();
|
||||
<FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, encrypted_info, &mut frame_bytes)?;
|
||||
ws_sink.send(Message::Binary(frame_bytes.to_vec().into())).await?;
|
||||
|
||||
info!("Client {} connected with IP {}", client_id, assigned_ip);
|
||||
|
||||
// Main packet loop
|
||||
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;
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("Decrypt 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);
|
||||
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);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Cleanup
|
||||
state.clients.write().await.remove(&client_id);
|
||||
state.ip_pool.lock().await.release(&assigned_ip);
|
||||
info!("Client {} disconnected, released IP {}", client_id, assigned_ip);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn uuid_v4() -> String {
|
||||
use rand::Rng;
|
||||
let mut rng = rand::thread_rng();
|
||||
let bytes: [u8; 16] = rng.gen();
|
||||
format!(
|
||||
"{:02x}{:02x}{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}",
|
||||
bytes[0], bytes[1], bytes[2], bytes[3],
|
||||
bytes[4], bytes[5],
|
||||
bytes[6], bytes[7],
|
||||
bytes[8], bytes[9],
|
||||
bytes[10], bytes[11], bytes[12], bytes[13], bytes[14], bytes[15],
|
||||
)
|
||||
}
|
||||
|
||||
fn timestamp_now() -> String {
|
||||
use std::time::SystemTime;
|
||||
let duration = SystemTime::now()
|
||||
.duration_since(SystemTime::UNIX_EPOCH)
|
||||
.unwrap_or_default();
|
||||
format!("{}", duration.as_secs())
|
||||
}
|
||||
Reference in New Issue
Block a user