viceroy_lib/

execute.rs

//! Guest code execution.

use {
    crate::{
        Error,
        acl::Acls,
        adapt,
        body::Body,
        body_tee::tee,
        cache::Cache,
        component as compute,
        config::{
            Backends, DeviceDetection, Dictionaries, ExperimentalModule, FakeValidFastlyKeys,
            Geolocation, UnknownImportBehavior,
        },
        downstream::{DownstreamMetadata, DownstreamRequest, DownstreamResponse, prepare_request},
        error::{ExecutionError, NonHttpResponse},
        handoff::{HandoffConfig, HandoffRequestInfo, HandoffTlsConfig, perform_handoff},
        http::framing::apply_response_framing,
        linking::{ComponentCtx, WasmCtx, create_store, link_host_functions},
        object_store::ObjectStores,
        sandbox::Sandbox,
        secret_store::SecretStores,
        shielding_site::ShieldingSites,
        upstream::TlsConfig,
    },
    futures::{
        Future,
        task::{Context, Poll},
    },
    http::StatusCode,
    hyper::{Request, Response},
    pin_project::pin_project,
    std::{
        collections::HashSet,
        fmt, fs,
        io::Write,
        net::{Ipv4Addr, SocketAddr},
        path::{Path, PathBuf},
        pin::Pin,
        sync::{
            Arc, Mutex,
            atomic::{AtomicBool, AtomicU64, Ordering},
        },
        thread::{self, JoinHandle},
        time::{Duration, Instant, SystemTime},
    },
    tokio::sync::Mutex as AsyncMutex,
    tokio::sync::mpsc,
    tokio::sync::oneshot::{self, Sender},
    tracing::{Instrument, Level, error, event, info, info_span},
    wasmtime::{
        Engine, GuestProfiler, InstancePre, Linker, Module, ProfilingStrategy,
        component::{self, Component},
    },
    wasmtime_wasi::I32Exit,
};

pub use wasmtime::WasmFeatures;

pub const DEFAULT_EPOCH_INTERRUPTION_PERIOD: Duration = Duration::from_micros(50);

const NEXT_REQ_PENDING_MAX: usize = 5;
const REGION_NONE: &str = "none";

enum Instance {
    Module(Module, InstancePre<WasmCtx>),
    Component(
        Component,
        compute::bindings::AdapterServicePre<ComponentCtx>,
    ),
}

impl Instance {
    fn unwrap_module(&self) -> (&Module, &InstancePre<WasmCtx>) {
        match self {
            Instance::Module(m, i) => (m, i),
            Instance::Component(_, _) => panic!("unwrap_module called on a component"),
        }
    }
}

/// Configuration for profiling Wasm execution.
///
/// Viceroy supports two types of profiling:
/// - **Guest profiling**: Profiles the WebAssembly guest code itself, producing Firefox profiler compatible JSON.
/// - **Native profiling**: Profiles wasmtime's JIT compiler for debugging across guest, viceroy and wasmtime.
#[derive(Clone)]
pub enum ProfilingConfig {
    /// No profiling enabled
    None,
    /// Profile the WebAssembly guest code
    Guest {
        /// Path to write profiling results. In serve mode, this should be a directory
        /// where per-request profiles are written. In run mode, this is a single file.
        path: PathBuf,
        /// Period at which the guest should be sampled (default: 50μs)
        sample_period: Duration,
    },
    /// Profile wasmtime's JIT compiler using native profiling tools
    Native(wasmtime::ProfilingStrategy),
}

impl ProfilingConfig {
    /// Get the native profiling strategy for wasmtime
    pub fn native_strategy(&self) -> wasmtime::ProfilingStrategy {
        match self {
            ProfilingConfig::Native(strategy) => *strategy,
            _ => wasmtime::ProfilingStrategy::None,
        }
    }

    /// Get the guest profile configuration if guest profiling is enabled
    fn guest_config(&self) -> Option<GuestProfileConfig> {
        match self {
            ProfilingConfig::Guest {
                path,
                sample_period,
            } => Some(GuestProfileConfig {
                path: path.clone(),
                sample_period: *sample_period,
            }),
            _ => None,
        }
    }
}

// Keep GuestProfileConfig internal for now to maintain backwards compatibility
#[derive(Clone)]
struct GuestProfileConfig {
    path: PathBuf,
    sample_period: Duration,
}

pub struct NextRequest(Option<(Box<DownstreamRequest>, Arc<ExecuteCtx>)>);

impl NextRequest {
    pub fn into_request(mut self) -> Option<DownstreamRequest> {
        self.0.take().map(|(r, _)| *r)
    }
}

impl fmt::Debug for NextRequest {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let debug = self.0.as_ref().map(|(r, _)| r);
        f.debug_tuple("NextRequest")
            .field(&debug)
            .finish_non_exhaustive()
    }
}

impl Drop for NextRequest {
    fn drop(&mut self) {
        let Some((req, ctx)) = self.0.take() else {
            return;
        };

        ctx.retry_request(*req);
    }
}

/// Execution context used by a [`ViceroyService`](struct.ViceroyService.html).
///
/// This is all of the state needed to instantiate a module, in order to respond to an HTTP
/// request. Note that it is very important that `ExecuteCtx` be cheaply clonable, as it is cloned
/// every time that a viceroy service handles an incoming connection.
pub struct ExecuteCtx {
    /// A reference to the global context for Wasm compilation.
    engine: Engine,
    /// An almost-linked Instance: each import function is linked, just needs a Store
    instance_pre: Arc<Instance>,
    /// Preloaded TLS certificates and configuration
    tls_config: TlsConfig,
    /// The ID to assign the next incoming request
    next_req_id: Arc<AtomicU64>,
    /// The cache for this service.
    cache: Arc<Cache>,
    /// Senders waiting for new requests for reusable sandboxes.
    pending_reuse: Arc<AsyncMutex<Vec<Sender<NextRequest>>>>,
    epoch_increment_thread: Option<JoinHandle<()>>,
    // `Arc` so that it can be tracked both by this context and `epoch_increment_thread`.
    epoch_increment_stop: Arc<AtomicBool>,
    /// Configuration for guest profiling if enabled
    guest_profile_config: Option<Arc<GuestProfileConfig>>,
    /// The settings supplied through [`ExecuteCtxBuilder`].
    settings: Settings,
}

/// The parts of an [`ExecuteCtx`] that can be configured through [`ExecuteCtxBuilder`].
///
/// These are held separately from the rest of the context so that the builder can collect them
/// before the Wasm engine has been created: building the engine is deferred to
/// [`ExecuteCtxBuilder::finish`], so that engine-affecting settings can be adjusted first.
struct Settings {
    /// The acls for this execution.
    acls: Acls,
    /// The backends for this execution.
    backends: Backends,
    /// The device detection mappings for this execution.
    device_detection: DeviceDetection,
    /// The geolocation mappings for this execution.
    geolocation: Geolocation,
    /// The dictionaries for this execution.
    dictionaries: Dictionaries,
    /// Path to the config, defaults to None
    config_path: Option<PathBuf>,
    /// Where to direct logging endpoint messages, defaults to stdout
    capture_logs: Arc<Mutex<dyn Write + Send>>,
    /// Whether to treat stdout as a logging endpoint
    log_stdout: bool,
    /// Whether to treat stderr as a logging endpoint
    log_stderr: bool,
    /// The local Pushpin proxy port
    local_pushpin_proxy_port: Option<u16>,
    /// Whether WebSocket passthrough is enabled for this service.
    ///
    /// Set this to `false` to simulate a service that does not have the WebSockets feature
    /// enabled, so that guests can exercise their handling of that state locally.
    enable_local_websocket_passthrough: bool,
    /// The ObjectStore associated with this instance of Viceroy
    object_store: ObjectStores,
    /// The secret stores for this execution.
    secret_stores: SecretStores,
    /// The shielding sites for this execution.
    shielding_sites: ShieldingSites,
    /// The valid mock Fastly API keys that should be considered valid.
    fake_valid_fastly_keys: FakeValidFastlyKeys,
}

impl Default for Settings {
    fn default() -> Self {
        Self {
            acls: Acls::new(),
            backends: Backends::default(),
            device_detection: DeviceDetection::default(),
            geolocation: Geolocation::default(),
            dictionaries: Dictionaries::default(),
            config_path: None,
            capture_logs: Arc::new(Mutex::new(std::io::stdout())),
            log_stdout: false,
            log_stderr: false,
            local_pushpin_proxy_port: None,
            enable_local_websocket_passthrough: true,
            object_store: ObjectStores::new(),
            secret_stores: SecretStores::new(),
            shielding_sites: ShieldingSites::new(),
            fake_valid_fastly_keys: FakeValidFastlyKeys::new(),
        }
    }
}

impl ExecuteCtx {
    /// Build a new execution context with unified profiling configuration.
    ///
    /// Note that the module is not read, compiled, or linked until
    /// [`ExecuteCtxBuilder::finish`] is called, so errors relating to the module itself are
    /// reported from there rather than here.
    pub fn build(
        module_path: impl AsRef<Path>,
        profiling: ProfilingConfig,
        wasi_modules: HashSet<ExperimentalModule>,
        unknown_import_behavior: UnknownImportBehavior,
        adapt_components: bool,
        wasm_features: WasmFeatures,
    ) -> Result<ExecuteCtxBuilder, Error> {
        Ok(ExecuteCtxBuilder {
            module_path: module_path.as_ref().to_owned(),
            profiling,
            wasi_modules,
            unknown_import_behavior,
            adapt_components,
            wasm_features,
            debug_info: false,
            settings: Settings::default(),
        })
    }

    /// Create a new execution context with unified profiling configuration.
    ///
    /// This is a convenience wrapper around `build().finish()`.
    pub fn new(
        module_path: impl AsRef<Path>,
        profiling: ProfilingConfig,
        wasi_modules: HashSet<ExperimentalModule>,
        unknown_import_behavior: UnknownImportBehavior,
        adapt_components: bool,
        wasm_features: WasmFeatures,
    ) -> Result<Arc<Self>, Error> {
        ExecuteCtx::build(
            module_path,
            profiling,
            wasi_modules,
            unknown_import_behavior,
            adapt_components,
            wasm_features,
        )?
        .finish()
    }

    /// Get the engine for this execution context.
    pub fn engine(&self) -> &Engine {
        &self.engine
    }

    /// Get the acls for this execution context.
    pub fn acls(&self) -> &Acls {
        &self.settings.acls
    }

    /// Get the backends for this execution context.
    pub fn backends(&self) -> &Backends {
        &self.settings.backends
    }

    /// Get the device detection mappings for this execution context.
    pub fn device_detection(&self) -> &DeviceDetection {
        &self.settings.device_detection
    }

    /// Get the geolocation mappings for this execution context.
    pub fn geolocation(&self) -> &Geolocation {
        &self.settings.geolocation
    }

    /// Get the dictionaries for this execution context.
    pub fn dictionaries(&self) -> &Dictionaries {
        &self.settings.dictionaries
    }

    /// Where to direct logging endpoint messages. Defaults to stdout.
    pub fn capture_logs(&self) -> Arc<Mutex<dyn Write + Send>> {
        self.settings.capture_logs.clone()
    }

    /// Whether to treat stdout as a logging endpoint.
    pub fn log_stdout(&self) -> bool {
        self.settings.log_stdout
    }

    /// Whether to treat stderr as a logging endpoint.
    pub fn log_stderr(&self) -> bool {
        self.settings.log_stderr
    }

    /// Gets the TLS configuration
    pub fn tls_config(&self) -> &TlsConfig {
        &self.tls_config
    }

    async fn maybe_receive_response(
        mut receiver: mpsc::Receiver<DownstreamResponse>,
    ) -> Option<(Response<Body>, Option<anyhow::Error>)> {
        loop {
            match receiver.recv().await? {
                DownstreamResponse::Http(mut resp) => {
                    apply_response_framing(&mut resp);

                    // Supporting 103 Early Hints responses is currently infeasible, as Hyper does
                    // not support sending multiple responses on a single connection. But we don't
                    // want to generate errors for them either. Early Hints will be dropped, but
                    // logged so that people will know they *did work*, even though they won't
                    // reach the client.
                    if resp.status().is_informational() {
                        // We'll do these at different log levels in case someone wants to squelch some.
                        tracing::warn!(
                            "Guest returned informational response ({}) which will not be sent to the client",
                            resp.status(),
                        );
                        tracing::info!("{:#?}", resp);
                        continue;
                    }

                    return Some((resp, None));
                }
                DownstreamResponse::Pending(pending) => {
                    let mut resp = pending
                        .recv_or_else(|e| {
                            let status = e.as_status_code();
                            let err = anyhow::Error::from(e);
                            anyhow_response_with_status(&err, status)
                        })
                        .await;

                    apply_response_framing(&mut resp);

                    return Some((resp, None));
                }
                DownstreamResponse::HandoffToPushpin(info) => {
                    let err = NonHttpResponse::HandoffToPushpin(info).into();
                    let resp = Response::new(Body::empty());
                    return Some((resp, Some(err)));
                }
                DownstreamResponse::HandoffToBackend(info) => {
                    let resp = Response::new(Body::empty());
                    let err = NonHttpResponse::HandoffToBackend(info).into();
                    return Some((resp, Some(err)));
                }
            }
        }
    }

    /// Asynchronously handle a request.
    ///
    /// This method fully instantiates the wasm module housed within the `ExecuteCtx`,
    /// including running the wasm start function. It then proceeds to execute the
    /// instantiated module's WASI entry point, running to completion. If execution
    /// results in an error, a response is still produced, but with a 500 status code.
    ///
    /// Build time: Before you build or test your code, we recommend to set the release flag
    /// e.g. `cargo test --release` otherwise the execution will be very slow. This has to do
    /// with the Cranelift compiler, which is extremely slow when compiled in debug mode.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use std::collections::HashSet;
    /// # use wasmtime::WasmFeatures;
    /// use hyper::{Body, http::Request};
    /// # use viceroy_lib::{Error, ExecuteCtx, ProfilingConfig, ViceroyService};
    /// # use viceroy_lib::config::UnknownImportBehavior;
    /// # async fn f() -> Result<(), Error> {
    /// # let req = Request::new(Body::from(""));
    /// let adapt_core_wasm = false;
    /// let wasm_features = WasmFeatures::default();
    /// let ctx = ExecuteCtx::new(
    ///     "path/to/a/file.wasm",
    ///     ProfilingConfig::None,
    ///     HashSet::new(),
    ///     UnknownImportBehavior::LinkError,
    ///     adapt_core_wasm,
    ///     wasm_features,
    /// )?;
    /// let local = "127.0.0.1:80".parse().unwrap();
    /// let remote = "127.0.0.1:0".parse().unwrap();
    /// let resp = ctx.handle_request(req, local, remote).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn handle_request(
        self: Arc<Self>,
        mut incoming_req: Request<hyper::Body>,
        local: SocketAddr,
        remote: SocketAddr,
    ) -> Result<(Response<Body>, Option<anyhow::Error>), Error> {
        let orig_req_on_upgrade = hyper::upgrade::on(&mut incoming_req);
        let (incoming_req_parts, incoming_req_body) = incoming_req.into_parts();
        let local_pushpin_proxy_port = self.settings.local_pushpin_proxy_port;

        let (body_for_wasm, orig_body_tee) = tee(incoming_req_body).await;
        let orig_request_info_for_pushpin = HandoffRequestInfo::from_parts(&incoming_req_parts);

        let original_headers = incoming_req_parts.headers.clone();
        let req = prepare_request(Request::from_parts(incoming_req_parts, body_for_wasm))?;

        let req_id = self
            .next_req_id
            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);

        let metadata = DownstreamMetadata {
            req_id,
            server_addr: local,
            client_addr: remote,
            compliance_region: String::from(REGION_NONE),
            original_headers,
        };

        let backends = self.settings.backends.clone();
        let tls_config = self.tls_config.clone();

        let (resp, mut err) = self.reuse_or_spawn_guest(req, metadata).await;

        let span = info_span!("request", id = req_id);
        let _span = span.enter();

        info!("response status: {:?}", resp.status());

        if let Some(e) = err {
            match e.downcast::<NonHttpResponse>() {
                Ok(NonHttpResponse::HandoffToPushpin(handoff_info)) => {
                    let backend_name = handoff_info.backend_name.clone();

                    info!("Pushpin handoff signaled to backend '{backend_name}'");

                    // Defensive only: `Sandbox::redirect_downstream_to_pushpin` is the sole
                    // producer of this handoff, and it now rejects the call with an
                    // `unsupported` error before signaling when no proxy port is configured.
                    let local_pushpin_proxy_port = match local_pushpin_proxy_port {
                        None => {
                            error!("Pushpin handoff signaled, but Pushpin mode not enabled.");
                            let mut resp = Response::new(Body::from(hyper::Body::from(
                                "Pushpin handoff signaled, but Pushpin mode not enabled.",
                            )));
                            *resp.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
                            return Ok((resp, None));
                        }
                        Some(port) => port,
                    };

                    let pushpin_addr =
                        SocketAddr::new(Ipv4Addr::LOCALHOST.into(), local_pushpin_proxy_port)
                            .to_string();

                    // To distinguish backends, the HTTP request header `pushpin-route: <backend_name>` is added
                    // to the request. Pushpin routes should be configured with `id=<backend_name>`.
                    let additional_headers =
                        vec![("pushpin-route".to_string(), backend_name.clone())];

                    let handoff_resp = perform_handoff(
                        handoff_info.request_info,
                        orig_request_info_for_pushpin,
                        orig_body_tee,
                        orig_req_on_upgrade,
                        HandoffConfig {
                            target_addr: pushpin_addr.clone(),
                            host_header: pushpin_addr, // Host header is the local proxy address
                            display_name: format!("Pushpin [{backend_name}]"),
                            path_prefix: None, // Path prefix is applied by Pushpin route
                            extra_headers: additional_headers,
                            tls_config: None, // Pushpin only runs locally, without https
                        },
                    )
                    .await;

                    let (p, hyper_body) = handoff_resp.into_parts();
                    return Ok((Response::from_parts(p, Body::from(hyper_body)), None));
                }
                Ok(NonHttpResponse::HandoffToBackend(handoff_info)) => {
                    let backend_name = handoff_info.backend_name.clone();

                    info!("Backend handoff signaled to backend '{backend_name}'");

                    let backend = backends.get(backend_name.as_str());
                    let backend = match backend {
                        None => {
                            error!("Backend handoff signaled to unknown backend '{backend_name}'.");
                            let mut resp = Response::new(Body::from(hyper::Body::from(format!(
                                "Backend handoff signaled to unknown backend '{backend_name}'."
                            ))));
                            *resp.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
                            return Ok((resp, None));
                        }
                        Some(backend) => backend,
                    };

                    let backend_uri = backend.uri.clone();
                    let tls_handoff_config =
                        if matches!(backend_uri.scheme_str(), Some("https" | "wss")) {
                            Some(HandoffTlsConfig {
                                ca_certs: backend.ca_certs.clone(),
                                client_cert: backend.client_cert.clone(),
                                use_sni: backend.use_sni,
                                cert_host: backend.cert_host.clone(),
                                dns_name_fallback: backend
                                    .uri
                                    .host()
                                    .unwrap_or_default()
                                    .to_string(),
                                is_grpc: backend.grpc,
                                base_tls_config: tls_config.clone(), // Pass the builder from self
                            })
                        } else {
                            None
                        };

                    let backend_host = backend_uri
                        .authority()
                        .map(|a| match (a.port(), backend_uri.scheme_str()) {
                            (None, Some("wss")) => format!("{}:443", a),
                            (None, Some("https")) => format!("{}:443", a),
                            (None, Some("ws")) => format!("{}:80", a),
                            (None, Some("http")) => format!("{}:80", a),
                            _ => a.to_string(), // already has a port, or scheme is unknown (unlikely)
                        })
                        .unwrap_or_default();
                    // Use override_host if present, otherwise fallback to the URI's authority
                    let host_header = backend
                        .override_host
                        .clone()
                        .map(|host| host.to_str()
                            .expect("`backend.override_host`, if provided, should be a valid header value")
                            .to_string())
                        .unwrap_or_else(|| backend_host.clone());

                    // Prepend a path if there's actually a path in the backend URI
                    let path_prefix = (!backend_uri.path().is_empty() && backend_uri.path() != "/")
                        .then(|| backend_uri.path().to_string());

                    let handoff_resp = perform_handoff(
                        handoff_info.request_info,
                        orig_request_info_for_pushpin,
                        orig_body_tee,
                        orig_req_on_upgrade,
                        HandoffConfig {
                            target_addr: backend_host,
                            host_header,
                            display_name: format!("Backend [{backend_name}]"),
                            path_prefix,
                            extra_headers: vec![], // Standard backends don't need extra headers
                            tls_config: tls_handoff_config,
                        },
                    )
                    .await;

                    let (p, hyper_body) = handoff_resp.into_parts();
                    return Ok((Response::from_parts(p, Body::from(hyper_body)), None));
                }
                Err(e) => {
                    err = Some(e);
                }
            }
        }

        Ok((resp, err))
    }

    /// Spawn a new guest to process a request whose processing was never attempted by
    /// a reused sandbox.
    pub(crate) fn retry_request(self: Arc<Self>, mut downstream: DownstreamRequest) {
        if downstream.sender.is_closed() {
            return;
        }

        tokio::task::spawn(async move {
            let (sender, receiver) = mpsc::channel(10);
            let original = std::mem::replace(&mut downstream.sender, sender);
            let (resp, err) = self.spawn_guest(downstream, receiver).await;
            let resp = guest_result_to_response(resp, err);
            let _ = original.send(DownstreamResponse::Http(resp)).await;
        });
    }

    pub async fn handle_request_with_runtime_error(
        self: Arc<Self>,
        incoming_req: Request<hyper::Body>,
        local: SocketAddr,
        remote: SocketAddr,
    ) -> Result<Response<Body>, Error> {
        let result = self.handle_request(incoming_req, local, remote).await?;
        let resp = guest_result_to_response(result.0, result.1);

        Ok(resp)
    }

    async fn reuse_or_spawn_guest(
        self: Arc<Self>,
        req: Request<Body>,
        metadata: DownstreamMetadata,
    ) -> (Response<Body>, Option<anyhow::Error>) {
        let (downstream, receiver) = DownstreamRequest::new(req, metadata);

        let mut next_req = NextRequest(Some((Box::new(downstream), self.clone())));
        let mut reusable = self.pending_reuse.lock().await;

        while let Some(pending) = reusable.pop() {
            match pending.send(next_req) {
                Ok(()) => {
                    // Drop lock and wait for the guest to process our request.
                    drop(reusable);

                    if let Some(response) = Self::maybe_receive_response(receiver).await {
                        return response;
                    }
                    return (Response::default(), None);
                }
                Err(nr) => next_req = nr,
            }
        }

        drop(reusable);

        let downstream = next_req
            .into_request()
            .expect("request should still be unprocessed");
        self.spawn_guest(downstream, receiver).await
    }

    async fn spawn_guest(
        self: Arc<Self>,
        downstream: DownstreamRequest,
        receiver: mpsc::Receiver<DownstreamResponse>,
    ) -> (Response<Body>, Option<anyhow::Error>) {
        let active_cpu_time_us = Arc::new(AtomicU64::new(0));

        // Spawn a separate task to run the guest code. That allows _this_ method to return a response early
        // if the guest sends one, while the guest continues to run afterward within its task.
        let req_id = downstream.metadata.req_id;
        let guest_handle = tokio::task::spawn(CpuTimeTracking::new(
            active_cpu_time_us.clone(),
            self.run_guest(downstream, active_cpu_time_us)
                .instrument(info_span!("request", id = req_id)),
        ));

        if let Some(response) = Self::maybe_receive_response(receiver).await {
            return response;
        }

        match guest_handle
            .await
            .expect("guest worker finished without panicking")
        {
            Ok(_) => (Response::new(Body::empty()), None),
            Err(ExecutionError::WasmTrap(e)) => {
                event!(
                    Level::ERROR,
                    "There was an error handling the request {}",
                    e.to_string()
                );
                (anyhow_response(&e), Some(e))
            }
            Err(e) => panic!("failed to run guest: {}", e),
        }
    }

    async fn run_guest(
        self: Arc<Self>,
        downstream: DownstreamRequest,
        active_cpu_time_us: Arc<AtomicU64>,
    ) -> Result<(), ExecutionError> {
        info!(
            "handling request {} {}",
            downstream.req.method(),
            downstream.req.uri()
        );
        let start_timestamp = Instant::now();
        let req_id = downstream.metadata.req_id;
        let sandbox = Sandbox::new(downstream, active_cpu_time_us, self.clone());

        let guest_profile_path = self.guest_profile_config.as_deref().map(|pcfg| {
            let now = SystemTime::now()
                .duration_since(SystemTime::UNIX_EPOCH)
                .unwrap()
                .as_secs();
            pcfg.path.join(format!("{}-{}.json", now, req_id))
        });

        match self.instance_pre.as_ref() {
            Instance::Component(component, instance_pre) => {
                let profiler = self.guest_profile_config.as_deref().map(|pcfg| {
                    let program_name = "main";
                    GuestProfiler::new_component(
                        program_name,
                        pcfg.sample_period,
                        component.clone(),
                        std::iter::empty(),
                    )
                });

                let req = sandbox.downstream_request();
                let body = sandbox.downstream_request_body();

                let mut store = ComponentCtx::create_store(&self, sandbox, profiler, |ctx| {
                    ctx.arg("compute-app");
                })
                .map_err(ExecutionError::Context)?;

                let compute = instance_pre
                    .instantiate_async(&mut store)
                    .await
                    .map_err(ExecutionError::Instantiation)?;

                let result = compute
                    .fastly_compute_http_incoming()
                    .call_handle(&mut store, req.into(), body.into())
                    .await;

                let outcome = match result {
                    Ok(Ok(())) => Ok(()),

                    Ok(Err(())) => {
                        event!(Level::ERROR, "WebAssembly exited with an error");
                        Err(ExecutionError::WasmTrap(anyhow::Error::msg("failed")))
                    }

                    Err(e) => {
                        if let Some(exit) = e.downcast_ref::<I32Exit>() {
                            if exit.0 == 0 {
                                Ok(())
                            } else {
                                event!(Level::ERROR, "WebAssembly exited with error: {:?}", e);
                                Err(ExecutionError::WasmTrap(e))
                            }
                        } else {
                            event!(Level::ERROR, "WebAssembly trapped: {:?}", e);
                            Err(ExecutionError::WasmTrap(e))
                        }
                    }
                };

                // If we collected a profile, write it to the file
                write_profile_component(&mut store, guest_profile_path.as_ref());

                // Ensure the downstream response channel is closed, whether or not a response was
                // sent during execution.
                let resp = outcome
                    .as_ref()
                    .err()
                    .map(exec_err_to_response)
                    .unwrap_or_default();
                store
                    .data_mut()
                    .sandbox
                    .close_downstream_response_sender(resp);

                let request_duration = Instant::now().duration_since(start_timestamp);

                info!(
                    "guest completed using {} of WebAssembly heap",
                    bytesize::ByteSize::b(store.data().limiter().memory_allocated as u64),
                );

                info!("guest completed in {:.0?}", request_duration);

                outcome
            }

            Instance::Module(module, instance_pre) => {
                let profiler = self.guest_profile_config.as_deref().map(|pcfg| {
                    let program_name = "main";
                    GuestProfiler::new(
                        program_name,
                        pcfg.sample_period,
                        vec![(program_name.to_string(), module.clone())],
                    )
                });

                // We currently have to postpone linking and instantiation to the guest task
                // due to wasmtime limitations, in particular the fact that `Instance` is not `Send`.
                // However, the fact that the module itself is created within `ExecuteCtx::new`
                // means that the heavy lifting happens only once.
                let mut store = create_store(&self, sandbox, profiler, |ctx| {
                    ctx.arg("compute-app");
                })
                .map_err(ExecutionError::Context)?;

                let instance = instance_pre
                    .instantiate_async(&mut store)
                    .await
                    .map_err(ExecutionError::Instantiation)?;

                // Pull out the `_start` function, which by convention with WASI is the main entry point for
                // an application.
                let main_func = instance
                    .get_typed_func::<(), ()>(&mut store, "_start")
                    .map_err(ExecutionError::Typechecking)?;

                // Invoke the entrypoint function, which may or may not send a downstream response.
                let outcome = match main_func.call_async(&mut store, ()).await {
                    Ok(_) => Ok(()),
                    Err(e) => {
                        if let Some(exit) = e.downcast_ref::<I32Exit>() {
                            if exit.0 == 0 {
                                Ok(())
                            } else {
                                event!(Level::ERROR, "WebAssembly exited with error: {:?}", e);
                                Err(ExecutionError::WasmTrap(e))
                            }
                        } else {
                            event!(Level::ERROR, "WebAssembly trapped: {:?}", e);
                            Err(ExecutionError::WasmTrap(e))
                        }
                    }
                };

                // If we collected a profile, write it to the file
                write_profile(&mut store, guest_profile_path.as_ref());

                // Ensure the downstream response channel is closed, whether or not a response was
                // sent during execution.
                let resp = outcome
                    .as_ref()
                    .err()
                    .map(exec_err_to_response)
                    .unwrap_or_default();
                store.data_mut().close_downstream_response_sender(resp);

                let request_duration = Instant::now().duration_since(start_timestamp);

                info!(
                    "request completed using {} of WebAssembly heap",
                    bytesize::ByteSize::b(store.data().limiter().memory_allocated as u64)
                );

                info!("request completed in {:.0?}", request_duration);

                outcome
            }
        }
    }

    pub async fn run_main(
        self: Arc<Self>,
        program_name: &str,
        args: &[String],
    ) -> Result<(), anyhow::Error> {
        // placeholders for request, result sender channel, and remote IP
        let req = Request::get("http://example.com/").body(Body::empty())?;
        let metadata = DownstreamMetadata {
            req_id: 0,
            server_addr: (Ipv4Addr::LOCALHOST, 80).into(),
            client_addr: (Ipv4Addr::LOCALHOST, 0).into(),
            compliance_region: String::from(REGION_NONE),
            original_headers: Default::default(),
        };
        let (downstream, receiver) = DownstreamRequest::new(req, metadata);
        let active_cpu_time_us = Arc::new(AtomicU64::new(0));

        let sandbox = Sandbox::new(downstream, active_cpu_time_us.clone(), self.clone());

        if let Instance::Component(_, _) = self.instance_pre.as_ref() {
            panic!("components not currently supported with `run`");
        }

        let (module, instance_pre) = self.instance_pre.unwrap_module();

        let profiler = self.guest_profile_config.as_deref().map(|pcfg| {
            GuestProfiler::new(
                program_name,
                pcfg.sample_period,
                vec![(program_name.to_string(), module.clone())],
            )
        });

        let mut store = create_store(&self, sandbox, profiler, |builder| {
            builder.arg(program_name);
            for arg in args {
                builder.arg(arg);
            }
        })
        .map_err(ExecutionError::Context)?;

        let instance = instance_pre
            .instantiate_async(&mut store)
            .await
            .map_err(ExecutionError::Instantiation)?;

        // Pull out the `_start` function, which by convention with WASI is the main entry point for
        // an application.
        let main_func = instance
            .get_typed_func::<(), ()>(&mut store, "_start")
            .map_err(ExecutionError::Typechecking)?;

        // Invoke the entrypoint function and collect its exit code
        let result =
            CpuTimeTracking::new(active_cpu_time_us, main_func.call_async(&mut store, ())).await;

        // If we collected a profile, write it to the file
        write_profile(
            &mut store,
            self.guest_profile_config.as_deref().map(|cfg| &cfg.path),
        );

        // Ensure the downstream response channel is closed, whether or not a response was
        // sent during execution.
        store
            .data_mut()
            .close_downstream_response_sender(Response::default());

        // We don't do anything with any response on the receiver, but
        // it's important to keep it alive until after the program has
        // finished.
        drop(receiver);

        result
    }

    pub fn cache(&self) -> &Arc<Cache> {
        &self.cache
    }

    pub fn config_path(&self) -> Option<&Path> {
        self.settings.config_path.as_deref()
    }

    pub fn object_store(&self) -> &ObjectStores {
        &self.settings.object_store
    }

    pub fn secret_stores(&self) -> &SecretStores {
        &self.settings.secret_stores
    }

    pub fn shielding_sites(&self) -> &ShieldingSites {
        &self.settings.shielding_sites
    }

    /// The local Pushpin proxy port, if Fanout is enabled for this execution context.
    pub fn local_pushpin_proxy_port(&self) -> Option<u16> {
        self.settings.local_pushpin_proxy_port
    }

    /// Whether WebSocket passthrough is enabled for this execution context.
    pub fn enable_local_websocket_passthrough(&self) -> bool {
        self.settings.enable_local_websocket_passthrough
    }

    /// Get the valid mock Fastly API keys for this execution context.
    pub fn fake_valid_fastly_keys(&self) -> &FakeValidFastlyKeys {
        &self.settings.fake_valid_fastly_keys
    }

    pub async fn register_pending_downstream(&self) -> Option<oneshot::Receiver<NextRequest>> {
        let mut pending = self.pending_reuse.lock().await;

        if pending.len() >= NEXT_REQ_PENDING_MAX {
            return None;
        }

        let (tx, rx) = oneshot::channel();
        pending.push(tx);

        Some(rx)
    }

    pub fn is_component(&self) -> bool {
        matches!(self.instance_pre.as_ref(), Instance::Component(_, _))
    }
}

pub struct ExecuteCtxBuilder {
    /// Path to the module to run.
    module_path: PathBuf,
    /// How execution of this module should be profiled, if at all.
    profiling: ProfilingConfig,
    /// The experimental wasi modules to link against.
    wasi_modules: HashSet<ExperimentalModule>,
    /// What to do with imports the host does not provide.
    unknown_import_behavior: UnknownImportBehavior,
    /// Whether to adapt core wasm modules to components before running them.
    adapt_components: bool,
    /// The wasm proposals to enable.
    wasm_features: WasmFeatures,
    /// Whether to compile the guest with native debug info, so that it can be inspected with a
    /// native debugger.
    debug_info: bool,
    /// The settings to hand to the finished context.
    settings: Settings,
}

impl ExecuteCtxBuilder {
    /// Read, compile, and link the module, producing an execution context ready to serve
    /// requests.
    pub fn finish(self) -> Result<Arc<ExecuteCtx>, Error> {
        let ExecuteCtxBuilder {
            module_path,
            profiling,
            wasi_modules,
            unknown_import_behavior,
            adapt_components,
            wasm_features,
            debug_info,
            settings,
        } = self;

        let guest_profile_config = profiling.guest_config();

        let input = fs::read(&module_path)?;

        let is_wat = module_path
            .extension()
            .map(|str| str == "wat")
            .unwrap_or(false);

        // When the input wasn't a component, but we're automatically adapting,
        // apply the component adapter.
        let is_component = adapt::is_component(&input);
        let (is_wat, is_component, input) = if !is_component && adapt_components {
            let input = if is_wat {
                let text = String::from_utf8(input)
                    .map_err(|_| anyhow::anyhow!("Failed to parse {}", module_path.display()))?;
                adapt::adapt_wat(&text)?
            } else {
                adapt::adapt_bytes(&input)?
            };

            (false, true, input)
        } else {
            (is_wat, is_component, input)
        };

        let config = &configure_wasmtime(wasm_features, profiling.native_strategy(), debug_info);
        let engine = Engine::new(config)?;
        let instance_pre = if is_component {
            let mut linker: component::Linker<ComponentCtx> = component::Linker::new(&engine);
            compute::link_host_functions(&mut linker)?;
            let component = if is_wat {
                Component::from_file(&engine, &module_path)?
            } else {
                Component::from_binary(&engine, &input)?
            };

            match unknown_import_behavior {
                UnknownImportBehavior::LinkError => (),
                UnknownImportBehavior::Trap => {
                    linker.define_unknown_imports_as_traps(&component)?
                }
            }

            let instance_pre = linker.instantiate_pre(&component)?;
            Instance::Component(
                component,
                compute::bindings::AdapterServicePre::new(instance_pre)?,
            )
        } else {
            let mut linker = Linker::new(&engine);
            link_host_functions(&mut linker, &wasi_modules)?;
            let module = if is_wat {
                Module::from_file(&engine, &module_path)?
            } else {
                Module::from_binary(&engine, &input)?
            };

            match unknown_import_behavior {
                UnknownImportBehavior::LinkError => (),
                UnknownImportBehavior::Trap => linker.define_unknown_imports_as_traps(&module)?,
            }

            let instance_pre = linker.instantiate_pre(&module)?;
            Instance::Module(module, instance_pre)
        };

        // Create the epoch-increment thread. Note that the period for epoch
        // interruptions is driven by the guest profiling sample period if
        // provided as guest stack sampling is done from the epoch
        // interruption callback.

        let epoch_increment_stop = Arc::new(AtomicBool::new(false));
        let engine_clone = engine.clone();
        let epoch_increment_stop_clone = epoch_increment_stop.clone();
        let sample_period = guest_profile_config
            .as_ref()
            .map(|c| c.sample_period)
            .unwrap_or(DEFAULT_EPOCH_INTERRUPTION_PERIOD);
        let epoch_increment_thread = Some(thread::spawn(move || {
            while !epoch_increment_stop_clone.load(Ordering::Relaxed) {
                thread::sleep(sample_period);
                engine_clone.increment_epoch();
            }
        }));

        Ok(Arc::new(ExecuteCtx {
            engine,
            instance_pre: Arc::new(instance_pre),
            tls_config: TlsConfig::new()?,
            next_req_id: Arc::new(AtomicU64::new(0)),
            cache: Arc::new(Cache::default()),
            pending_reuse: Arc::new(AsyncMutex::new(vec![])),
            epoch_increment_thread,
            epoch_increment_stop,
            guest_profile_config: guest_profile_config.map(Arc::new),
            settings,
        }))
    }

    /// Set whether the guest should be compiled with native debug info, so that it can be
    /// inspected with a native debugger.
    ///
    /// Defaults to `false`, as generating debug info both slows compilation down and requires the
    /// guest to have been built with debug info of its own.
    pub fn with_debug_info(mut self, debug_info: bool) -> Self {
        self.debug_info = debug_info;
        self
    }

    /// Set the acls for this execution context.
    pub fn with_acls(mut self, acls: Acls) -> Self {
        self.settings.acls = acls;
        self
    }

    /// Set the backends for this execution context.
    pub fn with_backends(mut self, backends: Backends) -> Self {
        self.settings.backends = backends;
        self
    }

    /// Set the device detection mappings for this execution context.
    pub fn with_device_detection(mut self, device_detection: DeviceDetection) -> Self {
        self.settings.device_detection = device_detection;
        self
    }

    /// Set the geolocation mappings for this execution context.
    pub fn with_geolocation(mut self, geolocation: Geolocation) -> Self {
        self.settings.geolocation = geolocation;
        self
    }

    /// Set the dictionaries for this execution context.
    pub fn with_dictionaries(mut self, dictionaries: Dictionaries) -> Self {
        self.settings.dictionaries = dictionaries;
        self
    }

    /// Set the object store for this execution context.
    pub fn with_object_stores(mut self, object_store: ObjectStores) -> Self {
        self.settings.object_store = object_store;
        self
    }

    /// Set the secret stores for this execution context.
    pub fn with_secret_stores(mut self, secret_stores: SecretStores) -> Self {
        self.settings.secret_stores = secret_stores;
        self
    }
    /// Set the shielding sites for this execution context.
    pub fn with_shielding_sites(mut self, shielding_sites: ShieldingSites) -> Self {
        self.settings.shielding_sites = shielding_sites;
        self
    }

    /// Set the valid mock Fastly API keys for this execution context.
    pub fn with_fake_valid_fastly_keys(
        mut self,
        fake_valid_fastly_keys: FakeValidFastlyKeys,
    ) -> Self {
        self.settings.fake_valid_fastly_keys = fake_valid_fastly_keys;
        self
    }

    /// Set the path to the config for this execution context.
    pub fn with_config_path(mut self, config_path: PathBuf) -> Self {
        self.settings.config_path = Some(config_path);
        self
    }

    /// Set where to direct logging endpoint messages for this execution
    /// context. Defaults to stdout.
    pub fn with_capture_logs(mut self, capture_logs: Arc<Mutex<dyn Write + Send>>) -> Self {
        self.settings.capture_logs = capture_logs;
        self
    }

    /// Set the stdout logging policy for this execution context.
    pub fn with_log_stdout(mut self, log_stdout: bool) -> Self {
        self.settings.log_stdout = log_stdout;
        self
    }

    /// Set the stderr logging policy for this execution context.
    pub fn with_log_stderr(mut self, log_stderr: bool) -> Self {
        self.settings.log_stderr = log_stderr;
        self
    }

    /// Set the local Pushpin proxy port
    pub fn with_local_pushpin_proxy_port(mut self, local_pushpin_proxy_port: Option<u16>) -> Self {
        self.settings.local_pushpin_proxy_port = local_pushpin_proxy_port;
        self
    }

    /// Set whether WebSocket passthrough is enabled for this execution context.
    pub fn with_enable_local_websocket_passthrough(
        mut self,
        enable_local_websocket_passthrough: bool,
    ) -> Self {
        self.settings.enable_local_websocket_passthrough = enable_local_websocket_passthrough;
        self
    }
}

fn write_profile_to_file(profile: Box<GuestProfiler>, path: &PathBuf) {
    match std::fs::File::create(path)
        .map_err(anyhow::Error::new)
        .and_then(|output| profile.finish(std::io::BufWriter::new(output)))
    {
        Err(e) => {
            event!(
                Level::ERROR,
                "failed writing profile at {}: {e:#}",
                path.display()
            );
        }
        _ => {
            event!(
                Level::INFO,
                "\nProfile written to: {}\nView this profile at https://profiler.firefox.com/.",
                path.display()
            );
        }
    }
}

fn write_profile(store: &mut wasmtime::Store<WasmCtx>, guest_profile_path: Option<&PathBuf>) {
    if let (Some(profile), Some(path)) =
        (store.data_mut().take_guest_profiler(), guest_profile_path)
    {
        write_profile_to_file(profile, path);
    }
}

fn write_profile_component(
    store: &mut wasmtime::Store<ComponentCtx>,
    guest_profile_path: Option<&PathBuf>,
) {
    if let (Some(profile), Some(path)) =
        (store.data_mut().take_guest_profiler(), guest_profile_path)
    {
        write_profile_to_file(profile, path);
    }
}

fn guest_result_to_response(resp: Response<Body>, err: Option<anyhow::Error>) -> Response<Body> {
    err.as_ref().map(anyhow_response).unwrap_or(resp)
}

fn exec_err_to_response(err: &ExecutionError) -> Response<Body> {
    if let ExecutionError::WasmTrap(e) = err {
        anyhow_response(e)
    } else {
        panic!("failed to run guest: {err}")
    }
}

fn anyhow_response(err: &anyhow::Error) -> Response<Body> {
    anyhow_response_with_status(err, hyper::StatusCode::INTERNAL_SERVER_ERROR)
}

fn anyhow_response_with_status(err: &anyhow::Error, status: hyper::StatusCode) -> Response<Body> {
    Response::builder()
        .status(status)
        .body(Body::from(format!("{err:?}").into_bytes()))
        .unwrap()
}

impl Drop for ExecuteCtx {
    fn drop(&mut self) {
        if let Some(join_handle) = self.epoch_increment_thread.take() {
            self.epoch_increment_stop.store(true, Ordering::Relaxed);
            join_handle.join().unwrap();
        }
    }
}

fn configure_wasmtime(
    wasm_features: WasmFeatures,
    profiling_strategy: ProfilingStrategy,
    debug_info: bool,
) -> wasmtime::Config {
    use wasmtime::{Config, InstanceAllocationStrategy, WasmBacktraceDetails};

    let mut config = Config::new();
    // Off by default: translating the guest's DWARF into native debug info costs compile time,
    // and is only useful when the guest was built with debug info to begin with.
    config.debug_info(debug_info);
    config.wasm_backtrace_details(WasmBacktraceDetails::Enable);
    config.async_support(true);
    config.epoch_interruption(true);
    config.profiler(profiling_strategy);

    config.allocation_strategy(InstanceAllocationStrategy::OnDemand);

    config.wasm_features(wasm_features, true);

    // Wasm permits the "relaxed" instructions to be nondeterministic
    // between runs, but requires them to be deterministic within runs.
    // Snapshotting a program's execution to avoid redundantly running
    // initialization code on each request is an important optimization,
    // so we enable deterministic lowerings for relaxed SIMD to ensure
    // that it works consistently even if the initialization runs on a
    // different host architecture.
    config.relaxed_simd_deterministic(true);

    // Enable the wide-arithmetic feature, which has now been voted
    // to phase-4.
    config.wasm_wide_arithmetic(true);

    config
}

#[pin_project]
struct CpuTimeTracking<F> {
    #[pin]
    future: F,
    time_spent: Arc<AtomicU64>,
}

impl<F> CpuTimeTracking<F> {
    fn new(time_spent: Arc<AtomicU64>, future: F) -> Self {
        CpuTimeTracking { future, time_spent }
    }
}

impl<E, F: Future<Output = Result<(), E>>> Future for CpuTimeTracking<F> {
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let me = self.project();

        let start = Instant::now();
        let result = me.future.poll(cx);
        // 2^64 microseconds is over half a million years, so I'm not terribly
        // worried about this cast.
        let runtime = start.elapsed().as_micros() as u64;
        let _ = me.time_spent.fetch_add(runtime, Ordering::SeqCst);
        result
    }
}