From 7bae575fa1634a47c86247b84947c5170ef2010d Mon Sep 17 00:00:00 2001
From: Roman Agureev <Agureev.RS@phystech.edu>
Date: Wed, 12 Jun 2024 16:47:40 +0300
Subject: [PATCH] fix: double spend on lock just in time

---
 contracts/FeeDistributor.vy                   | 12 ++++-----
 .../FeeDistribution/test_fee_distribution.py  | 25 +++++++++++++++++++
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/contracts/FeeDistributor.vy b/contracts/FeeDistributor.vy
index 2f7a38dc..9c4138b8 100644
--- a/contracts/FeeDistributor.vy
+++ b/contracts/FeeDistributor.vy
@@ -1,4 +1,4 @@
-# @version 0.2.7
+# @version 0.3.7
 """
 @title Curve Fee Distribution
 @author Curve Finance
@@ -80,7 +80,7 @@ def __init__(
     @notice Contract constructor
     @param _voting_escrow VotingEscrow contract address
     @param _start_time Epoch time for fee distribution to start
-    @param _token Fee token address (3CRV)
+    @param _token Fee token address (crvUSD)
     @param _admin Admin address
     @param _emergency_return Address to transfer `_token` balance to
                              if this contract is killed
@@ -193,7 +193,7 @@ def ve_for_at(_user: address, _timestamp: uint256) -> uint256:
 def _checkpoint_total_supply():
     ve: address = self.voting_escrow
     t: uint256 = self.time_cursor
-    rounded_timestamp: uint256 = block.timestamp / WEEK * WEEK
+    rounded_timestamp: uint256 = (block.timestamp - 1) / WEEK * WEEK
     VotingEscrow(ve).checkpoint()
 
     for i in range(20):
@@ -374,8 +374,8 @@ def claim_many(_receivers: address[20]) -> bool:
 @external
 def burn(_coin: address) -> bool:
     """
-    @notice Receive 3CRV into the contract and trigger a token checkpoint
-    @param _coin Address of the coin being received (must be 3CRV)
+    @notice Receive crvUSD into the contract and trigger a token checkpoint
+    @param _coin Address of the coin being received (must be crvUSD)
     @return bool success
     """
     assert _coin == self.token
@@ -428,7 +428,7 @@ def toggle_allow_checkpoint_token():
 def kill_me():
     """
     @notice Kill the contract
-    @dev Killing transfers the entire 3CRV balance to the emergency return address
+    @dev Killing transfers the entire crvUSD balance to the emergency return address
          and blocks the ability to claim or burn. The contract cannot be unkilled.
     """
     assert msg.sender == self.admin
diff --git a/tests/unitary/FeeDistribution/test_fee_distribution.py b/tests/unitary/FeeDistribution/test_fee_distribution.py
index d5a050f1..c711a9da 100644
--- a/tests/unitary/FeeDistribution/test_fee_distribution.py
+++ b/tests/unitary/FeeDistribution/test_fee_distribution.py
@@ -1,3 +1,6 @@
+import pytest
+
+
 DAY = 86400
 WEEK = 7 * DAY
 
@@ -139,3 +142,25 @@ def test_deposited_parallel(web3, chain, accounts, voting_escrow, fee_distributo
     balance_bob = coin_a.balanceOf(bob)
     assert balance_alice == balance_bob
     assert abs(balance_alice + balance_bob - 10 ** 19) < 20
+
+
+def test_checkpoint_and_lock(CheckpointLock, accounts, chain, fee_distributor, voting_escrow, token, coin_a):
+    distributor = fee_distributor()
+    distributor.toggle_allow_checkpoint_token()
+    contract = CheckpointLock.deploy(distributor, voting_escrow, token, {"from": accounts[0]})
+    voting_escrow.commit_smart_wallet_checker(contract, {"from": voting_escrow.admin()})
+    voting_escrow.apply_smart_wallet_checker({"from": voting_escrow.admin()})
+
+    amount0, amount1 = 10 ** 18, 10 ** 18
+    token.approve(contract, amount0 + amount1, {"from": accounts[0]})
+
+    rounded_ts = (chain.time() + WEEK - 1) // WEEK * WEEK
+    chain.mine(timestamp=rounded_ts - 1)
+    contract.checkpoint_and_lock(amount0, amount1, {"from": accounts[0]})
+    distributor.checkpoint_token()
+
+    coin_a._mint_for_testing(distributor, 10 ** 18, {"from": accounts[0]})
+    chain.mine(timestamp=rounded_ts + WEEK)
+
+    amount = distributor.claim(contract, {"from": accounts[0]}).return_value
+    assert int(amount) == pytest.approx(10 ** 18, rel=1e-5)