Fix AStETH rounding issue

README.md

@@ -125,6 +125,7 @@ To run tests for StETH integration use the following commands:
 
 ```
 npm install
+npm run compile
 npm run test:steth
 npm run test:steth-coverage # to run tests with coverage report
 ```

contracts/mocks/tokens/StETHMocked.sol

@@ -9,8 +9,9 @@ pragma solidity 0.6.12;
 contract StETHMocked {
   using LidoSafeMath for uint256;
 
-  uint256 internal _totalShares;
-  uint256 internal _pooledEther;
+  // use there values like real stETH has
+  uint256 internal _totalShares = 1608965089698263670456320;
+  uint256 internal _pooledEther = 1701398689820002221426255;
   mapping(address => uint256) private shares;
   mapping(address => mapping(address => uint256)) private allowances;
 
@@ -200,7 +201,7 @@ contract StETHMocked {
     }
 
     _mintShares(sender, sharesAmount);
-    _pooledEther = _pooledEther.add(sharesAmount);
+    _pooledEther = _pooledEther.add(deposit);
     return sharesAmount;
   }
 

contracts/protocol/tokenization/lido/AStETH.sol

@@ -105,7 +105,7 @@ contract AStETH is VersionedInitializable, IncentivizedERC20, IAToken {
     uint256 amount,
     uint256 index
   ) external override onlyLendingPool {
-    uint256 amountScaled = amount.rayDiv(_stEthRebasingIndex()).rayDiv(index);
+    uint256 amountScaled = _toInternalAmount(amount, _stEthRebasingIndex(), index);
     require(amountScaled != 0, Errors.CT_INVALID_BURN_AMOUNT);
     _burn(user, amountScaled);
 
@@ -130,7 +130,7 @@ contract AStETH is VersionedInitializable, IncentivizedERC20, IAToken {
   ) external override onlyLendingPool returns (bool) {
     uint256 previousBalance = super.balanceOf(user);
 
-    uint256 amountScaled = amount.rayDiv(_stEthRebasingIndex()).rayDiv(index);
+    uint256 amountScaled = _toInternalAmount(amount, _stEthRebasingIndex(), index);
     require(amountScaled != 0, Errors.CT_INVALID_MINT_AMOUNT);
     _mint(user, amountScaled);
 
@@ -152,10 +152,10 @@ contract AStETH is VersionedInitializable, IncentivizedERC20, IAToken {
     }
 
     // Compared to the normal mint, we don't check for rounding errors.
-    // The amount to mint can easily be very small since it is a fraction of the interest ccrued.
+    // The amount to mint can easily be very small since it is a fraction of the interest accrued.
     // In that case, the treasury will experience a (very small) loss, but it
-    // wont cause potentially valid transactions to fail.
-    _mint(RESERVE_TREASURY_ADDRESS, amount.rayDiv(_stEthRebasingIndex()).rayDiv(index));
+    // won't cause potentially valid transactions to fail.
+    _mint(RESERVE_TREASURY_ADDRESS, _toInternalAmount(amount, _stEthRebasingIndex(), index));
 
     emit Transfer(address(0), RESERVE_TREASURY_ADDRESS, amount);
     emit Mint(RESERVE_TREASURY_ADDRESS, amount, index);
@@ -207,6 +207,17 @@ contract AStETH is VersionedInitializable, IncentivizedERC20, IAToken {
     return _scaledBalanceOf(user, _stEthRebasingIndex());
   }
 
+  /**
+   * @dev Returns the internal balance of the user. The internal balance is the balance of
+   * the underlying asset of the user (sum of deposits of the user), divided by the current
+   * liquidity index at the moment of the update and by the current stETH rebasing index.
+   * @param user The user whose balance is calculated
+   * @return The internal balance of the user
+   **/
+  function internalBalanceOf(address user) external view returns (uint256) {
+    return super.balanceOf(user);
+  }
+
   /**
    * @dev Returns the scaled balance of the user and the scaled total supply.
    * @param user The address of the user
@@ -247,6 +258,15 @@ contract AStETH is VersionedInitializable, IncentivizedERC20, IAToken {
     return _scaledTotalSupply(_stEthRebasingIndex());
   }
 
+  /**
+   * @dev Returns the internal total supply of the token. Represents
+   * sum(debt/_stEthRebasingIndex/liquidityIndex).
+   * @return the internal total supply
+   */
+  function internalTotalSupply() external view returns (uint256) {
+    return super.totalSupply();
+  }
+
   /**
    * @dev Transfers the underlying asset to `target`. Used by the LendingPool to transfer
    * assets in borrow(), withdraw() and flashLoan()
@@ -315,15 +335,17 @@ contract AStETH is VersionedInitializable, IncentivizedERC20, IAToken {
     uint256 amount,
     bool validate
   ) internal {
-    uint256 index = POOL.getReserveNormalizedIncome(UNDERLYING_ASSET_ADDRESS);
+    uint256 aaveLiquidityIndex = POOL.getReserveNormalizedIncome(UNDERLYING_ASSET_ADDRESS);
+    uint256 stEthRebasingIndex = _stEthRebasingIndex();
 
-    uint256 rebasingIndex = _stEthRebasingIndex();
-    uint256 fromBalanceBefore = _scaledBalanceOf(from, rebasingIndex).rayMul(index);
-    uint256 toBalanceBefore = _scaledBalanceOf(to, rebasingIndex).rayMul(index);
+    uint256 fromBalanceBefore =
+      _scaledBalanceOf(from, stEthRebasingIndex).rayMul(aaveLiquidityIndex);
+    uint256 toBalanceBefore = _scaledBalanceOf(to, stEthRebasingIndex).rayMul(aaveLiquidityIndex);
 
-    super._transfer(from, to, amount.rayDiv(rebasingIndex).rayDiv(index));
+    super._transfer(from, to, _toInternalAmount(amount, stEthRebasingIndex, aaveLiquidityIndex));
 
     if (validate) {
+      require(fromBalanceBefore >= amount, 'ERC20: transfer amount exceeds balance');
       POOL.finalizeTransfer(
         UNDERLYING_ASSET_ADDRESS,
         from,
@@ -334,7 +356,7 @@ contract AStETH is VersionedInitializable, IncentivizedERC20, IAToken {
       );
     }
 
-    emit BalanceTransfer(from, to, amount, index);
+    emit BalanceTransfer(from, to, amount, aaveLiquidityIndex);
   }
 
   /**
@@ -351,42 +373,33 @@ contract AStETH is VersionedInitializable, IncentivizedERC20, IAToken {
     _transfer(from, to, amount, true);
   }
 
-  /**
-   * @return Current rebasin index of stETH in RAY
-   **/
-  function _stEthRebasingIndex() internal view returns (uint256) {
-    // Below expression returns how much Ether corresponds
-    // to 10 ** 27 shares. 10 ** 27 was taken  to provide
-    // same precision as AAVE's liquidity index, which
-    // counted in RAY's (decimals with 27 digits).
-    return ILido(UNDERLYING_ASSET_ADDRESS).getPooledEthByShares(1e27);
-  }
-
   function _scaledBalanceOf(address user, uint256 rebasingIndex) internal view returns (uint256) {
-    return super.balanceOf(user).rayMul(rebasingIndex);
+    return super.balanceOf(user).mul(rebasingIndex).div(WadRayMath.RAY);
   }
 
   function _scaledTotalSupply(uint256 rebasingIndex) internal view returns (uint256) {
-    return super.totalSupply().rayMul(rebasingIndex);
+    return super.totalSupply().mul(rebasingIndex).div(WadRayMath.RAY);
   }
 
   /**
-   * @dev Returns the internal balance of the user. The internal balance is the balance of
-   * the underlying asset of the user (sum of deposits of the user), divided by the current
-   * liquidity index at the moment of the update and by the current stETH rebasing index.
-   * @param user The user whose balance is calculated
-   * @return The internal balance of the user
+   * @return Current rebasing index of stETH in RAY
    **/
-  function internalBalanceOf(address user) external view returns (uint256) {
-    return super.balanceOf(user);
+  function _stEthRebasingIndex() internal view returns (uint256) {
+    // Returns amount of stETH corresponding to 10**27 stETH shares.
+    // The 10**27 is picked to provide the same precision as the AAVE
+    // liquidity index, which is in RAY (10**27).
+    return ILido(UNDERLYING_ASSET_ADDRESS).getPooledEthByShares(WadRayMath.RAY);
   }
 
   /**
-   * @dev Returns the internal total supply of the token. Represents
-   * sum(debt/_stEthRebasingIndex/liquidityIndex).
-   * @return the internal total supply
+   * @dev Converts amount of astETH to internal shares, based
+   *  on stEthRebasingIndex and aaveLiquidityIndex.
    */
-  function internalTotalSupply() external view returns (uint256) {
-    return super.totalSupply();
+  function _toInternalAmount(
+    uint256 amount,
+    uint256 stEthRebasingIndex,
+    uint256 aaveLiquidityIndex
+  ) internal view returns (uint256) {
+    return amount.mul(WadRayMath.RAY).div(stEthRebasingIndex).rayDiv(aaveLiquidityIndex);
   }
 }

contracts/protocol/tokenization/lido/README.md

@@ -57,7 +57,7 @@ function _stEthRebasingIndex() returns (uint256) {
   // to 10 ** 27 shares. 10 ** 27 was taken  to provide
   // same precision as AAVE's liquidity index, which
   // counted in RAY's (decimals with 27 digits).
-  return stETH.getPooledEthByShares(10**27);
+  return stETH.getPooledEthByShares(WadRayMath.RAY);
 }
 
 ```
@@ -68,27 +68,35 @@ With stETH rebasing index, `AStETH` allows to make rebases profit accountable, a
 function mint(address user, uint256 amount, uint256 liquidityIndex) {
     ...
     uint256 stEthRebasingIndex = _stEthRebasingIndex();
-    _mint(user, amount.rayDiv(stEthRebasingIndex).rayDiv(liquidityIndex));
+    _mint(user, _toInternalAmount(amount, stEthRebasingIndex, liquidityIndex));
     ...
 }
 
 function burn(address user, uint256 amount, uint256 liquidityIndex) {
     ...
     uint256 stEthRebasingIndex = _stEthRebasingIndex();
-    _burn(user, amount.rayDiv(stEthRebasingIndex)).rayDiv(liquidityIndex);
+    _burn(user, _toInternalAmount(amount, stEthRebasingIndex, liquidityIndex));
     ...
 }
+
+function _toInternalAmount(
+    uint256 amount,
+    uint256 stEthRebasingIndex,
+    uint256 aaveLiquidityIndex
+  ) internal view returns (uint256) {
+    return amount.mul(WadRayMath.RAY).div(stEthRebasingIndex).rayDiv(aaveLiquidityIndex);
+  }
 ```
 
 Then, according to AAVE's definitions, `scaledTotalSupply()` and `scaledBalanceOf()` might be calculated as:
 
 ```solidity=
 function scaledTotalSupply() returns (uint256) {
-  return _totalSupply.rayMul(_stEthRebasingIndex());
+  return _totalSupply.mul(_stEthRebasingIndex()).div(WadRayMath.RAY);
 }
 
 function scaledBalanceOf(address user) returns (uint256) {
-  return _balances[user].rayMul(_stEthRebasingIndex());
+  return _balances[user].mul(_stEthRebasingIndex()).div(WadRayMath.RAY);
 }
 
 ```

test/astETH/__setup.spec.ts

@@ -6,6 +6,7 @@ import bignumberChai from 'chai-bignumber';
 import { solidity } from 'ethereum-waffle';
 import { AstEthSetup } from './init';
 import '../helpers/utils/math';
+import { wei } from './helpers';
 
 chai.use(bignumberChai());
 chai.use(almostEqual());
@@ -15,6 +16,7 @@ let setup: AstEthSetup, evmSnapshotId;
 
 before(async () => {
   setup = await AstEthSetup.deploy();
+  await setup.priceFeed.setPrice(wei`0.99 ether`);
   evmSnapshotId = await hre.ethers.provider.send('evm_snapshot', []);
 });
 

test/astETH/asserts.ts

@@ -0,0 +1,88 @@
+import BigNumber from 'bignumber.js';
+import { expect } from 'chai';
+import { toWei } from './helpers';
+import { AstEthSetup, Lender } from './init';
+
+export function lt(actual: string, expected: string, message?: string) {
+  expect(actual).to.be.bignumber.lt(expected, message);
+}
+
+export function gt(actual: string, expected: string, message?: string) {
+  expect(actual).to.be.bignumber.gt(expected, message);
+}
+
+export function eq(actual: string, expected: string, message?: string) {
+  expect(actual).is.equal(expected, message);
+}
+
+export function almostEq(actual: string, expected: string, epsilon: string = '1') {
+  const lowerBound = new BigNumber(expected).minus(epsilon).toString();
+  const upperBound = new BigNumber(expected).plus(epsilon).toString();
+  expect(actual).to.be.bignumber.lte(upperBound);
+  expect(actual).to.be.bignumber.gte(lowerBound);
+}
+
+export function lte(actual: string, expected: string, epsilon: string = '1') {
+  const lowerBound = new BigNumber(expected).minus(epsilon).toString();
+  expect(actual).to.be.bignumber.lte(expected);
+  expect(actual).to.be.bignumber.gte(lowerBound);
+}
+
+export function gte(actual: string, expected: string, epsilon: string = '1') {
+  const upperBound = new BigNumber(expected).plus(epsilon).toString();
+  expect(actual).to.be.bignumber.gte(expected);
+  expect(actual).to.be.bignumber.lte(upperBound);
+}
+
+export async function astEthBalance(
+  lender: Lender,
+  expectedBalance: string,
+  epsilon: string = '1'
+) {
+  const [balance, internalBalance, liquidityIndex] = await Promise.all([
+    lender.astEthBalance(),
+    lender.astEthInternalBalance(),
+    lender.lendingPool.getReserveNormalizedIncome(lender.stETH.address).then(toWei),
+  ]);
+  lte(balance, expectedBalance, epsilon);
+  // to validate that amount of shares is correct
+  // we convert internal balance to stETH shares and assert with astETH balance
+  const fromInternalBalance = await lender.stETH.getPooledEthByShares(internalBalance).then(toWei);
+  eq(
+    new BigNumber(fromInternalBalance).rayMul(new BigNumber(liquidityIndex)).toFixed(0),
+    balance,
+    `Unexpected astETH.internalBalanceOf() value`
+  );
+}
+
+export async function astEthTotalSupply(
+  setup: AstEthSetup,
+  expectedValue: string,
+  epsilon: string = '1'
+) {
+  const [totalSupply, internalTotalSupply, stEthBalance, liquidityIndex] = await Promise.all([
+    setup.astEthTotalSupply(),
+    setup.astETH.internalTotalSupply().then(toWei),
+    setup.stETH.balanceOf(setup.astETH.address).then(toWei),
+    setup.aave.lendingPool.getReserveNormalizedIncome(setup.stETH.address).then(toWei),
+  ]);
+
+  lte(totalSupply, expectedValue, epsilon);
+  // to validate that internal number of shares is correct
+  // internal total supply converts to stETH and assert it with astETH total supply
+  const fromInternalTotalSupply = await setup.stETH
+    .getPooledEthByShares(internalTotalSupply)
+    .then(toWei);
+  eq(
+    new BigNumber(fromInternalTotalSupply).rayMul(new BigNumber(liquidityIndex)).toFixed(0),
+    totalSupply,
+    `Unexpected astETH.internalTotalSupply()`
+  );
+  eq(
+    totalSupply,
+    stEthBalance,
+    `astETH.totalSupply() is ${totalSupply}, but stETH.balanceOf(astETH) is ${stEthBalance}`
+  );
+}
+
+export default { lt, lte, eq, almostEq, gt, gte, astEthBalance, astEthTotalSupply };

test/astETH/astETH-allowance.spec.ts

@@ -1,38 +1,45 @@
-import { assertBalance, wei } from './helpers';
+import asserts from './asserts';
+import { wei } from './helpers';
 import { setup } from './__setup.spec';
 
 describe('AStETH Allowance:', function () {
   it('allowance', async () => {
     const { lenderA, lenderB } = setup.lenders;
+
     const allowanceBefore = await lenderA.astETH.allowance(lenderA.address, lenderB.address);
-    assertBalance(allowanceBefore.toString(), wei(0));
-    await lenderA.astETH.approve(lenderB.address, wei(10));
+    asserts.eq(allowanceBefore.toString(), wei`0`);
+
+    await lenderA.astETH.approve(lenderB.address, wei`10 ether`);
     const allowanceAfter = await lenderA.astETH.allowance(lenderA.address, lenderB.address);
-    assertBalance(allowanceAfter.toString(), wei(10));
+    asserts.eq(allowanceAfter.toString(), wei`10 ether`);
   });
   it('decreaseAllowance', async () => {
     const { lenderA, lenderB } = setup.lenders;
+
     const allowanceBefore = await lenderA.astETH.allowance(lenderA.address, lenderB.address);
-    assertBalance(allowanceBefore.toString(), wei(0));
-    await lenderA.astETH.approve(lenderB.address, wei(10));
+    asserts.eq(allowanceBefore.toString(), wei`0`);
+
+    await lenderA.astETH.approve(lenderB.address, wei`10 ether`);
     const allowanceAfter = await lenderA.astETH.allowance(lenderA.address, lenderB.address);
-    assertBalance(allowanceAfter.toString(), wei(10));
+    asserts.eq(allowanceAfter.toString(), wei`10 ether`);
 
-    await lenderA.astETH.decreaseAllowance(lenderB.address, wei(5));
+    await lenderA.astETH.decreaseAllowance(lenderB.address, wei`5 ether`);
     const allowanceAfterDecrease = await lenderA.astETH.allowance(lenderA.address, lenderB.address);
-    assertBalance(allowanceAfterDecrease.toString(), wei(5));
+    asserts.eq(allowanceAfterDecrease.toString(), wei`5 ether`);
   });
 
   it('increaseAllowance', async () => {
     const { lenderA, lenderB } = setup.lenders;
+
     const allowanceBefore = await lenderA.astETH.allowance(lenderA.address, lenderB.address);
-    assertBalance(allowanceBefore.toString(), wei(0));
-    await lenderA.astETH.approve(lenderB.address, wei(10));
+    asserts.eq(allowanceBefore.toString(), wei`0`);
+
+    await lenderA.astETH.approve(lenderB.address, wei`10 ether`);
     const allowanceAfter = await lenderA.astETH.allowance(lenderA.address, lenderB.address);
-    assertBalance(allowanceAfter.toString(), wei(10));
+    asserts.eq(allowanceAfter.toString(), wei`10 ether`);
 
-    await lenderA.astETH.increaseAllowance(lenderB.address, wei(5));
+    await lenderA.astETH.increaseAllowance(lenderB.address, wei`5 ether`);
     const allowanceAfterDecrease = await lenderA.astETH.allowance(lenderA.address, lenderB.address);
-    assertBalance(allowanceAfterDecrease.toString(), wei(15));
+    asserts.eq(allowanceAfterDecrease.toString(), wei`15 ether`);
   });
 });