This commit is contained in:
yly 2024-03-13 15:24:55 +08:00
parent 7be1bca385
commit 8a900fe7a7
8 changed files with 429 additions and 259 deletions

278
Cargo.lock generated
View File

@ -131,6 +131,7 @@ dependencies = [
"totp-rs",
"tower-cookies",
"tower-http",
"tower-sessions",
"tracing",
"tracing-subscriber",
"uuid",
@ -144,19 +145,19 @@ checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
[[package]]
name = "axum"
version = "0.6.20"
version = "0.7.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3b829e4e32b91e643de6eafe82b1d90675f5874230191a4ffbc1b336dec4d6bf"
checksum = "1236b4b292f6c4d6dc34604bb5120d85c3fe1d1aa596bd5cc52ca054d13e7b9e"
dependencies = [
"async-trait",
"axum-core",
"bitflags 1.3.2",
"bytes",
"futures-util",
"headers",
"http",
"http-body",
"http-body-util",
"hyper",
"hyper-util",
"itoa",
"matchit",
"memchr",
@ -173,23 +174,28 @@ dependencies = [
"tower",
"tower-layer",
"tower-service",
"tracing",
]
[[package]]
name = "axum-core"
version = "0.3.4"
version = "0.4.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "759fa577a247914fd3f7f76d62972792636412fbfd634cd452f6a385a74d2d2c"
checksum = "a15c63fd72d41492dc4f497196f5da1fb04fb7529e631d73630d1b491e47a2e3"
dependencies = [
"async-trait",
"bytes",
"futures-util",
"http",
"http-body",
"http-body-util",
"mime",
"pin-project-lite",
"rustversion",
"sync_wrapper",
"tower-layer",
"tower-service",
"tracing",
]
[[package]]
@ -219,6 +225,12 @@ version = "0.21.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "35636a1494ede3b646cc98f74f8e62c773a38a659ebc777a2cf26b9b74171df9"
[[package]]
name = "base64"
version = "0.22.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9475866fec1451be56a3c2400fd081ff546538961565ccb5b7142cbd22bc7a51"
[[package]]
name = "base64ct"
version = "1.6.0"
@ -302,9 +314,9 @@ checksum = "21a53c0a4d288377e7415b53dcfc3c04da5cdc2cc95c8d5ac178b58f0b861ad6"
[[package]]
name = "cookie"
version = "0.17.0"
version = "0.18.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7efb37c3e1ccb1ff97164ad95ac1606e8ccd35b3fa0a7d99a304c7f4a428cc24"
checksum = "3cd91cf61412820176e137621345ee43b3f4423e589e7ae4e50d601d93e35ef8"
dependencies = [
"percent-encoding",
"time",
@ -382,6 +394,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0f32d04922c60427da6f9fef14d042d9edddef64cb9d4ce0d64d0685fbeb1fd3"
dependencies = [
"powerfmt",
"serde",
]
[[package]]
@ -489,10 +502,24 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fed34cd105917e91daa4da6b3728c47b068749d6a62c59811f06ed2ac71d9da7"
[[package]]
name = "futures-channel"
version = "0.3.29"
name = "futures"
version = "0.3.30"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ff4dd66668b557604244583e3e1e1eada8c5c2e96a6d0d6653ede395b78bbacb"
checksum = "645c6916888f6cb6350d2550b80fb63e734897a8498abe35cfb732b6487804b0"
dependencies = [
"futures-channel",
"futures-core",
"futures-io",
"futures-sink",
"futures-task",
"futures-util",
]
[[package]]
name = "futures-channel"
version = "0.3.30"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eac8f7d7865dcb88bd4373ab671c8cf4508703796caa2b1985a9ca867b3fcb78"
dependencies = [
"futures-core",
"futures-sink",
@ -500,9 +527,9 @@ dependencies = [
[[package]]
name = "futures-core"
version = "0.3.29"
version = "0.3.30"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eb1d22c66e66d9d72e1758f0bd7d4fd0bee04cad842ee34587d68c07e45d088c"
checksum = "dfc6580bb841c5a68e9ef15c77ccc837b40a7504914d52e47b8b0e9bbda25a1d"
[[package]]
name = "futures-executor"
@ -528,15 +555,15 @@ dependencies = [
[[package]]
name = "futures-io"
version = "0.3.29"
version = "0.3.30"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8bf34a163b5c4c52d0478a4d757da8fb65cabef42ba90515efee0f6f9fa45aaa"
checksum = "a44623e20b9681a318efdd71c299b6b222ed6f231972bfe2f224ebad6311f0c1"
[[package]]
name = "futures-macro"
version = "0.3.29"
version = "0.3.30"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "53b153fd91e4b0147f4aced87be237c98248656bb01050b96bf3ee89220a8ddb"
checksum = "87750cf4b7a4c0625b1529e4c543c2182106e4dedc60a2a6455e00d212c489ac"
dependencies = [
"proc-macro2",
"quote",
@ -545,21 +572,21 @@ dependencies = [
[[package]]
name = "futures-sink"
version = "0.3.29"
version = "0.3.30"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e36d3378ee38c2a36ad710c5d30c2911d752cb941c00c72dbabfb786a7970817"
checksum = "9fb8e00e87438d937621c1c6269e53f536c14d3fbd6a042bb24879e57d474fb5"
[[package]]
name = "futures-task"
version = "0.3.29"
version = "0.3.30"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "efd193069b0ddadc69c46389b740bbccdd97203899b48d09c5f7969591d6bae2"
checksum = "38d84fa142264698cdce1a9f9172cf383a0c82de1bddcf3092901442c4097004"
[[package]]
name = "futures-util"
version = "0.3.29"
version = "0.3.30"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a19526d624e703a3179b3d322efec918b6246ea0fa51d41124525f00f1cc8104"
checksum = "3d6401deb83407ab3da39eba7e33987a73c3df0c82b4bb5813ee871c19c41d48"
dependencies = [
"futures-core",
"futures-io",
@ -599,6 +626,25 @@ version = "0.28.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6fb8d784f27acf97159b40fc4db5ecd8aa23b9ad5ef69cdd136d3bc80665f0c0"
[[package]]
name = "h2"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "31d030e59af851932b72ceebadf4a2b5986dba4c3b99dd2493f8273a0f151943"
dependencies = [
"bytes",
"fnv",
"futures-core",
"futures-sink",
"futures-util",
"http",
"indexmap",
"slab",
"tokio",
"tokio-util",
"tracing",
]
[[package]]
name = "hashbrown"
version = "0.14.2"
@ -618,30 +664,6 @@ dependencies = [
"hashbrown",
]
[[package]]
name = "headers"
version = "0.3.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "06683b93020a07e3dbcf5f8c0f6d40080d725bea7936fc01ad345c01b97dc270"
dependencies = [
"base64",
"bytes",
"headers-core",
"http",
"httpdate",
"mime",
"sha1",
]
[[package]]
name = "headers-core"
version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e7f66481bfee273957b1f20485a4ff3362987f85b2c236580d81b4eb7a326429"
dependencies = [
"http",
]
[[package]]
name = "heck"
version = "0.4.1"
@ -692,9 +714,9 @@ dependencies = [
[[package]]
name = "http"
version = "0.2.9"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bd6effc99afb63425aff9b05836f029929e345a6148a14b7ecd5ab67af944482"
checksum = "21b9ddb458710bc376481b842f5da65cdf31522de232c1ca8146abce2a358258"
dependencies = [
"bytes",
"fnv",
@ -703,20 +725,26 @@ dependencies = [
[[package]]
name = "http-body"
version = "0.4.5"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d5f38f16d184e36f2408a55281cd658ecbd3ca05cce6d6510a176eca393e26d1"
checksum = "1cac85db508abc24a2e48553ba12a996e87244a0395ce011e62b37158745d643"
dependencies = [
"bytes",
"http",
"pin-project-lite",
]
[[package]]
name = "http-range-header"
version = "0.3.1"
name = "http-body-util"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "add0ab9360ddbd88cfeb3bd9574a1d85cfdfa14db10b3e21d3700dbc4328758f"
checksum = "0475f8b2ac86659c21b64320d5d653f9efe42acd2a4e560073ec61a155a34f1d"
dependencies = [
"bytes",
"futures-core",
"http",
"http-body",
"pin-project-lite",
]
[[package]]
name = "httparse"
@ -738,25 +766,38 @@ checksum = "02296996cb8796d7c6e3bc2d9211b7802812d36999a51bb754123ead7d37d026"
[[package]]
name = "hyper"
version = "0.14.27"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ffb1cfd654a8219eaef89881fdb3bb3b1cdc5fa75ded05d6933b2b382e395468"
checksum = "186548d73ac615b32a73aafe38fb4f56c0d340e110e5a200bcadbaf2e199263a"
dependencies = [
"bytes",
"futures-channel",
"futures-core",
"futures-util",
"h2",
"http",
"http-body",
"httparse",
"httpdate",
"itoa",
"pin-project-lite",
"socket2 0.4.10",
"smallvec",
"tokio",
]
[[package]]
name = "hyper-util"
version = "0.1.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ca38ef113da30126bbff9cd1705f9273e15d45498615d138b0c20279ac7a76aa"
dependencies = [
"bytes",
"futures-util",
"http",
"http-body",
"hyper",
"pin-project-lite",
"socket2",
"tokio",
"tower-service",
"tracing",
"want",
]
[[package]]
@ -853,6 +894,7 @@ checksum = "3c168f8615b12bc01f9c17e2eb0cc07dcae1940121185446edc3744920e8ef45"
dependencies = [
"autocfg",
"scopeguard",
"serde",
]
[[package]]
@ -1438,19 +1480,9 @@ dependencies = [
[[package]]
name = "smallvec"
version = "1.11.1"
version = "1.13.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "942b4a808e05215192e39f4ab80813e599068285906cc91aa64f923db842bd5a"
[[package]]
name = "socket2"
version = "0.4.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9f7916fc008ca5542385b89a3d3ce689953c143e9304a9bf8beec1de48994c0d"
dependencies = [
"libc",
"winapi",
]
checksum = "e6ecd384b10a64542d77071bd64bd7b231f4ed5940fba55e98c3de13824cf3d7"
[[package]]
name = "socket2"
@ -1596,7 +1628,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "864b869fdf56263f4c95c45483191ea0af340f9f3e3e7b4d57a61c7c87a970db"
dependencies = [
"atoi",
"base64",
"base64 0.21.5",
"bitflags 2.4.1",
"byteorder",
"bytes",
@ -1638,7 +1670,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eb7ae0e6a97fb3ba33b23ac2671a5ce6e3cabe003f451abd5a56e7951d975624"
dependencies = [
"atoi",
"base64",
"base64 0.21.5",
"bitflags 2.4.1",
"byteorder",
"crc",
@ -1850,7 +1882,7 @@ dependencies = [
"parking_lot",
"pin-project-lite",
"signal-hook-registry",
"socket2 0.5.5",
"socket2",
"tokio-macros",
"windows-sys",
]
@ -1877,6 +1909,20 @@ dependencies = [
"tokio",
]
[[package]]
name = "tokio-util"
version = "0.7.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5419f34732d9eb6ee4c3578b7989078579b7f039cbbb9ca2c4da015749371e15"
dependencies = [
"bytes",
"futures-core",
"futures-sink",
"pin-project-lite",
"tokio",
"tracing",
]
[[package]]
name = "toml"
version = "0.5.11"
@ -1917,9 +1963,9 @@ dependencies = [
[[package]]
name = "tower-cookies"
version = "0.9.0"
version = "0.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "40f38d941a2ffd8402b36e02ae407637a9caceb693aaf2edc910437db0f36984"
checksum = "4fd0118512cf0b3768f7fcccf0bef1ae41d68f2b45edc1e77432b36c97c56c6d"
dependencies = [
"async-trait",
"axum-core",
@ -1934,17 +1980,15 @@ dependencies = [
[[package]]
name = "tower-http"
version = "0.4.4"
version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "61c5bb1d698276a2443e5ecfabc1008bf15a36c12e6a7176e7bf089ea9131140"
checksum = "1e9cd434a998747dd2c4276bc96ee2e0c7a2eadf3cae88e52be55a05fa9053f5"
dependencies = [
"bitflags 2.4.1",
"bytes",
"futures-core",
"futures-util",
"http",
"http-body",
"http-range-header",
"http-body-util",
"pin-project-lite",
"tower-layer",
"tower-service",
@ -1963,6 +2007,57 @@ version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b6bc1c9ce2b5135ac7f93c72918fc37feb872bdc6a5533a8b85eb4b86bfdae52"
[[package]]
name = "tower-sessions"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "989b4d77286a7fb96b094b9e62c4524cebe7bb720769b41588ebaac5d9a787ca"
dependencies = [
"async-trait",
"http",
"time",
"tokio",
"tower-cookies",
"tower-layer",
"tower-service",
"tower-sessions-core",
"tower-sessions-memory-store",
"tracing",
]
[[package]]
name = "tower-sessions-core"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c0a9049748900860b01f92d3decf5fec71b9d75008f07732956288b003a2282f"
dependencies = [
"async-trait",
"axum-core",
"base64 0.22.0",
"futures",
"http",
"parking_lot",
"rand",
"serde",
"serde_json",
"thiserror",
"time",
"tokio",
"tracing",
]
[[package]]
name = "tower-sessions-memory-store"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "36baf499920bb861ec9fa929a1f879cec0759af987c6360bd65226bc484ab846"
dependencies = [
"async-trait",
"time",
"tokio",
"tower-sessions-core",
]
[[package]]
name = "tracing"
version = "0.1.40"
@ -2025,12 +2120,6 @@ dependencies = [
"tracing-log",
]
[[package]]
name = "try-lock"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3528ecfd12c466c6f163363caf2d02a71161dd5e1cc6ae7b34207ea2d42d81ed"
[[package]]
name = "typenum"
version = "1.17.0"
@ -2109,15 +2198,6 @@ version = "0.9.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
[[package]]
name = "want"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e"
dependencies = [
"try-lock",
]
[[package]]
name = "wasi"
version = "0.11.0+wasi-snapshot-preview1"

View File

@ -6,16 +6,17 @@ edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
axum = {version = "0.6.20", features = [ "default", "headers" ]}
axum = {version = "0.7.4", features = [ "default", "form" ]}
tokio = { version = "1.0", features = ["full"] }
sqlx = { version = "0.7", features = ["runtime-tokio", "sqlite"] }
tower-cookies = "0.9.0"
tower-cookies = "0.10.0"
serde = "1.0.190"
uuid = { version = "1.5.0", features = ["v4", "fast-rng"] }
totp-rs = "5.4.0"
tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
tower-http = { version = "0.4.4", features = ["trace"] }
tower-http = { version = "0.5", features = ["trace"] }
tracing = "0.1.40"
askama = "0.10"
minijinja = "1.0.9"
once_cell = "1.18.0"
tower-sessions = "0.11.0"

BIN
auth.db

Binary file not shown.

0
src/config.rs Normal file
View File

0
src/controllers/mod.rs Normal file
View File

26
src/entities/mod.rs Normal file
View File

@ -0,0 +1,26 @@
use serde::Deserialize;
use std::time::Instant;
use std::collections::HashMap;
use tokio::sync::Mutex;
use uuid::Uuid;
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,
}
#[derive(Deserialize, sqlx::FromRow, Debug)]
pub struct User {
#[sqlx(rename = "NAME")]
pub username: String,
#[sqlx(rename = "KEY")]
pub uid: String,
}

View File

@ -2,32 +2,24 @@ use axum::extract::Query;
use axum::http::{HeaderMap, HeaderValue};
use axum::response::{Html, Redirect};
use axum::{extract::State, http::StatusCode, response::IntoResponse, routing::get, Form, Router};
use entities::*;
use minijinja::{context, Environment};
use serde::Deserialize;
use once_cell::sync::Lazy;
use services::{auth, gc_task, login, login_page, login_with_passkey};
use sqlx::sqlite::SqlitePool;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use std::{borrow::BorrowMut, env};
use std::{collections::HashMap, str::FromStr};
use tokio::sync::Mutex;
use tower_cookies::{Cookie, CookieManagerLayer, Cookies};
use tower_http::trace::{self, TraceLayer};
use tower_sessions::{Expiry, MemoryStore, Session, SessionManagerLayer};
use tracing::Level;
use uuid::Uuid;
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,
}
use once_cell::sync::Lazy;
pub mod controllers;
pub mod entities;
pub mod services;
static COOKIE_NAME: Lazy<String> =
Lazy::new(|| env::var("COOKIE_NAME").unwrap_or("aaron_auth".to_string()));
const SESSION_ACTIVE_TIME: Lazy<u64> = Lazy::new(|| {
@ -36,11 +28,17 @@ const SESSION_ACTIVE_TIME: Lazy<u64> = Lazy::new(|| {
.and_then(|value| value.parse().ok())
.unwrap_or(600)
});
const COOKIE_DOMAIN: Lazy<String> =
Lazy::new(|| env::var("DOMAIN").ok().unwrap_or(".aaronhu.cn".to_owned()));
const LOGIN_PAGE_HTML: &str = include_str!("../loginpage.html");
#[tokio::main]
async fn main() {
// 初始化日志记录器
tracing_subscriber::fmt::init();
let session_store = MemoryStore::default();
let session_layer = SessionManagerLayer::new(session_store)
.with_secure(false)
.with_expiry(Expiry::OnSessionEnd);
let pool = SqlitePool::connect(&env::var("DATABASE_URL").expect("DB URL NOT SPECIFIED"))
.await
.expect("DB OPEN FAILURE"); // our router
@ -52,155 +50,19 @@ async fn main() {
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)),
);
)
.with_state(state.clone())
.layer(CookieManagerLayer::new())
.layer(session_layer);
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 let std::collections::hash_map::Entry::Occupied(mut e) = locked.entry(s) {
// FIX, when accessed /auth with correct cookie, the cookie's expiration is delayed
let Some(v) = Some(e.insert(Instant::now() + Duration::from_secs(*SESSION_ACTIVE_TIME))) else {
tracing::info!("session:{} extended", session_token.value());
return StatusCode::UNAUTHORIZED;
};
if Instant::now() < v {
return StatusCode::OK;
}
}
}
StatusCode::UNAUTHORIZED
}
async fn login(
State(state): State<Arc<ServerState>>,
cookies: Cookies,
Query(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,
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"));
}
}
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()),
);
}
}
}
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;
}
}
false
}
async fn gc(state: Arc<ServerState>) -> Result<(), String> {
let mut locked = state.session.lock().await;
let current_time = Instant::now();
tracing::info!("before gc ,active Sessions {:?}", locked);
locked.borrow_mut().retain(|_, v| *v > current_time);
tracing::info!("gc fired,active Sessions {:?}", locked);
Ok(())
let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
axum::serve(listener, app).await.unwrap();
}

201
src/services/mod.rs Normal file
View File

@ -0,0 +1,201 @@
use axum::extract::Query;
use axum::http::{HeaderMap, HeaderValue};
use axum::response::{Html, Redirect};
use axum::{extract::State, http::StatusCode, response::IntoResponse, routing::get, Form, Router};
use minijinja::{context, Environment};
use serde::Deserialize;
use sqlx::sqlite::SqlitePool;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use std::{borrow::BorrowMut, env};
use std::{collections::HashMap, str::FromStr};
use tokio::sync::Mutex;
use tower_cookies::{Cookie, CookieManagerLayer, Cookies};
use tower_http::trace::{self, TraceLayer};
use tracing::Level;
use uuid::Uuid;
use crate::{ServerState, UserLoginForm, COOKIE_DOMAIN, COOKIE_NAME, LOGIN_PAGE_HTML, SESSION_ACTIVE_TIME};
pub 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
pub 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 let std::collections::hash_map::Entry::Occupied(mut e) = locked.entry(s) {
// FIX, when accessed /auth with correct cookie, the cookie's expiration is delayed
let Some(v) = Some(e.insert(Instant::now() + Duration::from_secs(*SESSION_ACTIVE_TIME))) else {
tracing::info!("session:{} extended", session_token.value());
return StatusCode::UNAUTHORIZED;
};
if Instant::now() < v {
return StatusCode::OK;
}
}
}
StatusCode::UNAUTHORIZED
}
pub async fn login(
State(state): State<Arc<ServerState>>,
cookies: Cookies,
Query(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,
Instant::now() + Duration::from_secs(*SESSION_ACTIVE_TIME),
);
let mut new_cookie = Cookie::new(&*COOKIE_NAME, s.to_string());
new_cookie.set_domain(COOKIE_DOMAIN.to_string());
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"));
}
}
Err((StatusCode::BAD_GATEWAY, "unreachable"))
}
pub 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()),
);
}
}
}
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;
}
}
false
}
pub async fn gc(state: Arc<ServerState>) -> Result<(), String> {
let mut locked = state.session.lock().await;
let current_time = Instant::now();
tracing::info!("before gc ,active Sessions {:?}", locked);
locked.borrow_mut().retain(|_, v| *v > current_time);
tracing::info!("gc fired,active Sessions {:?}", locked);
Ok(())
}
pub async fn login_with_passkey(
State(state): State<Arc<ServerState>>,
cookies: Cookies,
Query(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,
Instant::now() + Duration::from_secs(*SESSION_ACTIVE_TIME),
);
let mut new_cookie = Cookie::new(&*COOKIE_NAME, s.to_string());
new_cookie.set_domain(COOKIE_DOMAIN.to_string());
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"));
}
}
Err((StatusCode::BAD_GATEWAY, "unreachable"))
}