Fix kton staking v1 rewards distribution (#21)

* Fix kton staking v1 rewards distribution

* fmt

* Revert "fmt"

This reverts commit 77a56cb.

Makefile

@@ -6,4 +6,6 @@ fmt    :; @forge fmt
 dry-run:; @forge script script/Deploy.s.sol:DeployScript
 deploy :; @forge script script/Deploy.s.sol:DeployScript --broadcast --verify
 
+migrate:; @forge script script/Migrate.s.sol:MigrateScript --sender 0x08837De0Ae21C270383D9F2de4DB03c7b1314632 -vvvv
+
 .PHONY: all flat clean test salt deploy

foundry.toml

@@ -15,7 +15,7 @@ bytecode_hash = "ipfs"
 extra_output = ["storageLayout"]
 extra_output_files = ["metadata"]
 fs_permissions = [{ access = "read", path = "out" }]
-eth_rpc_url  = "https://rpc.darwinia.network"
+eth_rpc_url  = "https://crab-rpc.darwinia.network"
 etherscan_api_key = "xxx"
 
 [rpc_endpoints]

script/Migrate.s.sol

@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import {Script} from "forge-std/Script.sol";
+import {safeconsole} from "forge-std/safeconsole.sol";
+import {Upgrades} from "openzeppelin-foundry-upgrades/Upgrades.sol";
+import {Core} from "openzeppelin-foundry-upgrades/internal/Core.sol";
+
+import {KtonDAOVaultV2} from "../src/staking/KtonDAOVaultV2.sol";
+
+contract MigrateScript is Script {
+    address vault = 0x652182C6aBc0bBE41b5702b05a26d109A405EAcA;
+    address v2 = 0xC4784B3593fF0ace8773ec79EF4F8D8901a8DCfC;
+
+    function run() public {
+        vm.startBroadcast();
+
+        Core.upgradeProxyTo(vault, v2, abi.encodeCall(KtonDAOVaultV2.initializeV2, ()));
+
+        vm.stopBroadcast();
+    }
+}

src/staking/KtonDAOVaultV2.sol

@@ -0,0 +1,79 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.20;
+
+import "./interfaces/IRewardsDistributionRecipient.sol";
+import "./interfaces/IOldRewardsDistributionRecipient.sol";
+import "./interfaces/IStakingRewards.sol";
+import "./interfaces/IOldStakingRewards.sol";
+import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
+
+/// @custom:oz-upgrades-from KtonDAOVault
+contract KtonDAOVaultV2 is Initializable, Ownable2StepUpgradeable {
+    // "modlda/trsry" in bytes.
+    address public constant SYSTEM_PALLET = 0x6D6f646c64612f74727372790000000000000000;
+
+    address public constant OLD_KTON_STAKING_REWARDS = 0x000000000419683a1a03AbC21FC9da25fd2B4dD7;
+    address public constant OLD_KTON_REWARDS_DISTRIBUTION = 0x000000000Ae5DB7BDAf8D071e680452e33d91Dd5;
+
+    address public stakingRewards;
+
+    modifier onlySystem() {
+        require(msg.sender == SYSTEM_PALLET, "Caller is not RewardsDistribution contract");
+        _;
+    }
+
+    function initializeV2() public reinitializer(2) {
+        uint256 rewards = OLD_KTON_REWARDS_DISTRIBUTION.balance;
+        IOldRewardsDistributionRecipient(OLD_KTON_REWARDS_DISTRIBUTION).distributeRewards(
+            OLD_KTON_STAKING_REWARDS, rewards
+        );
+        emit RewardsDistributed(OLD_KTON_REWARDS_DISTRIBUTION, rewards);
+    }
+
+    /// @custom:oz-upgrades-unsafe-allow constructor
+    constructor() {
+        _disableInitializers();
+    }
+
+    receive() external payable {}
+
+    /// Runtime migration Step:
+    /// 1. Migrate OLD_KTON_REWARDS_DISTRIBUTION's owner to this contracts address.
+    /// 2. distributeRewards to this contract address.
+    /// Note: The amount of the reward must be passed in via msg.value.
+    function distributeRewards() external payable onlySystem returns (bool) {
+        uint256 reward = msg.value;
+        require(reward > 0, "Nothing to distribute");
+        require(
+            address(this).balance >= reward, "RewardsDistribution contract does not have enough tokens to distribute"
+        );
+
+        uint256 oldTotalSupply = IOldStakingRewards(OLD_KTON_STAKING_REWARDS).totalSupply();
+        uint256 newTotalSupply = IStakingRewards(stakingRewards).underlyingTotalSupply();
+        uint256 totalSupply = oldTotalSupply + newTotalSupply;
+
+        if (totalSupply == 0) {
+            return true;
+        }
+
+        uint256 oldReward = reward * oldTotalSupply / totalSupply;
+        uint256 newReward = reward - oldReward;
+
+        if (oldReward > 0) {
+            IOldRewardsDistributionRecipient(OLD_KTON_REWARDS_DISTRIBUTION).distributeRewards{value: oldReward}(
+                OLD_KTON_STAKING_REWARDS, oldReward
+            );
+            emit RewardsDistributed(OLD_KTON_REWARDS_DISTRIBUTION, oldReward);
+        }
+
+        if (newReward > 0) {
+            IRewardsDistributionRecipient(stakingRewards).notifyRewardAmount{value: newReward}();
+            emit RewardsDistributed(stakingRewards, newReward);
+        }
+
+        return true;
+    }
+
+    event RewardsDistributed(address stakingRewards, uint256 amount);
+}

src/staking/interfaces/IOldRewardsDistributionRecipient.sol

@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: MIT
+pragma solidity >=0.4.24;
+
+interface IOldRewardsDistributionRecipient {
+    function distributeRewards(address ktonStakingRewards, uint256 reward) external payable returns (bool);
+}