add tests, format files

TabbyML · Dec 2, 2023 · b6b5394 · b6b5394
1 parent 0ebba25
commit b6b5394
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 20 deletions.
diff --git a/ee/tabby-webserver/src/schema/auth.rs b/ee/tabby-webserver/src/schema/auth.rs
@@ -256,7 +256,10 @@ pub trait AuthenticationService: Send + Sync {
         password: String,
     ) -> std::result::Result<TokenAuthResponse, TokenAuthError>;
 
-    async fn refresh_token(&self, refresh_token: String) -> std::result::Result<RefreshTokenResponse, RefreshTokenError>;
+    async fn refresh_token(
+        &self,
+        refresh_token: String,
+    ) -> std::result::Result<RefreshTokenResponse, RefreshTokenError>;
     async fn verify_token(&self, access_token: String) -> Result<VerifyTokenResponse>;
     async fn is_admin_initialized(&self) -> Result<bool>;
 

diff --git a/ee/tabby-webserver/src/schema/mod.rs b/ee/tabby-webserver/src/schema/mod.rs
@@ -17,10 +17,12 @@ use self::{
     worker::WorkerService,
 };
 use crate::schema::{
-    auth::{RegisterResponse, TokenAuthResponse, VerifyTokenResponse},
+    auth::{
+        RefreshTokenError, RefreshTokenResponse, RegisterResponse, TokenAuthResponse,
+        VerifyTokenResponse,
+    },
     worker::Worker,
 };
-use crate::schema::auth::{RefreshTokenError, RefreshTokenResponse};
 
 pub trait ServiceLocator: Send + Sync {
     fn auth(&self) -> &dyn AuthenticationService;
@@ -138,7 +140,7 @@ impl Mutation {
 
     async fn refresh_token(
         ctx: &Context,
-        refresh_token: String
+        refresh_token: String,
     ) -> Result<RefreshTokenResponse, RefreshTokenError> {
         ctx.locator.auth().refresh_token(refresh_token).await
     }

diff --git a/ee/tabby-webserver/src/service/auth.rs b/ee/tabby-webserver/src/service/auth.rs
@@ -8,8 +8,11 @@ use async_trait::async_trait;
 use validator::Validate;
 
 use super::db::DbConn;
-use crate::schema::auth::{generate_jwt, validate_jwt, AuthenticationService, Claims, Invitation, RefreshTokenResponse, RegisterError, RegisterResponse, TokenAuthError, TokenAuthResponse, UserInfo, VerifyTokenResponse, RefreshTokenError};
-use crate::schema::auth::generate_refresh_token;
+use crate::schema::auth::{
+    generate_jwt, generate_refresh_token, validate_jwt, AuthenticationService, Claims, Invitation,
+    RefreshTokenError, RefreshTokenResponse, RegisterError, RegisterResponse, TokenAuthError,
+    TokenAuthResponse, UserInfo, VerifyTokenResponse,
+};
 
 /// Input parameters for register mutation
 /// `validate` attribute is used to validate the input parameters
@@ -143,10 +146,9 @@ impl AuthenticationService for DbConn {
             .await?;
         let user = self.get_user(id).await?.unwrap();
 
-        let (refresh_token, expires_in) = generate_refresh_token(
-            chrono::Utc::now().timestamp()
-        );
-        self.create_refresh_token(id, &refresh_token, expires_in).await?;
+        let (refresh_token, expires_in) = generate_refresh_token(chrono::Utc::now().timestamp());
+        self.create_refresh_token(id, &refresh_token, expires_in)
+            .await?;
 
         let Ok(access_token) = generate_jwt(Claims::new(UserInfo::new(
             user.email.clone(),
@@ -180,7 +182,8 @@ impl AuthenticationService for DbConn {
             Some(refresh_token) => refresh_token.token,
             None => {
                 let (token, expires_in) = generate_refresh_token(utc_ts);
-                self.create_refresh_token(user.id, &token, expires_in).await?;
+                self.create_refresh_token(user.id, &token, expires_in)
+                    .await?;
                 token
             }
         };
@@ -196,7 +199,10 @@ impl AuthenticationService for DbConn {
         Ok(resp)
     }
 
-    async fn refresh_token(&self, token: String) -> std::result::Result<RefreshTokenResponse, RefreshTokenError> {
+    async fn refresh_token(
+        &self,
+        token: String,
+    ) -> std::result::Result<RefreshTokenResponse, RefreshTokenError> {
         let Some(refresh_token) = self.get_refresh_token(&token).await? else {
             return Err(RefreshTokenError::InvalidRefreshToken);
         };
@@ -392,7 +398,10 @@ mod tests {
         let conn = DbConn::new_in_memory().await.unwrap();
         let reg_resp = register_admin_user(&conn).await;
 
-        let resp = conn.refresh_token(reg_resp.refresh_token.clone()).await.unwrap();
+        let resp = conn
+            .refresh_token(reg_resp.refresh_token.clone())
+            .await
+            .unwrap();
         // refreshed access token should be valid
         assert!(validate_jwt(&resp.access_token).is_ok());
         // refresh token should be no change
@@ -402,7 +411,8 @@ mod tests {
 
         let auth_resp = conn
             .token_auth(ADMIN_EMAIL.to_owned(), ADMIN_PASSWORD.to_owned())
-            .await.unwrap();
+            .await
+            .unwrap();
         assert_eq!(resp.refresh_token, auth_resp.refresh_token);
     }
 }
diff --git a/ee/tabby-webserver/src/service/db.rs b/ee/tabby-webserver/src/service/db.rs
@@ -333,10 +333,8 @@ pub struct RefreshToken {
 }
 
 impl RefreshToken {
-
     fn select(clause: &str) -> String {
-        r#"SELECT id, user_id, token, expires_in, created_at FROM refresh_tokens WHERE "#
-            .to_owned()
+        r#"SELECT id, user_id, token, expires_in, created_at FROM refresh_tokens WHERE "#.to_owned()
             + clause
     }
 
@@ -390,22 +388,28 @@ impl DbConn {
                     RefreshToken::select("token = ?").as_str(),
                     params![token],
                     RefreshToken::from_row,
-                ).optional()
+                )
+                .optional()
             })
             .await?;
 
         Ok(token)
     }
 
-    pub async fn get_user_unexpired_token(&self, user_id: i32, utc_ts: i64) -> Result<Option<RefreshToken>> {
+    pub async fn get_user_unexpired_token(
+        &self,
+        user_id: i32,
+        utc_ts: i64,
+    ) -> Result<Option<RefreshToken>> {
         let token = self
             .conn
             .call(move |c| {
                 c.query_row(
                     RefreshToken::select("user_id = ? AND expires_in > ?").as_str(),
                     params![user_id, utc_ts],
                     RefreshToken::from_row,
-                ).optional()
+                )
+                .optional()
             })
             .await?;
 
@@ -502,4 +506,35 @@ mod tests {
         let invitations = conn.list_invitations().await.unwrap();
         assert!(invitations.is_empty());
     }
+
+    #[tokio::test]
+    async fn test_create_refresh_token() {
+        let conn = DbConn::new_in_memory().await.unwrap();
+
+        conn.create_refresh_token(1, "test", 100).await.unwrap();
+
+        let token = conn.get_refresh_token("test").await.unwrap().unwrap();
+
+        assert_eq!(token.user_id, 1);
+        assert_eq!(token.token, "test");
+        assert_eq!(token.expires_in, 100);
+    }
+
+    #[tokio::test]
+    async fn test_get_user_unexpired_token() {
+        let conn = DbConn::new_in_memory().await.unwrap();
+
+        conn.create_refresh_token(1, "test", 100).await.unwrap();
+
+        // utc_ts is 10, expires_in is 100, so the token is not expired
+        let token = conn.get_user_unexpired_token(1, 10).await.unwrap().unwrap();
+
+        assert_eq!(token.user_id, 1);
+        assert_eq!(token.token, "test");
+        assert_eq!(token.expires_in, 100);
+
+        // utc_ts is 1000, expires_in is 100, so the token is expired
+        let token = conn.get_user_unexpired_token(1, 1000).await.unwrap();
+        assert!(token.is_none());
+    }
 }