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:
@@ -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,
|
||||
);
|
||||
|
||||
Reference in New Issue
Block a user