From 19d773e7ad185540a5cb8bef45ffcb3c1deb28cc Mon Sep 17 00:00:00 2001 From: Meng Zhang Date: Sat, 2 Dec 2023 00:14:29 +0800 Subject: [PATCH] feat: implement register api check with invitation code (#934) * feat(webserver): implement is_admin_initialized graphql api * refactor * add unit test * [autofix.ci] apply automated fixes * renaming * temp invitations * update * update * implement register check * test * update invitations --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- ee/tabby-webserver/graphql/schema.graphql | 13 ++- ee/tabby-webserver/src/schema/auth.rs | 15 +++ ee/tabby-webserver/src/schema/mod.rs | 42 +++++++- ee/tabby-webserver/src/service/auth.rs | 33 +++++- ee/tabby-webserver/src/service/db.rs | 121 +++++++++++++++++++++- 5 files changed, 214 insertions(+), 10 deletions(-) diff --git a/ee/tabby-webserver/graphql/schema.graphql b/ee/tabby-webserver/graphql/schema.graphql index 92d6e5625fd1..b30a8f327b58 100644 --- a/ee/tabby-webserver/graphql/schema.graphql +++ b/ee/tabby-webserver/graphql/schema.graphql @@ -10,9 +10,11 @@ enum WorkerKind { type Mutation { resetRegistrationToken: String! - register(email: String!, password1: String!, password2: String!): RegisterResponse! + register(email: String!, password1: String!, password2: String!, invitationCode: String): RegisterResponse! tokenAuth(email: String!, password: String!): TokenAuthResponse! verifyToken(token: String!): VerifyTokenResponse! + createInvitation(email: String!): Int! + deleteInvitation(id: Int!): Int! } type UserInfo { @@ -33,6 +35,15 @@ type Claims { type Query { workers: [Worker!]! registrationToken: String! + isAdminInitialized: Boolean! + invitations: [Invitation!]! +} + +type Invitation { + id: Int! + email: String! + code: String! + createdAt: String! } type Worker { diff --git a/ee/tabby-webserver/src/schema/auth.rs b/ee/tabby-webserver/src/schema/auth.rs index c23841b51f3a..f3ea36c4f815 100644 --- a/ee/tabby-webserver/src/schema/auth.rs +++ b/ee/tabby-webserver/src/schema/auth.rs @@ -1,5 +1,6 @@ use std::fmt::Debug; +use anyhow::Result; use async_trait::async_trait; use jsonwebtoken as jwt; use juniper::{FieldResult, GraphQLObject}; @@ -125,6 +126,15 @@ impl Claims { } } +#[derive(Debug, Default, Serialize, Deserialize, GraphQLObject)] +pub struct Invitation { + pub id: i32, + pub email: String, + pub code: String, + + pub created_at: String, +} + #[async_trait] pub trait AuthenticationService: Send + Sync { async fn register( @@ -132,11 +142,16 @@ pub trait AuthenticationService: Send + Sync { email: String, password1: String, password2: String, + invitation_code: Option, ) -> FieldResult; 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; + + async fn create_invitation(&self, email: String) -> Result; + async fn list_invitations(&self) -> Result>; + async fn delete_invitation(&self, id: i32) -> Result; } #[cfg(test)] diff --git a/ee/tabby-webserver/src/schema/mod.rs b/ee/tabby-webserver/src/schema/mod.rs index 1d305de5a544..aed2b0cbfc04 100644 --- a/ee/tabby-webserver/src/schema/mod.rs +++ b/ee/tabby-webserver/src/schema/mod.rs @@ -12,7 +12,10 @@ use juniper_axum::FromAuth; use tabby_common::api::{code::CodeSearch, event::RawEventLogger}; use validator::ValidationError; -use self::{auth::validate_jwt, worker::WorkerService}; +use self::{ + auth::{validate_jwt, Invitation}, + worker::WorkerService, +}; use crate::schema::{ auth::{RegisterResponse, TokenAuthResponse, VerifyTokenResponse}, worker::Worker, @@ -57,6 +60,15 @@ impl Query { async fn is_admin_initialized(ctx: &Context) -> FieldResult { ctx.locator.auth().is_admin_initialized().await } + + async fn invitations(ctx: &Context) -> FieldResult> { + if let Some(claims) = &ctx.claims { + if claims.user_info().is_admin() { + return Ok(ctx.locator.auth().list_invitations().await?); + } + } + Err(unauthorized("Only admin is able to query invitations")) + } } #[derive(Default)] @@ -71,9 +83,8 @@ impl Mutation { return Ok(reg_token); } } - Err(FieldError::new( + Err(unauthorized( "Only admin is able to reset registration token", - graphql_value!("Unauthorized"), )) } @@ -82,10 +93,11 @@ impl Mutation { email: String, password1: String, password2: String, + invitation_code: Option, ) -> FieldResult { ctx.locator .auth() - .register(email, password1, password2) + .register(email, password1, password2, invitation_code) .await } @@ -100,6 +112,24 @@ impl Mutation { async fn verify_token(ctx: &Context, token: String) -> FieldResult { ctx.locator.auth().verify_token(token).await } + + async fn create_invitation(ctx: &Context, email: String) -> FieldResult { + if let Some(claims) = &ctx.claims { + if claims.user_info().is_admin() { + return Ok(ctx.locator.auth().create_invitation(email).await?); + } + } + Err(unauthorized("Only admin is able to create invitation")) + } + + async fn delete_invitation(ctx: &Context, id: i32) -> FieldResult { + if let Some(claims) = &ctx.claims { + if claims.user_info().is_admin() { + return Ok(ctx.locator.auth().delete_invitation(id).await?); + } + } + Err(unauthorized("Only admin is able to delete invitation")) + } } #[derive(Debug)] @@ -135,3 +165,7 @@ pub type Schema = RootNode<'static, Query, Mutation, EmptySubscription> pub fn create_schema() -> Schema { Schema::new(Query, Mutation, EmptySubscription::new()) } + +fn unauthorized(msg: &str) -> FieldError { + FieldError::new(msg, graphql_value!("Unauthorized")) +} diff --git a/ee/tabby-webserver/src/service/auth.rs b/ee/tabby-webserver/src/service/auth.rs index 122e07d74443..3d0281d41788 100644 --- a/ee/tabby-webserver/src/service/auth.rs +++ b/ee/tabby-webserver/src/service/auth.rs @@ -1,3 +1,4 @@ +use anyhow::Result; use argon2::{ password_hash, password_hash::{rand_core::OsRng, SaltString}, @@ -10,8 +11,8 @@ use validator::Validate; use super::db::DbConn; use crate::schema::{ auth::{ - generate_jwt, validate_jwt, AuthenticationService, Claims, RefreshTokenResponse, - RegisterResponse, TokenAuthResponse, UserInfo, VerifyTokenResponse, + generate_jwt, validate_jwt, AuthenticationService, Claims, Invitation, + RefreshTokenResponse, RegisterResponse, TokenAuthResponse, UserInfo, VerifyTokenResponse, }, ValidationErrors, }; @@ -109,6 +110,7 @@ impl AuthenticationService for DbConn { email: String, password1: String, password2: String, + invitation_code: Option, ) -> FieldResult { let input = RegisterInput { email, @@ -126,6 +128,21 @@ impl AuthenticationService for DbConn { ValidationErrors { errors }.into_field_error() })?; + if self.is_admin_initialized().await? { + let err = Err("Invitation code is not valid".into()); + let Some(invitation_code) = invitation_code else { + return err; + }; + + let Some(invitation) = self.get_invitation_by_code(&invitation_code).await? else { + return err; + }; + + if invitation.email != input.email { + return err; + } + }; + // check if email exists if self.get_user_by_email(&input.email).await?.is_some() { return Err("Email already exists".into()); @@ -193,6 +210,18 @@ impl AuthenticationService for DbConn { let admin = self.list_admin_users().await?; Ok(!admin.is_empty()) } + + async fn create_invitation(&self, email: String) -> Result { + self.create_invitation(email).await + } + + async fn list_invitations(&self) -> Result> { + self.list_invitations().await + } + + async fn delete_invitation(&self, id: i32) -> Result { + self.delete_invitation(id).await + } } 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 20aaff44b682..a34cf41fc6f2 100644 --- a/ee/tabby-webserver/src/service/db.rs +++ b/ee/tabby-webserver/src/service/db.rs @@ -1,17 +1,20 @@ use std::{path::PathBuf, sync::Arc}; -use anyhow::Result; +use anyhow::{anyhow, Result}; use lazy_static::lazy_static; use rusqlite::{params, OptionalExtension, Row}; use rusqlite_migration::{AsyncMigrations, M}; use tabby_common::path::tabby_root; use tokio_rusqlite::Connection; +use uuid::Uuid; + +use crate::schema::auth::Invitation; lazy_static! { static ref MIGRATIONS: AsyncMigrations = AsyncMigrations::new(vec![ M::up( r#" - CREATE TABLE IF NOT EXISTS registration_token ( + CREATE TABLE registration_token ( id INTEGER PRIMARY KEY AUTOINCREMENT, token VARCHAR(255) NOT NULL, created_at TIMESTAMP DEFAULT (DATETIME('now')), @@ -22,7 +25,7 @@ lazy_static! { ), M::up( r#" - CREATE TABLE IF NOT EXISTS users ( + CREATE TABLE users ( id INTEGER PRIMARY KEY AUTOINCREMENT, email VARCHAR(150) NOT NULL COLLATE NOCASE, password_encrypted VARCHAR(128) NOT NULL, @@ -33,6 +36,18 @@ lazy_static! { ); "# ), + M::up( + r#" + CREATE TABLE invitations ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + email VARCHAR(150) NOT NULL COLLATE NOCASE, + code VARCHAR(36) NOT NULL, + created_at TIMESTAMP DEFAULT (DATETIME('now')), + CONSTRAINT `idx_email` UNIQUE (`email`) + CONSTRAINT `idx_code` UNIQUE (`code`) + ); + "# + ), ]); } @@ -200,6 +215,81 @@ impl DbConn { } } +impl Invitation { + fn from_row(row: &Row<'_>) -> std::result::Result { + Ok(Self { + id: row.get(0)?, + email: row.get(1)?, + code: row.get(2)?, + created_at: row.get(3)?, + }) + } +} + +/// db read/write operations for `invitations` table +impl DbConn { + pub async fn list_invitations(&self) -> Result> { + let invitations = self + .conn + .call(move |c| { + let mut stmt = + c.prepare(r#"SELECT id, email, code, created_at FROM invitations"#)?; + let iter = stmt.query_map([], Invitation::from_row)?; + Ok(iter.filter_map(|x| x.ok()).collect::>()) + }) + .await?; + + Ok(invitations) + } + + pub async fn get_invitation_by_code(&self, code: &str) -> Result> { + let code = code.to_owned(); + let token = self + .conn + .call(|conn| { + conn.query_row( + r#"SELECT id, email, code, created_at FROM invitations WHERE code = ?"#, + [code], + Invitation::from_row, + ) + .optional() + }) + .await?; + + Ok(token) + } + + pub async fn create_invitation(&self, email: String) -> Result { + let code = Uuid::new_v4().to_string(); + let res = self + .conn + .call(move |c| { + let mut stmt = + c.prepare(r#"INSERT INTO invitations (email, code) VALUES (?, ?)"#)?; + let rowid = stmt.insert((email, code))?; + Ok(rowid) + }) + .await?; + if res != 1 { + return Err(anyhow!("failed to create invitation")); + } + + Ok(res as i32) + } + + pub async fn delete_invitation(&self, id: i32) -> Result { + let res = self + .conn + .call(move |c| c.execute(r#"DELETE FROM invitations WHERE id = ?"#, params![id])) + .await?; + if res != 1 { + return Err(anyhow!("failed to delete invitation")); + } + + Ok(id) + } +} + #[cfg(test)] mod tests { use super::*; @@ -270,4 +360,29 @@ mod tests { create_admin_user(&conn).await; assert!(conn.is_admin_initialized().await.unwrap()); } + + #[tokio::test] + async fn test_invitations() { + let conn = new_in_memory().await.unwrap(); + + let email = "hello@example.com".to_owned(); + conn.create_invitation(email).await.unwrap(); + + let invitations = conn.list_invitations().await.unwrap(); + assert_eq!(1, invitations.len()); + + assert!(Uuid::parse_str(&invitations[0].code).is_ok()); + let invitation = conn + .get_invitation_by_code(&invitations[0].code) + .await + .ok() + .flatten() + .unwrap(); + assert_eq!(invitation.id, invitations[0].id); + + conn.delete_invitation(invitations[0].id).await.unwrap(); + + let invitations = conn.list_invitations().await.unwrap(); + assert!(invitations.is_empty()); + } }