From 82721e94982fb17c6f5a9dab6b90373faf7ae084 Mon Sep 17 00:00:00 2001
From: Pierre-Nicolas Watin-Augouard <pnwa@icloud.com>
Date: Sat, 6 Apr 2024 20:46:27 +0200
Subject: [PATCH] chore: refactor refresh token validation / family
 invalidation

---
 src/handlers/auth/mod.rs    | 36 ++++++++++++++++++++----------------
 src/handlers/auth/tokens.rs |  5 ++++-
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/src/handlers/auth/mod.rs b/src/handlers/auth/mod.rs
index 32726bd..2b4d97d 100644
--- a/src/handlers/auth/mod.rs
+++ b/src/handlers/auth/mod.rs
@@ -2,7 +2,6 @@ mod sign_in;
 mod sign_up;
 mod tokens;
 
-use anyhow::Context;
 use axum::{routing::post, Router};
 use chrono::{DateTime, Utc};
 use jsonwebtoken::{decode, encode, Header, Validation};
@@ -10,7 +9,7 @@ use serde::{de::DeserializeOwned, Deserialize, Serialize};
 use sqlx::PgPool;
 use uuid::Uuid;
 
-use crate::{error::AppError, settings::JWT_CONFIG};
+use crate::settings::JWT_CONFIG;
 
 pub fn auth_router() -> Router {
     Router::new()
@@ -117,7 +116,7 @@ impl RefreshToken {
         Ok(self)
     }
 
-    pub async fn validate(self, pool: &PgPool) -> Result<Self, AppError> {
+    pub async fn validate(self, pool: &PgPool) -> Result<Option<Self>, sqlx::Error> {
         let result = sqlx::query!(
             r#"
                 SELECT * FROM refresh_tokens WHERE jit = $1; 
@@ -125,23 +124,28 @@ impl RefreshToken {
             self.claims().jit
         )
         .fetch_optional(pool)
-        .await
-        .context("Failed to fetch execute query")?;
+        .await?;
 
         if result.is_none() {
-            sqlx::query!(
-                r#"
-                    DELETE FROM refresh_tokens WHERE family = $1;
-                "#,
-                self.claims().family
-            )
-            .execute(pool)
-            .await?;
-
-            return Err(AppError::InvalidRefreshToken);
+            self.invalidate_family(pool).await?;
+
+            return Ok(None);
         }
 
-        Ok(self)
+        Ok(Some(self))
+    }
+
+    pub async fn invalidate_family(&self, pool: &PgPool) -> Result<(), sqlx::Error> {
+        sqlx::query!(
+            r#"
+                DELETE FROM refresh_tokens WHERE family = $1;
+            "#,
+            self.claims().family
+        )
+        .execute(pool)
+        .await?;
+
+        Ok(())
     }
 }
 
diff --git a/src/handlers/auth/tokens.rs b/src/handlers/auth/tokens.rs
index c461608..0eb7187 100644
--- a/src/handlers/auth/tokens.rs
+++ b/src/handlers/auth/tokens.rs
@@ -26,9 +26,12 @@ async fn refresh_tokens(refresh_token: &str, pool: &PgPool) -> Result<TokensPair
 
     let user_id = refresh_token_claims.sub;
     let family = refresh_token_claims.family;
+
     RefreshToken::from(refresh_token_claims)
         .validate(pool)
-        .await?;
+        .await
+        .context("Failed to execute query.")?
+        .ok_or(AppError::InvalidRefreshToken)?;
 
     let refresh_token = RefreshToken::new(user_id, family).save(pool).await?;