feat(webserver): add admin user count limit (TabbyML#1548)

* feat(webserver): add admin user count limit

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes (attempt 2/3)

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

ee/tabby-db/src/users.rs

@@ -218,6 +218,13 @@ impl DbConn {
             .await?;
         Ok(users as usize)
     }
+
+    pub async fn count_active_admin_users(&self) -> Result<usize> {
+        let users = query_scalar!("SELECT COUNT(1) FROM users WHERE active and is_admin;")
+            .fetch_one(&self.pool)
+            .await?;
+        Ok(users as usize)
+    }
 }
 
 fn generate_auth_token() -> String {

ee/tabby-webserver/src/schema/license.rs

@@ -45,7 +45,7 @@ impl LicenseInfo {
     pub fn check_node_limit(&self, num_nodes: usize) -> bool {
         match self.r#type {
             LicenseType::Community => false,
-            LicenseType::Team => num_nodes < 2,
+            LicenseType::Team => num_nodes <= 2,
             LicenseType::Enterprise => true,
         }
     }
@@ -62,15 +62,33 @@ impl LicenseInfo {
 
         self
     }
-    pub fn ensure_available_seats(&self, num_seats: i32) -> Result<()> {
+
+    pub fn ensure_available_seats(&self, num_new_seats: usize) -> Result<()> {
         self.ensure_valid_license()?;
-        if (self.seats_used + num_seats) > self.seats {
+        if (self.seats_used as usize + num_new_seats) > self.seats as usize {
             return Err(CoreError::InvalidLicense(
                 "No sufficient seats under current license",
             ));
         }
         Ok(())
     }
+
+    pub fn ensure_admin_seats(&self, num_admins: usize) -> Result<()> {
+        self.ensure_valid_license()?;
+        let num_admin_seats = match self.r#type {
+            LicenseType::Community => 1,
+            LicenseType::Team => 3,
+            LicenseType::Enterprise => usize::MAX,
+        };
+
+        if num_admins > num_admin_seats {
+            return Err(CoreError::InvalidLicense(
+                "No sufficient admin seats under the license",
+            ));
+        }
+
+        Ok(())
+    }
 }
 
 #[async_trait]

ee/tabby-webserver/src/service/auth.rs

@@ -218,6 +218,12 @@ impl AuthenticationService for AuthenticationServiceImpl {
     }
 
     async fn update_user_role(&self, id: &ID, is_admin: bool) -> Result<()> {
+        if is_admin {
+            let license = self.license.read_license().await?;
+            let num_admins = self.db.count_active_admin_users().await?;
+            license.ensure_admin_seats(num_admins + 1)?;
+        }
+
         let id = id.as_rowid()?;
         let user = self.db.get_user(id).await?.context("User doesn't exits")?;
         if user.is_owner() {
@@ -382,9 +388,10 @@ impl AuthenticationService for AuthenticationServiceImpl {
     }
 
     async fn update_user_active(&self, id: &ID, active: bool) -> Result<()> {
+        let license = self.license.read_license().await?;
+
         if active {
-            // Check there's avaiable seat if switching user to active.
-            let license = self.license.read_license().await?;
+            // Check there's sufficient seat if switching user to active.
             license.ensure_available_seats(1)?;
         }
 
@@ -393,6 +400,11 @@ impl AuthenticationService for AuthenticationServiceImpl {
         if user.is_owner() {
             return Err(anyhow!("The owner's active status cannot be changed").into());
         }
+
+        if user.is_admin {
+            let num_admins = self.db.count_active_admin_users().await?;
+            license.ensure_admin_seats(num_admins + 1)?;
+        }
         Ok(self.db.update_user_active(id, active).await?)
     }
 }

ee/tabby-webserver/src/service/mod.rs

@@ -163,7 +163,7 @@ impl WorkerService for ServerContext {
             .await
             .map_err(|_| RegisterWorkerError::RequiresEnterpriseLicense)?;
 
-        if license.check_node_limit(count_workers) {
+        if license.check_node_limit(count_workers + 1) {
             return Err(RegisterWorkerError::RequiresEnterpriseLicense);
         }