From 55e7f6c3e24f54410004ea0b33602bcc6237d3e9 Mon Sep 17 00:00:00 2001
From: Sydhds <sylvain@status.im>
Date: Tue, 30 Sep 2025 11:55:43 +0200
Subject: [PATCH] Print an error if too many users are registered (#40)

---
 prover/src/registry_listener.rs | 12 +++++++++++-
 prover/src/user_db.rs           | 20 ++++++++++++++++++--
 prover/src/user_db_error.rs     |  2 ++
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/prover/src/registry_listener.rs b/prover/src/registry_listener.rs
index ff76dd3..6a66595 100644
--- a/prover/src/registry_listener.rs
+++ b/prover/src/registry_listener.rs
@@ -139,7 +139,17 @@ impl RegistryListener {
                 let id_commitment = self
                     .user_db
                     .on_new_user(&to_address)
-                    .map_err(HandleTransferError::Register)?;
+                    .map_err(HandleTransferError::Register);
+
+                // Don't stop the registry_listener if the user_db is full
+                // Prover will still be functional
+                if let Err(HandleTransferError::Register(RegisterError::TooManyUsers)) =
+                    id_commitment
+                {
+                    error!("Cannot register a new user: {:?}", id_commitment);
+                }
+
+                let id_commitment = id_commitment?;
 
                 let id_co =
                     U256::from_le_slice(BigUint::from(id_commitment).to_bytes_le().as_slice());
diff --git a/prover/src/user_db.rs b/prover/src/user_db.rs
index 6f06c35..857be46 100644
--- a/prover/src/user_db.rs
+++ b/prover/src/user_db.rs
@@ -17,6 +17,10 @@ use rocksdb::{
 };
 use serde::{Deserialize, Serialize};
 use tracing::error;
+use zerokit_utils::{
+    error::ZerokitMerkleTreeError,
+    pmtree::{PmtreeErrorKind, TreeErrorKind},
+};
 // internal
 use crate::epoch_service::{Epoch, EpochSlice};
 use crate::error::GetMerkleTreeProofError;
@@ -262,7 +266,19 @@ impl UserDb {
                 self.merkle_tree
                     .write()
                     .set(new_index.into(), rate_commit)
-                    .map_err(|e| RegisterError::TreeError(e.to_string()))?;
+                    .map_err(|e| {
+                        // Check Zerokit issue: https://github.com/vacp2p/zerokit/issues/343
+                        if matches!(
+                            e,
+                            ZerokitMerkleTreeError::PmtreeErrorKind(PmtreeErrorKind::TreeError(
+                                TreeErrorKind::IndexOutOfBounds
+                            ))
+                        ) {
+                            RegisterError::TooManyUsers
+                        } else {
+                            RegisterError::TreeError(e.to_string())
+                        }
+                    })?;
 
                 // Add index for user
                 merkle_index_serializer.serialize(&new_index, &mut buffer);
@@ -805,7 +821,7 @@ mod tests {
         assert_eq!(tree.read().leaves_set(), 2);
 
         let res = user_db.register(ADDR_2);
-        assert_matches!(res, Err(RegisterError::TreeError(_)));
+        assert_matches!(res, Err(RegisterError::TooManyUsers));
         assert_eq!(user_db.has_user(&ADDR_1), Ok(true));
         assert_eq!(user_db.has_user(&ADDR_2), Ok(false));
         assert_eq!(tree.read().leaves_set(), 2);
diff --git a/prover/src/user_db_error.rs b/prover/src/user_db_error.rs
index 879d679..d0c8299 100644
--- a/prover/src/user_db_error.rs
+++ b/prover/src/user_db_error.rs
@@ -27,6 +27,8 @@ pub enum RegisterError {
     AlreadyRegistered(Address),
     #[error(transparent)]
     Db(#[from] rocksdb::Error),
+    #[error("Too many users, exceeding merkle tree capacity...")]
+    TooManyUsers,
     #[error("Merkle tree error: {0}")]
     TreeError(String),
 }