Restricted snapBalances when consolidation is in progress

Author: naddison36

contracts/contracts/strategies/NativeStaking/CompoundingValidatorManager.sol

@@ -957,7 +957,7 @@ abstract contract CompoundingValidatorManager is Governable, Pausable {
     /// This behaviour is correct.
     ///
     /// The validator balances on the beacon chain can then be proved with `verifyBalances`.
-    function snapBalances() external {
+    function snapBalances() external onlyRegistrator {
         uint64 currentTimestamp = SafeCast.toUint64(block.timestamp);
         require(
             snappedBalance.timestamp + SNAP_BALANCES_DELAY < currentTimestamp,

contracts/contracts/strategies/NativeStaking/ConsolidationController.sol

@@ -287,6 +287,17 @@ contract ConsolidationController is Ownable {
      *
      */
 
+    /// @notice Forwards to the new Compounding Staking Strategy.
+    /// Is only callable by the validator registrator when a consolidation is in progress.
+    /// Anyone can call when there are no consolidations in progress.
+    function snapBalances() external {
+        if (consolidationCount > 0 && msg.sender != validatorRegistrator) {
+            revert("Consolidation in progress");
+        }
+
+        targetStrategy.snapBalances();
+    }
+
     /**
      * @notice Anyone can verify balances on the new Compounding Staking Strategy
      * as long as there are no consolidations in progress.

contracts/tasks/tasks.js

@@ -2448,10 +2448,14 @@ task("stakeValidator").setAction(async (_, __, runSuper) => {
   return runSuper();
 });
 
-subtask(
-  "snapBalances",
-  "Takes a snapshot of the staking strategy's balance"
-).setAction(snapBalances);
+subtask("snapBalances", "Takes a snapshot of the staking strategy's balance")
+  .addOptionalParam(
+    "consol",
+    "Call the consolidation controller instead of the strategy",
+    false,
+    types.boolean
+  )
+  .setAction(snapBalances);
 task("snapBalances").setAction(async (_, __, runSuper) => {
   return runSuper();
 });

contracts/tasks/validatorCompound.js

@@ -33,18 +33,20 @@ const {
 
 const log = require("../utils/logger")("task:validator:compounding");
 
-async function snapBalances() {
+async function snapBalances({ consol = false }) {
   const signer = await getSigner();
 
   // TODO check the slot of the first pending deposit is not zero
 
-  const strategy = await resolveContract(
-    "CompoundingStakingSSVStrategyProxy",
-    "CompoundingStakingSSVStrategy"
-  );
+  const contract = consol
+    ? await resolveContract("ConsolidationController")
+    : await resolveContract(
+        "CompoundingStakingSSVStrategyProxy",
+        "CompoundingStakingSSVStrategy"
+      );
 
-  log(`About to snap balances for strategy ${strategy.address}`);
-  const tx = await strategy.connect(signer).snapBalances();
+  log(`About to snap balances on ${contract.address}`);
+  const tx = await contract.connect(signer).snapBalances();
   await logTxDetails(tx, "snapBalances");
 
   const receipt = await tx.wait();

contracts/test/strategies/stakingConsolidation.mainnet.fork-test.js

@@ -201,7 +201,7 @@ describe("ForkTest: Consolidation of Staking Strategies", function () {
       "0x67d077e55a1f4af124a8df1e1ac720d1e1e557a5305d774c5ebf98be3479cf8d"
     );
 
-    await compoundingStakingStrategy.connect(registratorSigner).snapBalances();
+    await consolidationController.connect(registratorSigner).snapBalances();
 
     await consolidationController
       .connect(registratorSigner)
@@ -606,9 +606,17 @@ describe("ForkTest: Consolidation of Staking Strategies", function () {
 
       await expect(tx).to.be.revertedWith("No consolidation in progress");
     });
+    it("Should call snapBalance on the Consolidation Controller by anyone when no consolidation in progress", async () => {
+      const { josh } = fixture;
+      await advanceTime(12 * 40);
+
+      const tx = await consolidationController.connect(josh).snapBalances();
+
+      await expect(tx).to.emit(compoundingStakingStrategy, "BalancesSnapped");
+    });
     it("Fail to directly call verifyBalance on the Compounding Staking Strategy", async () => {
       await advanceTime(12 * 40);
-      await compoundingStakingStrategy.snapBalances();
+      await consolidationController.snapBalances();
 
       const tx = compoundingStakingStrategy
         .connect(registratorSigner)
@@ -795,17 +803,26 @@ describe("ForkTest: Consolidation of Staking Strategies", function () {
       await expect(tx).to.emit(compoundingStakingStrategy, "BalancesVerified");
     });
     // Balance proofs after the deposit to validator 13498458 has been verified.
-    it("Should call snapBalance on the Compounding Staking Strategy", async () => {
-      const { josh } = fixture;
+    it("Should call snapBalance on the Consolidation Controller by the Registrator after the consolidation has started", async () => {
       await advanceTime(12 * 40);
 
-      const tx = await compoundingStakingStrategy.connect(josh).snapBalances();
+      const tx = await consolidationController
+        .connect(registratorSigner)
+        .snapBalances();
 
       await expect(tx).to.emit(compoundingStakingStrategy, "BalancesSnapped");
     });
-    it("Fail to verifyBalance of a snapshot after the consolidation was started", async () => {
+    it("Fail snapBalance on the new compounding staking strategy when called by non-registrator after the consolidation has started", async () => {
+      const { josh } = fixture;
+      await advanceTime(12 * 40);
+
+      const tx = compoundingStakingStrategy.connect(josh).snapBalances();
+
+      await expect(tx).to.be.revertedWith("Not Registrator");
+    });
+    it("Fail to verifyBalance of a snapshot after the consolidation has started", async () => {
       await advanceTime(12 * 40);
-      await compoundingStakingStrategy.snapBalances();
+      await consolidationController.connect(registratorSigner).snapBalances();
 
       const tx = consolidationController
         .connect(registratorSigner)
@@ -1040,7 +1057,7 @@ describe("ForkTest: Consolidation of Staking Strategies", function () {
         beaconBlockRoot
       );
 
-      await compoundingStakingStrategy.connect(fixture.josh).snapBalances();
+      await consolidationController.connect(registratorSigner).snapBalances();
     });
 
     it("Should confirm consolidation", async () => {