block adding users over license limits in openid and ldap

crates/defguard_core/src/enterprise/directory_sync/mod.rs

@@ -28,6 +28,8 @@ use crate::{
             model::ldap_sync_allowed_for_user,
             utils::{ldap_add_users_to_groups, ldap_delete_users, ldap_remove_users_from_groups},
         },
+        license::get_cached_license,
+        limits::{get_counts, update_counts},
     },
     grpc::GatewayEvent,
     handlers::user::check_username,
@@ -644,6 +646,11 @@ async fn sync_all_users_state(
     let mut modified_users = Vec::new();
     let mut deleted_users = Vec::new();
     let mut created_users = Vec::new();
+    let mut user_count = get_counts().user();
+    let user_limit = get_cached_license()
+        .as_ref()
+        .and_then(|license| license.limits.as_ref().map(|limits| limits.users));
+    let mut blocked_import_notification_sent = false;
 
     sync_inactive_directory_users(
         &mut transaction,
@@ -725,7 +732,20 @@ async fn sync_all_users_state(
                         details.phone_number.clone(),
                     );
                     user.openid_sub.clone_from(&directory_user.id);
+                    if let Some(limit) = user_limit.filter(|limit| user_count >= *limit) {
+                        error!(
+                            "Skipping directory sync import of user {} (email: {}) because \
+                            license user limit has been reached ({}/{})",
+                            user.username, user.email, user_count, limit
+                        );
+                        if !blocked_import_notification_sent {
+                            blocked_import_notification_sent = true;
+                            // TODO: send emails
+                        }
+                        continue;
+                    }
                     let new_user = user.save(&mut *transaction).await?;
+                    user_count += 1;
                     created_users.push(new_user);
                 }
             }
@@ -860,6 +880,7 @@ async fn sync_all_users_state(
     debug!("Done processing missing users");
 
     transaction.commit().await?;
+    update_counts(pool).await?;
 
     // trigger LDAP sync
     ldap_delete_users(deleted_users.iter().collect::<Vec<_>>(), pool).await;

crates/defguard_core/src/enterprise/directory_sync/tests.rs

@@ -19,7 +19,14 @@ mod test {
     use tokio::sync::broadcast;
 
     use super::super::*;
-    use crate::enterprise::db::models::openid_provider::{DirectorySyncTarget, OpenIdProviderKind};
+    use crate::{
+        enterprise::{
+            db::models::openid_provider::{DirectorySyncTarget, OpenIdProviderKind},
+            license::{License, LicenseTier, set_cached_license},
+            limits::{get_counts, update_counts},
+        },
+        grpc::proto::enterprise::license::LicenseLimits,
+    };
 
     async fn get_test_network(pool: &PgPool) -> WireguardNetwork<Id> {
         WireguardNetwork::find_by_name(pool, "test")
@@ -855,4 +862,55 @@ mod test {
         // No events
         assert!(wg_rx.try_recv().is_err());
     }
+
+    #[sqlx::test]
+    async fn test_users_prefetch_respects_license_user_limit(
+        _: PgPoolOptions,
+        options: PgConnectOptions,
+    ) {
+        let pool = setup_pool(options).await;
+
+        let config = DefGuardConfig::new_test_config();
+        let _ = SERVER_CONFIG.set(config.clone());
+        let (wg_tx, mut wg_rx) = broadcast::channel::<GatewayEvent>(16);
+
+        // enable prefetching users
+        make_test_provider(
+            &pool,
+            DirectorySyncUserBehavior::Keep,
+            DirectorySyncUserBehavior::Keep,
+            DirectorySyncTarget::All,
+            true,
+        )
+        .await;
+
+        let user_limit = 1;
+        let license = License::new(
+            "test".to_string(),
+            false,
+            None,
+            Some(LicenseLimits {
+                users: user_limit,
+                devices: 100,
+                locations: 100,
+                network_devices: Some(100),
+            }),
+            None,
+            LicenseTier::Business,
+        );
+        set_cached_license(Some(license));
+        update_counts(&pool).await.unwrap();
+
+        do_directory_sync(&pool, &wg_tx).await.unwrap();
+        update_counts(&pool).await.unwrap();
+
+        let user_count = get_counts().user();
+        assert!(user_count <= user_limit);
+
+        let defguard_users = User::all(&pool).await.unwrap();
+        assert_eq!(defguard_users.len(), user_limit as usize);
+
+        // No events
+        assert!(wg_rx.try_recv().is_err());
+    }
 }

crates/defguard_core/src/enterprise/handlers/openid_login.rs

@@ -38,8 +38,10 @@ use crate::{
     appstate::AppState,
     enterprise::{
         db::models::openid_provider::OpenIdProvider,
-        directory_sync::sync_user_groups_if_configured, ldap::utils::ldap_update_user_state,
-        limits::update_counts,
+        directory_sync::sync_user_groups_if_configured,
+        ldap::utils::ldap_update_user_state,
+        license::get_cached_license,
+        limits::{get_counts, update_counts},
     },
     error::WebError,
     handlers::{
@@ -94,6 +96,16 @@ pub fn prune_username(username: &str, handling: OpenIdUsernameHandling) -> Strin
     result
 }
 
+fn reached_user_license_limit() -> Option<(u32, u32)> {
+    let user_count = get_counts().user();
+    let user_limit = get_cached_license()
+        .as_ref()
+        .and_then(|license| license.limits.as_ref().map(|limits| limits.users));
+    user_limit
+        .filter(|limit| user_count >= *limit)
+        .map(|limit| (user_count, limit))
+}
+
 /// Create HTTP client and prevent following redirects
 fn get_async_http_client() -> Result<reqwest::Client, WebError> {
     reqwest::Client::builder()
@@ -365,6 +377,19 @@ pub async fn user_from_claims(
                     )));
                 }
 
+                if let Some((user_count, limit)) = reached_user_license_limit() {
+                    error!(
+                        "Skipping OpenID account creation for user {} (email: {}) because \
+                        license user limit has been reached ({}/{})",
+                        username,
+                        email.as_str(),
+                        user_count,
+                        limit
+                    );
+                    // TODO: send emails
+                    return Err(WebError::Forbidden("License limit reached.".into()));
+                }
+
                 // Extract all necessary information from the token or call the userinfo endpoint.
                 let given_name = token_claims
                     .given_name()
@@ -644,6 +669,14 @@ pub(crate) async fn auth_callback(
 
 #[cfg(test)]
 mod test {
+    use crate::{
+        enterprise::{
+            license::{License, LicenseTier, set_cached_license},
+            limits::{Counts, set_counts},
+        },
+        grpc::proto::enterprise::license::LicenseLimits,
+    };
+
     use super::*;
 
     #[test]
@@ -762,4 +795,62 @@ mod test {
         let extracted = extract_state_data(&encoded);
         assert_eq!(extracted, Some("data.with.dots".to_string()));
     }
+
+    #[test]
+    fn test_reached_user_license_limit_reached() {
+        set_counts(Counts::new(2, 0, 0, 0));
+        let license = License::new(
+            "test".to_string(),
+            false,
+            None,
+            Some(LicenseLimits {
+                users: 2,
+                devices: 100,
+                locations: 100,
+                network_devices: Some(100),
+            }),
+            None,
+            LicenseTier::Business,
+        );
+        set_cached_license(Some(license));
+
+        assert_eq!(reached_user_license_limit(), Some((2, 2)));
+    }
+
+    #[test]
+    fn test_reached_user_license_limit_not_reached() {
+        set_counts(Counts::new(1, 0, 0, 0));
+        let license = License::new(
+            "test".to_string(),
+            false,
+            None,
+            Some(LicenseLimits {
+                users: 2,
+                devices: 100,
+                locations: 100,
+                network_devices: Some(100),
+            }),
+            None,
+            LicenseTier::Business,
+        );
+        set_cached_license(Some(license));
+
+        assert_eq!(reached_user_license_limit(), None);
+    }
+
+    #[test]
+    fn test_reached_user_license_limit_unlimited() {
+        set_counts(Counts::new(100, 0, 0, 0));
+        let license = License::new(
+            "test".to_string(),
+            false,
+            None,
+            None,
+            None,
+            LicenseTier::Business,
+        );
+        set_cached_license(Some(license));
+
+        assert_eq!(reached_user_license_limit(), None);
+    }
 }

crates/defguard_core/src/enterprise/ldap/error.rs

@@ -28,6 +28,8 @@ pub enum LdapError {
     InvalidUsername(String),
     #[error("LDAP object already exists: {0}")]
     ObjectAlreadyExists(String),
+    #[error("License user limit reached: {0}/{1}")]
+    LicenseUserLimitReached(u32, u32),
     #[error("User {0} does not belong to the defined synchronization groups in {1}")]
     UserNotInLDAPSyncGroups(String, &'static str),
 }

crates/defguard_core/src/enterprise/ldap/sync.rs

@@ -66,9 +66,13 @@ use sqlx::{PgConnection, PgPool};
 
 use super::{LDAPConfig, error::LdapError};
 use crate::{
-    enterprise::ldap::model::{
-        get_users_without_ldap_path, ldap_sync_allowed_for_user, update_from_ldap_user,
-        user_from_searchentry,
+    enterprise::{
+        ldap::model::{
+            get_users_without_ldap_path, ldap_sync_allowed_for_user, update_from_ldap_user,
+            user_from_searchentry,
+        },
+        license::get_cached_license,
+        limits::{get_counts, update_counts},
     },
     hashset,
 };
@@ -806,6 +810,13 @@ impl super::LDAPConnection {
     ) -> Result<(), LdapError> {
         let mut transaction = pool.begin().await?;
         let mut admin_count = User::find_admins(&mut *transaction).await?.len();
+        let mut user_count = get_counts().user();
+
+        let user_limit = get_cached_license()
+            .as_ref()
+            .and_then(|license| license.limits.as_ref().map(|limits| limits.users));
+        let mut blocked_import_notification_sent = false;
+
         for user in changes.delete_defguard {
             if user.is_admin(&mut *transaction).await? {
                 if admin_count == 1 {
@@ -849,12 +860,27 @@ impl super::LDAPConnection {
                     "LDAP user {} does not exist in Defguard yet, adding...",
                     user.username
                 );
+                if let Some(limit) = user_limit.filter(|limit| user_count >= *limit) {
+                    error!(
+                        "Skipping LDAP import of user {} (email: {}) because license user limit \
+                        has been reached ({}/{})",
+                        user.username, user.email, user_count, limit
+                    );
+                    if !blocked_import_notification_sent {
+                        blocked_import_notification_sent = true;
+                        // TODO: send emails
+                    }
+                    continue;
+                }
                 user.save(&mut *transaction).await?;
+                user_count += 1;
             }
         }
 
         transaction.commit().await?;
 
+        update_counts(pool).await?;
+
         for user in changes.delete_ldap {
             debug!("Deleting user {} from LDAP", user.username);
             self.delete_user(&user).await?;