fix: add hasReviewed with TRANSIENT_SIGNER_SLOT

EtherTW · Jan 22, 2025 · fff6ef0 · fff6ef0
1 parent 2d41ba4
commit fff6ef0
Show file tree

Hide file tree

Showing 2 changed files with 105 additions and 8 deletions.
diff --git a/src/RoyaltyAutoClaim.sol b/src/RoyaltyAutoClaim.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.8.27;
+pragma solidity ^0.8.28;
 
 import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
 import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
@@ -30,6 +30,7 @@ contract RoyaltyAutoClaim is UUPSUpgradeable, OwnableUpgradeable, IAccount, Reen
     error ForbiddenPaymaster();
     error ZeroRoyalty();
     error UnsupportSelector(bytes4 selector);
+    error AlreadyReviewed();
 
     uint8 public constant ROYALTY_LEVEL_20 = 20;
     uint8 public constant ROYALTY_LEVEL_40 = 40;
@@ -45,9 +46,9 @@ contract RoyaltyAutoClaim is UUPSUpgradeable, OwnableUpgradeable, IAccount, Reen
 
     struct Submission {
         address royaltyRecipient;
-        uint8 reviewCount;
         uint16 totalRoyaltyLevel;
         SubmissionStatus status;
+        uint8 reviewCount;
     }
 
     enum SubmissionStatus {
@@ -60,15 +61,20 @@ contract RoyaltyAutoClaim is UUPSUpgradeable, OwnableUpgradeable, IAccount, Reen
     struct MainStorage {
         Configs configs;
         mapping(string => Submission) submissions;
+        mapping(string => mapping(address => bool)) hasReviewed;
     }
 
+    // address transient userOpSigner;
+
     // keccak256(abi.encode(uint256(keccak256("royaltyautoclaim.storage.main")) - 1)) & ~bytes32(uint256(0xff));
     /// @dev cast index-erc7201 royaltyautoclaim.storage.main
-    bytes32 private constant MAIN_STORAGE_LOCATION = 0x41a2efc794119f946ab405955f96dacdfa298d25a3ae81c9a8cc1dea5771a900;
+    bytes32 private constant MAIN_STORAGE_SLOT = 0x41a2efc794119f946ab405955f96dacdfa298d25a3ae81c9a8cc1dea5771a900;
+    // @dev cast index-erc7201 royaltyautoclaim.storage.signer
+    bytes32 private constant TRANSIENT_SIGNER_SLOT = 0xbbc49793e8d16b6166d591f0a7a95f88efe9e6a08bf1603701d7f0fe05d7d600;
 
     function _getMainStorage() private pure returns (MainStorage storage $) {
         assembly {
-            $.slot := MAIN_STORAGE_LOCATION
+            $.slot := MAIN_STORAGE_SLOT
         }
     }
 
@@ -198,11 +204,28 @@ contract RoyaltyAutoClaim is UUPSUpgradeable, OwnableUpgradeable, IAccount, Reen
                 || royaltyLevel == ROYALTY_LEVEL_80,
             InvalidRoyaltyLevel(royaltyLevel)
         );
+
+        address reviewer = _getReviewer();
         MainStorage storage $ = _getMainStorage();
+        if ($.hasReviewed[title][reviewer]) {
+            revert AlreadyReviewed();
+        }
+        $.hasReviewed[title][reviewer] = true;
         $.submissions[title].reviewCount++;
         $.submissions[title].totalRoyaltyLevel += royaltyLevel;
     }
 
+    function _getReviewer() internal view returns (address) {
+        if (msg.sender == entryPoint()) {
+            address reviewer;
+            assembly {
+                reviewer := tload(TRANSIENT_SIGNER_SLOT)
+            }
+            return reviewer;
+        }
+        return msg.sender;
+    }
+
     // ================================ Submitter ================================
 
     function claimRoyalty(string memory title) public nonReentrant {
@@ -271,6 +294,11 @@ contract RoyaltyAutoClaim is UUPSUpgradeable, OwnableUpgradeable, IAccount, Reen
             if (!isReviewer(signer)) {
                 revert Unauthorized(signer);
             }
+
+            assembly {
+                tstore(TRANSIENT_SIGNER_SLOT, signer)
+            }
+
             return 0;
             // Anybody
         } else if (selector == this.claimRoyalty.selector) {

diff --git a/test/RoyaltyAutoClaim.t.sol b/test/RoyaltyAutoClaim.t.sol
@@ -403,20 +403,33 @@ contract RoyaltyAutoClaimTest is AATest {
         vm.prank(admin);
         RoyaltyAutoClaim(address(proxy)).registerSubmission("test", submitter);
 
-        // Test all valid royalty levels
+        // Add more reviewers for testing
+        address[] memory testReviewers = new address[](4);
+        bool[] memory status = new bool[](4);
+        for (uint256 i = 0; i < 4; i++) {
+            testReviewers[i] = vm.addr(10 + i);
+            status[i] = true;
+        }
+        vm.prank(admin);
+        RoyaltyAutoClaim(address(proxy)).updateReviewers(testReviewers, status);
+
+        // Test all valid royalty levels with different reviewers
         uint16[] memory validLevels = new uint16[](4);
         validLevels[0] = 20;
         validLevels[1] = 40;
         validLevels[2] = 60;
         validLevels[3] = 80;
 
+        uint256 expectedTotalLevel = 0;
         for (uint256 i = 0; i < validLevels.length; i++) {
-            vm.prank(reviewers[0]);
+            vm.prank(testReviewers[i]);
             RoyaltyAutoClaim(address(proxy)).reviewSubmission("test", validLevels[i]);
 
+            expectedTotalLevel += validLevels[i];
+
             RoyaltyAutoClaim.Submission memory submission = RoyaltyAutoClaim(address(proxy)).submissions("test");
             assertEq(submission.reviewCount, i + 1, "Review count should increment");
-            assertEq(submission.totalRoyaltyLevel, validLevels[i], "Total royalty level should match");
+            assertEq(submission.totalRoyaltyLevel, expectedTotalLevel, "Total royalty level should be cumulative");
         }
     }
 
@@ -462,10 +475,42 @@ contract RoyaltyAutoClaimTest is AATest {
 
         // Second review from same reviewer should fail
         vm.prank(reviewers[0]);
-        vm.expectRevert();
+        vm.expectRevert(abi.encodeWithSelector(RoyaltyAutoClaim.AlreadyReviewed.selector));
         RoyaltyAutoClaim(address(proxy)).reviewSubmission("test", 40);
     }
 
+    function testCannot_reviewSubmission_multiple_times_4337() public {
+        address submitter = vm.randomAddress();
+        address reviewer = vm.addr(0xbeef);
+
+        // Register reviewer
+        address[] memory testReviewers = new address[](1);
+        bool[] memory status = new bool[](1);
+        testReviewers[0] = reviewer;
+        status[0] = true;
+
+        vm.prank(admin);
+        RoyaltyAutoClaim(address(proxy)).updateReviewers(testReviewers, status);
+
+        // Register submission
+        vm.prank(admin);
+        RoyaltyAutoClaim(address(proxy)).registerSubmission("test", submitter);
+
+        // First review should succeed via UserOperation
+        PackedUserOperation memory userOp =
+            buildUserOp(0xbeef, abi.encodeCall(RoyaltyAutoClaim.reviewSubmission, ("test", 20)));
+        handleUserOp(userOp);
+
+        // Second review from same reviewer should fail
+        userOp = buildUserOp(0xbeef, abi.encodeCall(RoyaltyAutoClaim.reviewSubmission, ("test", 40)));
+
+        vm.expectEmit(false, true, true, true);
+        emit IEntryPoint.UserOperationRevertReason(
+            bytes32(0), address(proxy), userOp.nonce, abi.encodeWithSelector(RoyaltyAutoClaim.AlreadyReviewed.selector)
+        );
+        handleUserOp(userOp);
+    }
+
     // ======================================== Submitter Functions ========================================
 
     function test_claimRoyalty_with_erc20() public {
@@ -546,6 +591,30 @@ contract RoyaltyAutoClaimTest is AATest {
         RoyaltyAutoClaim(address(proxy)).claimRoyalty("test");
     }
 
+    function testCannot_claimRoyalty_if_already_claimed_4337() public {
+        address submitter = vm.addr(0xbeef);
+
+        // Setup submission and reviews
+        vm.prank(admin);
+        RoyaltyAutoClaim(address(proxy)).registerSubmission("test", submitter);
+        vm.prank(reviewers[0]);
+        RoyaltyAutoClaim(address(proxy)).reviewSubmission("test", 20);
+        vm.prank(reviewers[1]);
+        RoyaltyAutoClaim(address(proxy)).reviewSubmission("test", 40);
+
+        // First claim should succeed via UserOperation
+        PackedUserOperation memory userOp = buildUserOp(0xbeef, abi.encodeCall(RoyaltyAutoClaim.claimRoyalty, ("test")));
+        handleUserOp(userOp);
+
+        // Second claim should fail
+        userOp = buildUserOp(0xbeef, abi.encodeCall(RoyaltyAutoClaim.claimRoyalty, ("test")));
+        vm.expectEmit(false, true, true, true);
+        emit IEntryPoint.UserOperationRevertReason(
+            bytes32(0), address(proxy), userOp.nonce, abi.encodeWithSelector(RoyaltyAutoClaim.AlreadyClaimed.selector)
+        );
+        handleUserOp(userOp);
+    }
+
     function testCannot_claimRoyalty_if_not_enough_reviews() public {
         address submitter = vm.randomAddress();