feat: Treasury spoke compatible with multiple hubs

src/spoke/TreasurySpoke.sol

@@ -7,52 +7,45 @@ import {SafeERC20} from 'src/dependencies/openzeppelin/SafeERC20.sol';
 import {IERC20} from 'src/dependencies/openzeppelin/IERC20.sol';
 import {MathUtils} from 'src/libraries/math/MathUtils.sol';
 import {IHubBase} from 'src/hub/interfaces/IHubBase.sol';
-import {ITreasurySpoke, ISpokeBase} from 'src/spoke/interfaces/ITreasurySpoke.sol';
+import {ITreasurySpoke} from 'src/spoke/interfaces/ITreasurySpoke.sol';
 
 /// @title TreasurySpoke
 /// @author Aave Labs
 /// @notice Spoke contract used as a treasury where accumulated fees are treated as supplied assets.
 /// @dev Dedicated to a single user, controlled exclusively by the owner.
-/// @dev Utilizes all assets from the Hub without restrictions, making reserve and asset identifiers aligned.
-/// @dev Allows withdraw to claim fees and supply to invest back into the Hub via this dedicated spoke.
+/// @dev Allows withdraw to claim fees and supply to invest back into the corresponding hub via this dedicated spoke.
 contract TreasurySpoke is ITreasurySpoke, Ownable2Step {
   using SafeERC20 for IERC20;
 
-  /// @inheritdoc ITreasurySpoke
-  IHubBase public immutable HUB;
-
   /// @dev Constructor.
   /// @param owner_ The address of the owner.
-  /// @param hub_ The address of the Hub.
-  constructor(address owner_, address hub_) Ownable(owner_) {
-    require(hub_ != address(0), InvalidAddress());
-
-    HUB = IHubBase(hub_);
-  }
+  constructor(address owner_) Ownable(owner_) {}
 
   /// @inheritdoc ITreasurySpoke
   function supply(
-    uint256 reserveId,
+    address hub,
+    uint256 assetId,
     uint256 amount,
     address
   ) external onlyOwner returns (uint256, uint256) {
-    uint256 shares = HUB.add(reserveId, amount, msg.sender);
+    uint256 shares = IHubBase(hub).add(assetId, amount, msg.sender);
 
     return (shares, amount);
   }
 
   /// @inheritdoc ITreasurySpoke
   function withdraw(
-    uint256 reserveId,
+    address hub,
+    uint256 assetId,
     uint256 amount,
     address
   ) external onlyOwner returns (uint256, uint256) {
     // if amount to withdraw is greater than total supplied, withdraw all supplied assets
     uint256 withdrawnAmount = MathUtils.min(
       amount,
-      HUB.getSpokeAddedAssets(reserveId, address(this))
+      IHubBase(hub).getSpokeAddedAssets(assetId, address(this))
     );
-    uint256 withdrawnShares = HUB.remove(reserveId, withdrawnAmount, msg.sender);
+    uint256 withdrawnShares = IHubBase(hub).remove(assetId, withdrawnAmount, msg.sender);
 
     return (withdrawnShares, withdrawnAmount);
   }
@@ -63,55 +56,12 @@ contract TreasurySpoke is ITreasurySpoke, Ownable2Step {
   }
 
   /// @inheritdoc ITreasurySpoke
-  function getSuppliedAmount(uint256 reserveId) external view returns (uint256) {
-    return HUB.getSpokeAddedAssets(reserveId, address(this));
+  function getSuppliedAmount(address hub, uint256 assetId) external view returns (uint256) {
+    return IHubBase(hub).getSpokeAddedAssets(assetId, address(this));
   }
 
   /// @inheritdoc ITreasurySpoke
-  function getSuppliedShares(uint256 reserveId) external view returns (uint256) {
-    return HUB.getSpokeAddedShares(reserveId, address(this));
-  }
-
-  /// @inheritdoc ISpokeBase
-  function borrow(uint256, uint256, address) external pure returns (uint256, uint256) {
-    revert UnsupportedAction();
-  }
-
-  /// @inheritdoc ISpokeBase
-  function repay(uint256, uint256, address) external pure returns (uint256, uint256) {
-    revert UnsupportedAction();
-  }
-
-  /// @inheritdoc ISpokeBase
-  function liquidationCall(uint256, uint256, address, uint256, bool) external pure {
-    revert UnsupportedAction();
+  function getSuppliedShares(address hub, uint256 assetId) external view returns (uint256) {
+    return IHubBase(hub).getSpokeAddedShares(assetId, address(this));
   }
-
-  /// @inheritdoc ISpokeBase
-  function getUserDebt(uint256, address) external pure returns (uint256, uint256) {}
-
-  /// @inheritdoc ISpokeBase
-  function getUserTotalDebt(uint256, address) external pure returns (uint256) {}
-
-  /// @inheritdoc ISpokeBase
-  function getReserveSuppliedAssets(uint256 reserveId) external view returns (uint256) {
-    return HUB.getSpokeAddedAssets(reserveId, address(this));
-  }
-
-  /// @inheritdoc ISpokeBase
-  function getReserveSuppliedShares(uint256 reserveId) external view returns (uint256) {
-    return HUB.getSpokeAddedShares(reserveId, address(this));
-  }
-
-  /// @inheritdoc ISpokeBase
-  function getUserSuppliedAssets(uint256, address) external pure returns (uint256) {}
-
-  /// @inheritdoc ISpokeBase
-  function getUserSuppliedShares(uint256, address) external pure returns (uint256) {}
-
-  /// @inheritdoc ISpokeBase
-  function getReserveDebt(uint256) external pure returns (uint256, uint256) {}
-
-  /// @inheritdoc ISpokeBase
-  function getReserveTotalDebt(uint256) external pure returns (uint256) {}
 }

src/spoke/interfaces/ITreasurySpoke.sol

@@ -3,12 +3,11 @@
 pragma solidity ^0.8.0;
 
 import {IHubBase} from 'src/hub/interfaces/IHubBase.sol';
-import {ISpokeBase} from 'src/spoke/interfaces/ISpokeBase.sol';
 
 /// @title ITreasurySpoke
 /// @author Aave Labs
 /// @notice Interface for the TreasurySpoke.
-interface ITreasurySpoke is ISpokeBase {
+interface ITreasurySpoke {
   /// @notice Thrown when an unsupported action is attempted.
   error UnsupportedAction();
 
@@ -17,28 +16,32 @@ interface ITreasurySpoke is ISpokeBase {
 
   /// @notice Supplies a specified amount of the underlying asset to a given reserve.
   /// @dev The hub pulls the underlying asset from the caller, so prior approval is required.
-  /// @dev The reserve identifier must match the asset identifier in the hub.
-  /// @param reserveId The identifier of the reserve.
+  /// @dev The given asset identifier must match the asset identifier in the given hub.
+  /// @param hub The address of the hub contract.
+  /// @param assetId The identifier of the asset.
   /// @param amount The amount of asset to supply.
   /// @param onBehalfOf Unused parameter for this spoke.
   /// @return The amount of shares supplied.
   /// @return The amount of assets supplied.
   function supply(
-    uint256 reserveId,
+    address hub,
+    uint256 assetId,
     uint256 amount,
     address onBehalfOf
   ) external returns (uint256, uint256);
 
   /// @notice Withdraws a specified amount of underlying asset from the given reserve.
   /// @dev Providing an amount greater than the maximum withdrawable value signals a full withdrawal.
-  /// @dev The reserve identifier must match the asset identifier in the hub.
-  /// @param reserveId The identifier of the reserve.
+  /// @dev The given asset identifier must match the asset identifier in the given hub.
+  /// @param hub The address of the hub contract.
+  /// @param assetId The identifier of the asset.
   /// @param amount The amount of asset to withdraw.
   /// @param onBehalfOf Unused parameter for this spoke.
   /// @return The amount of shares withdrawn.
   /// @return The amount of assets withdrawn.
   function withdraw(
-    uint256 reserveId,
+    address hub,
+    uint256 assetId,
     uint256 amount,
     address onBehalfOf
   ) external returns (uint256, uint256);
@@ -50,19 +53,17 @@ interface ITreasurySpoke is ISpokeBase {
   function transfer(address token, address to, uint256 amount) external;
 
   /// @notice Returns the amount of assets supplied.
-  /// @dev The reserve identifier must match the asset identifier in the hub.
-  /// @param reserveId The identifier of the reserve.
+  /// @dev The given asset identifier must match the asset identifier in the given hub.
+  /// @param hub The address of the hub contract.
+  /// @param assetId The identifier of the asset.
   /// @return The amount of assets supplied.
-  function getSuppliedAmount(uint256 reserveId) external view returns (uint256);
+  function getSuppliedAmount(address hub, uint256 assetId) external view returns (uint256);
 
   /// @notice Returns the amount of shares supplied.
   /// @dev Shares are denominated relative to the supply side.
-  /// @dev The reserve identifier must match the asset identifier in the hub.
-  /// @param reserveId The identifier of the reserve.
+  /// @dev The asset identifier must match the asset identifier in the hub.
+  /// @param hub The address of the hub contract.
+  /// @param assetId The identifier of the asset.
   /// @return The amount of shares supplied.
-  function getSuppliedShares(uint256 reserveId) external view returns (uint256);
-
-  /// @notice Returns the interface of the associated hub.
-  /// @return The HubBase interface.
-  function HUB() external view returns (IHubBase);
+  function getSuppliedShares(address hub, uint256 assetId) external view returns (uint256);
 }

tests/Base.t.sol

@@ -283,7 +283,7 @@ abstract contract Base is Test {
     (spoke1, oracle1) = _deploySpokeWithOracle(ADMIN, address(accessManager), 'Spoke 1 (USD)');
     (spoke2, oracle2) = _deploySpokeWithOracle(ADMIN, address(accessManager), 'Spoke 2 (USD)');
     (spoke3, oracle3) = _deploySpokeWithOracle(ADMIN, address(accessManager), 'Spoke 3 (USD)');
-    treasurySpoke = ITreasurySpoke(new TreasurySpoke(TREASURY_ADMIN, address(hub1)));
+    treasurySpoke = ITreasurySpoke(new TreasurySpoke(TREASURY_ADMIN));
     dai = new MockERC20();
     eth = new MockERC20();
     usdc = new MockERC20();
@@ -1940,7 +1940,7 @@ abstract contract Base is Test {
       return; // nothing to withdraw
     }
     vm.prank(TREASURY_ADMIN);
-    treasurySpoke.withdraw(assetId, amount, address(treasurySpoke));
+    treasurySpoke.withdraw(address(hub1), assetId, amount, address(treasurySpoke));
   }
 
   function _assumeValidSupplier(address user) internal view {

tests/invariant/HubHandler.t.sol

@@ -59,7 +59,7 @@ contract HubHandler is Test {
       )
     );
     assertEq(address(spoke1), predictedSpoke, 'predictedSpoke');
-    treasurySpoke = new TreasurySpoke(hubAdmin, address(hub1));
+    treasurySpoke = new TreasurySpoke(hubAdmin);
     usdc = new MockERC20();
     dai = new MockERC20();
     usdt = new MockERC20();

tests/unit/HubConfigurator.t.sol

@@ -350,11 +350,11 @@ contract HubConfiguratorTest is HubBase {
       365 days
     );
 
-    assertGe(treasurySpoke.getSuppliedShares(daiAssetId), 0);
-    uint256 fees = treasurySpoke.getSuppliedAmount(daiAssetId);
+    assertGe(treasurySpoke.getSuppliedShares(address(hub1), daiAssetId), 0);
+    uint256 fees = treasurySpoke.getSuppliedAmount(address(hub1), daiAssetId);
 
     // Change the fee receiver
-    TreasurySpoke newTreasurySpoke = new TreasurySpoke(HUB_ADMIN, address(hub1));
+    TreasurySpoke newTreasurySpoke = new TreasurySpoke(HUB_ADMIN);
     vm.prank(HUB_CONFIGURATOR_ADMIN);
     hubConfigurator.updateFeeReceiver(address(hub1), daiAssetId, address(newTreasurySpoke));
 
@@ -369,25 +369,28 @@ contract HubConfiguratorTest is HubBase {
     );
 
     // Withdraw fees from the old treasury spoke
-    Utils.withdraw(
-      ISpoke(address(treasurySpoke)),
-      daiAssetId,
-      TREASURY_ADMIN,
-      fees,
-      address(treasurySpoke)
-    );
+    vm.prank(TREASURY_ADMIN);
+    treasurySpoke.withdraw(address(hub1), daiAssetId, fees, address(TREASURY_ADMIN));
 
-    assertEq(treasurySpoke.getSuppliedAmount(daiAssetId), 0, 'old treasury spoke should be empty');
+    assertEq(
+      treasurySpoke.getSuppliedAmount(address(hub1), daiAssetId),
+      0,
+      'old treasury spoke should be empty'
+    );
 
     // Accrue more fees, this time to new fee receiver
     skip(365 days);
 
     assertGt(
-      newTreasurySpoke.getSuppliedAmount(daiAssetId),
+      newTreasurySpoke.getSuppliedAmount(address(hub1), daiAssetId),
       0,
       'new fee receiver should have accrued fees'
     );
-    assertEq(treasurySpoke.getSuppliedAmount(daiAssetId), 0, 'old fee receiver should be empty');
+    assertEq(
+      treasurySpoke.getSuppliedAmount(address(hub1), daiAssetId),
+      0,
+      'old fee receiver should be empty'
+    );
   }
 
   /// @dev Test update fee receiver and old fee receiver still accrues fees
@@ -412,11 +415,11 @@ contract HubConfiguratorTest is HubBase {
       365 days
     );
 
-    assertGe(treasurySpoke.getSuppliedShares(daiAssetId), 0);
-    uint256 feeShares = treasurySpoke.getSuppliedShares(daiAssetId);
+    assertGe(treasurySpoke.getSuppliedShares(address(hub1), daiAssetId), 0);
+    uint256 feeShares = treasurySpoke.getSuppliedShares(address(hub1), daiAssetId);
 
     // Change the fee receiver
-    TreasurySpoke newTreasurySpoke = new TreasurySpoke(HUB_ADMIN, address(hub1));
+    TreasurySpoke newTreasurySpoke = new TreasurySpoke(HUB_ADMIN);
     vm.prank(HUB_CONFIGURATOR_ADMIN);
     hubConfigurator.updateFeeReceiver(address(hub1), daiAssetId, address(newTreasurySpoke));
 
@@ -434,41 +437,41 @@ contract HubConfiguratorTest is HubBase {
     );
 
     // Withdraw half the fee shares from the old treasury spoke
-    Utils.withdraw(
-      ISpoke(address(treasurySpoke)),
+    vm.startPrank(TREASURY_ADMIN);
+    treasurySpoke.withdraw(
+      address(hub1),
       daiAssetId,
-      TREASURY_ADMIN,
       hub1.previewRemoveByShares(daiAssetId, feeShares / 2),
-      address(treasurySpoke)
+      address(TREASURY_ADMIN)
     );
+    vm.stopPrank();
 
     // Get the remaining fee shares
-    feeShares = treasurySpoke.getSuppliedShares(daiAssetId);
+    feeShares = treasurySpoke.getSuppliedShares(address(hub1), daiAssetId);
 
     // Accrue more fees, this time to new fee receiver
     skip(365 days);
 
     // Check that new fee receiver is getting the fees, and not old treasury spoke
     assertGt(
-      newTreasurySpoke.getSuppliedAmount(daiAssetId),
+      newTreasurySpoke.getSuppliedAmount(address(hub1), daiAssetId),
       0,
       'new fee receiver should have accrued fees'
     );
     assertEq(
-      treasurySpoke.getSuppliedShares(daiAssetId),
+      treasurySpoke.getSuppliedShares(address(hub1), daiAssetId),
       feeShares,
       'old fee receiver should still have same share amount'
     );
 
     // Now withdraw remaining fee shares from old treasury spoke
-    Utils.withdraw(
-      ISpoke(address(treasurySpoke)),
-      daiAssetId,
-      TREASURY_ADMIN,
-      UINT256_MAX,
-      address(treasurySpoke)
+    vm.prank(TREASURY_ADMIN);
+    treasurySpoke.withdraw(address(hub1), daiAssetId, UINT256_MAX, address(TREASURY_ADMIN));
+    assertEq(
+      treasurySpoke.getSuppliedShares(address(hub1), daiAssetId),
+      0,
+      'old fee receiver should be empty'
     );
-    assertEq(treasurySpoke.getSuppliedShares(daiAssetId), 0, 'old fee receiver should be empty');
   }
 
   function test_updateFeeReceiver_Scenario() public {

tests/unit/Spoke/Spoke.MultipleHub.Base.t.sol

@@ -39,7 +39,7 @@ contract SpokeMultipleHubBase is SpokeBase {
     // Canonical hub and spoke
     hub1 = new Hub(address(accessManager));
     (spoke1, oracle1) = _deploySpokeWithOracle(ADMIN, address(accessManager), 'Spoke 1 (USD)');
-    treasurySpoke = new TreasurySpoke(ADMIN, address(hub1));
+    treasurySpoke = new TreasurySpoke(ADMIN);
     irStrategy = new AssetInterestRateStrategy(address(hub1));
 
     // New hub and spoke