kroki_rs/server/
mod.rs

1use crate::browser::BrowserManager;
2use crate::capabilities::Capabilities;
3use crate::config::Config;
4use crate::diagrams::registry::DiagramRegistry;
5use std::net::SocketAddr;
6use std::sync::Arc;
7
8pub mod admin;
9mod handlers;
10pub mod metrics;
11pub mod middleware;
12
13use metrics::PrometheusHandle;
14use middleware::circuit_breaker::CircuitBreakerManager;
15use middleware::rate_limit::RateLimiter;
16
17/// Shared application state injected into every Axum handler.
18#[derive(Clone)]
19pub struct AppState {
20    pub config: Config,
21    pub registry: Arc<DiagramRegistry>,
22    pub browser_manager: Option<Arc<BrowserManager>>,
23    pub rate_limiter: Option<RateLimiter>,
24    pub circuit_breaker: Option<CircuitBreakerManager>,
25    pub metrics_handle: Option<PrometheusHandle>,
26}
27
28/// Starts the Axum web server on the specified port.
29pub async fn run(config: Config) -> anyhow::Result<()> {
30    let capabilities = Capabilities::discover(&config);
31    let port = config.server.port;
32    let admin_port = config.server.admin_port;
33    let host = &config.server.host;
34
35    println!(
36        "\nšŸš€ Kroki-rs v{} is starting up...",
37        env!("CARGO_PKG_VERSION")
38    );
39    println!("------------------------------------------------------------");
40    println!("šŸ“” API Server:      http://{}:{}", host, port);
41    println!("šŸ› ļø  Admin Dashboard: http://{}:{}", host, admin_port);
42    println!("šŸ“– Documentation:   https://softmentor.github.io/kroki-rs/");
43    println!("------------------------------------------------------------\n");
44
45    tracing::info!("Capabilities discovered: {:?}", capabilities);
46
47    let browser_manager = match BrowserManager::start(
48        config.browser.pool_size,
49        config.browser.context_ttl_requests,
50    )
51    .await
52    {
53        Ok(manager) => Some(Arc::new(manager)),
54        Err(e) => {
55            tracing::warn!("Native Browser Backend failed to initialize: {}. Browser-based features (Mermaid, BPMN) will be disabled.", e);
56            None
57        }
58    };
59
60    let registry = Arc::new(DiagramRegistry::new(
61        &capabilities,
62        &config,
63        browser_manager.clone(),
64    ));
65
66    // Initialize rate limiter if enabled
67    let rate_limiter = if config.server.rate_limit.enabled {
68        tracing::info!(
69            "Rate limiting enabled: {} req/s, burst: {}",
70            config.server.rate_limit.requests_per_second,
71            config.server.rate_limit.burst_size
72        );
73        Some(RateLimiter::new(&config.server.rate_limit))
74    } else {
75        tracing::info!("Rate limiting disabled (dev mode)");
76        None
77    };
78
79    // Initialize circuit breaker if enabled
80    let circuit_breaker = if config.server.circuit_breaker.enabled {
81        tracing::info!(
82            "Circuit breaker enabled: threshold={}, reset={}s",
83            config.server.circuit_breaker.failure_threshold,
84            config.server.circuit_breaker.reset_timeout_secs
85        );
86        Some(CircuitBreakerManager::new(&config.server.circuit_breaker))
87    } else {
88        tracing::info!("Circuit breaker disabled");
89        None
90    };
91
92    // Initialize metrics if enabled
93    let metrics_handle = if config.server.metrics.enabled {
94        tracing::info!("Metrics collection enabled");
95        Some(metrics::init_metrics())
96    } else {
97        tracing::info!("Metrics collection disabled");
98        None
99    };
100
101    if config.server.auth.enabled {
102        tracing::info!(
103            "API key authentication enabled ({} key(s) configured)",
104            config.server.auth.api_keys.len()
105        );
106    } else {
107        tracing::info!("Authentication disabled (dev mode)");
108    }
109
110    let state = AppState {
111        config,
112        registry,
113        browser_manager,
114        rate_limiter,
115        circuit_breaker,
116        metrics_handle,
117    };
118
119    let admin_state = state.clone();
120    tokio::spawn(async move {
121        if let Err(e) = admin::run_admin_server(admin_state).await {
122            tracing::error!("Admin server failed: {}", e);
123        }
124    });
125
126    let listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{}", port)).await?;
127    let router = app(state.clone());
128    let service = router.into_make_service_with_connect_info::<SocketAddr>();
129    axum::serve(listener, service).await?;
130
131    Ok(())
132}
133
134fn app(state: AppState) -> axum::Router {
135    use axum::{middleware as mw, routing::get};
136
137    axum::Router::new()
138        .route("/", get(handlers::root))
139        .route("/{type}/{format}/{source}", get(handlers::get_diagram))
140        .route(
141            "/{type}/{format}",
142            axum::routing::post(handlers::post_render),
143        )
144        .layer(mw::from_fn_with_state(
145            state.clone(),
146            middleware::auth::auth_middleware,
147        ))
148        .layer(mw::from_fn_with_state(
149            state.clone(),
150            middleware::rate_limit::rate_limit_middleware,
151        ))
152        .with_state(state)
153}