Further fixing BOA-04 from CertiK audit

openfort-xyz · Dec 22, 2023 · 9597226 · 9597226
1 parent 89ad211
commit 9597226
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 4 deletions.
diff --git a/contracts/core/base/BaseOpenfortAccount.sol b/contracts/core/base/BaseOpenfortAccount.sol
@@ -141,16 +141,18 @@ abstract contract BaseOpenfortAccount is
             return false; // All other cases, deny
         } else if (funcSelector == EXECUTEBATCH_SELECTOR) {
             (address[] memory toContracts,,) = abi.decode(_callData[4:], (address[], uint256[], bytes[]));
+            uint256 numberOfInteractions = toContracts.length;
+            if (numberOfInteractions > 9) return false;
             if (!sessionKey.masterSessionKey) {
-                // Check if limit of transactions per sessionKey reached
-                if (sessionKey.limit < toContracts.length || toContracts.length > 9) return false;
+                // Check if limit of transactions per sessionKey is reached
+                if (sessionKey.limit < numberOfInteractions) return false;
                 unchecked {
-                    sessionKey.limit = sessionKey.limit - SafeCastUpgradeable.toUint48(toContracts.length);
+                    sessionKey.limit = sessionKey.limit - SafeCastUpgradeable.toUint48(numberOfInteractions);
                 }
             }
 
             uint256 i;
-            for (i; i < toContracts.length;) {
+            for (i; i < numberOfInteractions;) {
                 // Check if reenter, do not allow
                 if (toContracts[i] == address(this)) return false;
 

diff --git a/test/foundry/core/upgradeable/UpgradeableOpenfortAccountTest.t.sol b/test/foundry/core/upgradeable/UpgradeableOpenfortAccountTest.t.sol
@@ -281,6 +281,88 @@ contract UpgradeableOpenfortAccountTest is OpenfortBaseTest {
         assertEq(testCounter.counters(accountAddress), count);
     }
 
+    /*
+     * Use the executeBatch() with a big number
+     */
+    function testBatchingBigMasterSessionKey() public {
+        // Verify that the counter is still set to 0
+        assertEq(testCounter.counters(accountAddress), 0);
+
+        address sessionKey;
+        uint256 sessionKeyPrivKey;
+        (sessionKey, sessionKeyPrivKey) = makeAddrAndKey("sessionKey");
+        address[] memory emptyWhitelist;
+
+        // Register masterSessionKey
+        vm.prank(openfortAdmin);
+        IBaseRecoverableAccount(payable(accountAddress)).registerSessionKey(
+            sessionKey, 0, 2 ** 48 - 1, 2 ** 48 - 1, emptyWhitelist
+        );
+
+        uint256 count = 9;
+        address[] memory targets = new address[](count);
+        uint256[] memory values = new uint256[](count);
+        bytes[] memory callData = new bytes[](count);
+
+        for (uint256 i = 0; i < count; i += 1) {
+            targets[i] = address(testCounter);
+            values[i] = 0;
+            callData[i] = abi.encodeWithSignature("count()");
+        }
+
+        UserOperation[] memory userOp =
+            _setupUserOpExecuteBatch(accountAddress, sessionKeyPrivKey, bytes(""), targets, values, callData);
+
+        entryPoint.depositTo{value: 1 ether}(accountAddress);
+        vm.expectRevert();
+        entryPoint.simulateValidation(userOp[0]);
+        entryPoint.handleOps(userOp, beneficiary);
+
+        // Verify that the counter has increased
+        assertEq(testCounter.counters(accountAddress), count);
+    }
+
+    /*
+     * Use the executeBatch() with a "too big" number
+     */
+    function testFailBatchingBigMasterSessionKey() public {
+        // Verify that the counter is still set to 0
+        assertEq(testCounter.counters(accountAddress), 0);
+
+        address sessionKey;
+        uint256 sessionKeyPrivKey;
+        (sessionKey, sessionKeyPrivKey) = makeAddrAndKey("sessionKey");
+        address[] memory emptyWhitelist;
+
+        // Register masterSessionKey
+        vm.prank(openfortAdmin);
+        IBaseRecoverableAccount(payable(accountAddress)).registerSessionKey(
+            sessionKey, 0, 2 ** 48 - 1, 2 ** 48 - 1, emptyWhitelist
+        );
+
+        uint256 count = 10;
+        address[] memory targets = new address[](count);
+        uint256[] memory values = new uint256[](count);
+        bytes[] memory callData = new bytes[](count);
+
+        for (uint256 i = 0; i < count; i += 1) {
+            targets[i] = address(testCounter);
+            values[i] = 0;
+            callData[i] = abi.encodeWithSignature("count()");
+        }
+
+        UserOperation[] memory userOp =
+            _setupUserOpExecuteBatch(accountAddress, sessionKeyPrivKey, bytes(""), targets, values, callData);
+
+        entryPoint.depositTo{value: 1 ether}(accountAddress);
+        vm.expectRevert();
+        entryPoint.simulateValidation(userOp[0]);
+        entryPoint.handleOps(userOp, beneficiary);
+
+        // Verify that the counter has increased
+        assertEq(testCounter.counters(accountAddress), count);
+    }
+
     /*
      *  Should fail, try to use a sessionKey that is not registered.
      */