fix(wireguard): sync runtime peer management with client registration and derive the correct server public key from the WireGuard private key

This commit is contained in:
2026-03-31 02:11:29 +00:00
parent 42949b1233
commit 6e4cafe3c5
4 changed files with 152 additions and 15 deletions

View File

@@ -7,7 +7,7 @@ use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpListener;
use tokio::sync::{mpsc, Mutex, RwLock};
use tracing::{info, error, warn};
use tracing::{debug, info, error, warn};
use crate::acl;
use crate::client_registry::{ClientEntry, ClientRegistry};
@@ -594,6 +594,20 @@ impl VpnServer {
// Add to registry
state.client_registry.write().await.add(entry.clone())?;
// Register WG peer with the running WG listener (if active)
if self.wg_command_tx.is_some() {
let wg_peer_config = crate::wireguard::WgPeerConfig {
public_key: wg_pub.clone(),
preshared_key: None,
allowed_ips: vec![format!("{}/32", assigned_ip)],
endpoint: None,
persistent_keepalive: Some(25),
};
if let Err(e) = self.add_wg_peer(wg_peer_config).await {
warn!("Failed to register WG peer for client {}: {}", client_id, e);
}
}
// Build SmartVPN client config
let smartvpn_server_url = format!("wss://{}",
state.config.server_endpoint.as_deref()
@@ -610,6 +624,10 @@ impl VpnServer {
});
// Build WireGuard config string
let wg_server_pubkey = match &state.config.wg_private_key {
Some(wg_priv_key) => crate::wireguard::wg_public_key_from_private(wg_priv_key)?,
None => state.config.public_key.clone(),
};
let wg_endpoint = state.config.server_endpoint.as_deref()
.unwrap_or(&state.config.listen_addr);
let wg_allowed_ips = state.config.client_allowed_ips.as_ref()
@@ -622,7 +640,7 @@ impl VpnServer {
state.config.dns.as_ref()
.map(|d| format!("DNS = {}", d.join(", ")))
.unwrap_or_default(),
state.config.public_key,
wg_server_pubkey,
wg_allowed_ips,
wg_endpoint,
);
@@ -645,6 +663,14 @@ impl VpnServer {
let state = self.state.as_ref()
.ok_or_else(|| anyhow::anyhow!("Server not running"))?;
let entry = state.client_registry.write().await.remove(client_id)?;
// Remove WG peer from running listener
if self.wg_command_tx.is_some() {
if let Some(ref wg_key) = entry.wg_public_key {
if let Err(e) = self.remove_wg_peer(wg_key).await {
debug!("Failed to remove WG peer for client {}: {}", client_id, e);
}
}
}
// Release the IP if assigned
if let Some(ref ip_str) = entry.assigned_ip {
if let Ok(ip) = ip_str.parse::<Ipv4Addr>() {
@@ -731,6 +757,14 @@ impl VpnServer {
let state = self.state.as_ref()
.ok_or_else(|| anyhow::anyhow!("Server not running"))?;
// Capture old WG key before rotation (needed to remove from WG listener)
let old_wg_pub = {
let registry = state.client_registry.read().await;
let entry = registry.get_by_id(client_id)
.ok_or_else(|| anyhow::anyhow!("Client '{}' not found", client_id))?;
entry.wg_public_key.clone()
};
let (noise_pub, noise_priv) = crypto::generate_keypair()?;
let (wg_pub, wg_priv) = crate::wireguard::generate_wg_keypair();
@@ -749,6 +783,25 @@ impl VpnServer {
.and_then(|v| v.as_str())
.unwrap_or("0.0.0.0");
// Update WG listener: remove old peer, add new peer
if self.wg_command_tx.is_some() {
if let Some(ref old_key) = old_wg_pub {
if let Err(e) = self.remove_wg_peer(old_key).await {
debug!("Failed to remove old WG peer during rotation: {}", e);
}
}
let wg_peer_config = crate::wireguard::WgPeerConfig {
public_key: wg_pub.clone(),
preshared_key: None,
allowed_ips: vec![format!("{}/32", assigned_ip)],
endpoint: None,
persistent_keepalive: Some(25),
};
if let Err(e) = self.add_wg_peer(wg_peer_config).await {
warn!("Failed to register new WG peer during rotation: {}", e);
}
}
let smartvpn_server_url = format!("wss://{}",
state.config.server_endpoint.as_deref()
.unwrap_or(&state.config.listen_addr)
@@ -763,6 +816,10 @@ impl VpnServer {
"keepaliveIntervalSecs": state.config.keepalive_interval_secs,
});
let wg_server_pubkey = match &state.config.wg_private_key {
Some(wg_priv_key) => crate::wireguard::wg_public_key_from_private(wg_priv_key)?,
None => state.config.public_key.clone(),
};
let wg_endpoint = state.config.server_endpoint.as_deref()
.unwrap_or(&state.config.listen_addr);
let wg_allowed_ips = state.config.client_allowed_ips.as_ref()
@@ -774,7 +831,7 @@ impl VpnServer {
state.config.dns.as_ref()
.map(|d| format!("DNS = {}", d.join(", ")))
.unwrap_or_default(),
state.config.public_key,
wg_server_pubkey,
wg_allowed_ips,
wg_endpoint,
);
@@ -816,6 +873,10 @@ impl VpnServer {
}))
}
"wireguard" => {
let wg_server_pubkey = match &state.config.wg_private_key {
Some(wg_priv_key) => crate::wireguard::wg_public_key_from_private(wg_priv_key)?,
None => state.config.public_key.clone(),
};
let assigned_ip = entry.assigned_ip.as_deref().unwrap_or("0.0.0.0");
let wg_endpoint = state.config.server_endpoint.as_deref()
.unwrap_or(&state.config.listen_addr);
@@ -828,7 +889,7 @@ impl VpnServer {
state.config.dns.as_ref()
.map(|d| format!("DNS = {}", d.join(", ")))
.unwrap_or_default(),
state.config.public_key,
wg_server_pubkey,
wg_allowed_ips,
wg_endpoint,
);