rok-core 0.6.0

Core primitives for the rok ecosystem — errors, crypto, i18n, config, DI, and more
Documentation
pub mod error_handler;
pub mod provider;

use std::sync::Arc;

use axum::routing::Router;

use crate::container::Container;
use crate::error::{RokError, RokException};

use self::provider::ServiceProvider;

/// The rok application kernel.
///
/// Bootstraps the application: loads config, registers providers,
/// mounts routes, applies middleware, and starts the HTTP server.
///
/// # Example
///
/// ```rust,ignore
/// use rok_core::app::App;
///
/// App::boot()
///     .register(OrmServiceProvider::new("DATABASE_URL"))
///     .register(AuthServiceProvider::new(...))
///     .middleware(tower_http::trace::TraceLayer::new_for_http())
///     .routes(app_router)
///     .serve("127.0.0.1:3000")
///     .await?;
/// ```
pub struct App {
    /// Dependency injection container.
    pub container: Container,
    /// Registered service providers.
    providers: Vec<Box<dyn ServiceProvider>>,
    /// Middleware layers to apply to the router (outermost first).
    middleware: Vec<Box<dyn FnOnce(Router) -> Router + Send>>,
    /// Error handler.
    error_handler: Option<Arc<dyn Fn(Box<dyn RokException>) -> axum::response::Response + Send + Sync>>,
    /// Router builder.
    router_fn: Option<Box<dyn FnOnce(&App) -> Router + Send>>,
    /// Extra routers merged into the main router (e.g. WsRouter via `.merge()`).
    extra_routers: Vec<Router>,
}

impl App {
    /// Create a new application instance, loading `.env`.
    pub fn boot() -> Self {
        let _ = dotenvy::dotenv();
        // Discover config/ directory
        #[cfg(feature = "config")]
        {
            let discovered = crate::config::discover_configs();
            if !discovered.is_empty() {
                tracing::info!("Discovered {} config file(s)", discovered.len());
                for (key, _path) in &discovered {
                    tracing::debug!("  config/{}.toml", key);
                }
            }
        }
        Self {
            container: Container::new(),
            providers: Vec::new(),
            middleware: Vec::new(),
            error_handler: None,
            router_fn: None,
            extra_routers: Vec::new(),
        }
    }

    /// Register a service provider (will be booted during [`serve`](Self::serve)).
    pub fn register<P: ServiceProvider + 'static>(mut self, provider: P) -> Self {
        self.providers.push(Box::new(provider));
        self
    }

    /// Apply a middleware layer to the router.
    ///
    /// Accepts a closure that receives the [`Router`] and returns a new one.
    /// This avoids complex generic bounds on the method signature.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// .middleware(|router| router.layer(tower_http::trace::TraceLayer::new_for_http()))
    /// ```
    pub fn middleware<F>(mut self, f: F) -> Self
    where
        F: FnOnce(Router) -> Router + Send + 'static,
    {
        self.middleware.push(Box::new(f));
        self
    }

    /// Set the global error handler for unhandled exceptions.
    pub fn with_error_handler<F>(mut self, handler: F) -> Self
    where
        F: Fn(Box<dyn RokException>) -> axum::response::Response + Send + Sync + 'static,
    {
        self.error_handler = Some(Arc::new(handler));
        self
    }

    /// Mount route definitions.
    pub fn routes(mut self, f: fn(&App) -> Router) -> Self {
        self.router_fn = Some(Box::new(f));
        self
    }

    /// Merge an additional Axum [`Router`] into the final router.
    ///
    /// Useful for merging `WsRouter::into_router()`, admin panels, or any
    /// pre-built sub-router into the application without modifying the
    /// `routes()` function.
    pub fn merge(mut self, router: Router) -> Self {
        self.extra_routers.push(router);
        self
    }

    /// Resolve a service from the container.
    pub fn resolve<T>(&self) -> Result<Arc<T>, crate::container::ContainerError>
    where
        T: Send + Sync + 'static,
    {
        self.container.make::<T>()
    }

    /// Load a config struct — tries `config/{key}.toml` first,
    /// then falls back to environment variables.
    ///
    /// Requires `config` feature (enabled by `app`).
    pub fn config<T: crate::config::Configurable + crate::config::FromEnv>() -> T {
        crate::config::load_config::<T>()
            .unwrap_or_else(|| crate::config::Config::load::<T>())
    }

    /// Boot providers, build the router, and return it (no HTTP server started).
    ///
    /// Useful for testing or when you need to further customise the router
    /// before serving.
    pub async fn into_router(mut self) -> Result<Router, RokError> {
        let providers = std::mem::take(&mut self.providers);

        // 1. Register all providers
        for provider in &providers {
            let mut app_ref = AppRef(&mut self);
            provider
                .register(&mut app_ref)
                .await
                .map_err(|e| RokError::Internal(format!("{}::register: {e}", provider.name())))?;
        }

        // 2. Boot all providers
        for provider in &providers {
            provider
                .boot(&self)
                .await
                .map_err(|e| RokError::Internal(format!("{}::boot: {e}", provider.name())))?;
        }

        // 3. Build the router
        let mut router = match self.router_fn.take() {
            Some(f) => f(&self),
            None => Router::new(),
        };

        // 4. Merge extra routers (e.g. WsRouter, admin panels, sub-apps)
        for extra in self.extra_routers {
            router = router.merge(extra);
        }

        // 5. Apply middleware (reverse so first-added is outermost)
        let middleware = std::mem::take(&mut self.middleware);
        Ok(middleware.into_iter().rev().fold(router, |r, f| f(r)))
    }

    /// Boot providers, build the router, and start serving.
    pub async fn serve(self, addr: &str) -> Result<(), RokError> {
        let router = self.into_router().await?;

        let listener = tokio::net::TcpListener::bind(addr)
            .await
            .map_err(|e| RokError::Internal(format!("bind {addr}: {e}")))?;
        axum::serve(listener, router)
            .await
            .map_err(|e| RokError::Internal(format!("server: {e}")))?;

        Ok(())
    }
}

/// Mutable reference to App, passed to [`ServiceProvider::register`].
pub struct AppRef<'a>(pub &'a mut App);

impl<'a> AppRef<'a> {
    pub fn resolve<T>(&self) -> Result<Arc<T>, crate::container::ContainerError>
    where
        T: Send + Sync + 'static,
    {
        self.0.container.make::<T>()
    }

    pub fn singleton<T>(&mut self, instance: T)
    where
        T: Send + Sync + 'static,
    {
        self.0.container.singleton(instance);
    }

    pub fn bind<T, F>(&mut self, factory: F)
    where
        T: Send + Sync + 'static,
        F: Fn() -> T + Send + Sync + 'static,
    {
        self.0.container.bind::<T, F>(factory);
    }
}

pub use self::error_handler::CatchLayer;