Add a timer to revoke all expired tokens( Fix #279

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<usize, SqliteError> {
+        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<usize, PixivDownloaderDbError> {
+        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,

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<usize, PixivDownloaderDbError>;
+    #[cfg(feature = "server")]
     /// Set a user's information by ID
     /// * `id`: The user's ID
     /// * `name`: The user's name

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] = [

src/server/auth/token.rs

@@ -194,3 +194,8 @@ impl MatchRoute<Body, Body> for AuthTokenRoute {
         }
     }
 }
+
+pub async fn revoke_expired_tokens(ctx: Arc<ServerContext>) -> Result<(), PixivDownloaderError> {
+    ctx.db.revoke_expired_tokens().await?;
+    Ok(())
+}

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"))]

src/server/service.rs

@@ -64,9 +64,9 @@ pub struct PixivDownloaderMakeSvc {
 }
 
 impl PixivDownloaderMakeSvc {
-    pub async fn new() -> Self {
+    pub fn new(ctx: &Arc<ServerContext>) -> Self {
         Self {
-            context: Arc::new(ServerContext::default().await),
+            context: Arc::clone(ctx),
             routes: Arc::new(ServerRoutes::new()),
         }
     }
@@ -93,5 +93,8 @@ impl<T> Service<T> for PixivDownloaderMakeSvc {
 pub async fn start_server(
     addr: &SocketAddr,
 ) -> Result<Server<AddrIncoming, PixivDownloaderMakeSvc>, 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)
 }

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<ServerContext>) {
+    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);
+            }
+        }
+    }
+}

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<T> {
     /// Current running task