use axum::{
extract::Request,
http::StatusCode,
middleware::Next,
response::{IntoResponse, Response},
};
use std::sync::Arc;
use crate::store::KeyStore;
enum KeyRole { Admin, Tunnel }
async fn require_key(
req: Request,
next: Next,
header: &'static str,
role: KeyRole,
) -> Result<Response, Response> {
let store = req
.extensions()
.get::<Arc<KeyStore>>()
.cloned()
.ok_or_else(|| {
(StatusCode::INTERNAL_SERVER_ERROR, "store not found").into_response()
})?;
let key = req
.headers()
.get(header)
.and_then(|v| v.to_str().ok())
.unwrap_or("");
if key.is_empty() {
return Err((StatusCode::UNAUTHORIZED, format!("missing {header}")).into_response());
}
let valid = match role {
KeyRole::Admin => store.validate_admin(key).await.unwrap_or(false),
KeyRole::Tunnel => store.validate_tunnel(key).await.unwrap_or(false),
};
if valid {
Ok(next.run(req).await)
} else {
Err((StatusCode::UNAUTHORIZED, format!("invalid {header}")).into_response())
}
}
pub async fn require_admin_key(
req: Request,
next: Next,
) -> Result<Response, Response> {
require_key(req, next, "X-Admin-Key", KeyRole::Admin).await
}
pub async fn require_tunnel_key(
req: Request,
next: Next,
) -> Result<Response, Response> {
require_key(req, next, "X-Api-Key", KeyRole::Tunnel).await
}