Check in the NonceManager contract and its Foundry test

Author: forshtat

contracts/predeploys/NonceManager.sol

@@ -0,0 +1,72 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.15;
+
+/// @title NonceManager
+/// @notice The NonceManager manages nonce of smart accounts using RIP-7560.
+contract NonceManager {
+    /// @notice Semantic version.
+    /// @custom:semver 0.1.0
+    string public constant version = "0.1.0";
+
+    /// @notice The EntryPoint address defined at RIP-7560.
+    address internal constant AA_ENTRY_POINT = 0x0000000000000000000000000000000000007560;
+
+    fallback(bytes calldata data) external returns (bytes memory) {
+        if (msg.sender == AA_ENTRY_POINT) {
+            _validateIncrement(data);
+            return new bytes(0);
+        } else {
+            return abi.encodePacked(_get(data));
+        }
+    }
+
+    /// @notice Return the next nonce for this sender. Within a given key, the nonce values are sequenced
+    ///         (starting with zero, and incremented by one on each transaction).
+    ///         But transactions with different keys can come with arbitrary order.
+    /// @return nonce a full nonce to pass for next transaction with given sender and key.
+    function _get(bytes calldata /* data */) internal view returns (uint256 nonce) {
+        assembly {
+        // Check if calldata is 44 bytes long
+            if iszero(eq(calldatasize(), 44)) {
+                mstore(0x00, 0x947d5a84) // 'InvalidLength()'
+                revert(0x1c, 0x04)
+            }
+
+            let ptr := mload(0x40)
+            calldatacopy(ptr, 0, 44)
+
+        // Extract key and sender from calldata
+            let key := shr(64, mload(add(ptr, 20)))
+            mstore(0x00, key)
+            mstore(0x14, shr(96, mload(ptr)))
+
+        // Load nonce from storage
+            nonce := or(shl(64, key), sload(keccak256(0x04, 0x30)))
+        }
+    }
+
+    /// @notice validate nonce uniqueness for this account. Called by AA_ENTRY_POINT.
+    function _validateIncrement(bytes calldata /* data */) internal {
+        assembly {
+            let ptr := mload(0x40)
+            calldatacopy(ptr, 0, calldatasize())
+
+        // Store key and sender in memory
+            mstore(0x00, shr(64, mload(add(ptr, 20))))
+            mstore(0x14, shr(96, mload(ptr)))
+
+        // Calculate storage slot and load current nonce
+            let nonceSlot := keccak256(0x04, 0x30)
+            let currentNonce := sload(nonceSlot)
+
+        // Revert if nonce mismatch
+            if iszero(eq(shr(192, mload(add(ptr, 44))), currentNonce)) {
+                mstore(0, 0)
+                revert(0, 0) // Revert if nonce mismatch
+            }
+
+        // Increment nonce
+            sstore(nonceSlot, add(currentNonce, 1))
+        }
+    }
+}

contracts/test/NonceManager.t.sol

@@ -0,0 +1,73 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.15;
+
+// Testing utilities
+import { Test } from "forge-std/Test.sol";
+import "../predeploys/NonceManager.sol";
+
+contract NonceManagerTest is Test {
+    address private constant aaEntryPoint = 0x0000000000000000000000000000000000007560;
+    address private constant alice = 0x0000000000000000000000000000000000007777;
+    uint192 public aliceKey;
+    error InvalidLength();
+    NonceManager nonceManager;
+
+    /// @dev Sets up the test suite.
+    function setUp() public virtual {
+        nonceManager = new NonceManager();
+        aliceKey = uint192(3);
+    }
+
+    /// @dev Tests that increasing nonce is done properly.
+    function test_validateIncrement_succeeds() public {
+        vm.prank(aaEntryPoint);
+        (bool success, bytes memory returnData) =
+                                address(nonceManager).call(abi.encodePacked(alice, aliceKey, uint64(0)));
+        assertTrue(success);
+        assertEq(returnData, new bytes(0));
+    }
+
+    /// @dev Tests that increasing nonce with invalid nonce fails.
+    function test_validateIncrement_invalidNonce_fails() external {
+        vm.prank(aaEntryPoint);
+        (bool success,) = address(nonceManager).call(abi.encodePacked(alice, aliceKey, uint64(1)));
+        assertFalse(success);
+    }
+
+    /// @dev Tests that getting nonce through fallback is done properly.
+    function test_fallback_get_succeeds() external {
+        (bool success, bytes memory returnData) = address(nonceManager).call(abi.encodePacked(alice, aliceKey));
+        assertTrue(success);
+        assertEq(returnData, abi.encodePacked(uint192(aliceKey), uint64(0)));
+
+        test_validateIncrement_succeeds();
+
+        (success, returnData) = address(nonceManager).call(abi.encodePacked(alice, aliceKey));
+        assertTrue(success);
+        assertEq(returnData, abi.encodePacked(uint192(aliceKey), uint64(1)));
+    }
+
+    /// @dev Tests that getting nonce with invalid length fails.
+    function test_fallback_get_invalidLength_fails() external {
+        vm.expectRevert(InvalidLength.selector);
+        (bool success,) = address(nonceManager).call(abi.encodePacked(alice, aliceKey, uint64(1)));
+        (success); // silence unused variable warning
+    }
+
+    /// @dev Fuzz test for validateIncrement function.
+    function testFuzz_validateIncrement_succeeds(uint256 n) external {
+        bool success;
+        bytes memory returnData;
+
+        n = n % 1000; // limit the number of iterations to prevent OOG
+        for (uint256 i = 0; i < n; i++) {
+            vm.prank(aaEntryPoint);
+            (success,) = address(nonceManager).call(abi.encodePacked(alice, aliceKey, uint64(i)));
+            assertTrue(success);
+        }
+
+        (success, returnData) = address(nonceManager).call(abi.encodePacked(alice, aliceKey));
+        assertTrue(success);
+        assertEq(returnData, abi.encodePacked(uint192(aliceKey), uint64(n)));
+    }
+}