115 lines
3.3 KiB
Rust
115 lines
3.3 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;
|
|
#[allow(unused_imports)]
|
|
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, warn};
|
|
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("/", 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") {
|
|
match req.headers().get("Authorization") {
|
|
Some(key) if secret == *key => (),
|
|
Some(key) => {
|
|
warn!("Unauthorized request with key: {key:?}");
|
|
return Response::builder()
|
|
.status(StatusCode::UNAUTHORIZED)
|
|
.body(Body::empty())
|
|
.unwrap();
|
|
}
|
|
_ => {
|
|
warn!("Unauthorized request no key given");
|
|
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
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use tower::{self, ServiceExt};
|
|
|
|
#[tokio::test]
|
|
async fn test_authorization_disables_when_no_env_var_set() {
|
|
let (app, _) = app();
|
|
|
|
let resp = app
|
|
.oneshot(
|
|
axum::http::Request::builder()
|
|
.uri("/admin")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
}
|
|
}
|