From 51c9fe06a29b5ce3b591c79b502de3d0f7431679 Mon Sep 17 00:00:00 2001
From: tommyrharper <thomasrobertharper@gmail.com>
Date: Wed, 10 Jul 2024 19:36:05 +0100
Subject: [PATCH] =?UTF-8?q?=F0=9F=91=B7=F0=9F=8F=BB=20=F0=9F=A7=B9=20Add?=
 =?UTF-8?q?=20onlyEntryPoint=20modifier?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/MarginPaymaster.sol | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/MarginPaymaster.sol b/src/MarginPaymaster.sol
index 55bf81c..0926f3f 100644
--- a/src/MarginPaymaster.sol
+++ b/src/MarginPaymaster.sol
@@ -48,12 +48,20 @@ contract MarginPaymaster is IPaymaster, Zap {
         _USDC.approve(_uniV3Router, type(uint256).max);
     }
 
+    modifier onlyEntryPoint() {
+        if (msg.sender != entryPoint) revert InvalidEntryPoint();
+        _;
+    }
+
     function validatePaymasterUserOp(
         UserOperation calldata userOp,
         bytes32,
         uint256
-    ) external returns (bytes memory context, uint256 validationData) {
-        if (msg.sender != entryPoint) revert InvalidEntryPoint();
+    )
+        external
+        onlyEntryPoint
+        returns (bytes memory context, uint256 validationData)
+    {
         context = abi.encode(userOp.sender); // passed to the postOp method
         validationData = 0; // 0 means accept sponsorship, 1 means reject
     }
@@ -62,17 +70,15 @@ contract MarginPaymaster is IPaymaster, Zap {
         PostOpMode,
         bytes calldata context,
         uint256 actualGasCostInWei
-    ) external {
-        if (msg.sender != entryPoint) revert InvalidEntryPoint();
-
+    ) external onlyEntryPoint {
         (, int24 tick, , , , , ) = pool.slot0();
 
-        uint256 costOfGasInUSDC = OracleLibrary.getQuoteAtTick(
+        uint256 costOfGasInUSDC = (OracleLibrary.getQuoteAtTick(
             tick,
-            uint128(actualGasCostInWei), // TODO: account for gas costs of postOp func
+            uint128(actualGasCostInWei),
             address(weth),
             address(_USDC)
-        ) * 110 / 100;
+        ) * 110) / 100; // TODO: account for gas costs of postOp func
         uint256 costOfGasInsUSD = costOfGasInUSDC * 1e12;
 
         address sender = abi.decode(context, (address));