Skip to content

Commit

Permalink
fix: _checkDepositTokens _checkAssets...
Browse files Browse the repository at this point in the history
  • Loading branch information
@jakekidd
jakekidd committed Jul 19, 2024
1 parent c771773 commit 0a8a9a1
Show file tree
Hide file tree
Showing 5 changed files with 108 additions and 127 deletions.
56 changes: 29 additions & 27 deletions src/contracts/components/QWAaveV2.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,50 +3,53 @@ pragma solidity 0.8.23;

import {IERC20} from '@openzeppelin/contracts/token/ERC20/IERC20.sol';
import {IQWComponent} from 'interfaces/IQWComponent.sol';
import {QWComponentBase} from './QWComponentBase.sol';
import {ILendingPool} from 'interfaces/aave-v2/ILendingPool.sol';

/**
* @title AaveV2 Integration for Quant Wealth
* @notice This contract integrates with AaveV2 protocol for Quant Wealth management.
*/
contract QWAaveV2 is IQWComponent {
address public immutable QW_MANAGER;
contract QWAaveV2 is IQWComponent, QWComponentBase {
// Variables
address public immutable LENDING_POOL;

// Custom errors
error UnauthorizedAccess(); // Error for unauthorized caller

modifier onlyQwManager() {
if (msg.sender != QW_MANAGER) {
revert UnauthorizedAccess();
}
_;
}

/**
* @dev Constructor to initialize the contract with required addresses.
* @param _qwManager The address of the Quant Wealth Manager contract.
* @param _lendingPool The address of the AaveV2 pool contract.
*/
constructor(address _qwManager, address _lendingPool) {
QW_MANAGER = _qwManager;
constructor(
address _qwManager,
address _lendingPool
) QWComponentBase(_qwManager) {
LENDING_POOL = _lendingPool;
}

// Functions
/**
* @notice Executes a transaction on AaveV2 pool to deposit tokens.
* @dev This function is called by the parent contract to deposit tokens into the AaveV2 pool.
* @param _amount Amount of tokens to be deposited.
* @param _asset Address of the asset to be deposited.
* @param _asset The address of the asset token to be deposited.
* @return success boolean indicating the success of the transaction.
* @return assetAmountReceived Amount of asset tokens received.
*/
function open(uint256 _amount, address _asset) external override onlyQwManager returns (bool success, uint256 assetAmountReceived) {
IERC20(_asset).transferFrom(QW_MANAGER, address(this), _amount);
_checkDepositTokens(_amount, _asset);

// Approve the Aave lending pool to spend the tokens.
IERC20(_asset).approve(LENDING_POOL, _amount);

// Deposit tokens into Aave.
ILendingPool(LENDING_POOL).deposit(_asset, _amount, address(this), 0);
assetAmountReceived = IERC20(_asset).balanceOf(address(this));

// Get the balance of aTokens, which will reflect the principal investment(s) + interest.
// Check to ensure we have received the target asset.
assetAmountReceived = _checkAssetsAny(_asset);

// Transfer assets to QWManager.
IERC20(_asset).transfer(QW_MANAGER, assetAmountReceived);

success = true;
}
Expand All @@ -55,23 +58,22 @@ contract QWAaveV2 is IQWComponent {
* @notice Executes a transaction on AaveV2 pool to withdraw tokens.
* @dev This function is called by the parent contract to withdraw tokens from the AaveV2 pool.
* @param _amount Amount of holdings to be withdrawn.
* @param _asset Address of the asset to be withdrawn.
* @param _asset The address of the asset token to be withdrawn.
* @return success boolean indicating the success of the transaction.
* @return tokenAmountReceived Number of tokens to be returned to the user in exchange for the withdrawn ratio.
*/
function close(uint256 _amount, address _asset) external override onlyQwManager returns (bool success, uint256 tokenAmountReceived) {
_checkAssets(_amount, _asset);

// Withdraw the tokens from Aave.
ILendingPool(LENDING_POOL).withdraw(_asset, _amount, address(this));
tokenAmountReceived = IERC20(_asset).balanceOf(address(this));

// Check the balance of the investment token received.
tokenAmountReceived = _checkDepositTokensAny(_asset);

// Transfer tokens to QWManager.
IERC20(_asset).transfer(QW_MANAGER, tokenAmountReceived);
success = true;
}

/**
* @notice Gets the address of the Quant Wealth Manager contract.
* @return address The address of the Quant Wealth Manager contract.
*/
function getQWManager() external view override returns (address) {
return QW_MANAGER;
success = true;
}
}
52 changes: 22 additions & 30 deletions src/contracts/components/QWAaveV3.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,49 +4,47 @@ pragma solidity 0.8.23;
import {IPool} from '@aave/core-v3/contracts/interfaces/IPool.sol';
import {IERC20} from '@openzeppelin/contracts/token/ERC20/IERC20.sol';
import {IQWComponent} from 'interfaces/IQWComponent.sol';
import {QWComponentBase} from './QWComponentBase.sol';

/**
* @title AaveV3 Integration for Quant Wealth
* @notice This contract integrates with AaveV3 protocol for Quant Wealth management.
*/
contract QWAaveV3 is IQWComponent {
address public immutable QW_MANAGER;
contract QWAaveV3 is IQWComponent, QWComponentBase {
// Variables
address public immutable POOL;

// Custom errors
error UnauthorizedAccess(); // Error for unauthorized caller

modifier onlyQwManager() {
if (msg.sender != QW_MANAGER) {
revert UnauthorizedAccess();
}
_;
}

/**
* @dev Constructor to initialize the contract with required addresses.
* @param _qwManager The address of the Quant Wealth Manager contract.
* @param _pool The address of the AaveV3 pool contract.
*/
constructor(address _qwManager, address _pool) {
QW_MANAGER = _qwManager;
constructor(
address _qwManager,
address _pool
) QWComponentBase(_qwManager) {
POOL = _pool;
}

// Functions
/**
* @notice Executes a transaction on AaveV3 pool to deposit tokens.
* @dev This function is called by the parent contract to deposit tokens into the AaveV3 pool.
* @param _amount Amount of tokens to be deposited.
* @param _asset Address of the asset to be deposited.
* @param _asset The address of the asset token.
* @return success boolean indicating the success of the transaction.
* @return assetAmountReceived Amount of assets received from the deposit.
*/
function open(uint256 _amount, address _asset) external override onlyQwManager returns (bool success, uint256 assetAmountReceived) {
IERC20(_asset).transferFrom(QW_MANAGER, address(this), _amount);
IERC20(_asset).approve(POOL, _amount);
_checkDepositTokens(_amount, _asset);

IERC20(_asset).approve(POOL, _amount);
IPool(POOL).supply(_asset, _amount, address(this), 0);
assetAmountReceived = IERC20(_asset).balanceOf(address(this));

assetAmountReceived = _checkAssetsAny(_asset);

// Transfer assets to QWManager.
IERC20(_asset).transfer(QW_MANAGER, assetAmountReceived);

success = true;
}
Expand All @@ -55,23 +53,17 @@ contract QWAaveV3 is IQWComponent {
* @notice Executes a transaction on AaveV3 pool to withdraw tokens.
* @dev This function is called by the parent contract to withdraw tokens from the AaveV3 pool.
* @param _amount Amount to withdraw.
* @param _asset Address of the asset to be withdrawn.
* @param _asset The address of the asset token.
* @return success boolean indicating the success of the transaction.
* @return tokenAmountReceived Amount of tokens received from the withdrawal.
*/
function close(uint256 _amount, address _asset) external override onlyQwManager returns (bool success, uint256 tokenAmountReceived) {
IPool(POOL).withdraw(_asset, _amount, address(this));
tokenAmountReceived = IERC20(_asset).balanceOf(address(this));
_checkAssets(_amount, _asset);

IERC20(_asset).transfer(QW_MANAGER, tokenAmountReceived);
success = true;
}
IPool(POOL).withdraw(_asset, _amount, QW_MANAGER);

/**
* @notice Gets the address of the Quant Wealth Manager contract.
* @return address The address of the Quant Wealth Manager contract.
*/
function getQWManager() external view override returns (address) {
return QW_MANAGER;
tokenAmountReceived = _checkDepositTokensAny(_asset);

success = true;
}
}
52 changes: 17 additions & 35 deletions src/contracts/components/QWComponentBase.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,11 @@ import {IQWComponent} from 'interfaces/IQWComponent.sol';

abstract contract QWComponentBase is IQWComponent {
address public immutable QW_MANAGER;
address public immutable INVESTMENT_TOKEN;
address public immutable ASSET_TOKEN;

// Custom errors
error InvalidCallData(); // Error for invalid call data
error UnauthorizedAccess(); // Error for unauthorized caller
error IncorrectInvestmentTokensReceived(uint256 amount);
error IncorrectDepositTokensReceived(uint256 amount);
error IncorrectAssetTokensReceived(uint256 amount);

modifier onlyQwManager() {
Expand All @@ -24,10 +22,8 @@ abstract contract QWComponentBase is IQWComponent {
_;
}

constructor(address _qwManager, address _investmentToken, address _assetToken) {
constructor(address _qwManager) {
QW_MANAGER = _qwManager;
INVESTMENT_TOKEN = _investmentToken;
ASSET_TOKEN = _assetToken;
}

/**
Expand All @@ -38,60 +34,46 @@ abstract contract QWComponentBase is IQWComponent {
}

/**
* @notice Gets the address of the investment token, the token input and output for the contract.
* @notice Checks whether we've received the proper amount of deposit tokens.
*/
function getInvestmentToken() external view override returns (address) {
return INVESTMENT_TOKEN;
}

/**
* @notice Gets the address of the asset token, the token that is purchased and sold using the investment token.
*/
function getAssetToken() external view override returns (address) {
return ASSET_TOKEN;
}

/**
* @notice Checks whether we've received the proper amount of investment tokens.
*/
function _checkInvestment(uint256 _expectedAmount) internal view returns (uint256) {
uint256 balance = IERC20(INVESTMENT_TOKEN).balanceOf(address(this));
function _checkDepositTokens(uint256 _expectedAmount, address _depositToken) internal view returns (uint256) {
uint256 balance = IERC20(_depositToken).balanceOf(address(this));
if (balance != _expectedAmount) {
revert IncorrectInvestmentTokensReceived(balance);
revert IncorrectDepositTokensReceived(balance);
}
return balance;
}

/**
* @notice Checks whether we've received any amount of investment tokens.
* @notice Checks whether we've received any amount of deposit tokens.
*/
function _checkInvestmentAny() internal view returns (uint256) {
uint256 balance = IERC20(INVESTMENT_TOKEN).balanceOf(address(this));
function _checkDepositTokensAny(address _depositToken) internal view returns (uint256) {
uint256 balance = IERC20(_depositToken).balanceOf(address(this));
if (balance == 0) {
revert IncorrectInvestmentTokensReceived(balance);
revert IncorrectDepositTokensReceived(balance);
}
return balance;
}

/**
* @notice Check how much assets we have received upon closing a position.
*/
function _checkAssets(uint256 _expectedAmount) internal view returns (uint256) {
uint256 balance = IERC20(ASSET_TOKEN).balanceOf(address(this));
function _checkAssets(uint256 _expectedAmount, address _asset) internal view returns (uint256) {
uint256 balance = IERC20(_asset).balanceOf(address(this));
if (balance != _expectedAmount) {
revert IncorrectInvestmentTokensReceived(balance);
revert IncorrectAssetTokensReceived(balance);
}
return balance;
}

/**
* @notice Checks whether we've received any amount of investment tokens.
* @notice Checks whether we've received any amount of asset tokens.
*/
function _checkAssetsAny() internal view returns (uint256) {
uint256 balance = IERC20(ASSET_TOKEN).balanceOf(address(this));
function _checkAssetsAny(address _asset) internal view returns (uint256) {
uint256 balance = IERC20(_asset).balanceOf(address(this));
if (balance == 0) {
revert IncorrectAssetTokensReceived(balance);
}
return balance;
}
}
}
34 changes: 18 additions & 16 deletions src/contracts/components/QWCompound.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,16 +17,12 @@ contract QWCompound is IQWComponent, QWComponentBase {
/**
* @dev Constructor to initialize the contract with required addresses.
* @param _qwManager The address of the Quant Wealth Manager contract.
* @param _investmentToken The address of the investment token (e.g., USDC).
* @param _assetToken The address of the asset token received from Compound (e.g., cUSDC).
* @param _comet The address of the Compound comet contract.
*/
constructor(
address _qwManager,
address _investmentToken,
address _assetToken,
address _comet
) QWComponentBase(_qwManager, _investmentToken, _assetToken) {
) QWComponentBase(_qwManager) {
COMET = _comet;
}

Expand All @@ -35,43 +31,49 @@ contract QWCompound is IQWComponent, QWComponentBase {
* @notice Executes a transaction on Compound comet to deposit tokens.
* @dev This function is called by the parent contract to deposit tokens into the Compound comet.
* @param _amount Amount of tokens to be deposited.
* @param _asset The address of the asset token to be deposited.
* @return success boolean indicating the success of the transaction.
* @return assetAmountReceived The amount of asset tokens received from the deposit.
*/
function open(
uint256 _amount
uint256 _amount,
address _asset
) external override onlyQwManager returns (bool success, uint256 assetAmountReceived) {
_checkInvestment(_amount);
_checkDepositTokens(_amount, _asset);

IERC20(INVESTMENT_TOKEN).approve(COMET, _amount);
IERC20(_asset).approve(COMET, _amount);

// Perform the supply to Compound.
IComet(COMET).supplyTo(address(this), INVESTMENT_TOKEN, _amount);
IComet(COMET).supplyTo(address(this), _asset, _amount);

assetAmountReceived = _checkAssetsAny();
assetAmountReceived = _checkAssetsAny(_asset);
success = true;

// TODO: Transfer tokens to QWManager
// Transfer assets to QWManager.
IERC20(_asset).transfer(QW_MANAGER, assetAmountReceived);
}

/**
* @notice Executes a transaction on Compound comet to withdraw tokens.
* @dev This function is called by the parent contract to withdraw tokens from the Compound comet.
* @param _amount Amount of asset tokens to withdraw.
* @param _asset The address of the asset token.
* @return success boolean indicating the success of the transaction.
* @return tokenAmountReceived The amount of tokens received from the withdrawal.
*/
function close(
uint256 _amount
uint256 _amount,
address _asset
) external override onlyQwManager returns (bool success, uint256 tokenAmountReceived) {
_checkAssets(_amount);
_checkAssets(_amount, _asset);

// Perform the withdraw from Compound.
IComet(COMET).withdrawTo(address(this), INVESTMENT_TOKEN, _amount);
IComet(COMET).withdrawTo(address(this), _asset, _amount);

tokenAmountReceived = _checkInvestmentAny();
tokenAmountReceived = _checkDepositTokensAny(_asset);
success = true;

// TODO: transfer tokens to QWManager
// Transfer tokens to QWManager.
IERC20(_asset).transfer(QW_MANAGER, tokenAmountReceived);
}
}
Loading

0 comments on commit 0a8a9a1

Please sign in to comment.