`_aTokenShares` is not correctly updated in case of `withdraw()`

[hats-bug-reporter[bot]]

Github username: @0xRizwan Twitter username: 0xRizwann Submission hash (on-chain): 0x7f8e6e398793186bebb18313c0813ea29f50601bd32b9c19a8a3554945c1659c Severity: high

Description: Description\

OrigamiAaveV3BorrowAndLend.sol is an Origami abstraction over a borrow/lend money market for a single supplyToken and a single borrowToken. This contract is an Aave V3 specific interface, borrowing using variable debt only.

The contracts has supply() and withdraw() functions to supply and withdraw collateral tokens by approved position owner.

For supply, it calls _supply() function,

    function supply(
        uint256 supplyAmount
    ) external override onlyPositionOwner {
        _supply(supplyAmount);
    }

For supply, it calls _supply() function,

    function _supply(uint256 supplyAmount) internal {
        uint256 sharesBefore = aaveAToken.scaledBalanceOf(address(this));
        aavePool.supply(supplyToken, supplyAmount, address(this), referralCode);
        _aTokenShares += aaveAToken.scaledBalanceOf(address(this)) - sharesBefore;
    }

This has used _aTokenShares variable which is used to manually track tThe number of Aave aToken shares rather than relying on balanceOf function.

Similarly,

    function withdraw(
        uint256 withdrawAmount, 
        address recipient
    ) external override onlyPositionOwner returns (uint256 amountWithdrawn) {
        amountWithdrawn = _withdraw(withdrawAmount, recipient);
    }

For withdraw, it calls _withdraw() function,

    function _withdraw(uint256 withdrawAmount, address recipient) internal returns (uint256 amountWithdrawn) {
        uint256 sharesBefore = aaveAToken.scaledBalanceOf(address(this));
        amountWithdrawn = aavePool.withdraw(supplyToken, withdrawAmount, recipient);
        _aTokenShares = _aTokenShares + aaveAToken.scaledBalanceOf(address(this)) - sharesBefore;
    }

Notice, here that _aTokenShares are further increased instead of decrease as the while calling the supply() function _aTokenShares were increased by the difference of aaveAToken.scaledBalanceOf(address(this)) - sharesBefore. Therefore, while calling withdraw() functions, _aTokenShares should decreased by aaveAToken.scaledBalanceOf(address(this)) - sharesBefore

Otherwise, the _aTokenShares will be tracked incorrectly.

Recommendations\

    function _withdraw(uint256 withdrawAmount, address recipient) internal returns (uint256 amountWithdrawn) {
        uint256 sharesBefore = aaveAToken.scaledBalanceOf(address(this));
        amountWithdrawn = aavePool.withdraw(supplyToken, withdrawAmount, recipient);
-        _aTokenShares = _aTokenShares + aaveAToken.scaledBalanceOf(address(this)) - sharesBefore;
+        _aTokenShares -= aaveAToken.scaledBalanceOf(address(this)) - sharesBefore;
    }

[frontier159]

The maths as it stands are correct as is.

Assume _aTokenShares=80, but someone donated/sent 20 more aTokens than we were tracking manually, so:

_aTokenShares == 80
aaveAToken.scaledBalanceOf(address(this)) == 100

Then supply() is called for the equivalent of 9 aTokenShares:

// after the aavePool.supply():
aaveAToken.scaledBalanceOf(address(this)) == 109 // 100 + 9
_aTokenShares = 89 // 80 + (109 - 100)

Now, withdraw() is called for the equivalent of 50 aTokenShares:

// after the aavePool.withdraw():
aaveAToken.scaledBalanceOf(address(this)) == 59
_aTokenShares = 39 // _aTokenShares + aaveAToken.scaledBalanceOf(address(this)) - sharesBefore = 89 + 59 - 109

Doing as you suggest is actually incorrect as it would lead to an underflow panic:

_aTokenShares -= aaveAToken.scaledBalanceOf(address(this)) - sharesBefore;
_aTokenShares = _aTokenShares - (aaveAToken.scaledBalanceOf(address(this)) - sharesBefore)
_aTokenShares = 89 - (59 - 109)

Perhaps you meant:

_aTokenShares -= sharesBefore - aaveAToken.scaledBalanceOf(address(this));

But this is mathematically equivalent to what we already have.

If you think there's more to it, please create a PoC/test