initial
This commit is contained in:
324
rust/src/client.rs
Normal file
324
rust/src/client.rs
Normal file
@@ -0,0 +1,324 @@
|
||||
use anyhow::Result;
|
||||
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_tungstenite::tungstenite::Message;
|
||||
use tracing::{info, error, warn};
|
||||
|
||||
use crate::codec::{Frame, FrameCodec, PacketType};
|
||||
use crate::crypto;
|
||||
use crate::transport;
|
||||
|
||||
/// Client configuration (matches TS IVpnClientConfig).
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ClientConfig {
|
||||
pub server_url: String,
|
||||
pub server_public_key: String,
|
||||
pub dns: Option<Vec<String>>,
|
||||
pub mtu: Option<u16>,
|
||||
pub keepalive_interval_secs: Option<u64>,
|
||||
}
|
||||
|
||||
/// Client statistics.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ClientStatistics {
|
||||
pub bytes_sent: u64,
|
||||
pub bytes_received: u64,
|
||||
pub packets_sent: u64,
|
||||
pub packets_received: u64,
|
||||
pub keepalives_sent: u64,
|
||||
pub keepalives_received: u64,
|
||||
}
|
||||
|
||||
/// Client connection state.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum ClientState {
|
||||
Disconnected,
|
||||
Connecting,
|
||||
Handshaking,
|
||||
Connected,
|
||||
Reconnecting,
|
||||
Error(String),
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ClientState {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Disconnected => write!(f, "disconnected"),
|
||||
Self::Connecting => write!(f, "connecting"),
|
||||
Self::Handshaking => write!(f, "handshaking"),
|
||||
Self::Connected => write!(f, "connected"),
|
||||
Self::Reconnecting => write!(f, "reconnecting"),
|
||||
Self::Error(e) => write!(f, "error: {}", e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The VPN client.
|
||||
pub struct VpnClient {
|
||||
state: Arc<RwLock<ClientState>>,
|
||||
stats: Arc<RwLock<ClientStatistics>>,
|
||||
assigned_ip: Arc<RwLock<Option<String>>>,
|
||||
shutdown_tx: Option<mpsc::Sender<()>>,
|
||||
connected_since: Arc<RwLock<Option<std::time::Instant>>>,
|
||||
}
|
||||
|
||||
impl VpnClient {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
state: Arc::new(RwLock::new(ClientState::Disconnected)),
|
||||
stats: Arc::new(RwLock::new(ClientStatistics::default())),
|
||||
assigned_ip: Arc::new(RwLock::new(None)),
|
||||
shutdown_tx: None,
|
||||
connected_since: Arc::new(RwLock::new(None)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Connect to the VPN server.
|
||||
pub async fn connect(&mut self, config: ClientConfig) -> Result<String> {
|
||||
if *self.state.read().await != ClientState::Disconnected {
|
||||
anyhow::bail!("Client is not disconnected");
|
||||
}
|
||||
|
||||
*self.state.write().await = ClientState::Connecting;
|
||||
|
||||
let (shutdown_tx, shutdown_rx) = mpsc::channel::<()>(1);
|
||||
self.shutdown_tx = Some(shutdown_tx);
|
||||
|
||||
let state = self.state.clone();
|
||||
let stats = self.stats.clone();
|
||||
let assigned_ip_ref = self.assigned_ip.clone();
|
||||
let connected_since = self.connected_since.clone();
|
||||
|
||||
// Decode server public key
|
||||
let server_pub_key = base64::Engine::decode(
|
||||
&base64::engine::general_purpose::STANDARD,
|
||||
&config.server_public_key,
|
||||
)?;
|
||||
|
||||
// Connect to WebSocket server
|
||||
let ws = transport::connect_to_server(&config.server_url).await?;
|
||||
let (mut ws_sink, mut ws_stream) = ws.split();
|
||||
|
||||
// Noise NK handshake (client side = initiator)
|
||||
*state.write().await = ClientState::Handshaking;
|
||||
let mut initiator = crypto::create_initiator(&server_pub_key)?;
|
||||
let mut buf = vec![0u8; 65535];
|
||||
|
||||
// -> e, es
|
||||
let len = initiator.write_message(&[], &mut buf)?;
|
||||
let init_frame = Frame {
|
||||
packet_type: PacketType::HandshakeInit,
|
||||
payload: buf[..len].to_vec(),
|
||||
};
|
||||
let mut frame_bytes = BytesMut::new();
|
||||
<FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, init_frame, &mut frame_bytes)?;
|
||||
ws_sink.send(Message::Binary(frame_bytes.to_vec().into())).await?;
|
||||
|
||||
// <- e, ee
|
||||
let resp_msg = match ws_stream.next().await {
|
||||
Some(Ok(Message::Binary(data))) => data.to_vec(),
|
||||
Some(Ok(_)) => anyhow::bail!("Expected binary handshake response"),
|
||||
Some(Err(e)) => anyhow::bail!("WebSocket error during handshake: {}", e),
|
||||
None => anyhow::bail!("Connection closed during handshake"),
|
||||
};
|
||||
|
||||
let mut frame_buf = BytesMut::from(&resp_msg[..]);
|
||||
let frame = <FrameCodec as tokio_util::codec::Decoder>::decode(&mut FrameCodec, &mut frame_buf)?
|
||||
.ok_or_else(|| anyhow::anyhow!("Incomplete handshake response frame"))?;
|
||||
|
||||
if frame.packet_type != PacketType::HandshakeResp {
|
||||
anyhow::bail!("Expected HandshakeResp, got {:?}", frame.packet_type);
|
||||
}
|
||||
|
||||
initiator.read_message(&frame.payload, &mut buf)?;
|
||||
let mut noise_transport = initiator.into_transport_mode()?;
|
||||
|
||||
// Receive assigned IP info (encrypted)
|
||||
let info_msg = match ws_stream.next().await {
|
||||
Some(Ok(Message::Binary(data))) => data.to_vec(),
|
||||
_ => anyhow::bail!("Expected IP info message"),
|
||||
};
|
||||
|
||||
let mut frame_buf = BytesMut::from(&info_msg[..]);
|
||||
let frame = <FrameCodec as tokio_util::codec::Decoder>::decode(&mut FrameCodec, &mut frame_buf)?
|
||||
.ok_or_else(|| anyhow::anyhow!("Incomplete IP info frame"))?;
|
||||
|
||||
let len = noise_transport.read_message(&frame.payload, &mut buf)?;
|
||||
let ip_info: serde_json::Value = serde_json::from_slice(&buf[..len])?;
|
||||
let assigned_ip = ip_info["assignedIp"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("Missing assignedIp in server response"))?
|
||||
.to_string();
|
||||
|
||||
*assigned_ip_ref.write().await = Some(assigned_ip.clone());
|
||||
*connected_since.write().await = Some(std::time::Instant::now());
|
||||
*state.write().await = ClientState::Connected;
|
||||
|
||||
info!("Connected to VPN, assigned IP: {}", assigned_ip);
|
||||
|
||||
// Spawn packet forwarding loop
|
||||
let assigned_ip_clone = assigned_ip.clone();
|
||||
tokio::spawn(client_loop(
|
||||
ws_sink,
|
||||
ws_stream,
|
||||
noise_transport,
|
||||
state,
|
||||
stats,
|
||||
shutdown_rx,
|
||||
config.keepalive_interval_secs.unwrap_or(30),
|
||||
));
|
||||
|
||||
Ok(assigned_ip_clone)
|
||||
}
|
||||
|
||||
/// Disconnect from the VPN server.
|
||||
pub async fn disconnect(&mut self) -> Result<()> {
|
||||
if let Some(tx) = self.shutdown_tx.take() {
|
||||
let _ = tx.send(()).await;
|
||||
}
|
||||
*self.assigned_ip.write().await = None;
|
||||
*self.connected_since.write().await = None;
|
||||
*self.state.write().await = ClientState::Disconnected;
|
||||
info!("Disconnected from VPN");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get current status.
|
||||
pub async fn get_status(&self) -> serde_json::Value {
|
||||
let state = self.state.read().await;
|
||||
let ip = self.assigned_ip.read().await;
|
||||
let since = self.connected_since.read().await;
|
||||
|
||||
let mut status = serde_json::json!({
|
||||
"state": format!("{}", *state),
|
||||
});
|
||||
|
||||
if let Some(ref ip) = *ip {
|
||||
status["assignedIp"] = serde_json::json!(ip);
|
||||
}
|
||||
if let Some(instant) = *since {
|
||||
status["uptimeSeconds"] = serde_json::json!(instant.elapsed().as_secs());
|
||||
}
|
||||
|
||||
status
|
||||
}
|
||||
|
||||
/// Get traffic statistics.
|
||||
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);
|
||||
|
||||
serde_json::json!({
|
||||
"bytesSent": stats.bytes_sent,
|
||||
"bytesReceived": stats.bytes_received,
|
||||
"packetsSent": stats.packets_sent,
|
||||
"packetsReceived": stats.packets_received,
|
||||
"keepalivesSent": stats.keepalives_sent,
|
||||
"keepalivesReceived": stats.keepalives_received,
|
||||
"uptimeSeconds": uptime,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// The main client packet forwarding loop (runs in a spawned task).
|
||||
async fn client_loop(
|
||||
mut ws_sink: futures_util::stream::SplitSink<transport::WsStream, Message>,
|
||||
mut ws_stream: futures_util::stream::SplitStream<transport::WsStream>,
|
||||
mut noise_transport: snow::TransportState,
|
||||
state: Arc<RwLock<ClientState>>,
|
||||
stats: Arc<RwLock<ClientStatistics>>,
|
||||
mut shutdown_rx: mpsc::Receiver<()>,
|
||||
keepalive_secs: u64,
|
||||
) {
|
||||
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! {
|
||||
msg = ws_stream.next() => {
|
||||
match msg {
|
||||
Some(Ok(Message::Binary(data))) => {
|
||||
let mut frame_buf = BytesMut::from(&data[..][..]);
|
||||
if let Ok(Some(frame)) = <FrameCodec as tokio_util::codec::Decoder>::decode(&mut FrameCodec, &mut frame_buf) {
|
||||
match frame.packet_type {
|
||||
PacketType::IpPacket => {
|
||||
match noise_transport.read_message(&frame.payload, &mut buf) {
|
||||
Ok(len) => {
|
||||
let mut s = stats.write().await;
|
||||
s.bytes_received += len as u64;
|
||||
s.packets_received += 1;
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("Decrypt error: {}", e);
|
||||
*state.write().await = ClientState::Error(e.to_string());
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
PacketType::KeepaliveAck => {
|
||||
stats.write().await.keepalives_received += 1;
|
||||
}
|
||||
PacketType::Disconnect => {
|
||||
info!("Server sent disconnect");
|
||||
*state.write().await = ClientState::Disconnected;
|
||||
break;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(Ok(Message::Close(_))) | None => {
|
||||
info!("Connection closed");
|
||||
*state.write().await = ClientState::Disconnected;
|
||||
break;
|
||||
}
|
||||
Some(Ok(Message::Ping(data))) => {
|
||||
let _ = ws_sink.send(Message::Pong(data)).await;
|
||||
}
|
||||
Some(Ok(_)) => continue,
|
||||
Some(Err(e)) => {
|
||||
error!("WebSocket error: {}", e);
|
||||
*state.write().await = ClientState::Error(e.to_string());
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = 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");
|
||||
*state.write().await = ClientState::Disconnected;
|
||||
break;
|
||||
}
|
||||
stats.write().await.keepalives_sent += 1;
|
||||
}
|
||||
}
|
||||
_ = shutdown_rx.recv() => {
|
||||
// Send disconnect frame
|
||||
let dc_frame = Frame {
|
||||
packet_type: PacketType::Disconnect,
|
||||
payload: vec![],
|
||||
};
|
||||
let mut frame_bytes = BytesMut::new();
|
||||
if <FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, dc_frame, &mut frame_bytes).is_ok() {
|
||||
let _ = ws_sink.send(Message::Binary(frame_bytes.to_vec().into())).await;
|
||||
}
|
||||
let _ = ws_sink.close().await;
|
||||
*state.write().await = ClientState::Disconnected;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
186
rust/src/codec.rs
Normal file
186
rust/src/codec.rs
Normal file
@@ -0,0 +1,186 @@
|
||||
use bytes::{Buf, BufMut, BytesMut};
|
||||
use tokio_util::codec::{Decoder, Encoder};
|
||||
|
||||
/// Packet types for the VPN binary protocol.
|
||||
#[repr(u8)]
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum PacketType {
|
||||
HandshakeInit = 0x01,
|
||||
HandshakeResp = 0x02,
|
||||
IpPacket = 0x10,
|
||||
Keepalive = 0x20,
|
||||
KeepaliveAck = 0x21,
|
||||
SessionResume = 0x30,
|
||||
SessionResumeOk = 0x31,
|
||||
SessionResumeErr = 0x32,
|
||||
Disconnect = 0x3F,
|
||||
}
|
||||
|
||||
impl PacketType {
|
||||
pub fn from_u8(v: u8) -> Option<Self> {
|
||||
match v {
|
||||
0x01 => Some(Self::HandshakeInit),
|
||||
0x02 => Some(Self::HandshakeResp),
|
||||
0x10 => Some(Self::IpPacket),
|
||||
0x20 => Some(Self::Keepalive),
|
||||
0x21 => Some(Self::KeepaliveAck),
|
||||
0x30 => Some(Self::SessionResume),
|
||||
0x31 => Some(Self::SessionResumeOk),
|
||||
0x32 => Some(Self::SessionResumeErr),
|
||||
0x3F => Some(Self::Disconnect),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A framed packet: [type:1B][length:4B][payload:NB]
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Frame {
|
||||
pub packet_type: PacketType,
|
||||
pub payload: Vec<u8>,
|
||||
}
|
||||
|
||||
/// Maximum frame payload size (64 KB).
|
||||
pub const MAX_FRAME_PAYLOAD: usize = 65536;
|
||||
|
||||
/// Header size: 1 byte type + 4 bytes length.
|
||||
pub const HEADER_SIZE: usize = 5;
|
||||
|
||||
/// tokio_util codec for Frame encode/decode over byte streams.
|
||||
pub struct FrameCodec;
|
||||
|
||||
impl Decoder for FrameCodec {
|
||||
type Item = Frame;
|
||||
type Error = std::io::Error;
|
||||
|
||||
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Frame>, Self::Error> {
|
||||
if src.len() < HEADER_SIZE {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let packet_type_byte = src[0];
|
||||
let length = u32::from_be_bytes([src[1], src[2], src[3], src[4]]) as usize;
|
||||
|
||||
if length > MAX_FRAME_PAYLOAD {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
format!("Frame payload too large: {} bytes", length),
|
||||
));
|
||||
}
|
||||
|
||||
if src.len() < HEADER_SIZE + length {
|
||||
// Reserve capacity for the remaining bytes
|
||||
src.reserve(HEADER_SIZE + length - src.len());
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let packet_type = PacketType::from_u8(packet_type_byte).ok_or_else(|| {
|
||||
std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
format!("Unknown packet type: 0x{:02x}", packet_type_byte),
|
||||
)
|
||||
})?;
|
||||
|
||||
src.advance(HEADER_SIZE);
|
||||
let payload = src.split_to(length).to_vec();
|
||||
|
||||
Ok(Some(Frame {
|
||||
packet_type,
|
||||
payload,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
impl Encoder<Frame> for FrameCodec {
|
||||
type Error = std::io::Error;
|
||||
|
||||
fn encode(&mut self, item: Frame, dst: &mut BytesMut) -> Result<(), Self::Error> {
|
||||
if item.payload.len() > MAX_FRAME_PAYLOAD {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
format!("Payload too large: {} bytes", item.payload.len()),
|
||||
));
|
||||
}
|
||||
|
||||
dst.reserve(HEADER_SIZE + item.payload.len());
|
||||
dst.put_u8(item.packet_type as u8);
|
||||
dst.put_u32(item.payload.len() as u32);
|
||||
dst.put_slice(&item.payload);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn roundtrip_frame() {
|
||||
let frame = Frame {
|
||||
packet_type: PacketType::IpPacket,
|
||||
payload: vec![1, 2, 3, 4, 5],
|
||||
};
|
||||
|
||||
let mut buf = BytesMut::new();
|
||||
let mut codec = FrameCodec;
|
||||
codec.encode(frame.clone(), &mut buf).unwrap();
|
||||
|
||||
let decoded = codec.decode(&mut buf).unwrap().unwrap();
|
||||
assert_eq!(decoded.packet_type, PacketType::IpPacket);
|
||||
assert_eq!(decoded.payload, vec![1, 2, 3, 4, 5]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn partial_frame() {
|
||||
let mut buf = BytesMut::from(&[0x10, 0x00, 0x00][..]);
|
||||
let mut codec = FrameCodec;
|
||||
// Not enough bytes for header
|
||||
assert!(codec.decode(&mut buf).unwrap().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reject_oversized_frame() {
|
||||
let mut buf = BytesMut::new();
|
||||
buf.put_u8(0x10); // IpPacket
|
||||
buf.put_u32(MAX_FRAME_PAYLOAD as u32 + 1);
|
||||
let mut codec = FrameCodec;
|
||||
assert!(codec.decode(&mut buf).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reject_unknown_packet_type() {
|
||||
let mut buf = BytesMut::new();
|
||||
buf.put_u8(0xFF);
|
||||
buf.put_u32(0);
|
||||
let mut codec = FrameCodec;
|
||||
assert!(codec.decode(&mut buf).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_packet_types_roundtrip() {
|
||||
let types = [
|
||||
PacketType::HandshakeInit,
|
||||
PacketType::HandshakeResp,
|
||||
PacketType::IpPacket,
|
||||
PacketType::Keepalive,
|
||||
PacketType::KeepaliveAck,
|
||||
PacketType::SessionResume,
|
||||
PacketType::SessionResumeOk,
|
||||
PacketType::SessionResumeErr,
|
||||
PacketType::Disconnect,
|
||||
];
|
||||
|
||||
for pt in types {
|
||||
let frame = Frame {
|
||||
packet_type: pt,
|
||||
payload: vec![42],
|
||||
};
|
||||
let mut buf = BytesMut::new();
|
||||
let mut codec = FrameCodec;
|
||||
codec.encode(frame, &mut buf).unwrap();
|
||||
let decoded = codec.decode(&mut buf).unwrap().unwrap();
|
||||
assert_eq!(decoded.packet_type, pt);
|
||||
assert_eq!(decoded.payload, vec![42]);
|
||||
}
|
||||
}
|
||||
}
|
||||
203
rust/src/crypto.rs
Normal file
203
rust/src/crypto.rs
Normal file
@@ -0,0 +1,203 @@
|
||||
use anyhow::Result;
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::STANDARD as BASE64;
|
||||
use snow::Builder;
|
||||
|
||||
/// Noise protocol pattern: NK (client knows server pubkey, no client auth at Noise level)
|
||||
const NOISE_PATTERN: &str = "Noise_NK_25519_ChaChaPoly_BLAKE2s";
|
||||
|
||||
/// Generate a new Noise static keypair.
|
||||
/// Returns (public_key_base64, private_key_base64).
|
||||
pub fn generate_keypair() -> Result<(String, String)> {
|
||||
let builder = Builder::new(NOISE_PATTERN.parse()?);
|
||||
let keypair = builder.generate_keypair()?;
|
||||
let public_key = BASE64.encode(&keypair.public);
|
||||
let private_key = BASE64.encode(&keypair.private);
|
||||
Ok((public_key, private_key))
|
||||
}
|
||||
|
||||
/// Generate a raw Noise static keypair (not base64 encoded).
|
||||
pub fn generate_keypair_raw() -> Result<snow::Keypair> {
|
||||
let builder = Builder::new(NOISE_PATTERN.parse()?);
|
||||
Ok(builder.generate_keypair()?)
|
||||
}
|
||||
|
||||
/// Create a Noise NK initiator (client side).
|
||||
/// The client knows the server's static public key.
|
||||
pub fn create_initiator(server_public_key: &[u8]) -> Result<snow::HandshakeState> {
|
||||
let builder = Builder::new(NOISE_PATTERN.parse()?);
|
||||
let state = builder
|
||||
.remote_public_key(server_public_key)
|
||||
.build_initiator()?;
|
||||
Ok(state)
|
||||
}
|
||||
|
||||
/// Create a Noise NK responder (server side).
|
||||
/// The server uses its static private key.
|
||||
pub fn create_responder(private_key: &[u8]) -> Result<snow::HandshakeState> {
|
||||
let builder = Builder::new(NOISE_PATTERN.parse()?);
|
||||
let state = builder
|
||||
.local_private_key(private_key)
|
||||
.build_responder()?;
|
||||
Ok(state)
|
||||
}
|
||||
|
||||
/// Perform the full Noise NK handshake between initiator and responder.
|
||||
/// Returns (initiator_transport, responder_transport).
|
||||
pub fn perform_handshake(
|
||||
mut initiator: snow::HandshakeState,
|
||||
mut responder: snow::HandshakeState,
|
||||
) -> Result<(snow::TransportState, snow::TransportState)> {
|
||||
let mut buf = vec![0u8; 65535];
|
||||
|
||||
// -> e, es (initiator sends)
|
||||
let len = initiator.write_message(&[], &mut buf)?;
|
||||
let msg1 = buf[..len].to_vec();
|
||||
|
||||
// <- e, ee (responder reads and responds)
|
||||
responder.read_message(&msg1, &mut buf)?;
|
||||
let len = responder.write_message(&[], &mut buf)?;
|
||||
let msg2 = buf[..len].to_vec();
|
||||
|
||||
// Initiator reads response
|
||||
initiator.read_message(&msg2, &mut buf)?;
|
||||
|
||||
let i_transport = initiator.into_transport_mode()?;
|
||||
let r_transport = responder.into_transport_mode()?;
|
||||
|
||||
Ok((i_transport, r_transport))
|
||||
}
|
||||
|
||||
/// XChaCha20-Poly1305 encryption for post-handshake data.
|
||||
/// Uses random 24-byte nonces (safe due to large nonce space).
|
||||
pub mod xchacha {
|
||||
use anyhow::Result;
|
||||
use chacha20poly1305::{
|
||||
XChaCha20Poly1305, XNonce,
|
||||
aead::{Aead, KeyInit},
|
||||
};
|
||||
use rand::RngCore;
|
||||
|
||||
pub const NONCE_SIZE: usize = 24;
|
||||
pub const TAG_SIZE: usize = 16;
|
||||
|
||||
/// Encrypt plaintext with XChaCha20-Poly1305.
|
||||
/// Returns: nonce (24 bytes) + ciphertext + tag (16 bytes).
|
||||
pub fn encrypt(key: &[u8; 32], plaintext: &[u8]) -> Result<Vec<u8>> {
|
||||
let cipher = XChaCha20Poly1305::new(key.into());
|
||||
let mut nonce_bytes = [0u8; NONCE_SIZE];
|
||||
rand::thread_rng().fill_bytes(&mut nonce_bytes);
|
||||
let nonce = XNonce::from_slice(&nonce_bytes);
|
||||
|
||||
let ciphertext = cipher
|
||||
.encrypt(nonce, plaintext)
|
||||
.map_err(|e| anyhow::anyhow!("Encryption failed: {}", e))?;
|
||||
|
||||
let mut output = Vec::with_capacity(NONCE_SIZE + ciphertext.len());
|
||||
output.extend_from_slice(&nonce_bytes);
|
||||
output.extend_from_slice(&ciphertext);
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
/// Decrypt data encrypted with `encrypt()`.
|
||||
/// Input: nonce (24 bytes) + ciphertext + tag (16 bytes).
|
||||
pub fn decrypt(key: &[u8; 32], data: &[u8]) -> Result<Vec<u8>> {
|
||||
if data.len() < NONCE_SIZE + TAG_SIZE {
|
||||
anyhow::bail!("Ciphertext too short: {} bytes", data.len());
|
||||
}
|
||||
let (nonce_bytes, ciphertext) = data.split_at(NONCE_SIZE);
|
||||
let nonce = XNonce::from_slice(nonce_bytes);
|
||||
let cipher = XChaCha20Poly1305::new(key.into());
|
||||
|
||||
let plaintext = cipher
|
||||
.decrypt(nonce, ciphertext)
|
||||
.map_err(|e| anyhow::anyhow!("Decryption failed: {}", e))?;
|
||||
Ok(plaintext)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn keypair_generation() {
|
||||
let (pub_key, priv_key) = generate_keypair().unwrap();
|
||||
// Base64-encoded 32-byte keys = 44 chars
|
||||
assert_eq!(pub_key.len(), 44);
|
||||
assert_eq!(priv_key.len(), 44);
|
||||
|
||||
// Verify they decode back to 32 bytes
|
||||
let pub_bytes = BASE64.decode(&pub_key).unwrap();
|
||||
let priv_bytes = BASE64.decode(&priv_key).unwrap();
|
||||
assert_eq!(pub_bytes.len(), 32);
|
||||
assert_eq!(priv_bytes.len(), 32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn noise_handshake() {
|
||||
let server_kp = generate_keypair_raw().unwrap();
|
||||
|
||||
let initiator = create_initiator(&server_kp.public).unwrap();
|
||||
let responder = create_responder(&server_kp.private).unwrap();
|
||||
|
||||
let (mut i_transport, mut r_transport) =
|
||||
perform_handshake(initiator, responder).unwrap();
|
||||
|
||||
// Test encrypted communication
|
||||
let mut buf = vec![0u8; 65535];
|
||||
let plaintext = b"hello from client";
|
||||
let len = i_transport.write_message(plaintext, &mut buf).unwrap();
|
||||
let mut out = vec![0u8; 65535];
|
||||
let len = r_transport.read_message(&buf[..len], &mut out).unwrap();
|
||||
assert_eq!(&out[..len], plaintext);
|
||||
|
||||
// Reverse direction
|
||||
let plaintext = b"hello from server";
|
||||
let len = r_transport.write_message(plaintext, &mut buf).unwrap();
|
||||
let len = i_transport.read_message(&buf[..len], &mut out).unwrap();
|
||||
assert_eq!(&out[..len], plaintext);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn xchacha_encrypt_decrypt() {
|
||||
let key = [42u8; 32];
|
||||
let plaintext = b"secret VPN payload data";
|
||||
|
||||
let encrypted = xchacha::encrypt(&key, plaintext).unwrap();
|
||||
// encrypted = nonce(24) + ciphertext + tag(16)
|
||||
assert_eq!(encrypted.len(), 24 + plaintext.len() + 16);
|
||||
|
||||
let decrypted = xchacha::decrypt(&key, &encrypted).unwrap();
|
||||
assert_eq!(decrypted, plaintext);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn xchacha_wrong_key_fails() {
|
||||
let key = [42u8; 32];
|
||||
let wrong_key = [43u8; 32];
|
||||
let plaintext = b"secret data";
|
||||
|
||||
let encrypted = xchacha::encrypt(&key, plaintext).unwrap();
|
||||
assert!(xchacha::decrypt(&wrong_key, &encrypted).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn xchacha_too_short_fails() {
|
||||
let key = [42u8; 32];
|
||||
let short = vec![0u8; 30]; // less than nonce + tag
|
||||
assert!(xchacha::decrypt(&key, &short).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn xchacha_tampered_fails() {
|
||||
let key = [42u8; 32];
|
||||
let plaintext = b"secret data";
|
||||
|
||||
let mut encrypted = xchacha::encrypt(&key, plaintext).unwrap();
|
||||
// Tamper with ciphertext
|
||||
let last = encrypted.len() - 1;
|
||||
encrypted[last] ^= 0xFF;
|
||||
assert!(xchacha::decrypt(&key, &encrypted).is_err());
|
||||
}
|
||||
}
|
||||
87
rust/src/keepalive.rs
Normal file
87
rust/src/keepalive.rs
Normal file
@@ -0,0 +1,87 @@
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::time::{interval, timeout};
|
||||
use tracing::{debug, warn};
|
||||
|
||||
/// Default keepalive interval (30 seconds).
|
||||
pub const DEFAULT_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30);
|
||||
|
||||
/// Default keepalive ACK timeout (10 seconds).
|
||||
pub const DEFAULT_KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
|
||||
/// Signals from the keepalive monitor.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum KeepaliveSignal {
|
||||
/// Time to send a keepalive ping.
|
||||
SendPing,
|
||||
/// Peer is considered dead (no ACK received within timeout).
|
||||
PeerDead,
|
||||
}
|
||||
|
||||
/// A keepalive monitor that emits signals on a channel.
|
||||
pub struct KeepaliveMonitor {
|
||||
interval: Duration,
|
||||
timeout_duration: Duration,
|
||||
signal_tx: mpsc::Sender<KeepaliveSignal>,
|
||||
ack_rx: mpsc::Receiver<()>,
|
||||
}
|
||||
|
||||
/// Handle returned to the caller to send ACKs and receive signals.
|
||||
pub struct KeepaliveHandle {
|
||||
pub signal_rx: mpsc::Receiver<KeepaliveSignal>,
|
||||
pub ack_tx: mpsc::Sender<()>,
|
||||
}
|
||||
|
||||
/// Create a keepalive monitor and its handle.
|
||||
pub fn create_keepalive(
|
||||
keepalive_interval: Option<Duration>,
|
||||
keepalive_timeout: Option<Duration>,
|
||||
) -> (KeepaliveMonitor, KeepaliveHandle) {
|
||||
let (signal_tx, signal_rx) = mpsc::channel(8);
|
||||
let (ack_tx, ack_rx) = mpsc::channel(8);
|
||||
|
||||
let monitor = KeepaliveMonitor {
|
||||
interval: keepalive_interval.unwrap_or(DEFAULT_KEEPALIVE_INTERVAL),
|
||||
timeout_duration: keepalive_timeout.unwrap_or(DEFAULT_KEEPALIVE_TIMEOUT),
|
||||
signal_tx,
|
||||
ack_rx,
|
||||
};
|
||||
|
||||
let handle = KeepaliveHandle { signal_rx, ack_tx };
|
||||
|
||||
(monitor, handle)
|
||||
}
|
||||
|
||||
impl KeepaliveMonitor {
|
||||
/// Run the keepalive loop. Blocks until the peer is dead or channels close.
|
||||
pub async fn run(mut self) {
|
||||
let mut ticker = interval(self.interval);
|
||||
ticker.tick().await; // skip first immediate tick
|
||||
|
||||
loop {
|
||||
ticker.tick().await;
|
||||
debug!("Sending keepalive ping signal");
|
||||
|
||||
if self.signal_tx.send(KeepaliveSignal::SendPing).await.is_err() {
|
||||
// Channel closed
|
||||
break;
|
||||
}
|
||||
|
||||
// Wait for ACK within timeout
|
||||
match timeout(self.timeout_duration, self.ack_rx.recv()).await {
|
||||
Ok(Some(())) => {
|
||||
debug!("Keepalive ACK received");
|
||||
}
|
||||
Ok(None) => {
|
||||
// Channel closed
|
||||
break;
|
||||
}
|
||||
Err(_) => {
|
||||
warn!("Keepalive ACK timeout — peer considered dead");
|
||||
let _ = self.signal_tx.send(KeepaliveSignal::PeerDead).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
13
rust/src/lib.rs
Normal file
13
rust/src/lib.rs
Normal file
@@ -0,0 +1,13 @@
|
||||
// Module declarations — each module is in its own file.
|
||||
// This file exists for library-level re-exports if needed.
|
||||
|
||||
pub mod management;
|
||||
pub mod codec;
|
||||
pub mod crypto;
|
||||
pub mod transport;
|
||||
pub mod keepalive;
|
||||
pub mod tunnel;
|
||||
pub mod network;
|
||||
pub mod server;
|
||||
pub mod client;
|
||||
pub mod reconnect;
|
||||
69
rust/src/main.rs
Normal file
69
rust/src/main.rs
Normal file
@@ -0,0 +1,69 @@
|
||||
use clap::Parser;
|
||||
use tracing::info;
|
||||
|
||||
#[global_allocator]
|
||||
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
|
||||
|
||||
use smartvpn_daemon::{management, crypto};
|
||||
|
||||
/// SmartVPN daemon — data plane for the @push.rocks/smartvpn TypeScript control plane.
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(name = "smartvpn_daemon", version, about)]
|
||||
struct Cli {
|
||||
/// Run in management mode (stdio: JSON lines on stdin/stdout)
|
||||
#[arg(long)]
|
||||
management: bool,
|
||||
|
||||
/// Run in management mode with Unix socket at the given path
|
||||
#[arg(long, value_name = "PATH")]
|
||||
management_socket: Option<String>,
|
||||
|
||||
/// Daemon mode: client or server
|
||||
#[arg(long, value_name = "MODE")]
|
||||
mode: Option<String>,
|
||||
|
||||
/// Generate a Noise keypair and print to stdout, then exit
|
||||
#[arg(long)]
|
||||
generate_keypair: bool,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
// Initialize tracing (logs go to stderr so stdout is clean for IPC)
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
|
||||
)
|
||||
.with_writer(std::io::stderr)
|
||||
.init();
|
||||
|
||||
let cli = Cli::parse();
|
||||
|
||||
if cli.generate_keypair {
|
||||
let keypair = crypto::generate_keypair()?;
|
||||
let output = serde_json::json!({
|
||||
"publicKey": keypair.0,
|
||||
"privateKey": keypair.1,
|
||||
});
|
||||
println!("{}", serde_json::to_string_pretty(&output)?);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let mode = cli.mode.unwrap_or_else(|| "client".to_string());
|
||||
if mode != "client" && mode != "server" {
|
||||
anyhow::bail!("Invalid mode '{}': must be 'client' or 'server'", mode);
|
||||
}
|
||||
|
||||
if let Some(socket_path) = cli.management_socket {
|
||||
info!("Starting management loop (socket mode) on {} for {} mode", socket_path, mode);
|
||||
management::management_loop_socket(&socket_path, &mode).await?;
|
||||
} else if cli.management {
|
||||
info!("Starting management loop (stdio mode) for {} mode", mode);
|
||||
management::management_loop_stdio(&mode).await?;
|
||||
} else {
|
||||
anyhow::bail!("Must specify --management or --management-socket <path>");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
364
rust/src/management.rs
Normal file
364
rust/src/management.rs
Normal file
@@ -0,0 +1,364 @@
|
||||
use anyhow::Result;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::{info, error, warn};
|
||||
|
||||
use crate::client::{ClientConfig, VpnClient};
|
||||
use crate::crypto;
|
||||
use crate::server::{ServerConfig, VpnServer};
|
||||
|
||||
// ============================================================================
|
||||
// IPC protocol types
|
||||
// ============================================================================
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ManagementRequest {
|
||||
pub id: String,
|
||||
pub method: String,
|
||||
#[serde(default)]
|
||||
pub params: serde_json::Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct ManagementResponse {
|
||||
pub id: String,
|
||||
pub success: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub result: Option<serde_json::Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct ManagementEvent {
|
||||
pub event: String,
|
||||
pub data: serde_json::Value,
|
||||
}
|
||||
|
||||
impl ManagementResponse {
|
||||
fn ok(id: String, result: serde_json::Value) -> Self {
|
||||
Self {
|
||||
id,
|
||||
success: true,
|
||||
result: Some(result),
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn err(id: String, message: String) -> Self {
|
||||
Self {
|
||||
id,
|
||||
success: false,
|
||||
result: None,
|
||||
error: Some(message),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Stdio management mode
|
||||
// ============================================================================
|
||||
|
||||
fn send_line_stdout(line: &str) {
|
||||
use std::io::Write;
|
||||
let stdout = std::io::stdout();
|
||||
let mut handle = stdout.lock();
|
||||
let _ = handle.write_all(line.as_bytes());
|
||||
let _ = handle.write_all(b"\n");
|
||||
let _ = handle.flush();
|
||||
}
|
||||
|
||||
fn send_response_stdout(response: &ManagementResponse) {
|
||||
match serde_json::to_string(response) {
|
||||
Ok(json) => send_line_stdout(&json),
|
||||
Err(e) => error!("Failed to serialize management response: {}", e),
|
||||
}
|
||||
}
|
||||
|
||||
fn send_event_stdout(event: &str, data: serde_json::Value) {
|
||||
let evt = ManagementEvent {
|
||||
event: event.to_string(),
|
||||
data,
|
||||
};
|
||||
match serde_json::to_string(&evt) {
|
||||
Ok(json) => send_line_stdout(&json),
|
||||
Err(e) => error!("Failed to serialize management event: {}", e),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn management_loop_stdio(mode: &str) -> Result<()> {
|
||||
let stdin = BufReader::new(tokio::io::stdin());
|
||||
let mut lines = stdin.lines();
|
||||
|
||||
let mut vpn_client = VpnClient::new();
|
||||
let mut vpn_server = VpnServer::new();
|
||||
|
||||
send_event_stdout("ready", serde_json::json!({ "mode": mode }));
|
||||
|
||||
loop {
|
||||
let line = match lines.next_line().await {
|
||||
Ok(Some(line)) => line,
|
||||
Ok(None) => {
|
||||
info!("Management stdin closed, shutting down");
|
||||
break;
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Error reading management stdin: {}", e);
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
let line = line.trim().to_string();
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let request: ManagementRequest = match serde_json::from_str(&line) {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
error!("Failed to parse management request: {}", e);
|
||||
send_response_stdout(&ManagementResponse::err(
|
||||
"unknown".to_string(),
|
||||
format!("Failed to parse request: {}", e),
|
||||
));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let response = match mode {
|
||||
"client" => handle_client_request(&request, &mut vpn_client).await,
|
||||
"server" => handle_server_request(&request, &mut vpn_server).await,
|
||||
_ => ManagementResponse::err(request.id.clone(), format!("Unknown mode: {}", mode)),
|
||||
};
|
||||
send_response_stdout(&response);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Socket management mode
|
||||
// ============================================================================
|
||||
|
||||
pub async fn management_loop_socket(socket_path: &str, mode: &str) -> Result<()> {
|
||||
let _ = tokio::fs::remove_file(socket_path).await;
|
||||
|
||||
let listener = tokio::net::UnixListener::bind(socket_path)?;
|
||||
info!("Management socket listening on {}", socket_path);
|
||||
|
||||
// Shared state behind Mutex for socket mode (multiple connections)
|
||||
let vpn_client = std::sync::Arc::new(Mutex::new(VpnClient::new()));
|
||||
let vpn_server = std::sync::Arc::new(Mutex::new(VpnServer::new()));
|
||||
|
||||
loop {
|
||||
match listener.accept().await {
|
||||
Ok((stream, _addr)) => {
|
||||
let mode = mode.to_string();
|
||||
let client = vpn_client.clone();
|
||||
let server = vpn_server.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) =
|
||||
handle_socket_connection(stream, &mode, client, server).await
|
||||
{
|
||||
warn!("Socket connection error: {}", e);
|
||||
}
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Failed to accept socket connection: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_socket_connection(
|
||||
stream: tokio::net::UnixStream,
|
||||
mode: &str,
|
||||
vpn_client: std::sync::Arc<Mutex<VpnClient>>,
|
||||
vpn_server: std::sync::Arc<Mutex<VpnServer>>,
|
||||
) -> Result<()> {
|
||||
let (reader, mut writer) = stream.into_split();
|
||||
let buf_reader = BufReader::new(reader);
|
||||
let mut lines = buf_reader.lines();
|
||||
|
||||
let ready_event = ManagementEvent {
|
||||
event: "ready".to_string(),
|
||||
data: serde_json::json!({ "mode": mode }),
|
||||
};
|
||||
let ready_json = serde_json::to_string(&ready_event)?;
|
||||
writer.write_all(ready_json.as_bytes()).await?;
|
||||
writer.write_all(b"\n").await?;
|
||||
writer.flush().await?;
|
||||
|
||||
loop {
|
||||
let line = match lines.next_line().await {
|
||||
Ok(Some(line)) => line,
|
||||
Ok(None) => {
|
||||
info!("Socket client disconnected");
|
||||
break;
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Error reading from socket client: {}", e);
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
let line = line.trim().to_string();
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let request: ManagementRequest = match serde_json::from_str(&line) {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
let resp = ManagementResponse::err(
|
||||
"unknown".to_string(),
|
||||
format!("Failed to parse request: {}", e),
|
||||
);
|
||||
let json = serde_json::to_string(&resp)?;
|
||||
writer.write_all(json.as_bytes()).await?;
|
||||
writer.write_all(b"\n").await?;
|
||||
writer.flush().await?;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let response = match mode {
|
||||
"client" => {
|
||||
let mut client = vpn_client.lock().await;
|
||||
handle_client_request(&request, &mut client).await
|
||||
}
|
||||
"server" => {
|
||||
let mut server = vpn_server.lock().await;
|
||||
handle_server_request(&request, &mut server).await
|
||||
}
|
||||
_ => ManagementResponse::err(request.id.clone(), format!("Unknown mode: {}", mode)),
|
||||
};
|
||||
|
||||
let json = serde_json::to_string(&response)?;
|
||||
writer.write_all(json.as_bytes()).await?;
|
||||
writer.write_all(b"\n").await?;
|
||||
writer.flush().await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Client command handlers
|
||||
// ============================================================================
|
||||
|
||||
async fn handle_client_request(
|
||||
request: &ManagementRequest,
|
||||
vpn_client: &mut VpnClient,
|
||||
) -> ManagementResponse {
|
||||
let id = request.id.clone();
|
||||
|
||||
match request.method.as_str() {
|
||||
"connect" => {
|
||||
let config: ClientConfig = match serde_json::from_value(
|
||||
request.params.get("config").cloned().unwrap_or_default(),
|
||||
) {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
return ManagementResponse::err(id, format!("Invalid config: {}", e));
|
||||
}
|
||||
};
|
||||
|
||||
match vpn_client.connect(config).await {
|
||||
Ok(assigned_ip) => {
|
||||
ManagementResponse::ok(id, serde_json::json!({ "assignedIp": assigned_ip }))
|
||||
}
|
||||
Err(e) => ManagementResponse::err(id, format!("Connect failed: {}", e)),
|
||||
}
|
||||
}
|
||||
"disconnect" => match vpn_client.disconnect().await {
|
||||
Ok(()) => ManagementResponse::ok(id, serde_json::json!({})),
|
||||
Err(e) => ManagementResponse::err(id, format!("Disconnect failed: {}", e)),
|
||||
},
|
||||
"getStatus" => {
|
||||
let status = vpn_client.get_status().await;
|
||||
ManagementResponse::ok(id, status)
|
||||
}
|
||||
"getStatistics" => {
|
||||
let stats = vpn_client.get_statistics().await;
|
||||
ManagementResponse::ok(id, stats)
|
||||
}
|
||||
_ => ManagementResponse::err(id, format!("Unknown client method: {}", request.method)),
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Server command handlers
|
||||
// ============================================================================
|
||||
|
||||
async fn handle_server_request(
|
||||
request: &ManagementRequest,
|
||||
vpn_server: &mut VpnServer,
|
||||
) -> ManagementResponse {
|
||||
let id = request.id.clone();
|
||||
|
||||
match request.method.as_str() {
|
||||
"start" => {
|
||||
let config: ServerConfig = match serde_json::from_value(
|
||||
request.params.get("config").cloned().unwrap_or_default(),
|
||||
) {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
return ManagementResponse::err(id, format!("Invalid config: {}", e));
|
||||
}
|
||||
};
|
||||
|
||||
match vpn_server.start(config).await {
|
||||
Ok(()) => ManagementResponse::ok(id, serde_json::json!({})),
|
||||
Err(e) => ManagementResponse::err(id, format!("Start failed: {}", e)),
|
||||
}
|
||||
}
|
||||
"stop" => match vpn_server.stop().await {
|
||||
Ok(()) => ManagementResponse::ok(id, serde_json::json!({})),
|
||||
Err(e) => ManagementResponse::err(id, format!("Stop failed: {}", e)),
|
||||
},
|
||||
"getStatus" => {
|
||||
let status = vpn_server.get_status();
|
||||
ManagementResponse::ok(id, status)
|
||||
}
|
||||
"getStatistics" => {
|
||||
let stats = vpn_server.get_statistics().await;
|
||||
match serde_json::to_value(&stats) {
|
||||
Ok(v) => ManagementResponse::ok(id, v),
|
||||
Err(e) => ManagementResponse::err(id, format!("Serialize error: {}", e)),
|
||||
}
|
||||
}
|
||||
"listClients" => {
|
||||
let clients = vpn_server.list_clients().await;
|
||||
match serde_json::to_value(&clients) {
|
||||
Ok(v) => ManagementResponse::ok(id, serde_json::json!({ "clients": v })),
|
||||
Err(e) => ManagementResponse::err(id, format!("Serialize error: {}", e)),
|
||||
}
|
||||
}
|
||||
"disconnectClient" => {
|
||||
let client_id = match request.params.get("clientId").and_then(|v| v.as_str()) {
|
||||
Some(id) => id.to_string(),
|
||||
None => {
|
||||
return ManagementResponse::err(id, "Missing clientId parameter".to_string())
|
||||
}
|
||||
};
|
||||
match vpn_server.disconnect_client(&client_id).await {
|
||||
Ok(()) => ManagementResponse::ok(id, serde_json::json!({})),
|
||||
Err(e) => ManagementResponse::err(id, format!("Disconnect client failed: {}", e)),
|
||||
}
|
||||
}
|
||||
"generateKeypair" => match crypto::generate_keypair() {
|
||||
Ok((public_key, private_key)) => ManagementResponse::ok(
|
||||
id,
|
||||
serde_json::json!({
|
||||
"publicKey": public_key,
|
||||
"privateKey": private_key,
|
||||
}),
|
||||
),
|
||||
Err(e) => ManagementResponse::err(id, format!("Keypair generation failed: {}", e)),
|
||||
},
|
||||
_ => ManagementResponse::err(id, format!("Unknown server method: {}", request.method)),
|
||||
}
|
||||
}
|
||||
195
rust/src/network.rs
Normal file
195
rust/src/network.rs
Normal file
@@ -0,0 +1,195 @@
|
||||
use anyhow::Result;
|
||||
use std::collections::HashMap;
|
||||
use std::net::Ipv4Addr;
|
||||
use tracing::{info, warn};
|
||||
|
||||
/// IP pool manager for allocating VPN client addresses from a subnet.
|
||||
pub struct IpPool {
|
||||
/// Network address (e.g., 10.8.0.0)
|
||||
network: Ipv4Addr,
|
||||
/// Prefix length (e.g., 24)
|
||||
prefix_len: u8,
|
||||
/// Allocated IPs: IP -> client_id
|
||||
allocated: HashMap<Ipv4Addr, String>,
|
||||
/// Next candidate offset (skipping .0 network and .1 gateway)
|
||||
next_offset: u32,
|
||||
}
|
||||
|
||||
impl IpPool {
|
||||
/// Create a new IP pool from a CIDR subnet string (e.g., "10.8.0.0/24").
|
||||
pub fn new(subnet: &str) -> Result<Self> {
|
||||
let parts: Vec<&str> = subnet.split('/').collect();
|
||||
if parts.len() != 2 {
|
||||
anyhow::bail!("Invalid subnet format: {}", subnet);
|
||||
}
|
||||
let network: Ipv4Addr = parts[0].parse()?;
|
||||
let prefix_len: u8 = parts[1].parse()?;
|
||||
if prefix_len > 30 {
|
||||
anyhow::bail!("Prefix too long for VPN pool: /{}", prefix_len);
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
network,
|
||||
prefix_len,
|
||||
allocated: HashMap::new(),
|
||||
next_offset: 2, // Skip .0 (network) and .1 (server/gateway)
|
||||
})
|
||||
}
|
||||
|
||||
/// Get the gateway/server address (first usable IP, e.g., 10.8.0.1).
|
||||
pub fn gateway_addr(&self) -> Ipv4Addr {
|
||||
let net_u32 = u32::from(self.network);
|
||||
Ipv4Addr::from(net_u32 + 1)
|
||||
}
|
||||
|
||||
/// Total number of usable client addresses in the pool.
|
||||
pub fn capacity(&self) -> u32 {
|
||||
let host_bits = 32 - self.prefix_len as u32;
|
||||
let total = 1u32 << host_bits;
|
||||
total.saturating_sub(3) // minus network, gateway, broadcast
|
||||
}
|
||||
|
||||
/// Allocate an IP for a client. Returns the assigned IP.
|
||||
pub fn allocate(&mut self, client_id: &str) -> Result<Ipv4Addr> {
|
||||
let host_bits = 32 - self.prefix_len as u32;
|
||||
let max_offset = (1u32 << host_bits) - 1; // broadcast offset
|
||||
|
||||
// Try to find a free IP starting from next_offset
|
||||
let start = self.next_offset;
|
||||
let mut offset = start;
|
||||
loop {
|
||||
if offset >= max_offset {
|
||||
offset = 2; // wrap around
|
||||
}
|
||||
|
||||
let ip = Ipv4Addr::from(u32::from(self.network) + offset);
|
||||
if !self.allocated.contains_key(&ip) {
|
||||
self.allocated.insert(ip, client_id.to_string());
|
||||
self.next_offset = offset + 1;
|
||||
info!("Allocated IP {} for client {}", ip, client_id);
|
||||
return Ok(ip);
|
||||
}
|
||||
|
||||
offset += 1;
|
||||
if offset == start {
|
||||
anyhow::bail!("IP pool exhausted");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Release an IP back to the pool.
|
||||
pub fn release(&mut self, ip: &Ipv4Addr) -> Option<String> {
|
||||
let client_id = self.allocated.remove(ip);
|
||||
if let Some(ref id) = client_id {
|
||||
info!("Released IP {} from client {}", ip, id);
|
||||
}
|
||||
client_id
|
||||
}
|
||||
|
||||
/// Number of currently allocated IPs.
|
||||
pub fn allocated_count(&self) -> usize {
|
||||
self.allocated.len()
|
||||
}
|
||||
}
|
||||
|
||||
/// Enable IP forwarding on Linux.
|
||||
pub fn enable_ip_forwarding() -> Result<()> {
|
||||
std::fs::write("/proc/sys/net/ipv4/ip_forward", "1")?;
|
||||
info!("Enabled IPv4 forwarding");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Set up NAT/masquerade using iptables for a given subnet and outbound interface.
|
||||
pub async fn setup_nat(subnet: &str, interface: &str) -> Result<()> {
|
||||
let output = tokio::process::Command::new("iptables")
|
||||
.args([
|
||||
"-t", "nat", "-A", "POSTROUTING",
|
||||
"-s", subnet,
|
||||
"-o", interface,
|
||||
"-j", "MASQUERADE",
|
||||
])
|
||||
.output()
|
||||
.await?;
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
anyhow::bail!("iptables NAT setup failed: {}", stderr);
|
||||
}
|
||||
|
||||
info!("NAT masquerade set up for {} via {}", subnet, interface);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Remove NAT/masquerade rule.
|
||||
pub async fn remove_nat(subnet: &str, interface: &str) -> Result<()> {
|
||||
let output = tokio::process::Command::new("iptables")
|
||||
.args([
|
||||
"-t", "nat", "-D", "POSTROUTING",
|
||||
"-s", subnet,
|
||||
"-o", interface,
|
||||
"-j", "MASQUERADE",
|
||||
])
|
||||
.output()
|
||||
.await?;
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
warn!("iptables NAT removal failed (may not exist): {}", stderr);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get the default outbound network interface name.
|
||||
pub fn get_default_interface() -> Result<String> {
|
||||
// Parse /proc/net/route for the default route
|
||||
let content = std::fs::read_to_string("/proc/net/route")?;
|
||||
for line in content.lines().skip(1) {
|
||||
let fields: Vec<&str> = line.split_whitespace().collect();
|
||||
if fields.len() >= 2 && fields[1] == "00000000" {
|
||||
return Ok(fields[0].to_string());
|
||||
}
|
||||
}
|
||||
anyhow::bail!("Could not determine default network interface")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn ip_pool_basic() {
|
||||
let mut pool = IpPool::new("10.8.0.0/24").unwrap();
|
||||
assert_eq!(pool.gateway_addr(), Ipv4Addr::new(10, 8, 0, 1));
|
||||
assert_eq!(pool.capacity(), 253); // 256 - 3 (net, gw, broadcast)
|
||||
|
||||
let ip1 = pool.allocate("client1").unwrap();
|
||||
assert_eq!(ip1, Ipv4Addr::new(10, 8, 0, 2));
|
||||
|
||||
let ip2 = pool.allocate("client2").unwrap();
|
||||
assert_eq!(ip2, Ipv4Addr::new(10, 8, 0, 3));
|
||||
|
||||
assert_eq!(pool.allocated_count(), 2);
|
||||
|
||||
pool.release(&ip1);
|
||||
assert_eq!(pool.allocated_count(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ip_pool_small_subnet() {
|
||||
let mut pool = IpPool::new("192.168.1.0/30").unwrap();
|
||||
// /30 = 4 addresses: .0 net, .1 gw, .2 client, .3 broadcast
|
||||
assert_eq!(pool.capacity(), 1);
|
||||
|
||||
let ip = pool.allocate("client1").unwrap();
|
||||
assert_eq!(ip, Ipv4Addr::new(192, 168, 1, 2));
|
||||
|
||||
// Pool should be exhausted
|
||||
assert!(pool.allocate("client2").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ip_pool_invalid_subnet() {
|
||||
assert!(IpPool::new("invalid").is_err());
|
||||
assert!(IpPool::new("10.8.0.0/31").is_err());
|
||||
}
|
||||
}
|
||||
149
rust/src/reconnect.rs
Normal file
149
rust/src/reconnect.rs
Normal file
@@ -0,0 +1,149 @@
|
||||
use std::time::Duration;
|
||||
use rand::Rng;
|
||||
use tracing::{info, warn};
|
||||
|
||||
/// Reconnection strategy with exponential backoff and jitter.
|
||||
pub struct ReconnectStrategy {
|
||||
/// Base delay (default: 1 second).
|
||||
pub base_delay: Duration,
|
||||
/// Maximum delay cap (default: 30 seconds).
|
||||
pub max_delay: Duration,
|
||||
/// Maximum number of attempts before giving up (0 = infinite).
|
||||
pub max_attempts: u32,
|
||||
/// Current attempt counter.
|
||||
attempts: u32,
|
||||
}
|
||||
|
||||
impl Default for ReconnectStrategy {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
base_delay: Duration::from_secs(1),
|
||||
max_delay: Duration::from_secs(30),
|
||||
max_attempts: 0,
|
||||
attempts: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ReconnectStrategy {
|
||||
pub fn new(base_delay: Duration, max_delay: Duration, max_attempts: u32) -> Self {
|
||||
Self {
|
||||
base_delay,
|
||||
max_delay,
|
||||
max_attempts,
|
||||
attempts: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the next backoff delay, or None if max attempts exceeded.
|
||||
pub fn next_delay(&mut self) -> Option<Duration> {
|
||||
if self.max_attempts > 0 && self.attempts >= self.max_attempts {
|
||||
warn!("Max reconnection attempts ({}) exceeded", self.max_attempts);
|
||||
return None;
|
||||
}
|
||||
|
||||
let base_ms = self.base_delay.as_millis() as u64;
|
||||
let exp_ms = base_ms.saturating_mul(1u64 << self.attempts.min(20));
|
||||
let max_ms = self.max_delay.as_millis() as u64;
|
||||
let capped_ms = exp_ms.min(max_ms);
|
||||
|
||||
// Add jitter: ±25%
|
||||
let jitter_range = capped_ms / 4;
|
||||
let jitter = if jitter_range > 0 {
|
||||
rand::thread_rng().gen_range(0..jitter_range * 2) as i64 - jitter_range as i64
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let final_ms = (capped_ms as i64 + jitter).max(0) as u64;
|
||||
|
||||
self.attempts += 1;
|
||||
let delay = Duration::from_millis(final_ms);
|
||||
info!(
|
||||
"Reconnect attempt {} in {:?}",
|
||||
self.attempts, delay
|
||||
);
|
||||
Some(delay)
|
||||
}
|
||||
|
||||
/// Reset the attempt counter (on successful connection).
|
||||
pub fn reset(&mut self) {
|
||||
self.attempts = 0;
|
||||
}
|
||||
|
||||
/// Current attempt number.
|
||||
pub fn attempts(&self) -> u32 {
|
||||
self.attempts
|
||||
}
|
||||
}
|
||||
|
||||
/// Session resume token — opaque blob the client sends to resume a session.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SessionToken {
|
||||
pub token: Vec<u8>,
|
||||
}
|
||||
|
||||
impl SessionToken {
|
||||
/// Generate a random session token.
|
||||
pub fn generate() -> Self {
|
||||
let mut token = vec![0u8; 32];
|
||||
rand::thread_rng().fill(&mut token[..]);
|
||||
Self { token }
|
||||
}
|
||||
|
||||
pub fn from_bytes(data: Vec<u8>) -> Self {
|
||||
Self { token: data }
|
||||
}
|
||||
|
||||
pub fn as_bytes(&self) -> &[u8] {
|
||||
&self.token
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn exponential_backoff() {
|
||||
let mut strategy = ReconnectStrategy::new(
|
||||
Duration::from_millis(100),
|
||||
Duration::from_secs(5),
|
||||
5,
|
||||
);
|
||||
|
||||
// Should get 5 delays
|
||||
for i in 0..5 {
|
||||
let delay = strategy.next_delay();
|
||||
assert!(delay.is_some(), "attempt {} should succeed", i);
|
||||
}
|
||||
|
||||
// 6th should fail
|
||||
assert!(strategy.next_delay().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_restores_attempts() {
|
||||
let mut strategy = ReconnectStrategy::new(
|
||||
Duration::from_millis(100),
|
||||
Duration::from_secs(5),
|
||||
2,
|
||||
);
|
||||
|
||||
strategy.next_delay();
|
||||
strategy.next_delay();
|
||||
assert!(strategy.next_delay().is_none());
|
||||
|
||||
strategy.reset();
|
||||
assert_eq!(strategy.attempts(), 0);
|
||||
assert!(strategy.next_delay().is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_token_generation() {
|
||||
let token = SessionToken::generate();
|
||||
assert_eq!(token.as_bytes().len(), 32);
|
||||
|
||||
let token2 = SessionToken::generate();
|
||||
assert_ne!(token.as_bytes(), token2.as_bytes()); // extremely unlikely to be equal
|
||||
}
|
||||
}
|
||||
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())
|
||||
}
|
||||
55
rust/src/transport.rs
Normal file
55
rust/src/transport.rs
Normal file
@@ -0,0 +1,55 @@
|
||||
use anyhow::Result;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use tokio::net::TcpStream;
|
||||
use tokio_tungstenite::{
|
||||
connect_async, tungstenite::Message, MaybeTlsStream, WebSocketStream,
|
||||
};
|
||||
use tracing::info;
|
||||
|
||||
/// A WebSocket connection (either client or server side).
|
||||
pub type WsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
|
||||
|
||||
/// Connect to a WebSocket server as a client.
|
||||
pub async fn connect_to_server(url: &str) -> Result<WsStream> {
|
||||
info!("Connecting to WebSocket server: {}", url);
|
||||
let (ws_stream, response) = connect_async(url).await?;
|
||||
info!("WebSocket connected, status: {}", response.status());
|
||||
Ok(ws_stream)
|
||||
}
|
||||
|
||||
/// Send a binary message over the WebSocket.
|
||||
pub async fn send_binary(ws: &mut WsStream, data: Vec<u8>) -> Result<()> {
|
||||
ws.send(Message::Binary(data.into())).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Receive the next binary message from the WebSocket.
|
||||
/// Returns None if the connection is closed.
|
||||
pub async fn recv_binary(ws: &mut WsStream) -> Result<Option<Vec<u8>>> {
|
||||
loop {
|
||||
match ws.next().await {
|
||||
Some(Ok(Message::Binary(data))) => return Ok(Some(data.to_vec())),
|
||||
Some(Ok(Message::Close(_))) => return Ok(None),
|
||||
Some(Ok(Message::Ping(data))) => {
|
||||
ws.send(Message::Pong(data)).await?;
|
||||
}
|
||||
Some(Ok(_)) => continue,
|
||||
Some(Err(e)) => return Err(anyhow::anyhow!("WebSocket error: {}", e)),
|
||||
None => return Ok(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a close frame.
|
||||
pub async fn close(ws: &mut WsStream) -> Result<()> {
|
||||
ws.close(None).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// WebSocket server acceptor — accepts a TcpStream and performs the WebSocket upgrade.
|
||||
pub async fn accept_connection(
|
||||
stream: TcpStream,
|
||||
) -> Result<WebSocketStream<MaybeTlsStream<TcpStream>>> {
|
||||
let ws = tokio_tungstenite::accept_async(MaybeTlsStream::Plain(stream)).await?;
|
||||
Ok(ws)
|
||||
}
|
||||
79
rust/src/tunnel.rs
Normal file
79
rust/src/tunnel.rs
Normal file
@@ -0,0 +1,79 @@
|
||||
use anyhow::Result;
|
||||
use std::net::Ipv4Addr;
|
||||
use tracing::info;
|
||||
|
||||
/// Configuration for creating a TUN device.
|
||||
pub struct TunConfig {
|
||||
pub name: String,
|
||||
pub address: Ipv4Addr,
|
||||
pub netmask: Ipv4Addr,
|
||||
pub mtu: u16,
|
||||
}
|
||||
|
||||
impl Default for TunConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
name: "smartvpn0".to_string(),
|
||||
address: Ipv4Addr::new(10, 8, 0, 1),
|
||||
netmask: Ipv4Addr::new(255, 255, 255, 0),
|
||||
mtu: 1420,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Create and configure a TUN device.
|
||||
/// Returns an async TUN device handle.
|
||||
pub fn create_tun(config: &TunConfig) -> Result<tun::AsyncDevice> {
|
||||
let mut tun_config = tun::Configuration::default();
|
||||
tun_config
|
||||
.tun_name(&config.name)
|
||||
.address(config.address)
|
||||
.netmask(config.netmask)
|
||||
.mtu(config.mtu as u16)
|
||||
.up();
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
tun_config.platform_config(|p| {
|
||||
p.ensure_root_privileges(true);
|
||||
});
|
||||
|
||||
let device = tun::create_as_async(&tun_config)?;
|
||||
info!(
|
||||
"TUN device {} created: addr={}, mtu={}",
|
||||
config.name, config.address, config.mtu
|
||||
);
|
||||
Ok(device)
|
||||
}
|
||||
|
||||
/// Set up routing: add a route for the VPN subnet through the TUN device.
|
||||
pub async fn add_route(subnet: &str, device_name: &str) -> Result<()> {
|
||||
let output = tokio::process::Command::new("ip")
|
||||
.args(["route", "add", subnet, "dev", device_name])
|
||||
.output()
|
||||
.await?;
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
// Ignore "File exists" errors (route already set)
|
||||
if !stderr.contains("File exists") {
|
||||
anyhow::bail!("Failed to add route: {}", stderr);
|
||||
}
|
||||
}
|
||||
|
||||
info!("Added route {} via {}", subnet, device_name);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Remove a route.
|
||||
pub async fn remove_route(subnet: &str, device_name: &str) -> Result<()> {
|
||||
let output = tokio::process::Command::new("ip")
|
||||
.args(["route", "del", subnet, "dev", device_name])
|
||||
.output()
|
||||
.await?;
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
tracing::warn!("Failed to remove route (may not exist): {}", stderr);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user