From c5fd3b5ab4046f429239ea984a00f3be92509755 Mon Sep 17 00:00:00 2001 From: "Eugene Y. Q. Shen" Date: Thu, 19 Sep 2024 08:01:33 -0700 Subject: [PATCH 1/2] [NES-211] validate that there is enough AssetToken to create new yield distributions --- smart-wallets/src/extensions/AssetVault.sol | 46 ++++++++++++-------- smart-wallets/src/interfaces/IAssetVault.sol | 18 ++++---- 2 files changed, 38 insertions(+), 26 deletions(-) diff --git a/smart-wallets/src/extensions/AssetVault.sol b/smart-wallets/src/extensions/AssetVault.sol index b556aef..ec0800a 100644 --- a/smart-wallets/src/extensions/AssetVault.sol +++ b/smart-wallets/src/extensions/AssetVault.sol @@ -1,10 +1,10 @@ // SPDX-License-Identifier: MIT pragma solidity ^0.8.25; -import { ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol"; +import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; +import { IAssetToken } from "../interfaces/IAssetToken.sol"; import { IAssetVault } from "../interfaces/IAssetVault.sol"; -import { AssetToken } from "../token/AssetToken.sol"; /** * @title AssetVault @@ -43,9 +43,9 @@ contract AssetVault is IAssetVault { /// @custom:storage-location erc7201:plume.storage.AssetVault struct AssetVaultStorage { /// @dev Mapping of yield allowances for each asset token and beneficiary - mapping(AssetToken assetToken => mapping(address beneficiary => Yield allowance)) yieldAllowances; + mapping(IAssetToken assetToken => mapping(address beneficiary => Yield allowance)) yieldAllowances; /// @dev Mapping of the yield distribution list for each asset token - mapping(AssetToken assetToken => YieldDistributionListItem distribution) yieldDistributions; + mapping(IAssetToken assetToken => YieldDistributionListItem distribution) yieldDistributions; } // keccak256(abi.encode(uint256(keccak256("plume.storage.AssetVault")) - 1)) & ~bytes32(uint256(0xff)) @@ -80,7 +80,7 @@ contract AssetVault is IAssetVault { * @param expiration Timestamp at which the yield expires */ event YieldAllowanceUpdated( - AssetToken indexed assetToken, address indexed beneficiary, uint256 amount, uint256 expiration + IAssetToken indexed assetToken, address indexed beneficiary, uint256 amount, uint256 expiration ); /** @@ -91,7 +91,7 @@ contract AssetVault is IAssetVault { * @param yieldShare Amount of currencyToken that was redistributed to the beneficiary */ event YieldRedistributed( - AssetToken indexed assetToken, address indexed beneficiary, ERC20 indexed currencyToken, uint256 yieldShare + IAssetToken indexed assetToken, address indexed beneficiary, IERC20 indexed currencyToken, uint256 yieldShare ); /** @@ -102,7 +102,7 @@ contract AssetVault is IAssetVault { * @param expiration Timestamp at which the yield expires */ event YieldDistributionCreated( - AssetToken indexed assetToken, address indexed beneficiary, uint256 amount, uint256 expiration + IAssetToken indexed assetToken, address indexed beneficiary, uint256 amount, uint256 expiration ); /** @@ -111,14 +111,14 @@ contract AssetVault is IAssetVault { * @param beneficiary Address of the beneficiary of the yield distribution * @param amount Amount of AssetTokens that are renounced from the yield distributions of the beneficiary */ - event YieldDistributionRenounced(AssetToken indexed assetToken, address indexed beneficiary, uint256 amount); + event YieldDistributionRenounced(IAssetToken indexed assetToken, address indexed beneficiary, uint256 amount); /** * @notice Emitted when anyone clears expired yield distributions from the linked list * @param assetToken AssetToken from which the yield is to be redistributed * @param amountCleared Amount of AssetTokens that were cleared from the yield distributions */ - event YieldDistributionsCleared(AssetToken indexed assetToken, uint256 amountCleared); + event YieldDistributionsCleared(IAssetToken indexed assetToken, uint256 amountCleared); // Errors @@ -150,9 +150,16 @@ contract AssetVault is IAssetVault { * @param amount Amount of assetTokens that the beneficiary tried to accept the yield of */ error InsufficientYieldAllowance( - AssetToken assetToken, address beneficiary, uint256 allowanceAmount, uint256 amount + IAssetToken assetToken, address beneficiary, uint256 allowanceAmount, uint256 amount ); + /** + * @notice Indicates a failure because the user wallet does not have enough AssetTokens + * @param assetToken AssetToken for which a new yield distribution is to be made + * @param amount Amount of assetTokens that the user wallet tried to add to the distribution + */ + error InsufficientBalance(IAssetToken assetToken, uint256 amount); + /** * @notice Indicates a failure because the beneficiary does not have enough yield distributions * @param assetToken AssetToken from which the yield is to be redistributed @@ -161,7 +168,7 @@ contract AssetVault is IAssetVault { * @param amountRenounced Amount of assetTokens that the beneficiary tried to renounce the yield of */ error InsufficientYieldDistributions( - AssetToken assetToken, address beneficiary, uint256 amount, uint256 amountRenounced + IAssetToken assetToken, address beneficiary, uint256 amount, uint256 amountRenounced ); /** @@ -203,7 +210,7 @@ contract AssetVault is IAssetVault { * @param expiration Timestamp at which the yield expires */ function updateYieldAllowance( - AssetToken assetToken, + IAssetToken assetToken, address beneficiary, uint256 amount, uint256 expiration @@ -234,8 +241,8 @@ contract AssetVault is IAssetVault { * @param currencyTokenAmount Amount of currencyToken to redistribute */ function redistributeYield( - AssetToken assetToken, - ERC20 currencyToken, + IAssetToken assetToken, + IERC20 currencyToken, uint256 currencyTokenAmount ) external onlyWallet { if (currencyTokenAmount == 0) { @@ -265,7 +272,7 @@ contract AssetVault is IAssetVault { * @notice Get the number of AssetTokens that are currently locked in the AssetVault * @param assetToken AssetToken from which the yield is to be redistributed */ - function getBalanceLocked(AssetToken assetToken) public view returns (uint256 balanceLocked) { + function getBalanceLocked(IAssetToken assetToken) public view returns (uint256 balanceLocked) { // Iterate through the list and sum up the locked balance across all yield distributions YieldDistributionListItem storage distribution = _getAssetVaultStorage().yieldDistributions[assetToken]; uint256 amountLocked = distribution.yield.amount; @@ -285,7 +292,7 @@ contract AssetVault is IAssetVault { * @param amount Amount of AssetTokens included in this yield allowance * @param expiration Timestamp at which the yield expires */ - function acceptYieldAllowance(AssetToken assetToken, uint256 amount, uint256 expiration) external { + function acceptYieldAllowance(IAssetToken assetToken, uint256 amount, uint256 expiration) external { AssetVaultStorage storage $ = _getAssetVaultStorage(); address beneficiary = msg.sender; Yield storage allowance = $.yieldAllowances[assetToken][beneficiary]; @@ -302,6 +309,9 @@ contract AssetVault is IAssetVault { if (allowance.amount < amount) { revert InsufficientYieldAllowance(assetToken, beneficiary, allowance.amount, amount); } + if (assetToken.getBalanceAvailable(address(this)) < amount) { + revert InsufficientBalance(assetToken, amount); + } allowance.amount -= amount; @@ -327,7 +337,7 @@ contract AssetVault is IAssetVault { * @param expiration Timestamp at which the yield expires */ function renounceYieldDistribution( - AssetToken assetToken, + IAssetToken assetToken, uint256 amount, uint256 expiration ) external returns (uint256 amountRenounced) { @@ -376,7 +386,7 @@ contract AssetVault is IAssetVault { * reaching the gas limit, and the caller must call the function again to clear more. * @param assetToken AssetToken from which the yield is to be redistributed */ - function clearYieldDistributions(AssetToken assetToken) external { + function clearYieldDistributions(IAssetToken assetToken) external { uint256 amountCleared = 0; // Iterate through the list and delete all expired yield distributions diff --git a/smart-wallets/src/interfaces/IAssetVault.sol b/smart-wallets/src/interfaces/IAssetVault.sol index 7cba9da..08971c8 100644 --- a/smart-wallets/src/interfaces/IAssetVault.sol +++ b/smart-wallets/src/interfaces/IAssetVault.sol @@ -1,26 +1,28 @@ // SPDX-License-Identifier: MIT pragma solidity ^0.8.25; -import { AssetToken } from "../token/AssetToken.sol"; -import { ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol"; +import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; + +import { IAssetToken } from "../interfaces/IAssetToken.sol"; interface IAssetVault { function updateYieldAllowance( - AssetToken assetToken, + IAssetToken assetToken, address beneficiary, uint256 amount, uint256 expiration ) external; - function redistributeYield(AssetToken assetToken, ERC20 currencyToken, uint256 currencyTokenAmount) external; + function redistributeYield(IAssetToken assetToken, IERC20 currencyToken, uint256 currencyTokenAmount) external; - function getBalanceLocked(AssetToken assetToken) external view returns (uint256 balanceLocked); - function acceptYieldAllowance(AssetToken assetToken, uint256 amount, uint256 expiration) external; + function wallet() external view returns (address wallet); + function getBalanceLocked(IAssetToken assetToken) external view returns (uint256 balanceLocked); + function acceptYieldAllowance(IAssetToken assetToken, uint256 amount, uint256 expiration) external; function renounceYieldDistribution( - AssetToken assetToken, + IAssetToken assetToken, uint256 amount, uint256 expiration ) external returns (uint256 amountRenounced); - function clearYieldDistributions(AssetToken assetToken) external; + function clearYieldDistributions(IAssetToken assetToken) external; } From 6963b7e262f2e21ff776c91928b8f6fe55440850 Mon Sep 17 00:00:00 2001 From: "Eugene Y. Q. Shen" Date: Thu, 19 Sep 2024 08:10:23 -0700 Subject: [PATCH 2/2] check if the distribution already exists --- smart-wallets/src/extensions/AssetVault.sol | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/smart-wallets/src/extensions/AssetVault.sol b/smart-wallets/src/extensions/AssetVault.sol index ec0800a..ae7c1ae 100644 --- a/smart-wallets/src/extensions/AssetVault.sol +++ b/smart-wallets/src/extensions/AssetVault.sol @@ -315,9 +315,14 @@ contract AssetVault is IAssetVault { allowance.amount -= amount; - // Append the newly created yield distribution to the end of the linked list + // Either update the existing distribution with the same expiration or append a new one YieldDistributionListItem storage distribution = $.yieldDistributions[assetToken]; while (distribution.yield.amount > 0) { + if (distribution.beneficiary == beneficiary && distribution.yield.expiration == expiration) { + distribution.yield.amount += amount; + emit YieldDistributionCreated(assetToken, beneficiary, amount, expiration); + return; + } distribution = distribution.next[0]; } distribution.beneficiary = beneficiary;