camel-component-wasm 0.46.0

//! WasmSourceConsumer — Consumer trait implementation for the `source` world.
//!
//! The guest IS the source: it owns an async `run()` loop that awaits
//! host-provided `accept-http` / `submit-exchange`. The host bridges the
//! async guest to the pipeline via bounded tokio channels, and uses the
//! cancel token (raced in each import) + epoch deadline for cancellation.

use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use camel_api::CamelError;
use camel_api::security_policy::{AccessMode, RouteSecurityPlan, TransportId};
use camel_component_api::{
    ComponentContext, ConcurrencyModel, Consumer, ConsumerContext, SecurityContext,
};
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tracing::{debug, warn};
use wasmtime::component::{Component, Linker};
use wasmtime::{AsContextMut, Config, Engine, Store};

use crate::config::WasmConfig;
use crate::source_bindings::Source;
use crate::source_bindings::camel::plugin::source_host::CapabilityRequest;
use crate::source_host::{
    DEFAULT_MAX_REQUEST_BODY_BYTES, HttpListenerHandle, SourceChannels, SourceHostState,
    WasmSourceKernelAuth, add_to_linker, run_http_listener, run_pipeline_bridge,
};
use crate::staged_listener;

/// Epoch deadline (in ticks) granted to the guest's `configure()` call.
///
/// No epoch ticker runs during `configure()`, so this budget is never
/// decremented; it exists solely to move the store off its default
/// already-expired deadline (0) that would otherwise trap the guest at
/// the first epoch check when `epoch_interruption(true)` is enabled.
const CONFIGURE_EPOCH_DEADLINE: u64 = u64::MAX;

/// Consumer backed by a WASM `source` world component.
pub struct WasmSourceConsumer {
    module_path: PathBuf,
    /// Full `wasm:...` endpoint URI — the naming key for bind-gate and
    /// bind-conflict errors (the consumer has no other route identity).
    endpoint_uri: String,
    guest_config: Vec<(String, String)>,
    config: WasmConfig,
    #[allow(dead_code)]
    registry: Arc<dyn ComponentContext>,
    cancel_token: CancellationToken,
    engine: Option<Arc<Engine>>,
    run_task: Option<JoinHandle<Result<(), CamelError>>>,
    listener_task: Option<JoinHandle<()>>,
    bridge_task: Option<JoinHandle<()>>,
    /// Kernel authentication state (plan + providers) captured from the
    /// route's security context. `None` until `set_security_context`
    /// delivers a context carrying both pieces.
    kernel: Option<Arc<WasmSourceKernelAuth>>,
    /// The whole classification plan captured from the security context —
    /// not just the access mode. Task 1.4's bind-gate snapshot needs the
    /// full plan even when providers are absent (plan-only contexts).
    plan: Option<RouteSecurityPlan>,
}

impl WasmSourceConsumer {
    pub fn new(
        module_path: PathBuf,
        uri: impl Into<String>,
        config: WasmConfig,
        guest_config: Vec<(String, String)>,
        registry: Arc<dyn ComponentContext>,
    ) -> Self {
        Self {
            module_path,
            endpoint_uri: uri.into(),
            guest_config,
            config,
            registry,
            cancel_token: CancellationToken::new(),
            engine: None,
            run_task: None,
            listener_task: None,
            bridge_task: None,
            kernel: None,
            plan: None,
        }
    }

    /// Test accessor for the captured classification mode.
    ///
    /// Integration tests are off-crate and cannot read the `pub(crate)`
    /// kernel fields; this exposes only what the capture tests pin.
    #[doc(hidden)]
    pub fn plan_access_mode(&self) -> Option<AccessMode> {
        self.plan.as_ref().map(|p| p.access_mode.clone())
    }
}

#[async_trait::async_trait]
impl Consumer for WasmSourceConsumer {
    async fn start(&mut self, ctx: ConsumerContext) -> Result<(), CamelError> {
        // Adopt the Runtime's context token: stop() must observe it, and the
        // clone into SourceHostState below must carry it into every host
        // import's cancel select!. The new() token is only a placeholder.
        self.cancel_token = ctx.cancel_token();

        // 1. Engine config: component model + async + epoch interruption.
        //    concurrency_support is REQUIRED for Store::run_concurrent /
        //    call_run_async (async source world). Mirrors the plugin engine
        //    (runtime.rs).
        let mut wasm_config = Config::new();
        wasm_config.wasm_component_model(true);
        wasm_config.epoch_interruption(true);
        wasm_config.concurrency_support(true);
        let engine = Arc::new(
            Engine::new(&wasm_config)
                .map_err(|e| CamelError::ProcessorError(format!("wasmtime engine: {e}")))?,
        );

        // 2. Size cap: reject oversized modules before compilation (R4-H3)
        crate::config::validate_wasm_size(&self.module_path, self.config.max_wasm_size_bytes)
            .map_err(CamelError::ProcessorError)?;

        // 3. Load component
        let component = Component::from_file(&engine, &self.module_path)
            .map_err(|e| CamelError::ProcessorError(format!("component load: {e}")))?;

        // 3. Create bounded channels
        let channels = SourceChannels::new();

        // 4. Create WASI context (guest targets wasm32-wasip2)
        let wasi = wasmtime_wasi::WasiCtxBuilder::new().build();

        // 5. Create store with host state
        let cancel = self.cancel_token.clone();
        let mut store = Store::new(
            &engine,
            SourceHostState {
                table: wasmtime::component::ResourceTable::new(),
                wasi,
                request_rx: channels.request_rx,
                exchange_tx: channels.exchange_tx,
                cancel_token: cancel.clone(),
                max_request_body_bytes: DEFAULT_MAX_REQUEST_BODY_BYTES,
                pending_principal: None,
                accept_outstanding: false,
            },
        );

        // 6. Create linker + register host functions
        let mut linker: Linker<SourceHostState> = Linker::new(&engine);
        add_to_linker(&mut linker)
            .map_err(|e| CamelError::ProcessorError(format!("linker: {e}")))?;

        // 7. Set a large epoch deadline before any guest code runs. With
        // epoch_interruption(true) the store's default deadline is 0 (already
        // expired); instantiate_async and configure both run guest WASM, so
        // without this they trap immediately with "wasm trap: interrupt". No
        // epoch ticker runs during setup, so a u64::MAX budget disables
        // interruption for both calls. The run task later resets this to 1
        // (a pure stop tripwire). Mirrors the plugin path (runtime.rs:175).
        store.set_epoch_deadline(CONFIGURE_EPOCH_DEADLINE);

        // 8. Instantiate component (async bindings — see source_bindings.rs).
        let source = Source::instantiate_async(&mut store, &component, &linker)
            .await
            .map_err(|e| CamelError::ProcessorError(format!("instantiate: {e}")))?;

        // 9. Call guest configure() → SourcePlan
        //
        // configure() is a SYNC export (only run/accept-http/submit-exchange
        // are async), but concurrency_support(true) forbids the sync
        // `TypedFunc::call` path — so dispatch it via `call_async` on the
        // underlying typed func. No async imports are needed during configure,
        // so this stays out of run_concurrent. The large deadline set above
        // covers it.
        let (plan_result,) = source
            .func_configure()
            .call_async(&mut store, (self.guest_config.as_slice(),))
            .await
            .map_err(|e| CamelError::ProcessorError(format!("configure: {e}")))?;
        let plan = plan_result
            .map_err(|e| CamelError::ProcessorError(format!("configure guest error: {e:?}")))?;

        // 10. Validate source-plan
        if plan.capabilities.len() != 1 {
            return Err(CamelError::EndpointCreationFailed(
                "source-plan must have exactly one capability".into(),
            ));
        }
        let CapabilityRequest::HttpListener(listener_spec) = &plan.capabilities[0];

        // 10b. Reject unsupported concurrency. The host drains the guest
        // strictly sequentially through capacity-1 channels, so silently
        // degrading `concurrent(N)` to sequential would violate the guest's
        // declared contract. Reject explicitly until concurrent mode exists.
        use crate::source_bindings::camel::plugin::source_host::ConcurrencyModel as PlanConcurrency;
        match plan.concurrency {
            PlanConcurrency::Sequential => {}
            PlanConcurrency::Concurrent(max) => {
                return Err(CamelError::EndpointCreationFailed(format!(
                    "WASM source does not support concurrent({max}); only sequential is implemented"
                )));
            }
        }

        // 11. Create http-listener resource handle in the table
        let listener = store
            .data_mut()
            .table
            .push(HttpListenerHandle)
            .map_err(|e| CamelError::ProcessorError(format!("resource table: {e}")))?;

        // 12. Parse bind address from listener spec
        let bind_addr: std::net::SocketAddr = listener_spec.bind.parse().map_err(|e| {
            CamelError::EndpointCreationFailed(format!(
                "invalid bind address '{}': {}",
                listener_spec.bind, e
            ))
        })?;
        let path_filter = listener_spec.path.clone();

        // 12b. Resolve the effective bind (Task 1.4): the operator `bind`
        // (guest-config entry captured by `parse_guest_config` from the URI
        // query) must agree with the guest's declared listener bind. Equal
        // after SocketAddr parse-and-compare → operator bind (canonical
        // form); absent → guest bind; different → refuse to start, naming
        // the route URI, the operator bind, and the guest bind. Never bind
        // silently on either side of a conflict.
        let operator_bind = self
            .guest_config
            .iter()
            .find(|(key, _)| key == "bind")
            .map(|(_, value)| value.clone());
        let bind_addr = if let Some(operator_bind) = operator_bind.as_deref() {
            // An unparseable operator bind can never equal the (already
            // parsed) guest bind — it falls into the same conflict refusal.
            match operator_bind.parse::<std::net::SocketAddr>() {
                Ok(operator_addr) if operator_addr == bind_addr => bind_addr,
                _ => {
                    return Err(CamelError::RouteError(format!(
                        "wasm source '{}': operator bind '{operator_bind}' conflicts with \
                         guest-declared bind '{}'",
                        self.endpoint_uri, listener_spec.bind
                    )));
                }
            }
        } else {
            bind_addr
        };

        // 12c. ADR-0061 per-bind exposure gate, immediately before the
        // socket is bound — a refused route never spawns a listener. The
        // gate plan is the EXACT retained classification: the kernel plan
        // (plan + providers captured), else the plan-only classification,
        // else the Public default `compile_route_security_plan` builds for
        // undeclared routes (a non-Public classification without wiring
        // yields NO Public routes, so the gate passes — the fail-closed
        // denial happens per-request, Task 2.2).
        let gate_plan = if let Some(kernel) = self.kernel.as_ref() {
            kernel.plan.clone()
        } else if let Some(plan) = self.plan.as_ref() {
            plan.clone()
        } else {
            RouteSecurityPlan {
                access_mode: AccessMode::Public,
                provider_ref: None,
                transport: TransportId::Wasm,
                credential_sources: Vec::new(),
                audience_binding: None,
            }
        };
        let bind_key = bind_addr.to_string();
        camel_auth::enforce_bind_exposure_gate(
            &bind_key,
            bind_addr.ip().is_loopback(),
            &[(self.endpoint_uri.as_str(), &gate_plan)],
            crate::WasmSourceBindAcks::global().acknowledged(&bind_key),
        )?;

        // 13. Bind the TCP listener synchronously BEFORE spawning, so a bind
        // failure (port in use, permission denied) surfaces as a start()
        // error rather than a background warning. Without this the guest
        // exits cleanly and the route looks healthy with nothing accepting.
        //
        // Staged consumption first (wasm-bound-address): a test-tier helper
        // may have parked a pre-bound listener under this exact address;
        // taking it is one-shot. Ordering is load-bearing — this runs AFTER
        // the 12b agreement and the 12c exposure gate, so a refused route
        // never consumes a staged slot. `take`'s Err (same-port
        // different-host conflict) is already `EndpointCreationFailed` and
        // propagates via `?`.
        let tcp_listener = match staged_listener::take(bind_addr)? {
            Some(listener) => listener,
            None => tokio::net::TcpListener::bind(bind_addr)
                .await
                .map_err(|e| {
                    CamelError::Io(format!(
                        "failed to bind source HTTP listener {bind_addr}: {e}"
                    ))
                })?,
        };
        tracing::info!(%bind_addr, "source HTTP listener bound");

        // 14. Spawn HTTP listener task (serve on the already-bound listener).
        // Task 2.2: clone the kernel + derived classification out of the
        // consumer before the task spawns — the listener owns the host-edge
        // handshake (deny 401 before the request channel is touched).
        let kernel = self.kernel.clone();
        let plan_access = self.plan.as_ref().map(|p| p.access_mode.clone());
        let listener_cancel = cancel.clone();
        let request_tx = channels.request_tx;
        let lt = tokio::spawn(async move {
            if let Err(e) = run_http_listener(
                tcp_listener,
                path_filter,
                request_tx,
                listener_cancel,
                kernel,
                plan_access,
            )
            .await
            {
                warn!("WASM source HTTP listener exited: {e:?}");
            }
        });

        // 14. Spawn pipeline bridge task
        let bridge_ctx = ctx;
        let exchange_rx = channels.exchange_rx;
        let bt = tokio::spawn(async move {
            if let Err(e) = run_pipeline_bridge(exchange_rx, bridge_ctx).await {
                warn!("WASM source pipeline bridge exited: {e:?}");
            }
        });

        // 15. Spawn the async guest run() loop on a tokio task — THE CRITICAL STEP.
        //
        // The guest's `run` export is async; its host imports (accept-http,
        // submit-exchange) are async too. The whole thing is driven under
        // `Store::run_concurrent`, which is what makes the accessor-based
        // async imports dispatchable (concurrency_support(true) is required).
        //
        // Epoch: deadline(1) is a pure stop tripwire. No epoch ticker runs
        // during normal operation (each source consumer owns a private engine),
        // so the deadline never advances until stop()'s single
        // engine.increment_epoch(). That one increment traps a guest stuck in
        // a CPU-bound loop; a guest parked in an import is unblocked first by
        // the import's cancel_token select!.
        //
        // No watchdog: the source `run()` is an unbounded loop that
        // legitimately parks in `accept-http` for arbitrary durations between
        // requests. A fixed no-progress timeout would kill idle webhook
        // sources. The existing safeguards cover all failure modes:
        //   - Epoch interruption (deadline(1) + increment_epoch from stop())
        //     catches CPU-bound guest loops.
        //   - Cancel-token select! inside each import unblocks park-on-stop.
        //   - Route-level supervision catches "guest hangs forever" at the
        //     route manager level.
        let run_handle = tokio::spawn(async move {
            store.set_epoch_deadline(1);

            // 3-layer result: run_concurrent Result< trappable Result<
            // WIT Result > >. peel_concurrent flattens the outer two
            // wasmtime::Error layers.
            let result = crate::error::peel_concurrent(
                store
                    .as_context_mut()
                    .run_concurrent(async |accessor| source.call_run(accessor, listener).await)
                    .await,
                |e| CamelError::ProcessorError(format!("guest trap: {e}")),
                |e| CamelError::ProcessorError(format!("guest trap: {e}")),
            );

            match result {
                Ok(Ok(())) => {
                    debug!("WASM source guest run() exited normally");
                    Ok(())
                }
                Ok(Err(e)) => {
                    warn!("WASM source guest run() returned error: {e:?}");
                    Err(CamelError::ProcessorError(format!("guest run: {e:?}")))
                }
                Err(e) => {
                    warn!("WASM source guest trapped: {e}");
                    Err(e)
                }
            }
        });

        // 16. Store engine + JoinHandles
        self.engine = Some(engine);
        self.run_task = Some(run_handle);
        self.listener_task = Some(lt);
        self.bridge_task = Some(bt);

        Ok(())
    }

    async fn stop(&mut self) -> Result<(), CamelError> {
        // Idempotent — safe on any exit path
        self.cancel_token.cancel();

        // Epoch tick — interrupts CPU-bound guest loops
        if let Some(ref engine) = self.engine {
            engine.increment_epoch();
        }

        // The Runtime owns and monitors run_task via background_task_handle().
        // If the Runtime has already taken the handle, run_task is None here
        // and the Runtime is responsible for joining it. If we still own it
        // (stop() called without background_task_handle, or a test path),
        // join with a grace timeout and abort on timeout — dropping a
        // JoinHandle only detaches the task (keeps it running forever).
        let grace = std::time::Duration::from_secs(5);
        if let Some(task) = self.run_task.take() {
            join_or_abort(task, "source run", grace).await;
        }

        // Join listener + bridge tasks. Graceful shutdown via the cancel token
        // is raced against a timeout; on timeout we abort() the task.
        if let Some(task) = self.listener_task.take() {
            join_or_abort(task, "listener", grace).await;
        }
        if let Some(task) = self.bridge_task.take() {
            join_or_abort(task, "bridge", grace).await;
        }

        Ok(())
    }

    fn concurrency_model(&self) -> ConcurrencyModel {
        ConcurrencyModel::Sequential
    }

    /// Capture the route's security classification before `start()`.
    ///
    /// The whole compiled plan is retained regardless of provider presence
    /// (plan-only contexts carry classification without wiring). The kernel
    /// state additionally requires providers: without them no principal can
    /// be minted, so the handshake stays unwired and non-Public requests
    /// fail closed at the host edge (Task 2.2). Credential-source capability
    /// for `wasm:` is enforced centrally at plan compilation
    /// (`credential_source_allowed`) — not re-validated here.
    fn set_security_context(&mut self, ctx: SecurityContext) {
        self.plan = ctx.plan.clone();
        self.kernel = WasmSourceKernelAuth::from_security_context(&ctx).map(Arc::new);
    }

    fn background_task_handle(&mut self) -> Option<JoinHandle<Result<(), CamelError>>> {
        self.run_task.take()
    }
}

/// Race joining `task` against a `grace` timeout; on timeout, `abort()` it
/// and await the abort so the task is gone before returning.
///
/// Dropping a [`JoinHandle`] only *detaches* the task — it keeps running in
/// the background. Only [`JoinHandle::abort`] actually cancels it, so this is
/// what makes `stop()` deterministic: once it returns, the task is not still
/// running.
async fn join_or_abort<T: Send + 'static>(task: JoinHandle<T>, label: &str, grace: Duration) {
    tokio::pin!(task);
    tokio::select! {
        result = &mut task => {
            if let Err(e) = result {
                warn!("{label} task panicked on shutdown: {e}");
            }
        }
        _ = tokio::time::sleep(grace) => {
            warn!("{label} task did not exit within {grace:?}, aborting");
            task.abort();
            // Await the cancellation so the task is truly gone.
            if let Err(e) = task.await {
                debug!("{label} task aborted: {e}");
            }
        }
    }
}

/// Extract guest config (bind, path, etc.) from URI query params.
/// WasmConfig::from_uri drops unknown params, so we parse them separately.
pub fn parse_guest_config(uri: &str) -> Vec<(String, String)> {
    let query = match uri.find('?') {
        Some(i) => &uri[i + 1..],
        None => return Vec::new(),
    };

    query
        .split('&')
        .filter_map(|pair| {
            let (k, v) = pair.split_once('=')?;
            // Capture all params the guest might need
            if matches!(k, "bind" | "path" | "method") {
                Some((k.to_string(), v.to_string()))
            } else {
                None
            }
        })
        .collect()
}

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

    #[tokio::test]
    async fn stop_does_not_wait_for_runtime_owned_run_task() {
        let config = WasmConfig {
            timeout_secs: 1,
            ..WasmConfig::default()
        };
        let mut consumer = WasmSourceConsumer::new(
            PathBuf::from("unused.wasm"),
            "wasm:unused.wasm",
            config,
            Vec::new(),
            Arc::new(camel_component_api::NoOpComponentContext),
        );
        // Mirrors the real run task shape (tokio::spawn of an async future).
        consumer.run_task = Some(tokio::spawn(async {
            tokio::time::sleep(std::time::Duration::from_secs(2)).await;
            Ok(())
        }));

        // Simulate the Runtime taking ownership of the run task handle.
        // Once taken, stop() no longer sees run_task and must not wait for it.
        let _runtime_owned = consumer.background_task_handle();

        tokio::time::timeout(std::time::Duration::from_millis(100), consumer.stop())
            .await
            .expect("stop() must not wait for a runtime-owned run_task")
            .expect("stop() should succeed");
    }

    /// When stop() still owns the run task (Runtime hasn't taken it), it must
    /// abort it after the grace timeout rather than detaching (leaking) it.
    #[tokio::test]
    async fn stop_aborts_owned_run_task_after_grace() {
        let config = WasmConfig {
            timeout_secs: 1,
            ..WasmConfig::default()
        };
        let mut consumer = WasmSourceConsumer::new(
            PathBuf::from("unused.wasm"),
            "wasm:unused.wasm",
            config,
            Vec::new(),
            Arc::new(camel_component_api::NoOpComponentContext),
        );
        // A run task that never exits on its own — only abort will stop it.
        consumer.run_task = Some(tokio::spawn(async {
            tokio::time::sleep(std::time::Duration::from_secs(60)).await;
            Ok(())
        }));

        // stop() should abort the stuck task within the grace window (5s) +
        // margin. Asserting 15s ceiling; the abort should fire at ~5s.
        let result =
            tokio::time::timeout(std::time::Duration::from_secs(15), consumer.stop()).await;
        assert!(
            result.is_ok(),
            "stop() must abort a stuck run task within the grace window, not hang"
        );
        result.unwrap().expect("stop() should succeed");
    }
}