feat(wip): completed removeMember() and test case (#18)

* updated * added removeMember test
jimmychu0807 · Jan 13, 2025 · 13e6672 · 13e6672
1 parent 5e6ae75
commit 13e6672
Show file tree

Hide file tree

Showing 5 changed files with 118 additions and 63 deletions.
diff --git a/package.json b/package.json
@@ -18,9 +18,9 @@
   "devDependencies": {
     "@rhinestone/modulekit": "~0.5.4",
     "@semaphore-protocol/contracts": "github:jimmychu0807/semaphore#identity-cli&path:/packages/contracts/contracts",
-    "@semaphore-protocol/core": "github:jimmychu0807/semaphore#identity-cli&path:/packages/core",
     "@semaphore-protocol/identity": "github:jimmychu0807/semaphore#identity-cli&path:/packages/identity",
     "@semaphore-protocol/proof": "github:jimmychu0807/semaphore#identity-cli&path:/packages/proof",
+    "@semaphore-protocol/group": "github:jimmychu0807/semaphore#identity-cli&path:/packages/group",
     "poseidon-solidity": "github:chancehudson/poseidon-solidity#main",
     "rimraf": "^5.0.5",
     "solady": "^0.0.287"
@@ -47,7 +47,7 @@
     "prepack": "pnpm install && bash ./shell/prepare-artifacts.sh",
     "prettier:check": "prettier --no-error-on-unmatched-pattern -c \"{src,test,script}/**/*.{json,md,svg,yml}\"",
     "prettier:write": "prettier --no-error-on-unmatched-pattern -w \"{src,test,script}/**/*.{json,md,svg,yml}\"",
-    "test": "COMPLIANCE=true forge test --ffi",
+    "test": "forge test --ffi",
     "test:lite": "FOUNDRY_PROFILE=lite forge test",
     "test:optimized": "pnpm run build:optimized && FOUNDRY_PROFILE=test-optimized forge test"
   },

diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/src/SemaphoreMSAValidator.sol b/src/SemaphoreMSAValidator.sol
@@ -68,6 +68,7 @@ contract SemaphoreMSAValidator is ERC7579ValidatorBase {
     ISemaphoreGroups public groups;
     mapping(address account => uint256 groupId) public groupMapping;
     mapping(address account => uint8 threshold) public thresholds;
+    mapping(address account => uint8 count) public memberCount;
 
     // smart account -> hash(call(params)) -> valid proof count
     mapping(address account => mapping(bytes32 txHash => ExtCallCount callDataCount)) public
@@ -133,6 +134,7 @@ contract SemaphoreMSAValidator is ERC7579ValidatorBase {
 
         // Add members to the group
         semaphore.addMembers(groupId, cmts);
+        memberCount[account] = uint8(cmts.length);
 
         emit ModuleInitialized(account);
     }
@@ -143,6 +145,7 @@ contract SemaphoreMSAValidator is ERC7579ValidatorBase {
         delete thresholds[account];
         delete groupMapping[account];
         delete acctSeqNum[account];
+        delete memberCount[account];
 
         //TODO: what is a good way to delete entries associated with `acctTxCount[account]`,
         //   The following line will make the compiler fail.
@@ -151,15 +154,9 @@ contract SemaphoreMSAValidator is ERC7579ValidatorBase {
         emit ModuleUninitialized(account);
     }
 
-    function memberCount(address account) public view returns (uint8 cnt) {
-        // account doesn't belong to a semaphore group. We return 0
-        if (thresholds[account] == 0) return 0;
-        cnt = uint8(groups.getMerkleTreeSize(groupMapping[account]));
-    }
-
     function setThreshold(uint8 newThreshold) external moduleInstalled {
         address account = msg.sender;
-        if (newThreshold == 0 || newThreshold > memberCount(account)) {
+        if (newThreshold == 0 || newThreshold > memberCount[account]) {
             revert InvalidThreshold(account);
         }
 
@@ -171,14 +168,16 @@ contract SemaphoreMSAValidator is ERC7579ValidatorBase {
         address account = msg.sender;
         uint256 groupId = groupMapping[account];
 
-        if (memberCount(account) + cmts.length > MAX_MEMBERS) revert MaxMemberReached(account);
+        if (memberCount[account] + cmts.length > MAX_MEMBERS) revert MaxMemberReached(account);
 
         for (uint256 i = 0; i < cmts.length; ++i) {
             if (cmts[i] == uint256(0)) revert InvalidCommitment(account);
             if (groups.hasMember(groupId, cmts[i])) revert IsMemberAlready(account, cmts[i]);
         }
 
         semaphore.addMembers(groupId, cmts);
+        memberCount[account] += uint8(cmts.length);
+
         emit AddedMembers(account, cmts.length);
     }
 
@@ -191,12 +190,13 @@ contract SemaphoreMSAValidator is ERC7579ValidatorBase {
     {
         address account = msg.sender;
 
-        if (memberCount(account) == thresholds[account]) revert MemberCntReachesThreshold(account);
+        if (memberCount[account] == thresholds[account]) revert MemberCntReachesThreshold(account);
 
         uint256 groupId = groupMapping[account];
         if (!groups.hasMember(groupId, cmt)) revert MemberNotExists(account, cmt);
 
         semaphore.removeMember(groupId, cmt, merkleProofSiblings);
+        memberCount[account] -= 1;
 
         emit RemovedMember(account, cmt);
     }
@@ -347,7 +347,7 @@ contract SemaphoreMSAValidator is ERC7579ValidatorBase {
         uint256 cmt = Identity.getCommitment(pubKey);
         if (!groups.hasMember(groupId, cmt)) revert MemberNotExists(account, cmt);
 
-        // We don't allow call to other contract.
+        // We don't allow call to other contracts.
         address targetAddr = address(bytes20(userOp.callData[100:120]));
         if (targetAddr != address(this)) revert NonValidatorCallBanned(targetAddr, address(this));
 
@@ -356,10 +356,9 @@ contract SemaphoreMSAValidator is ERC7579ValidatorBase {
         bytes memory valAndCallData = userOp.callData[120:];
         bytes4 funcSel = bytes4(LibBytes.slice(valAndCallData, 32, 36));
 
-        // Allow only these few types on function calls to pass, and reject all other on-chain
-        //   calls. They must be executed via `executeTx()` function.
+        // We only allow calls to `initiateTx()`, `signTx()`, and `executeTx()` to pass,
+        //   and reject the rest.
         if (_isAllowedSelector(funcSel)) return VALIDATION_SUCCESS;
-
         revert NonAllowedSelector(account, funcSel);
     }
 

diff --git a/test/SemaphoreMSAValidator.t.sol b/test/SemaphoreMSAValidator.t.sol
@@ -31,6 +31,7 @@ import { SemaphoreMSAValidator, ERC7579ValidatorBase } from "../src/SemaphoreMSA
 import {
     getEmptyUserOperation,
     getEmptySemaphoreProof,
+    getGroupRmMerkleProof,
     getTestUserOpCallData,
     Identity,
     IdentityLib
@@ -236,8 +237,14 @@ contract SemaphoreValidatorUnitTest is RhinestoneModuleKit, Test {
         uint256[] memory newMembers = new uint256[](1);
         newMembers[0] = newCommitment;
 
-        vm.prank(smartAcct.account);
+        // Test: addMembers() is successfully executed
+        vm.startPrank(smartAcct.account);
+        vm.expectEmit(true, true, true, true, address(semaphoreValidator));
+        emit SemaphoreMSAValidator.AddedMembers(smartAcct.account, uint256(1));
         semaphoreValidator.addMembers(newMembers);
+        vm.stopPrank();
+
+        assertEq(semaphoreValidator.memberCount(smartAcct.account), 2);
 
         // Test: the userOp should pass
         uint256 validationData = ERC7579ValidatorBase.ValidationData.unwrap(
@@ -246,8 +253,40 @@ contract SemaphoreValidatorUnitTest is RhinestoneModuleKit, Test {
         assertEq(validationData, VALIDATION_SUCCESS);
     }
 
-    function test_removeMember() public setupSmartAcctWithMembersThreshold(2, 1) {
-        revert("to be implemented");
+    function test_removeMember() public setupSmartAcctWithMembersThreshold(MEMBER_NUM, 1) {
+        uint256[] memory cmts = _getMemberCmts(MEMBER_NUM);
+        User storage rmUser = $users[0];
+        uint256 rmCmt = rmUser.identity.commitment();
+
+        (uint256[] memory merkleProof,) = getGroupRmMerkleProof(cmts, rmCmt);
+
+        // Test: remove member
+        vm.startPrank(smartAcct.account);
+        vm.expectEmit(true, true, true, true, address(semaphoreValidator));
+        emit SemaphoreMSAValidator.RemovedMember(smartAcct.account, rmCmt);
+        semaphoreValidator.removeMember(rmCmt, merkleProof);
+        vm.stopPrank();
+
+        assertEq(semaphoreValidator.memberCount(smartAcct.account), MEMBER_NUM - 1);
+
+        // Compose a UserOp
+        PackedUserOperation memory userOp = getEmptyUserOperation();
+        userOp.sender = smartAcct.account;
+        userOp.callData = getTestUserOpCallData(
+            0,
+            address(semaphoreValidator),
+            abi.encodeWithSelector(SemaphoreMSAValidator.initiateTx.selector)
+        );
+        bytes32 userOpHash = bytes32(keccak256("userOpHash"));
+        userOp.signature = rmUser.identity.signHash(userOpHash);
+
+        // Test: the userOp should fail and revert
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                SemaphoreMSAValidator.MemberNotExists.selector, smartAcct.account, rmCmt
+            )
+        );
+        semaphoreValidator.validateUserOp(userOp, userOpHash);
     }
 
     function _getSemaphoreValidatorUserOpData(

diff --git a/test/utils/TestUtils.sol b/test/utils/TestUtils.sol
@@ -7,6 +7,10 @@ import { ISemaphore } from "../../src/utils/Semaphore.sol";
 // import { console } from "forge-std/console.sol";
 import { LibString } from "solady/Milady.sol";
 
+// https://github.com/foundry-rs/forge-std/blob/master/src/Base.sol#L9
+address constant VM_ADDRESS = 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D;
+Vm constant vm = Vm(VM_ADDRESS);
+
 struct ValidationData {
     address aggregator;
     uint48 validAfter;
@@ -49,13 +53,47 @@ function getTestUserOpCallData(
     callData = bytes.concat(new bytes(100), bytes20(targetAddr), bytes32(value), txCallData);
 }
 
+function getGroupRmMerkleProof(
+    uint256[] memory members,
+    uint256 removal
+)
+    returns (uint256[] memory merkleProof, uint256 root)
+{
+    string[] memory cmd = new string[](5);
+    cmd[0] = "pnpm";
+    cmd[1] = "semaphore-group";
+    cmd[2] = "remove-member";
+    cmd[3] = _join(members);
+    cmd[4] = LibString.toString(removal);
+
+    bytes memory outBytes = vm.ffi(cmd);
+    string memory outStr = string(outBytes);
+    string[] memory retStr = LibString.split(outStr, " ");
+
+    merkleProof = _splitToUint(retStr[0]);
+    root = vm.parseUint(retStr[1]);
+}
+
+function _splitToUint(string memory str) pure returns (uint256[] memory retArr) {
+    string[] memory arr = LibString.split(str, ",");
+    retArr = new uint256[](arr.length);
+    for (uint256 i = 0; i < arr.length; i++) {
+        retArr[i] = vm.parseUint(arr[i]);
+    }
+}
+
+function _join(uint256[] memory members) pure returns (string memory retStr) {
+    for (uint256 i = 0; i < members.length; i++) {
+        retStr = string.concat(retStr, LibString.toString(members[i]));
+        if (i < members.length - 1) {
+            retStr = string.concat(retStr, ",");
+        }
+    }
+}
+
 type Identity is bytes32;
 
 library IdentityLib {
-    // https://github.com/foundry-rs/forge-std/blob/master/src/Base.sol#L9
-    address internal constant VM_ADDRESS = 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D;
-    Vm internal constant vm = Vm(VM_ADDRESS);
-
     function genIdentity(uint256 seed) public view returns (Identity) {
         return Identity.wrap(keccak256(abi.encodePacked(seed, address(this))));
     }