maestro/crates/maestro-e2e/src/daemon.rs

//! Subprocess wrapper around `maestrod`.

use std::net::{Ipv4Addr, SocketAddr};
use std::path::PathBuf;
use std::process::{Child, Command, Stdio};
use std::time::Duration;

use tempfile::TempDir;

use crate::build::workspace_target_bin;
use crate::util::{free_port, wait_for_async};

/// A live `maestrod` running in a child process. Kill the child on drop.
///
/// `state_dir`, `home`, and `tmux_tmpdir` are kept as `TempDir`s so they
/// outlive the process and get cleaned up automatically.
#[derive(Debug)]
pub struct Daemon {
    pub daemon_id: String,
    pub bind_addr: SocketAddr,
    pub daemon_url: String,
    child: Option<Child>,
    home: TempDir,
    state: TempDir,
    tmux: TempDir,
    registry_url: String,
    token: String,
}

/// Optional knobs for `Daemon::spawn_with` — perimeter + allow_inbound for
/// NAP-247 broker tests.
#[derive(Debug, Default, Clone)]
pub struct DaemonOptions {
    pub perimeter_id: Option<String>,
    pub private_addr: Option<String>,
    pub allow_inbound: Vec<String>,
}

impl Daemon {
    /// Spawn a fresh daemon pointing at `registry_base_url` (the registry's
    /// HTTP base, e.g. `http://127.0.0.1:1234`). The daemon derives its WS
    /// endpoint internally via `derive_daemon_ws_url`.
    pub async fn spawn(daemon_id: &str, registry_base_url: &str, token: &str) -> Self {
        Self::spawn_with(
            daemon_id,
            registry_base_url,
            token,
            DaemonOptions::default(),
        )
        .await
    }

    /// Spawn with NAP-247 options (perimeter advertisement, allow_inbound).
    pub async fn spawn_with(
        daemon_id: &str,
        registry_base_url: &str,
        token: &str,
        opts: DaemonOptions,
    ) -> Self {
        let port = free_port();
        let bind = SocketAddr::from((Ipv4Addr::LOCALHOST, port));
        Self::spawn_with_port_and_opts(daemon_id, registry_base_url, token, bind, opts).await
    }

    /// Spawn on a caller-supplied bind address. Used by NAP-247 broker tests
    /// where the daemon's advertised `private_addr` must equal the actual
    /// bind address so the broker can dial it back.
    pub async fn spawn_at(
        daemon_id: &str,
        registry_base_url: &str,
        token: &str,
        bind: SocketAddr,
        opts: DaemonOptions,
    ) -> Self {
        Self::spawn_with_port_and_opts(daemon_id, registry_base_url, token, bind, opts).await
    }

    /// Stop the underlying process, but keep its state dirs around so we can
    /// `restart()` against the same `$HOME` / `state` / `TMUX_TMPDIR`.
    pub fn stop(&mut self) {
        if let Some(mut c) = self.child.take() {
            let _ = c.kill();
            let _ = c.wait();
        }
    }

    /// Re-spawn the daemon on the same port + state. The previous tmux
    /// server (under `TMUX_TMPDIR`) is left running so `recover()` can adopt
    /// its sessions. Restarts inherit the original options via the on-disk
    /// config file written at `spawn_with_port_and_opts` time.
    pub async fn restart(&mut self) {
        self.stop();
        let child = launch(
            &self.daemon_id,
            self.bind_addr,
            &self.daemon_url,
            &self.registry_url,
            &self.token,
            self.home.path(),
            self.state.path(),
            self.tmux.path(),
            &DaemonOptions::default(),
        );
        self.child = Some(child);
        self.wait_until_attach_ready().await;
    }

    async fn spawn_with_port_and_opts(
        daemon_id: &str,
        registry_url: &str,
        token: &str,
        bind: SocketAddr,
        opts: DaemonOptions,
    ) -> Self {
        let home = TempDir::new().expect("home tmpdir");
        let state = TempDir::new().expect("state tmpdir");
        let tmux = TempDir::new().expect("tmux tmpdir");

        // Write a minimal config.toml so allow_inbound (and any future
        // file-only fields) reach the daemon. Required fields keep coming
        // through env vars set by `launch`.
        write_config(home.path(), &opts);

        let daemon_url = format!("ws://{bind}");
        let child = launch(
            daemon_id,
            bind,
            &daemon_url,
            registry_url,
            token,
            home.path(),
            state.path(),
            tmux.path(),
            &opts,
        );
        let me = Self {
            daemon_id: daemon_id.to_string(),
            bind_addr: bind,
            daemon_url,
            child: Some(child),
            home,
            state,
            tmux,
            registry_url: registry_url.to_string(),
            token: token.to_string(),
        };
        me.wait_until_attach_ready().await;
        me
    }

    async fn wait_until_attach_ready(&self) {
        let addr = self.bind_addr;
        wait_for_async(
            "daemon attach port",
            Duration::from_secs(5),
            || async move { tokio::net::TcpStream::connect(addr).await.is_ok() },
        )
        .await;
    }

    pub fn home(&self) -> &std::path::Path {
        self.home.path()
    }

    pub fn state_dir(&self) -> PathBuf {
        // Matches maestro_daemon::default_state_dir which uses XDG_STATE_HOME.
        self.home.path().join(".local/state/maestrod")
    }
}

impl Drop for Daemon {
    fn drop(&mut self) {
        if let Some(mut c) = self.child.take() {
            let _ = c.kill();
            let _ = c.wait();
        }
    }
}

#[allow(clippy::too_many_arguments)]
fn launch(
    daemon_id: &str,
    bind: SocketAddr,
    daemon_url: &str,
    registry_url: &str,
    token: &str,
    home: &std::path::Path,
    _state: &std::path::Path,
    tmux_tmpdir: &std::path::Path,
    opts: &DaemonOptions,
) -> Child {
    let bin = workspace_target_bin("maestro", "maestro");
    let mut cmd = Command::new(&bin);
    cmd.arg("daemon");
    cmd.env_clear()
        .env(
            "PATH",
            std::env::var("PATH").unwrap_or_else(|_| "/usr/bin:/bin".into()),
        )
        .env("HOME", home)
        .env("XDG_CONFIG_HOME", home.join(".config"))
        .env("XDG_STATE_HOME", home.join(".local/state"))
        .env("TMUX_TMPDIR", tmux_tmpdir)
        .env("MAESTRO_TOKEN", token)
        .env("MAESTRO_DAEMON_ID", daemon_id)
        .env("MAESTRO_REGISTRY_URL", registry_url)
        .env("MAESTRO_DAEMON_BIND_ADDR", bind.to_string())
        .env("MAESTRO_DAEMON_URL", daemon_url)
        .env("RUST_LOG", "warn")
        .stdout(Stdio::null())
        .stderr(Stdio::inherit())
        .arg("run");
    if let Some(p) = &opts.perimeter_id {
        cmd.env("MAESTRO_DAEMON_PERIMETER_ID", p);
    }
    if let Some(a) = &opts.private_addr {
        cmd.env("MAESTRO_DAEMON_PRIVATE_ADDR", a);
    }
    cmd.spawn()
        .unwrap_or_else(|e| panic!("spawn {} failed: {e}", bin.display()))
}

/// Write a minimal `~/.config/maestro/config.toml` carrying the NAP-247
/// `[daemon.policy] allow_inbound` list. Everything else still comes from
/// env vars set by `launch`.
fn write_config(home: &std::path::Path, opts: &DaemonOptions) {
    if opts.allow_inbound.is_empty() {
        return;
    }
    let dir = home.join(".config/maestro");
    std::fs::create_dir_all(&dir).unwrap();
    let mut allow = String::new();
    for (i, id) in opts.allow_inbound.iter().enumerate() {
        if i > 0 {
            allow.push_str(", ");
        }
        allow.push('"');
        allow.push_str(id);
        allow.push('"');
    }
    let body = format!("[daemon.policy]\nallow_inbound = [{allow}]\n");
    std::fs::write(dir.join("config.toml"), body).unwrap();
}