From 1d3fe5cc2b8c6be6154040d5f5156c1f628455de Mon Sep 17 00:00:00 2001
From: Meng Zhang <meng@tabbyml.com>
Date: Sun, 10 Dec 2023 11:16:44 +0800
Subject: [PATCH] =?UTF-8?q?feat(tabby-webserver):=20when=20user=20is=20reg?=
 =?UTF-8?q?istered=20with=20invitation=20code,=20=E2=80=A6=20(#999)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* feat(tabby-webserver): when user is registered with invitation code, delete invitation when user is created

* update cmts

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
---
 ee/tabby-webserver/src/service/auth.rs     | 27 ++++++++++++---
 ee/tabby-webserver/src/service/db/users.rs | 40 +++++++++++++++++++---
 2 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/ee/tabby-webserver/src/service/auth.rs b/ee/tabby-webserver/src/service/auth.rs
index ee238ac..9e1b84b 100644
--- a/ee/tabby-webserver/src/service/auth.rs
+++ b/ee/tabby-webserver/src/service/auth.rs
@@ -150,7 +150,7 @@ impl AuthenticationService for DbConn {
         input.validate()?;
 
         let is_admin_initialized = self.is_admin_initialized().await?;
-        if is_admin_initialized {
+        let invitation = if is_admin_initialized {
             let err = Err(RegisterError::InvalidInvitationCode);
             let Some(invitation_code) = invitation_code else {
                 return err;
@@ -163,6 +163,10 @@ impl AuthenticationService for DbConn {
             if invitation.email != input.email {
                 return err;
             }
+
+            Some(invitation)
+        } else {
+            None
         };
 
         // check if email exists
@@ -174,9 +178,19 @@ impl AuthenticationService for DbConn {
             return Err(RegisterError::Unknown);
         };
 
-        let id = self
-            .create_user(input.email.clone(), pwd_hash, !is_admin_initialized)
-            .await?;
+        let id = if let Some(invitation) = invitation {
+            self.create_user_with_invitation(
+                input.email.clone(),
+                pwd_hash,
+                !is_admin_initialized,
+                invitation.id,
+            )
+            .await?
+        } else {
+            self.create_user(input.email.clone(), pwd_hash, !is_admin_initialized)
+                .await?
+        };
+
         let user = self.get_user(id).await?.unwrap();
 
         let refresh_token = generate_refresh_token();
@@ -421,8 +435,11 @@ mod tests {
                 Some(invitation.code.clone())
             )
             .await,
-            Err(RegisterError::DuplicateEmail)
+            Err(RegisterError::InvalidInvitationCode)
         );
+
+        // Used invitation should have been deleted,  following delete attempt should fail.
+        assert!(conn.delete_invitation(invitation.id).await.is_err());
     }
 
     #[tokio::test]
diff --git a/ee/tabby-webserver/src/service/db/users.rs b/ee/tabby-webserver/src/service/db/users.rs
index a4afa86..160395a 100644
--- a/ee/tabby-webserver/src/service/db/users.rs
+++ b/ee/tabby-webserver/src/service/db/users.rs
@@ -58,14 +58,46 @@ impl DbConn {
         email: String,
         password_encrypted: String,
         is_admin: bool,
+    ) -> Result<i32> {
+        self.create_user_impl(email, password_encrypted, is_admin, None)
+            .await
+    }
+
+    pub async fn create_user_with_invitation(
+        &self,
+        email: String,
+        password_encrypted: String,
+        is_admin: bool,
+        invitation_id: i32,
+    ) -> Result<i32> {
+        self.create_user_impl(email, password_encrypted, is_admin, Some(invitation_id))
+            .await
+    }
+
+    async fn create_user_impl(
+        &self,
+        email: String,
+        password_encrypted: String,
+        is_admin: bool,
+        invitation_id: Option<i32>,
     ) -> Result<i32> {
         let res = self
             .conn
             .call(move |c| {
-                let mut stmt = c.prepare(
-                    r#"INSERT INTO users (email, password_encrypted, is_admin, auth_token) VALUES (?, ?, ?, ?)"#,
-                )?;
-                let id = stmt.insert((email, password_encrypted, is_admin, generate_auth_token()))?;
+                let tx = c.transaction()?;
+
+                if let Some(invitation_id) = invitation_id {
+                    tx.execute("DELETE FROM invitations WHERE id = ?", params![invitation_id])?;
+                }
+
+                let id = {
+                    let mut stmt = tx.prepare(
+                        r#"INSERT INTO users (email, password_encrypted, is_admin, auth_token) VALUES (?, ?, ?, ?)"#,
+                    )?;
+                    stmt.insert((email, password_encrypted, is_admin, generate_auth_token()))?
+                };
+
+                tx.commit()?;
                 Ok(id)
             })
             .await?;