From e7d9dd80a9b515bdea3d06a06a778a7fae57fc9c Mon Sep 17 00:00:00 2001
From: Benjamin <benjaminxh+github@gmail.com>
Date: Fri, 26 Jul 2024 08:45:50 +0200
Subject: [PATCH] add startCheckpoint on module manager

---
 mainnet-contracts/src/PufferModule.sol                |  9 +++++++++
 mainnet-contracts/src/PufferModuleManager.sol         | 11 +++++++++++
 mainnet-contracts/src/interface/IPufferModule.sol     |  5 +++++
 .../src/interface/IPufferModuleManager.sol            |  5 +++++
 4 files changed, 30 insertions(+)

diff --git a/mainnet-contracts/src/PufferModule.sol b/mainnet-contracts/src/PufferModule.sol
index c2f2946..0cf2298 100644
--- a/mainnet-contracts/src/PufferModule.sol
+++ b/mainnet-contracts/src/PufferModule.sol
@@ -189,6 +189,15 @@ contract PufferModule is IPufferModule, Initializable, AccessManagedUpgradeable
         });
     }
 
+    /**
+     * @inheritdoc IPufferModule
+     * @dev Restricted to PufferModuleManager
+     */
+    function startCheckpoint() external virtual onlyPufferModuleManager {
+        ModuleStorage storage $ = _getPufferModuleStorage();
+        $.eigenPod.startCheckpoint({ revertIfNoBalance: true });
+    }
+
     /**
      * @dev Restricted to PufferProtocol
      */
diff --git a/mainnet-contracts/src/PufferModuleManager.sol b/mainnet-contracts/src/PufferModuleManager.sol
index f945680..2307ff7 100644
--- a/mainnet-contracts/src/PufferModuleManager.sol
+++ b/mainnet-contracts/src/PufferModuleManager.sol
@@ -342,6 +342,17 @@ contract PufferModuleManager is IPufferModuleManager, AccessManagedUpgradeable,
         emit CustomCallSucceeded(address(restakingOperator), target, customCalldata, response);
     }
 
+    /**
+     * @inheritdoc IPufferModuleManager
+     * @dev Restricted to the DAO
+     */
+    function callStartCheckpoint(address[] calldata moduleAddresses) external virtual restricted {
+        for (uint256 i = 0; i < moduleAddresses.length; i++) {
+            // reverts if supplied with a duplicate module address
+            IPufferModule(moduleAddresses[i]).startCheckpoint();
+        }
+    }
+
     /**
      * @inheritdoc IPufferModuleManager
      * @dev Restricted to the DAO
diff --git a/mainnet-contracts/src/interface/IPufferModule.sol b/mainnet-contracts/src/interface/IPufferModule.sol
index 7ec96a6..4b1de3d 100644
--- a/mainnet-contracts/src/interface/IPufferModule.sol
+++ b/mainnet-contracts/src/interface/IPufferModule.sol
@@ -113,6 +113,11 @@ interface IPufferModule {
         bool[] calldata receiveAsTokens
     ) external;
 
+    /**
+     * @notice Starts the checkpoint for the module
+     */
+    function startCheckpoint() external;
+
     /**
      * @notice Function callable only by PufferProtocol
      * @param to is the destination address
diff --git a/mainnet-contracts/src/interface/IPufferModuleManager.sol b/mainnet-contracts/src/interface/IPufferModuleManager.sol
index e2e23d0..c163f5c 100644
--- a/mainnet-contracts/src/interface/IPufferModuleManager.sol
+++ b/mainnet-contracts/src/interface/IPufferModuleManager.sol
@@ -219,6 +219,11 @@ interface IPufferModuleManager {
      */
     function callSetProofSubmitter(bytes32 moduleName, address proofSubmitter) external;
 
+    /**
+     * @notice Starts the checkpointing on puffer modules
+     */
+    function callStartCheckpoint(address[] calldata moduleAddresses) external;
+
     /**
      * @notice Calls the modifyOperatorDetails function on the restaking operator
      * @param restakingOperator is the address of the restaking operator