From 3912fd00845af625e3f189cf2d9f74b1c3559f62 Mon Sep 17 00:00:00 2001
From: mdecimus <mauro@stalw.art>
Date: Sun, 8 Sep 2024 17:24:33 +0200
Subject: [PATCH] Fix: Account count is incorrect

---
 crates/common/src/enterprise/config.rs | 18 ++++++--------
 crates/common/src/lib.rs               | 34 ++++++++++++++++++++------
 crates/trc/src/event/level.rs          |  3 +--
 3 files changed, 35 insertions(+), 20 deletions(-)

diff --git a/crates/common/src/enterprise/config.rs b/crates/common/src/enterprise/config.rs
index f9c44103..6374156a 100644
--- a/crates/common/src/enterprise/config.rs
+++ b/crates/common/src/enterprise/config.rs
@@ -10,8 +10,7 @@
 
 use std::time::Duration;
 
-use jmap_proto::types::collection::Collection;
-use store::{BitmapKey, Store, Stores};
+use store::{Store, Stores};
 use trc::{EventType, MetricType, TOTAL_EVENT_COUNT};
 use utils::config::{
     cron::SimpleCron,
@@ -19,7 +18,10 @@ use utils::config::{
     Config,
 };
 
-use crate::expr::{tokenizer::TokenMap, Expression};
+use crate::{
+    expr::{tokenizer::TokenMap, Expression},
+    total_accounts,
+};
 
 use super::{
     license::LicenseValidator, AlertContent, AlertContentToken, AlertMethod, Enterprise,
@@ -40,17 +42,13 @@ impl Enterprise {
             }
         };
 
-        match data
-            .get_bitmap(BitmapKey::document_ids(u32::MAX, Collection::Principal))
-            .await
-        {
-            Ok(Some(bitmap)) if bitmap.len() > license.accounts as u64 => {
+        match total_accounts(data).await {
+            Ok(total) if total > license.accounts as u64 => {
                 config.new_build_warning(
                     "enterprise.license-key",
                     format!(
                         "License key is valid but only allows {} accounts, found {}.",
-                        license.accounts,
-                        bitmap.len()
+                        license.accounts, total
                     ),
                 );
                 return None;
diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs
index 339ab224..efe2f250 100644
--- a/crates/common/src/lib.rs
+++ b/crates/common/src/lib.rs
@@ -21,7 +21,6 @@ use config::{
 };
 use directory::{core::secret::verify_secret_hash, Directory, Principal, QueryBy, Type};
 use expr::if_block::IfBlock;
-use jmap_proto::types::collection::Collection;
 use listener::{
     blocked::{AllowedIps, BlockedIps},
     tls::TlsManager,
@@ -31,7 +30,7 @@ use mail_send::Credentials;
 use sieve::Sieve;
 use store::{
     write::{DirectoryClass, QueueClass, ValueClass},
-    BitmapKey, IterateParams, LookupStore, ValueKey,
+    IterateParams, LookupStore, ValueKey,
 };
 use tokio::sync::{mpsc, oneshot};
 use trc::AddContext;
@@ -343,12 +342,7 @@ impl Core {
     }
 
     pub async fn total_accounts(&self) -> trc::Result<u64> {
-        self.storage
-            .data
-            .get_bitmap(BitmapKey::document_ids(u32::MAX, Collection::Principal))
-            .await
-            .caused_by(trc::location!())
-            .map(|bitmap| bitmap.map_or(0, |b| b.len()))
+        total_accounts(&self.storage.data).await
     }
 
     pub async fn total_domains(&self) -> trc::Result<u64> {
@@ -376,6 +370,30 @@ impl Core {
     }
 }
 
+pub(crate) async fn total_accounts(store: &store::Store) -> trc::Result<u64> {
+    let mut total = 0;
+    store
+        .iterate(
+            IterateParams::new(
+                ValueKey::from(ValueClass::Directory(DirectoryClass::NameToId(vec![]))),
+                ValueKey::from(ValueClass::Directory(DirectoryClass::NameToId(vec![
+                    u8::MAX;
+                    10
+                ]))),
+            )
+            .ascending(),
+            |_, value| {
+                if matches!(value.last(), Some(0u8 | 4u8)) {
+                    total += 1;
+                }
+                Ok(true)
+            },
+        )
+        .await
+        .caused_by(trc::location!())
+        .map(|_| total)
+}
+
 trait CredentialsUsername {
     fn login(&self) -> &str;
 }
diff --git a/crates/trc/src/event/level.rs b/crates/trc/src/event/level.rs
index 5e553255..aafdd912 100644
--- a/crates/trc/src/event/level.rs
+++ b/crates/trc/src/event/level.rs
@@ -239,10 +239,9 @@ impl EventType {
                 ConfigEvent::DefaultApplied
                 | ConfigEvent::MissingSetting
                 | ConfigEvent::UnusedSetting
-                | ConfigEvent::ParseWarning
-                | ConfigEvent::BuildWarning
                 | ConfigEvent::AlreadyUpToDate
                 | ConfigEvent::ExternalKeyIgnored => Level::Debug,
+                ConfigEvent::ParseWarning | ConfigEvent::BuildWarning => Level::Warn,
                 ConfigEvent::ImportExternal => Level::Info,
             },
             EventType::Resource(cause) => match cause {