auth/src/main.rs

use axum::extract::Query;
use axum::response::{Html, Redirect};
use axum::http::{Uri, HeaderMap, HeaderValue};
use axum::{
    extract::State,
    http::StatusCode,
    response::IntoResponse,
    routing::get,
    Form, Router,
};
use serde::Deserialize;
use sqlx::sqlite::SqlitePool;
use std::{env, borrow::BorrowMut};
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use std::{collections::HashMap, str::FromStr};
use tokio::sync::Mutex;
use tower_cookies::{Cookie, CookieManagerLayer, Cookies};
use uuid::Uuid;
use tower_http::trace::{self, TraceLayer};
use tracing::Level;
use minijinja::{Environment, context};
pub struct ServerState {
    pub db: sqlx::Pool<sqlx::Sqlite>,
    pub session: Mutex<HashMap<Uuid, Instant>>,
    // started: Instant,
}
#[derive(Deserialize, sqlx::FromRow,Debug)]
pub struct UserLoginForm {
    #[sqlx(rename = "NAME")]
    pub username: String,
    #[sqlx(rename = "KEY")]
    pub otp: String,
}
const COOKIE_NAME: &str = "aaron_auth";
const SESSION_ACTIVE_TIME: u64 = 600;
const LOGIN_PAGE_HTML: &str = include_str!("../loginpage.html");
#[tokio::main]
async fn main() {
    // 初始化日志记录器
    tracing_subscriber::fmt::init();
    let pool = SqlitePool::connect(&env::var("DATABASE_URL").expect("DB URL NOT SPECIFIED"))
        .await
        .expect("DB OPEN FAILURE"); // our router
    let state = Arc::new(ServerState {
        db: pool,
        session: HashMap::new().into(),
        // started: std::time::Instant::now(),
    });
    let app = Router::new()
        .route("/auth", get(auth)) // http://127.0.0.1:3000
        .route("/login", get(login_page).post(login))
        .with_state(state.clone())
        .layer(CookieManagerLayer::new())
        .layer(
            TraceLayer::new_for_http()
                .make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
                .on_response(trace::DefaultOnResponse::new().level(Level::INFO)),
        );
    let port = env::var("PORT").unwrap_or("3000".to_string());
    let port = port.parse::<u16>().unwrap_or(3000);
    let addr = SocketAddr::from(([127, 0, 0, 1], port));
    // run it with hyper on localhost:3000
    tokio::spawn(gc_task(state.clone()));
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}


async fn gc_task(state:Arc<ServerState>){
    let mut interval = tokio::time::interval(Duration::from_secs(SESSION_ACTIVE_TIME));
    loop {
        interval.tick().await;
        let res = gc(state.clone()).await;
        match res {
            Ok(_) => tracing::info!("gc completed"),
            Err(s) => tracing::error!("gc failed:{}",s),
        }
    }
}

// 处理/auth
async fn auth(
    State(state): State<Arc<ServerState>>,
    // Form(frm): Form<UserLoginForm>,
    cookies: Cookies,
) -> StatusCode {
    if let Some(session_token) = cookies.get(COOKIE_NAME) {
        tracing::info!("session:{}",session_token.value());
        let Ok(s) = uuid::Uuid::from_str(session_token.value()) else {
            return StatusCode::UNAUTHORIZED;
        };
        let mut locked = state.session.lock().await;
        if locked.contains_key(&s) { // FIX, when accessed /auth with correct cookie, the cookie's expiration is delayed
            let Some(v) = locked.insert(s,Instant::now()+Duration::from_secs(SESSION_ACTIVE_TIME)) else {
                return StatusCode::UNAUTHORIZED;
            };
            if Instant::now() < v {
                return StatusCode::OK;
            }
        }
    }
    return StatusCode::UNAUTHORIZED;
}

async fn login(
    State(state): State<Arc<ServerState>>,
    cookies: Cookies,
    Query(mut params): Query<HashMap<String, String>>,
    Form(frm): Form<UserLoginForm>,
) -> Result<Redirect, (StatusCode, &'static str)> {
    let conn = state.db.acquire().await;
    let Ok(mut conn) = conn else {
        return Err((StatusCode::BAD_GATEWAY, "db连接错误"));
    };
    tracing::info!("{:?}",&frm);
    let target = sqlx::query_as::<_, UserLoginForm>(
        r#"
        SELECT NAME, KEY FROM USERS WHERE NAME = ?
        "#,
    )
    .bind(frm.username)
    .fetch_optional(&mut *conn)
    .await;
    tracing::info!("{:?}",&target);

    if let Ok(Some(target)) = target {
        if check_otp(target.otp, frm.otp) {
            let s = Uuid::new_v4();
            let mut locked = state.session.lock().await;
            locked.insert(s.clone(), Instant::now() + Duration::from_secs(SESSION_ACTIVE_TIME));
            let mut new_cookie = Cookie::new(COOKIE_NAME, s.to_string());
            new_cookie.set_domain(".aaronhu.cn");
            cookies.add(new_cookie);
            if let Some(original_uri) = params.get("original_url") {
                return Ok(Redirect::to(&original_uri));
            }

            return Err((StatusCode::ACCEPTED, "ok"));
        } else {
            return Err((StatusCode::UNAUTHORIZED,"wrong password"));
        }
    }
    return Err((StatusCode::BAD_GATEWAY, "unreachable"));
}

async fn login_page(
    headers: HeaderMap<HeaderValue>,
)-> impl IntoResponse{
    tracing::info!("Headers: {:#?}",headers);
    let mut env = Environment::new();
    env.add_template("login.html", LOGIN_PAGE_HTML).unwrap();
    let template = env.get_template("login.html").unwrap();
    if let Some(original_uri) = headers.get("X-Original-URI") {
        if let Ok(uri) = original_uri.to_str() {
            tracing::info!("redirect to {}",uri);
            if !uri.is_empty() {
                let uri = "?original_url=".to_owned()+uri;
                return Html(template.render(context! { url => uri }).unwrap_or("Error".to_string()))
            }
        }
    }
    return Html(template.render(context! { url => String::new() }).unwrap_or("Error".to_string()))
}

pub fn check_otp(key_from_db: String, user_input_otp: String) -> bool {
    use totp_rs::{Algorithm, Secret, TOTP};
    let totp = TOTP::new(
        Algorithm::SHA1,
        6,
        1,
        30,
        Secret::Raw(key_from_db.as_bytes().to_vec())
            .to_bytes()
            .unwrap(),
    );
    if let Ok(otp) = totp {
        if let Ok(token) = otp.generate_current() {
            return token == user_input_otp;
        }
    }
    return false;
}

async fn gc(state:Arc<ServerState>)->Result<(),String>{
    let mut locked = state.session.lock().await;
    let current_time = Instant::now();
    locked.borrow_mut().retain(|_,v| *v < current_time);
    tracing::info!("gc fired,active Sessions {:?}",locked);
    Ok(())
}