diff --git a/packages/contracts-bedrock/interfaces/periphery/staking/IPolicyEngineStaking.sol b/packages/contracts-bedrock/interfaces/periphery/staking/IPolicyEngineStaking.sol index 288af54c3b9..cdcf3dcf250 100644 --- a/packages/contracts-bedrock/interfaces/periphery/staking/IPolicyEngineStaking.sol +++ b/packages/contracts-bedrock/interfaces/periphery/staking/IPolicyEngineStaking.sol @@ -34,8 +34,14 @@ interface IPolicyEngineStaking is ISemver { /// @notice Emitted when ownership is transferred. event OwnershipTransferred(address indexed previousOwner, address indexed newOwner); - /// @notice Thrown when the caller is not the owner. - error PolicyEngineStaking_OnlyOwner(); + /// @notice Emitted when a new owner is nominated via `transferOwnership`. + event OwnershipTransferStarted(address indexed previousOwner, address indexed newOwner); + + /// @notice Thrown when the caller is not authorized (OZ Ownable). + error OwnableUnauthorizedAccount(address account); + + /// @notice Thrown when an invalid owner is provided (OZ Ownable). + error OwnableInvalidOwner(address owner); /// @notice Thrown when the staking is paused. error PolicyEngineStaking_Paused(); @@ -69,10 +75,24 @@ interface IPolicyEngineStaking is ISemver { /// @notice Returns the contract owner. function owner() external view returns (address); - /// @notice Transfers ownership of the contract to a new account. Only callable by owner. - /// @param _newOwner The address of the new owner. + /// @notice Returns the pending owner nominated via `transferOwnership`. + function pendingOwner() external view returns (address); + + /// @notice Starts the ownership transfer of the contract to a new account. Replaces the + /// pending transfer if there is one. The nominated address must call + /// `acceptOwnership` to finalize the transfer (two-step pattern). Only callable + /// by owner. + /// + /// @param _newOwner The address of the nominated owner. function transferOwnership(address _newOwner) external; + /// @notice Accepts ownership after being nominated via `transferOwnership`. + /// Only callable by the pending owner. + function acceptOwnership() external; + + /// @notice Renounces ownership of the contract. Only callable by owner. + function renounceOwnership() external; + /// @notice Returns whether the contract is paused. function paused() external view returns (bool); @@ -114,6 +134,7 @@ interface IPolicyEngineStaking is ISemver { /// @notice Sets whether a staker can set the caller as beneficiary. When disallowing, /// if the staker's current beneficiary is the caller, their stake attribution is /// moved back to the staker (beneficiary reset to self). + /// This function is callable even when the contract is paused. /// /// @param _staker The staker address. /// @param _allowed Whether the staker is allowed to set the caller as beneficiary. diff --git a/packages/contracts-bedrock/interfaces/universal/IOwnable.sol b/packages/contracts-bedrock/interfaces/universal/IOwnable.sol index b6f48de59b2..78a2a0632a3 100644 --- a/packages/contracts-bedrock/interfaces/universal/IOwnable.sol +++ b/packages/contracts-bedrock/interfaces/universal/IOwnable.sol @@ -6,6 +6,12 @@ pragma solidity ^0.8.0; interface IOwnable { event OwnershipTransferred(address indexed previousOwner, address indexed newOwner); + /// @dev The caller account is not authorized to perform an operation. + error OwnableUnauthorizedAccount(address account); + + /// @dev The owner is not a valid owner account. (eg. `address(0)`) + error OwnableInvalidOwner(address owner); + function owner() external view returns (address); function renounceOwnership() external; function transferOwnership(address newOwner) external; // nosemgrep diff --git a/packages/contracts-bedrock/scripts/checks/interfaces/main.go b/packages/contracts-bedrock/scripts/checks/interfaces/main.go index b2ed175cc15..3c87ab8df7d 100644 --- a/packages/contracts-bedrock/scripts/checks/interfaces/main.go +++ b/packages/contracts-bedrock/scripts/checks/interfaces/main.go @@ -33,6 +33,9 @@ var excludeContracts = []string{ // Constructor inheritance differences "IL2ProxyAdmin", + // OZ v4/v5 Ownable mismatch: IOwnable has v5 errors, AddressManager uses v4 Ownable + "IAddressManager", + // TODO: Interfaces that need to be fixed "IInitializable", "IOptimismMintableERC20", "ILegacyMintableERC20", "KontrolCheatsBase", "IResolvedDelegateProxy", @@ -41,7 +44,7 @@ var excludeContracts = []string{ // excludeSourceContracts is a list of contracts that are allowed to not have interfaces var excludeSourceContracts = []string{ // Base contracts with no external functions - "CrossDomainOwnable", "CrossDomainOwnable2", "CrossDomainOwnable3", "CrossDomainMessengerLegacySpacer0", "CrossDomainMessengerLegacySpacer1", + "CrossDomainOwnable", "CrossDomainOwnable2", "CrossDomainOwnable3", "CrossDomainMessengerLegacySpacer0", "CrossDomainMessengerLegacySpacer1", "PolicyEngineStakingMapping", // Helper contracts "SafeSend", "EventLogger", "StorageSetter", "DisputeMonitorHelper", "GameHelper", diff --git a/packages/contracts-bedrock/snapshots/abi/PolicyEngineStaking.json b/packages/contracts-bedrock/snapshots/abi/PolicyEngineStaking.json index 826dad96564..392f7250fc7 100644 --- a/packages/contracts-bedrock/snapshots/abi/PolicyEngineStaking.json +++ b/packages/contracts-bedrock/snapshots/abi/PolicyEngineStaking.json @@ -28,6 +28,13 @@ "stateMutability": "view", "type": "function" }, + { + "inputs": [], + "name": "acceptOwnership", + "outputs": [], + "stateMutability": "nonpayable", + "type": "function" + }, { "inputs": [ { @@ -122,6 +129,26 @@ "stateMutability": "view", "type": "function" }, + { + "inputs": [], + "name": "pendingOwner", + "outputs": [ + { + "internalType": "address", + "name": "", + "type": "address" + } + ], + "stateMutability": "view", + "type": "function" + }, + { + "inputs": [], + "name": "renounceOwnership", + "outputs": [], + "stateMutability": "nonpayable", + "type": "function" + }, { "inputs": [ { @@ -217,7 +244,7 @@ "inputs": [ { "internalType": "address", - "name": "_newOwner", + "name": "newOwner", "type": "address" } ], @@ -341,6 +368,25 @@ "name": "EffectiveStakeChanged", "type": "event" }, + { + "anonymous": false, + "inputs": [ + { + "indexed": true, + "internalType": "address", + "name": "previousOwner", + "type": "address" + }, + { + "indexed": true, + "internalType": "address", + "name": "newOwner", + "type": "address" + } + ], + "name": "OwnershipTransferStarted", + "type": "event" + }, { "anonymous": false, "inputs": [ @@ -411,23 +457,40 @@ "type": "event" }, { - "inputs": [], - "name": "PolicyEngineStaking_InsufficientStake", + "inputs": [ + { + "internalType": "address", + "name": "owner", + "type": "address" + } + ], + "name": "OwnableInvalidOwner", + "type": "error" + }, + { + "inputs": [ + { + "internalType": "address", + "name": "account", + "type": "address" + } + ], + "name": "OwnableUnauthorizedAccount", "type": "error" }, { "inputs": [], - "name": "PolicyEngineStaking_NoStake", + "name": "PolicyEngineStaking_InsufficientStake", "type": "error" }, { "inputs": [], - "name": "PolicyEngineStaking_NotAllowedToSetBeneficiary", + "name": "PolicyEngineStaking_NoStake", "type": "error" }, { "inputs": [], - "name": "PolicyEngineStaking_OnlyOwner", + "name": "PolicyEngineStaking_NotAllowedToSetBeneficiary", "type": "error" }, { diff --git a/packages/contracts-bedrock/snapshots/abi/PolicyEngineStakingMapping.json b/packages/contracts-bedrock/snapshots/abi/PolicyEngineStakingMapping.json new file mode 100644 index 00000000000..7c0eeab87e9 --- /dev/null +++ b/packages/contracts-bedrock/snapshots/abi/PolicyEngineStakingMapping.json @@ -0,0 +1,39 @@ +[ + { + "inputs": [], + "name": "PE_DATA_SLOT", + "outputs": [ + { + "internalType": "bytes32", + "name": "", + "type": "bytes32" + } + ], + "stateMutability": "view", + "type": "function" + }, + { + "inputs": [ + { + "internalType": "address", + "name": "account", + "type": "address" + } + ], + "name": "peData", + "outputs": [ + { + "internalType": "uint128", + "name": "effectiveStake", + "type": "uint128" + }, + { + "internalType": "uint128", + "name": "lastUpdate", + "type": "uint128" + } + ], + "stateMutability": "view", + "type": "function" + } +] \ No newline at end of file diff --git a/packages/contracts-bedrock/snapshots/semver-lock.json b/packages/contracts-bedrock/snapshots/semver-lock.json index c3dd415e65d..255854101f0 100644 --- a/packages/contracts-bedrock/snapshots/semver-lock.json +++ b/packages/contracts-bedrock/snapshots/semver-lock.json @@ -240,8 +240,8 @@ "sourceCodeHash": "0x955bd0c9b47e43219865e4e92abf28d916c96de20cbdf2f94c8ab14d02083759" }, "src/periphery/staking/PolicyEngineStaking.sol:PolicyEngineStaking": { - "initCodeHash": "0xc0c04b0dddbf7831bf5abfccc6a569d92f9b7ab0ec53278f6d1cf7113041d59d", - "sourceCodeHash": "0x998ddc9f24e3c85b1d588c263838261442edaad1dde5424fd55c2d4e1243592a" + "initCodeHash": "0x0dd679c6e955a4bf4af7d98da8c42ac7fd48c5c28af4c61eaeddb782690d6223", + "sourceCodeHash": "0x24fe77ae2b0dacad1f61f6dd65026ce3a34df7c4d0331eadc661b2e062d6bb54" }, "src/safe/DeputyPauseModule.sol:DeputyPauseModule": { "initCodeHash": "0x18422b48c4901ed6fd9338d76d3c5aecfff9a7add34b05c6e21c23d0011ed6bf", diff --git a/packages/contracts-bedrock/snapshots/storageLayout/PolicyEngineStaking.json b/packages/contracts-bedrock/snapshots/storageLayout/PolicyEngineStaking.json index 186637cbf9d..639009332cb 100644 --- a/packages/contracts-bedrock/snapshots/storageLayout/PolicyEngineStaking.json +++ b/packages/contracts-bedrock/snapshots/storageLayout/PolicyEngineStaking.json @@ -4,34 +4,41 @@ "label": "peData", "offset": 0, "slot": "0", - "type": "mapping(address => struct PolicyEngineStaking.PEData)" + "type": "mapping(address => struct PolicyEngineStakingMapping.PEData)" + }, + { + "bytes": "20", + "label": "_owner", + "offset": 0, + "slot": "1", + "type": "address" + }, + { + "bytes": "20", + "label": "_pendingOwner", + "offset": 0, + "slot": "2", + "type": "address" }, { "bytes": "32", "label": "allowlist", "offset": 0, - "slot": "1", + "slot": "3", "type": "mapping(address => mapping(address => bool))" }, { "bytes": "32", "label": "stakingData", "offset": 0, - "slot": "2", + "slot": "4", "type": "mapping(address => struct PolicyEngineStaking.StakedData)" }, { "bytes": "1", "label": "paused", "offset": 0, - "slot": "3", + "slot": "5", "type": "bool" - }, - { - "bytes": "20", - "label": "_owner", - "offset": 1, - "slot": "3", - "type": "address" } ] \ No newline at end of file diff --git a/packages/contracts-bedrock/snapshots/storageLayout/PolicyEngineStakingMapping.json b/packages/contracts-bedrock/snapshots/storageLayout/PolicyEngineStakingMapping.json new file mode 100644 index 00000000000..75ab9c8c1ba --- /dev/null +++ b/packages/contracts-bedrock/snapshots/storageLayout/PolicyEngineStakingMapping.json @@ -0,0 +1,9 @@ +[ + { + "bytes": "32", + "label": "peData", + "offset": 0, + "slot": "0", + "type": "mapping(address => struct PolicyEngineStakingMapping.PEData)" + } +] \ No newline at end of file diff --git a/packages/contracts-bedrock/src/periphery/staking/PolicyEngineStaking.sol b/packages/contracts-bedrock/src/periphery/staking/PolicyEngineStaking.sol index 890a77d2898..a5490d608c0 100644 --- a/packages/contracts-bedrock/src/periphery/staking/PolicyEngineStaking.sol +++ b/packages/contracts-bedrock/src/periphery/staking/PolicyEngineStaking.sol @@ -8,11 +8,36 @@ import { ISemver } from "interfaces/universal/ISemver.sol"; // Libraries import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; +// Inheritance +import { Ownable } from "@openzeppelin/contracts-v5/access/Ownable.sol"; +import { Ownable2Step } from "@openzeppelin/contracts-v5/access/Ownable2Step.sol"; + +/// @title PolicyEngineStakingMapping +/// @notice Holds the `peData` mapping at storage slot 0 so that `op-rbuilder` can read +/// effective-stake data at a known, stable location. Inherited first by +/// `PolicyEngineStaking` to guarantee the slot assignment. +contract PolicyEngineStakingMapping { + /// @notice Policy Engine data per account. Packed in one slot for PE reads. + /// @custom:field effectiveStake The exact stake amount used for ordering. + /// @custom:field lastUpdate The timestamp of the latest change on their effective stake. + struct PEData { + uint128 effectiveStake; + uint128 lastUpdate; + } + + /// @notice Base storage slot for PE data mapping. Policy Engine reads from + /// keccak256(abi.encode(account, PE_DATA_SLOT)). + bytes32 public constant PE_DATA_SLOT = 0; + + /// @notice Slot 0: PE data mapping. + mapping(address account => PEData) public peData; +} + /// @title PolicyEngineStaking /// @notice Periphery contract for stake-based transaction ordering in op-rbuilder. Users stake governance tokens /// and optionally link to a beneficiary who receives ordering power. Supports partial unstake. /// Invariant: every staked token has a beneficiary (self or linked). No receivedStake tracking or unlink(). -contract PolicyEngineStaking is ISemver { +contract PolicyEngineStaking is PolicyEngineStakingMapping, Ownable2Step, ISemver { using SafeERC20 for IERC20; /// @notice Staking stakingData per account. @@ -23,29 +48,14 @@ contract PolicyEngineStaking is ISemver { address beneficiary; } - /// @notice Policy Engine stakingData per account. Packed in one slot for PE reads. - /// @custom:field effectiveStake The exact stake amount used for ordering. - /// @custom:field lastUpdate The timestamp of the latest change on their effective stake. - struct PEData { - uint128 effectiveStake; - uint128 lastUpdate; - } - /// @notice Semantic version. /// @custom:semver 1.0.0 string public constant version = "1.0.0"; - /// @notice Base storage slot for PE stakingData mapping. Policy Engine reads from - /// keccak256(abi.encode(account, PE_DATA_SLOT)). - bytes32 public constant PE_DATA_SLOT = 0; - /// @notice The ERC20 token used for staking. // nosemgrep: sol-safety-no-immutable-variables IERC20 internal immutable STAKING_TOKEN; - /// @notice Slot 0: PE stakingData mapping. - mapping(address account => PEData) public peData; - /// @notice Allowlist: beneficiary => staker => allowed. mapping(address beneficiary => mapping(address staker => bool allowed)) public allowlist; @@ -55,9 +65,6 @@ contract PolicyEngineStaking is ISemver { /// @notice Paused state. bool public paused; - /// @notice The owner of the contract. Can pause, unpause, and transfer ownership. - address private _owner; - /// @notice Emitted when a user stakes OP tokens. /// @param account The address that staked tokens. /// @param amount The amount of tokens staked. @@ -95,14 +102,6 @@ contract PolicyEngineStaking is ISemver { /// @notice Emitted when the staking is unpaused. event Unpaused(); - /// @notice Emitted when ownership is transferred. - /// @param previousOwner The address of the previous owner. - /// @param newOwner The address of the new owner. - event OwnershipTransferred(address indexed previousOwner, address indexed newOwner); - - /// @notice Thrown when the caller is not the owner. - error PolicyEngineStaking_OnlyOwner(); - /// @notice Thrown when the staking is paused. error PolicyEngineStaking_Paused(); @@ -133,10 +132,8 @@ contract PolicyEngineStaking is ISemver { /// @notice Constructs the PolicyEngineStaking contract. /// @param _ownerAddr The address that can pause and unpause staking. /// @param _token The ERC20 token used for staking. - constructor(address _ownerAddr, address _token) { - if (_ownerAddr == address(0)) revert PolicyEngineStaking_ZeroAddress(); + constructor(address _ownerAddr, address _token) Ownable(_ownerAddr) { if (_token == address(0)) revert PolicyEngineStaking_ZeroAddress(); - _owner = _ownerAddr; STAKING_TOKEN = IERC20(_token); } @@ -146,17 +143,6 @@ contract PolicyEngineStaking is ISemver { _; } - /// @notice Modifier that reverts when the caller is not the owner. - modifier onlyOwner() { - if (msg.sender != _owner) revert PolicyEngineStaking_OnlyOwner(); - _; - } - - /// @notice Returns the owner address. - function owner() external view returns (address) { - return _owner; - } - /// @notice Returns the staking token address. /// /// @return The ERC20 token used for staking. @@ -164,14 +150,6 @@ contract PolicyEngineStaking is ISemver { return STAKING_TOKEN; } - /// @notice Transfers ownership of the contract to a new account. - /// @param _newOwner The address of the new owner. - function transferOwnership(address _newOwner) external onlyOwner { - if (_newOwner == address(0)) revert PolicyEngineStaking_ZeroAddress(); - emit OwnershipTransferred(_owner, _newOwner); - _owner = _newOwner; - } - /// @notice Pauses the contract. Stake is disabled while paused. function pause() external onlyOwner { paused = true; @@ -275,7 +253,15 @@ contract PolicyEngineStaking is ISemver { /// @notice Sets whether a staker can set the caller as beneficiary. When disallowing, /// if the staker's current beneficiary is the caller, their stake attribution is /// moved back to the staker (beneficiary reset to self). - /// + /// @dev This function is intentionally NOT gated by `whenNotPaused`. Allowlist + /// mutations remain available during pause so that beneficiaries can revoke + /// stakers at any time. Note that disallowing a staker who is currently + /// delegated to the caller will move effective stake attribution back to the + /// staker, changing ordering-power state even while the contract is paused. + /// @dev Trust assumption: stakers who delegate to a beneficiary implicitly trust + /// that the beneficiary will not remove them from the allowlist at a + /// disadvantageous time. Removal triggers `_increasePeData` on the staker, + /// which resets their `lastUpdate` and thus their accumulated staking weight. /// @param _staker The staker to allow or deny. /// @param _allowed The allowed state. function setAllowedStaker(address _staker, bool _allowed) public { @@ -318,13 +304,18 @@ contract PolicyEngineStaking is ISemver { emit EffectiveStakeChanged(_account, pe.effectiveStake); } - /// @notice Decreases effective stake for an account and updates timestamp. + /// @notice Decreases effective stake for an account. Only resets `lastUpdate` when + /// the effective stake reaches zero to avoid stale timestamps; otherwise the + /// existing timestamp is preserved so remaining stake keeps its staking weight. + /// /// @param _account The account address. /// @param _amount The amount to subtract. function _decreasePeData(address _account, uint128 _amount) internal { PEData storage pe = peData[_account]; pe.effectiveStake -= _amount; - pe.lastUpdate = uint128(block.timestamp); + if (pe.effectiveStake == 0) { + pe.lastUpdate = uint128(block.timestamp); + } emit EffectiveStakeChanged(_account, pe.effectiveStake); } } diff --git a/packages/contracts-bedrock/test/periphery/staking/PolicyEngineStaking.t.sol b/packages/contracts-bedrock/test/periphery/staking/PolicyEngineStaking.t.sol index 2125c28730d..178afe64e10 100644 --- a/packages/contracts-bedrock/test/periphery/staking/PolicyEngineStaking.t.sol +++ b/packages/contracts-bedrock/test/periphery/staking/PolicyEngineStaking.t.sol @@ -28,6 +28,7 @@ abstract contract PolicyEngineStaking_TestInit is CommonTest { event BeneficiaryAllowlistUpdated(address indexed beneficiary, address indexed staker, bool allowed); event Paused(); event Unpaused(); + event OwnershipTransferStarted(address indexed previousOwner, address indexed newOwner); function setUp() public virtual override { super.setUp(); @@ -61,59 +62,113 @@ abstract contract PolicyEngineStaking_TestInit is CommonTest { } /// @title PolicyEngineStaking_TransferOwnership_Test -/// @notice Tests the `transferOwnership` function. +/// @notice Tests the two-step `transferOwnership` / `acceptOwnership` pattern. contract PolicyEngineStaking_TransferOwnership_Test is PolicyEngineStaking_TestInit { - /// @notice Tests that owner can transfer ownership. + /// @notice Tests that transferOwnership nominates a pending owner without changing owner. function testFuzz_transferOwnership_succeeds(address _newOwner) external { vm.assume(_newOwner != address(0)); vm.expectEmit(address(staking)); - emit OwnershipTransferred(owner, _newOwner); + emit OwnershipTransferStarted(owner, _newOwner); vm.prank(owner); staking.transferOwnership(_newOwner); - assertEq(staking.owner(), _newOwner); + // Owner should NOT change yet + assertEq(staking.owner(), owner); + assertEq(staking.pendingOwner(), _newOwner); + } + + /// @notice Tests that pendingOwner can accept ownership. + function test_acceptOwnership_succeeds() external { + address newOwner = makeAddr("newOwner"); + + vm.prank(owner); + staking.transferOwnership(newOwner); + + vm.expectEmit(address(staking)); + emit OwnershipTransferred(owner, newOwner); + + vm.prank(newOwner); + staking.acceptOwnership(); + + assertEq(staking.owner(), newOwner); + assertEq(staking.pendingOwner(), address(0)); } - /// @notice Tests that new owner can exercise ownership after transfer. + /// @notice Tests that new owner can exercise ownership after accepting. function test_transferOwnership_newOwnerCanPause_succeeds() external { address newOwner = makeAddr("newOwner"); vm.prank(owner); staking.transferOwnership(newOwner); + vm.prank(newOwner); + staking.acceptOwnership(); + vm.prank(newOwner); staking.pause(); assertTrue(staking.paused()); } - /// @notice Tests that old owner loses ownership after transfer. + /// @notice Tests that old owner loses ownership after transfer is accepted. function test_transferOwnership_oldOwnerReverts_reverts() external { address newOwner = makeAddr("newOwner"); vm.prank(owner); staking.transferOwnership(newOwner); + vm.prank(newOwner); + staking.acceptOwnership(); + vm.prank(owner); - vm.expectRevert(IPolicyEngineStaking.PolicyEngineStaking_OnlyOwner.selector); + vm.expectRevert(abi.encodeWithSelector(IPolicyEngineStaking.OwnableUnauthorizedAccount.selector, owner)); staking.pause(); } - /// @notice Tests that non-owner cannot transfer ownership. + /// @notice Tests that non-owner cannot call transferOwnership. function testFuzz_transferOwnership_notOwner_reverts(address _caller) external { vm.assume(_caller != owner && _caller != address(0)); - vm.expectRevert(IPolicyEngineStaking.PolicyEngineStaking_OnlyOwner.selector); + vm.expectRevert(abi.encodeWithSelector(IPolicyEngineStaking.OwnableUnauthorizedAccount.selector, _caller)); vm.prank(_caller); staking.transferOwnership(alice); } - /// @notice Tests that transferring to zero address reverts. - function test_transferOwnership_zeroAddress_reverts() external { + /// @notice Tests that transferring to zero address cancels pending transfer. + function test_transferOwnership_zeroAddressCancels_succeeds() external { + address newOwner = makeAddr("newOwner"); + + vm.prank(owner); + staking.transferOwnership(newOwner); + assertEq(staking.pendingOwner(), newOwner); + vm.prank(owner); - vm.expectRevert(IPolicyEngineStaking.PolicyEngineStaking_ZeroAddress.selector); staking.transferOwnership(address(0)); + assertEq(staking.pendingOwner(), address(0)); + } + + /// @notice Tests that non-pending-owner cannot accept ownership. + function test_acceptOwnership_notPendingOwner_reverts() external { + address newOwner = makeAddr("newOwner"); + + vm.prank(owner); + staking.transferOwnership(newOwner); + + vm.prank(alice); + vm.expectRevert(abi.encodeWithSelector(IPolicyEngineStaking.OwnableUnauthorizedAccount.selector, alice)); + staking.acceptOwnership(); + } + + /// @notice Tests that pendingOwner view returns the correct value. + function test_pendingOwner_succeeds() external { + assertEq(staking.pendingOwner(), address(0)); + + address newOwner = makeAddr("newOwner"); + vm.prank(owner); + staking.transferOwnership(newOwner); + + assertEq(staking.pendingOwner(), newOwner); } } @@ -142,7 +197,7 @@ contract PolicyEngineStaking_Pause_Test is PolicyEngineStaking_TestInit { /// @notice Tests that non-owner cannot pause. function testFuzz_pause_notOwner_reverts(address _caller) external { vm.assume(_caller != owner && _caller != address(0)); - vm.expectRevert(IPolicyEngineStaking.PolicyEngineStaking_OnlyOwner.selector); + vm.expectRevert(abi.encodeWithSelector(IPolicyEngineStaking.OwnableUnauthorizedAccount.selector, _caller)); vm.prank(_caller); staking.pause(); } @@ -153,7 +208,7 @@ contract PolicyEngineStaking_Pause_Test is PolicyEngineStaking_TestInit { staking.pause(); vm.assume(_caller != owner && _caller != address(0)); - vm.expectRevert(IPolicyEngineStaking.PolicyEngineStaking_OnlyOwner.selector); + vm.expectRevert(abi.encodeWithSelector(IPolicyEngineStaking.OwnableUnauthorizedAccount.selector, _caller)); vm.prank(_caller); staking.unpause(); } @@ -624,7 +679,7 @@ contract PolicyEngineStaking_Constructor_Test is PolicyEngineStaking_TestInit { /// @notice Tests that the constructor reverts when owner is zero address. function test_constructor_zeroOwner_reverts() external { - vm.expectRevert(IPolicyEngineStaking.PolicyEngineStaking_ZeroAddress.selector); + vm.expectRevert(abi.encodeWithSelector(IPolicyEngineStaking.OwnableInvalidOwner.selector, address(0))); vm.deployCode( "PolicyEngineStaking.sol:PolicyEngineStaking", abi.encode(address(0), Predeploys.GOVERNANCE_TOKEN) ); @@ -1009,3 +1064,47 @@ contract PolicyEngineStaking_Integration_Test is PolicyEngineStaking_TestInit { return a; } } + +/// @title PolicyEngineStaking_Unstake_LastUpdate_Test +/// @notice Tests lastUpdate behavior on partial vs full unstake. +contract PolicyEngineStaking_Unstake_LastUpdate_Test is PolicyEngineStaking_TestInit { + /// @notice Tests that partial unstake preserves lastUpdate (does not reset it). + function test_unstake_partialPreservesLastUpdate_succeeds() external { + uint128 stakeAmount = uint128(100 ether); + uint128 unstakeAmount = uint128(40 ether); + + vm.prank(alice); + staking.stake(stakeAmount, alice); + (, uint128 lastUpdateAfterStake) = staking.peData(alice); + + // Warp time forward + vm.warp(block.timestamp + 1000); + + // Partial unstake — lastUpdate should NOT change + vm.prank(alice); + staking.unstake(unstakeAmount); + + (uint128 effectiveStake, uint128 lastUpdateAfterUnstake) = staking.peData(alice); + assertEq(effectiveStake, stakeAmount - unstakeAmount); + assertEq(lastUpdateAfterUnstake, lastUpdateAfterStake, "lastUpdate should be preserved on partial unstake"); + } + + /// @notice Tests that full unstake resets lastUpdate to block.timestamp. + function test_unstake_fullResetsLastUpdate_succeeds() external { + uint128 stakeAmount = uint128(100 ether); + + vm.prank(alice); + staking.stake(stakeAmount, alice); + + // Warp time forward + uint256 newTimestamp = block.timestamp + 1000; + vm.warp(newTimestamp); + + // Full unstake — lastUpdate should reset to block.timestamp + vm.prank(alice); + staking.unstake(stakeAmount); + + (, uint128 lastUpdateAfterUnstake) = staking.peData(alice); + assertEq(lastUpdateAfterUnstake, newTimestamp, "lastUpdate should reset on full unstake"); + } +}