From f544e1f94a79e11f4d4495c63908c8fc914e9701 Mon Sep 17 00:00:00 2001 From: "ultrasecr.eth" <241804+ultrasecreth@users.noreply.github.com> Date: Wed, 14 Feb 2024 16:47:33 +0000 Subject: [PATCH] Permissioned AaveWrapper --- src/aave/AaveWrapper.sol | 4 +- src/aave/PermissionedAaveWrapper.sol | 24 ++++ .../interfaces/IPoolAddressesProviderV3.sol | 2 + test/PermissionedAaveWrapper.v3.t.sol | 128 ++++++++++++++++++ 4 files changed, 156 insertions(+), 2 deletions(-) create mode 100644 src/aave/PermissionedAaveWrapper.sol create mode 100644 test/PermissionedAaveWrapper.v3.t.sol diff --git a/src/aave/AaveWrapper.sol b/src/aave/AaveWrapper.sol index b66d242..ded8884 100644 --- a/src/aave/AaveWrapper.sol +++ b/src/aave/AaveWrapper.sol @@ -72,7 +72,7 @@ contract AaveWrapper is BaseWrapper, IFlashLoanReceiverV2V3 { return true; } - function _flashLoan(address asset, uint256 amount, bytes memory data) internal override { + function _flashLoan(address asset, uint256 amount, bytes memory data) internal virtual override { IPool(POOL).flashLoan({ receiverAddress: address(this), assets: asset.toArray(), @@ -92,7 +92,7 @@ contract AaveWrapper is BaseWrapper, IFlashLoanReceiverV2V3 { max = !isFrozen && isActive && isFlashLoanEnabled ? IERC20(asset).balanceOf(aTokenAddress) : 0; } - function _flashFee(uint256 amount) internal view returns (uint256) { + function _flashFee(uint256 amount) internal view virtual returns (uint256) { return Math.mulDiv(amount, IPool(POOL).FLASHLOAN_PREMIUM_TOTAL() * 0.0001e18, WAD, Math.Rounding.Ceil); } } diff --git a/src/aave/PermissionedAaveWrapper.sol b/src/aave/PermissionedAaveWrapper.sol new file mode 100644 index 0000000..910d254 --- /dev/null +++ b/src/aave/PermissionedAaveWrapper.sol @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: MIT +// Thanks to ultrasecr.eth +pragma solidity ^0.8.19; + +import { Registry } from "lib/registry/src/Registry.sol"; +import { AccessControl } from "@openzeppelin/contracts/access/AccessControl.sol"; +import { AaveWrapper } from "./AaveWrapper.sol"; + +contract PermissionedAaveWrapper is AaveWrapper, AccessControl { + bytes32 public constant BORROWER = keccak256("BORROWER"); + + constructor(address owner, Registry reg, string memory name) AaveWrapper(reg, name) { + _grantRole(DEFAULT_ADMIN_ROLE, owner); + } + + function _flashLoan(address asset, uint256 amount, bytes memory data) internal override onlyRole(BORROWER) { + super._flashLoan(asset, amount, data); + } + + // This contract will be whitelisted in Aave so it pays 0 fees + function _flashFee(uint256) internal pure override returns (uint256) { + return 0; + } +} diff --git a/src/aave/interfaces/IPoolAddressesProviderV3.sol b/src/aave/interfaces/IPoolAddressesProviderV3.sol index a55fffc..1355904 100644 --- a/src/aave/interfaces/IPoolAddressesProviderV3.sol +++ b/src/aave/interfaces/IPoolAddressesProviderV3.sol @@ -7,4 +7,6 @@ interface IPoolAddressesProviderV3 { function getPool() external view returns (address); function getPoolDataProvider() external view returns (IPoolDataProvider); + + function getACLManager() external view returns (address); } diff --git a/test/PermissionedAaveWrapper.v3.t.sol b/test/PermissionedAaveWrapper.v3.t.sol new file mode 100644 index 0000000..9dafb18 --- /dev/null +++ b/test/PermissionedAaveWrapper.v3.t.sol @@ -0,0 +1,128 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity >=0.8.19 <0.9.0; + +import { PRBTest } from "@prb/test/PRBTest.sol"; +import { console2 } from "forge-std/console2.sol"; +import { StdCheats } from "forge-std/StdCheats.sol"; + +import { Registry } from "lib/registry/src/Registry.sol"; + +import { IERC20Metadata as IERC20 } from "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol"; +import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; + +import { MockBorrower } from "./MockBorrower.sol"; +import { IAccessControl } from "@openzeppelin/contracts/access/IAccessControl.sol"; +import { PermissionedAaveWrapper, AaveWrapper } from "../src/aave/PermissionedAaveWrapper.sol"; +import { Arrays } from "src/utils/Arrays.sol"; +import { IPoolAddressesProviderV3 } from "../src/aave/interfaces/IPoolAddressesProviderV3.sol"; + +/// @dev If this is your first time with Forge, read this tutorial in the Foundry Book: +/// https://book.getfoundry.sh/forge/writing-tests +contract PermissionedAaveWrapperTest is PRBTest, StdCheats { + using Arrays for *; + using SafeERC20 for IERC20; + + PermissionedAaveWrapper internal wrapper; + MockBorrower internal borrower; + address internal dai; + IPoolAddressesProviderV3 internal provider; + + address internal owner = makeAddr("owner"); + + /// @dev A function invoked before each test case is run. + function setUp() public virtual { + // Revert if there is no API key. + string memory alchemyApiKey = vm.envOr("API_KEY_ALCHEMY", string("")); + if (bytes(alchemyApiKey).length == 0) { + revert("API_KEY_ALCHEMY variable missing"); + } + + vm.createSelectFork({ urlOrAlias: "arbitrum_one", blockNumber: 98_674_994 }); + provider = IPoolAddressesProviderV3(0xa97684ead0e402dC232d5A977953DF7ECBaB3CDb); + dai = 0xDA10009cBd5D07dd0CeCc66161FC93D7c9000da1; + + Registry registry = new Registry(address(this)); + registry.set( + "AaveV3Wrapper", abi.encode(provider.getPool(), address(provider), provider.getPoolDataProvider(), false) + ); + + wrapper = new PermissionedAaveWrapper(owner, registry, "AaveV3"); + borrower = new MockBorrower(wrapper); + } + + /// @dev Basic test. Run it with `forge test -vvv` to see the console log. + function test_flashFee() external { + console2.log("test_flashFee"); + assertEq(wrapper.flashFee(dai, 1e18), 0, "Fee not right"); + } + + function test_maxFlashLoan() external { + console2.log("test_maxFlashLoan"); + assertEq(wrapper.maxFlashLoan(dai), 3_258_387.712396344524653246e18, "Max flash loan not right"); + } + + function test_flashLoan() external { + console2.log("test_flashLoan"); + bytes32 role = wrapper.BORROWER(); + vm.prank(owner); + wrapper.grantRole(role, address(borrower)); + + vm.mockCall( + provider.getACLManager(), + abi.encodeWithSelector(IACLManager.isFlashBorrower.selector, wrapper), + abi.encode(true) + ); + + uint256 loan = 1e18; + uint256 fee = wrapper.flashFee(dai, loan); + IERC20(dai).safeTransfer(address(borrower), fee); + bytes memory result = borrower.flashBorrow(dai, loan); + + // Test the return values passed through the wrapper + (bytes32 callbackReturn) = abi.decode(result, (bytes32)); + assertEq(uint256(callbackReturn), uint256(borrower.ERC3156PP_CALLBACK_SUCCESS()), "Callback failed"); + + // Test the borrower state during the callback + assertEq(borrower.flashInitiator(), address(borrower)); + assertEq(address(borrower.flashAsset()), address(dai)); + assertEq(borrower.flashAmount(), loan); + assertEq(borrower.flashBalance(), loan + fee); // The amount we transferred to pay for fees, plus the amount we + // borrowed + assertEq(borrower.flashFee(), fee); + } + + function test_executeOperation_permissions() public { + vm.expectRevert(AaveWrapper.NotPool.selector); + wrapper.executeOperation({ + assets: address(dai).toArray(), + amounts: 1e18.toArray(), + fees: 0.toArray(), + initiator: address(wrapper), + params: "" + }); + + vm.prank(provider.getPool()); + vm.expectRevert(AaveWrapper.NotInitiator.selector); + wrapper.executeOperation({ + assets: address(dai).toArray(), + amounts: 1e18.toArray(), + fees: 0.toArray(), + initiator: address(0x666), + params: "" + }); + } + + function test_flashLoan_permissions() external { + console2.log("test_flashLoan_permissions"); + vm.expectRevert( + abi.encodeWithSelector( + IAccessControl.AccessControlUnauthorizedAccount.selector, borrower, wrapper.BORROWER() + ) + ); + borrower.flashBorrow(dai, 1e18); + } +} + +interface IACLManager { + function isFlashBorrower(address) external view returns (bool); +}