smartvpn/rust/src/management.rs

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)),
    }
}