Fix/agent auto assignment and vpn tunnel

agent/src/client.rs

@@ -117,6 +117,12 @@ pub struct ResourceGroupPeer {
     pub wg_tunnel_ip: String,
 }
 
+#[derive(Debug, Clone, serde::Deserialize)]
+pub struct VpnPeer {
+    pub client_public_key: String,
+    pub client_tunnel_ip: String,
+}
+
 pub struct ApiClient {
     client: Client,
     gateway_url: String,
@@ -327,6 +333,38 @@ impl ApiClient {
             .context("Failed to parse resource group peers response")
     }
 
+    pub async fn fetch_resource_group_vpn_peers(
+        &self,
+        api_key: &str,
+        resource_group_id: &str,
+    ) -> Result<Vec<VpnPeer>> {
+        let url = format!(
+            "{}/api/resource-groups/{}/vpn-peers",
+            self.gateway_url, resource_group_id
+        );
+
+        let resp = self
+            .client
+            .get(&url)
+            .header("X-API-Key", api_key)
+            .send()
+            .await
+            .context("Failed to fetch resource group vpn peers")?;
+
+        if !resp.status().is_success() {
+            let status = resp.status();
+            anyhow::bail!(
+                "Failed to fetch resource group vpn peers status={} {}",
+                status,
+                resp.text().await.unwrap_or_default()
+            );
+        }
+
+        resp.json::<Vec<VpnPeer>>()
+            .await
+            .context("Failed to parse resource group vpn peers response")
+    }
+
     pub async fn fetch_active_resource_group_ids(&self, api_key: &str) -> Result<Vec<String>> {
         let url = format!("{}/api/resource-groups/agent/active-ids", self.gateway_url);
 

agent/src/main.rs

@@ -57,7 +57,9 @@ async fn main() -> Result<()> {
         .and_then(|v| v.parse().ok())
         .unwrap_or(60);
 
-    let wg_endpoint = std::env::var("CSFX_WG_ENDPOINT").ok();
+    let wg_endpoint = std::env::var("CSFX_WG_ENDPOINT")
+        .ok()
+        .or_else(detect_wg_endpoint);
     let wg_tunnel_ip = if config::is_registered() {
         config::load_config().ok().and_then(|cfg| cfg.wg_tunnel_ip)
     } else {
@@ -109,6 +111,18 @@ async fn main() -> Result<()> {
 
     info!(agent_id = %agent_id, "Agent registered, starting heartbeat loop");
 
+    let mgmt_tunnel_ip = config::load_config().ok().and_then(|cfg| cfg.wg_tunnel_ip);
+    if let Some(ref tunnel_ip) = mgmt_tunnel_ip {
+        if let Err(e) =
+            wireguard::ensure_mgmt_interface(&wg_identity.private_key_b64, MGMT_WG_PORT, tunnel_ip)
+                .await
+        {
+            warn!(error = %e, "Failed to bring up management WireGuard interface");
+        }
+    } else {
+        warn!("No management tunnel IP available, VPN peering disabled for this agent");
+    }
+
     if let Err(e) = nftables::ensure_table_and_chain().await {
         warn!(error = %e, "Failed to initialize nftables resource group isolation");
     }
@@ -252,6 +266,7 @@ async fn run_heartbeat_loop(
                 process_volumes(client, agent_id, api_key, &mounted_volumes).await;
                 let resource_group_ids = process_workloads(client, api_key, &docker, &running_containers, &workload_phases, &mounted_volumes, &restart_counts, wg_private_key_b64).await;
                 sync_wireguard_peers(client, api_key, agent_id, &resource_group_ids).await;
+                sync_vpn_peers(client, api_key, &resource_group_ids).await;
                 cleanup_stale_resource_groups(client, api_key, wg_private_key_b64).await;
 
                 let statuses = build_container_statuses(&docker, &running_containers, &workload_phases).await;
@@ -680,6 +695,32 @@ async fn sync_wireguard_peers(
     }
 }
 
+async fn sync_vpn_peers(client: &client::ApiClient, api_key: &str, resource_group_ids: &[String]) {
+    let mut wg_peers: Vec<wireguard::Peer> = Vec::new();
+
+    for resource_group_id in resource_group_ids {
+        match client
+            .fetch_resource_group_vpn_peers(api_key, resource_group_id)
+            .await
+        {
+            Ok(peers) => {
+                wg_peers.extend(peers.into_iter().map(|p| wireguard::Peer {
+                    public_key: p.client_public_key,
+                    endpoint: None,
+                    allowed_ips: format!("{}/32", p.client_tunnel_ip),
+                }));
+            }
+            Err(e) => {
+                warn!(resource_group_id = %resource_group_id, error = %e, "Failed to fetch resource group vpn peers");
+            }
+        }
+    }
+
+    if let Err(e) = wireguard::reconcile_peers(wireguard::MGMT_INTERFACE_NAME, &wg_peers).await {
+        warn!(error = %e, "Failed to sync VPN client peers");
+    }
+}
+
 async fn cleanup_stale_resource_groups(
     client: &client::ApiClient,
     api_key: &str,
@@ -806,3 +847,12 @@ async fn push_workload_stats(
         warn!(error = %e, "Failed to push workload stats");
     }
 }
+
+const MGMT_WG_PORT: u16 = 51820;
+
+fn detect_wg_endpoint() -> Option<String> {
+    let socket = std::net::UdpSocket::bind("0.0.0.0:0").ok()?;
+    socket.connect("8.8.8.8:80").ok()?;
+    let local_ip = socket.local_addr().ok()?.ip();
+    Some(format!("{}:{}", local_ip, MGMT_WG_PORT))
+}

agent/src/wireguard.rs

@@ -9,6 +9,28 @@ pub struct Peer {
     pub allowed_ips: String,
 }
 
+pub const MGMT_INTERFACE_NAME: &str = "wgmgmt0";
+
+pub async fn ensure_mgmt_interface(
+    private_key_b64: &str,
+    listen_port: u16,
+    tunnel_ip: &str,
+) -> Result<()> {
+    let iface = MGMT_INTERFACE_NAME;
+
+    if !interface_exists(iface).await? {
+        run_ip(&["link", "add", "dev", iface, "type", "wireguard"]).await?;
+        set_private_key(iface, private_key_b64).await?;
+        run_wg(&["set", iface, "listen-port", &listen_port.to_string()]).await?;
+        run_ip(&["address", "add", &format!("{}/32", tunnel_ip), "dev", iface]).await?;
+        run_ip(&["link", "set", "up", "dev", iface]).await?;
+
+        info!(iface = %iface, tunnel_ip = %tunnel_ip, "Management WireGuard interface ready");
+    }
+
+    Ok(())
+}
+
 pub fn rg_interface_name(resource_group_id: &str) -> String {
     const FNV_OFFSET_BASIS: u32 = 0x811c9dc5;
     const FNV_PRIME: u32 = 0x01000193;
@@ -111,6 +133,38 @@ pub async fn set_peers(iface: &str, peers: &[Peer]) -> Result<()> {
     Ok(())
 }
 
+pub async fn reconcile_peers(iface: &str, peers: &[Peer]) -> Result<()> {
+    let current = list_peer_public_keys(iface).await?;
+    let desired: std::collections::HashSet<&str> =
+        peers.iter().map(|p| p.public_key.as_str()).collect();
+
+    for stale_key in current.iter().filter(|k| !desired.contains(k.as_str())) {
+        run_wg(&["set", iface, "peer", stale_key, "remove"]).await?;
+        info!(iface = %iface, public_key = %stale_key, "Removed stale WireGuard peer");
+    }
+
+    set_peers(iface, peers).await
+}
+
+async fn list_peer_public_keys(iface: &str) -> Result<Vec<String>> {
+    let output = Command::new("wg")
+        .args(["show", iface, "peers"])
+        .output()
+        .await
+        .context("failed to execute wg show peers")?;
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        anyhow::bail!("wg show peers failed iface={} stderr={}", iface, stderr);
+    }
+
+    Ok(String::from_utf8_lossy(&output.stdout)
+        .lines()
+        .map(|l| l.trim().to_string())
+        .filter(|l| !l.is_empty())
+        .collect())
+}
+
 async fn interface_exists(iface: &str) -> Result<bool> {
     let output = Command::new("ip")
         .args(["link", "show", iface])

control-plane/Dockerfile

@@ -21,6 +21,7 @@ FROM base AS dependencies
 COPY Cargo.toml Cargo.lock ./
 
 COPY agent/Cargo.toml ./agent/
+COPY agent/csfx-guest-init/Cargo.toml ./agent/csfx-guest-init/
 COPY control-plane/api-gateway/Cargo.toml ./control-plane/api-gateway/
 COPY control-plane/scheduler/Cargo.toml ./control-plane/scheduler/
 COPY control-plane/failover-controller/Cargo.toml ./control-plane/failover-controller/
@@ -34,6 +35,7 @@ COPY control-plane/csfx-migrate/Cargo.toml ./control-plane/csfx-migrate/
 COPY control-plane/csfx-updater/Cargo.toml ./control-plane/csfx-updater/
 
 RUN mkdir -p agent/src \
+    agent/csfx-guest-init/src \
     control-plane/api-gateway/src \
     control-plane/scheduler/src \
     control-plane/failover-controller/src \
@@ -46,6 +48,7 @@ RUN mkdir -p agent/src \
     control-plane/csfx-migrate/src \
     control-plane/csfx-updater/src \
     && echo "fn main() {}" > agent/src/main.rs \
+    && echo "fn main() {}" > agent/csfx-guest-init/src/main.rs \
     && echo "fn main() {}" > control-plane/api-gateway/src/main.rs \
     && echo "fn main() {}" > control-plane/scheduler/src/main.rs \
     && echo "fn main() {}" > control-plane/failover-controller/src/main.rs \

control-plane/api-gateway/Cargo.toml

@@ -31,6 +31,7 @@ http = { workspace = true }
 rcgen = { workspace = true }
 rustls = { workspace = true }
 ring = { workspace = true }
+x25519-dalek = { workspace = true }
 
 # Serialization
 serde = { workspace = true }