Improve RecoveryRequest

zkemail · Sep 21, 2024 · 1e94cc1 · 1e94cc1
1 parent b2ee542
commit 1e94cc1
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 10 deletions.
diff --git a/src/EmailRecoveryManager.sol b/src/EmailRecoveryManager.sol
@@ -35,7 +35,7 @@ abstract contract EmailRecoveryManager is
     /*                    CONSTANTS & STORAGE                     */
     /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
     using EnumerableMap for EnumerableMap.Bytes32ToUintMap;
-
+    using EnumerableMap for EnumerableMap.AddressToUintMap;
     /**
      * Minimum required time window between when a recovery attempt becomes valid and when it
      * becomes invalid
@@ -131,7 +131,10 @@ abstract contract EmailRecoveryManager is
         view
         returns (bool)
     {
-        return recoveryRequests[account].guardianVoted[guardian];
+        if(!recoveryRequests[account].guardianVoted.contains(guardian)) {
+            return false;
+        }
+        return recoveryRequests[account].guardianVoted.get(guardian) == 1;
     }
 
     /**
@@ -379,15 +382,18 @@ abstract contract EmailRecoveryManager is
         bytes32 recoveryDataHash = IEmailRecoveryCommandHandler(commandHandler)
             .parseRecoveryDataHash(templateIdx, commandParams);
 
-        if (recoveryRequest.guardianVoted[guardian]) {
+        if (hasGuardianVoted(account, guardian)) {
             revert GuardianAlreadyVoted();
         }
 
         uint256 executeBefore = block.timestamp + recoveryConfigs[account].expiry;
         recoveryRequest.executeBefore = executeBefore;
-        recoveryRequest.guardianVoted[guardian] = true;
+        recoveryRequest.guardianVoted.set(guardian,1);
 
-        uint256 currentWeight = recoveryRequest.recoveryDataHashWeight.get(recoveryDataHash);
+        uint256 currentWeight = 0;
+        if(recoveryRequest.recoveryDataHashWeight.contains(recoveryDataHash)) {
+            currentWeight = recoveryRequest.recoveryDataHashWeight.get(recoveryDataHash);
+        }
         recoveryRequest.recoveryDataHashWeight.set(recoveryDataHash, currentWeight += guardianStorage.weight);
 
         if (recoveryRequest.recoveryDataHashWeight.get(recoveryDataHash) >= guardianConfig.threshold) {
@@ -431,7 +437,10 @@ abstract contract EmailRecoveryManager is
         }
 
         bytes32 recoveryDataHash = keccak256(recoveryData);
-        uint256 currentWeight = recoveryRequest.recoveryDataHashWeight.get(recoveryDataHash);
+        uint256 currentWeight = 0;
+        if(recoveryRequest.recoveryDataHashWeight.contains(recoveryDataHash)) {
+            currentWeight = recoveryRequest.recoveryDataHashWeight.get(recoveryDataHash);
+        }
         if (currentWeight < threshold) {
             revert NotEnoughApprovals(currentWeight, threshold);
         }
@@ -443,9 +452,8 @@ abstract contract EmailRecoveryManager is
         if (block.timestamp >= recoveryRequest.executeBefore) {
             revert RecoveryRequestExpired(block.timestamp, recoveryRequest.executeBefore);
         }
-
-        delete recoveryRequests[account];
-        // TODO: delete mapping values
+
+        clearRecoveryRequest(account);
 
         recover(account, recoveryData);
 
@@ -529,4 +537,27 @@ abstract contract EmailRecoveryManager is
 
         emit RecoveryDeInitialized(account);
     }
+
+    function clearRecoveryRequest(address account) internal {
+        RecoveryRequest storage recoveryRequest = recoveryRequests[account];
+        uint256 lengthWeights = recoveryRequest.recoveryDataHashWeight.length();
+        bytes32[] memory keysWeights = new bytes32[](lengthWeights);
+        for (uint256 i = 0; i < lengthWeights; i++) {
+            (bytes32 key, ) = recoveryRequest.recoveryDataHashWeight.at(i);
+            keysWeights[i] = key;
+        }
+        for(uint256 i = 0; i < lengthWeights; i++) {
+            recoveryRequest.recoveryDataHashWeight.remove(keysWeights[i]);
+        }
+        uint256 lengthVotes = recoveryRequest.guardianVoted.length();
+        address[] memory keysVotes = new address[](lengthVotes);
+        for (uint256 i = 0; i < lengthVotes; i++) {
+            (address key, ) = recoveryRequest.guardianVoted.at(i);
+            keysVotes[i] = key;
+        }
+        for(uint256 i = 0; i < lengthVotes; i++) {
+            recoveryRequest.guardianVoted.remove(keysVotes[i]);
+        }
+        delete recoveryRequests[account];
+    }
 }
diff --git a/src/interfaces/IEmailRecoveryManager.sol b/src/interfaces/IEmailRecoveryManager.sol
@@ -33,7 +33,7 @@ interface IEmailRecoveryManager {
         uint256 executeBefore; // the timestamp from which the recovery request becomes invalid
         // mapping(bytes32 recoveryDataHash => uint256 currentWeight) recoveryDataHashWeight; // total weight of all guardian approvals for the recovery request
         EnumerableMap.Bytes32ToUintMap recoveryDataHashWeight;
-        mapping(address guardian => bool isVoted) guardianVoted; // the keccak256 hash of the recovery data used to execute the recovery attempt
+        EnumerableMap.AddressToUintMap guardianVoted;
     }
 
     /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/