Add more invariants

src/PufferProtocol.sol

@@ -134,9 +134,10 @@ contract PufferProtocol is IPufferProtocol, AccessManagedUpgradeable, UUPSUpgrad
         validator.signature = data.signature;
         validator.status = Status.PENDING;
         validator.strategy = address($.strategies[strategyName]);
-        validator.bond = SafeCastLib.toUint64(pufETHReceived);
-        validator.monthsCommited = SafeCastLib.toUint40(numberOfMonths);
-        validator.commitmentAmount = SafeCastLib.toUint64(msg.value - validatorBond);
+        // No need for Safecast because of the validations inside of _checkValidatorRegistrationInputs
+        validator.bond = uint64(pufETHReceived);
+        validator.monthsCommitted = uint40(numberOfMonths);
+        validator.lastCommitmentPayment = uint64(block.timestamp);
         validator.node = msg.sender;
 
         uint256 validatorIndex = $.pendingValidatorIndicies[strategyName];
@@ -203,23 +204,24 @@ contract PufferProtocol is IPufferProtocol, AccessManagedUpgradeable, UUPSUpgrad
         strategy.callStake({ pubKey: validator.pubKey, signature: validator.signature, depositDataRoot: depositDataRoot });
     }
 
+    /**
+     * @inheritdoc IPufferProtocol
+     */
     function extendCommitment(bytes32 strategyName, uint256 validatorIndex, uint256 numberOfMonths) external payable {
         ProtocolStorage storage $ = _getPufferProtocolStorage();
         Validator storage validator = $.validators[strategyName][validatorIndex];
 
-        if (numberOfMonths > 13) {
-            revert InvalidNumberOfMonths();
-        }
-
+        // Causes panic for invalid numberOfMonths
         uint256 smoothingCommitment = $.smoothingCommitments[numberOfMonths - 1];
 
         // Node operator can purchase commitment for multiple months
         if ((msg.value != smoothingCommitment)) {
             revert InvalidETHAmount();
         }
 
-        validator.monthsCommited += uint40(numberOfMonths);
-        validator.commitmentAmount = uint64(msg.value);
+        // No need for Safecast because of the validatrions above
+        validator.monthsCommitted = uint40(numberOfMonths);
+        validator.lastCommitmentPayment = uint64(block.timestamp);
 
         emit SmoothingCommitmentPaid(validator.pubKey, block.timestamp, msg.value);
 
@@ -294,7 +296,7 @@ contract PufferProtocol is IPufferProtocol, AccessManagedUpgradeable, UUPSUpgrad
         // Remove what we don't
         delete validator.strategy;
         delete validator.node;
-        delete validator.monthsCommited;
+        delete validator.monthsCommitted;
         delete validator.bond;
         delete validator.pubKey;
         delete validator.signature;
@@ -789,6 +791,7 @@ contract PufferProtocol is IPufferProtocol, AccessManagedUpgradeable, UUPSUpgrad
             revert InvalidBLSPublicKeySet();
         }
 
+        // panic for invalid `numberOfMonths`
         uint256 smoothingCommitment = $.smoothingCommitments[numberOfMonths - 1];
 
         if (msg.value != (smoothingCommitment + validatorBond)) {

src/TokenRescuer.sol

@@ -7,7 +7,6 @@ import { IERC1155Receiver } from "openzeppelin/token/ERC1155/IERC1155Receiver.so
 import { IERC1155 } from "openzeppelin/token/ERC1155/IERC1155.sol";
 import { IERC721Receiver } from "openzeppelin/token/ERC721/IERC721Receiver.sol";
 import { IERC721 } from "openzeppelin/token/ERC721/ERC721.sol";
-import { Unauthorized } from "puffer/Errors.sol";
 
 /**
  * @title TokenRescuer

src/interface/IBeaconDepositContract.sol

@@ -6,12 +6,17 @@ pragma solidity >=0.8.0 <0.9.0;
  * https://etherscan.io/address/0x00000000219ab540356cbb839cbe05303d7705fa#code
  */
 interface IBeaconDepositContract {
+    //solhint-disable-next-line func-param-name-mixedcase
     event DepositEvent(bytes pubkey, bytes withdrawal_credentials, bytes amount, bytes signature, bytes index);
 
     function deposit(
+        //solhint-disable-next-line func-param-name-mixedcase
         bytes memory pubkey,
+        //solhint-disable-next-line func-param-name-mixedcase
         bytes memory withdrawal_credentials,
+        //solhint-disable-next-line func-param-name-mixedcase
         bytes memory signature,
+        //solhint-disable-next-line func-param-name-mixedcase
         bytes32 deposit_data_root
     ) external payable;
     function get_deposit_count() external view returns (bytes memory);

src/interface/IEnclaveVerifier.sol

@@ -2,7 +2,6 @@
 pragma solidity >=0.8.0 <0.9.0;
 
 import { RaveEvidence } from "puffer/struct/RaveEvidence.sol";
-import { Unauthorized } from "puffer/Errors.sol";
 
 /**
  * @title IEnclaveVerifier interface

src/interface/IPufferPool.sol

@@ -2,7 +2,6 @@
 pragma solidity >=0.8.0 <0.9.0;
 
 import { IERC20 } from "openzeppelin/token/ERC20/IERC20.sol";
-import { IERC20Permit } from "openzeppelin/token/ERC20/extensions/IERC20Permit.sol";
 
 /**
  * @title IPufferPool

src/interface/IPufferProtocol.sol

@@ -399,6 +399,14 @@ interface IPufferProtocol is IPufferProtocolStorage {
         external
         payable;
 
+    /**
+     * @notice Extends the commitment for a validator in a specific strategy
+     * @param strategyName The name of the strategy
+     * @param validatorIndex The index of the validator in the strategy
+     * @param numberOfMonths The number of months to extend the commitment for
+     */
+    function extendCommitment(bytes32 strategyName, uint256 validatorIndex, uint256 numberOfMonths) external payable;
+
     /**
      * @notice Returns the pending validator index for `strategyName`
      */

src/struct/Validator.sol

@@ -9,8 +9,8 @@ import { Status } from "puffer/struct/Status.sol";
 struct Validator {
     address node; // Address of the Node operator
     address strategy; // In which strategy is the Validator participating
-    uint40 monthsCommited; // Date when the smoothing commitment was paid
-    uint64 commitmentAmount; // ETH amount that the Node operator paid
+    uint40 monthsCommitted; // Number of months
+    uint64 lastCommitmentPayment; // block.timestamp of the last payment
     uint64 bond; // Validator bond (in pufETH)
     Status status; // Validator status
     bytes pubKey; // Validator public key

test/handlers/PufferProtocolHandler.sol

@@ -48,6 +48,8 @@ contract PufferProtocolHandler is Test {
     uint256 public ghost_locked_amount;
     uint256 public ghost_eth_rewards_amount;
     uint256 public ghost_block_number = 1;
+    uint256 public ghost_validators = 0;
+    uint256 public ghost_pufETH_bond_amount = 0; // bond amount that should be in puffer protocol
 
     // Previous ETH balance of PufferPool
     uint256 public previousBalance;
@@ -258,24 +260,28 @@ contract PufferProtocolHandler is Test {
         bytes32[] memory strategyWeights = pufferProtocol.getStrategyWeights();
         uint256 strategyIndex = strategySelectorSeed % strategyWeights.length;
 
-        bytes32 startegyName = strategyWeights[strategyIndex];
+        bytes32 strategyName = strategyWeights[strategyIndex];
 
         vm.deal(nodeOperator, 5 ether);
         vm.startPrank(nodeOperator);
 
-        uint256 depositedETHAmount = _registerValidatorKey(pubKeyPart, startegyName);
+        uint256 validatorIndex = pufferProtocol.getPendingValidatorIndex(strategyName);
+
+        uint256 depositedETHAmount = _registerValidatorKey(pubKeyPart, strategyName);
 
         // Store data and push to queue
         ProvisioningData memory validator;
         validator.status = Status.PENDING;
         validator.pubKeypart = pubKeyPart;
 
-        _validatorQueue[startegyName].push(validator);
+        _validatorQueue[strategyName].push(validator);
 
         vm.stopPrank();
 
         // Account for that deposited eth in ghost variable
         ghost_eth_deposited_amount += depositedETHAmount;
+        ghost_validators += 1;
+        ghost_pufETH_bond_amount += pool.calculateETHToPufETHAmount(1 ether);
 
         // Add node operator to the set
         _nodeOperators.add(nodeOperator);

test/invariant/PufferProtocolInvariants.sol

@@ -19,7 +19,7 @@ contract PufferProtocolInvariants is TestHelper {
         targetContract(address(handler));
     }
 
-    // // Guardian multisig is not supposed to change
+    // Guardian multisig is not supposed to change
     function invariant_guardiansCanNeverChange() public {
         assertTrue(address(guardiansSafe) == address(pufferProtocol.GUARDIANS()));
     }
@@ -33,6 +33,33 @@ contract PufferProtocolInvariants is TestHelper {
         }
     }
 
+    // Make sure that the pufETH doesn't disappear
+    function invariant_pufferProtocolBond() public {
+        // Validate against ghost variable
+        uint256 pufETHinProtocol = pool.balanceOf(address(pufferProtocol));
+        assertEq(handler.ghost_pufETH_bond_amount(), pufETHinProtocol, "missing bond from the protocol");
+
+        // Validate by calculating eth
+        uint256 ethAmount = pool.calculatePufETHtoETHAmount(pufETHinProtocol);
+        uint256 originalETHAmountDeposited = (handler.ghost_validators() * 1 ether);
+
+        // If the eth amount is lower than the original eth deposited, it is because of the rounding down (calculation for when we pay out users)
+        if (ethAmount < originalETHAmountDeposited) {
+            assertApproxEqRel(
+                ethAmount,
+                originalETHAmountDeposited,
+                0.01e18,
+                "bond should be worth more than the number of validators depositing"
+            );
+        } else {
+            assertGe(
+                ethAmount,
+                originalETHAmountDeposited,
+                "bond should be worth more than the number of validators depositing"
+            );
+        }
+    }
+
     // pufETH should always be worth more than ETH
     function invariant_pufEthToETHRate() public {
         // Exchange rate should always be bigger than 1:1, we are not supposed to be losing anything with this setup ATM

test/unit/PufferProtocol.t.sol

@@ -58,6 +58,7 @@ contract PufferProtocolTest is TestHelper {
     // Setup
     function testSetup() public {
         assertTrue(address(pufferProtocol.getWithdrawalPool()) != address(0), "non zero address");
+        assertTrue(address(pufferProtocol.getPufferPool()) != address(0), "pufer pool address");
         address strategy = pufferProtocol.getStrategyAddress(NO_RESTAKING);
         assertEq(PufferStrategy(payable(strategy)).NAME(), NO_RESTAKING, "bad name");
     }
@@ -77,9 +78,11 @@ contract PufferProtocolTest is TestHelper {
         _registerValidatorKey(bytes32("bob"), NO_RESTAKING);
 
         (bytes32 strategyName, uint256 idx) = pufferProtocol.getNextValidatorToProvision();
+        uint256 strategySelecitonIdx = pufferProtocol.getStrategySelectIndex();
 
         assertEq(strategyName, NO_RESTAKING, "strategy");
         assertEq(idx, 0, "idx");
+        assertEq(strategySelecitonIdx, 0, "strategy selection idx");
 
         assertTrue(pool.balanceOf(address(this)) == 0, "zero pufETH");
 
@@ -100,6 +103,8 @@ contract PufferProtocolTest is TestHelper {
         vm.expectEmit(true, true, true, true);
         emit SuccesfullyProvisioned(_getPubKey(bytes32("bob")), 1, NO_RESTAKING);
         pufferProtocol.provisionNode(_getGuardianSignatures(_getPubKey(bytes32("bob"))));
+        strategySelecitonIdx = pufferProtocol.getStrategySelectIndex();
+        assertEq(strategySelecitonIdx, 1, "strategy idx changed");
     }
 
     // Create an existing strategy should revert
@@ -202,30 +207,27 @@ contract PufferProtocolTest is TestHelper {
         assertEq(idx, 8, "idx");
     }
 
-    // // Test extending validator commitment
-    // function testExtendCommitment() public {
-    //     _registerValidatorKey(bytes32("alice"), NO_RESTAKING);
-
-    //     Validator memory validator = pufferProtocol.getValidatorInfo(NO_RESTAKING, 0);
-    //     assertTrue(validator.node == address(this), "node operator");
-
-    //     uint256 firstPayment = validator.commitmentExpiration;
-    //     assertEq(firstPayment, block.timestamp + 30 days, "lastPayment");
+    // Test extending validator commitment
+    function testExtendCommitment() public {
+        _registerValidatorKey(bytes32("alice"), NO_RESTAKING);
 
-    //     vm.warp(1000);
+        Validator memory validator = pufferProtocol.getValidatorInfo(NO_RESTAKING, 0);
+        assertTrue(validator.node == address(this), "node operator");
 
-    //     vm.expectRevert();
-    //     pufferProtocol.extendCommitment{ value: 0 }(NO_RESTAKING, 0);
+        vm.warp(1000);
 
-    //     vm.expectRevert(IPufferProtocol.InvalidETHAmount.selector);
-    //     pufferProtocol.extendCommitment{ value: 5 ether }(NO_RESTAKING, 0);
+        // Amounts dont match
+        vm.expectRevert(IPufferProtocol.InvalidETHAmount.selector);
+        pufferProtocol.extendCommitment{ value: 5 ether }(NO_RESTAKING, 0, 5);
 
-    //     pufferProtocol.extendCommitment{ value: pufferProtocol.getSmoothingCommitment(NO_RESTAKING) }(NO_RESTAKING, 0);
+        // Should extend
+        pufferProtocol.extendCommitment{ value: pufferProtocol.getSmoothingCommitment(5) }(NO_RESTAKING, 0, 5);
 
-    //     validator = pufferProtocol.getValidatorInfo(NO_RESTAKING, 0);
+        validator = pufferProtocol.getValidatorInfo(NO_RESTAKING, 0);
 
-    //     assertTrue(validator.commitmentExpiration == block.timestamp + 30 days, "lastPayment");
-    // }
+        assertTrue(validator.monthsCommitted == 5, "lastPayment");
+        assertTrue(validator.lastCommitmentPayment == block.timestamp, "lastPayment");
+    }
 
     // Try updating for future block
     function testProofOfReserve() external {
@@ -683,6 +685,17 @@ contract PufferProtocolTest is TestHelper {
             merkleProof: aliceProof
         });
 
+        // Try again, now the validator will be in invalid state
+        vm.expectRevert(abi.encodeWithSelector(IPufferProtocol.InvalidValidatorState.selector, Status.EXITED));
+        pufferProtocol.stopValidator({
+            strategyName: NO_RESTAKING,
+            validatorIndex: 0,
+            blockNumber: 100,
+            withdrawalAmount: 32.14 ether,
+            wasSlashed: false,
+            merkleProof: aliceProof
+        });
+
         assertEq(pool.balanceOf(alice), 1 ether, "alice received back the bond in pufETH");
 
         bytes32[] memory bobProof = new bytes32[](1);
@@ -811,7 +824,7 @@ contract PufferProtocolTest is TestHelper {
             }),
             blsEncryptedPrivKeyShares: new bytes[](3),
             blsPubKeySet: new bytes(48),
-            raveEvidence: new bytes(1) // Guardians are checking it off chain
+            raveEvidence: bytes("mock rave") // Guardians are checking it off chain
          });
 
         return validatorData;