refactor(webserver): remove IDWrapper (#1957)

TabbyML · Apr 25, 2024 · ce657ef · ce657ef
1 parent 2292147
commit ce657ef
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 20 deletions.
diff --git a/ee/tabby-webserver/src/schema/auth.rs b/ee/tabby-webserver/src/schema/auth.rs
@@ -197,11 +197,6 @@ impl RefreshTokenResponse {
     }
 }
 
-// IDWrapper to used as a type guard for refactoring, can be removed in a follow up PR.
-// FIXME(meng): refactor out IDWrapper.
-#[derive(Serialize, Deserialize, Debug)]
-pub struct IDWrapper(pub ID);
-
 #[derive(Debug, Serialize, Deserialize)]
 pub struct JWTPayload {
     /// Expiration time (as UTC timestamp)
@@ -211,7 +206,7 @@ pub struct JWTPayload {
     iat: i64,
 
     /// User id string
-    pub sub: IDWrapper,
+    pub sub: ID,
 }
 
 impl JWTPayload {
@@ -220,7 +215,7 @@ impl JWTPayload {
         Self {
             iat: now as i64,
             exp: (now + *JWT_DEFAULT_EXP) as i64,
-            sub: IDWrapper(id),
+            sub: id,
         }
     }
 }
@@ -525,6 +520,6 @@ mod tests {
         let claims = JWTPayload::new(ID::from("test".to_owned()));
         let token = generate_jwt(claims).unwrap();
         let claims = validate_jwt(&token).unwrap();
-        assert_eq!(claims.sub.0.to_string(), "test");
+        assert_eq!(claims.sub.to_string(), "test");
     }
 }
diff --git a/ee/tabby-webserver/src/schema/mod.rs b/ee/tabby-webserver/src/schema/mod.rs
@@ -139,7 +139,7 @@ async fn check_admin(ctx: &Context) -> Result<(), CoreError> {
 
 async fn check_user(ctx: &Context) -> Result<User, CoreError> {
     let claims = check_claims(ctx)?;
-    let user = ctx.locator.auth().get_user(&claims.sub.0).await?;
+    let user = ctx.locator.auth().get_user(&claims.sub).await?;
     Ok(user)
 }
 
@@ -173,7 +173,7 @@ impl Query {
 
     async fn me(ctx: &Context) -> Result<User> {
         let claims = check_claims(ctx)?;
-        ctx.locator.auth().get_user(&claims.sub.0).await
+        ctx.locator.auth().get_user(&claims.sub).await
     }
 
     async fn users(
@@ -512,7 +512,7 @@ impl Mutation {
         ctx.locator
             .auth()
             .update_user_password(
-                &claims.sub.0,
+                &claims.sub,
                 input.old_password.as_deref(),
                 &input.new_password1,
             )
@@ -524,23 +524,20 @@ impl Mutation {
         let claims = check_claims(ctx)?;
         ctx.locator
             .auth()
-            .reset_user_auth_token(&claims.sub.0)
+            .reset_user_auth_token(&claims.sub)
             .await?;
         Ok(true)
     }
 
     async fn logout_all_sessions(ctx: &Context) -> Result<bool> {
         let claims = check_claims(ctx)?;
-        ctx.locator
-            .auth()
-            .logout_all_sessions(&claims.sub.0)
-            .await?;
+        ctx.locator.auth().logout_all_sessions(&claims.sub).await?;
         Ok(true)
     }
 
     async fn update_user_active(ctx: &Context, id: ID, active: bool) -> Result<bool> {
         check_admin(ctx).await?;
-        if ctx.claims.as_ref().is_some_and(|c| c.sub.0 == id) {
+        if ctx.claims.as_ref().is_some_and(|c| c.sub == id) {
             return Err(CoreError::Forbidden(
                 "You cannot change your own active status",
             ));
@@ -551,7 +548,7 @@ impl Mutation {
 
     async fn update_user_role(ctx: &Context, id: ID, is_admin: bool) -> Result<bool> {
         check_admin(ctx).await?;
-        if ctx.claims.as_ref().is_some_and(|c| c.sub.0 == id) {
+        if ctx.claims.as_ref().is_some_and(|c| c.sub == id) {
             return Err(CoreError::Forbidden("You cannot update your own role"));
         }
         ctx.locator.auth().update_user_role(&id, is_admin).await?;
@@ -564,7 +561,7 @@ impl Mutation {
         avatar_base64: Option<String>,
     ) -> Result<bool> {
         let claims = check_claims(ctx)?;
-        if claims.sub.0 != id {
+        if claims.sub != id {
             return Err(CoreError::Unauthorized(
                 "You cannot change another user's avatar",
             ));

diff --git a/ee/tabby-webserver/src/service/mod.rs b/ee/tabby-webserver/src/service/mod.rs
@@ -137,7 +137,7 @@ impl ServerContext {
 
         // Allow JWT based access (from web browser), regardless of the license status.
         if let Ok(jwt) = self.auth.verify_access_token(token).await {
-            return (true, Some(jwt.sub.0));
+            return (true, Some(jwt.sub));
         }
 
         let is_license_valid = self.license.read().await.ensure_valid_license().is_ok();