From 25d07528b1107682674bfe0bed56523238fcacb1 Mon Sep 17 00:00:00 2001
From: Pablo Carranza Velez <pablo@edgeandnode.com>
Date: Fri, 26 Jan 2024 13:14:26 -0300
Subject: [PATCH] fix: ensure thawing period when unstaking always rounds up

---
 .../contracts/staking/libs/MathUtils.sol      | 10 ++++--
 .../contracts/staking/libs/Stakes.sol         |  2 +-
 .../test/unit/staking/staking.test.ts         | 34 +++++++++++++++++++
 3 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/packages/contracts/contracts/staking/libs/MathUtils.sol b/packages/contracts/contracts/staking/libs/MathUtils.sol
index 4194b6da8..c20a1c3b0 100644
--- a/packages/contracts/contracts/staking/libs/MathUtils.sol
+++ b/packages/contracts/contracts/staking/libs/MathUtils.sol
@@ -14,19 +14,23 @@ library MathUtils {
     /**
      * @dev Calculates the weighted average of two values pondering each of these
      * values based on configured weights. The contribution of each value N is
-     * weightN/(weightA + weightB).
+     * weightN/(weightA + weightB). The calculation rounds up to ensure the result
+     * is always greater than the smallest of the two values.
      * @param valueA The amount for value A
      * @param weightA The weight to use for value A
      * @param valueB The amount for value B
      * @param weightB The weight to use for value B
      */
-    function weightedAverage(
+    function weightedAverageRoundingUp(
         uint256 valueA,
         uint256 weightA,
         uint256 valueB,
         uint256 weightB
     ) internal pure returns (uint256) {
-        return valueA.mul(weightA).add(valueB.mul(weightB)).div(weightA.add(weightB));
+        return
+            valueA.mul(weightA).add(valueB.mul(weightB)).add(weightA.add(weightB).sub(1)).div(
+                weightA.add(weightB)
+            );
     }
 
     /**
diff --git a/packages/contracts/contracts/staking/libs/Stakes.sol b/packages/contracts/contracts/staking/libs/Stakes.sol
index 5882944d8..38b5f5daa 100644
--- a/packages/contracts/contracts/staking/libs/Stakes.sol
+++ b/packages/contracts/contracts/staking/libs/Stakes.sol
@@ -73,7 +73,7 @@ library Stakes {
         // Take into account period averaging for multiple unstake requests
         uint256 lockingPeriod = _period;
         if (stake.tokensLocked > 0) {
-            lockingPeriod = MathUtils.weightedAverage(
+            lockingPeriod = MathUtils.weightedAverageRoundingUp(
                 MathUtils.diffOrZero(stake.tokensLockedUntil, block.number), // Remaining thawing period
                 stake.tokensLocked, // Weighted by remaining unstaked tokens
                 _period, // Thawing period
diff --git a/packages/contracts/test/unit/staking/staking.test.ts b/packages/contracts/test/unit/staking/staking.test.ts
index 6e88d23db..04775a272 100644
--- a/packages/contracts/test/unit/staking/staking.test.ts
+++ b/packages/contracts/test/unit/staking/staking.test.ts
@@ -239,6 +239,40 @@ describe('Staking:Stakes', () => {
         expect(afterIndexerStake.tokensLockedUntil).eq(expectedLockedUntil)
       })
 
+      it('should always increase the thawing period on subsequent unstakes', async function () {
+        const tokensToUnstake = toGRT('10')
+        const tokensToUnstakeSecondTime = toGRT('0.000001')
+        const thawingPeriod = toBN(await staking.thawingPeriod())
+
+        // Unstake (1)
+        const tx1 = await staking.connect(indexer).unstake(tokensToUnstake)
+        const receipt1 = await tx1.wait()
+        const event1: Event = receipt1.events.pop()
+        const tokensLockedUntil1 = event1.args['until']
+
+        // Move forward before the tokens are unlocked for withdrawal
+        await helpers.mineUpTo(tokensLockedUntil1.sub(5))
+
+        // Calculate locking time for tokens taking into account the previous unstake request
+        const currentBlock = await helpers.latestBlock()
+
+        // Ensure at least 1 block is added (i.e. the weighted average rounds up)
+        const expectedLockedUntil = tokensLockedUntil1.add(1)
+
+        // Unstake (2)
+        const tx2 = await staking.connect(indexer).unstake(tokensToUnstakeSecondTime)
+        const receipt2 = await tx2.wait()
+
+        // Verify events
+        const event2: Event = receipt2.events.pop()
+        expect(event2.args['until']).eq(expectedLockedUntil)
+
+        // Verify state
+        const afterIndexerStake = await staking.stakes(indexer.address)
+        expect(afterIndexerStake.tokensLocked).eq(tokensToUnstake.add(tokensToUnstakeSecondTime)) // we unstaked two times
+        expect(afterIndexerStake.tokensLockedUntil).eq(expectedLockedUntil)
+      })
+
       it('should unstake and withdraw if some tokens are unthawed', async function () {
         const tokensToUnstake = toGRT('10')
         const thawingPeriod = toBN(await staking.thawingPeriod())