feat(tabby-webserver): when user is registered with invitation code, … (#999)

* 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

@@ -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]

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?;