From c138fdbdcfb1d5514226b6ec2c2b0acc22b8a56d Mon Sep 17 00:00:00 2001 From: r4bbit <445106+0x-r4bbit@users.noreply.github.com> Date: Thu, 17 Oct 2024 20:53:41 +0200 Subject: [PATCH] feat: add `TrustedCodehashAccess` contract and interface This enables whitelisting in the staking manager contract using the `TrustedCodehashAccess` interface. These changes were largely taken from #39 with few changes: - Due to different OZ versions, `Ownable()` constructor needs to be called with owner Closes #15 --- .gas-report | 26 ++++----- .gas-snapshot | 64 +++++++++++------------ src/RewardsStreamerMP.sol | 9 ++-- src/TrustedCodehashAccess.sol | 49 +++++++++++++++++ src/interfaces/IStakeManager.sol | 3 +- src/interfaces/ITrustedCodehashAccess.sol | 29 ++++++++++ test/RewardsStreamerMP.t.sol | 6 ++- 7 files changed, 136 insertions(+), 50 deletions(-) create mode 100644 src/TrustedCodehashAccess.sol create mode 100644 src/interfaces/ITrustedCodehashAccess.sol diff --git a/.gas-report b/.gas-report index 3e3f0bb..6a056dc 100644 --- a/.gas-report +++ b/.gas-report @@ -15,22 +15,24 @@ | src/RewardsStreamerMP.sol:RewardsStreamerMP contract | | | | | | |------------------------------------------------------|-----------------|-------|--------|-------|---------| | Deployment Cost | Deployment Size | | | | | -| 1010922 | 4521 | | | | | +| 1216069 | 5542 | | | | | | Function Name | min | avg | median | max | # calls | -| MAX_LOCKUP_PERIOD | 273 | 273 | 273 | 273 | 22 | -| MAX_MULTIPLIER | 274 | 274 | 274 | 274 | 28 | +| MAX_LOCKUP_PERIOD | 228 | 228 | 228 | 228 | 22 | +| MAX_MULTIPLIER | 229 | 229 | 229 | 229 | 28 | | MIN_LOCKUP_PERIOD | 253 | 253 | 253 | 253 | 11 | | MP_RATE_PER_YEAR | 231 | 231 | 231 | 231 | 3 | | SCALE_FACTOR | 229 | 229 | 229 | 229 | 39 | | STAKE_TOKEN | 250 | 250 | 250 | 250 | 128 | -| accountedRewards | 373 | 931 | 373 | 2373 | 68 | -| getAccount | 1574 | 1574 | 1574 | 1574 | 65 | +| accountedRewards | 395 | 953 | 395 | 2395 | 68 | +| getAccount | 1596 | 1596 | 1596 | 1596 | 65 | +| isTrustedCodehash | 541 | 1041 | 541 | 2541 | 128 | | rewardIndex | 351 | 380 | 351 | 2351 | 68 | -| totalMP | 330 | 330 | 330 | 330 | 71 | -| totalMaxMP | 373 | 373 | 373 | 373 | 71 | -| totalStaked | 352 | 352 | 352 | 352 | 71 | -| updateAccountMP | 34632 | 36870 | 37134 | 37134 | 19 | -| updateGlobalState | 29986 | 55567 | 47365 | 80313 | 25 | +| setTrustedCodehash | 47926 | 47926 | 47926 | 47926 | 32 | +| totalMP | 352 | 352 | 352 | 352 | 71 | +| totalMaxMP | 395 | 395 | 395 | 395 | 71 | +| totalStaked | 374 | 374 | 374 | 374 | 71 | +| updateAccountMP | 34654 | 36892 | 37156 | 37156 | 19 | +| updateGlobalState | 30008 | 55588 | 47387 | 80334 | 25 | | src/StakeVault.sol:StakeVault contract | | | | | | @@ -38,8 +40,8 @@ | Deployment Cost | Deployment Size | | | | | | 857099 | 4070 | | | | | | Function Name | min | avg | median | max | # calls | -| stake | 194660 | 231720 | 238353 | 258837 | 46 | -| unstake | 81452 | 109772 | 99015 | 141680 | 13 | +| stake | 196978 | 234038 | 240671 | 261155 | 46 | +| unstake | 83287 | 111971 | 101308 | 144002 | 13 | | src/XPNFTToken.sol:XPNFTToken contract | | | | | | diff --git a/.gas-snapshot b/.gas-snapshot index 0521fe7..6dae3d7 100644 --- a/.gas-snapshot +++ b/.gas-snapshot @@ -1,4 +1,4 @@ -IntegrationTest:testStakeFoo() (gas: 1459151) +IntegrationTest:testStakeFoo() (gas: 1471979) NFTMetadataGeneratorSVGTest:testGenerateMetadata() (gas: 92874) NFTMetadataGeneratorSVGTest:testSetImageStrings() (gas: 60081) NFTMetadataGeneratorSVGTest:testSetImageStringsRevert() (gas: 35818) @@ -6,37 +6,37 @@ NFTMetadataGeneratorURLTest:testGenerateMetadata() (gas: 109345) NFTMetadataGeneratorURLTest:testSetBaseURL() (gas: 50653) NFTMetadataGeneratorURLTest:testSetBaseURLRevert() (gas: 35993) RewardsStreamerTest:testStake() (gas: 869874) -StakeTest:test_StakeMultipleAccounts() (gas: 484173) -StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 629575) -StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 796295) -StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 490078) -StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 511480) -StakeTest:test_StakeOneAccount() (gas: 279745) -StakeTest:test_StakeOneAccountAndRewards() (gas: 425143) -StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 485909) -StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 480997) -StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 293513) -StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 293524) -StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 293635) -UnstakeTest:test_StakeMultipleAccounts() (gas: 484195) -UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 629552) -UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 796317) -UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 490055) -UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 511502) -UnstakeTest:test_StakeOneAccount() (gas: 279768) -UnstakeTest:test_StakeOneAccountAndRewards() (gas: 425165) -UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 485931) -UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 480977) -UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 293558) -UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 293524) -UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 293613) -UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 495152) -UnstakeTest:test_UnstakeMultipleAccounts() (gas: 672552) -UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 991588) -UnstakeTest:test_UnstakeOneAccount() (gas: 468592) -UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 483480) -UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 574902) -UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 505088) +StakeTest:test_StakeMultipleAccounts() (gas: 488941) +StakeTest:test_StakeMultipleAccountsAndRewards() (gas: 634452) +StakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 801369) +StakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 494622) +StakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 515979) +StakeTest:test_StakeOneAccount() (gas: 282173) +StakeTest:test_StakeOneAccountAndRewards() (gas: 427680) +StakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 488578) +StakeTest:test_StakeOneAccountReachingMPLimit() (gas: 483621) +StakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 295784) +StakeTest:test_StakeOneAccountWithMinLockUp() (gas: 295840) +StakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 295951) +UnstakeTest:test_StakeMultipleAccounts() (gas: 488963) +UnstakeTest:test_StakeMultipleAccountsAndRewards() (gas: 634429) +UnstakeTest:test_StakeMultipleAccountsMPIncreasesMaxMPDoesNotChange() (gas: 801391) +UnstakeTest:test_StakeMultipleAccountsWithMinLockUp() (gas: 494599) +UnstakeTest:test_StakeMultipleAccountsWithRandomLockUp() (gas: 516001) +UnstakeTest:test_StakeOneAccount() (gas: 282196) +UnstakeTest:test_StakeOneAccountAndRewards() (gas: 427702) +UnstakeTest:test_StakeOneAccountMPIncreasesMaxMPDoesNotChange() (gas: 488600) +UnstakeTest:test_StakeOneAccountReachingMPLimit() (gas: 483601) +UnstakeTest:test_StakeOneAccountWithMaxLockUp() (gas: 295829) +UnstakeTest:test_StakeOneAccountWithMinLockUp() (gas: 295840) +UnstakeTest:test_StakeOneAccountWithRandomLockUp() (gas: 295929) +UnstakeTest:test_UnstakeBonusMPAndAccuredMP() (gas: 499981) +UnstakeTest:test_UnstakeMultipleAccounts() (gas: 681122) +UnstakeTest:test_UnstakeMultipleAccountsAndRewards() (gas: 1002814) +UnstakeTest:test_UnstakeOneAccount() (gas: 474888) +UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 488421) +UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 579871) +UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 509737) XPNFTTokenTest:testApproveNotAllowed() (gas: 10507) XPNFTTokenTest:testGetApproved() (gas: 10531) XPNFTTokenTest:testIsApprovedForAll() (gas: 10705) diff --git a/src/RewardsStreamerMP.sol b/src/RewardsStreamerMP.sol index 3f30fd7..c52f443 100644 --- a/src/RewardsStreamerMP.sol +++ b/src/RewardsStreamerMP.sol @@ -4,9 +4,10 @@ pragma solidity ^0.8.26; import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; import { ReentrancyGuard } from "@openzeppelin/contracts/utils/ReentrancyGuard.sol"; import { IStakeManager } from "./interfaces/IStakeManager.sol"; +import { TrustedCodehashAccess } from "./TrustedCodehashAccess.sol"; // Rewards Streamer with Multiplier Points -contract RewardsStreamerMP is ReentrancyGuard, IStakeManager { +contract RewardsStreamerMP is IStakeManager, TrustedCodehashAccess, ReentrancyGuard { error StakingManager__AmountCannotBeZero(); error StakingManager__TransferFailed(); error StakingManager__InsufficientBalance(); @@ -42,13 +43,13 @@ contract RewardsStreamerMP is ReentrancyGuard, IStakeManager { mapping(address account => Account data) public accounts; - constructor(address _stakingToken, address _rewardToken) { + constructor(address _owner, address _stakingToken, address _rewardToken) TrustedCodehashAccess(_owner) { STAKE_TOKEN = IERC20(_stakingToken); REWARD_TOKEN = IERC20(_rewardToken); lastMPUpdatedTime = block.timestamp; } - function stake(uint256 amount, uint256 lockPeriod) external nonReentrant { + function stake(uint256 amount, uint256 lockPeriod) external onlyTrustedCodehash nonReentrant { if (amount == 0) { revert StakingManager__AmountCannotBeZero(); } @@ -98,7 +99,7 @@ contract RewardsStreamerMP is ReentrancyGuard, IStakeManager { account.lastMPUpdateTime = block.timestamp; } - function unstake(uint256 amount) external nonReentrant { + function unstake(uint256 amount) external onlyTrustedCodehash nonReentrant { Account storage account = accounts[msg.sender]; if (amount > account.stakedBalance) { revert StakingManager__InsufficientBalance(); diff --git a/src/TrustedCodehashAccess.sol b/src/TrustedCodehashAccess.sol new file mode 100644 index 0000000..d1579a4 --- /dev/null +++ b/src/TrustedCodehashAccess.sol @@ -0,0 +1,49 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.26; + +import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol"; +import { ITrustedCodehashAccess } from "./interfaces/ITrustedCodehashAccess.sol"; +/** + * @title TrustedCodehashAccess + * @author Ricardo Guilherme Schmidt + * @notice Ensures that only specific contract bytecode hashes are trusted to + * interact with the functions using the `onlyTrustedCodehash` modifier. + */ + +contract TrustedCodehashAccess is ITrustedCodehashAccess, Ownable { + mapping(bytes32 codehash => bool permission) private trustedCodehashes; + + /** + * @notice Restricts access based on the codehash of the caller. + * Only contracts with trusted codehashes can execute functions using this modifier. + */ + modifier onlyTrustedCodehash() { + bytes32 codehash = msg.sender.codehash; + if (!trustedCodehashes[codehash]) { + revert TrustedCodehashAccess__UnauthorizedCodehash(); + } + _; + } + + constructor(address _owner) Ownable(_owner) { } + + /** + * @notice Allows the owner to set or update the trust status for a contract's codehash. + * @dev Emits the `TrustedCodehashUpdated` event whenever a codehash is updated. + * @param _codehash The bytecode hash of the contract. + * @param _trusted Boolean flag to designate the contract as trusted or not. + */ + function setTrustedCodehash(bytes32 _codehash, bool _trusted) external onlyOwner { + trustedCodehashes[_codehash] = _trusted; + emit TrustedCodehashUpdated(_codehash, _trusted); + } + + /** + * @notice Checks if a contract's codehash is trusted to interact with protected functions. + * @param _codehash The bytecode hash of the contract. + * @return bool True if the codehash is trusted, false otherwise. + */ + function isTrustedCodehash(bytes32 _codehash) external view returns (bool) { + return trustedCodehashes[_codehash]; + } +} diff --git a/src/interfaces/IStakeManager.sol b/src/interfaces/IStakeManager.sol index 2104e5d..c213e9c 100644 --- a/src/interfaces/IStakeManager.sol +++ b/src/interfaces/IStakeManager.sol @@ -2,8 +2,9 @@ pragma solidity ^0.8.26; import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; +import { ITrustedCodehashAccess } from "./ITrustedCodehashAccess.sol"; -interface IStakeManager { +interface IStakeManager is ITrustedCodehashAccess { error StakeManager__FundsLocked(); error StakeManager__InvalidLockTime(); error StakeManager__InsufficientFunds(); diff --git a/src/interfaces/ITrustedCodehashAccess.sol b/src/interfaces/ITrustedCodehashAccess.sol new file mode 100644 index 0000000..303b002 --- /dev/null +++ b/src/interfaces/ITrustedCodehashAccess.sol @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.26; + +/** + * @title TrustedCodehashAccess + * @author Ricardo Guilherme Schmidt + * @notice Ensures that only specific contract bytecode hashes are trusted to + * interact with the functions using the `onlyTrustedCodehash` modifier. + */ +interface ITrustedCodehashAccess { + error TrustedCodehashAccess__UnauthorizedCodehash(); + + event TrustedCodehashUpdated(bytes32 indexed codehash, bool trusted); + + /** + * @notice Allows the owner to set or update the trust status for a contract's codehash. + * @dev Emits the `TrustedCodehashUpdated` event whenever a codehash is updated. + * @param _codehash The bytecode hash of the contract. + * @param _trusted Boolean flag to designate the contract as trusted or not. + */ + function setTrustedCodehash(bytes32 _codehash, bool _trusted) external; + + /** + * @notice Checks if a contract's codehash is trusted to interact with protected functions. + * @param _codehash The bytecode hash of the contract. + * @return bool True if the codehash is trusted, false otherwise. + */ + function isTrustedCodehash(bytes32 _codehash) external view returns (bool); +} diff --git a/test/RewardsStreamerMP.t.sol b/test/RewardsStreamerMP.t.sol index 0f90945..2fe79df 100644 --- a/test/RewardsStreamerMP.t.sol +++ b/test/RewardsStreamerMP.t.sol @@ -22,7 +22,7 @@ contract RewardsStreamerMPTest is Test { function setUp() public virtual { rewardToken = new MockToken("Reward Token", "RT"); stakingToken = new MockToken("Staking Token", "ST"); - streamer = new RewardsStreamerMP(address(stakingToken), address(rewardToken)); + streamer = new RewardsStreamerMP(address(this), address(stakingToken), address(rewardToken)); address[4] memory accounts = [alice, bob, charlie, dave]; for (uint256 i = 0; i < accounts.length; i++) { @@ -85,6 +85,10 @@ contract RewardsStreamerMPTest is Test { function _createTestVault(address owner) internal returns (StakeVault vault) { vm.prank(owner); vault = new StakeVault(owner, streamer); + + if (!streamer.isTrustedCodehash(address(vault).codehash)) { + streamer.setTrustedCodehash(address(vault).codehash, true); + } } function _stake(address account, uint256 amount, uint256 lockupTime) public {