From fcff540220d42eb5370d6fb7ffec38057ac0f9aa Mon Sep 17 00:00:00 2001 From: lifegpc Date: Thu, 29 Sep 2022 01:30:03 +0000 Subject: [PATCH] Add a timer to revoke all expired tokens( Fix #279 --- src/db/sqlite/db.rs | 23 ++++++++++++++++++++++- src/db/traits.rs | 4 ++++ src/server/auth/mod.rs | 2 +- src/server/auth/token.rs | 5 +++++ src/server/mod.rs | 2 ++ src/server/service.rs | 9 ++++++--- src/server/timer.rs | 23 +++++++++++++++++++++++ src/task_manager.rs | 17 +++++++++++++++++ 8 files changed, 80 insertions(+), 5 deletions(-) create mode 100644 src/server/timer.rs diff --git a/src/db/sqlite/db.rs b/src/db/sqlite/db.rs index e179efc..e9e9aa1 100644 --- a/src/db/sqlite/db.rs +++ b/src/db/sqlite/db.rs @@ -99,6 +99,7 @@ impl PixivDownloaderSqlite { Ok(()) } + #[cfg(feature = "server")] fn _add_token( tx: &Transaction, user_id: u64, @@ -113,6 +114,7 @@ impl PixivDownloaderSqlite { Ok(()) } + #[cfg(feature = "server")] fn _add_user( tx: &Transaction, name: &str, @@ -159,7 +161,11 @@ impl PixivDownloaderSqlite { tx.execute(TOKEN_TABLE, [])?; } if db_version < [1, 0, 0, 4] { - tx.execute("DROP TABLE authors, files, tags, token, users;", [])?; + tx.execute("DROP TABLE authors;", [])?; + tx.execute("DROP TABLE files;", [])?; + tx.execute("DROP TABLE tags;", [])?; + tx.execute("DROP TABLE token;", [])?; + tx.execute("DROP TABLE users;", [])?; tx.execute(AUTHORS_TABLE, [])?; tx.execute(FILES_TABLE, [])?; tx.execute(TAGS_TABLE, [])?; @@ -317,6 +323,12 @@ impl PixivDownloaderSqlite { } } + #[cfg(feature = "server")] + fn _revoke_expired_tokens(ts: &Transaction) -> Result { + let now = Utc::now(); + Ok(ts.execute("DELETE FROM token WHERE expired_at < ?;", [now])?) + } + #[cfg(feature = "server")] async fn _set_user( &self, @@ -515,6 +527,15 @@ impl PixivDownloaderDb for PixivDownloaderSqlite { Ok(()) } + #[cfg(feature = "server")] + async fn revoke_expired_tokens(&self) -> Result { + let mut db = self.db.lock().await; + let mut tx = db.transaction()?; + let size = Self::_revoke_expired_tokens(&mut tx)?; + tx.commit()?; + Ok(size) + } + #[cfg(feature = "server")] async fn set_user( &self, diff --git a/src/db/traits.rs b/src/db/traits.rs index 3252954..68577e7 100644 --- a/src/db/traits.rs +++ b/src/db/traits.rs @@ -72,6 +72,10 @@ pub trait PixivDownloaderDb { /// Initialize the database (create tables, migrate data, etc.) async fn init(&self) -> Result<(), PixivDownloaderDbError>; #[cfg(feature = "server")] + /// Remove all expired tokens + /// Return the number of removed tokens + async fn revoke_expired_tokens(&self) -> Result; + #[cfg(feature = "server")] /// Set a user's information by ID /// * `id`: The user's ID /// * `name`: The user's name diff --git a/src/server/auth/mod.rs b/src/server/auth/mod.rs index cc26ca8..bf77148 100644 --- a/src/server/auth/mod.rs +++ b/src/server/auth/mod.rs @@ -5,7 +5,7 @@ pub mod user; pub use pubkey::{AuthPubkeyContext, AuthPubkeyRoute, RSAKey}; pub use status::{AuthStatusContext, AuthStatusRoute}; -pub use token::{AuthTokenContext, AuthTokenRoute}; +pub use token::{revoke_expired_tokens, AuthTokenContext, AuthTokenRoute}; pub use user::{AuthUserContext, AuthUserRoute}; const PASSWORD_SALT: [u8; 64] = [ diff --git a/src/server/auth/token.rs b/src/server/auth/token.rs index 513bb3a..ee61543 100644 --- a/src/server/auth/token.rs +++ b/src/server/auth/token.rs @@ -194,3 +194,8 @@ impl MatchRoute for AuthTokenRoute { } } } + +pub async fn revoke_expired_tokens(ctx: Arc) -> Result<(), PixivDownloaderError> { + ctx.db.revoke_expired_tokens().await?; + Ok(()) +} diff --git a/src/server/mod.rs b/src/server/mod.rs index f9a7633..46ecebc 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -13,6 +13,8 @@ pub mod result; pub mod route; /// Services pub mod service; +/// Tasks invoked by timer +pub mod timer; /// Traits pub mod traits; #[cfg(all(test, feature = "db_sqlite"))] diff --git a/src/server/service.rs b/src/server/service.rs index 64dfeba..adec2ed 100644 --- a/src/server/service.rs +++ b/src/server/service.rs @@ -64,9 +64,9 @@ pub struct PixivDownloaderMakeSvc { } impl PixivDownloaderMakeSvc { - pub async fn new() -> Self { + pub fn new(ctx: &Arc) -> Self { Self { - context: Arc::new(ServerContext::default().await), + context: Arc::clone(ctx), routes: Arc::new(ServerRoutes::new()), } } @@ -93,5 +93,8 @@ impl Service for PixivDownloaderMakeSvc { pub async fn start_server( addr: &SocketAddr, ) -> Result, hyper::Error> { - Ok(Server::try_bind(addr)?.serve(PixivDownloaderMakeSvc::new().await)) + let ctx = Arc::new(ServerContext::default().await); + let ser = Server::try_bind(addr)?.serve(PixivDownloaderMakeSvc::new(&ctx)); + tokio::spawn(super::timer::start_timer(ctx)); + Ok(ser) } diff --git a/src/server/timer.rs b/src/server/timer.rs new file mode 100644 index 0000000..70a7c8a --- /dev/null +++ b/src/server/timer.rs @@ -0,0 +1,23 @@ +use super::auth::*; +use super::context::ServerContext; +use crate::task_manager::{MaxCount, TaskManager}; +use std::sync::Arc; +use tokio::time::{interval_at, Duration, Instant}; + +pub async fn start_timer(ctx: Arc) { + let mut interval = interval_at(Instant::now(), Duration::from_secs(60)); + let task_count = Arc::new(futures_util::lock::Mutex::new(0usize)); + let max_count = MaxCount::new(8); + let tasks = TaskManager::new(task_count, max_count); + loop { + interval.tick().await; + tasks.add_task(revoke_expired_tokens(ctx.clone())).await; + tasks.join().await; + for task in tasks.take_finished_tasks() { + let re = task.await; + if let Ok(Err(e)) = re { + println!("Timer task error: {}", e); + } + } + } +} diff --git a/src/task_manager.rs b/src/task_manager.rs index 493f9c4..8361286 100644 --- a/src/task_manager.rs +++ b/src/task_manager.rs @@ -55,6 +55,23 @@ impl GetMaxCount for MaxDownloadPostTasks { } } +#[derive(Clone, Debug)] +pub struct MaxCount { + max_count: usize, +} + +impl MaxCount { + pub fn new(max_count: usize) -> Self { + MaxCount { max_count } + } +} + +impl GetMaxCount for MaxCount { + fn get_max_count(&self) -> usize { + self.max_count + } +} + /// Task manager pub struct TaskManager { /// Current running task