From 0b00b05b5f2e0fd8abf5b678669e8b401544fdb5 Mon Sep 17 00:00:00 2001
From: Harry | Seungmin Jeon <94497407+sm-stack@users.noreply.github.com>
Date: Tue, 23 Jul 2024 14:23:18 +0900
Subject: [PATCH] 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
---
 .../contracts/contracts/L1/AssetManager.sol   | 53 ++++++++++++++++++-
 .../contracts/L1/interfaces/IAssetManager.sol | 17 +++---
 2 files changed, 61 insertions(+), 9 deletions(-)

diff --git a/packages/contracts/contracts/L1/AssetManager.sol b/packages/contracts/contracts/L1/AssetManager.sol
index ef0786395..39beebadc 100644
--- 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,11 +697,12 @@ 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) {
@@ -676,6 +710,21 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager {
         }
     }
 
+    /**
+     * @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
index 2401be5d1..a983fd4b0 100644
--- 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.
      *