AI-Studio/runtime/src/dotnet.rs

use std::collections::HashMap;
use std::path::Path;
use std::sync::{Arc, Mutex};
use base64::Engine;
use base64::prelude::BASE64_STANDARD;
use log::{error, info, warn};
use once_cell::sync::Lazy;
use tauri::{Manager, Url};
use tauri_plugin_shell::process::{CommandChild, CommandEvent, TerminatedPayload};
use tauri_plugin_shell::ShellExt;
use crate::runtime_api_token::API_TOKEN;
use crate::app_window::change_location_to;
use crate::runtime_certificate::CERTIFICATE_FINGERPRINT;
use crate::encryption::ENCRYPTION;
use crate::environment::{
    is_dev, resolve_external_http_custom_root_certificate_policy, DATA_DIRECTORY,
    DOTNET_ENV_CUSTOM_ROOT_CERTIFICATE_ALLOWED_HOSTS,
    DOTNET_ENV_CUSTOM_ROOT_CERTIFICATE_BUNDLE_PATH,
    DOTNET_ENV_CUSTOM_ROOT_CERTIFICATE_POLICY_CONFIGURED,
    DOTNET_ENV_CUSTOM_ROOT_CERTIFICATES_ENABLED,
};
use crate::network::get_available_port;
use crate::runtime_api::API_SERVER_PORT;
use crate::stale_process_cleanup::{kill_stale_process, log_potential_stale_process};
use crate::sidecar_types::SidecarType;

// The .NET server is started in a separate process and communicates with this
// runtime process via IPC. However, we do net start the .NET server in
// the development environment.
static DOTNET_SERVER: Lazy<Arc<Mutex<Option<CommandChild>>>> = Lazy::new(|| Arc::new(Mutex::new(None)));

// The .NET server port is relevant for the production environment only, sine we
// do not start the server in the development environment.
static DOTNET_SERVER_PORT: Lazy<u16> = Lazy::new(|| get_available_port().unwrap());

static DOTNET_INITIALIZED: Lazy<Mutex<bool>> = Lazy::new(|| Mutex::new(false));

pub const PID_FILE_NAME: &str = "mindwork_ai_studio.pid";
const SIDECAR_TYPE:SidecarType = SidecarType::Dotnet;

/// Tells the .NET host where to extract the native libraries bundled into the server.
const DOTNET_ENV_BUNDLE_EXTRACT_BASE_DIR: &str = "DOTNET_BUNDLE_EXTRACT_BASE_DIR";

/// The directory below the app's cache directory where the .NET host extracts them.
const DOTNET_BUNDLE_EXTRACTION_DIRECTORY_NAME: &str = "dotnet";

/// Removes ANSI escape sequences and non-printable control chars from stdout and stderr lines.
fn sanitize_output_line(line: &str) -> String {
    let mut sanitized = String::with_capacity(line.len());
    let mut chars = line.chars().peekable();

    while let Some(ch) = chars.next() {
        if ch == '\u{1B}' {
            if let Some(next) = chars.peek().copied() {
                // CSI sequence: ESC [ ... <final>
                if next == '[' {
                    chars.next();
                    for csi_char in chars.by_ref() {
                        let code = csi_char as u32;
                        if (0x40..=0x7E).contains(&code) {
                            break;
                        }
                    }
                    continue;
                }

                // OSC sequence: ESC ] ... (BEL or ESC \)
                if next == ']' {
                    chars.next();
                    let mut previous_was_escape = false;
                    for osc_char in chars.by_ref() {
                        if osc_char == '\u{07}' {
                            break;
                        }

                        if previous_was_escape && osc_char == '\\' {
                            break;
                        }

                        previous_was_escape = osc_char == '\u{1B}';
                    }
                    continue;
                }
            }

            // Unknown escape sequence: ignore the escape char itself.
            continue;
        }

        if ch.is_control() && ch != '\t' {
            continue;
        }

        sanitized.push(ch);
    }

    sanitized
}

/// Describes how a process ended, for the log.
fn describe_termination(payload: &TerminatedPayload) -> String {
    match (payload.code, payload.signal) {
        (Some(code), _) => format!("exit code {code}"),
        (None, Some(signal)) => format!("signal {signal}"),
        (None, None) => String::from("neither an exit code nor a signal"),
    }
}

/// Returns the desired port of the .NET server. Our .NET app calls this endpoint to get
/// the port where the .NET server should listen to.
pub async fn dotnet_port() -> String {
    let dotnet_server_port = *DOTNET_SERVER_PORT;
    format!("{dotnet_server_port}")
}

fn external_http_custom_root_certificate_policy_environment() -> Vec<(String, String)> {
    let policy = resolve_external_http_custom_root_certificate_policy();
    if !policy.is_configured {
        return Vec::new();
    }

    vec![
        (String::from(DOTNET_ENV_CUSTOM_ROOT_CERTIFICATE_POLICY_CONFIGURED), String::from("true")),
        (String::from(DOTNET_ENV_CUSTOM_ROOT_CERTIFICATES_ENABLED), policy.enabled.to_string()),
        (String::from(DOTNET_ENV_CUSTOM_ROOT_CERTIFICATE_BUNDLE_PATH), policy.bundle_path),
        (String::from(DOTNET_ENV_CUSTOM_ROOT_CERTIFICATE_ALLOWED_HOSTS), policy.allowed_hosts),
    ]
}

/// Returns the directory where the .NET host extracts the native libraries of the server.
///
/// The server is a single-file bundle, and the native libraries inside it -- SQLite above all --
/// have to be written to disk before the .NET host can load them. That happens before any of our
/// own .NET code runs. Left to itself, the host picks `$HOME/.net` or `%TEMP%\.net` and exits at
/// once when that location is not writable. The Flatpak sandbox mounts the home directory
/// read-only, so the server never started there. Our own cache directory is writable wherever
/// the app runs, which is why we choose the location on every platform instead of trusting the
/// default.
///
/// Returns `None` when the directory is not available. The host then falls back to its default,
/// and should that fail as well, its message shows up in our log as stderr of the .NET server.
fn dotnet_bundle_extraction_directory<R: tauri::Runtime>(app_handle: &tauri::AppHandle<R>) -> Option<String> {
    let cache_directory = match app_handle.path().app_cache_dir() {
        Ok(path) => path,
        Err(e) => {
            error!(Source = "Bootloader .NET"; "Failed to resolve the app cache directory for extracting the native libraries of the .NET server: {e}");
            return None;
        }
    };

    let extraction_directory = cache_directory.join(DOTNET_BUNDLE_EXTRACTION_DIRECTORY_NAME);
    if let Err(e) = std::fs::create_dir_all(&extraction_directory) {
        error!(Source = "Bootloader .NET"; "Failed to create the directory '{}' for extracting the native libraries of the .NET server: {e}", extraction_directory.display());
        return None;
    }

    match extraction_directory.to_str() {
        Some(path) => {
            info!(Source = "Bootloader .NET"; "The .NET server extracts its native libraries to '{path}'.");
            Some(path.to_string())
        }
        None => {
            error!(Source = "Bootloader .NET"; "The directory '{}' for extracting the native libraries of the .NET server is not valid UTF-8.", extraction_directory.display());
            None
        }
    }
}

/// Creates the startup environment file for the .NET server in the development
/// environment. The file is created in the root directory of the repository.
/// Creating that env file on a production environment would be a security
/// issue, since it contains the secret password and salt in plain text.
/// Anyone could read that file and decrypt the secret communication
/// between the .NET server and the Tauri app.
///
/// Therefore, we not only create the file in the development environment
/// but also remove that code from any production build.
#[cfg(debug_assertions)]
pub fn create_startup_env_file() {

    // Get the secret password & salt and convert it to a base64 string:
    let secret_password = BASE64_STANDARD.encode(ENCRYPTION.secret_password);
    let secret_key_salt = BASE64_STANDARD.encode(ENCRYPTION.secret_key_salt);
    let api_port = *API_SERVER_PORT;

    warn!(Source = "Bootloader .NET"; "Development environment detected; create the startup env file at '../startup.env'.");
    let env_file_path = std::path::PathBuf::from("..").join("startup.env");
    let mut env_file = std::fs::File::create(env_file_path).unwrap();
    let mut env_file_lines = vec![
        format!("AI_STUDIO_SECRET_PASSWORD={secret_password}"),
        format!("AI_STUDIO_SECRET_KEY_SALT={secret_key_salt}"),
        format!("AI_STUDIO_CERTIFICATE_FINGERPRINT={}", CERTIFICATE_FINGERPRINT.get().unwrap()),
        format!("AI_STUDIO_API_PORT={api_port}"),
        format!("AI_STUDIO_API_TOKEN={}", API_TOKEN.to_hex_text()),
    ];
    for (key, value) in external_http_custom_root_certificate_policy_environment() {
        env_file_lines.push(format!("{key}={value}"));
    }

    std::io::Write::write_all(&mut env_file, env_file_lines.join("\n").as_bytes()).unwrap();
    info!(Source = "Bootloader .NET"; "The startup env file was created successfully.");
}

/// Starts the .NET server in a separate process.
pub fn start_dotnet_server<R: tauri::Runtime>(app_handle: tauri::AppHandle<R>) {

    // Get the secret password & salt and convert it to a base64 string:
    let secret_password = BASE64_STANDARD.encode(ENCRYPTION.secret_password);
    let secret_key_salt = BASE64_STANDARD.encode(ENCRYPTION.secret_key_salt);
    let api_port = *API_SERVER_PORT;

    let mut dotnet_server_environment: HashMap<String, String> = HashMap::from_iter([
        (String::from("AI_STUDIO_SECRET_PASSWORD"), secret_password),
        (String::from("AI_STUDIO_SECRET_KEY_SALT"), secret_key_salt),
        (String::from("AI_STUDIO_CERTIFICATE_FINGERPRINT"), CERTIFICATE_FINGERPRINT.get().unwrap().to_string()),
        (String::from("AI_STUDIO_API_PORT"), format!("{api_port}")),
        (String::from("AI_STUDIO_API_TOKEN"), API_TOKEN.to_hex_text().to_string()),
    ]);
    dotnet_server_environment.extend(external_http_custom_root_certificate_policy_environment());
    if let Some(extraction_directory) = dotnet_bundle_extraction_directory(&app_handle) {
        dotnet_server_environment.insert(String::from(DOTNET_ENV_BUNDLE_EXTRACT_BASE_DIR), extraction_directory);
    }

    info!("Try to start the .NET server...");
    let server_spawn_clone = DOTNET_SERVER.clone();
    tauri::async_runtime::spawn(async move {
        let shell = app_handle.shell();
        let (mut rx, child) = shell
            .sidecar("mindworkAIStudioServer")
            .expect("Failed to create sidecar")
            .envs(dotnet_server_environment)
            .spawn()
            .expect("Failed to spawn .NET server process.");
        let server_pid = child.pid();
        info!(Source = "Bootloader .NET"; "The .NET server process started with PID={server_pid}.");
        log_potential_stale_process(Path::new(DATA_DIRECTORY.get().unwrap()).join(PID_FILE_NAME), server_pid, SIDECAR_TYPE);

        // Save the server process to stop it later:
        *server_spawn_clone.lock().unwrap() = Some(child);

        // Log the output of the .NET server:
        // NOTE: Log events are sent via structured HTTP API calls.
        // This loop serves for fundamental output (e.g., startup errors).
        //
        // The .NET host writes its own failures to stderr, before any of our .NET code runs.
        // Without stderr and the termination, a server that cannot even start leaves no trace.
        while let Some(event) = rx.recv().await {
            match event {
                CommandEvent::Stdout(line) => {
                    let line_utf8 = String::from_utf8_lossy(&line).to_string();
                    let line = sanitize_output_line(line_utf8.trim_end());
                    if !line.trim().is_empty() {
                        info!(Source = ".NET Server (stdout)"; "{line}");
                    }
                }

                CommandEvent::Stderr(line) => {
                    let line_utf8 = String::from_utf8_lossy(&line).to_string();
                    let line = sanitize_output_line(line_utf8.trim_end());
                    if !line.trim().is_empty() {
                        error!(Source = ".NET Server (stderr)"; "{line}");
                    }
                }

                CommandEvent::Error(e) => {
                    error!(Source = "Bootloader .NET"; "Failed to read the output of the .NET server: {e}");
                }

                CommandEvent::Terminated(payload) => {
                    //
                    // stop_dotnet_server() takes the child out before killing it. When the child is
                    // still here, nobody asked the server to stop.
                    //
                    let termination = describe_termination(&payload);
                    if server_spawn_clone.lock().unwrap().is_some() {
                        error!(Source = "Bootloader .NET"; "The .NET server process terminated unexpectedly with {termination}.");
                    } else {
                        info!(Source = "Bootloader .NET"; "The .NET server process terminated with {termination}.");
                    }
                }

                _ => (),
            }
        }
    });
}

/// This endpoint is called by the .NET server to signal that the server is ready.
pub async fn dotnet_ready() {

    // We create a manual scope for the lock to be released as soon as possible.
    // This is necessary because we cannot await any function while the lock is
    // held.
    {
        let mut initialized = DOTNET_INITIALIZED.lock().unwrap();
        if !is_dev() && *initialized {
            error!("Anyone tried to initialize the runtime twice. This is not intended.");
            return;
        }

        info!("The .NET server was booted successfully.");
        *initialized = true;
    }

    let dotnet_server_port = *DOTNET_SERVER_PORT;
    let url = match Url::parse(format!("http://localhost:{dotnet_server_port}").as_str())
    {
        Ok(url) => url,
        Err(msg) => {
            error!("Error while parsing URL for navigating to the app: {msg}");
            return;
        }
    };

    change_location_to(url.as_str()).await;
}

/// Stops the .NET server process.
pub fn stop_dotnet_server() {
    if let Some(server_process) = DOTNET_SERVER.lock().unwrap().take() {
        let server_kill_result = server_process.kill();
        match server_kill_result {
            Ok(_) => info!("The .NET server process was stopped."),
            Err(e) => error!("Failed to stop the .NET server process: {e}."),
        }
    } else {
        warn!("The .NET server process was not started or is already stopped.");
    }
    info!("Start dotnet server cleanup");
    cleanup_dotnet_server();
}

/// Remove old Pid files and kill the corresponding processes
pub fn cleanup_dotnet_server() {
    let pid_path = Path::new(DATA_DIRECTORY.get().unwrap()).join(PID_FILE_NAME);
    if let Err(e) = kill_stale_process(pid_path, SIDECAR_TYPE) {
        warn!(Source = ".NET"; "Error during the cleanup of .NET: {}", e);
    }
}