use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use axum::body::Body;
use axum::extract::State;
use axum::http::{header, Request, StatusCode};
use axum::middleware::Next;
use axum::response::{Html, IntoResponse, Redirect, Response};
use axum::Form;
use serde::Deserialize;
const SESSION_TTL: Duration = Duration::from_secs(12 * 3600);
const COOKIE: &str = "dtmrs_session";
pub struct Auth {
user: String,
password: String,
sessions: Mutex<HashMap<String, Instant>>,
}
impl Auth {
pub fn from_env() -> Option<Arc<Self>> {
let password = std::env::var("DTMRS_ADMIN_PASSWORD").unwrap_or_default();
if password.is_empty() {
return None;
}
Some(Arc::new(Self {
user: std::env::var("DTMRS_ADMIN_USER").unwrap_or_else(|_| "admin".into()),
password,
sessions: Mutex::new(HashMap::new()),
}))
}
fn matches(&self, user: &str, password: &str) -> bool {
let a = user.as_bytes();
let b = self.user.as_bytes();
let c = password.as_bytes();
let d = self.password.as_bytes();
let mut diff = (a.len() ^ b.len()) | (c.len() ^ d.len());
for i in 0..a.len().max(b.len()) {
diff |= usize::from(a.get(i).copied().unwrap_or(0) ^ b.get(i).copied().unwrap_or(1));
}
for i in 0..c.len().max(d.len()) {
diff |= usize::from(c.get(i).copied().unwrap_or(0) ^ d.get(i).copied().unwrap_or(1));
}
diff == 0
}
fn issue(&self) -> String {
use rand::Rng;
let raw: [u8; 32] = rand::thread_rng().gen();
let token: String = raw.iter().map(|b| format!("{b:02x}")).collect();
let mut s = self.sessions.lock().unwrap();
let now = Instant::now();
s.retain(|_, exp| *exp > now); s.insert(token.clone(), now + SESSION_TTL);
token
}
fn valid(&self, token: &str) -> bool {
let mut s = self.sessions.lock().unwrap();
match s.get(token) {
Some(exp) if *exp > Instant::now() => true,
Some(_) => {
s.remove(token);
false
}
None => false,
}
}
fn revoke(&self, token: &str) {
self.sessions.lock().unwrap().remove(token);
}
}
fn cookie_of(req: &Request<Body>) -> Option<String> {
req.headers()
.get(header::COOKIE)?
.to_str()
.ok()?
.split(';')
.filter_map(|kv| kv.split_once('='))
.find(|(k, _)| k.trim() == COOKIE)
.map(|(_, v)| v.trim().to_string())
}
fn is_https(req_headers: &axum::http::HeaderMap) -> bool {
req_headers
.get("x-forwarded-proto")
.and_then(|v| v.to_str().ok())
.map(|v| v.eq_ignore_ascii_case("https"))
.unwrap_or(false)
}
pub async fn guard(
State(auth): State<Arc<Auth>>,
req: Request<Body>,
next: Next,
) -> Response {
let path = req.uri().path();
if path == "/health" || path == "/login" || path == "/logout" {
return next.run(req).await;
}
if cookie_of(&req).is_some_and(|t| auth.valid(&t)) {
return next.run(req).await;
}
let wants_html = req
.headers()
.get(header::ACCEPT)
.and_then(|v| v.to_str().ok())
.is_some_and(|v| v.contains("text/html"));
if wants_html {
Redirect::to("/login").into_response()
} else {
(StatusCode::UNAUTHORIZED, "需要登录").into_response()
}
}
#[derive(Deserialize)]
pub struct LoginForm {
user: String,
password: String,
}
pub async fn login_page() -> Html<&'static str> {
Html(include_str!("login.html"))
}
pub async fn login_submit(
State(auth): State<Arc<Auth>>,
headers: axum::http::HeaderMap,
Form(f): Form<LoginForm>,
) -> Response {
if !auth.matches(&f.user, &f.password) {
return (
StatusCode::UNAUTHORIZED,
Html(include_str!("login.html").replace(
"<!--ERR-->",
r#"<p class="err">用户名或密码不对</p>"#,
)),
)
.into_response();
}
let token = auth.issue();
let secure = if is_https(&headers) { "; Secure" } else { "" };
let cookie = format!(
"{COOKIE}={token}; Path=/; HttpOnly; SameSite=Lax; Max-Age={}{secure}",
SESSION_TTL.as_secs()
);
([(header::SET_COOKIE, cookie)], Redirect::to("/console")).into_response()
}
pub async fn logout(State(auth): State<Arc<Auth>>, req: Request<Body>) -> Response {
if let Some(t) = cookie_of(&req) {
auth.revoke(&t);
}
(
[(
header::SET_COOKIE,
format!("{COOKIE}=; Path=/; HttpOnly; Max-Age=0"),
)],
Redirect::to("/login"),
)
.into_response()
}