From 358450117f450273286f0cfe78ee046a3a47e1f7 Mon Sep 17 00:00:00 2001
From: 0xchin <alanracciatti1220@gmail.com>
Date: Thu, 17 Oct 2024 10:52:43 -0300
Subject: [PATCH] fix: payments with yield reverting in token without vault

---
 src/contracts/Grateful.sol | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/contracts/Grateful.sol b/src/contracts/Grateful.sol
index 283e888..5bd423e 100644
--- a/src/contracts/Grateful.sol
+++ b/src/contracts/Grateful.sol
@@ -115,6 +115,7 @@ contract Grateful is IGrateful, Ownable2Step {
     return _amount - feeAmount;
   }
 
+  /// @inheritdoc IGrateful
   function owner() public view override(IGrateful, Ownable) returns (address) {
     return super.owner();
   }
@@ -359,10 +360,13 @@ contract Grateful is IGrateful, Ownable2Step {
         if (yieldingFunds[recipient]) {
           AaveV3ERC4626 vault = vaults[_token];
           if (address(vault) == address(0)) {
-            revert Grateful_VaultNotSet();
+            if (!IERC20(_token).transfer(recipient, recipientShare)) {
+              revert Grateful_TransferFailed();
+            }
+          } else {
+            uint256 _shares = vault.deposit(recipientShare, address(this));
+            shares[recipient][_token] += _shares;
           }
-          uint256 _shares = vault.deposit(recipientShare, address(this));
-          shares[recipient][_token] += _shares;
         } else {
           // Transfer tokens to recipient
           if (!IERC20(_token).transfer(recipient, recipientShare)) {
@@ -375,10 +379,13 @@ contract Grateful is IGrateful, Ownable2Step {
       if (yieldingFunds[_merchant]) {
         AaveV3ERC4626 vault = vaults[_token];
         if (address(vault) == address(0)) {
-          revert Grateful_VaultNotSet();
+          if (!IERC20(_token).transfer(_merchant, amountWithFee)) {
+            revert Grateful_TransferFailed();
+          }
+        } else {
+          uint256 _shares = vault.deposit(amountWithFee, address(this));
+          shares[_merchant][_token] += _shares;
         }
-        uint256 _shares = vault.deposit(amountWithFee, address(this));
-        shares[_merchant][_token] += _shares;
       } else {
         // Transfer tokens to merchant
         if (!IERC20(_token).transfer(_merchant, amountWithFee)) {