feat(db, webserver): Support user avatars (TabbyML#1691)

* feat(db, webserver): Support user avatars

* [autofix.ci] apply automated fixes

* Fix tests

* Don't return avatars from graphql

* Rename method

* Remove avatar field from UserDAO and User

* Apply suggestions

* Fix errors

* Fix avatar endpoint

* Add unit test

* Don't allow admins to change other users' avatars

* Apply suggestions

* [autofix.ci] apply automated fixes

* update

* update

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Meng Zhang <[email protected]>

.github/workflows/ast-grep-lint.yml

@@ -10,4 +10,6 @@ jobs:
         uses: actions/checkout@v4
 
       - name: ast-grep lint step
-        uses: ast-grep/[email protected]
+        uses: ast-grep/[email protected]
+        with:
+          version: 0.20.1

ee/tabby-db/migrations/0019_user-avatar.down.sql

@@ -0,0 +1 @@
+ALTER TABLE users DROP COLUMN avatar;

ee/tabby-db/migrations/0019_user-avatar.up.sql

@@ -0,0 +1 @@
+ALTER TABLE users ADD COLUMN avatar BLOB DEFAULT NULL;

ee/tabby-db/src/users.rs

@@ -221,6 +221,20 @@ impl DbConn {
         Ok(())
     }
 
+    pub async fn update_user_avatar(&self, id: i32, avatar: Option<Box<[u8]>>) -> Result<()> {
+        query!("UPDATE users SET avatar = ? WHERE id = ?;", avatar, id)
+            .execute(&self.pool)
+            .await?;
+        Ok(())
+    }
+
+    pub async fn get_user_avatar(&self, id: i32) -> Result<Option<Box<[u8]>>> {
+        let avatar = query_scalar!("SELECT avatar FROM users WHERE id = ?", id)
+            .fetch_one(&self.pool)
+            .await?;
+        Ok(avatar.map(Vec::into_boxed_slice))
+    }
+
     pub async fn count_active_users(&self) -> Result<usize> {
         self.cache
             .active_user_count

ee/tabby-webserver/Cargo.toml

@@ -13,6 +13,7 @@ anyhow.workspace = true
 argon2 = "0.5.1"
 async-trait.workspace = true
 axum = { workspace = true, features = ["ws", "headers"] }
+base64 = "0.22.0"
 bincode = "1.3.3"
 chrono = { workspace = true, features = ["serde"] }
 futures.workspace = true

ee/tabby-webserver/graphql/schema.graphql

@@ -22,6 +22,7 @@ type Mutation {
   logoutAllSessions: Boolean!
   updateUserActive(id: ID!, active: Boolean!): Boolean!
   updateUserRole(id: ID!, isAdmin: Boolean!): Boolean!
+  uploadUserAvatarBase64(id: ID!, avatarBase64: String): Boolean!
   register(email: String!, password1: String!, password2: String!, invitationCode: String): RegisterResponse!
   tokenAuth(email: String!, password: String!): TokenAuthResponse!
   verifyToken(token: String!): Boolean!

ee/tabby-webserver/src/handler.rs

@@ -1,21 +1,23 @@
 use std::sync::Arc;
 
 use axum::{
-    extract::State,
+    extract::{Path, State},
     http::Request,
     middleware::{from_fn_with_state, Next},
+    response::{IntoResponse, Response},
     routing, Extension, Json, Router,
 };
-use hyper::{Body, StatusCode};
-use juniper_axum::{graphiql, graphql, playground};
+use hyper::{header::CONTENT_TYPE, Body, StatusCode};
+use juniper::ID;
+use juniper_axum::{extract::AuthBearer, graphiql, graphql, playground};
 use tabby_common::api::{code::CodeSearch, event::EventLogger, server_setting::ServerSetting};
 use tabby_db::DbConn;
-use tracing::warn;
+use tracing::{error, warn};
 
 use crate::{
     cron, hub, oauth,
     repositories::{self, RepositoryCache},
-    schema::{create_schema, Schema, ServiceLocator},
+    schema::{auth::AuthenticationService, create_schema, Schema, ServiceLocator},
     service::{create_service_locator, event_logger::new_event_logger},
     ui,
 };
@@ -77,6 +79,12 @@ impl WebserverHandle {
                 "/repositories",
                 repositories::routes(rs.clone(), ctx.auth()),
             )
+            .route(
+                "/avatar/:id",
+                routing::get(avatar)
+                    .with_state(ctx.auth())
+                    .layer(from_fn_with_state(ctx.auth(), require_login_middleware)),
+            )
             .nest("/oauth", oauth::routes(ctx.auth()));
 
         let ui = ui.route("/graphiql", routing::get(graphiql("/graphql", None)));
@@ -87,6 +95,29 @@ impl WebserverHandle {
     }
 }
 
+pub(crate) async fn require_login_middleware(
+    State(auth): State<Arc<dyn AuthenticationService>>,
+    AuthBearer(token): AuthBearer,
+    request: Request<Body>,
+    next: Next<Body>,
+) -> axum::response::Response {
+    let unauthorized = axum::response::Response::builder()
+        .status(StatusCode::UNAUTHORIZED)
+        .body(Body::empty())
+        .unwrap()
+        .into_response();
+
+    let Some(token) = token else {
+        return unauthorized;
+    };
+
+    let Ok(_) = auth.verify_access_token(&token).await else {
+        return unauthorized;
+    };
+
+    next.run(request).await
+}
+
 async fn distributed_tabby_layer(
     State(ws): State<Arc<dyn ServiceLocator>>,
     request: Request<Body>,
@@ -110,3 +141,22 @@ async fn server_setting(
         disable_client_side_telemetry: security_setting.disable_client_side_telemetry,
     }))
 }
+
+async fn avatar(
+    State(state): State<Arc<dyn AuthenticationService>>,
+    Path(id): Path<ID>,
+) -> Result<Response<Body>, StatusCode> {
+    let avatar = state
+        .get_user_avatar(&id)
+        .await
+        .map_err(|e| {
+            error!("Failed to retrieve avatar: {e}");
+            StatusCode::INTERNAL_SERVER_ERROR
+        })?
+        .ok_or(StatusCode::NOT_FOUND)?;
+    let mut response = Response::new(Body::from(avatar.into_vec()));
+    response
+        .headers_mut()
+        .insert(CONTENT_TYPE, "image/*".parse().unwrap());
+    Ok(response)
+}

ee/tabby-webserver/src/lib.rs

@@ -9,6 +9,9 @@ mod schema;
 mod service;
 mod ui;
 
+#[cfg(test)]
+pub use service::*;
+
 pub mod public {
 
     pub use super::{

ee/tabby-webserver/src/repositories/mod.rs

@@ -5,17 +5,16 @@ use std::sync::Arc;
 use anyhow::Result;
 use axum::{
     extract::{Path, State},
-    http::{Request, StatusCode},
-    middleware::{from_fn_with_state, Next},
-    response::{IntoResponse, Response},
+    http::StatusCode,
+    middleware::from_fn_with_state,
+    response::Response,
     routing, Json, Router,
 };
-use hyper::Body;
-use juniper_axum::extract::AuthBearer;
 pub use resolve::RepositoryCache;
 use tracing::{instrument, warn};
 
 use crate::{
+    handler::require_login_middleware,
     repositories::resolve::{RepositoryMeta, ResolveParams},
     schema::auth::AuthenticationService,
 };
@@ -37,29 +36,6 @@ pub fn routes(rs: Arc<ResolveState>, auth: Arc<dyn AuthenticationService>) -> Ro
         .layer(from_fn_with_state(auth, require_login_middleware))
 }
 
-async fn require_login_middleware(
-    State(auth): State<Arc<dyn AuthenticationService>>,
-    AuthBearer(token): AuthBearer,
-    request: Request<Body>,
-    next: Next<Body>,
-) -> axum::response::Response {
-    let unauthorized = axum::response::Response::builder()
-        .status(StatusCode::UNAUTHORIZED)
-        .body(Body::empty())
-        .unwrap()
-        .into_response();
-
-    let Some(token) = token else {
-        return unauthorized;
-    };
-
-    let Ok(_) = auth.verify_access_token(&token).await else {
-        return unauthorized;
-    };
-
-    next.run(request).await
-}
-
 async fn not_found() -> StatusCode {
     StatusCode::NOT_FOUND
 }

ee/tabby-webserver/src/schema/auth.rs

@@ -463,6 +463,8 @@ pub trait AuthenticationService: Send + Sync {
     async fn delete_oauth_credential(&self, provider: OAuthProvider) -> Result<()>;
     async fn update_user_active(&self, id: &ID, active: bool) -> Result<()>;
     async fn update_user_role(&self, id: &ID, is_admin: bool) -> Result<()>;
+    async fn update_user_avatar(&self, id: &ID, avatar: Option<Box<[u8]>>) -> Result<()>;
+    async fn get_user_avatar(&self, id: &ID) -> Result<Option<Box<[u8]>>>;
 }
 
 fn validate_password(value: &str) -> Result<(), validator::ValidationError> {

ee/tabby-webserver/src/schema/mod.rs

@@ -12,6 +12,7 @@ use auth::{
     validate_jwt, AuthenticationService, Invitation, RefreshTokenResponse, RegisterResponse,
     TokenAuthResponse, User,
 };
+use base64::Engine;
 use job::{JobRun, JobService};
 use juniper::{
     graphql_object, graphql_value, EmptySubscription, FieldError, FieldResult, GraphQLObject,
@@ -420,6 +421,28 @@ impl Mutation {
         Ok(true)
     }
 
+    async fn upload_user_avatar_base64(
+        ctx: &Context,
+        id: ID,
+        avatar_base64: Option<String>,
+    ) -> Result<bool> {
+        let claims = check_claims(ctx)?;
+        if claims.sub.0 != id {
+            return Err(CoreError::Unauthorized(
+                "You cannot change another user's avatar",
+            ));
+        }
+        // ast-grep-ignore: use-schema-result
+        use anyhow::Context;
+        let avatar = avatar_base64
+            .map(|avatar| base64::prelude::BASE64_STANDARD.decode(avatar.as_bytes()))
+            .transpose()
+            .context("avatar is not valid base64 string")?
+            .map(Vec::into_boxed_slice);
+        ctx.locator.auth().update_user_avatar(&id, avatar).await?;
+        Ok(true)
+    }
+
     async fn register(
         ctx: &Context,
         email: String,

ee/tabby-webserver/src/service/auth.rs

@@ -190,6 +190,16 @@ impl AuthenticationService for AuthenticationServiceImpl {
         Ok(())
     }
 
+    async fn update_user_avatar(&self, id: &ID, avatar: Option<Box<[u8]>>) -> Result<()> {
+        let id = id.as_rowid()?;
+        self.db.update_user_avatar(id, avatar).await?;
+        Ok(())
+    }
+
+    async fn get_user_avatar(&self, id: &ID) -> Result<Option<Box<[u8]>>> {
+        Ok(self.db.get_user_avatar(id.as_rowid()?).await?)
+    }
+
     async fn token_auth(&self, email: String, password: String) -> Result<TokenAuthResponse> {
         let Some(user) = self.db.get_user_by_email(&email).await? else {
             return Err(anyhow!("Invalid email address or password").into());