warpgate 0.35.3

use crate::plugin_error::WarpgatePluginError;
use extism::{Error, Function, Manifest, Plugin};
use scc::hash_map::Entry;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use starbase_styles::{apply_style_tags, color};
use starbase_utils::{
    envx::{bool_var, is_ci},
    hash,
};
use std::fmt::Debug;
use std::path::{Path, PathBuf};
use std::sync::{Arc, OnceLock};
use std::time::Instant;
use system_env::{SystemArch, SystemLibc, SystemOS};
use tokio::sync::RwLock;
use tokio::task::spawn_blocking;
use tracing::{instrument, trace};
use warpgate_api::{
    HostEnvironment, Id, RealPath, VirtualPath, convert_to_real_path, convert_to_virtual_path,
    sort_paths_list,
};

fn is_incompatible_runtime(error: &Error) -> bool {
    let check = |message: String| {
        // unknown import: `env::exec_command` has not been defined
        message.contains("unknown import") && message.contains("env::")
    };

    if let Some(source) = error.source()
        && check(source.to_string())
    {
        return true;
    }

    check(error.to_string())
}

/// Inject our default configuration into the provided plugin [`Manifest`].
/// This will set `plugin_id`, `host_environment`, and `virtual_paths`
/// for use within PDKs.
#[instrument(skip(manifest))]
pub fn inject_default_manifest_config(
    id: &Id,
    home_dir: &Path,
    manifest: &mut Manifest,
) -> Result<(), WarpgatePluginError> {
    if !manifest.config.contains_key("plugin_id") {
        trace!(id = id.as_str(), "Storing plugin identifier");

        manifest.config.insert("plugin_id".into(), id.to_string());
    }

    if !manifest.config.contains_key("host_environment") {
        let os = SystemOS::from_env();

        let env = serde_json::to_string(&HostEnvironment {
            arch: SystemArch::from_env(),
            ci: is_ci(),
            libc: SystemLibc::detect(os),
            os,
            home_dir: VirtualPath::new("/userhome"),
        })
        .map_err(|error| WarpgatePluginError::InvalidInput {
            id: id.to_owned(),
            func: "host_environment".into(),
            error: Box::new(error),
        })?;

        trace!(id = id.as_str(), env = %env, "Storing host environment");

        manifest.config.insert("host_environment".into(), env);
    }

    if !manifest.config.contains_key("virtual_paths") {
        let mut paths = vec![];

        if let Some(allowed_paths) = &manifest.allowed_paths {
            for (host, guest) in allowed_paths {
                paths.push((PathBuf::from(host), guest.to_owned()));
            }
        }

        sort_paths_list(&mut paths);

        let paths =
            serde_json::to_string(&paths).map_err(|error| WarpgatePluginError::InvalidInput {
                id: id.to_owned(),
                func: "virtual_paths".into(),
                error: Box::new(error),
            })?;

        trace!(id = id.as_str(), paths = %paths, "Storing virtual paths");

        manifest.config.insert("virtual_paths".into(), paths);
    }

    Ok(())
}

/// A callback that is executed before and after a plugin function call.
/// Receives the function name, the input (before only), and the output (after only).
pub type OnCallFn = Arc<dyn Fn(&str, Option<&str>, Option<&str>) + Send + Sync>;

/// A container around Extism's [`Plugin`] and [`Manifest`] types that provides convenience
/// methods for calling and caching functions from the WASM plugin. It also provides
/// additional methods for easily working with WASI and virtual paths.
pub struct PluginContainer {
    /// Unique identifier of the plugin.
    pub id: Id,

    /// The [`Manifest`] that the plugin was created with.
    pub manifest: Manifest,

    /// Mapping of virtual paths, from host to guest paths.
    pub virtual_paths: Vec<(PathBuf, PathBuf)>,

    cache: bool,
    debug: bool,

    func_cache: Arc<scc::HashMap<String, Vec<u8>>>,
    on_call_func: Arc<OnceLock<OnCallFn>>,
    plugin: Arc<RwLock<Plugin>>,
}

impl PluginContainer {
    /// Create a new container with the provided manifest and host functions.
    #[instrument(name = "new_plugin", skip(manifest, functions))]
    pub fn new(
        id: Id,
        manifest: Manifest,
        functions: impl IntoIterator<Item = Function>,
    ) -> Result<PluginContainer, WarpgatePluginError> {
        trace!(id = id.as_str(), "Creating plugin container");

        let plugin = Plugin::new(&manifest, functions, true).map_err(|error| {
            if is_incompatible_runtime(&error) {
                WarpgatePluginError::IncompatibleRuntime { id: id.clone() }
            } else {
                WarpgatePluginError::FailedContainer {
                    id: id.clone(),
                    error: Box::new(error),
                }
            }
        })?;

        trace!(
            id = id.as_str(),
            plugin = plugin.id.to_string(),
            "Created plugin container",
        );

        let mut virtual_paths = match manifest.allowed_paths.as_ref() {
            Some(paths) => paths
                .iter()
                .map(|(host, guest)| (PathBuf::from(host), guest.to_owned()))
                .collect(),
            None => Vec::new(),
        };

        sort_paths_list(&mut virtual_paths);

        Ok(PluginContainer {
            virtual_paths,
            manifest,
            plugin: Arc::new(RwLock::new(plugin)),
            id,
            func_cache: Arc::new(scc::HashMap::new()),
            on_call_func: Arc::new(OnceLock::new()),
            cache: !bool_var("WARPGATE_NO_FUNC_CACHE"),
            debug: bool_var("WARPGATE_DEBUG_CALL"),
        })
    }

    /// Set a callback handler to be executed when calling a plugin function.
    pub fn set_on_call(&self, func: OnCallFn) {
        let _ = self.on_call_func.set(func);
    }

    /// Call a function on the plugin with no input and cache the output before returning it.
    /// Subsequent calls will read from the cache.
    pub async fn cache_func<F, O>(&self, func: F) -> Result<O, WarpgatePluginError>
    where
        F: Debug + AsRef<str>,
        O: Debug + DeserializeOwned,
    {
        self.cache_func_with(func, Empty::default()).await
    }

    /// Call a function on the plugin with the given input and cache the output
    /// before returning it. Subsequent calls with the same input will read from the cache.
    #[instrument(skip(self))]
    pub async fn cache_func_with<F, I, O>(
        &self,
        func: F,
        input: I,
    ) -> Result<O, WarpgatePluginError>
    where
        F: Debug + AsRef<str>,
        I: Debug + Serialize,
        O: Debug + DeserializeOwned,
    {
        if !self.cache {
            return self.call_func_with(func, input).await;
        }

        let func = func.as_ref();
        let input = self.format_input(func, input)?;
        let cache_key = format!("{func}-{}", hash::base64::from_bytes(&input));

        match self.func_cache.entry_async(cache_key).await {
            // Check if cache exists already
            Entry::Occupied(entry) => self.parse_output(func, entry.get()),
            // Otherwise call the function and cache the result
            Entry::Vacant(entry) => {
                let data = self.call(func, input).await?;
                let output: O = self.parse_output(func, &data)?;

                entry.insert_entry(data);

                Ok(output)
            }
        }
    }

    /// Call a function on the plugin with no input and return the output.
    pub async fn call_func<F, O>(&self, func: F) -> Result<O, WarpgatePluginError>
    where
        F: Debug + AsRef<str>,
        O: Debug + DeserializeOwned,
    {
        self.call_func_with(func, Empty::default()).await
    }

    /// Call a function on the plugin with the given input and return the output.
    #[instrument(skip(self))]
    pub async fn call_func_with<F, I, O>(&self, func: F, input: I) -> Result<O, WarpgatePluginError>
    where
        F: Debug + AsRef<str>,
        I: Debug + Serialize,
        O: Debug + DeserializeOwned,
    {
        let func = func.as_ref();

        self.parse_output(
            func,
            &self.call(func, self.format_input(func, input)?).await?,
        )
    }

    /// Call a function on the plugin with the given input and ignore the output.
    #[instrument(skip(self))]
    pub async fn call_func_without_output<F, I>(
        &self,
        func: F,
        input: I,
    ) -> Result<(), WarpgatePluginError>
    where
        F: Debug + AsRef<str>,
        I: Debug + Serialize,
    {
        let func = func.as_ref();

        self.call(func, self.format_input(func, input)?).await?;

        Ok(())
    }

    /// Call a function on the plugin with the given raw input and return the raw output.
    #[instrument(skip(self, input))]
    pub async fn call(&self, func: &str, input: String) -> Result<Vec<u8>, WarpgatePluginError> {
        let mut instance = Arc::clone(&self.plugin).write_owned().await;
        let uuid = instance.id.to_string(); // Copy
        let instant = Instant::now();
        let truncate_size = 5000;

        trace!(
            id = self.id.as_str(),
            plugin = &uuid,
            input = %(if input.len() > truncate_size && !self.debug {
                "(truncated)"
            } else {
                &input
            }),
            "Calling guest function {}",
            color::property(func),
        );

        if let Some(callback) = self.on_call_func.get() {
            callback(func, Some(&input), None);
        }

        let func_name = func.to_string();

        let result = spawn_blocking(move || instance.call::<String, Vec<u8>>(func_name, input))
            .await
            .map_err(|error| WarpgatePluginError::FailedPluginCall {
                id: self.id.clone(),
                func: func.to_owned(),
                error: error.to_string(),
            })?;

        let output = result.map_err(|error| {
            if is_incompatible_runtime(&error) {
                return WarpgatePluginError::IncompatibleRuntime {
                    id: self.id.clone(),
                };
            }

            let message = apply_style_tags(
                error
                    .source()
                    .map(|src| src.to_string())
                    .unwrap_or_else(|| error.to_string())
                    .replace("\\\\n", "\n")
                    .replace("\\n", "\n")
                    .trim(),
            );

            // When in debug mode, include more information around errors.
            #[cfg(debug_assertions)]
            {
                WarpgatePluginError::FailedPluginCall {
                    id: self.id.clone(),
                    func: func.to_owned(),
                    error: message,
                }
            }

            // When in release mode, errors don't render properly with the
            // previous variant, so this is a special variant that renders as-is.
            #[cfg(not(debug_assertions))]
            {
                WarpgatePluginError::FailedPluginCallRelease { error: message }
            }
        })?;

        trace!(
            id = self.id.as_str(),
            plugin = &uuid,
            output = %(if output.len() > truncate_size && !self.debug {
                "(truncated)".to_string()
            } else {
                String::from_utf8_lossy(&output).to_string()
            }),
            elapsed = ?instant.elapsed(),
            "Called guest function {}",
            color::property(func),
        );

        if let Some(callback) = self.on_call_func.get() {
            callback(func, None, Some(String::from_utf8_lossy(&output).as_ref()));
        }

        Ok(output)
    }

    /// Return true if the plugin has a function with the given id.
    #[instrument(skip(self))]
    pub async fn has_func<T>(&self, func: T) -> bool
    where
        T: AsRef<str> + Debug,
    {
        let func = func.as_ref();

        match self.func_cache.entry_async(func.into()).await {
            Entry::Occupied(entry) => entry.get()[0] == 1,
            Entry::Vacant(entry) => {
                let exists = self.plugin.read().await.function_exists(func);
                entry.insert_entry(vec![exists as u8]);
                exists
            }
        }
    }

    /// Convert the provided virtual guest path to an absolute host path.
    pub fn to_real_path<P>(&self, path: P) -> RealPath
    where
        P: AsRef<Path>,
    {
        convert_to_real_path(&path, &self.virtual_paths)
            .unwrap_or_else(|| RealPath::new(path.as_ref().as_os_str()))
    }

    /// Convert the provided list of virtual guest paths to an absolute host paths.
    pub fn to_real_paths<I, P>(&self, paths: I) -> Vec<RealPath>
    where
        I: IntoIterator<Item = P>,
        P: AsRef<Path>,
    {
        paths
            .into_iter()
            .map(|path| self.to_real_path(path))
            .collect()
    }

    /// Convert the provided absolute host path to a virtual guest path.
    pub fn to_virtual_path<P>(&self, path: P) -> VirtualPath
    where
        P: AsRef<Path>,
    {
        convert_to_virtual_path(&path, &self.virtual_paths)
            .unwrap_or_else(|| VirtualPath::new(path.as_ref().as_os_str()))
    }

    /// Convert the provided list of absolute host paths to virtual guest paths.
    pub fn to_virtual_paths<I, P>(&self, paths: I) -> Vec<VirtualPath>
    where
        I: IntoIterator<Item = P>,
        P: AsRef<Path>,
    {
        paths
            .into_iter()
            .map(|path| self.to_virtual_path(path))
            .collect()
    }

    fn format_input<I: Serialize>(
        &self,
        func: &str,
        input: I,
    ) -> Result<String, WarpgatePluginError> {
        serde_json::to_string(&input).map_err(|error| WarpgatePluginError::InvalidInput {
            id: self.id.clone(),
            func: func.to_owned(),
            error: Box::new(error),
        })
    }

    fn parse_output<O: DeserializeOwned>(
        &self,
        func: &str,
        data: &[u8],
    ) -> Result<O, WarpgatePluginError> {
        serde_json::from_slice(data).map_err(|error| WarpgatePluginError::InvalidOutput {
            id: self.id.clone(),
            func: func.to_owned(),
            error: Box::new(error),
        })
    }
}

#[derive(Debug, Default, Deserialize, Serialize)]
struct Empty {}