[VPD-808] Add internal cash to prevent donation attack

contracts/Tokens/VTokens/VBep20.sol

@@ -4,7 +4,6 @@ pragma solidity 0.8.25;
 
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
-
 import { ComptrollerInterface } from "../../Comptroller/ComptrollerInterface.sol";
 import { InterestRateModelV8 } from "../../InterestRateModels/InterestRateModelV8.sol";
 import { VBep20Interface, VTokenInterface } from "./VTokenInterfaces.sol";
@@ -223,8 +222,10 @@ contract VBep20 is VToken, VBep20Interface {
         uint256 balanceBefore = token.balanceOf(address(this));
         token.safeTransferFrom(from, address(this), amount);
         uint256 balanceAfter = token.balanceOf(address(this));
+        uint256 actualAmount = balanceAfter - balanceBefore;
+        internalCash += actualAmount;
         // Return the amount that was *actually* transferred
-        return balanceAfter - balanceBefore;
+        return actualAmount;
     }
 
     /**
@@ -233,6 +234,7 @@ contract VBep20 is VToken, VBep20Interface {
      * @param amount Amount of underlying to transfer
      */
     function doTransferOut(address payable to, uint256 amount) internal virtual override {
+        internalCash -= amount;
         IERC20 token = IERC20(underlying);
         token.safeTransfer(to, amount);
     }
@@ -243,6 +245,61 @@ contract VBep20 is VToken, VBep20Interface {
      * @return The quantity of underlying tokens owned by this contract
      */
     function getCashPrior() internal view override returns (uint) {
-        return IERC20(underlying).balanceOf(address(this));
+        return internalCash;
+    }
+
+    /**
+     * @notice Admin function to sync internalCash with actual token balance
+     * @dev Used for one-time migration after upgrading existing deployed markets
+     */
+    function syncCash() external {
+        require(msg.sender == admin, "only admin");
+        uint256 oldInternalCash = internalCash;
+        internalCash = IERC20(underlying).balanceOf(address(this));
+        emit CashSynced(oldInternalCash, internalCash);
     }
+
+    /**
+     * @notice Recover excess underlying tokens sent directly to the contract (e.g. donation attack)
+     * @dev Only recovers the difference between actual balance and internalCash. ACM controlled.
+     * @param token The address of the token to sweep (must be the underlying)
+     */
+    function sweepToken(address token) external {
+        uint256 excess;
+
+        assembly {
+            // ACM check: accessControlManager.isAllowedToCall(msg.sender, "sweepToken(address)")
+            let ptr := mload(0x40)
+            mstore(ptr, 0x7e13143600000000000000000000000000000000000000000000000000000000)
+            mstore(add(ptr, 0x04), caller())
+            mstore(add(ptr, 0x24), 0x40)
+            mstore(add(ptr, 0x44), 21)
+            mstore(add(ptr, 0x64), "sweepToken(address)\x00\x00")
+            let _acm := sload(accessControlManager.slot)
+            if iszero(staticcall(gas(), _acm, ptr, 0x84, ptr, 0x20)) { revert(0, 0) }
+            if iszero(mload(ptr)) { revert(0, 0) }
+
+            // require(token == underlying)
+            if iszero(eq(token, sload(underlying.slot))) { revert(0, 0) }
+
+            // excess = balanceOf(this) - internalCash
+            mstore(0x00, 0x70a0823100000000000000000000000000000000000000000000000000000000)
+            mstore(0x04, address())
+            if iszero(staticcall(gas(), token, 0x00, 0x24, 0x00, 0x20)) { revert(0, 0) }
+            excess := sub(mload(0x00), sload(internalCash.slot))
+            if iszero(excess) { revert(0, 0) }
+
+            // transfer(msg.sender, excess)
+            mstore(0x00, 0xa9059cbb00000000000000000000000000000000000000000000000000000000)
+            mstore(0x04, caller())
+            mstore(0x24, excess)
+            if iszero(call(gas(), token, 0, 0x00, 0x44, 0x00, 0x20)) { revert(0, 0) }
+            if returndatasize() {
+                if iszero(mload(0x00)) { revert(0, 0) }
+            }
+        }
+
+        emit TokenSwept(token, msg.sender, excess);
+    }
+
 }

contracts/Tokens/VTokens/VTokenInterfaces.sol

@@ -172,12 +172,18 @@ contract VTokenStorage is VTokenStorageBase {
      */
     uint256 public flashLoanAmount;
 
+    /**
+     * @notice Tracked internal cash balance, immune to direct token transfers (donation attacks)
+     * @dev Updated only via doTransferIn/doTransferOut. Must be initialized via syncCash() after upgrade.
+     */
+    uint256 public internalCash;
+
     /**
      * @dev This empty reserved space is put in place to allow future versions to add new
      * variables without shifting down storage in the inheritance chain.
      * See https://docs.openzeppelin.com/contracts/4.x/upgradeable#storage_gaps
      */
-    uint256[46] private __gap;
+    uint256[45] private __gap;
 }
 
 abstract contract VTokenInterface is VTokenStorage {
@@ -320,6 +326,16 @@ abstract contract VTokenInterface is VTokenStorage {
         uint256 protocolFee
     );
 
+    /**
+     * @notice Event emitted when internalCash is synced with actual token balance
+     */
+    event CashSynced(uint256 oldInternalCash, uint256 newInternalCash);
+
+    /**
+     * @notice Event emitted when excess tokens are swept by admin
+     */
+    event TokenSwept(address indexed token, address indexed recipient, uint256 amount);
+
     /**
      * @notice Event emitted when flashLoan fee mantissa is updated
      */

contracts/test/EvilXToken.sol

@@ -206,4 +206,8 @@ contract EvilXToken is VBep20Delegate {
     function harnessCallBorrowAllowed(uint amount) public returns (uint) {
         return comptroller.borrowAllowed(address(this), msg.sender, amount);
     }
+
+    function harnessSetInternalCash(uint256 _internalCash) public {
+        internalCash = _internalCash;
+    }
 }

contracts/test/VBep20Harness.sol

@@ -166,6 +166,10 @@ contract VBep20Harness is VBep20Immutable {
     function harnessCallBorrowAllowed(uint amount) public returns (uint) {
         return comptroller.borrowAllowed(address(this), msg.sender, amount);
     }
+
+    function harnessSetInternalCash(uint256 _internalCash) public {
+        internalCash = _internalCash;
+    }
 }
 
 contract VBep20Scenario is VBep20Immutable {
@@ -415,6 +419,10 @@ contract VBep20DelegateHarness is VBep20Delegate {
     function harnessCallBorrowAllowed(uint amount) public returns (uint) {
         return comptroller.borrowAllowed(address(this), msg.sender, amount);
     }
+
+    function harnessSetInternalCash(uint256 _internalCash) public {
+        internalCash = _internalCash;
+    }
 }
 
 contract VBep20DelegateScenario is VBep20Delegate {

contracts/test/VBep20MockDelegate.sol

@@ -187,8 +187,7 @@ contract VBep20MockDelegate is VToken, VBep20Interface {
      * @return The quantity of underlying tokens owned by this contract
      */
     function getCashPrior() internal view override returns (uint) {
-        IERC20 token = IERC20(underlying);
-        return token.balanceOf(address(this));
+        return internalCash;
     }
 
     /**
@@ -206,7 +205,9 @@ contract VBep20MockDelegate is VToken, VBep20Interface {
         token.safeTransferFrom(from, address(this), amount);
         // Calculate the amount that was *actually* transferred
         uint balanceAfter = IERC20(underlying).balanceOf(address(this));
-        return balanceAfter - balanceBefore;
+        uint actualAmount = balanceAfter - balanceBefore;
+        internalCash += actualAmount;
+        return actualAmount;
     }
 
     /**
@@ -219,7 +220,12 @@ contract VBep20MockDelegate is VToken, VBep20Interface {
      *            See here: https://medium.com/coinmonks/missing-return-value-bug-at-least-130-tokens-affected-d67bf08521ca
      */
     function doTransferOut(address payable to, uint amount) internal override {
+        internalCash -= amount;
         IERC20 token = IERC20(underlying);
         token.safeTransfer(to, amount);
     }
+
+    function harnessSetInternalCash(uint256 _internalCash) public {
+        internalCash = _internalCash;
+    }
 }

tests/hardhat/Comptroller/Diamond/flashLoan.ts

@@ -315,6 +315,8 @@ describe("FlashLoan", async () => {
 
       await underlyingA.harnessSetBalance(vTokenA.address, parseUnits("50", 18));
       await underlyingB.harnessSetBalance(vTokenB.address, parseUnits("50", 18));
+      await vTokenA.harnessSetInternalCash(parseUnits("50", 18));
+      await vTokenB.harnessSetInternalCash(parseUnits("50", 18));
 
       // Execute the flashLoan from the mockReceiverContract
       await expect(
@@ -349,6 +351,8 @@ describe("FlashLoan", async () => {
 
       await underlyingA.harnessSetBalance(vTokenA.address, parseUnits("50", 18));
       await underlyingB.harnessSetBalance(vTokenB.address, parseUnits("50", 18));
+      await vTokenA.harnessSetInternalCash(parseUnits("50", 18));
+      await vTokenB.harnessSetInternalCash(parseUnits("50", 18));
 
       await expect(
         badReceiver
@@ -415,6 +419,8 @@ describe("FlashLoan", async () => {
 
       await underlyingA.harnessSetBalance(vTokenA.address, parseUnits("60", 18));
       await underlyingB.harnessSetBalance(vTokenB.address, parseUnits("60", 18));
+      await vTokenA.harnessSetInternalCash(parseUnits("60", 18));
+      await vTokenB.harnessSetInternalCash(parseUnits("60", 18));
 
       await expect(
         mockReceiverContract
@@ -446,6 +452,8 @@ describe("FlashLoan", async () => {
 
       await underlyingA.harnessSetBalance(vTokenA.address, parseUnits("50", 18));
       await underlyingB.harnessSetBalance(vTokenB.address, parseUnits("50", 18));
+      await vTokenA.harnessSetInternalCash(parseUnits("50", 18));
+      await vTokenB.harnessSetInternalCash(parseUnits("50", 18));
 
       await comptroller.connect(alice).updateDelegate(mockReceiverContract.address, true);
 
@@ -489,6 +497,8 @@ describe("FlashLoan", async () => {
 
       await underlyingA.harnessSetBalance(vTokenA.address, parseUnits("60", 18));
       await underlyingB.harnessSetBalance(vTokenB.address, parseUnits("60", 18));
+      await vTokenA.harnessSetInternalCash(parseUnits("60", 18));
+      await vTokenB.harnessSetInternalCash(parseUnits("60", 18));
 
       await underlyingA.harnessSetBalance(alice.address, parseUnits("1000", 18));
       await underlyingB.harnessSetBalance(alice.address, parseUnits("1000", 18));
@@ -574,6 +584,8 @@ describe("FlashLoan", async () => {
 
       await underlyingA.harnessSetBalance(vTokenA.address, parseUnits("60", 18));
       await underlyingB.harnessSetBalance(vTokenB.address, parseUnits("60", 18));
+      await vTokenA.harnessSetInternalCash(parseUnits("60", 18));
+      await vTokenB.harnessSetInternalCash(parseUnits("60", 18));
 
       // Calculate expected protocol fees
       const expectedProtocolFeeA = flashLoanAmount1
@@ -708,6 +720,8 @@ describe("FlashLoan", async () => {
 
       await underlyingA.harnessSetBalance(vTokenA.address, parseUnits("60", 18));
       await underlyingB.harnessSetBalance(vTokenB.address, parseUnits("60", 18));
+      await vTokenA.harnessSetInternalCash(parseUnits("60", 18));
+      await vTokenB.harnessSetInternalCash(parseUnits("60", 18));
 
       // Calculate expected fees
       const expectedFeeA = flashLoanAmount1.mul(totalFeeMantissaTokenA).div(parseUnits("1", 18));

tests/hardhat/EvilXToken.ts

@@ -191,6 +191,7 @@ describe("Evil Token test", async () => {
     await vToken1.setReduceReservesBlockDelta(10000);
     await vToken1.connect(user).mint(convertToUnit(1, 4));
     await underlying3.harnessSetBalance(vToken3.address, convertToUnit(1, 8));
+    await vToken3.harnessSetInternalCash(convertToUnit(1, 8));
 
     const protocolShareReserve = await smock.fake<IProtocolShareReserve>("IProtocolShareReserve");
     protocolShareReserve.updateAssetsState.returns(true);