feat(DrawCalculatorTimelock: accept timestamp parameter in lock to dy…

…namically adjust unlock time
pooltogether · Oct 13, 2021 · c0ab310 · c0ab310
1 parent 1a7b03c
commit c0ab310
Show file tree

Hide file tree

Showing 8 changed files with 54 additions and 26 deletions.
diff --git a/contracts/DrawCalculatorTimelock.sol b/contracts/DrawCalculatorTimelock.sol
@@ -76,13 +76,11 @@ contract DrawCalculatorTimelock is IDrawCalculatorTimelock, Manageable {
     }
 
     /// @inheritdoc IDrawCalculatorTimelock
-    function lock(uint32 _drawId) external override onlyManagerOrOwner returns (bool) {
+    function lock(uint32 _drawId, uint32 _timestamp) external override onlyManagerOrOwner returns (bool) {
         Timelock memory _timelock = timelock;
         require(_drawId == _timelock.drawId + 1, "OM/not-drawid-plus-one");
 
         _requireTimelockElapsed(_timelock);
-
-        uint128 _timestamp = uint128(block.timestamp);
         timelock = Timelock({ drawId: _drawId, timestamp: _timestamp });
         emit LockedDraw(_drawId, uint32(_timestamp));
 

diff --git a/contracts/L1TimelockTrigger.sol b/contracts/L1TimelockTrigger.sol
@@ -62,15 +62,16 @@ contract L1TimelockTrigger is Manageable {
     /**
      * @notice Push Draw onto draws ring buffer history.
      * @dev    Restricts new draws by forcing a push timelock.
-     * @param _drawId draw id
-     * @param _prizeDistribution PrizeDistribution parameters
+     * @param _draw Draw struct
+     * @param _prizeDistribution PrizeDistribution struct
      */
-    function push(uint32 _drawId, IPrizeDistributionBuffer.PrizeDistribution memory _prizeDistribution)
+    function push(IDrawBeacon.Draw calldata _draw, IPrizeDistributionBuffer.PrizeDistribution memory _prizeDistribution)
         external
         onlyManagerOrOwner
-    {
-        timelock.lock(_drawId);
-        prizeDistributionBuffer.pushPrizeDistribution(_drawId, _prizeDistribution);
-        emit PrizeDistributionPushed(_drawId, _prizeDistribution);
+    {   
+        // Locks the new PrizeDistribution according to the Draw endtime.
+        timelock.lock(_draw.drawId, uint32(_draw.timestamp + _draw.beaconPeriodSeconds));
+        prizeDistributionBuffer.pushPrizeDistribution(_draw.drawId, _prizeDistribution);
+        emit PrizeDistributionPushed(_draw.drawId, _prizeDistribution);
     }
 }
diff --git a/contracts/L2TimelockTrigger.sol b/contracts/L2TimelockTrigger.sol
@@ -76,7 +76,7 @@ contract L2TimelockTrigger is Manageable {
         external
         onlyManagerOrOwner
     {
-        timelock.lock(_draw.drawId);
+        timelock.lock(_draw.drawId, uint32(_draw.timestamp + _draw.beaconPeriodSeconds));
         drawBuffer.pushDraw(_draw);
         prizeDistributionBuffer.pushPrizeDistribution(_draw.drawId, _prizeDistribution);
         emit DrawAndPrizeDistributionPushed(_draw.drawId, _draw, _prizeDistribution);

diff --git a/contracts/interfaces/IDrawCalculatorTimelock.sol b/contracts/interfaces/IDrawCalculatorTimelock.sol
@@ -5,6 +5,12 @@ pragma solidity 0.8.6;
 import "@pooltogether/v4-core/contracts/interfaces/IDrawCalculator.sol";
 
 interface IDrawCalculatorTimelock {
+
+     /**
+        * @notice Emitted when target draw id is locked.
+        * @param timestamp The epoch timestamp to unlock the current locked Draw
+        * @param drawId    The Draw to unlock
+     */
     struct Timelock {
         uint128 timestamp;
         uint32 drawId;
@@ -47,9 +53,10 @@ interface IDrawCalculatorTimelock {
      * @notice Lock passed draw id for `timelockDuration` seconds.
      * @dev    Restricts new draws by forcing a push timelock.
      * @param _drawId Draw id to lock.
+     * @param _timestamp Epoch timestamp to unlock the draw.
      * @return True if operation was successful.
      */
-    function lock(uint32 _drawId) external returns (bool);
+    function lock(uint32 _drawId, uint32 _timestamp) external returns (bool);
 
     /**
      * @notice Read internal DrawCalculator variable.

diff --git a/test/DrawCalculatorTimelock.test.ts b/test/DrawCalculatorTimelock.test.ts
@@ -102,7 +102,7 @@ describe('DrawCalculatorTimelock', () => {
 
         it('should lock next draw id', async () => {
             await increaseTime(timelockDuration + 1);
-            await expect(drawCalculatorTimelock.lock(2))
+            await expect(drawCalculatorTimelock.lock(2, (await getBlock('latest')).timestamp + 1))
                 .to.emit(drawCalculatorTimelock, 'LockedDraw')
 
             const timelock = await drawCalculatorTimelock.getTimelock();
@@ -116,7 +116,7 @@ describe('DrawCalculatorTimelock', () => {
             await drawCalculatorTimelock.setManager(wallet2.address);
 
             await increaseTime(timelockDuration + 1);
-            await drawCalculatorTimelock.connect(wallet2).lock(2);
+            await drawCalculatorTimelock.connect(wallet2).lock(2, (await getBlock('latest')).timestamp + 1);
 
             const timelock = await drawCalculatorTimelock.getTimelock();
             const currentTimestamp = (await getBlock('latest')).timestamp;
@@ -126,13 +126,13 @@ describe('DrawCalculatorTimelock', () => {
         });
 
         it('should fail if not called by the owner or manager', async () => {
-            await expect(drawCalculatorTimelock.connect(wallet2).lock(1)).to.be.revertedWith(
+            await expect(drawCalculatorTimelock.connect(wallet2).lock(1, (await getBlock('latest')).timestamp)).to.be.revertedWith(
                 'Manageable/caller-not-manager-or-owner',
             );
         });
 
         it('should fail to lock if trying to lock current or previous draw id', async () => {
-            await expect(drawCalculatorTimelock.lock(1)).to.be.revertedWith(
+            await expect(drawCalculatorTimelock.lock(1, (await getBlock('latest')).timestamp)).to.be.revertedWith(
                 'OM/not-drawid-plus-one',
             );
         });

diff --git a/test/L1TimelockTrigger.test.ts b/test/L1TimelockTrigger.test.ts
@@ -54,26 +54,49 @@ describe('L1TimelockTrigger', () => {
 
     describe('push()', () => {
         it('should allow a push when no push has happened', async () => {
+            const Draw = {
+                drawId: 1,
+                timestamp: 22,
+                winningRandomNumber: 333,
+                beaconPeriodStartedAt: 4444,
+                beaconPeriodSeconds: 55555,
+            }
+
             await prizeDistributionBuffer.mock.pushPrizeDistribution.returns(0);
-            await drawCalculatorTimelock.mock.lock.withArgs(0).returns(true);
-            expect(l1TimelockTrigger.push(0, newPrizeDistribution()))
+            await drawCalculatorTimelock.mock.lock.withArgs(Draw.drawId, Draw.timestamp + Draw.beaconPeriodSeconds).returns(true);
+            expect(l1TimelockTrigger.push(Draw, newPrizeDistribution()))
                 .to.emit(l1TimelockTrigger, 'PrizeDistributionPushed');
         });
 
         it('should not allow a push from a non-owner', async () => {
+            const Draw = {
+                drawId: 1,
+                timestamp: 22,
+                winningRandomNumber: 333,
+                beaconPeriodStartedAt: 4444,
+                beaconPeriodSeconds: 55555,
+            }
             await expect(
-                l1TimelockTrigger.connect(wallet2).push(0, newPrizeDistribution()),
+                l1TimelockTrigger.connect(wallet2).push(Draw, newPrizeDistribution()),
             ).to.be.revertedWith('Manageable/caller-not-manager-or-owner');
         });
 
         it('should not allow a push if a draw is still timelocked', async () => {
+            const Draw = {
+                drawId: 1,
+                timestamp: 22,
+                winningRandomNumber: 333,
+                beaconPeriodStartedAt: 4444,
+                beaconPeriodSeconds: 55555,
+            }
+
             await drawCalculatorTimelock.mock.lock
-                .withArgs(0)
+                .withArgs(Draw.drawId, Draw.timestamp + Draw.beaconPeriodSeconds)
                 .revertsWithReason('OM/timelock-not-expired');
 
             await prizeDistributionBuffer.mock.pushPrizeDistribution.returns(0);
 
-            await expect(l1TimelockTrigger.push(0, newPrizeDistribution())).to.be.revertedWith(
+            await expect(l1TimelockTrigger.push(Draw, newPrizeDistribution())).to.be.revertedWith(
                 'OM/timelock-not-expired',
             );
         });

diff --git a/test/L2TimelockTrigger.test.ts b/test/L2TimelockTrigger.test.ts
@@ -72,7 +72,7 @@ describe('L2TimelockTrigger', () => {
         it('should allow a push when no push has happened', async () => {
             await drawBuffer.mock.pushDraw.returns(draw.drawId);
             await prizeDistributionBuffer.mock.pushPrizeDistribution.returns(true);
-            await drawCalculatorTimelock.mock.lock.withArgs(0).returns(true);
+            await drawCalculatorTimelock.mock.lock.returns(true);
             await expect(l2TimelockTrigger.push(draw, newPrizeDistribution()))
                 .to.emit(l2TimelockTrigger, 'DrawAndPrizeDistributionPushed');
         });
@@ -85,7 +85,6 @@ describe('L2TimelockTrigger', () => {
 
         it('should not allow a push if a draw is still timelocked', async () => {
             await drawCalculatorTimelock.mock.lock
-                .withArgs(0)
                 .revertsWithReason('OM/timelock-not-expired');
 
             await drawBuffer.mock.pushDraw.returns(draw.drawId);

diff --git a/test/helpers/prizeDistribution.ts b/test/helpers/prizeDistribution.ts
@@ -26,12 +26,12 @@ const prizeDistribution: PrizeDistribution = {
 };
 
 export const newPrizeDistribution = (cardinality: number = 5): any => {
-    const distributions = [...ZERO_DISTRIBUTIONS];
-    distributions[0] = ethers.utils.parseUnits("0.5", 9);
+    const tiers = [...ZERO_DISTRIBUTIONS];
+    tiers[0] = ethers.utils.parseUnits("0.5", 9);
 
     return {
         ...prizeDistribution,
-        distributions,
+        tiers,
         matchCardinality: BigNumber.from(cardinality),
     };
 };