diff --git a/contracts/MigrationZapper.sol b/contracts/MigrationZapper.sol index f4c98890..79b90d9a 100644 --- a/contracts/MigrationZapper.sol +++ b/contracts/MigrationZapper.sol @@ -12,11 +12,17 @@ contract MigrationZapper { IMigrator public immutable migrator; IStaking public immutable ognStaking; - constructor(address _ogv, address _ogn, address _migrator, address _ognStaking) { + address public immutable governor; + + error NotGovernor(); + + constructor(address _ogv, address _ogn, address _migrator, address _ognStaking, address _governor) { ogv = IMintableERC20(_ogv); ogn = IMintableERC20(_ogn); migrator = IMigrator(_migrator); ognStaking = IStaking(_ognStaking); + + governor = _governor; } function initialize() external { @@ -60,5 +66,23 @@ contract MigrationZapper { false, -1 // New stake ); + + // Transfer remaining OGN to the receiver + if (ognReceived > newStakeAmount) { + ogn.transfer(msg.sender, ognReceived - newStakeAmount); + } + } + + /** + * Transfers any tokens sent by mistake out of the contract + * @param token Token address + * @param amount Amount of token to transfer + */ + function transferTokens(address token, uint256 amount) external { + if (msg.sender != governor) { + revert NotGovernor(); + } + + IMintableERC20(token).transfer(governor, amount); } } diff --git a/script/deploy/mainnet/012_MigrationZapperScript.sol b/script/deploy/mainnet/012_MigrationZapperScript.sol index 4e4e8074..dc5f2aca 100644 --- a/script/deploy/mainnet/012_MigrationZapperScript.sol +++ b/script/deploy/mainnet/012_MigrationZapperScript.sol @@ -33,8 +33,9 @@ contract MigrationZapperScript is BaseMainnetScript { console.log("Deploy:"); console.log("------------"); - MigrationZapper zapper = - new MigrationZapper(Addresses.OGV, Addresses.OGN, deployedContracts["MIGRATOR"], deployedContracts["XOGN"]); + MigrationZapper zapper = new MigrationZapper( + Addresses.OGV, Addresses.OGN, deployedContracts["MIGRATOR"], deployedContracts["XOGN"], Addresses.TIMELOCK + ); _recordDeploy("MIGRATION_ZAPPER", address(zapper)); // Make sure Migrator can move OGV and OGN diff --git a/tests/staking/ZapperForkTest.t.sol b/tests/staking/ZapperForkTest.t.sol index 6677b96e..19842784 100644 --- a/tests/staking/ZapperForkTest.t.sol +++ b/tests/staking/ZapperForkTest.t.sol @@ -84,4 +84,38 @@ contract ZapperForkTest is Test { vm.stopPrank(); } + + function testMigrateBalanceAndPartialStake() public { + vm.startPrank(ogvWhale); + + uint256[] memory lockupIds = new uint256[](0); + uint256 stakeAmount = (50 ether * 0.09137 ether) / 1 ether; + + uint256 ognBalanceBefore = ogn.balanceOf(ogvWhale); + + zapper.migrate(100 ether, stakeAmount, 300 days); + + // Should have it in a single OGN lockup + (uint128 amount, uint128 end, uint256 points) = xogn.lockups(ogvWhale, 0); + assertEq(amount, stakeAmount, "Balance not staked"); + + assertEq(ogn.balanceOf(ogvWhale), ognBalanceBefore + stakeAmount, "Less OGN received"); + + vm.stopPrank(); + } + + function testTransferTokens() public { + vm.startPrank(ogvWhale); + ogn.transfer(address(zapper), 100 ether); + vm.stopPrank(); + + vm.startPrank(Addresses.TIMELOCK); + zapper.transferTokens(address(ogn), 100 ether); + vm.stopPrank(); + + vm.startPrank(ogvWhale); + vm.expectRevert(bytes4(keccak256("NotGovernor()"))); + zapper.transferTokens(address(ogn), 100 ether); + vm.stopPrank(); + } }