churn-v2/crates/churn/src/agent/plugins.rs

use anyhow::Context;
use futures::StreamExt;
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use tokio::io::AsyncWriteExt;
use tokio::sync::Mutex;
use wasmtime::component::*;
use wasmtime::{Config, Engine, Store};
use wasmtime_wasi::{DirPerms, FilePerms, WasiCtx, WasiCtxBuilder, WasiView};

wasmtime::component::bindgen!({
    path: "wit/world.wit",
    world: "churn",
    async: true
});

#[derive(Clone)]
pub struct PluginStore {
    inner: Arc<Mutex<InnerPluginStore>>,
}

impl PluginStore {
    pub fn new() -> anyhow::Result<Self> {
        Ok(Self {
            inner: Arc::new(Mutex::new(InnerPluginStore::new()?)),
        })
    }

    pub async fn id(&self, plugin: &str) -> anyhow::Result<String> {
        let mut inner = self.inner.lock().await;
        inner.id(plugin).await
    }

    pub async fn execute(&self, plugin: &str) -> anyhow::Result<()> {
        let mut inner = self.inner.lock().await;
        inner.execute(plugin).await
    }
}

pub struct InnerPluginStore {
    store: wasmtime::Store<ServerWasiView>,
    linker: wasmtime::component::Linker<ServerWasiView>,
    engine: wasmtime::Engine,
}

impl InnerPluginStore {
    pub fn new() -> anyhow::Result<Self> {
        let mut config = Config::default();
        config.wasm_component_model(true);
        config.async_support(true);
        let engine = Engine::new(&config)?;
        let mut linker: wasmtime::component::Linker<ServerWasiView> = Linker::new(&engine);

        // Add the command world (aka WASI CLI) to the linker
        wasmtime_wasi::add_to_linker_async(&mut linker).context("Failed to link command world")?;
        let wasi_view = ServerWasiView::new();
        let store = Store::new(&engine, wasi_view);

        Ok(Self {
            store,
            linker,
            engine,
        })
    }

    pub async fn id(&mut self, plugin: &str) -> anyhow::Result<String> {
        let plugin = self.ensure_plugin(plugin).await?;

        plugin
            .interface0
            .call_id(&mut self.store)
            .await
            .context("Failed to call add function")
    }

    pub async fn execute(&mut self, plugin: &str) -> anyhow::Result<()> {
        let plugin = self.ensure_plugin(plugin).await?;

        plugin
            .interface0
            .call_execute(&mut self.store)
            .await
            .context("Failed to call add function")
    }

    async fn ensure_plugin(&mut self, plugin: &str) -> anyhow::Result<Churn> {
        let cache = dirs::cache_dir()
            .ok_or(anyhow::anyhow!("failed to find cache dir"))?
            .join("io.kjuulh.churn");

        let (plugin_name, plugin_version) = plugin.split_once("@").unwrap_or((plugin, "latest"));

        let plugin_path = cache
            .join("plugins")
            .join(plugin_name)
            .join(plugin_version)
            .join(format!("{plugin_name}.wasm"));

        let no_cache: bool = std::env::var("CHURN_NO_CACHE")
            .unwrap_or("false".into())
            .parse()?;

        if !plugin_path.exists() || no_cache {
            tracing::info!(
                plugin_name = plugin_name,
                plugin_version = plugin_version,
                "downloading plugin"
            );
            if let Some(parent) = plugin_path.parent() {
                tokio::fs::create_dir_all(parent).await?;
            }

            let req = reqwest::get(format!("https://api-minio.front.kjuulh.io/churn-registry/{plugin_name}/{plugin_version}/{plugin_name}.wasm")).await.context("failed to get plugin from registry")?;
            let mut stream = req.bytes_stream();
            let mut file = tokio::fs::File::create(&plugin_path).await?;
            while let Some(chunk) = stream.next().await {
                let chunk = chunk?;
                file.write_all(&chunk).await?;
            }
            file.flush().await?;
        }

        let component =
            Component::from_file(&self.engine, plugin_path).context("Component file not found")?;

        tracing::debug!(
            plugin_name = plugin_name,
            plugin_version = plugin_version,
            "instantiating plugin"
        );
        let instance = Churn::instantiate_async(&mut self.store, &component, &self.linker)
            .await
            .context("Failed to instantiate the example world")?;

        Ok(instance)
    }
}

struct ServerWasiView {
    table: ResourceTable,
    ctx: WasiCtx,
}

impl ServerWasiView {
    fn new() -> Self {
        let table = ResourceTable::new();

        let ctx = WasiCtxBuilder::new()
            .inherit_stdio()
            .inherit_env()
            .inherit_network()
            .preopened_dir("/", "/", DirPerms::all(), FilePerms::all())
            .expect("to be able to open root")
            .build();

        Self { table, ctx }
    }
}

impl WasiView for ServerWasiView {
    fn table(&mut self) -> &mut ResourceTable {
        &mut self.table
    }

    fn ctx(&mut self) -> &mut WasiCtx {
        &mut self.ctx
    }
}