feat: add custom withdrwaw

Grateful-protocol · Nov 1, 2024 · 95ccc81 · 95ccc81
1 parent c9f1519
commit 95ccc81
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/src/contracts/Grateful.sol b/src/contracts/Grateful.sol
@@ -281,6 +281,19 @@ contract Grateful is IGrateful, Ownable2Step {
     vault.redeem(_shares, msg.sender, address(this));
   }
 
+  function withdraw(address _token, uint256 _assets) external onlyWhenTokenWhitelisted(_token) {
+    AaveV3ERC4626 vault = vaults[_token];
+    if (address(vault) == address(0)) {
+      revert Grateful_VaultNotSet();
+    }
+    uint256 _shares = shares[msg.sender][_token];
+    if (vault.convertToShares(_assets) < _shares) {
+      revert Grateful_WithdrawExceedsShares();
+    }
+    shares[msg.sender][_token] = 0;
+    vault.withdraw(_assets, msg.sender, address(this));
+  }
+
   /// @inheritdoc IGrateful
   function switchYieldingFunds() external {
     yieldingFunds[msg.sender] = !yieldingFunds[msg.sender];

diff --git a/src/interfaces/IGrateful.sol b/src/interfaces/IGrateful.sol
@@ -75,6 +75,9 @@ interface IGrateful {
   /// @notice Thrown when the payment id has been used.
   error Grateful_PaymentIdAlreadyUsed();
 
+  /// @notice Thrown when the user tries to withdraw more than he has
+  error Grateful_WithdrawExceedsShares();
+
   /*///////////////////////////////////////////////////////////////
                                VARIABLES
     //////////////////////////////////////////////////////////////*/