feat: add transfer from a designated vault when increasing reward

kroma-network · Jul 21, 2024 · 6d34e1f · 6d34e1f
1 parent bf82155
commit 6d34e1f
Showing 1 changed file with 22 additions and 12 deletions.
diff --git a/packages/contracts/contracts/L1/AssetManager.sol b/packages/contracts/contracts/L1/AssetManager.sol
@@ -56,6 +56,11 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager {
      */
     address public immutable SECURITY_COUNCIL;
 
+    /**
+     * @notice The address of Validator Reward Vault. Can be updated via upgrade.
+     */
+    address public immutable VALIDATOR_REWARD_VAULT;
+
     /**
      * @notice Address of ValidatorManager contract. Can be updated via upgrade.
      */
@@ -104,19 +109,21 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager {
     /**
      * @notice Constructs the AssetManager contract.
      *
-     * @param _assetToken         Address of the KRO token.
-     * @param _kgh                Address of the KGH token.
-     * @param _kghManager         Address of the KGHManager contract.
-     * @param _securityCouncil    Address of the SecurityCouncil contract.
-     * @param _validatorManager   Address of the ValidatorManager contract.
-     * @param _undelegationPeriod Period that should wait to finalize the undelegation.
-     * @param _bondAmount         Amount to bond.
+     * @param _assetToken           Address of the KRO token.
+     * @param _kgh                  Address of the KGH token.
+     * @param _kghManager           Address of the KGHManager contract.
+     * @param _securityCouncil      Address of the SecurityCouncil contract.
+     * @param _validatorRewardVault Address of the Validator Reward Vault.
+     * @param _validatorManager     Address of the ValidatorManager contract.
+     * @param _undelegationPeriod   Period that should wait to finalize the undelegation.
+     * @param _bondAmount           Amount to bond.
      */
     constructor(
         IERC20 _assetToken,
         IERC721 _kgh,
         IKGHManager _kghManager,
         address _securityCouncil,
+        address _validatorRewardVault,
         IValidatorManager _validatorManager,
         uint128 _undelegationPeriod,
         uint128 _bondAmount
@@ -125,6 +132,7 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager {
         KGH = _kgh;
         KGH_MANAGER = _kghManager;
         SECURITY_COUNCIL = _securityCouncil;
+        VALIDATOR_REWARD_VAULT = _validatorRewardVault;
         VALIDATOR_MANAGER = _validatorManager;
         UNDELEGATION_PERIOD = _undelegationPeriod;
         BOND_AMOUNT = _bondAmount;
@@ -719,21 +727,23 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager {
         uint128 boostedReward,
         uint128 validatorReward
     ) external onlyValidatorManager {
+        // Distribute the reward from a designated vault to the AssetManager contract.
+        ASSET_TOKEN.transferFrom(VALIDATOR_REWARD_VAULT, address(this), baseReward + boostedReward + validatorReward);
+
         // If reward is distributed to SECURITY_COUNCIL, transfer it directly.
         if (validator == SECURITY_COUNCIL) {
             ASSET_TOKEN.safeTransfer(
                 SECURITY_COUNCIL,
                 baseReward + boostedReward + validatorReward
             );
         } else {
-            Vault storage vault = _vaults[validator];
+            Asset storage asset = _vaults[validator].asset;
             unchecked {
-                vault.asset.totalKro += baseReward;
-                // TODO: handle reward for boosted reward
+                asset.totalKro += baseReward;
+                asset.validatorKro += validatorReward;
+                asset.rewardPerKghStored += boostedReward / asset.totalKgh;
             }
         }
-
-        // TODO - Distribute the reward from a designated vault to the AssetManager contract.
     }
 
     /**