Prevent CELO Transfer to Unreleased Treasury

packages/protocol/contracts-0.8/common/CeloUnreleasedTreasury.sol

@@ -5,15 +5,27 @@ import "@openzeppelin/contracts8/security/ReentrancyGuard.sol";
 import "@openzeppelin/contracts8/utils/math/Math.sol";
 
 import "./UsingRegistry.sol";
-import "../common/IsL2Check.sol";
 
 import "../../contracts/common/Initializable.sol";
 import "./interfaces/ICeloUnreleasedTreasuryInitializer.sol";
 
 /**
  * @title Contract for unreleased Celo tokens.
+ * @notice This contract is not allowed to receive transfers of CELO,
+ * to avoid miscalculating the epoch rewards and to prevent any malicious actor
+ * from routing stolen fund through the epoch reward distribution.
  */
-contract CeloUnreleasedTreasury is UsingRegistry, ReentrancyGuard, Initializable, IsL2Check {
+contract CeloUnreleasedTreasury is
+  ICeloUnreleasedTreasuryInitializer,
+  UsingRegistry,
+  ReentrancyGuard,
+  Initializable
+{
+  bool internal hasAlreadyReleased;
+
+  // Remaining epoch rewards to distribute.
+  uint256 internal remainingBalanceToRelease;
+
   event Released(address indexed to, uint256 amount);
 
   modifier onlyEpochManager() {
@@ -46,11 +58,31 @@ contract CeloUnreleasedTreasury is UsingRegistry, ReentrancyGuard, Initializable
    * @param amount The amount to release.
    */
   function release(address to, uint256 amount) external onlyEpochManager {
-    require(address(this).balance >= amount, "Insufficient balance.");
+    if (!hasAlreadyReleased) {
+      remainingBalanceToRelease = address(this).balance;
+      hasAlreadyReleased = true;
+    }
+
+    require(remainingBalanceToRelease >= amount, "Insufficient balance.");
+    remainingBalanceToRelease -= amount;
     require(getCeloToken().transfer(to, amount), "CELO transfer failed.");
+
     emit Released(to, amount);
   }
 
+  /**
+   * @notice Returns the remaining balance this contract has left to release.
+   * @dev This uses internal accounting of the released balance,
+   * to avoid recounting CELO that was transferred back to this contract.
+   */
+  function getRemainingBalanceToRelease() external view returns (uint256) {
+    if (!hasAlreadyReleased) {
+      return address(this).balance;
+    } else {
+      return remainingBalanceToRelease;
+    }
+  }
+
   /**
    * @notice Returns the storage, major, minor, and patch version of the contract.
    * @return Storage version of the contract.

packages/protocol/contracts-0.8/common/test/MockCeloUnreleasedTreasury.sol

@@ -8,8 +8,28 @@ import "../UsingRegistry.sol";
  * @title A mock CeloUnreleasedTreasury for testing.
  */
 contract MockCeloUnreleasedTreasury is ICeloUnreleasedTreasury, UsingRegistry {
+  bool internal hasAlreadyReleased;
+  uint256 internal remainingTreasure;
   function release(address to, uint256 amount) external {
-    require(address(this).balance >= amount, "Insufficient balance.");
+    if (!hasAlreadyReleased) {
+      remainingTreasure = address(this).balance;
+      hasAlreadyReleased = true;
+    }
+
+    require(remainingTreasure >= amount, "Insufficient balance.");
     require(getCeloToken().transfer(to, amount), "CELO transfer failed.");
+    remainingTreasure -= amount;
+  }
+
+  function getRemainingBalanceToRelease() external view returns (uint256) {
+    remainingTreasure;
+  }
+
+  function setRemainingTreasure(uint256 _amount) public {
+    remainingTreasure = _amount;
+  }
+
+  function setFirstRelease(bool _hasAlreadyReleased) public {
+    hasAlreadyReleased = _hasAlreadyReleased;
   }
 }

packages/protocol/contracts/common/GoldToken.sol

@@ -10,7 +10,6 @@ import "./Initializable.sol";
 import "./interfaces/ICeloToken.sol";
 import "./interfaces/ICeloTokenInitializer.sol";
 import "./interfaces/ICeloVersionedContract.sol";
-import "./interfaces/ICeloUnreleasedTreasury.sol";
 import "../../contracts-0.8/common/IsL2Check.sol";
 
 contract GoldToken is
@@ -270,8 +269,7 @@ contract GoldToken is
    */
   function allocatedSupply() public view returns (uint256) {
     if (isL2()) {
-      return
-        CELO_SUPPLY_CAP - registry.getAddressForOrDie(CELO_UNRELEASED_TREASURY_REGISTRY_ID).balance;
+      return CELO_SUPPLY_CAP - getCeloUnreleasedTreasury().getRemainingBalanceToRelease();
     } else {
       return totalSupply();
     }

packages/protocol/contracts/common/interfaces/ICeloUnreleasedTreasury.sol

@@ -8,4 +8,6 @@ interface ICeloUnreleasedTreasury {
    * @param amount The amount to release.
    */
   function release(address to, uint256 amount) external;
+
+  function getRemainingBalanceToRelease() external view returns (uint256);
 }

packages/protocol/test-sol/devchain/migration/Migration.t.sol

@@ -22,6 +22,8 @@ import "@celo-contracts/governance/interfaces/IValidators.sol";
 import "@celo-contracts-8/common/interfaces/IPrecompiles.sol";
 import "@celo-contracts-8/common/interfaces/IScoreManager.sol";
 
+import "@openzeppelin/contracts8/token/ERC20/IERC20.sol";
+
 contract IntegrationTest is Test, TestConstants, Utils08 {
   IRegistry registry = IRegistry(REGISTRY_ADDRESS);
 
@@ -236,6 +238,19 @@ contract EpochManagerIntegrationTest is IntegrationTest, MigrationsConstants {
     epochManager.initializeSystem(100, block.number, firstElected);
   }
 
+  function test_Reverts_whenTransferingCeloToUnreleasedTreasury() public {
+    _MockL2Migration(validatorsList);
+
+    blockTravel(vm, 43200);
+    timeTravel(vm, DAY);
+
+    IERC20 _celoToken = IERC20(address(celoToken));
+    vm.prank(randomAddress);
+
+    (bool success, ) = address(unreleasedTreasury).call{ value: 50000 ether }("");
+    assertFalse(success);
+  }
+
   function test_SetsCurrentRewardBlock() public {
     _MockL2Migration(validatorsList);
 
@@ -245,8 +260,9 @@ contract EpochManagerIntegrationTest is IntegrationTest, MigrationsConstants {
     epochManager.startNextEpochProcess();
 
     (, , , uint256 _currentRewardsBlock) = epochManager.getCurrentEpoch();
-
+    (uint256 status, , , , ) = epochManager.getEpochProcessingState();
     assertEq(_currentRewardsBlock, block.number);
+    assertEq(status, 1);
   }
 
   function _MockL2Migration(address[] memory _validatorsList) internal {

packages/protocol/test-sol/unit/common/CeloUnreleasedTreasury.t.sol

@@ -62,8 +62,9 @@ contract CeloUnreleasedTreasuryTest is Test, TestConstants, IsL2Check {
     deployCodeTo("Registry.sol", abi.encode(false), REGISTRY_ADDRESS);
     registry = IRegistry(REGISTRY_ADDRESS);
 
-    deployCodeTo("GoldToken.sol", abi.encode(false), celoTokenAddress);
+    deployCodeTo("GoldToken.sol", abi.encode(true), celoTokenAddress);
     celoToken = ICeloToken(celoTokenAddress);
+    celoToken.initialize(REGISTRY_ADDRESS);
     // Using a mock contract, as foundry does not allow for library linking when using deployCodeTo
     governance = new MockGovernance();
 
@@ -156,3 +157,39 @@ contract CeloUnreleasedTreasuryTest_release is CeloUnreleasedTreasuryTest {
     celoUnreleasedTreasury.release(randomAddress, 4);
   }
 }
+contract CeloUnreleasedTreasuryTest_getRemainingBalanceToRelease is CeloUnreleasedTreasuryTest {
+  uint256 _startingBalance;
+  function setUp() public override {
+    super.setUp();
+    newCeloUnreleasedTreasury();
+    _startingBalance = address(celoUnreleasedTreasury).balance;
+  }
+
+  function test_ShouldReturnContractBalanceBeforeFirstRelease() public {
+    uint256 _remainingBalance = celoUnreleasedTreasury.getRemainingBalanceToRelease();
+
+    assertEq(_startingBalance, _remainingBalance);
+  }
+
+  function test_ShouldReturnRemainingBalanceToReleaseAfterFirstRelease() public {
+    vm.prank(epochManagerAddress);
+
+    celoUnreleasedTreasury.release(randomAddress, 4);
+    uint256 _remainingBalance = celoUnreleasedTreasury.getRemainingBalanceToRelease();
+    assertEq(_remainingBalance, _startingBalance - 4);
+  }
+
+  function test_RemainingBalanceToReleaseShouldRemainUnchangedAfterCeloTransferBackToContract()
+    public
+  {
+    vm.prank(epochManagerAddress);
+
+    celoUnreleasedTreasury.release(randomAddress, 4);
+    uint256 _remainingBalanceBeforeTransfer = celoUnreleasedTreasury.getRemainingBalanceToRelease();
+    assertEq(_remainingBalanceBeforeTransfer, _startingBalance - 4);
+    // set the contract balance to mock a CELO token transfer
+    vm.deal(address(celoUnreleasedTreasury), L2_INITIAL_STASH_BALANCE);
+    uint256 _remainingBalanceAfterTransfer = celoUnreleasedTreasury.getRemainingBalanceToRelease();
+    assertEq(_remainingBalanceAfterTransfer, _remainingBalanceBeforeTransfer);
+  }
+}

packages/protocol/test-sol/unit/common/EpochManager.t.sol

@@ -7,7 +7,6 @@ import "@celo-contracts-8/stability/test/MockStableToken.sol";
 import "@celo-contracts-8/common/test/MockCeloToken.sol";
 import "@celo-contracts/common/interfaces/ICeloToken.sol";
 import "@celo-contracts-8/common/ScoreManager.sol";
-import { CeloUnreleasedTreasury } from "@celo-contracts-8/common/CeloUnreleasedTreasury.sol";
 import { ICeloUnreleasedTreasury } from "@celo-contracts/common/interfaces/ICeloUnreleasedTreasury.sol";
 
 import { TestConstants } from "@test-sol/constants.sol";

packages/protocol/test-sol/unit/common/GoldToken.t.sol

@@ -15,7 +15,7 @@ contract GoldTokenTest is Test, TestConstants, IsL2Check {
   address receiver;
   address sender;
   address celoTokenOwner;
-  address celoTokenDistributionSchedule;
+  address celoUnreleasedTreasuryAddress;
 
   event Transfer(address indexed from, address indexed to, uint256 value);
   event TransferComment(string comment);
@@ -26,16 +26,17 @@ contract GoldTokenTest is Test, TestConstants, IsL2Check {
   }
 
   function setUp() public {
+    celoTokenOwner = actor("celoTokenOwner");
+    celoUnreleasedTreasuryAddress = actor("celoUnreleasedTreasury");
     deployCodeTo("Registry.sol", abi.encode(false), REGISTRY_ADDRESS);
+    deployCodeTo("CeloUnreleasedTreasury.sol", abi.encode(false), celoUnreleasedTreasuryAddress);
     registry = IRegistry(REGISTRY_ADDRESS);
 
-    celoTokenOwner = actor("celoTokenOwner");
-    celoTokenDistributionSchedule = actor("celoTokenDistributionSchedule");
     vm.prank(celoTokenOwner);
     celoToken = new GoldToken(true);
     vm.prank(celoTokenOwner);
     celoToken.setRegistry(REGISTRY_ADDRESS);
-    registry.setAddressFor("CeloUnreleasedTreasury", celoTokenDistributionSchedule);
+    registry.setAddressFor("CeloUnreleasedTreasury", celoUnreleasedTreasuryAddress);
     receiver = actor("receiver");
     sender = actor("sender");
     vm.deal(receiver, ONE_CELOTOKEN);
@@ -126,6 +127,26 @@ contract GoldTokenTest_transfer is GoldTokenTest {
     vm.expectRevert();
     celoToken.transfer(address(0), ONE_CELOTOKEN);
   }
+
+  function test_Succeeds_whenTransferingToCeloUnreleasedTreasury() public {
+    vm.prank(sender);
+    uint256 balanceBefore = celoToken.balanceOf(celoUnreleasedTreasuryAddress);
+
+    celoToken.transfer(celoUnreleasedTreasuryAddress, ONE_CELOTOKEN);
+    uint256 balanceAfter = celoToken.balanceOf(celoUnreleasedTreasuryAddress);
+    assertGt(balanceAfter, balanceBefore);
+  }
+
+  function test_FailsWhenNativeTransferingToCeloUnreleasedTreasury() public payable {
+    (bool success, ) = address(uint160(celoUnreleasedTreasuryAddress)).call.value(ONE_CELOTOKEN)(
+      ""
+    );
+
+    assertFalse(success);
+
+    bool sent = address(uint160(celoUnreleasedTreasuryAddress)).send(ONE_CELOTOKEN);
+    assertFalse(sent);
+  }
 }
 
 contract GoldTokenTest_transferFrom is GoldTokenTest {
@@ -150,6 +171,14 @@ contract GoldTokenTest_transferFrom is GoldTokenTest {
     celoToken.transferFrom(sender, address(0), ONE_CELOTOKEN);
   }
 
+  function test_Succeeds_whenTransferingToCeloUnreleasedTreasury() public {
+    uint256 balanceBefore = celoToken.balanceOf(celoUnreleasedTreasuryAddress);
+    vm.prank(receiver);
+    celoToken.transferFrom(sender, celoUnreleasedTreasuryAddress, ONE_CELOTOKEN);
+    uint256 balanceAfter = celoToken.balanceOf(celoUnreleasedTreasuryAddress);
+    assertGt(balanceAfter, balanceBefore);
+  }
+
   function test_Reverts_WhenTransferMoreThanSenderHas() public {
     uint256 value = sender.balance + ONE_CELOTOKEN * 4;
 
@@ -198,7 +227,7 @@ contract GoldTokenTest_mint is GoldTokenTest {
     vm.expectRevert("Only VM can call");
     celoToken.mint(receiver, ONE_CELOTOKEN);
 
-    vm.prank(celoTokenDistributionSchedule);
+    vm.prank(celoUnreleasedTreasuryAddress);
     vm.expectRevert("Only VM can call");
     celoToken.mint(receiver, ONE_CELOTOKEN);
   }
@@ -220,7 +249,7 @@ contract GoldTokenTest_mint is GoldTokenTest {
 
   function test_Reverts_whenL2() public _whenL2 {
     vm.expectRevert("This method is no longer supported in L2.");
-    vm.prank(celoTokenDistributionSchedule);
+    vm.prank(celoUnreleasedTreasuryAddress);
     celoToken.mint(receiver, ONE_CELOTOKEN);
     vm.expectRevert("This method is no longer supported in L2.");
     vm.prank(address(0));
@@ -249,23 +278,24 @@ contract CeloTokenMockTest is Test, TestConstants {
   GoldTokenMock mockCeloToken;
   uint256 ONE_CELOTOKEN = 1000000000000000000;
   address burnAddress = address(0x000000000000000000000000000000000000dEaD);
-  address celoUnreleasedTreasury;
+  address celoUnreleasedTreasuryAddress = actor("CeloUnreleasedTreasury");
 
   modifier _whenL2() {
     deployCodeTo("Registry.sol", abi.encode(false), PROXY_ADMIN_ADDRESS);
-    vm.deal(celoUnreleasedTreasury, L2_INITIAL_STASH_BALANCE);
+    vm.deal(celoUnreleasedTreasuryAddress, L2_INITIAL_STASH_BALANCE);
     _;
   }
 
   function setUp() public {
     deployCodeTo("Registry.sol", abi.encode(false), REGISTRY_ADDRESS);
+    deployCodeTo("CeloUnreleasedTreasury.sol", abi.encode(false), celoUnreleasedTreasuryAddress);
     registry = IRegistry(REGISTRY_ADDRESS);
 
     mockCeloToken = new GoldTokenMock();
     mockCeloToken.setRegistry(REGISTRY_ADDRESS);
     mockCeloToken.setTotalSupply(L1_MINTED_CELO_SUPPLY);
-    celoUnreleasedTreasury = actor("CeloUnreleasedTreasury");
-    registry.setAddressFor("CeloUnreleasedTreasury", celoUnreleasedTreasury);
+    vm.deal(celoUnreleasedTreasuryAddress, L2_INITIAL_STASH_BALANCE);
+    registry.setAddressFor("CeloUnreleasedTreasury", celoUnreleasedTreasuryAddress);
   }
 }
 
@@ -304,7 +334,7 @@ contract GoldTokenTest_AllocatedSupply is CeloTokenMockTest {
   }
 
   function test_ShouldReturn_WhenWithdrawn_WhenInL2() public _whenL2 {
-    deal(address(celoUnreleasedTreasury), ONE_CELOTOKEN);
+    deal(celoUnreleasedTreasuryAddress, ONE_CELOTOKEN);
     assertEq(mockCeloToken.allocatedSupply(), mockCeloToken.totalSupply() - ONE_CELOTOKEN);
   }
 }