Permissioned AaveWrapper

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);
     }
 }

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;
+    }
+}

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);
 }

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);
+}