`totalAssets()` in `OsToken` incorrectly excludes treasury assets

Github username: @milotruck Submission hash (on-chain): 0x786d7a93c9d40afe7674d19aff3460c615ed5e222bb9b87ae4989a3eaf492c81 Severity: low

Description:

Bug Description

In OsToken.sol, the totalAssets() function calculates the current total amount of assets as follows:

OsToken.sol#L84-L90

  function totalAssets() public view override returns (uint256) {
    uint256 profitAccrued = _unclaimedAssets();
    if (profitAccrued == 0) return _totalAssets;

    uint256 treasuryAssets = Math.mulDiv(profitAccrued, feePercent, _maxFeePercent);
    return _totalAssets + profitAccrued - treasuryAssets;
  }

As seen from above, the treasuryAssets is subtracted from _totalAssets + profitAccrued to calculate the total amount of assets.

However, this contradicts the updateState() function, which does not subtract treasuryAssets from _totalAssets:

OsToken.sol#L261-L262

    // calculate treasury assets
    uint256 newTotalAssets = _totalAssets + profitAccrued;

OsToken.sol#L299

    _totalAssets = SafeCast.toUint128(newTotalAssets);

Therefore, as long as treasuryAssets is non-zero, the value returned by totalAssets() will actually be smaller than the total asset amount after updateState() is called, making it incorrect.

Attack Scenario

Assume that a vault has the following state:

_totalAssets = 10 ether
profitAccrued = 1 ether
feePercent = 500, which is 5%.

When totalAssets() is called, it will return:

treasuryAssets = 1 ether * 500 / 10_000 = 5e16
_totalAssets + profitAccrued - treasuryAssets = 10 ether + 1 ether - 5e16 = 11e18 - 5e16

When updateState() is called, _totalAssets is updated to:

_totalAssets + profitAccrued = 11 ether
Note that profitAccrued is also reset to 0 after updateState() is called.

Now, since profitAccrued = 0, totalAssets() will return a different value:

treasuryAssets = 0 * 500 / 10_000 = 0
_totalAssets + profitAccrued - treasuryAssets = 11 ether + 0 - 0 = 11e18

As totalAssets() incorrectly excludes treasuryAssets, its return value was 5e16 smaller before updateState() is called.

Impact

As totalAssets() incorrectly excludes treasuryAssets, it cannot be used to accurately determine the total amount of assets for OsToken.

This could potentially have more severe impacts if other contracts rely on totalAssets() for calculations, as its return value can change in a single transaction by calling updateState().

Recommended Mitigation

Consider including treasuryAssets in the value returned by totalAssets(), which is consistent with the logic in the updateState() function:

OsToken.sol#L84-L90

  function totalAssets() public view override returns (uint256) {
    uint256 profitAccrued = _unclaimedAssets();
+   return _totalAssets + profitAccrued;
-   if (profitAccrued == 0) return _totalAssets;
-
-   uint256 treasuryAssets = Math.mulDiv(profitAccrued, feePercent, _maxFeePercent);
-   return _totalAssets + profitAccrued - treasuryAssets;
  }

hats-finance / StakeWise-0xd91cd6ed6c9a112fdc112b1a3c66e47697f522cd