feat(contracts): add withdraw method for validator (#357)

* feat: add withdraw method for validator * feat: apply PR reviews * feat: review applied * feat: review applied * feat: review applied
kroma-network · Jul 23, 2024 · 0b00b05 · 0b00b05
1 parent f8bbed7
commit 0b00b05
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 9 deletions.
diff --git a/packages/contracts/contracts/L1/AssetManager.sol b/packages/contracts/contracts/L1/AssetManager.sol
@@ -94,6 +94,14 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager {
         _;
     }
 
+    /**
+     * @notice Modifier to check if the caller is the withdraw account of the validator.
+     */
+    modifier onlyWithdrawAccount(address validator) {
+        if (msg.sender != _vaults[validator].withdrawAccount) revert NotAllowedCaller();
+        _;
+    }
+
     /**
      * @notice Semantic version.
      * @custom:semver 1.0.0
@@ -190,6 +198,13 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager {
             MIN_DELEGATION_PERIOD;
     }
 
+    /**
+     * @inheritdoc IAssetManager
+     */
+    function canWithdrawAt(address validator) public view returns (uint128) {
+        return _vaults[validator].lastDepositedAt + MIN_DELEGATION_PERIOD;
+    }
+
     /**
      * @inheritdoc IAssetManager
      */
@@ -284,6 +299,24 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager {
         emit Deposited(msg.sender, assets);
     }
 
+    /**
+     * @inheritdoc IAssetManager
+     */
+    function withdraw(address validator, uint128 assets) external onlyWithdrawAccount(validator) {
+        if (assets == 0) revert NotAllowedZeroInput();
+        if (canWithdrawAt(validator) > block.timestamp) {
+            revert NotElapsedMinDelegationPeriod();
+        }
+
+        _withdraw(validator, assets);
+
+        VALIDATOR_MANAGER.updateValidatorTree(validator, true);
+
+        ASSET_TOKEN.safeTransfer(_vaults[validator].withdrawAccount, assets);
+
+        emit Withdrawn(validator, assets);
+    }
+
     /**
      * @inheritdoc IAssetManager
      */
@@ -664,18 +697,34 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager {
      * @param updateTree Flag to update the validator tree.
      */
     function _deposit(address validator, uint128 assets, bool updateTree) internal {
-        Asset storage asset = _vaults[validator].asset;
+        Vault storage vault = _vaults[validator];
         ASSET_TOKEN.safeTransferFrom(validator, address(this), assets);
 
         unchecked {
-            asset.validatorKro += assets;
+            vault.asset.validatorKro += assets;
+            vault.lastDepositedAt = uint128(block.timestamp);
         }
 
         if (updateTree) {
             VALIDATOR_MANAGER.updateValidatorTree(validator, false);
         }
     }
 
+    /**
+     * @notice Internal function to withdraw KRO by the validator.
+     *
+     * @param validator Address of the validator.
+     * @param assets    The amount of KRO to withdraw.
+     */
+    function _withdraw(address validator, uint128 assets) internal {
+        Asset storage asset = _vaults[validator].asset;
+        if (assets > asset.validatorKro - asset.validatorKroBonded) revert InsufficientAsset();
+
+        unchecked {
+            asset.validatorKro -= assets;
+        }
+    }
+
     /**
      * @notice Internal function to delegate KRO to the validator.
      *

diff --git a/packages/contracts/contracts/L1/interfaces/IAssetManager.sol b/packages/contracts/contracts/L1/interfaces/IAssetManager.sol
@@ -400,6 +400,16 @@ interface IAssetManager {
      */
     function deposit(uint128 assets) external;
 
+    /**
+     * @notice Withdraw KRO. To withdraw KRO, the validator should be initiated and MIN_DELEGATION_PERIOD
+     *         should be passed after the last deposit time. Only withdrawAccount of the validator can call
+     *         this function.
+     *
+     * @param validator Address of the validator.
+     * @param assets    The amount of KRO to withdraw.
+     */
+    function withdraw(address validator, uint128 assets) external;
+
     /**
      * @notice Delegate KRO to the validator and returns the amount of shares that the vault would
      *         exchange.
@@ -427,13 +437,6 @@ interface IAssetManager {
      */
     function delegateKghBatch(address validator, uint256[] calldata tokenIds) external;
 
-    /**
-     * @notice Withdraw KRO by the validator.
-     *
-     * @param assets The amount of KRO to withdraw.
-     */
-    function withdraw(uint128 assets) external;
-
     /**
      * @notice Undelegate the KRO of given assets for the given validator.
      *