From 88e5187b57a083261ac70ae0da00c594ac2e0697 Mon Sep 17 00:00:00 2001 From: Meng Zhang Date: Fri, 1 Dec 2023 22:35:02 +0800 Subject: [PATCH] feat(webserver): implement is_admin_initialized graphql api (#929) * feat(webserver): implement is_admin_initialized graphql api * refactor * add unit test * [autofix.ci] apply automated fixes * renaming * refactor: server -> locator * fix unused --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- ee/tabby-webserver/src/lib.rs | 2 +- ee/tabby-webserver/src/schema/auth.rs | 1 + ee/tabby-webserver/src/schema/mod.rs | 22 ++++--- ee/tabby-webserver/src/service/auth.rs | 5 ++ ee/tabby-webserver/src/service/db.rs | 80 +++++++++++++++++++------- 5 files changed, 79 insertions(+), 31 deletions(-) diff --git a/ee/tabby-webserver/src/lib.rs b/ee/tabby-webserver/src/lib.rs index 2331bd5..ec2fedc 100644 --- a/ee/tabby-webserver/src/lib.rs +++ b/ee/tabby-webserver/src/lib.rs @@ -28,7 +28,7 @@ use axum::{ use hyper::Body; use juniper_axum::{graphiql, graphql, playground}; use schema::{ - worker::{RegisterWorkerError, Worker, WorkerKind, WorkerService}, + worker::{RegisterWorkerError, Worker, WorkerKind}, Schema, ServiceLocator, }; use service::create_service_locator; diff --git a/ee/tabby-webserver/src/schema/auth.rs b/ee/tabby-webserver/src/schema/auth.rs index 48e79ea..c23841b 100644 --- a/ee/tabby-webserver/src/schema/auth.rs +++ b/ee/tabby-webserver/src/schema/auth.rs @@ -136,6 +136,7 @@ pub trait AuthenticationService: Send + Sync { async fn token_auth(&self, email: String, password: String) -> FieldResult; async fn refresh_token(&self, refresh_token: String) -> FieldResult; async fn verify_token(&self, access_token: String) -> FieldResult; + async fn is_admin_initialized(&self) -> FieldResult; } #[cfg(test)] diff --git a/ee/tabby-webserver/src/schema/mod.rs b/ee/tabby-webserver/src/schema/mod.rs index 0df47fa..1d305de 100644 --- a/ee/tabby-webserver/src/schema/mod.rs +++ b/ee/tabby-webserver/src/schema/mod.rs @@ -27,13 +27,13 @@ pub trait ServiceLocator: Send + Sync { pub struct Context { claims: Option, - server: Arc, + locator: Arc, } impl FromAuth> for Context { - fn build(server: Arc, bearer: Option) -> Self { + fn build(locator: Arc, bearer: Option) -> Self { let claims = bearer.and_then(|token| validate_jwt(&token).ok()); - Self { claims, server } + Self { claims, locator } } } @@ -46,13 +46,17 @@ pub struct Query; #[graphql_object(context = Context)] impl Query { async fn workers(ctx: &Context) -> Vec { - ctx.server.worker().list_workers().await + ctx.locator.worker().list_workers().await } async fn registration_token(ctx: &Context) -> FieldResult { - let token = ctx.server.worker().read_registration_token().await?; + let token = ctx.locator.worker().read_registration_token().await?; Ok(token) } + + async fn is_admin_initialized(ctx: &Context) -> FieldResult { + ctx.locator.auth().is_admin_initialized().await + } } #[derive(Default)] @@ -63,7 +67,7 @@ impl Mutation { async fn reset_registration_token(ctx: &Context) -> FieldResult { if let Some(claims) = &ctx.claims { if claims.user_info().is_admin() { - let reg_token = ctx.server.worker().reset_registration_token().await?; + let reg_token = ctx.locator.worker().reset_registration_token().await?; return Ok(reg_token); } } @@ -79,7 +83,7 @@ impl Mutation { password1: String, password2: String, ) -> FieldResult { - ctx.server + ctx.locator .auth() .register(email, password1, password2) .await @@ -90,11 +94,11 @@ impl Mutation { email: String, password: String, ) -> FieldResult { - ctx.server.auth().token_auth(email, password).await + ctx.locator.auth().token_auth(email, password).await } async fn verify_token(ctx: &Context, token: String) -> FieldResult { - ctx.server.auth().verify_token(token).await + ctx.locator.auth().verify_token(token).await } } diff --git a/ee/tabby-webserver/src/service/auth.rs b/ee/tabby-webserver/src/service/auth.rs index 447d184..122e07d 100644 --- a/ee/tabby-webserver/src/service/auth.rs +++ b/ee/tabby-webserver/src/service/auth.rs @@ -188,6 +188,11 @@ impl AuthenticationService for DbConn { let resp = VerifyTokenResponse::new(claims); Ok(resp) } + + async fn is_admin_initialized(&self) -> FieldResult { + let admin = self.list_admin_users().await?; + Ok(!admin.is_empty()) + } } fn password_hash(raw: &str) -> password_hash::Result { diff --git a/ee/tabby-webserver/src/service/db.rs b/ee/tabby-webserver/src/service/db.rs index d3e1045..20aaff4 100644 --- a/ee/tabby-webserver/src/service/db.rs +++ b/ee/tabby-webserver/src/service/db.rs @@ -2,7 +2,7 @@ use std::{path::PathBuf, sync::Arc}; use anyhow::Result; use lazy_static::lazy_static; -use rusqlite::{params, OptionalExtension}; +use rusqlite::{params, OptionalExtension, Row}; use rusqlite_migration::{AsyncMigrations, M}; use tabby_common::path::tabby_root; use tokio_rusqlite::Connection; @@ -47,6 +47,25 @@ pub struct User { pub is_admin: bool, } +impl User { + fn select(clause: &str) -> String { + r#"SELECT id, email, password_encrypted, is_admin, created_at, updated_at FROM users WHERE "# + .to_owned() + + clause + } + + fn from_row(row: &Row<'_>) -> std::result::Result { + Ok(User { + id: row.get(0)?, + email: row.get(1)?, + password_encrypted: row.get(2)?, + is_admin: row.get(3)?, + created_at: row.get(4)?, + updated_at: row.get(5)?, + }) + } +} + async fn db_path() -> Result { let db_dir = tabby_root().join("ee"); tokio::fs::create_dir_all(db_dir.clone()).await?; @@ -156,35 +175,51 @@ impl DbConn { .conn .call(move |c| { c.query_row( - r#"SELECT id, email, password_encrypted, is_admin, created_at, updated_at FROM users WHERE email = ?"#, + User::select("email = ?").as_str(), params![email], - |row| { - Ok(User { - id: row.get(0)?, - email: row.get(1)?, - password_encrypted: row.get(2)?, - is_admin: row.get(3)?, - created_at: row.get(4)?, - updated_at: row.get(5)?, - }) - }, - ).optional() + User::from_row, + ) + .optional() }) .await?; Ok(user) } + + pub async fn list_admin_users(&self) -> Result> { + let users = self + .conn + .call(move |c| { + let mut stmt = c.prepare(&User::select("is_admin"))?; + let user_iter = stmt.query_map([], User::from_row)?; + Ok(user_iter.filter_map(|x| x.ok()).collect::>()) + }) + .await?; + + Ok(users) + } } #[cfg(test)] mod tests { use super::*; + use crate::schema::auth::AuthenticationService; async fn new_in_memory() -> Result { let conn = Connection::open_in_memory().await?; DbConn::init_db(conn).await } + async fn create_admin_user(conn: &DbConn) -> String { + let email = "test@example.com"; + let passwd = "123456"; + let is_admin = true; + conn.create_user(email.to_string(), passwd.to_string(), is_admin) + .await + .unwrap(); + email.to_owned() + } + #[tokio::test] async fn migrations_test() { assert!(MIGRATIONS.validate().await.is_ok()); @@ -212,14 +247,8 @@ mod tests { async fn test_create_user() { let conn = new_in_memory().await.unwrap(); - let email = "test@example.com"; - let passwd = "123456"; - let is_admin = true; - conn.create_user(email.to_string(), passwd.to_string(), is_admin) - .await - .unwrap(); - - let user = conn.get_user_by_email(email).await.unwrap().unwrap(); + let email = create_admin_user(&conn).await; + let user = conn.get_user_by_email(&email).await.unwrap().unwrap(); assert_eq!(user.id, 1); } @@ -232,4 +261,13 @@ mod tests { assert!(user.is_none()); } + + #[tokio::test] + async fn test_is_admin_initialized() { + let conn = new_in_memory().await.unwrap(); + + assert!(!conn.is_admin_initialized().await.unwrap()); + create_admin_user(&conn).await; + assert!(conn.is_admin_initialized().await.unwrap()); + } }