From 9c7eb60a2d4b4415fac694c42fa666a38a14fb42 Mon Sep 17 00:00:00 2001 From: rayeaster Date: Tue, 19 Dec 2023 13:47:36 +0800 Subject: [PATCH] fix governor related issue reported by C4 --- .../contracts/Dependencies/RolesAuthority.sol | 6 +- packages/contracts/contracts/Governor.sol | 30 +++++-- .../foundry_test/Governor.fundamentals.t.sol | 90 +++++++++++++++++++ 3 files changed, 120 insertions(+), 6 deletions(-) diff --git a/packages/contracts/contracts/Dependencies/RolesAuthority.sol b/packages/contracts/contracts/Dependencies/RolesAuthority.sol index 14fbbf382..6ddf2a6ca 100644 --- a/packages/contracts/contracts/Dependencies/RolesAuthority.sol +++ b/packages/contracts/contracts/Dependencies/RolesAuthority.sol @@ -118,7 +118,11 @@ contract RolesAuthority is IRolesAuthority, Auth, Authority { } } else { getRolesWithCapability[target][functionSig] &= ~bytes32(1 << role); - enabledFunctionSigsByTarget[target].remove(bytes32(functionSig)); + + // If no role exist for this target & functionSig, mark it as disabled + if (getRolesWithCapability[target][functionSig] == bytes32(0)) { + enabledFunctionSigsByTarget[target].remove(bytes32(functionSig)); + } // If no enabled function signatures exist for this target, remove target if (enabledFunctionSigsByTarget[target].length() == 0) { diff --git a/packages/contracts/contracts/Governor.sol b/packages/contracts/contracts/Governor.sol index 019595251..68593b87a 100644 --- a/packages/contracts/contracts/Governor.sol +++ b/packages/contracts/contracts/Governor.sol @@ -73,19 +73,29 @@ contract Governor is RolesAuthority { function getRolesForUser(address user) external view returns (uint8[] memory rolesForUser) { // Enumerate over all possible roles and check if enabled uint256 count; - for (uint8 i = 0; i < type(uint8).max; i++) { + for (uint8 i = 0; i <= type(uint8).max; ) { if (doesUserHaveRole(user, i)) { count += 1; } + if (i < type(uint8).max) { + i = i + 1; + } else { + break; + } } if (count > 0) { uint256 j = 0; rolesForUser = new uint8[](count); - for (uint8 i = 0; i < type(uint8).max; i++) { + for (uint8 i = 0; i <= type(uint8).max; ) { if (doesUserHaveRole(user, i)) { rolesForUser[j] = i; j++; } + if (i < type(uint8).max) { + i = i + 1; + } else { + break; + } } } } @@ -95,21 +105,31 @@ contract Governor is RolesAuthority { /// @return roleIds An array of role IDs extracted from the byte map. function getRolesFromByteMap(bytes32 byteMap) public pure returns (uint8[] memory roleIds) { uint256 count; - for (uint8 i = 0; i < type(uint8).max; i++) { + for (uint8 i = 0; i <= type(uint8).max; ) { bool roleEnabled = (uint256(byteMap >> i) & 1) != 0; if (roleEnabled) { count += 1; } + if (i < type(uint8).max) { + i = i + 1; + } else { + break; + } } if (count > 0) { uint256 j = 0; roleIds = new uint8[](count); - for (uint8 i = 0; i < type(uint8).max; i++) { + for (uint8 i = 0; i <= type(uint8).max; ) { bool roleEnabled = (uint256(byteMap >> i) & 1) != 0; if (roleEnabled) { roleIds[j] = i; j++; } + if (i < type(uint8).max) { + i = i + 1; + } else { + break; + } } } } @@ -119,7 +139,7 @@ contract Governor is RolesAuthority { /// @return A bytes32 value encoding the roles. function getByteMapFromRoles(uint8[] memory roleIds) public pure returns (bytes32) { bytes32 _data; - for (uint8 i = 0; i < roleIds.length; i++) { + for (uint256 i = 0; i < roleIds.length; i++) { _data |= bytes32(1 << uint256(roleIds[i])); } return _data; diff --git a/packages/contracts/foundry_test/Governor.fundamentals.t.sol b/packages/contracts/foundry_test/Governor.fundamentals.t.sol index 5c511cc1b..72930b19f 100644 --- a/packages/contracts/foundry_test/Governor.fundamentals.t.sol +++ b/packages/contracts/foundry_test/Governor.fundamentals.t.sol @@ -26,6 +26,60 @@ contract GovernorFundamentalsTest is eBTCBaseFixture { _testRoleWithCapabilitiesCanCallGovernorFunctions(testUser, testRole); } + // Test: The address with the granted roles should be able to retrieve all roles + function test_RolesRetrievalWithMaxAllowed() public { + address testUser = address(0x1); + + vm.startPrank(defaultGovernance); + for (uint8 i = 0; i <= type(uint8).max; ) { + authority.setUserRole(testUser, i, true); + authority.setRoleName(i, "TestRole"); + if (i < type(uint8).max) { + i = i + 1; + } else { + break; + } + } + vm.stopPrank(); + + uint8[] memory rolesForUser = authority.getRolesForUser(testUser); + for (uint8 i = 0; i <= type(uint8).max; ) { + assertEq(rolesForUser[i], i, "!retrieved role msimatch"); + if (i < type(uint8).max) { + i = i + 1; + } else { + break; + } + } + } + + // Test: The address with the granted roles should be able to retrieve all roles via bitmap + function test_RolesRetrievalWithMaxAllowedViaMap() public { + uint8[] memory rolesForMapGiven = new uint8[](256); + vm.startPrank(defaultGovernance); + for (uint8 i = 0; i <= type(uint8).max; ) { + authority.setRoleName(i, "TestRole"); + rolesForMapGiven[i] = i; + if (i < type(uint8).max) { + i = i + 1; + } else { + break; + } + } + vm.stopPrank(); + + bytes32 _roleMap = authority.getByteMapFromRoles(rolesForMapGiven); + uint8[] memory rolesForMap = authority.getRolesFromByteMap(_roleMap); + for (uint8 i = 0; i <= type(uint8).max; ) { + assertEq(rolesForMap[i], i, "!retrieved role msimatch"); + if (i < type(uint8).max) { + i = i + 1; + } else { + break; + } + } + } + // Setup: Assign capabilities for all functions to the contract owner and then renounce ownership // Test: The previous owner should still be able to call all Governor functions because of the granted capabilities function test_OwnerCanCallAllFunctionsAfterGainingProperCapabilitiesAndRenouncingOwnership( @@ -114,6 +168,42 @@ contract GovernorFundamentalsTest is eBTCBaseFixture { public {} + function test_GetEnabledFunctionsInTargetReturnsExpectedAfterRoleModification() public { + // set all capabilities to two roles + _grantAllGovernorCapabilitiesToRole(1); + _grantAllGovernorCapabilitiesToRole(2); + + // retrieve all funcs enabled + bytes4[] memory funcSigsOriginal = authority.getEnabledFunctionsInTarget(address(authority)); + bytes4 _firstFunc = bytes4(keccak256("setRoleName(uint8,string)")); + bytes4 _lastFunc = bytes4(keccak256("burnCapability(address,bytes4)")); + assertEq(funcSigsOriginal[0], _firstFunc, "!mismatch first enabled function sig"); + + // now revoke capability from one role + vm.startPrank(defaultGovernance); + authority.setRoleCapability(1, address(authority), _firstFunc, false); + vm.stopPrank(); + + // check again the function, should still enabled since there is a second role 2 could call + bytes4[] memory funcSigsAgain = authority.getEnabledFunctionsInTarget(address(authority)); + assertEq(funcSigsAgain[0], _firstFunc, "!!mismatch first enabled function sig"); + assertEq(funcSigsAgain.length, funcSigsOriginal.length, "!!mismatch enabled function count"); + + // now revoke capability from the second role + vm.startPrank(defaultGovernance); + authority.setRoleCapability(2, address(authority), _firstFunc, false); + vm.stopPrank(); + + // check again the function, should disabled since there is no role could call + bytes4[] memory funcSigsLast = authority.getEnabledFunctionsInTarget(address(authority)); + assertEq(funcSigsLast[0], _lastFunc, "!!!mismatch first enabled function sig"); + assertEq( + funcSigsLast.length, + funcSigsOriginal.length - 1, + "!!!mismatch enabled function count" + ); + } + /// @dev Helper function to grant all Governor setter capabilities to a specific role /// @dev Assumes default governance still has ownerships function _grantAllGovernorCapabilitiesToRole(uint8 role) internal {