Don't return avatars from graphql

TabbyML · Mar 20, 2024 · 22d2db4 · 22d2db4
1 parent 87c59e6
commit 22d2db4
Show file tree

Hide file tree

Showing 13 changed files with 99 additions and 41 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/ee/tabby-db/migrations/0019_user-avatar.up.sql b/ee/tabby-db/migrations/0019_user-avatar.up.sql
@@ -1 +1 @@
-ALTER TABLE users ADD COLUMN avatar TEXT DEFAULT NULL;
+ALTER TABLE users ADD COLUMN avatar BLOB DEFAULT NULL;
diff --git a/ee/tabby-db/schema.sqlite b/ee/tabby-db/schema.sqlite
diff --git a/ee/tabby-db/src/users.rs b/ee/tabby-db/src/users.rs
@@ -223,17 +223,20 @@ impl DbConn {
         Ok(())
     }
 
-    pub async fn update_user_avatar(&self, id: i32, avatar_base64: Option<String>) -> Result<()> {
-        query!(
-            "UPDATE users SET avatar = ? WHERE id = ?;",
-            avatar_base64,
-            id
-        )
-        .execute(&self.pool)
-        .await?;
+    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

diff --git a/ee/tabby-webserver/Cargo.toml b/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

diff --git a/ee/tabby-webserver/src/avatar.rs b/ee/tabby-webserver/src/avatar.rs
@@ -0,0 +1,30 @@
+use std::sync::Arc;
+
+use axum::{
+    extract::{Path, State},
+    middleware::from_fn_with_state,
+    response::Response,
+    routing, Router,
+};
+use hyper::{Body, StatusCode};
+
+use crate::{hub::require_login_middleware, schema::auth::AuthenticationService, service::AsID};
+
+pub fn routes(auth: Arc<dyn AuthenticationService>) -> Router {
+    Router::new()
+        .route("/:id", routing::get(avatar))
+        .with_state(auth.clone())
+        .layer(from_fn_with_state(auth, require_login_middleware))
+}
+
+pub async fn avatar(
+    State(state): State<Arc<dyn AuthenticationService>>,
+    Path(id): Path<i64>,
+) -> Result<Response<Body>, StatusCode> {
+    let avatar = state
+        .get_user_avatar(&id.as_id())
+        .await
+        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
+        .ok_or(StatusCode::NOT_FOUND)?;
+    Ok(Response::new(Body::from(avatar.into_vec())))
+}
diff --git a/ee/tabby-webserver/src/handler.rs b/ee/tabby-webserver/src/handler.rs
@@ -13,7 +13,7 @@ use tabby_db::DbConn;
 use tracing::warn;
 
 use crate::{
-    cron, hub, oauth,
+    avatar, cron, hub, oauth,
     repositories::{self, RepositoryCache},
     schema::{create_schema, Schema, ServiceLocator},
     service::{create_service_locator, event_logger::new_event_logger},
@@ -77,6 +77,7 @@ impl WebserverHandle {
                 "/repositories",
                 repositories::routes(rs.clone(), ctx.auth()),
             )
+            .nest("/avatar", avatar::routes(ctx.auth()))
             .nest("/oauth", oauth::routes(ctx.auth()));
 
         let ui = ui.route("/graphiql", routing::get(graphiql("/graphql", None)));

diff --git a/ee/tabby-webserver/src/hub/mod.rs b/ee/tabby-webserver/src/hub/mod.rs
@@ -9,17 +9,18 @@ use std::{
 use api::{ConnectHubRequest, Hub};
 use axum::{
     extract::{ws::WebSocket, ConnectInfo, State, WebSocketUpgrade},
+    middleware::Next,
     response::IntoResponse,
     TypedHeader,
 };
-use hyper::{Body, StatusCode};
+use hyper::{Body, Request, StatusCode};
 use juniper_axum::extract::AuthBearer;
 use tabby_common::{api::code::SearchResponse, config::RepositoryConfig};
 use tarpc::server::{BaseChannel, Channel};
 use tracing::warn;
 use websocket::WebSocketTransport;
 
-use crate::schema::ServiceLocator;
+use crate::schema::{auth::AuthenticationService, ServiceLocator};
 
 pub(crate) async fn ws_handler(
     ws: WebSocketUpgrade,
@@ -50,6 +51,29 @@ pub(crate) async fn ws_handler(
         .into_response()
 }
 
+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 handle_socket(
     state: Arc<dyn ServiceLocator>,
     socket: WebSocket,

diff --git a/ee/tabby-webserver/src/lib.rs b/ee/tabby-webserver/src/lib.rs
@@ -1,5 +1,6 @@
 //! Defines behavior for the tabby webserver which allows users to interact with enterprise features.
 //! Using the web interface (e.g chat playground) requires using this module with the `--webserver` flag on the command line.
+mod avatar;
 mod cron;
 mod handler;
 mod hub;

diff --git a/ee/tabby-webserver/src/repositories/mod.rs b/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::{
+    hub::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
 }

diff --git a/ee/tabby-webserver/src/schema/auth.rs b/ee/tabby-webserver/src/schema/auth.rs
@@ -469,6 +469,7 @@ pub trait AuthenticationService: Send + Sync {
     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<String>) -> Result<()>;
+    async fn get_user_avatar(&self, id: &ID) -> Result<Option<Box<[u8]>>>;
 }
 
 fn validate_password(value: &str) -> Result<(), validator::ValidationError> {

diff --git a/ee/tabby-webserver/src/service/auth.rs b/ee/tabby-webserver/src/service/auth.rs
@@ -7,6 +7,7 @@ use argon2::{
     Argon2, PasswordHasher, PasswordVerifier,
 };
 use async_trait::async_trait;
+use base64::Engine;
 use chrono::{Duration, Utc};
 use juniper::ID;
 use tabby_db::{DbConn, InvitationDAO};
@@ -194,10 +195,23 @@ impl AuthenticationService for AuthenticationServiceImpl {
 
     async fn update_user_avatar(&self, id: &ID, avatar: Option<String>) -> Result<()> {
         let id = id.as_rowid()?;
+        let avatar = match avatar {
+            Some(avatar) => Some(
+                base64::prelude::BASE64_STANDARD
+                    .decode(avatar.as_bytes())
+                    .map_err(anyhow::Error::from)?
+                    .into_boxed_slice(),
+            ),
+            None => None,
+        };
         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());

diff --git a/ee/tabby-webserver/src/service/mod.rs b/ee/tabby-webserver/src/service/mod.rs
@@ -17,7 +17,7 @@ use axum::{
     middleware::Next,
     response::IntoResponse,
 };
-pub(in crate::service) use dao::{AsID, AsRowid};
+pub(crate) use dao::{AsID, AsRowid};
 use hyper::{client::HttpConnector, Body, Client, StatusCode};
 use juniper::ID;
 use tabby_common::{
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		ALTER TABLE users ADD COLUMN avatar TEXT DEFAULT NULL;
		ALTER TABLE users ADD COLUMN avatar BLOB DEFAULT NULL;