rivet_core/

application.rs

use std::any::Any;
use std::collections::BTreeMap;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::{Arc, OnceLock};

use rivet_foundation::{ConfigValue, ContainerResolveExt, FromConfigValue};
use rivet_http::{Method, Request, Response};
use rivet_logger::handlers::build_handler_from_config;
use rivet_logger::{
    init_default_tracing, set_channel_handler_build_config, set_channel_handlers, set_handler, Log,
};
use rivet_routing::Match;
use tracing::{debug, warn};

use crate::dispatch::Dispatcher;
use crate::error::DispatchError;
use crate::middleware::Middleware;
use crate::state::Runtime;

/// Built rivet runtime that implements the dispatch contract.
pub struct Application {
    state: Runtime,
}

static CURRENT_APP: OnceLock<Arc<dyn Dispatcher>> = OnceLock::new();

impl Application {
    pub(crate) fn from_state(state: Runtime) -> Self {
        Self { state }
    }

    /// Number of routes registered in the runtime.
    pub fn route_count(&self) -> usize {
        self.state.route_count()
    }

    /// Number of middleware instances in the runtime chain.
    pub fn middleware_count(&self) -> usize {
        self.state.middleware_count()
    }

    /// Resolve a typed service instance from the application container.
    pub fn make<T: Any + Send + Sync + 'static>(
        &self,
    ) -> Result<Arc<T>, crate::error::RivetError> {
        self.state.container().resolve::<T>().map_err(Into::into)
    }

    /// Read a configuration value by dot-notation key.
    pub fn config(&self, key: &str) -> Option<ConfigValue> {
        self.state.config().get(key)
    }

    /// Read and convert a configuration value by dot-notation key.
    pub fn config_typed<T: FromConfigValue>(&self, key: &str) -> Option<T> {
        self.config(key).and_then(T::from_config_value)
    }

    /// Initialize rivet logging runtime and configured handlers.
    pub fn init_logging(&self) {
        init_default_tracing();
        self.init_log_handler();
    }

    /// Log that the application is starting to serve requests.
    pub fn log_serving(&self, addr: SocketAddr) {
        self.init_logging();
        Log::channel("stdout").info(format!("Listening on http://{addr}"));
    }

    fn init_log_handler(&self) {
        let Some(log_config) = self.config_typed::<ConfigValue>("log") else {
            return;
        };

        let base_path = self
            .config_typed::<String>("app.base_path")
            .map(PathBuf::from)
            .unwrap_or_else(|| PathBuf::from("."));

        match build_handler_from_config(&log_config, &base_path) {
            Ok(handler) => set_handler(handler),
            Err(err) => warn!(
                target: "rivet",
                error = %err,
                "failed to initialize configured log handler"
            ),
        }

        set_channel_handlers(BTreeMap::new());
        set_channel_handler_build_config(log_config, base_path);
    }
}

pub fn set_current(app: Arc<dyn Dispatcher>) -> Result<(), crate::error::RivetError> {
    CURRENT_APP.set(app).map_err(|_| {
        crate::error::RivetError::Build("application global is already initialized".to_string())
    })
}

pub fn app() -> Arc<dyn Dispatcher> {
    CURRENT_APP
        .get()
        .cloned()
        .expect("application global not initialized")
}

pub fn try_app() -> Option<Arc<dyn Dispatcher>> {
    CURRENT_APP.get().cloned()
}

pub fn config<T: FromConfigValue>(key: &str) -> Option<T> {
    app().config_value(key).and_then(T::from_config_value)
}

impl Dispatcher for Application {
    fn dispatch(&self, req: Request) -> Response {
        let _ = self.state.container();
        let _ = self.state.config();
        let method = req.method.clone();
        let route = req.path.clone();

        debug!(
            method = ?req.method,
            path = %req.path,
            middleware_count = self.state.middleware_count(),
            "dispatch started"
        );

        let terminal = self.route_terminal(&req);

        let response = Self::run_middleware_chain(self.state.middleware(), req, terminal.as_ref());
        Log::info(format!(
            "{} {} {}",
            method_label(&method),
            route,
            response.status
        ));
        debug!(status = response.status, "dispatch completed");
        response
    }

    fn config_value(&self, key: &str) -> Option<ConfigValue> {
        self.config(key)
    }

    fn resolve_any(
        &self,
        type_name: &'static str,
    ) -> Result<Arc<dyn Any + Send + Sync>, crate::error::RivetError> {
        self.state
            .container()
            .resolve_any(type_name)
            .map_err(Into::into)
    }
}

fn method_label(method: &Method) -> &'static str {
    match method {
        Method::Get => "GET",
        Method::Post => "POST",
        Method::Put => "PUT",
        Method::Patch => "PATCH",
        Method::Delete => "DELETE",
        Method::Head => "HEAD",
        Method::Options => "OPTIONS",
    }
}

impl Application {
    fn route_terminal(
        &self,
        req: &Request,
    ) -> Box<dyn Fn(Request) -> Response + Send + Sync + 'static> {
        match self.state.routes().match_request(&req.method, &req.path) {
            Match::Matched {
                route,
                head_fallback,
            } => {
                debug!(
                    route_method = ?route.method,
                    route_path = %route.path,
                    head_fallback,
                    "route matched"
                );

                let route = route.clone();
                Box::new(move |req: Request| {
                    let mut response = route.invoke(req);
                    if head_fallback {
                        response.body.clear();
                    }
                    response
                })
            }
            Match::MethodNotAllowed { allow } => {
                warn!(allow = ?allow, "route method not allowed");
                let err = DispatchError::MethodNotAllowed { allow };
                Box::new(move |_req: Request| {
                    warn!(error = %err, "dispatch terminal returned error");
                    err.clone().into_response()
                })
            }
            Match::NotFound => {
                let err = DispatchError::NotFound;
                debug!("route not found");
                Box::new(move |_req: Request| {
                    warn!(error = %err, "dispatch terminal returned error");
                    err.clone().into_response()
                })
            }
        }
    }

    fn run_middleware_chain(
        middlewares: &[Arc<dyn Middleware>],
        req: Request,
        terminal: &(dyn Fn(Request) -> Response + Send + Sync),
    ) -> Response {
        debug!(
            middleware_count = middlewares.len(),
            "executing middleware chain"
        );
        Self::run_middleware_at(0, middlewares, req, terminal)
    }

    fn run_middleware_at(
        index: usize,
        middlewares: &[Arc<dyn Middleware>],
        req: Request,
        terminal: &(dyn Fn(Request) -> Response + Send + Sync),
    ) -> Response {
        if index >= middlewares.len() {
            debug!("reached terminal middleware boundary");
            return terminal(req);
        }

        let middleware = Arc::clone(&middlewares[index]);
        debug!(middleware_index = index, "running middleware");
        let next =
            |next_req: Request| Self::run_middleware_at(index + 1, middlewares, next_req, terminal);

        middleware.handle(req, &next)
    }
}