hello-agent/src/main.rs

// hello-agent: a headless RustDesk-protocol-compatible support agent.
//
// One binary, two run modes: a console / installer entry point that handles
// --install / --uninstall / --config, and a "--service" entry registered with
// the Windows SCM that spawns the actual worker into the active console
// session as "--server".
//
// The protocol stack (rendezvous, NAT punch, screen capture, input, login
// flow) is reused unchanged from `librustdesk`. This crate is just the
// thin shell that gives us a different CLI surface, our own service install
// path, and a native approval popup in place of the Flutter CM.
//
// We override `hbb_common`'s default `APP_NAME` ("RustDesk") with our own
// product name as the very first thing every process does. APP_NAME is read
// lazily from a `RwLock<String>` whenever any path is computed (config dir,
// log dir, named-pipe namespace, …), so setting it before any of those
// initializers fire is enough to redirect all hbb_common state under
// `%APPDATA%\hello-agent\` and the matching LocalService path. Identical
// to the `read_custom_client` write path the upstream Flutter build uses
// for OEM rebrands.

#![cfg_attr(not(target_os = "windows"), allow(dead_code, unused_imports))]

mod cli;
mod config_import;
mod inventory;

#[cfg(target_os = "windows")]
mod cm_popup;
#[cfg(target_os = "windows")]
mod exec;
mod login_events;
mod perf;
mod perf_events;
#[cfg(target_os = "windows")]
mod service;
#[cfg(target_os = "windows")]
mod unattended_password;
#[cfg(target_os = "windows")]
mod wifi_native;

use cli::{Action, ParsedArgs};

/// Product name used to namespace all on-disk state and the IPC pipe path.
/// Written into `hbb_common::config::APP_NAME` at the top of `main` so
/// every subsequent path computation (config dir, log dir, named pipe)
/// targets `%APPDATA%\hello-agent\` rather than the upstream default of
/// `%APPDATA%\RustDesk\`. Must be set before any code touches a path —
/// `hbb_common` initializes path globals lazily on first read.
///
/// Important: this value also drives upstream's installer lookup paths.
/// `librustdesk::platform::get_install_info` computes the expected install
/// dir as `%ProgramFiles%\<APP_NAME>` and the expected exe filename as
/// `<APP_NAME>.exe`. Keeping `APP_NAME` aligned with the lowercase-hyphenated
/// install path (`%ProgramFiles%\hello-agent\hello-agent.exe`) is what
/// makes `--update` (which delegates to `librustdesk::platform::update_me`)
/// find the binary it needs to replace, kill the right process by image
/// name, and rename the staged exe to `hello-agent.exe` after the copy.
/// Renaming this constant without renaming the install dir / exe will
/// silently break self-update.
pub const APP_NAME: &str = "hello-agent";

/// Set up logging. We delegate to `hbb_common::init_log`, which:
///   * In **debug** builds: installs `env_logger` writing to stderr.
///   * In **release** builds: installs `flexi_logger` writing to a rolling
///     file under `<config_dir>/log/<mode>/` — the SYSTEM service log ends
///     up at `%SystemRoot%\ServiceProfiles\LocalService\AppData\Roaming\hello-agent\log\<mode>\`
///     and the user-mode log at `%APPDATA%\hello-agent\log\<mode>\`.
///
/// The `mode` label segregates per-run-mode log files so service worker
/// chatter doesn't tangle with --install diagnostics. `init_log` is
/// `Once`-guarded internally so calling it twice is harmless.
fn init_logging(mode: &str) {
    let _ = hbb_common::init_log(false, mode);
}

fn main() {
    // MUST be the very first line. See the doc-comment on `APP_NAME` —
    // anything that lazily reads a config / log / pipe path before this
    // runs would cache `"RustDesk"` in `hbb_common`'s path globals and
    // we'd never recover.
    *hbb_common::config::APP_NAME.write().unwrap() = APP_NAME.to_owned();
    // Identify ourselves to the rustdesk-server's /api/sysinfo endpoint
    // so the admin Devices page can show "hello-agent 0.1.0" instead of
    // the embedded rustdesk core version. These RwLocks are read once
    // per sysinfo upload by hbbs_http::sync; setting them here (before
    // start_server) ensures the very first upload carries the identity.
    *hbb_common::config::AGENT_NAME.write().unwrap() = APP_NAME.to_owned();
    *hbb_common::config::AGENT_VERSION.write().unwrap() = env!("CARGO_PKG_VERSION").to_owned();

    let parsed = match ParsedArgs::from_argv(std::env::args().skip(1)) {
        Ok(p) => p,
        Err(e) => {
            eprintln!("hello-agent: {e}");
            eprintln!();
            cli::print_usage();
            std::process::exit(2);
        }
    };

    // Initialize logging *after* arg parsing so the per-mode log file path
    // is deterministic. `init_log` is Once-guarded internally.
    let mode = match parsed.action {
        Action::Install => "install",
        Action::Uninstall => "uninstall",
        Action::Service => "service",
        Action::Server => "server",
        Action::Cm => "cm",
        Action::Update => "update",
        Action::ConfigOnly | Action::None => "hello-agent",
    };
    init_logging(mode);

    // --update is the self-replacement re-entry: the running service's
    // updater downloads a new hello-agent.exe to %TEMP%, verifies its
    // SHA256, then launches `<temp>\hello-agent.exe --update` as an
    // elevated child. We are that child — `current_exe()` is the staged
    // new binary, and our only job is to copy ourselves over the
    // installed location and restart the service. Do it before the
    // config-import dance below so a corrupt-on-disk config can't block
    // an update from going through.
    if parsed.action == Action::Update {
        #[cfg(target_os = "windows")]
        {
            match librustdesk::platform::update_me(false) {
                Ok(()) => {
                    log::info!("hello-agent: --update completed");
                }
                Err(e) => {
                    log::error!("hello-agent: --update failed: {e:#}");
                    std::process::exit(1);
                }
            }
        }
        #[cfg(not(target_os = "windows"))]
        {
            eprintln!("hello-agent: --update is Windows-only.");
            std::process::exit(1);
        }
        return;
    }

    // --config is allowed to combine with --install (one-line MDM deploy)
    // but on its own is a separate operation. Apply it first so --install
    // sees the populated config.
    if let Some(blob) = parsed.config_blob.as_deref() {
        if let Err(e) = config_import::apply(blob) {
            eprintln!("hello-agent: --config failed: {e:#}");
            std::process::exit(2);
        }
    }

    // Bake in fallback rendezvous defaults. Idempotent — if --config above
    // (or a prior install) already set custom-rendezvous-server, this is a
    // no-op. Without this, a bare `hello-agent.exe --install` would land
    // at an unconfigured agent that can't reach any server.
    //
    // Skipped for `--uninstall`: an uninstall flow has no business mutating
    // the calling user's config, and otherwise we'd write defaults into
    // %APPDATA% right before tearing the agent down. (`--update` is
    // dispatched in the early-return block above and never reaches here.)
    if parsed.action != Action::Uninstall {
        config_import::apply_defaults_if_empty();
    }

    match parsed.action {
        Action::Install => {
            #[cfg(target_os = "windows")]
            {
                if let Err(e) = service::install() {
                    eprintln!("hello-agent: install failed: {e:#}");
                    std::process::exit(1);
                }
                println!("hello-agent: installed and started.");
            }
            #[cfg(not(target_os = "windows"))]
            {
                eprintln!("hello-agent: --install is Windows-only for now.");
                std::process::exit(1);
            }
        }
        Action::Uninstall => {
            #[cfg(target_os = "windows")]
            {
                if let Err(e) = service::uninstall() {
                    eprintln!("hello-agent: uninstall failed: {e:#}");
                    std::process::exit(1);
                }
                println!("hello-agent: uninstalled.");
            }
            #[cfg(not(target_os = "windows"))]
            {
                eprintln!("hello-agent: --uninstall is Windows-only for now.");
                std::process::exit(1);
            }
        }
        Action::Service => {
            #[cfg(target_os = "windows")]
            {
                if let Err(e) = service::run_as_service() {
                    eprintln!("hello-agent: service dispatcher failed: {e:#}");
                    std::process::exit(1);
                }
            }
            #[cfg(not(target_os = "windows"))]
            {
                eprintln!("hello-agent: --service is Windows-only.");
                std::process::exit(1);
            }
        }
        Action::Server => run_server(),
        Action::Cm => {
            // Spawned by the SYSTEM-token --server worker (via librustdesk's
            // run_as_user) when the rustdesk core wants a CM. Runs as the
            // logged-in user, binds the `_cm` IPC pipe, services one Login
            // request with a MessageBoxW, replies, exits.
            #[cfg(target_os = "windows")]
            cm_popup::run_blocking();
        }
        Action::ConfigOnly => {
            // --config without --install or --service: just persist and exit.
        }
        Action::None => {
            // No flags: dev mode. Run as a foreground server so the operator
            // can watch logs. Production deployments use --install + --service.
            run_server();
        }
        Action::Update => {
            // Handled in the early-return block above (before config-import).
            // The match has to cover this variant for exhaustiveness.
            unreachable!("Action::Update is dispatched before this match");
        }
    }
}

fn run_server() {
    // Clear any stale `approve-mode = click` left by older hello-agent
    // versions. ApproveMode comes from `password_security::approve_mode`:
    // "password" → password only, "click" → popup only, anything else →
    // both (try password first, fall back to popup). We want both so
    // that (a) attended sessions still go through the cm_popup approval,
    // and (b) unattended sessions can authenticate with the per-boot
    // password we report to the admin UI. Setting to "" is idempotent
    // and overrides any leftover "click" value on disk.
    hbb_common::config::Config::set_option("approve-mode".into(), "".into());

    // Pre-spawn the --cm child *on the user's interactive desktop* before
    // start_server boots. librustdesk's start_ipc has its own
    // run_as_user(["--cm"]) fallback, but it goes through C-side
    // LaunchProcessWin with show=FALSE → lpDesktop=NULL → child inherits
    // the parent's desktop, which (because we were spawned by the Session-0
    // service) is the invisible Session 0 service desktop. Our spawn
    // helper sets lpDesktop = winsta0\\default explicitly, putting the
    // popup on the user's screen. Once our --cm is bound to `_cm`,
    // start_ipc's first ipc::connect("_cm") succeeds and rustdesk's
    // built-in fallback never fires.
    //
    // We target *our own* session (whichever the supervisor placed us in
    // — physical console, RDP, multi-user) rather than the physical
    // console specifically. WTSGetActiveConsoleSessionId would point at
    // the empty / lock-screen console session in RDP-only scenarios.
    #[cfg(target_os = "windows")]
    match service::spawn_cm_in_my_session() {
        Ok(pid) => log::info!("spawned --cm child pid={pid} on winsta0\\default"),
        Err(e) => log::warn!(
            "could not pre-spawn --cm child ({e:#}); rustdesk's start_ipc fallback may be invisible"
        ),
    }

    // Kick off CMDB inventory collection on a background thread before
    // start_server boots. PowerShell's first-run cost (a few hundred ms
    // to a few seconds) shouldn't delay the rendezvous heartbeat — the
    // sysinfo upload loop only fires every TIME_CONN seconds, so the
    // inventory will be ready in time for the very first /api/sysinfo
    // POST whether collection finishes in 50ms or 5s. We deliberately
    // don't retry on failure: a transient PowerShell hiccup leaves the
    // INVENTORY global empty, and sync.rs simply omits the `inventory`
    // key from the upload. Next agent restart re-tries.
    std::thread::spawn(|| {
        let inv = inventory::collect_inventory();
        if !inv.is_empty() {
            *hbb_common::config::INVENTORY.write().unwrap() = inv;
        }
    });

    // Start the PowerShell remote-exec worker. Subscribes to the
    // broadcast channel in the vendored sync layer; the channel is
    // shared in-process so the worker MUST run in this --server process
    // (where sync.rs lives), not the --service supervisor. The worker
    // is idle until an admin dispatches an exec from the dashboard.
    // Gated server-side on peer.managed=1 + strategy.enable-remote-exec.
    #[cfg(target_os = "windows")]
    exec::start();

    // `start_server` is `#[tokio::main]` and runs forever. (is_server=true,
    // no_server=false). It boots the default IPC server, input service,
    // rendezvous mediator, and heartbeat sync.
    librustdesk::start_server(true, false);
}