added method for withdrawing surplus

Author: d10r

packages/ethereum-contracts/contracts/gov/SuperfluidGovernanceBase.sol

@@ -141,13 +141,6 @@ abstract contract SuperfluidGovernanceBase is ISuperfluidGovernance
         host.changeSuperTokenAdmin(token, newAdmin);
     }
 
-    function setSuperTokenYieldBackend(ISuperfluid host, ISuperToken token, address yieldBackend)
-        external
-        onlyAuthorized(host)
-    {
-        token.setYieldBackend(yieldBackend);
-    }
-
     function batchChangeSuperTokenAdmin(ISuperfluid host, ISuperToken[] calldata token, address[] calldata newAdmins)
         external
         onlyAuthorized(host)

packages/ethereum-contracts/contracts/interfaces/superfluid/ISuperToken.sol

@@ -71,13 +71,6 @@ interface ISuperToken is ISuperfluidToken, IERC20Metadata, IERC777, IERC20Permit
      */
     function getAdmin() external view returns (address admin);
 
-    /**
-     * @notice Sets the yield backend for the SuperToken
-     * @dev Only the admin can call this function
-     * @param yieldBackend Address of the yield backend contract, or address(0) to disable the yield backend
-     */
-    function setYieldBackend(address yieldBackend) external;
-
     function getYieldBackend() external view returns (address yieldBackend);
 
     /**************************************************************************

packages/ethereum-contracts/contracts/interfaces/superfluid/IYieldBackend.sol

@@ -11,11 +11,19 @@ pragma solidity ^0.8.23;
  * one -> no means withdraw not in the context of a downgrade.
  */
 interface IYieldBackend {
-    function init() external;
-    function deinit() external;
+    /// Invoked by `SuperToken.enableYieldBackend()` as delegatecall.
+    /// Sets up the SuperToken as needed, e.g. by giving required approvals.
+    function enable() external;
+
+    /// Invoked by `SuperToken.disableYieldBackend()` as delegatecall.
+    /// Restores the prior state, e.g. by revoking given approvals
+    function disable() external;
 
     function deposit(uint256 amount) external;
     function depositMax() external;
     function withdraw(uint256 amount) external;
     function withdrawMax() external;
+
+    /// tranfers the deposited asset exceeding the required underlying to the preset treasury account
+    function withdrawSurplus(uint256 totalSupply) external;
 }

packages/ethereum-contracts/contracts/superfluid/AaveYieldBackend.sol

@@ -2,8 +2,8 @@
 pragma solidity ^0.8.23;
 
 import { IYieldBackend } from "../interfaces/superfluid/IYieldBackend.sol";
+import { IERC20, ISuperToken } from "../interfaces/superfluid/ISuperfluid.sol";
 import { IPool } from "aave-v3/interfaces/IPool.sol";
-import { IERC20 } from "../interfaces/superfluid/ISuperfluid.sol";
 
 
 /**
@@ -17,6 +17,8 @@ contract AaveYieldBackend is IYieldBackend {
     IERC20 public immutable ASSET_TOKEN;
     IPool public immutable AAVE_POOL;
     IERC20 public immutable A_TOKEN;
+    // TODO: make an immutable
+    address constant SURPLUS_RECEIVER = 0xac808840f02c47C05507f48165d2222FF28EF4e1; // dao.superfluid.eth
 
     // THIS CONTRACT CANNOT HAVE STATE VARIABLES!
     // IF STATE IS NEEDED, USE NAMESPACED STORAGE LAYOUT (EIP-7201)
@@ -33,12 +35,12 @@ contract AaveYieldBackend is IYieldBackend {
         A_TOKEN = IERC20(aavePool.getReserveAToken(address(assetToken)));
     }
 
-    function init() external {
+    function enable() external {
         // approve Aave pool to fetch asset
         ASSET_TOKEN.approve(address(AAVE_POOL), type(uint256).max);
     }
 
-    function deinit() external {
+    function disable() external {
         // Revoke approval
         ASSET_TOKEN.approve(address(AAVE_POOL), 0);
     }
@@ -66,4 +68,12 @@ contract AaveYieldBackend is IYieldBackend {
         // We can delegate the max calculation to the Aave pool by setting amount to type(uint256).max
         AAVE_POOL.withdraw(address(ASSET_TOKEN), type(uint256).max, address(this));
     }
+
+    function withdrawSurplus(uint256 totalSupply) external {
+        // totalSupply is always 18 decimals while assetToken and aToken may not
+        (uint256 normalizedTotalSupply,) = ISuperToken(address(this)).toUnderlyingAmount(totalSupply);
+        // decrement by 1 in order to offset Aave's rounding up
+        uint256 surplusAmount = A_TOKEN.balanceOf(address(this)) - normalizedTotalSupply - 1;
+        AAVE_POOL.withdraw(address(ASSET_TOKEN), surplusAmount, SURPLUS_RECEIVER);
+    }
 }

packages/ethereum-contracts/contracts/superfluid/SuperToken.sol

@@ -198,20 +198,11 @@ contract SuperToken is
         }
     }
 
-    function setYieldBackend(address newYieldBackend) external onlyAdmin {
-        if (address(_yieldBackend) != address(0)) {
-            _disableYieldBackend();
-        }
-        if (address(newYieldBackend) != address(0)) {
-            _enableYieldBackend(IYieldBackend(newYieldBackend));
-        }
-    }
-
-    function _enableYieldBackend(IYieldBackend newYieldBackend) internal {
+    function enableYieldBackend(IYieldBackend newYieldBackend) external onlyAdmin {
         require(address(_yieldBackend) == address(0), "yield backend already set");
         _yieldBackend = newYieldBackend;
         (bool success, ) = address(_yieldBackend).delegatecall(
-            abi.encodeCall(IYieldBackend.initSuperToken, ())
+            abi.encodeCall(IYieldBackend.enable, ())
         );
         require(success, "delegatecall failed");
         (success, ) = address(_yieldBackend).delegatecall(
@@ -222,13 +213,14 @@ contract SuperToken is
     }
 
     // withdraws everything and removes allowances
-    function _disableYieldBackend() internal {
+    function disableYieldBackend() external onlyAdmin {
+        require(address(_yieldBackend) != address(0), "yield backend not set");
         (bool success, ) = address(_yieldBackend).delegatecall(
             abi.encodeCall(IYieldBackend.withdrawMax, ())
         );
         require(success, "delegatecall failed");
         (success, ) = address(_yieldBackend).delegatecall(
-            abi.encodeCall(IYieldBackend.deinitSuperToken, ())
+            abi.encodeCall(IYieldBackend.disable, ())
         );
         // TODO: should this be allowed to fail?
         require(success, "delegatecall failed");
@@ -240,6 +232,14 @@ contract SuperToken is
         return address(_yieldBackend);
     }
 
+    function withdrawSurplusFromYieldBackend() external onlyAdmin {
+        require(address(_yieldBackend) != address(0), "yield backend not set");
+        (bool success, ) = address(_yieldBackend).delegatecall(
+            abi.encodeCall(IYieldBackend.withdrawSurplus, (_totalSupply))
+        );
+        require(success, "delegatecall failed");
+    }
+
     /**************************************************************************
      * ERC20 Token Info
      *************************************************************************/

packages/ethereum-contracts/contracts/superfluid/Superfluid.sol

@@ -19,7 +19,6 @@ import {
     ISuperTokenFactory,
     IAccessControl
 } from "../interfaces/superfluid/ISuperfluid.sol";
-import { IYieldBackend } from "../interfaces/superfluid/IYieldBackend.sol";
 import { GeneralDistributionAgreementV1 } from "../agreements/gdav1/GeneralDistributionAgreementV1.sol";
 import { SuperfluidUpgradeableBeacon } from "../upgradability/SuperfluidUpgradeableBeacon.sol";
 import { CallUtils } from "../libs/CallUtils.sol";
@@ -333,10 +332,6 @@ contract Superfluid is
         token.changeAdmin(newAdmin);
     }
 
-    function setSuperTokenYieldBackend(ISuperToken token, address yieldBackend) external onlyGovernance {
-        token.setYieldBackend(yieldBackend);
-    }
-
     ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
     // Superfluid Upgradeable Beacon
     ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

packages/ethereum-contracts/test/foundry/superfluid/SuperTokenYield.t.sol

@@ -15,6 +15,7 @@ import { IPool } from "aave-v3/interfaces/IPool.sol";
  */
 contract SuperTokenYieldForkTest is Test {
     address constant ALICE = address(0x420);
+    address constant ADMIN = address(0xAAA);
 
     // Base network constants
     uint256 internal constant CHAIN_ID = 8453;
@@ -37,11 +38,6 @@ contract SuperTokenYieldForkTest is Test {
     /// @notice Aave V3 Pool contract
     IPool public aavePool;
 
-    /// @notice Admin address (this contract)
-    address public admin;
-    /// @notice Test user address
-    address public user;
-    
     /// @notice Set up the test environment by forking Base and deploying AaveYieldBackend
     function setUp() public {
         // Fork Base using public RPC
@@ -50,10 +46,6 @@ contract SuperTokenYieldForkTest is Test {
         // Verify we're on Base
         assertEq(block.chainid, CHAIN_ID, "Chainid mismatch");
 
-        // Initialize test accounts
-        admin = address(this);
-        user = address(0x1234);
-        
         // Get Aave Pool
         aavePool = IPool(AAVE_POOL);
 
@@ -77,14 +69,33 @@ contract SuperTokenYieldForkTest is Test {
         superToken.updateCode(address(newSuperTokenLogic));
         vm.stopPrank();
 
+        // designate an admin for the SuperToken
+        vm.startPrank(address(superToken.getHost()));
+        superToken.changeAdmin(ADMIN);
+        vm.stopPrank();
+
+        // provide ALICE with underlying and let her approve for upgrade
+        deal(USDC, ALICE, type(uint128).max);
+        vm.startPrank(ALICE);
+        IERC20(USDC).approve(address(superToken), type(uint256).max);
+
         console.log("aaveBackend address", address(aaveBackend));
     }
 
     function _enableYieldBackend() public {
-        vm.startPrank(address(superToken.getHost()));
-        superToken.setYieldBackend(address(aaveBackend));
+        vm.startPrank(ADMIN);
+        superToken.enableYieldBackend(aaveBackend);
         vm.stopPrank();
     }
+
+    function _verifyInvariants() internal view {
+        // underlyingBalance + aTokenBalance >= superToken.supply()
+        uint256 underlyingBalance = IERC20(USDC).balanceOf(address(superToken));
+        uint256 aTokenBalance = IERC20(aUSDC).balanceOf(address(superToken));
+        (uint256 superTokenNormalizedSupply,) = superToken.toUnderlyingAmount(superToken.totalSupply());
+
+        assertGe(underlyingBalance + aTokenBalance, superTokenNormalizedSupply, "invariant failed: underlyingBalance + aTokenBalance insufficient");
+    }
 
     /// @notice Test that we're forking the correct Base network
     function testForkBaseNetwork() public view {
@@ -116,37 +127,30 @@ contract SuperTokenYieldForkTest is Test {
         console.log("aUSDC balance of SuperToken", IERC20(aUSDC).balanceOf(address(superToken)));
         // TODO: We'd want asset balance to equal aToken balance. But that's not exactly the case.
         // what else shall be require?
+        _verifyInvariants();
     }
 
     function testDisableYieldBackend() public {
-        // store underlying balance before enabling yield backend
-        uint256 underlyingBalanceBefore = IERC20(USDC).balanceOf(address(superToken));
-
         _enableYieldBackend();
 
-        vm.startPrank(address(superToken.getHost()));
-        superToken.setYieldBackend(address(0));
+        vm.startPrank(ADMIN);
+        superToken.disableYieldBackend();
         vm.stopPrank();
         assertEq(address(superToken.getYieldBackend()), address(0), "Yield backend mismatch");
 
         // the SuperToken should now have a non-zero USDC balance and a zero aUSDC balance
         assertGt(IERC20(USDC).balanceOf(address(superToken)), 0, "USDC balance should be non-zero");
         assertEq(IERC20(aUSDC).balanceOf(address(superToken)), 0, "aUSDC balance should be zero");
 
-        // get underlying balance after disabling yield backend
-        uint256 underlyingBalanceAfter = IERC20(USDC).balanceOf(address(superToken));
-        //assertEq(underlyingBalanceAfter, underlyingBalanceBefore, "Underlying balance should be the same");
+        _verifyInvariants();
     }
 
     // TODO: bool fuzz arg for disabled/enabled backend
     function testUpgradeDowngrade() public {
         _enableYieldBackend();
 
-        deal(USDC, ALICE, 1000 ether);
-
         uint256 aTokenBalanceBefore = IERC20(aUSDC).balanceOf(address(superToken));
         vm.startPrank(ALICE);
-        IERC20(USDC).approve(address(superToken), type(uint256).max);
         superToken.upgrade(1 ether);
         vm.stopPrank();
 
@@ -167,7 +171,7 @@ contract SuperTokenYieldForkTest is Test {
         superToken.downgrade(1 ether);
         vm.stopPrank();
 
-        uint256 aTokenBalanceAfterDowngrade = IERC20(aUSDC).balanceOf(address(superToken));
+        _verifyInvariants();
     }
 
     // ============ Gas Benchmarking Tests ============
@@ -176,19 +180,13 @@ contract SuperTokenYieldForkTest is Test {
     /// @dev Separate test function to avoid cold/warm storage slot interference
     function testGasUpgrade_WithoutYieldBackend() public {
         // Ensure yield backend is NOT set
-        vm.startPrank(address(superToken.getHost()));
-        superToken.setYieldBackend(address(0));
-        vm.stopPrank();
         assertEq(address(superToken.getYieldBackend()), address(0), "Yield backend should not be set");
 
         // Prepare test state
         // 1000 USDC = 1000 * 1e6 (USDC has 6 decimals)
         // In SuperToken units (18 decimals), this is 1000 * 1e18
         uint256 upgradeAmount = 1000 * 1e18;
-        deal(USDC, ALICE, 1000 * 1e6);
         vm.startPrank(ALICE);
-        IERC20(USDC).approve(address(superToken), type(uint256).max);
-        
         // Measure gas for upgrade
         uint256 gasBefore = gasleft();
         superToken.upgrade(upgradeAmount);
@@ -211,10 +209,7 @@ contract SuperTokenYieldForkTest is Test {
         // 1000 USDC = 1000 * 1e6 (USDC has 6 decimals)
         // In SuperToken units (18 decimals), this is 1000 * 1e18
         uint256 upgradeAmount = 1000 * 1e18;
-        deal(USDC, ALICE, 1000 * 1e6);
         vm.startPrank(ALICE);
-        IERC20(USDC).approve(address(superToken), type(uint256).max);
-        
         // Measure gas for upgrade
         uint256 gasBefore = gasleft();
         superToken.upgrade(upgradeAmount);
@@ -230,17 +225,13 @@ contract SuperTokenYieldForkTest is Test {
     /// @dev Separate test function to avoid cold/warm storage slot interference
     function testGasDowngrade_WithoutYieldBackend() public {
         // Ensure yield backend is NOT set
-        vm.startPrank(address(superToken.getHost()));
-        superToken.setYieldBackend(address(0));
-        vm.stopPrank();
+        assertEq(address(superToken.getYieldBackend()), address(0), "Yield backend should not be set");
 
         // First, upgrade some tokens for ALICE to downgrade later
         // 1000 USDC = 1000 * 1e6 (USDC has 6 decimals)
         // In SuperToken units (18 decimals), this is 1000 * 1e18
         uint256 initialUpgradeAmount = 1000 * 1e18;
-        deal(USDC, ALICE, 1000 * 1e6);
         vm.startPrank(ALICE);
-        IERC20(USDC).approve(address(superToken), type(uint256).max);
         superToken.upgrade(initialUpgradeAmount);
         vm.stopPrank();
 
@@ -270,9 +261,7 @@ contract SuperTokenYieldForkTest is Test {
         // 1000 USDC = 1000 * 1e6 (USDC has 6 decimals)
         // In SuperToken units (18 decimals), this is 1000 * 1e18
         uint256 initialUpgradeAmount = 1000 * 1e18;
-        deal(USDC, ALICE, 1000 * 1e6);
         vm.startPrank(ALICE);
-        IERC20(USDC).approve(address(superToken), type(uint256).max);
         superToken.upgrade(initialUpgradeAmount);
         vm.stopPrank();
 
@@ -291,5 +280,45 @@ contract SuperTokenYieldForkTest is Test {
         console.log("Gas used", gasUsed);
         console.log("Amount downgraded", amountToDowngrade);
     }
+
+    function testWithdrawSurplusFromYieldBackend() public {
+        address SURPLUS_RECEIVER = 0xac808840f02c47C05507f48165d2222FF28EF4e1;
+        
+        // Simulate yield accumulation by transferring extra underlying to SuperToken
+        uint256 surplusAmount = 100 * 1e6; // 100 USDC
+        deal(USDC, address(this), surplusAmount);
+
+        _enableYieldBackend();
+        
+        // Upgrade tokens to create supply
+        uint256 upgradeAmount = 1000 * 1e18;
+        vm.startPrank(ALICE);
+        superToken.upgrade(upgradeAmount);
+        vm.stopPrank();
+
+        uint256 receiverBalanceBefore = IERC20(USDC).balanceOf(SURPLUS_RECEIVER);
+        uint256 aTokenBalanceBefore = IERC20(aUSDC).balanceOf(address(superToken));
+
+        // log USDC and aUSDC balances of SuperToken
+        console.log("USDC balance of SuperToken", IERC20(USDC).balanceOf(address(superToken)));
+        console.log("aUSDC balance of SuperToken", IERC20(aUSDC).balanceOf(address(superToken)));
+        // log normalized total supply
+        (uint256 normalizedTotalSupply, uint256 adjustedAmount) = superToken.toUnderlyingAmount(superToken.totalSupply());
+        console.log("normalized total supply", normalizedTotalSupply);
+        console.log("adjusted amount", adjustedAmount);
+        
+        vm.startPrank(ADMIN);
+        superToken.withdrawSurplusFromYieldBackend();
+        vm.stopPrank();
+        
+        uint256 receiverBalanceAfter = IERC20(USDC).balanceOf(SURPLUS_RECEIVER);
+        uint256 aTokenBalanceAfter = IERC20(aUSDC).balanceOf(address(superToken));
+        console.log("aToken balance after", aTokenBalanceAfter);
+        console.log("aToken balance diff", aTokenBalanceBefore - aTokenBalanceAfter);
+        
+        assertGt(receiverBalanceAfter, receiverBalanceBefore, "Surplus should be withdrawn to receiver");
+        assertLt(aTokenBalanceAfter, aTokenBalanceBefore, "aToken balance should decrease");
+        _verifyInvariants();
+    }
 }