use axum::Router;
use axum::http::{HeaderMap, StatusCode, header};
use axum::response::{IntoResponse, Redirect, Response};
use axum::routing::get;
use std::sync::Arc;
use tracing::info;
use assay_domain::events::EngineEventBus;
use assay_workflow::events::WorkflowEventBus;
use assay_workflow::{WorkflowCtx, WorkflowStore};
use crate::state::EngineState;
pub fn build_app<S: WorkflowStore + Clone + 'static>(state: EngineState<S>) -> Router {
let operator_ui_enabled = state.engine_config.dashboard.operator_enabled();
let auth_ui_enabled = state.engine_config.dashboard.auth_ui_enabled();
let state_for_workflow = state.clone();
let workflow_router = assay_workflow::api::router(Arc::clone(&state.workflow), |r| {
r.layer(axum::middleware::from_fn_with_state(
state_for_workflow,
admin_bearer_middleware::<S>,
))
});
let healthz = Router::new().route(
"/healthz",
get(|| async { Redirect::permanent("/api/v1/engine/core/health") }),
);
let engine_api_router = crate::engine_api::router::<S>().with_state(state.clone());
let mut app = workflow_router.merge(healthz).merge(engine_api_router);
if operator_ui_enabled {
let dashboard_router =
assay_dashboard::workflow_router().with_state(Arc::clone(&state.dashboard));
let engine_console_router = assay_dashboard::engine_router();
app = app.merge(dashboard_router).merge(engine_console_router);
}
if state.auth.is_some() {
let spec_router =
assay_auth::oidc_spec_router::<EngineState<S>>().with_state(state.clone());
app = app.nest("/auth", spec_router);
let engine_auth_router =
assay_auth::engine_auth_router::<EngineState<S>>().with_state(state.clone());
app = app.nest("/api/v1/engine/auth", engine_auth_router);
if auth_ui_enabled {
app = app.merge(assay_dashboard::auth_public_router());
}
if operator_ui_enabled {
app = app.merge(assay_dashboard::auth_console_router());
}
}
#[cfg(feature = "vault")]
if state.vault.is_some() {
let state_for_vault = state.clone();
let vault = assay_vault::router::vault_router::<EngineState<S>, _>(|r| {
r.layer(axum::middleware::from_fn_with_state(
state_for_vault,
admin_bearer_middleware::<S>,
))
})
.with_state(state.clone());
app = app.nest("/api/v1/vault", vault);
if operator_ui_enabled {
app = app.merge(assay_dashboard::vault_router());
}
}
#[cfg(all(feature = "vault", feature = "vault-bitwarden-compat"))]
if state.vault.is_some() && state.auth.is_some() {
let bw =
assay_vault::bitwarden_compat::router::<EngineState<S>>().with_state(state.clone());
app = app.merge(bw);
}
if auth_ui_enabled && !operator_ui_enabled {
let auth_url = state
.engine_config
.auth
.public_url
.as_deref()
.unwrap_or(&state.engine_config.server.public_url);
let auth_host = url::Url::parse(auth_url)
.ok()
.and_then(|url| url.host_str().map(str::to_owned));
let root = Router::new()
.route("/", get(auth_origin_root))
.with_state(auth_host);
app = app.merge(root);
}
if !state.engine_config.server.allowed_hosts.is_empty() {
app = app.layer(axum::middleware::from_fn_with_state(
state,
allowed_host_middleware::<S>,
));
}
app
}
async fn auth_origin_root(
axum::extract::State(auth_host): axum::extract::State<Option<String>>,
headers: HeaderMap,
) -> Response {
let request_host = request_host(&headers);
if auth_host
.as_deref()
.zip(request_host.as_deref())
.is_some_and(|(expected, actual)| expected.eq_ignore_ascii_case(actual))
{
return Redirect::temporary("/auth/landing").into_response();
}
StatusCode::NOT_FOUND.into_response()
}
async fn allowed_host_middleware<S: WorkflowStore + Clone + 'static>(
axum::extract::State(state): axum::extract::State<EngineState<S>>,
request: axum::extract::Request,
next: axum::middleware::Next,
) -> Response {
if request.uri().path() == "/api/v1/engine/core/health"
|| host_is_allowed(request.headers(), &state.engine_config.server.allowed_hosts)
{
return next.run(request).await;
}
StatusCode::MISDIRECTED_REQUEST.into_response()
}
fn host_is_allowed(headers: &HeaderMap, allowed_hosts: &[String]) -> bool {
if allowed_hosts.is_empty() {
return true;
}
let Some(host) = request_host(headers) else {
return false;
};
allowed_hosts
.iter()
.any(|allowed| allowed.eq_ignore_ascii_case(&host))
}
fn request_host(headers: &HeaderMap) -> Option<String> {
let value = headers.get(header::HOST)?.to_str().ok()?;
value
.parse::<axum::http::uri::Authority>()
.ok()
.map(|authority| authority.host().to_owned())
}
async fn admin_bearer_middleware<S: WorkflowStore + Clone + 'static>(
axum::extract::State(state): axum::extract::State<EngineState<S>>,
request: axum::extract::Request,
next: axum::middleware::Next,
) -> axum::response::Response {
let path = request.uri().path();
if (path.starts_with("/share/") && path != "/share/revoke")
|| (path.starts_with("/api/v1/vault/share/") && path != "/api/v1/vault/share/revoke")
{
return next.run(request).await;
}
let keys = crate::state::AdminApiKeys(Arc::clone(&state.admin_api_keys));
let outcome = match state.auth.as_ref() {
Some(auth) => assay_auth::gate::require_admin_or_jwt(request.headers(), auth, &keys)
.await
.map(|_| ()),
None => assay_auth::gate::require_admin_bearer(request.headers(), &keys),
};
if let Err(r) = outcome {
return *r;
}
next.run(request).await
}
pub async fn serve<S: WorkflowStore + Clone + 'static>(
bind_addr: &str,
state: EngineState<S>,
) -> anyhow::Result<()> {
let app = build_app(state);
bind_and_serve(bind_addr, app).await
}
pub async fn bind_and_serve(bind_addr: &str, app: axum::Router) -> anyhow::Result<()> {
let listener = tokio::net::TcpListener::bind(bind_addr)
.await
.map_err(|e| anyhow::anyhow!("bind {bind_addr}: {e}"))?;
let actual = listener.local_addr()?;
info!(target: "assay-engine", %actual, "listening");
axum::serve(listener, app).await?;
Ok(())
}
pub fn build_workflow_ctx<S: WorkflowStore + 'static>(store: S) -> Arc<WorkflowCtx<S>> {
let ctx = WorkflowCtx::start(Arc::new(store)).with_binary_version(env!("CARGO_PKG_VERSION"));
Arc::new(ctx)
}
pub fn build_workflow_ctx_with_bus<S: WorkflowStore + 'static>(
store: S,
bus: Arc<dyn EngineEventBus>,
) -> Arc<WorkflowCtx<S>> {
let ctx = WorkflowCtx::start(Arc::new(store))
.with_binary_version(env!("CARGO_PKG_VERSION"))
.with_event_bus(WorkflowEventBus::new(bus));
Arc::new(ctx)
}
#[cfg(test)]
mod host_boundary_tests {
use axum::extract::State;
use axum::http::{HeaderMap, HeaderValue, StatusCode, header};
use super::{auth_origin_root, host_is_allowed};
#[tokio::test]
async fn auth_origin_root_enters_public_auth_while_engine_root_stays_hidden() {
let mut auth_headers = HeaderMap::new();
auth_headers.insert(header::HOST, HeaderValue::from_static("auth.assay.rs"));
let auth = auth_origin_root(State(Some("auth.assay.rs".to_string())), auth_headers).await;
assert_eq!(auth.status(), StatusCode::TEMPORARY_REDIRECT);
assert_eq!(auth.headers()[header::LOCATION], "/auth/landing");
let mut engine_headers = HeaderMap::new();
engine_headers.insert(header::HOST, HeaderValue::from_static("engine.assay.rs"));
let engine =
auth_origin_root(State(Some("auth.assay.rs".to_string())), engine_headers).await;
assert_eq!(engine.status(), StatusCode::NOT_FOUND);
}
#[test]
fn configured_hosts_are_case_insensitive_and_port_agnostic() {
let allowed = vec!["auth.assay.rs".to_string(), "engine.assay.rs".to_string()];
let mut headers = HeaderMap::new();
headers.insert(header::HOST, HeaderValue::from_static("AUTH.ASSAY.RS:443"));
assert!(host_is_allowed(&headers, &allowed));
}
#[test]
fn unknown_and_missing_hosts_are_rejected_when_the_allowlist_is_configured() {
let allowed = vec!["auth.assay.rs".to_string(), "engine.assay.rs".to_string()];
let mut headers = HeaderMap::new();
headers.insert(header::HOST, HeaderValue::from_static("assay-auth.fly.dev"));
assert!(!host_is_allowed(&headers, &allowed));
assert!(!host_is_allowed(&HeaderMap::new(), &allowed));
}
#[test]
fn an_empty_allowlist_preserves_embedded_and_local_callers() {
assert!(host_is_allowed(&HeaderMap::new(), &[]));
}
}