Menu
Public repository / master
AkurAI-IDP / src/main.rs
Latest change 88d40a597e3b6cb9debe2a5ecc194d3da7a17d5e - Wire akurai-log Monitor shipper into tracing setup by Ólafur Búi Ólafsson
#![allow(clippy::result_large_err)]
#![allow(clippy::too_many_arguments)]
#![allow(special_module_name)]
pub mod config;
pub mod db;
pub mod lib;
pub mod middleware;
pub mod routes;
use axum::{
Extension, Json, Router,
body::Body,
extract::Request,
http::{Response, StatusCode, header},
middleware as axum_mw,
response::{IntoResponse, Redirect},
routing::get,
};
use serde_json::json;
use subtle::ConstantTimeEq;
use tower_http::trace::TraceLayer;
use crate::middleware::rate_limit::RateLimitState;
use akurai_telemetry::TelemetryDelivery;
// ---------------------------------------------------------------------------
// Embedded static assets (compiled into the binary via include_bytes!)
// ---------------------------------------------------------------------------
const FAVICON_SVG: &[u8] = include_bytes!("../assets/favicon.svg");
const APPLE_TOUCH_ICON_PNG: &[u8] = include_bytes!("../assets/apple-touch-icon.png");
const THEME_JS: &[u8] = include_bytes!("../assets/theme.js");
const THEMES_CSS: &[u8] = include_bytes!("../assets/themes.css");
const THEMES_JSON: &[u8] = include_bytes!("../assets/themes.json");
const LANG_JS: &[u8] = include_bytes!("../assets/lang.js");
const AGENT_JS: &[u8] = include_bytes!("../assets/agent.js");
#[tokio::main]
async fn main() {
// stdout logging, plus shipping to AkurAI-Monitor when MONITOR_INGEST_URL/TOKEN are set
akurai_log::init("akurai-idp");
tracing::info!("akurai-idp starting up");
// Ensure DB is initialized (triggers LazyLock)
db::with_db(|_conn| {});
// Ensure signing key exists
ensure_signing_key();
let cfg = config::get();
let telemetry_delivery = TelemetryDelivery::new(lib::telemetry::TracingTelemetryReceiver, 256)
.expect("valid telemetry delivery capacity");
let telemetry = lib::telemetry::OrganizationTelemetry::from_sender(telemetry_delivery.sender());
// --- Build the router ---
// Admin routes (token-auth'd)
let admin_router = admin_routes().layer(axum_mw::from_fn(admin_auth_middleware));
// Build the full app using each module's router()
let app = Router::new()
// Auth pages (login/mfa/logout) with CSRF
.merge(
routes::auth_pages::router().layer(axum_mw::from_fn(middleware::csrf::csrf_protection)),
)
// OIDC discovery
.merge(routes::well_known::router())
// OIDC endpoints
.merge(
routes::authorize::router().layer(axum_mw::from_fn(middleware::auth::session_optional)),
)
.merge(routes::token::router())
.merge(routes::userinfo::router().layer(axum_mw::from_fn(middleware::auth::bearer_auth)))
.merge(routes::introspect::router())
// Canonical organization membership and workspace linkage API.
.nest(
"/organizations",
routes::organizations::router().layer(axum_mw::from_fn(middleware::auth::bearer_auth)),
)
.merge(routes::revoke::router())
// Authenticated agent console
.merge(
routes::agent::router()
.layer(axum_mw::from_fn(middleware::auth::session_auth))
.layer(axum_mw::from_fn(middleware::csrf::csrf_protection)),
)
// Account: the account router declares full `/account/...` paths, so it must
// be merged at the root — nesting under "/account" double-prefixes every route
// to `/account/account/...` and 404s the real ones (incl. the "Account" nav link).
.merge(
routes::account::router()
.layer(axum_mw::from_fn(middleware::auth::session_auth))
.layer(axum_mw::from_fn(middleware::csrf::csrf_protection)),
)
// Bare root: apps link to the issuer's bare domain for "Manage your account",
// so send `/` to the account page (which itself redirects to /login if needed).
.route("/", get(root_redirect))
// Admin API
.nest("/admin", admin_router)
// Admin MCP server for agents (same admin bearer token)
.merge(routes::mcp::router(admin_routes()).layer(axum_mw::from_fn(admin_auth_middleware)))
// Static assets
.route("/favicon.svg", get(serve_favicon))
.route("/apple-touch-icon.png", get(serve_apple_touch_icon))
.route("/theme.js", get(serve_theme_js))
.route("/themes.css", get(serve_themes_css))
.route("/themes.json", get(serve_themes_json))
.route("/lang.js", get(serve_lang_js))
.route("/agent.js", get(serve_agent_js))
// Health check
.route("/health", get(health))
.route("/api/health", get(health))
// 404 catch-all
.fallback(not_found)
// Global middleware: secure headers + default rate limit + request logging
.layer(axum_mw::from_fn(middleware::secure_headers::secure_headers))
.layer(axum_mw::from_fn(middleware::rate_limit::rate_limit))
.layer(TraceLayer::new_for_http())
.layer(Extension(telemetry))
.layer(Extension(RateLimitState::new(60_000, 100)));
tracing::info!("listening on {}", cfg.listen_addr);
let listener = tokio::net::TcpListener::bind(&cfg.listen_addr)
.await
.expect("failed to bind listener");
axum::serve(listener, app).await.expect("server error");
}
/// Admin API routes, unauthenticated; callers add `admin_auth_middleware`.
fn admin_routes() -> Router {
Router::new()
.nest("/users", routes::admin::users::router())
.nest("/groups", routes::admin::groups::router())
.nest("/tenants", routes::admin::tenants::router())
.nest("/clients", routes::admin::clients::router())
.nest("/audit", routes::admin::audit::router())
}
/// Admin auth middleware: requires Bearer token matching IDP_ADMIN_TOKEN (constant-time compare).
async fn admin_auth_middleware(req: Request, next: axum::middleware::Next) -> impl IntoResponse {
let cfg = config::get();
if cfg.admin_token.is_empty() {
return (
StatusCode::SERVICE_UNAVAILABLE,
Json(json!({"error": "admin API not configured (IDP_ADMIN_TOKEN not set)"})),
)
.into_response();
}
let auth_header = req
.headers()
.get(header::AUTHORIZATION)
.and_then(|v| v.to_str().ok());
let token = match auth_header {
Some(h) if h.starts_with("Bearer ") => &h[7..],
_ => {
return (
StatusCode::UNAUTHORIZED,
Json(json!({"error": "missing or invalid Authorization header"})),
)
.into_response();
}
};
// Constant-time comparison
let expected = cfg.admin_token.as_bytes();
let provided = token.as_bytes();
if expected.len() != provided.len() || expected.ct_eq(provided).unwrap_u8() != 1 {
return (
StatusCode::UNAUTHORIZED,
Json(json!({"error": "invalid admin token"})),
)
.into_response();
}
next.run(req).await.into_response()
}
/// Ensure at least one active signing key exists. Generate one if the table is empty.
fn ensure_signing_key() {
let has_active = db::with_db(|conn| {
let count: i64 = conn
.query_row(
"SELECT COUNT(*) FROM signing_keys WHERE active = 1",
[],
|row| row.get(0),
)
.unwrap_or(0);
count > 0
});
if !has_active {
tracing::info!("no active signing key found, generating one");
let (kid, pub_pem, priv_pem) = lib::jwt::generate_signing_key();
let now = chrono::Utc::now().timestamp();
db::with_db(|conn| {
conn.execute(
"INSERT INTO signing_keys (kid, alg, public_key, private_key_enc, active, created_at) VALUES (?1, ?2, ?3, ?4, 1, ?5)",
rusqlite::params![kid, "EdDSA", pub_pem, priv_pem, now],
)
.expect("failed to insert signing key");
});
tracing::info!("signing key generated: kid={kid}");
}
}
async fn serve_favicon() -> Response<Body> {
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "image/svg+xml")
.header(header::CACHE_CONTROL, "public, max-age=86400")
.body(Body::from(FAVICON_SVG.to_vec()))
.unwrap()
}
async fn serve_apple_touch_icon() -> Response<Body> {
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "image/png")
.header(header::CACHE_CONTROL, "public, max-age=86400")
.body(Body::from(APPLE_TOUCH_ICON_PNG.to_vec()))
.unwrap()
}
async fn serve_theme_js() -> Response<Body> {
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "text/javascript; charset=utf-8")
.header(header::CACHE_CONTROL, "public, max-age=3600")
.body(Body::from(THEME_JS.to_vec()))
.unwrap()
}
async fn serve_themes_css() -> Response<Body> {
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "text/css; charset=utf-8")
.header(header::CACHE_CONTROL, "public, max-age=3600")
.body(Body::from(THEMES_CSS.to_vec()))
.unwrap()
}
async fn serve_themes_json() -> Response<Body> {
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "application/json; charset=utf-8")
.header(header::CACHE_CONTROL, "public, max-age=3600")
.body(Body::from(THEMES_JSON.to_vec()))
.unwrap()
}
async fn serve_lang_js() -> Response<Body> {
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "text/javascript; charset=utf-8")
.header(header::CACHE_CONTROL, "public, max-age=3600")
.body(Body::from(LANG_JS.to_vec()))
.unwrap()
}
async fn serve_agent_js() -> Response<Body> {
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "text/javascript; charset=utf-8")
.header(header::CACHE_CONTROL, "no-cache")
.body(Body::from(AGENT_JS.to_vec()))
.unwrap()
}
/// AkurAI-Framework health contract version this app reports.
const FRAMEWORK_CONTRACT_VERSION: &str = "0.8.4";
pub(crate) async fn health() -> impl IntoResponse {
Json(health_document())
}
/// Framework health contract plus the legacy `{ok, service}` fields.
fn health_document() -> serde_json::Value {
json!({
"app": "akurai-idp",
"version": env!("CARGO_PKG_VERSION"),
"status": "ok",
"framework": "AkurAI-Framework",
"framework_version": FRAMEWORK_CONTRACT_VERSION,
"ok": true,
"service": "akurai-idp",
})
}
/// Bare root `/` → the account page. Relying-party apps point their "Manage your
/// account" links at the issuer's bare domain; without this the IDP 404s there.
async fn root_redirect() -> Redirect {
Redirect::to("/account")
}
async fn not_found() -> impl IntoResponse {
(StatusCode::NOT_FOUND, Json(json!({"error": "not found"})))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn health_meets_framework_contract_and_keeps_legacy_fields() {
let doc = health_document();
assert_eq!(doc["app"], "akurai-idp");
assert_eq!(doc["status"], "ok");
assert_eq!(doc["framework"], "AkurAI-Framework");
assert_eq!(doc["version"], env!("CARGO_PKG_VERSION"));
assert_eq!(doc["ok"], true);
assert_eq!(doc["service"], "akurai-idp");
}
}