refactor(auth): Session revocation

auth/src/features/api/v1/auth/reset_password.rs

@@ -50,10 +50,6 @@ pub async fn controller_with_password(
     jar: CookieJar,
     Json(req): Json<ResetPasswordWithPasswordService>,
 ) -> (CookieJar, Response) {
-    if let Err(err) = req.reset_password(&state.db, &login.user.0).await {
-        return (jar, app_error_to_response(err));
-    }
-
     let mut redis = match state.redis.get_multiplexed_tokio_connection().await {
         Ok(redis) => redis,
         Err(err) => {
@@ -71,12 +67,34 @@ pub async fn controller_with_password(
         },
     };
 
-    if let Err(err) = SessionService::delete(&mut redis, &login.session).await {
+    if let Err(err) = req.reset_password(&login.user.0).await {
+        return (jar, app_error_to_response(err));
+    }
+
+    if let Err(err) = SessionService::delete_all_by_uid(&mut redis, login.user.0.uid).await {
+        let jar = jar.remove_session_cookie();
         return (
             jar,
             app_error_to_response(AppError::infra(
                 AppErrorKind::InternalError,
-                "auth.controller.reset_password_password.delete_session",
+                "auth.controller.reset_password_password.delete_all_sessions",
+                err,
+            )),
+        );
+    }
+
+    if let Err(err) = login
+        .user
+        .0
+        .update_password(state.db.as_ref(), req.new_password.clone())
+        .await
+    {
+        let jar = jar.remove_session_cookie();
+        return (
+            jar,
+            app_error_to_response(AppError::infra(
+                AppErrorKind::InternalError,
+                "auth.reset_password.password.update_password",
                 err,
             )),
         );
@@ -158,6 +176,14 @@ impl ResetPasswordWithTokenService {
             ));
         }
 
+        if let Err(err) = SessionService::delete_all_by_uid(redis, uid).await {
+            return Err(AppError::infra(
+                AppErrorKind::InternalError,
+                "auth.reset_password.token.delete_all_sessions",
+                err,
+            ));
+        }
+
         if let Err(err) = user.update_password(conn, self.new_password.clone()).await {
             return Err(AppError::infra(
                 AppErrorKind::InternalError,
@@ -171,11 +197,7 @@ impl ResetPasswordWithTokenService {
 }
 
 impl ResetPasswordWithPasswordService {
-    pub async fn reset_password(
-        &self,
-        conn: &DatabaseConnection,
-        user: &users::Model,
-    ) -> Result<(), AppError> {
+    pub async fn reset_password(&self, user: &users::Model) -> Result<(), AppError> {
         if !user.check_password(self.old_password.clone()) {
             return Err(AppError::biz(
                 AppErrorKind::Unauthorized,
@@ -199,14 +221,6 @@ impl ResetPasswordWithPasswordService {
             .with_detail(err.to_string())
         })?;
 
-        if let Err(err) = user.update_password(conn, self.new_password.clone()).await {
-            return Err(AppError::infra(
-                AppErrorKind::InternalError,
-                "auth.reset_password.password.update_password",
-                err,
-            ));
-        }
-
         Ok(())
     }
 }

auth/src/features/api/v1/users/delete.rs

@@ -41,18 +41,18 @@ pub async fn controller(
         },
     };
 
-    let session = login.session;
+    let uid = login.user.0.uid;
 
     if let Err(err) = req.delete(&state.db, login.user.0).await {
         return (jar, app_error_to_response(err));
     }
 
-    if let Err(err) = SessionService::delete(&mut redis, &session).await {
+    if let Err(err) = SessionService::delete_all_by_uid(&mut redis, uid).await {
         return (
             jar,
             app_error_to_response(AppError::infra(
                 AppErrorKind::InternalError,
-                "users.controller.delete.delete_session",
+                "users.controller.delete.delete_all_sessions",
                 err,
             )),
         );
@@ -68,10 +68,34 @@ pub async fn controller_by_uid(
     State(state): State<AppState>,
     Path(uid): Path<i32>,
 ) -> Response {
+    let mut redis = match state.redis.get_multiplexed_tokio_connection().await {
+        Ok(redis) => redis,
+        Err(err) => {
+            return app_error_to_response(
+                AppError::infra(
+                    AppErrorKind::InternalError,
+                    "users.controller.admin_delete.redis",
+                    err,
+                )
+                .with_detail("Unable to connect to redis"),
+            );
+        },
+    };
+
     let service = AdminDeleteService;
 
     match service.delete(&state.db, uid, login.level).await {
-        Ok(()) => ResponseCode::OK.into(),
+        Ok(()) => {
+            if let Err(err) = SessionService::delete_all_by_uid(&mut redis, uid).await {
+                return app_error_to_response(AppError::infra(
+                    AppErrorKind::InternalError,
+                    "users.controller.admin_delete.delete_all_sessions",
+                    err,
+                ));
+            }
+
+            ResponseCode::OK.into()
+        },
         Err(err) => app_error_to_response(err),
     }
 }

crates/token/src/services/password_reset.rs

@@ -15,7 +15,7 @@ pub struct PasswordResetTokenService;
 impl TokenStore for PasswordResetTokenService {
     type Payload = PasswordResetToken;
 
-    const PREFIX: &'static str = "password-reset";
+    const PREFIX: &'static str = "madoka::auth::password-reset";
 }
 
 impl PasswordResetTokenService {

crates/token/src/services/register_token.rs

@@ -92,11 +92,11 @@ pub struct RegisterTokenLease {
 impl TokenStore for RegisterTokenService {
     type Payload = RegisterToken;
 
-    const PREFIX: &'static str = "register-token";
+    const PREFIX: &'static str = "madoka::auth::register-token";
 }
 
 impl RegisterTokenService {
-    const INDEX_PREFIX: &'static str = "register-token-index";
+    const INDEX_PREFIX: &'static str = "madoka::auth::register-token-user";
 
     fn index_key(uid: i32) -> String {
         format!("{}::{uid}", Self::INDEX_PREFIX)

crates/token/src/services/session.rs

@@ -4,6 +4,68 @@ use serde::{Deserialize, Serialize};
 
 use crate::{TokenError, TokenStore, backend::RedisTokenBackend};
 
+const SESSION_PREFIX: &str = "madoka::auth::session";
+const USER_SESSION_INDEX_PREFIX: &str = "madoka::auth::session-user";
+const CREATE_SESSION_SCRIPT: &str = r#"
+local session_key = KEYS[1]
+local index_key = KEYS[2]
+local payload = ARGV[1]
+local ttl_secs = tonumber(ARGV[2])
+local sid = ARGV[3]
+local exp = tonumber(ARGV[4])
+local now = tonumber(ARGV[5])
+
+redis.call('ZADD', index_key, exp, sid)
+redis.call('ZREMRANGEBYSCORE', index_key, '-inf', now)
+
+if redis.call('ZCARD', index_key) == 0 then
+    redis.call('DEL', index_key)
+else
+    local current_ttl = redis.call('TTL', index_key)
+    if current_ttl == -1 or (current_ttl >= 0 and current_ttl < ttl_secs) then
+        redis.call('EXPIRE', index_key, ttl_secs)
+    end
+end
+
+redis.call('SETEX', session_key, ttl_secs, payload)
+
+return 1
+"#;
+const DELETE_SESSION_SCRIPT: &str = r#"
+local session_key = KEYS[1]
+local sid = ARGV[1]
+local index_prefix = ARGV[2]
+
+local payload = redis.call('GETDEL', session_key)
+if not payload then
+    return 0
+end
+
+local ok, decoded = pcall(cjson.decode, payload)
+if ok and decoded and decoded.uid ~= nil then
+    local index_key = index_prefix .. '::' .. tostring(decoded.uid)
+    redis.call('ZREM', index_key, sid)
+end
+
+return 1
+"#;
+const DELETE_ALL_BY_UID_SCRIPT: &str = r#"
+local index_key = KEYS[1]
+local session_prefix = ARGV[1]
+local now = tonumber(ARGV[2])
+
+redis.call('ZREMRANGEBYSCORE', index_key, '-inf', now)
+local sids = redis.call('ZRANGE', index_key, 0, -1)
+
+for _, sid in ipairs(sids) do
+    local session_key = session_prefix .. '::' .. sid
+    redis.call('DEL', session_key)
+end
+
+redis.call('DEL', index_key)
+return #sids
+"#;
+
 #[derive(Serialize, Deserialize, Debug, Clone, Copy)]
 pub struct Session {
     pub uid: i32,
@@ -38,7 +100,7 @@ pub struct SessionService;
 impl TokenStore for SessionService {
     type Payload = Session;
 
-    const PREFIX: &'static str = "session";
+    const PREFIX: &'static str = SESSION_PREFIX;
 }
 
 impl SessionService {
@@ -48,14 +110,55 @@ impl SessionService {
         ttl_secs: u64,
     ) -> Result<String, TokenError> {
         let session = generate_session(uid, ttl_secs);
-        <Self as TokenStore>::issue(redis, &session, ttl_secs).await
+        let session_id = uuid::Uuid::new_v4().to_string();
+        let session_payload = serde_json::to_string(&session)?;
+        let session_key = format!("{SESSION_PREFIX}::{session_id}");
+        let index_key = format!("{USER_SESSION_INDEX_PREFIX}::{uid}");
+        let now = Utc::now().timestamp() as usize;
+
+        let script = redis::Script::new(CREATE_SESSION_SCRIPT);
+        let mut invocation = script.prepare_invoke();
+        invocation
+            .key(session_key)
+            .key(index_key)
+            .arg(session_payload)
+            .arg(ttl_secs)
+            .arg(&session_id)
+            .arg(session.exp)
+            .arg(now);
+        let _: i32 = invocation.invoke_async(redis).await?;
+
+        Ok(session_id)
     }
 
     pub async fn delete(
         redis: &mut MultiplexedConnection,
         session_id: &str,
     ) -> Result<(), TokenError> {
-        <Self as TokenStore>::revoke(redis, session_id).await
+        let session_key = format!("{SESSION_PREFIX}::{session_id}");
+        let script = redis::Script::new(DELETE_SESSION_SCRIPT);
+        let mut invocation = script.prepare_invoke();
+        invocation
+            .key(session_key)
+            .arg(session_id)
+            .arg(USER_SESSION_INDEX_PREFIX);
+        let _: i32 = invocation.invoke_async(redis).await?;
+
+        Ok(())
+    }
+
+    pub async fn delete_all_by_uid(
+        redis: &mut MultiplexedConnection,
+        uid: i32,
+    ) -> Result<(), TokenError> {
+        let index_key = format!("{USER_SESSION_INDEX_PREFIX}::{uid}");
+        let now = Utc::now().timestamp() as usize;
+
+        let script = redis::Script::new(DELETE_ALL_BY_UID_SCRIPT);
+        let mut invocation = script.prepare_invoke();
+        invocation.key(index_key).arg(SESSION_PREFIX).arg(now);
+        let _: i32 = invocation.invoke_async(redis).await?;
+        Ok(())
     }
 
     pub async fn resolve(