fix: add allowlist for exchange venues in DepositWrapper

enzymefinance · Oct 11, 2023 · 3f40c0b · 3f40c0b
1 parent 32b84c8
commit 3f40c0b
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 29 deletions.
diff --git a/contracts/release/peripheral/DepositWrapper.sol b/contracts/release/peripheral/DepositWrapper.sol
@@ -14,20 +14,24 @@ pragma solidity 0.8.19;
 import {Address} from "openzeppelin-solc-0.8/utils/Address.sol";
 import {ERC20} from "openzeppelin-solc-0.8/token/ERC20/ERC20.sol";
 import {SafeERC20} from "openzeppelin-solc-0.8/token/ERC20/utils/SafeERC20.sol";
-import {IComptroller} from "../core/fund/comptroller/IComptroller.sol";
 import {IWETH} from "../../external-interfaces/IWETH.sol";
+import {IAddressListRegistry} from "../../persistent/address-list-registry/IAddressListRegistry.sol";
 import {AssetHelpers} from "../../utils/0.8.19/AssetHelpers.sol";
+import {IComptroller} from "../core/fund/comptroller/IComptroller.sol";
 
 /// @title DepositWrapper Contract
 /// @author Enzyme Council <[email protected]>
 /// @notice Logic related to wrapping deposit actions
 contract DepositWrapper is AssetHelpers {
     using SafeERC20 for ERC20;
 
-    bytes4 private constant BUY_SHARES_ON_BEHALF_SELECTOR = 0x877fd894;
+    IAddressListRegistry private immutable ADDRESS_LIST_REGISTRY;
+    uint256 private immutable ALLOWED_EXCHANGES_LIST_ID;
     IWETH private immutable WRAPPED_NATIVE_ASSET;
 
-    constructor(IWETH _wrappedNativeAsset) {
+    constructor(address _addressListRegistryAddress, uint256 _allowedExchangesListId, IWETH _wrappedNativeAsset) {
+        ADDRESS_LIST_REGISTRY = IAddressListRegistry(_addressListRegistryAddress);
+        ALLOWED_EXCHANGES_LIST_ID = _allowedExchangesListId;
         WRAPPED_NATIVE_ASSET = _wrappedNativeAsset;
     }
 
@@ -180,11 +184,10 @@ contract DepositWrapper is AssetHelpers {
         ERC20 _inputAsset,
         uint256 _maxInputAssetAmount
     ) private returns (uint256 sharesReceived_) {
-        // Deny access to privileged core calls originating from this contract
-        {
-            bytes4 exchangeSelector = bytes4(_exchangeData[:4]);
-            require(exchangeSelector != BUY_SHARES_ON_BEHALF_SELECTOR, "__exchangeAndBuyShares: Disallowed selector");
-        }
+        require(
+            ADDRESS_LIST_REGISTRY.isInList({_id: ALLOWED_EXCHANGES_LIST_ID, _item: _exchange}),
+            "__exchangeAndBuyShares: Unallowed _exchange"
+        );
 
         // Exchange the _inputAsset to the fund's denomination asset
         __approveAssetMaxAsNeeded({

diff --git a/tests/tests/peripheral/DepositWrapper.t.sol b/tests/tests/peripheral/DepositWrapper.t.sol
@@ -5,6 +5,7 @@ import {SafeERC20} from "openzeppelin-solc-0.8/token/ERC20/utils/SafeERC20.sol";
 
 import {IntegrationTest} from "tests/bases/IntegrationTest.sol";
 
+import {UpdateType as AddressListUpdateType} from "tests/utils/core/ListRegistryUtils.sol";
 import {
     ETHEREUM_SWAP_ROUTER as ETHEREUM_UNISWAP_ROUTER,
     POLYGON_SWAP_ROUTER as POLYGON_UNISWAP_ROUTER,
@@ -13,23 +14,33 @@ import {
 
 import {IERC20} from "tests/interfaces/external/IERC20.sol";
 
+import {IAddressListRegistry} from "tests/interfaces/internal/IAddressListRegistry.sol";
 import {IDepositWrapper} from "tests/interfaces/internal/IDepositWrapper.sol";
 import {IComptrollerLib} from "tests/interfaces/internal/IComptrollerLib.sol";
 import {IVaultLib} from "tests/interfaces/internal/IVaultLib.sol";
 
 abstract contract TestBase is IntegrationTest, UniswapV3Utils {
+    uint256 internal allowedExchangesListId;
+    IDepositWrapper internal depositWrapper;
+
     address internal sharesBuyer = makeAddr("SharesBuyer");
     address internal vaultOwner = makeAddr("VaultOwner");
     IComptrollerLib internal comptrollerProxy;
     IVaultLib internal vaultProxy;
 
     // Defined by parent contract
-    IDepositWrapper internal depositWrapper;
     IERC20 internal denominationAsset;
     address internal exchangeAddress;
     address internal exchangeApprovalTargetAddress;
 
     function setUp() public virtual override {
+        // Create an allowedExchanges list, with Uniswap router as the only allowed exchange
+        allowedExchangesListId = core.persistent.addressListRegistry.createList({
+            _owner: address(this),
+            _updateType: uint8(AddressListUpdateType.AddAndRemove),
+            _initialItems: toArray(exchangeAddress)
+        });
+
         // Deploy deposit wrapper
         depositWrapper = __deployDepositWrapper();
 
@@ -49,7 +60,9 @@ abstract contract TestBase is IntegrationTest, UniswapV3Utils {
     // DEPLOYMENT HELPERS
 
     function __deployDepositWrapper() internal returns (IDepositWrapper) {
-        bytes memory args = abi.encode(wrappedNativeToken);
+        require(allowedExchangesListId != 0, "__deployDepositWrapper: Unset listId");
+
+        bytes memory args = abi.encode(core.persistent.addressListRegistry, allowedExchangesListId, wrappedNativeToken);
 
         return IDepositWrapper(deployCode("DepositWrapper.sol", args));
     }
@@ -116,25 +129,6 @@ abstract contract ExchangeErc20AndBuySharesTest is TestBase {
         vm.stopPrank();
     }
 
-    function test_failWithDisallowedSelector() public {
-        bytes memory badExchangeData =
-            abi.encodeWithSelector(IComptrollerLib.buySharesOnBehalf.selector, sharesBuyer, inputAssetAmount, 1);
-
-        // Should fail with a disallowed selector
-        vm.expectRevert("__exchangeAndBuyShares: Disallowed selector");
-        vm.prank(sharesBuyer);
-        depositWrapper.exchangeErc20AndBuyShares({
-            _comptrollerProxy: address(comptrollerProxy),
-            _minSharesQuantity: minExpectedShares,
-            _inputAsset: address(inputAsset),
-            _maxInputAssetAmount: inputAssetAmount,
-            _exchange: exchangeAddress,
-            _exchangeApproveTarget: exchangeApprovalTargetAddress,
-            _exchangeData: badExchangeData,
-            _exchangeMinReceived: 0
-        });
-    }
-
     function test_failWithExchangeMinNotReceived() public {
         uint256 denominationAssetValue = core.release.valueInterpreter.calcCanonicalAssetValue({
             _baseAsset: address(inputAsset),
@@ -172,6 +166,24 @@ abstract contract ExchangeErc20AndBuySharesTest is TestBase {
         });
     }
 
+    function test_failWithUnallowedExchange() public {
+        address badExchange = makeAddr("BadExchange");
+
+        // Should fail with a disallowed selector
+        vm.expectRevert("__exchangeAndBuyShares: Unallowed _exchange");
+        vm.prank(sharesBuyer);
+        depositWrapper.exchangeErc20AndBuyShares({
+            _comptrollerProxy: address(comptrollerProxy),
+            _minSharesQuantity: minExpectedShares,
+            _inputAsset: address(inputAsset),
+            _maxInputAssetAmount: inputAssetAmount,
+            _exchange: badExchange,
+            _exchangeApproveTarget: exchangeApprovalTargetAddress,
+            _exchangeData: exchangeData,
+            _exchangeMinReceived: 0
+        });
+    }
+
     function test_success() public {
         vm.prank(sharesBuyer);
         depositWrapper.exchangeErc20AndBuyShares({