cardinal_plugins/

container.rs

use crate::builtin::restricted_route_middleware::RestrictedRouteMiddleware;
use crate::request_context::RequestContext;
use crate::runner::{DynRequestMiddleware, DynResponseMiddleware, MiddlewareResult};
use cardinal_base::context::CardinalContext;
use cardinal_base::provider::Provider;
use cardinal_config::Plugin;
use cardinal_errors::CardinalError;
use cardinal_wasm_plugins::host::{HostFunctionBuilder, HostImportHandle};
use cardinal_wasm_plugins::plugin::WasmPlugin;
use cardinal_wasm_plugins::runner::{host_import_from_builder, ExecutionPhase, WasmRunner};
use cardinal_wasm_plugins::wasmer::{Function, FunctionEnv, Store};
use cardinal_wasm_plugins::{ResponseState, SharedExecutionContext};
use pingora::http::ResponseHeader;
use pingora::prelude::Session;
use std::collections::HashMap;
use std::sync::Arc;
use tracing::{error, warn};

pub enum PluginBuiltInType {
    Inbound(Arc<DynRequestMiddleware>),
    Outbound(Arc<DynResponseMiddleware>),
}

pub enum PluginHandler {
    Builtin(PluginBuiltInType),
    Wasm(Arc<WasmPlugin>),
}

pub struct PluginContainer {
    plugins: HashMap<String, Arc<PluginHandler>>,
    host_imports: Vec<HostImportHandle>,
}

impl PluginContainer {
    pub fn new() -> Self {
        Self {
            plugins: HashMap::from_iter(Self::builtin_plugins()),
            host_imports: Vec::new(),
        }
    }

    pub fn new_empty() -> Self {
        Self {
            plugins: HashMap::new(),
            host_imports: Vec::new(),
        }
    }

    pub fn builtin_plugins() -> Vec<(String, Arc<PluginHandler>)> {
        vec![(
            "RestrictedRouteMiddleware".to_string(),
            Arc::new(PluginHandler::Builtin(PluginBuiltInType::Inbound(
                Arc::new(RestrictedRouteMiddleware),
            ))),
        )]
    }

    pub fn add_plugin(&mut self, name: String, plugin: PluginHandler) {
        self.plugins.insert(name, Arc::new(plugin));
    }

    pub fn remove_plugin(&mut self, name: &str) {
        self.plugins.remove(name);
    }

    pub fn add_host_function<F>(
        &mut self,
        namespace: impl Into<String>,
        name: impl Into<String>,
        builder: F,
    ) where
        F: Fn(&mut Store, &FunctionEnv<SharedExecutionContext>) -> Function + Send + Sync + 'static,
    {
        let builder: HostFunctionBuilder = Arc::new(builder);
        let import = host_import_from_builder(namespace, name, builder);
        self.host_imports.push(import);
    }

    pub fn extend_host_functions<I>(&mut self, functions: I)
    where
        I: IntoIterator<Item = HostImportHandle>,
    {
        self.host_imports.extend(functions);
    }

    fn host_imports(&self) -> Option<&[HostImportHandle]> {
        if self.host_imports.is_empty() {
            None
        } else {
            Some(&self.host_imports)
        }
    }

    pub async fn run_request_filter(
        &self,
        name: &str,
        session: &mut Session,
        req_ctx: &mut RequestContext,
    ) -> Result<MiddlewareResult, CardinalError> {
        let plugin = self
            .plugins
            .get(name)
            .ok_or_else(|| CardinalError::Other(format!("Plugin {name} does not exist")))?;

        let (result, shared_ctx, force_respond) = match plugin.as_ref() {
            PluginHandler::Builtin(builtin) => match builtin {
                PluginBuiltInType::Inbound(filter) => {
                    let result = filter
                        .on_request(session, req_ctx, req_ctx.cardinal_context.clone())
                        .await?;

                    Ok((result, req_ctx.shared_context(), false))
                }
                PluginBuiltInType::Outbound(_) => Err(CardinalError::Other(format!(
                    "The filter {name} is not a request filter"
                ))),
            },
            PluginHandler::Wasm(wasm) => {
                let runner = WasmRunner::new(wasm, ExecutionPhase::Inbound, self.host_imports());

                let exec = runner.run(req_ctx.shared_context())?;
                let result = if exec.should_continue {
                    MiddlewareResult::Continue(HashMap::new())
                } else {
                    MiddlewareResult::Responded
                };

                Ok((
                    result,
                    exec.execution_context.clone(),
                    !exec.should_continue,
                ))
            }
        }?;

        let (header_updates, response_snapshot) = Self::snapshot_execution_context(&shared_ctx);

        Self::apply_request_header_updates(session, header_updates);

        Ok(Self::finalize_request_result(result, response_snapshot, session, force_respond).await)
    }

    fn snapshot_execution_context(
        shared_ctx: &SharedExecutionContext,
    ) -> (Vec<(String, String)>, ResponseState) {
        let guard = shared_ctx.read();

        let request_headers: Vec<(String, String)> = guard
            .request()
            .headers()
            .iter()
            .filter_map(|(key, value)| {
                value
                    .to_str()
                    .ok()
                    .map(|v| (key.as_str().to_string(), v.to_string()))
            })
            .collect();

        let response_state = guard.response().clone();
        (request_headers, response_state)
    }

    fn apply_request_header_updates(session: &mut Session, header_updates: Vec<(String, String)>) {
        if header_updates.is_empty() {
            return;
        }

        for (key, val) in header_updates {
            let _ = session.req_header_mut().insert_header(key, val);
        }
    }

    async fn finalize_request_result(
        result: MiddlewareResult,
        response_snapshot: ResponseState,
        session: &mut Session,
        force_respond: bool,
    ) -> MiddlewareResult {
        if force_respond || response_snapshot.status_override().is_some() {
            let state = Self::build_response_header(&response_snapshot);
            return Self::respond_from_response_state(state, response_snapshot.status(), session)
                .await;
        }

        match result {
            MiddlewareResult::Responded => MiddlewareResult::Responded,
            MiddlewareResult::Continue(mut headers) => {
                let context_headers: HashMap<String, String> = response_snapshot
                    .headers()
                    .iter()
                    .filter_map(|(key, value)| {
                        value
                            .to_str()
                            .ok()
                            .map(|v| (key.as_str().to_string(), v.to_string()))
                    })
                    .collect();

                headers.extend(context_headers);
                MiddlewareResult::Continue(headers)
            }
        }
    }

    pub async fn run_response_filter(
        &self,
        name: &str,
        session: &mut Session,
        req_ctx: &mut RequestContext,
        response: &mut pingora::http::ResponseHeader,
    ) {
        let plugin = self
            .plugins
            .get(name)
            .ok_or_else(|| CardinalError::Other(format!("Plugin {name} does not exist")));

        if let Ok(plugin) = plugin {
            match plugin.as_ref() {
                PluginHandler::Builtin(builtin) => match builtin {
                    PluginBuiltInType::Inbound(_) => {
                        error!("The filter {name} is not a response filter");
                    }
                    PluginBuiltInType::Outbound(filter) => {
                        filter
                            .on_response(
                                session,
                                req_ctx,
                                response,
                                req_ctx.cardinal_context.clone(),
                            )
                            .await
                    }
                },
                PluginHandler::Wasm(wasm) => {
                    let runner =
                        WasmRunner::new(wasm, ExecutionPhase::Outbound, self.host_imports());

                    match runner.run(req_ctx.shared_context()) {
                        Ok(exec) => {
                            let snapshot = {
                                let guard = exec.execution_context.read();
                                guard.response().clone()
                            };

                            for (key, val) in snapshot.headers().iter() {
                                let _ = response.insert_header(key.clone(), val.clone());
                            }

                            if let Some(status) = snapshot.status_override() {
                                let _ = response.set_status(status);
                            }
                        }
                        Err(e) => {
                            error!("Failed to run plugin {}: {}", name, e);
                        }
                    }
                }
            }
        }
    }

    pub fn build_response_header(response: &ResponseState) -> ResponseHeader {
        let mut header = ResponseHeader::build(response.status(), None)
            .expect("failed to build response header");

        for (key, value) in response.headers().iter() {
            let _ = header.insert_header(key.clone(), value.clone());
        }

        header
    }

    pub async fn respond_from_response_state(
        response_header: ResponseHeader,
        status: u16,
        session: &mut Session,
    ) -> MiddlewareResult {
        let _ = session
            .write_response_header(Box::new(response_header), false)
            .await;
        let _ = session.respond_error(status).await;

        MiddlewareResult::Responded
    }
}

impl Default for PluginContainer {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait::async_trait]
impl Provider for PluginContainer {
    async fn provide(ctx: &CardinalContext) -> Result<Self, CardinalError> {
        let preloaded_plugins = ctx.config.plugins.clone();
        let mut plugin_container = PluginContainer::new();

        for plugin in preloaded_plugins {
            let plugin_name = plugin.name();
            let plugin_exists = plugin_container.plugins.contains_key(plugin_name);

            if plugin_exists {
                warn!("Plugin {} already exists, skipping", plugin_name);
                continue;
            }

            match plugin {
                Plugin::Builtin(_) => continue,
                Plugin::Wasm(wasm_config) => {
                    let wasm_plugin = WasmPlugin::from_path(&wasm_config.path).map_err(|e| {
                        CardinalError::Other(format!(
                            "Failed to load plugin {}: {}",
                            wasm_config.name, e
                        ))
                    })?;
                    plugin_container.plugins.insert(
                        wasm_config.name.clone(),
                        Arc::new(PluginHandler::Wasm(Arc::new(wasm_plugin))),
                    );
                }
            }
        }

        Ok(plugin_container)
    }
}