ados_sdk/

lifecycle.rs

//! Plugin lifecycle trait and the runner entry.
//!
//! Ports `ados.plugins.runner` for the `runtime: rust` case. A Rust plugin's
//! binary implements [`Plugin`] and calls [`run_plugin`] from `main`; the
//! runner reads `--socket` / `--token` / `--agent-id`, falling back to the
//! `ADOS_PLUGIN_*` values of the token credential systemd loads into
//! `$CREDENTIALS_DIRECTORY` (or, without one, the environment — the exact
//! contract `runner.py` follows), connects the
//! [`PluginIpcClient`], builds a [`PluginContext`], and drives the lifecycle
//! hooks until SIGTERM/SIGINT, then runs the teardown hooks. A lost host
//! session (a plugin-host restart) also ends the run, as an error, so the
//! binary exits non-zero and its unit restarts it with a fresh session.
//!
//! **Entry: function, not proc-macro.** A `#[ados_plugin]` attribute macro
//! would expand to nothing more than `fn main() { run_plugin::<P>() }` — it
//! saves one line and costs a separate proc-macro crate, a syn/quote dependency
//! tree, and an opaque codegen step. The plain generic function is the cleaner
//! contract: a plugin author writes a three-line `main`, the call site is
//! explicit, and there is no macro to debug. So the SDK ships [`run_plugin`]
//! and no proc-macro sub-crate.
//!
//! Hook order matches the Python runner exactly:
//! `on_install` -> `on_enable` -> `on_configure` -> `on_start` -> (wait for
//! shutdown) -> `on_stop` -> `on_disable`. Every hook has a default no-op so a
//! plugin overrides only what it needs.

use std::collections::BTreeMap;
use std::future::Future;
use std::sync::Arc;

use async_trait::async_trait;
use rmpv::Value;
use thiserror::Error;

use crate::client::{ClientError, PluginIpcClient};
use crate::context::PluginContext;

/// The lifecycle hook trait a Rust plugin implements. Mirrors the duck-typed
/// `on_install` / `on_enable` / `on_configure` / `on_start` / `on_stop` /
/// `on_disable` methods the Python runner calls via `getattr`.
///
/// Every hook defaults to a no-op. `on_configure` receives the resolved config
/// map (the Python runner passes `{}` today; the SDK passes the same shape and
/// will carry the live config when the host wires it).
#[async_trait]
pub trait Plugin: Send + Sync {
    /// Construct the plugin. Called once before any hook. The runner uses this
    /// rather than `Default` so a plugin can fail construction explicitly.
    fn new() -> Self
    where
        Self: Sized;

    async fn on_install(&mut self, _ctx: &PluginContext) -> Result<(), ClientError> {
        Ok(())
    }
    async fn on_enable(&mut self, _ctx: &PluginContext) -> Result<(), ClientError> {
        Ok(())
    }
    async fn on_configure(
        &mut self,
        _ctx: &PluginContext,
        _config: &BTreeMap<String, Value>,
    ) -> Result<(), ClientError> {
        Ok(())
    }
    async fn on_start(&mut self, _ctx: &PluginContext) -> Result<(), ClientError> {
        Ok(())
    }
    async fn on_stop(&mut self, _ctx: &PluginContext) -> Result<(), ClientError> {
        Ok(())
    }
    async fn on_disable(&mut self, _ctx: &PluginContext) -> Result<(), ClientError> {
        Ok(())
    }
}

/// Errors raised while bootstrapping or running a plugin via [`run_plugin`].
#[derive(Debug, Error)]
pub enum RunnerError {
    /// The plugin id positional argument was missing.
    #[error("missing plugin id argument")]
    MissingPluginId,
    /// `--socket` / `ADOS_PLUGIN_SOCKET` or `--token` / `ADOS_PLUGIN_TOKEN` was
    /// absent, so there is no host to connect to. Mirrors the Python runner's
    /// fall-back-to-null-ipc path; here it is an explicit error because a Rust
    /// `runtime: rust` plugin is always launched with the bridge wired.
    #[error("no supervisor socket/token supplied (set --socket and --token)")]
    NoBridge,
    /// The client failed to connect or handshake.
    #[error("ipc connect failed: {0}")]
    Connect(#[from] ClientError),
    /// The plugin host closed this plugin's session (a plugin-host restart) or
    /// the stream broke. Nothing reconnects it: the restarted host re-mints the
    /// socket and token for a fresh process, so the runner stops and the
    /// binary exits non-zero for its unit's `Restart=` to bring it back.
    #[error(
        "plugin host session lost; exiting so the unit restarts with a fresh socket and token"
    )]
    SessionLost,
}

/// The systemd credential id a plugin unit loads its token file under. It
/// appears as `$CREDENTIALS_DIRECTORY/<this>` inside the plugin process.
pub const TOKEN_CREDENTIAL_NAME: &str = "ados-plugin-token";

/// The parsed runner arguments. Mirrors the `plugin_id` positional plus the
/// `--socket` / `--token` / `--agent-id` options the Python runner reads, with
/// the `ADOS_PLUGIN_SOCKET` / `ADOS_PLUGIN_TOKEN` / `ADOS_PLUGIN_AGENT_ID`
/// fallbacks read from the token credential.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RunnerArgs {
    pub plugin_id: String,
    pub socket_path: Option<String>,
    pub token: Option<String>,
    pub agent_id: String,
    /// The plugin's per-drone data directory, from `ADOS_PLUGIN_DATA_DIR`. The
    /// host delivers it in the token credential; a plugin reads it through
    /// `ctx.data_dir` rather than re-deriving a path it might get wrong.
    pub data_dir: Option<String>,
}

impl RunnerArgs {
    /// Parse from a raw argv (excluding `argv[0]`) and an env lookup. The env
    /// lookup is injected so the parse is unit-testable without touching the
    /// process environment.
    ///
    /// A value an option does not give comes from the `KEY=VALUE` lines of
    /// `$CREDENTIALS_DIRECTORY/ados-plugin-token`, the file a systemd unit
    /// loads with `LoadCredential=` (the token never sits in the process
    /// environment there). Without a credential the environment supplies it,
    /// which is how the launchd entry script delivers the same keys.
    pub fn parse<F>(args: &[String], env: F) -> Result<Self, RunnerError>
    where
        F: Fn(&str) -> Option<String>,
    {
        let mut plugin_id: Option<String> = None;
        let mut socket_path: Option<String> = None;
        let mut token: Option<String> = None;
        let mut agent_id: Option<String> = None;

        let mut i = 0;
        while i < args.len() {
            let arg = &args[i];
            match arg.as_str() {
                "--socket" => {
                    socket_path = args.get(i + 1).cloned();
                    i += 2;
                }
                "--token" => {
                    token = args.get(i + 1).cloned();
                    i += 2;
                }
                "--agent-id" => {
                    agent_id = args.get(i + 1).cloned();
                    i += 2;
                }
                other if other.starts_with("--socket=") => {
                    socket_path = Some(other["--socket=".len()..].to_string());
                    i += 1;
                }
                other if other.starts_with("--token=") => {
                    token = Some(other["--token=".len()..].to_string());
                    i += 1;
                }
                other if other.starts_with("--agent-id=") => {
                    agent_id = Some(other["--agent-id=".len()..].to_string());
                    i += 1;
                }
                _ if plugin_id.is_none() => {
                    plugin_id = Some(arg.clone());
                    i += 1;
                }
                _ => {
                    i += 1;
                }
            }
        }

        let credential = read_credential(&env);
        let launch_value = |key: &str| {
            credential
                .get(key)
                .filter(|v| !v.is_empty())
                .cloned()
                .or_else(|| env(key).filter(|v| !v.is_empty()))
        };
        Ok(Self {
            plugin_id: plugin_id.ok_or(RunnerError::MissingPluginId)?,
            socket_path: socket_path.or_else(|| launch_value("ADOS_PLUGIN_SOCKET")),
            token: token.or_else(|| launch_value("ADOS_PLUGIN_TOKEN")),
            agent_id: agent_id
                .or_else(|| launch_value("ADOS_PLUGIN_AGENT_ID"))
                .unwrap_or_default(),
            data_dir: launch_value("ADOS_PLUGIN_DATA_DIR"),
        })
    }
}

/// The `KEY=VALUE` lines of the token credential under the
/// `CREDENTIALS_DIRECTORY` the lookup names; empty when there is none or it
/// cannot be read. Values are taken verbatim after the first `=`.
fn read_credential<F>(env: &F) -> BTreeMap<String, String>
where
    F: Fn(&str) -> Option<String>,
{
    let Some(dir) = env("CREDENTIALS_DIRECTORY").filter(|d| !d.is_empty()) else {
        return BTreeMap::new();
    };
    let path = std::path::Path::new(&dir).join(TOKEN_CREDENTIAL_NAME);
    let Ok(body) = std::fs::read_to_string(path) else {
        return BTreeMap::new();
    };
    body.lines()
        .filter_map(|line| line.split_once('='))
        .filter(|(key, _)| !key.is_empty())
        .map(|(key, value)| (key.to_string(), value.to_string()))
        .collect()
}

/// Run a plugin to completion: parse argv/env, connect, build the context,
/// drive the hooks, and tear down on the shutdown signal. This is the
/// `runtime: rust` plugin's `main` body.
///
/// `plugin_version` is the version the host installed; a real binary embeds its
/// own `env!("CARGO_PKG_VERSION")`. `static_config` is the manifest-supplied
/// config the runner would have read off disk; the SDK takes it as an argument
/// so the binary stays in control of where it loads config from.
///
/// Returns when the shutdown future resolves (SIGTERM/SIGINT in a real binary)
/// or a hook errors. Returns [`RunnerError::SessionLost`] when the host session
/// drops while the plugin runs; the binary MUST exit non-zero on any error so
/// systemd restarts it with a fresh session.
pub async fn run_plugin<P, S>(
    plugin_version: impl Into<String>,
    static_config: BTreeMap<String, Value>,
    shutdown: S,
) -> Result<(), RunnerError>
where
    P: Plugin,
    S: Future<Output = ()>,
{
    let argv: Vec<String> = std::env::args().skip(1).collect();
    let args = RunnerArgs::parse(&argv, |k| std::env::var(k).ok())?;
    run_plugin_with::<P, S>(args, plugin_version, static_config, shutdown).await
}

/// Run a plugin against pre-parsed [`RunnerArgs`]. Splits out from
/// [`run_plugin`] so a host harness or a test can supply args directly without
/// touching the process environment.
pub async fn run_plugin_with<P, S>(
    args: RunnerArgs,
    plugin_version: impl Into<String>,
    static_config: BTreeMap<String, Value>,
    shutdown: S,
) -> Result<(), RunnerError>
where
    P: Plugin,
    S: Future<Output = ()>,
{
    let (Some(socket_path), Some(token)) = (args.socket_path.clone(), args.token.clone()) else {
        return Err(RunnerError::NoBridge);
    };

    let ipc = Arc::new(PluginIpcClient::new(
        args.plugin_id.clone(),
        token,
        socket_path,
    ));
    ipc.connect().await?;

    ensure_data_dir(args.data_dir.as_deref());

    let ctx = PluginContext::new(
        ipc.clone(),
        plugin_version,
        args.agent_id,
        args.data_dir,
        static_config.clone(),
    );

    let mut plugin = P::new();
    let result = drive(
        &mut plugin,
        &ctx,
        &static_config,
        shutdown,
        ipc.session_lost(),
    )
    .await;

    // Always close the client, success or failure.
    ipc.close().await;
    result
}

/// Create the plugin's data directory so it exists before the plugin writes to
/// it. The host creates the plugin's data dir (`.../plugin-data/<id>`, owned by
/// the plugin's account and bound writable into its unit) and the per-drone
/// leaf it points `ctx.data_dir` at when that dir is new; a leaf for a device
/// id that appeared later, or a run outside the host, is made here, inside the
/// plugin's own dir. Best-effort: a failure is logged and left to surface as
/// the plugin's own write error rather than aborting the runner. `None` is a
/// no-op.
fn ensure_data_dir(data_dir: Option<&str>) {
    if let Some(dir) = data_dir {
        if let Err(e) = std::fs::create_dir_all(dir) {
            tracing::warn!(dir, error = %e, "could not create plugin data dir");
        }
    }
}

/// Drive the hook sequence. Separated so the teardown hooks run even when a
/// startup hook errors. Mirrors the Python runner's try/finally shape.
///
/// The run phase ends on the shutdown signal or on the loss of the host
/// session, whichever comes first. A lost session still runs the teardown
/// hooks so the plugin releases its local resources, but their errors are
/// only logged (their host calls cannot succeed any more) and the result is
/// [`RunnerError::SessionLost`].
async fn drive<P, S, L>(
    plugin: &mut P,
    ctx: &PluginContext,
    config: &BTreeMap<String, Value>,
    shutdown: S,
    session_lost: L,
) -> Result<(), RunnerError>
where
    P: Plugin,
    S: Future<Output = ()>,
    L: Future<Output = ()>,
{
    plugin.on_install(ctx).await?;
    plugin.on_enable(ctx).await?;
    plugin.on_configure(ctx, config).await?;
    plugin.on_start(ctx).await?;

    tracing::info!(plugin_id = %ctx.plugin_id, "plugin ready");
    let lost = tokio::select! {
        () = shutdown => false,
        () = session_lost => true,
    };

    if lost {
        tracing::error!(
            plugin_id = %ctx.plugin_id,
            "plugin host session lost; stopping so the unit restarts with a fresh session"
        );
        for (hook, result) in [
            ("on_stop", plugin.on_stop(ctx).await),
            ("on_disable", plugin.on_disable(ctx).await),
        ] {
            if let Err(e) = result {
                tracing::debug!(plugin_id = %ctx.plugin_id, hook, error = %e, "teardown hook failed after session loss");
            }
        }
        return Err(RunnerError::SessionLost);
    }

    plugin.on_stop(ctx).await?;
    plugin.on_disable(ctx).await?;
    tracing::info!(plugin_id = %ctx.plugin_id, "plugin clean exit");
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn no_env(_: &str) -> Option<String> {
        None
    }

    #[test]
    fn parses_positional_id_and_space_separated_options() {
        let argv = vec![
            "com.example.demo".to_string(),
            "--socket".to_string(),
            "/run/ados/plugins/x.sock".to_string(),
            "--token".to_string(),
            "v1|p|s|0|600||sig".to_string(),
            "--agent-id".to_string(),
            "agent-42".to_string(),
        ];
        let args = RunnerArgs::parse(&argv, no_env).unwrap();
        assert_eq!(args.plugin_id, "com.example.demo");
        assert_eq!(
            args.socket_path.as_deref(),
            Some("/run/ados/plugins/x.sock")
        );
        assert_eq!(args.token.as_deref(), Some("v1|p|s|0|600||sig"));
        assert_eq!(args.agent_id, "agent-42");
    }

    #[test]
    fn parses_equals_form_options() {
        let argv = vec![
            "com.example.demo".to_string(),
            "--socket=/tmp/y.sock".to_string(),
            "--token=tok".to_string(),
        ];
        let args = RunnerArgs::parse(&argv, no_env).unwrap();
        assert_eq!(args.socket_path.as_deref(), Some("/tmp/y.sock"));
        assert_eq!(args.token.as_deref(), Some("tok"));
        // agent-id defaults to empty, matching the Python `default=""`.
        assert_eq!(args.agent_id, "");
    }

    #[test]
    fn falls_back_to_env_for_socket_token_agent() {
        let argv = vec!["com.example.demo".to_string()];
        let env = |k: &str| match k {
            "ADOS_PLUGIN_SOCKET" => Some("/env/sock".to_string()),
            "ADOS_PLUGIN_TOKEN" => Some("env-tok".to_string()),
            "ADOS_PLUGIN_AGENT_ID" => Some("env-agent".to_string()),
            _ => None,
        };
        let args = RunnerArgs::parse(&argv, env).unwrap();
        assert_eq!(args.socket_path.as_deref(), Some("/env/sock"));
        assert_eq!(args.token.as_deref(), Some("env-tok"));
        assert_eq!(args.agent_id, "env-agent");
    }

    #[test]
    fn explicit_option_overrides_env() {
        let argv = vec![
            "com.example.demo".to_string(),
            "--socket".to_string(),
            "/explicit/sock".to_string(),
        ];
        let env = |k: &str| match k {
            "ADOS_PLUGIN_SOCKET" => Some("/env/sock".to_string()),
            _ => None,
        };
        let args = RunnerArgs::parse(&argv, env).unwrap();
        assert_eq!(args.socket_path.as_deref(), Some("/explicit/sock"));
    }

    #[test]
    fn missing_plugin_id_errors() {
        let argv = vec!["--socket".to_string(), "/tmp/x".to_string()];
        let err = RunnerArgs::parse(&argv, no_env).unwrap_err();
        assert!(matches!(err, RunnerError::MissingPluginId));
    }

    #[test]
    fn the_token_credential_supplies_what_the_options_do_not() {
        // Under systemd the launch values arrive only in the credential file;
        // a stale environment value never wins over it, and an option does.
        let dir = tempfile::tempdir().unwrap();
        std::fs::write(
            dir.path().join(TOKEN_CREDENTIAL_NAME),
            "ADOS_PLUGIN_TOKEN=pay=load|sig==\n\
             ADOS_PLUGIN_SOCKET=/run/ados/plugins/com.example.demo/host.sock\n\
             ADOS_PLUGIN_AGENT_ID=drone-abc\n\
             ADOS_PLUGIN_DATA_DIR=/var/ados/plugin-data/com.example.demo/drones/drone-abc\n",
        )
        .unwrap();
        let creds = dir.path().to_string_lossy().into_owned();
        let env = move |k: &str| match k {
            "CREDENTIALS_DIRECTORY" => Some(creds.clone()),
            "ADOS_PLUGIN_TOKEN" => Some("stale-env".to_string()),
            _ => None,
        };
        let argv = vec![
            "com.example.demo".to_string(),
            "--agent-id=explicit".to_string(),
        ];
        let args = RunnerArgs::parse(&argv, env).unwrap();
        assert_eq!(args.token.as_deref(), Some("pay=load|sig=="));
        assert_eq!(
            args.socket_path.as_deref(),
            Some("/run/ados/plugins/com.example.demo/host.sock")
        );
        assert_eq!(args.agent_id, "explicit");
        assert_eq!(
            args.data_dir.as_deref(),
            Some("/var/ados/plugin-data/com.example.demo/drones/drone-abc")
        );
    }

    #[test]
    fn an_unreadable_credential_falls_back_to_the_environment() {
        let dir = tempfile::tempdir().unwrap();
        let missing = dir.path().join("absent").to_string_lossy().into_owned();
        let env = move |k: &str| match k {
            "CREDENTIALS_DIRECTORY" => Some(missing.clone()),
            "ADOS_PLUGIN_TOKEN" => Some("env-tok".to_string()),
            _ => None,
        };
        let args = RunnerArgs::parse(&["com.example.demo".to_string()], env).unwrap();
        assert_eq!(args.token.as_deref(), Some("env-tok"));
        assert_eq!(args.socket_path, None);
    }

    /// A dummy plugin proves the trait's default-no-op hooks compile and that a
    /// minimal implementation only needs `new`.
    struct DummyPlugin {
        started: bool,
    }

    #[async_trait]
    impl Plugin for DummyPlugin {
        fn new() -> Self {
            Self { started: false }
        }
        async fn on_start(&mut self, _ctx: &PluginContext) -> Result<(), ClientError> {
            self.started = true;
            Ok(())
        }
    }

    #[test]
    fn dummy_plugin_constructs_with_only_new() {
        let p = DummyPlugin::new();
        assert!(!p.started);
    }

    #[test]
    fn ensure_data_dir_creates_the_nested_leaf() {
        // The bug this closes: the per-drone leaf never existed, so a plugin's
        // first write failed. Prove the nested path is created and a file can be
        // written under it — not just that a path string was produced.
        let base = tempfile::tempdir().unwrap();
        let leaf = base.path().join("com.example.plugin/drones/drone-xyz");
        assert!(!leaf.exists());

        ensure_data_dir(Some(leaf.to_str().unwrap()));
        assert!(
            leaf.is_dir(),
            "the nested data dir must exist after the call"
        );

        // A plugin can now actually write there.
        std::fs::write(leaf.join("state.json"), b"{}").expect("write under data dir");

        // Idempotent: a second call on an existing dir is fine.
        ensure_data_dir(Some(leaf.to_str().unwrap()));
        // None is a no-op (does not panic).
        ensure_data_dir(None);
    }

    #[tokio::test]
    async fn no_bridge_when_socket_or_token_absent() {
        // plugin id present but no socket/token -> NoBridge, before any connect.
        let args = RunnerArgs {
            plugin_id: "com.example.demo".to_string(),
            socket_path: None,
            token: None,
            agent_id: String::new(),
            data_dir: None,
        };
        let err = run_plugin_with::<DummyPlugin, _>(
            args,
            "1.0.0",
            BTreeMap::new(),
            std::future::ready(()),
        )
        .await
        .unwrap_err();
        assert!(matches!(err, RunnerError::NoBridge));
    }

    #[tokio::test]
    async fn a_dropped_host_session_ends_the_runner_with_an_error() {
        use ados_protocol::plugin::{Envelope, PROTOCOL_VERSION};
        use tokio::io::{AsyncReadExt, AsyncWriteExt};

        // A plugin-host restart: the host answers the handshake, then its end
        // of the session goes away. The runner must stop with an error (the
        // binary then exits non-zero and systemd restarts it) instead of
        // running on against a dead socket until an operator intervenes.
        let dir = tempfile::tempdir().expect("tempdir");
        let sock = dir.path().join("p.sock");
        let listener = tokio::net::UnixListener::bind(&sock).expect("bind");
        let host = tokio::spawn(async move {
            let (mut stream, _) = listener.accept().await.expect("accept");
            let len = stream.read_u32().await.expect("hello header") as usize;
            let mut body = vec![0u8; len];
            stream.read_exact(&mut body).await.expect("hello body");
            let hello = Envelope::from_msgpack(&body).expect("hello envelope");
            let ready = Envelope {
                version: PROTOCOL_VERSION,
                kind: "response".to_string(),
                method: String::new(),
                capability: String::new(),
                args: Value::Map(vec![(Value::from("ready"), Value::Boolean(true))]),
                request_id: hello.request_id,
                token: String::new(),
                error: None,
            };
            stream
                .write_all(&ready.encode_frame().unwrap())
                .await
                .expect("ready");
            // The host restarts: this session is gone.
            drop(stream);
        });

        let args = RunnerArgs {
            plugin_id: "com.example.demo".to_string(),
            socket_path: Some(sock.to_string_lossy().into_owned()),
            token: Some("tok".to_string()),
            agent_id: String::new(),
            data_dir: None,
        };
        // No shutdown signal ever arrives; only the session loss can end it.
        let result = tokio::time::timeout(
            std::time::Duration::from_secs(5),
            run_plugin_with::<DummyPlugin, _>(
                args,
                "1.0.0",
                BTreeMap::new(),
                std::future::pending(),
            ),
        )
        .await
        .expect("the runner must stop once the host session is lost");
        host.await.expect("fake host");
        assert!(
            matches!(result, Err(RunnerError::SessionLost)),
            "got {result:?}"
        );
    }
}