84 lines
2.5 KiB
Rust
84 lines
2.5 KiB
Rust
mod healthcheck;
|
|
use std::sync::mpsc;
|
|
use std::time::Duration;
|
|
|
|
use axum::extract::Request;
|
|
use axum::middleware::{Next, from_fn};
|
|
use axum::response::Response;
|
|
use axum::routing::{get, post};
|
|
use axum::{Router, body::Body};
|
|
use http::StatusCode;
|
|
use tower::ServiceBuilder;
|
|
use tower_http::timeout::TimeoutLayer;
|
|
use tower_http::trace::TraceLayer;
|
|
use tracing::Level;
|
|
use uuid::Uuid;
|
|
|
|
pub fn app() -> (Router, mpsc::Receiver<bool>) {
|
|
let (tx, rx) = mpsc::channel();
|
|
(
|
|
Router::new()
|
|
.route("/healthcheck", get(healthcheck::healthcheck))
|
|
.route("/forever", get(std::future::pending::<()>))
|
|
.nest("/admin", admin(tx))
|
|
.layer(
|
|
ServiceBuilder::new()
|
|
.layer(
|
|
TraceLayer::new_for_http().make_span_with(|req: &Request<Body>| {
|
|
tracing::span!(
|
|
Level::DEBUG,
|
|
"request",
|
|
trace_id = Uuid::now_v7().to_string(),
|
|
method = format!("{}", req.method()),
|
|
uri = format!("{}", req.uri()),
|
|
)
|
|
}),
|
|
)
|
|
.layer(TimeoutLayer::new(Duration::from_secs(10))),
|
|
),
|
|
rx,
|
|
)
|
|
}
|
|
|
|
fn admin(tx: mpsc::Sender<bool>) -> Router {
|
|
let r = Router::new().route("/test", get(async || StatusCode::OK));
|
|
|
|
let r = add_shutdown_endpoint(r, tx);
|
|
r.layer(from_fn(async |req: Request, next: Next| {
|
|
if let Ok(secret) = std::env::var("ADMIN_SECRET") {
|
|
println!("ADMIN_SECRET: {secret}");
|
|
match req.headers().get("Authorization") {
|
|
Some(key) if secret == key.to_owned() => (),
|
|
_ => {
|
|
return Response::builder()
|
|
.status(StatusCode::UNAUTHORIZED)
|
|
.body(Body::empty())
|
|
.unwrap();
|
|
}
|
|
}
|
|
}
|
|
|
|
next.run(req).await
|
|
}))
|
|
}
|
|
|
|
#[cfg(feature = "shutdown")]
|
|
fn add_shutdown_endpoint(r: Router, tx: mpsc::Sender<bool>) -> Router {
|
|
r.route(
|
|
"/shutdown",
|
|
post(async move || {
|
|
let res = tx.send(true);
|
|
if res.is_ok() {
|
|
StatusCode::OK
|
|
} else {
|
|
StatusCode::INTERNAL_SERVER_ERROR
|
|
}
|
|
}),
|
|
)
|
|
}
|
|
|
|
#[cfg(not(feature = "shutdown"))]
|
|
fn add_shutdown_endpoint(r: Router, _: mpsc::Sender<bool>) -> Router {
|
|
r
|
|
}
|