NexaVPN/desktop-client/src-tauri/src/lib.rs

mod tunnel_manager;

use std::{fs, net::TcpListener, path::PathBuf, sync::Mutex};

use base64::{engine::general_purpose::STANDARD, Engine as _};
use rand_core::OsRng;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use tauri::{
    menu::{MenuBuilder, MenuItem, MenuItemBuilder, PredefinedMenuItem},
    tray::{MouseButton, MouseButtonState, TrayIconBuilder, TrayIconEvent},
    AppHandle, Manager, State, WebviewWindow, Window, WindowEvent, Wry,
};
use x25519_dalek::{PublicKey, StaticSecret};

const PROFILE_NAME: &str = "NexaVPN";
const MAIN_WINDOW_LABEL: &str = "main";
const SINGLE_INSTANCE_ADDR: &str = "127.0.0.1:53190";

struct AppState {
    session: Mutex<Option<SessionState>>,
    tray: Mutex<Option<TrayState>>,
    single_instance_lock: TcpListener,
}

struct TrayState {
    status_item: MenuItem<Wry>,
    received_item: MenuItem<Wry>,
    sent_item: MenuItem<Wry>,
    toggle_item: MenuItem<Wry>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct SessionState {
    access_token: String,
    refresh_token: String,
    server_url: String,
    profile_path: String,
    private_key: String,
    enrollment: EnrollmentResult,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct EnrollmentPayload {
    server_url: String,
    username: String,
    password: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct EnrollmentResult {
    assigned_ip: String,
    resources: Vec<String>,
    profile_revision: u32,
    gateway_endpoint: String,
    profile_path: String,
    last_sync_time: String,
    tunnel_strategy: String,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct TunnelMetrics {
    active: bool,
    rx_bytes: u64,
    tx_bytes: u64,
}

#[derive(Debug, Serialize)]
struct LoginRequest<'a> {
    username: &'a str,
    password: &'a str,
}

#[derive(Debug, Serialize)]
struct RefreshRequest<'a> {
    refresh_token: &'a str,
}

#[derive(Debug, Deserialize)]
struct LoginResponse {
    #[serde(rename = "access_token")]
    access_token: String,
    #[serde(rename = "refresh_token")]
    refresh_token: String,
}

#[derive(Debug, Serialize)]
struct EnrollRequest<'a> {
    name: &'a str,
    platform: &'a str,
    #[serde(rename = "os_version")]
    os_version: &'a str,
    #[serde(rename = "app_version")]
    app_version: &'a str,
    #[serde(rename = "device_fingerprint")]
    device_fingerprint: String,
    #[serde(rename = "public_key")]
    public_key: String,
}

#[derive(Debug, Deserialize)]
struct EnrollResponse {
    device: DeviceView,
    peer: PeerView,
    profile: ProfileView,
    resources: Vec<ResourceView>,
}

#[derive(Debug, Deserialize)]
struct DeviceView {
    id: String,
}

#[derive(Debug, Deserialize)]
struct PeerView {
    #[serde(rename = "assigned_ip")]
    assigned_ip: String,
    gateway: GatewayView,
    #[serde(rename = "profile_revision")]
    profile_revision: u32,
}

#[derive(Debug, Deserialize)]
struct GatewayView {
    endpoint: String,
}

#[derive(Debug, Deserialize)]
struct ResourceView {
    value: String,
}

#[derive(Debug, Deserialize)]
struct ProfileView {
    content: String,
}

#[tauri::command]
async fn enroll_device(
    app: AppHandle,
    payload: EnrollmentPayload,
    state: State<'_, AppState>,
) -> Result<EnrollmentResult, String> {
    if payload.server_url.trim().is_empty() || payload.username.trim().is_empty() || payload.password.trim().is_empty() {
        return Err("Server URL, username, and password are required".into());
    }

    let client = Client::builder()
        .use_rustls_tls()
        .build()
        .map_err(|err| err.to_string())?;

    let login_response = client
        .post(format!("{}/api/v1/auth/login", payload.server_url.trim_end_matches('/')))
        .json(&LoginRequest {
            username: &payload.username,
            password: &payload.password,
        })
        .send()
        .await
        .map_err(|err| format!("Login failed: {}", err))?;

    if !login_response.status().is_success() {
        return Err(format!("Login failed with status {}", login_response.status()));
    }

    let login = login_response
        .json::<LoginResponse>()
        .await
        .map_err(|err| format!("Unable to decode login response: {}", err))?;

    let (private_key, public_key) = generate_keypair();
    let enroll_response = client
        .post(format!("{}/api/v1/devices/enroll", payload.server_url.trim_end_matches('/')))
        .bearer_auth(&login.access_token)
        .json(&EnrollRequest {
            name: "NexaVPN Desktop",
            platform: if cfg!(target_os = "macos") { "macos" } else { "windows" },
            os_version: std::env::consts::OS,
            app_version: env!("CARGO_PKG_VERSION"),
            device_fingerprint: build_fingerprint(&payload.server_url, &payload.username, &public_key),
            public_key,
        })
        .send()
        .await
        .map_err(|err| format!("Enrollment failed: {}", err))?;

    if !enroll_response.status().is_success() {
        let status = enroll_response.status();
        let body = enroll_response
            .text()
            .await
            .unwrap_or_else(|_| "<unable to read response body>".into());
        return Err(format!("Enrollment failed with status {}: {}", status, body));
    }

    let enroll = enroll_response
        .json::<EnrollResponse>()
        .await
        .map_err(|err| format!("Unable to decode enrollment response: {}", err))?;

    let profile_content = materialize_profile(&enroll.profile.content, &private_key);
    let profile_path = write_profile(&app, &profile_content)?;
    let result = EnrollmentResult {
        assigned_ip: enroll.peer.assigned_ip,
        resources: enroll.resources.into_iter().map(|resource| resource.value).collect(),
        profile_revision: enroll.peer.profile_revision,
        gateway_endpoint: enroll.peer.gateway.endpoint,
        profile_path: profile_path.display().to_string(),
        last_sync_time: now_label(),
        tunnel_strategy: tunnel_manager::current_tunnel_strategy().into(),
    };

    let session_state = SessionState {
        access_token: login.access_token,
        refresh_token: login.refresh_token,
        server_url: payload.server_url,
        profile_path: result.profile_path.clone(),
        private_key,
        enrollment: result.clone(),
    };

    write_session_state(&app, &session_state)?;
    let mut session = state.session.lock().map_err(|_| "Unable to store client state".to_string())?;
    *session = Some(session_state);
    drop(session);
    refresh_tray_menu(&app);

    Ok(result)
}

#[tauri::command]
fn load_state(app: AppHandle, state: State<'_, AppState>) -> Result<Option<EnrollmentResult>, String> {
    let loaded = read_session_state(&app)?;
    let mut session = state.session.lock().map_err(|_| "Unable to read client state".to_string())?;
    *session = loaded.clone();
    drop(session);
    refresh_tray_menu(&app);
    Ok(loaded.map(|value| value.enrollment))
}

#[tauri::command]
fn clear_session(app: AppHandle, state: State<'_, AppState>) -> Result<(), String> {
    let app_dir = ensure_app_dir(&app)?;
    let session_path = app_dir.join("session.json");
    let profile_path = app_dir.join(format!("{}.conf", PROFILE_NAME));

    if session_path.exists() {
        fs::remove_file(&session_path).map_err(|err| format!("Unable to remove session state: {}", err))?;
    }
    if profile_path.exists() {
        fs::remove_file(&profile_path).map_err(|err| format!("Unable to remove profile: {}", err))?;
    }

    let mut session = state.session.lock().map_err(|_| "Unable to clear client state".to_string())?;
    *session = None;
    drop(session);
    refresh_tray_menu(&app);
    Ok(())
}

#[tauri::command]
async fn sync_profile(app: AppHandle, state: State<'_, AppState>) -> Result<EnrollmentResult, String> {
    let session_state = sync_current_session(&app).await?;
    refresh_tray_menu(&app);
    Ok(session_state.enrollment)
}

#[tauri::command]
async fn connect_tunnel(app: AppHandle) -> Result<EnrollmentResult, String> {
    let session_state = sync_current_session(&app).await?;
    let result = tunnel_manager::connect(&app, std::path::Path::new(&session_state.profile_path));
    refresh_tray_menu(&app);
    result?;
    Ok(session_state.enrollment)
}

#[tauri::command]
fn disconnect_tunnel(app: AppHandle, state: State<'_, AppState>) -> Result<(), String> {
    let profile_path = {
        let session = state.session.lock().map_err(|_| "Unable to read client state".to_string())?;
        let session = session.as_ref().ok_or_else(|| "No active session is available".to_string())?;
        session.profile_path.clone()
    };
    let result = tunnel_manager::disconnect(&app, std::path::Path::new(&profile_path));
    refresh_tray_menu(&app);
    result
}

#[tauri::command]
fn tunnel_status(app: AppHandle, state: State<'_, AppState>) -> Result<bool, String> {
    let profile_path = {
        let session = state.session.lock().map_err(|_| "Unable to read client state".to_string())?;
        let session = session.as_ref().ok_or_else(|| "No active session is available".to_string())?;
        session.profile_path.clone()
    };
    tunnel_manager::is_active(&app, std::path::Path::new(&profile_path))
}

#[tauri::command]
fn tunnel_metrics(app: AppHandle, state: State<'_, AppState>) -> Result<TunnelMetrics, String> {
    let profile_path = {
        let session = state.session.lock().map_err(|_| "Unable to read client state".to_string())?;
        let session = session.as_ref().ok_or_else(|| "No active session is available".to_string())?;
        session.profile_path.clone()
    };
    let metrics = tunnel_manager::metrics(&app, std::path::Path::new(&profile_path))?;
    let mapped = TunnelMetrics {
        active: metrics.active,
        rx_bytes: metrics.rx_bytes,
        tx_bytes: metrics.tx_bytes,
    };
    let _ = update_tray_menu(&app, mapped);
    Ok(mapped)
}

fn generate_keypair() -> (String, String) {
    let private = StaticSecret::random_from_rng(OsRng);
    let public = PublicKey::from(&private);
    (
        STANDARD.encode(private.to_bytes()),
        STANDARD.encode(public.to_bytes()),
    )
}

fn now_label() -> String {
    "Just now".into()
}

fn build_fingerprint(server_url: &str, username: &str, public_key: &str) -> String {
    format!("nexavpn:{}:{}:{}", server_url, username, public_key)
}

fn materialize_profile(profile_content: &str, private_key: &str) -> String {
    profile_content
        .replace("__CLIENT_GENERATED_PRIVATE_KEY__", private_key)
        .replace("__CLIENT_PRIVATE_KEY_REQUIRED__", private_key)
}

fn write_profile(app: &AppHandle, profile_content: &str) -> Result<PathBuf, String> {
    let app_dir = ensure_app_dir(app)?;
    let profile_path = app_dir.join(format!("{}.conf", PROFILE_NAME));
    fs::write(&profile_path, profile_content).map_err(|err| format!("Unable to store profile: {}", err))?;
    Ok(profile_path)
}

fn write_session_state(app: &AppHandle, session: &SessionState) -> Result<(), String> {
    let app_dir = ensure_app_dir(app)?;
    let session_path = app_dir.join("session.json");
    let json = serde_json::to_vec_pretty(session).map_err(|err| err.to_string())?;
    fs::write(session_path, json).map_err(|err| format!("Unable to persist session state: {}", err))
}

fn read_session_state(app: &AppHandle) -> Result<Option<SessionState>, String> {
    let app_dir = ensure_app_dir(app)?;
    let session_path = app_dir.join("session.json");
    if !session_path.exists() {
        return Ok(None);
    }
    let raw = fs::read(session_path).map_err(|err| format!("Unable to read session state: {}", err))?;
    let value = serde_json::from_slice::<SessionState>(&raw).map_err(|err| format!("Unable to parse session state: {}", err))?;
    Ok(Some(value))
}

fn ensure_app_dir(app: &AppHandle) -> Result<PathBuf, String> {
    let dir = app
        .path()
        .app_config_dir()
        .map_err(|err| format!("Unable to resolve app config dir: {}", err))?;
    fs::create_dir_all(&dir).map_err(|err| format!("Unable to create app dir: {}", err))?;
    Ok(dir)
}

fn format_data_size(bytes: u64) -> String {
    if bytes == 0 {
        return "0 B".into();
    }

    let units = ["B", "KB", "MB", "GB", "TB"];
    let mut value = bytes as f64;
    let mut unit_index = 0;

    while value >= 1024.0 && unit_index < units.len() - 1 {
        value /= 1024.0;
        unit_index += 1;
    }

    if value >= 100.0 || unit_index == 0 {
        format!("{value:.0} {}", units[unit_index])
    } else {
        format!("{value:.1} {}", units[unit_index])
    }
}

fn current_metrics(app: &AppHandle) -> Result<TunnelMetrics, String> {
    let state = app.state::<AppState>();
    let profile_path = {
        let session = state.session.lock().map_err(|_| "Unable to read client state".to_string())?;
        let session = session.as_ref().ok_or_else(|| "No active session is available".to_string())?;
        session.profile_path.clone()
    };
    match tunnel_manager::metrics(app, std::path::Path::new(&profile_path)) {
        Ok(metrics) => Ok(TunnelMetrics {
            active: metrics.active,
            rx_bytes: metrics.rx_bytes,
            tx_bytes: metrics.tx_bytes,
        }),
        Err(_) => {
            let active = tunnel_manager::is_active(app, std::path::Path::new(&profile_path))?;
            Ok(TunnelMetrics {
                active,
                rx_bytes: 0,
                tx_bytes: 0,
            })
        }
    }
}

async fn sync_current_session(app: &AppHandle) -> Result<SessionState, String> {
    let mut existing = {
        let state = app.state::<AppState>();
        let session = state.session.lock().map_err(|_| "Unable to read client state".to_string())?;
        session.clone().ok_or_else(|| "No enrolled profile is available yet".to_string())?
    };

    let client = Client::builder()
        .use_rustls_tls()
        .build()
        .map_err(|err| err.to_string())?;

    let profile_url = format!("{}/api/v1/me/profile", existing.server_url.trim_end_matches('/'));
    let mut response = client
        .get(&profile_url)
        .bearer_auth(&existing.access_token)
        .send()
        .await
        .map_err(|err| format!("Profile sync failed: {}", err))?;

    if response.status().as_u16() == 401 {
        let refresh = client
            .post(format!("{}/api/v1/auth/refresh", existing.server_url.trim_end_matches('/')))
            .json(&RefreshRequest {
                refresh_token: &existing.refresh_token,
            })
            .send()
            .await
            .map_err(|err| format!("Session refresh failed: {}", err))?;

        if !refresh.status().is_success() {
            let status = refresh.status();
            let body = refresh
                .text()
                .await
                .unwrap_or_else(|_| "<unable to read response body>".into());
            return Err(format!("Session refresh failed with status {}: {}", status, body));
        }

        let refreshed = refresh
            .json::<LoginResponse>()
            .await
            .map_err(|err| format!("Unable to decode refresh response: {}", err))?;

        existing.access_token = refreshed.access_token;
        existing.refresh_token = refreshed.refresh_token;

        response = client
            .get(&profile_url)
            .bearer_auth(&existing.access_token)
            .send()
            .await
            .map_err(|err| format!("Profile sync failed: {}", err))?;
    }

    if !response.status().is_success() {
        let status = response.status();
        let body = response
            .text()
            .await
            .unwrap_or_else(|_| "<unable to read response body>".into());
        return Err(format!("Profile sync failed with status {}: {}", status, body));
    }

    let enroll = response
        .json::<EnrollResponse>()
        .await
        .map_err(|err| format!("Unable to decode profile sync response: {}", err))?;

    let profile_content = materialize_profile(&enroll.profile.content, &existing.private_key);
    let profile_path = write_profile(app, &profile_content)?;
    let result = EnrollmentResult {
        assigned_ip: enroll.peer.assigned_ip,
        resources: enroll.resources.into_iter().map(|resource| resource.value).collect(),
        profile_revision: enroll.peer.profile_revision,
        gateway_endpoint: enroll.peer.gateway.endpoint,
        profile_path: profile_path.display().to_string(),
        last_sync_time: now_label(),
        tunnel_strategy: tunnel_manager::current_tunnel_strategy().into(),
    };

    let session_state = SessionState {
        access_token: existing.access_token,
        refresh_token: existing.refresh_token,
        server_url: existing.server_url,
        profile_path: result.profile_path.clone(),
        private_key: existing.private_key,
        enrollment: result.clone(),
    };

    write_session_state(app, &session_state)?;
    let state = app.state::<AppState>();
    let mut session = state.session.lock().map_err(|_| "Unable to store client state".to_string())?;
    *session = Some(session_state.clone());
    Ok(session_state)
}

fn update_tray_menu(app: &AppHandle, metrics: TunnelMetrics) -> Result<(), String> {
    let state = app.state::<AppState>();
    let tray = state.tray.lock().map_err(|_| "Unable to update tray state".to_string())?;
    let Some(tray) = tray.as_ref() else {
        return Ok(());
    };

    let status_text = if metrics.active {
        "Status: Connected"
    } else {
        "Status: Disconnected"
    };
    let toggle_text = if metrics.active { "Disconnect NexaVPN" } else { "Connect NexaVPN" };

    let _ = tray.status_item.set_text(status_text);
    let _ = tray.received_item.set_text(format!("Received: {}", format_data_size(metrics.rx_bytes)));
    let _ = tray.sent_item.set_text(format!("Sent: {}", format_data_size(metrics.tx_bytes)));
    let _ = tray.toggle_item.set_text(toggle_text);
    Ok(())
}

fn refresh_tray_menu(app: &AppHandle) {
    let metrics = current_metrics(app).unwrap_or(TunnelMetrics {
        active: false,
        rx_bytes: 0,
        tx_bytes: 0,
    });
    let _ = update_tray_menu(app, metrics);
}

fn toggle_tray_connection(app: &AppHandle) {
    let metrics = current_metrics(app).unwrap_or(TunnelMetrics {
        active: false,
        rx_bytes: 0,
        tx_bytes: 0,
    });

    let state = app.state::<AppState>();
    let profile_path = match state.session.lock() {
        Ok(session) => match session.as_ref() {
            Some(session) => session.profile_path.clone(),
            None => return,
        },
        Err(_) => return,
    };

    if metrics.active {
        if tunnel_manager::disconnect(app, std::path::Path::new(&profile_path)).is_ok() {
            refresh_tray_menu(app);
        }
        return;
    }

    let app_handle = app.clone();
    tauri::async_runtime::spawn(async move {
        if let Ok(session_state) = sync_current_session(&app_handle).await {
            let _ = tunnel_manager::connect(&app_handle, std::path::Path::new(&session_state.profile_path));
            refresh_tray_menu(&app_handle);
        }
    });
}

fn restore_webview_window(window: &WebviewWindow) {
    let _ = window.show();
    let _ = window.unminimize();
    let _ = window.set_focus();
}

fn hide_main_window(window: &Window) {
    let _ = window.hide();
}

#[cfg_attr(mobile, tauri::mobile_entry_point)]
pub fn run() {
    tauri::Builder::default()
        .setup(|app| {
            let single_instance_lock = TcpListener::bind(SINGLE_INSTANCE_ADDR)
                .map_err(|_| format!("{} is already running.", app.package_info().name))?;

            let status_item = MenuItemBuilder::with_id("status", "Status: Disconnected").build(app)?;
            let received_item = MenuItemBuilder::with_id("received", "Received: 0 B").build(app)?;
            let sent_item = MenuItemBuilder::with_id("sent", "Sent: 0 B").build(app)?;
            let toggle_item = MenuItemBuilder::with_id("toggle", "Connect NexaVPN").build(app)?;
            let separator_top = PredefinedMenuItem::separator(app)?;
            let separator_bottom = PredefinedMenuItem::separator(app)?;
            let open_item = MenuItemBuilder::with_id("open", "Open NexaVPN Client").build(app)?;
            let quit_item = MenuItemBuilder::with_id("quit", "Quit NexaVPN Client").build(app)?;
            let menu = MenuBuilder::new(app)
                .items(&[
                    &status_item,
                    &received_item,
                    &sent_item,
                    &separator_top,
                    &toggle_item,
                    &separator_bottom,
                    &open_item,
                    &quit_item,
                ])
                .build()?;

            let mut tray = TrayIconBuilder::new().menu(&menu).show_menu_on_left_click(false);
            if let Some(icon) = app.default_window_icon() {
                tray = tray.icon(icon.clone());
            }

            tray
                .on_menu_event(|app, event| match event.id().as_ref() {
                    "status" | "received" | "sent" => {}
                    "toggle" => {
                        toggle_tray_connection(&app);
                    }
                    "open" => {
                        if let Some(window) = app.get_webview_window(MAIN_WINDOW_LABEL) {
                            restore_webview_window(&window);
                        }
                    }
                    "quit" => {
                        app.exit(0);
                    }
                    _ => {}
                })
                .on_tray_icon_event(|tray, event| {
                    if let TrayIconEvent::Click {
                        button: MouseButton::Left,
                        button_state: MouseButtonState::Up,
                        ..
                    } = event
                    {
                        let app = tray.app_handle();
                        if let Some(window) = app.get_webview_window(MAIN_WINDOW_LABEL) {
                            restore_webview_window(&window);
                        }
                    }
                })
                .build(app)?;

            app.manage(AppState {
                session: Mutex::new(None),
                tray: Mutex::new(Some(TrayState {
                    status_item,
                    received_item,
                    sent_item,
                    toggle_item,
                })),
                single_instance_lock,
            });
            refresh_tray_menu(app.handle());

            Ok(())
        })
        .on_window_event(|window, event| match event {
            WindowEvent::CloseRequested { .. } => {
                window.app_handle().exit(0);
            }
            WindowEvent::Resized(_) => {
                if window.is_minimized().unwrap_or(false) {
                    hide_main_window(window);
                }
            }
            _ => {}
        })
        .invoke_handler(tauri::generate_handler![load_state, clear_session, enroll_device, sync_profile, connect_tunnel, disconnect_tunnel, tunnel_status, tunnel_metrics])
        .run(tauri::generate_context!())
        .expect("error while running tauri application");
}