From 2644aaca7f0a285c86d86ad5387b47551f61f111 Mon Sep 17 00:00:00 2001
From: Lasse Herskind <16536249+LHerskind@users.noreply.github.com>
Date: Mon, 21 Jun 2021 16:07:44 +0200
Subject: [PATCH] fix: Remove block number constraint from
 collectAndUpdateRewards() + add test

---
 .../protocol/tokenization/StaticATokenLM.sol  | 38 +++++-----
 ...ic-atoken-liquidity-mining-rewards.spec.ts | 71 +++++++++++++++++++
 2 files changed, 89 insertions(+), 20 deletions(-)

diff --git a/contracts/protocol/tokenization/StaticATokenLM.sol b/contracts/protocol/tokenization/StaticATokenLM.sol
index 1c1b35b8..05ef9656 100644
--- a/contracts/protocol/tokenization/StaticATokenLM.sol
+++ b/contracts/protocol/tokenization/StaticATokenLM.sol
@@ -464,30 +464,28 @@ contract StaticATokenLM is ERC20 {
    * @dev Claims rewards from `INCENTIVES_CONTROLLER` and updates internal accounting of rewards.
    */
   function collectAndUpdateRewards() public {
-    if (block.number > _lastRewardBlock) {
-      _lastRewardBlock = block.number;
-      uint256 supply = totalSupply();
+    _lastRewardBlock = block.number;
+    uint256 supply = totalSupply();
 
-      address[] memory assets = new address[](1);
-      assets[0] = address(ATOKEN);
+    address[] memory assets = new address[](1);
+    assets[0] = address(ATOKEN);
 
-      uint256 freshlyClaimed =
-        INCENTIVES_CONTROLLER.claimRewards(assets, type(uint256).max, address(this));
-      uint256 lifetimeRewards = _lifetimeRewardsClaimed.add(freshlyClaimed);
-      uint256 rewardsAccrued = lifetimeRewards.sub(_lifetimeRewards).wadToRay();
+    uint256 freshlyClaimed =
+      INCENTIVES_CONTROLLER.claimRewards(assets, type(uint256).max, address(this));
+    uint256 lifetimeRewards = _lifetimeRewardsClaimed.add(freshlyClaimed);
+    uint256 rewardsAccrued = lifetimeRewards.sub(_lifetimeRewards).wadToRay();
 
-      if (supply > 0 && rewardsAccrued > 0) {
-        _accRewardsPerToken = _accRewardsPerToken.add(
-          (rewardsAccrued).rayDivNoRounding(supply.wadToRay())
-        );
-      }
-
-      if (rewardsAccrued > 0) {
-        _lifetimeRewards = lifetimeRewards;
-      }
-      // Unsure if we can also move this in
-      _lifetimeRewardsClaimed = lifetimeRewards;
+    if (supply > 0 && rewardsAccrued > 0) {
+      _accRewardsPerToken = _accRewardsPerToken.add(
+        (rewardsAccrued).rayDivNoRounding(supply.wadToRay())
+      );
     }
+
+    if (rewardsAccrued > 0) {
+      _lifetimeRewards = lifetimeRewards;
+    }
+
+    _lifetimeRewardsClaimed = lifetimeRewards;
   }
 
   /**
diff --git a/test-suites/test-aave/mainnet/static-atoken-lm/static-atoken-liquidity-mining-rewards.spec.ts b/test-suites/test-aave/mainnet/static-atoken-lm/static-atoken-liquidity-mining-rewards.spec.ts
index 99a85540..c510b76b 100644
--- a/test-suites/test-aave/mainnet/static-atoken-lm/static-atoken-liquidity-mining-rewards.spec.ts
+++ b/test-suites/test-aave/mainnet/static-atoken-lm/static-atoken-liquidity-mining-rewards.spec.ts
@@ -839,4 +839,75 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
       expect(await stkAave.balanceOf(await users[i].getAddress())).to.be.eq(pendingReward);
     }
   });
+
+  it('Checks that withdraw and collect in different blocks updates _lifetimeRewardsClaimed as expected', async () => {
+    const users = await DRE.ethers.getSigners();
+    const user = users[0];
+    const depositAmount = utils.parseEther('1');
+
+    // Preparation
+    await waitForTx(await weth.connect(user).deposit({ value: depositAmount }));
+    await waitForTx(
+      await weth.connect(user).approve(staticAToken.address, depositAmount, defaultTxParams)
+    );
+
+    // Deposit
+    await waitForTx(
+      await staticAToken
+        .connect(user)
+        .deposit(await user.getAddress(), depositAmount, 0, true, defaultTxParams)
+    );
+
+    await advanceTimeAndBlock(60);
+
+    expect(await staticAToken.getLifetimeRewardsClaimed()).to.be.eq(0);
+    expect(await staticAToken.getClaimableRewards(user.address)).to.be.gt(0);
+    expect(await stkAave.balanceOf(user.address)).to.be.eq(0);
+
+    await waitForTx(await staticAToken.connect(user).withdraw(user.address, MAX_UINT_AMOUNT, true));
+    await staticAToken.collectAndUpdateRewards();
+    await staticAToken.connect(user).claimRewards(user.address, false);
+
+    expect(await staticAToken.getLifetimeRewardsClaimed()).to.be.gt(0);
+    expect(await staticAToken.getClaimableRewards(user.address)).to.be.eq(0);
+    expect(await stkAave.balanceOf(user.address)).to.be.gt(0);
+  });
+
+  it('Checks that withdraw and collect in the same block updates _lifetimeRewardsClaimed as expected', async () => {
+    const users = await DRE.ethers.getSigners();
+    const user = users[0];
+    const depositAmount = utils.parseEther('1');
+
+    // Preparation
+    await waitForTx(await weth.connect(user).deposit({ value: depositAmount }));
+    await waitForTx(
+      await weth.connect(user).approve(staticAToken.address, depositAmount, defaultTxParams)
+    );
+
+    // Deposit
+    await waitForTx(
+      await staticAToken
+        .connect(user)
+        .deposit(await user.getAddress(), depositAmount, 0, true, defaultTxParams)
+    );
+
+    await advanceTimeAndBlock(60);
+
+    expect(await staticAToken.getLifetimeRewardsClaimed()).to.be.eq(0);
+    expect(await staticAToken.getClaimableRewards(user.address)).to.be.gt(0);
+    expect(await stkAave.balanceOf(user.address)).to.be.eq(0);
+
+    await DRE.network.provider.send('evm_setAutomine', [false]);
+
+    await staticAToken.connect(user).withdraw(user.address, MAX_UINT_AMOUNT, true);
+    await staticAToken.collectAndUpdateRewards();
+    await staticAToken.connect(user).claimRewards(user.address, false);
+
+    await DRE.network.provider.send('evm_mine', []);
+    await DRE.network.provider.send('evm_setAutomine', [true]);
+
+    expect(await staticAToken.getLifetimeRewardsClaimed()).to.be.gt(0);
+    expect(await staticAToken.getClaimableRewards(user.address)).to.be.eq(0);
+    expect(await stkAave.balanceOf(user.address)).to.be.gt(0);
+  });
 });