From b8c730f931de4917abc3cc2d3e4b94a3dba87e0f Mon Sep 17 00:00:00 2001
From: Alex Forshtat <forshtat1@gmail.com>
Date: Mon, 30 Oct 2023 19:05:23 +0200
Subject: [PATCH] AA-221: Prevent Paymaster 'postOp' function from reverting
 the entire bundle

---
 contracts/core/EntryPoint.sol                 |  7 ++++--
 .../test/TestPaymasterRevertCustomError.sol   | 24 +++++++++++++++++--
 test/entrypoint.test.ts                       | 20 ++++++++++++++++
 3 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/contracts/core/EntryPoint.sol b/contracts/core/EntryPoint.sol
index 805b5f7c..8e587074 100644
--- a/contracts/core/EntryPoint.sol
+++ b/contracts/core/EntryPoint.sol
@@ -662,9 +662,12 @@ contract EntryPoint is IEntryPoint, StakeManager, NonceManager, ReentrancyGuard,
                 if (context.length > 0) {
                     actualGasCost = actualGas * gasPrice;
                     if (mode != IPaymaster.PostOpMode.postOpReverted) {
-                        IPaymaster(paymaster).postOp{
+                        try IPaymaster(paymaster).postOp{
                             gas: mUserOp.verificationGasLimit
-                        }(mode, context, actualGasCost);
+                        }(mode, context, actualGasCost)
+                        {} catch (bytes memory reason) {
+                            revert(string(abi.encodePacked("AA96 postOp reverted", reason)));
+                        }
                     }
                 }
             }
diff --git a/contracts/test/TestPaymasterRevertCustomError.sol b/contracts/test/TestPaymasterRevertCustomError.sol
index 1ceac5e0..0517d594 100644
--- a/contracts/test/TestPaymasterRevertCustomError.sol
+++ b/contracts/test/TestPaymasterRevertCustomError.sol
@@ -9,6 +9,15 @@ import "../core/BasePaymaster.sol";
 error CustomError();
 
 contract TestPaymasterRevertCustomError is BasePaymaster {
+    bytes32 private constant INNER_OUT_OF_GAS = hex"deaddead";
+
+    enum RevertType {
+        customError,
+        entryPointError
+    }
+
+    RevertType private revertType;
+
     // solhint-disable no-empty-blocks
     constructor(IEntryPoint _entryPoint) BasePaymaster(_entryPoint)
     {}
@@ -20,8 +29,19 @@ contract TestPaymasterRevertCustomError is BasePaymaster {
         context = abi.encodePacked(userOp.sender);
     }
 
-    function _postOp(PostOpMode, bytes calldata, uint256) internal pure override {
+    function setRevertType(RevertType _revertType) external {
+        revertType = _revertType;
+    }
 
-        revert CustomError();
+    function _postOp(PostOpMode, bytes calldata, uint256) internal view override {
+        if (revertType == RevertType.customError){
+            revert CustomError();
+        }
+        else if (revertType == RevertType.entryPointError){
+            assembly {
+                mstore(0, INNER_OUT_OF_GAS)
+                revert(0, 32)
+            }
+        }
     }
 }
diff --git a/test/entrypoint.test.ts b/test/entrypoint.test.ts
index 0b420246..01588fe6 100644
--- a/test/entrypoint.test.ts
+++ b/test/entrypoint.test.ts
@@ -1075,6 +1075,26 @@ describe('EntryPoint', function () {
       it('should not revert when paymaster reverts with custom error on postOp', async function () {
         const account3Owner = createAccountOwner()
         const errorPostOp = await new TestPaymasterRevertCustomError__factory(ethersSigner).deploy(entryPoint.address)
+        await errorPostOp.setRevertType(0)
+        await errorPostOp.addStake(globalUnstakeDelaySec, { value: paymasterStake })
+        await errorPostOp.deposit({ value: ONE_ETH })
+
+        const op = await fillAndSign({
+          paymasterAndData: errorPostOp.address,
+          callData: accountExecFromEntryPoint.data,
+          initCode: getAccountInitCode(account3Owner.address, simpleAccountFactory),
+
+          verificationGasLimit: 3e6,
+          callGasLimit: 1e6
+        }, account3Owner, entryPoint)
+        const beneficiaryAddress = createAddress()
+        await entryPoint.handleOps([op], beneficiaryAddress)
+      })
+
+      it('should not revert when paymaster reverts with known EntryPoint error in postOp', async function () {
+        const account3Owner = createAccountOwner()
+        const errorPostOp = await new TestPaymasterRevertCustomError__factory(ethersSigner).deploy(entryPoint.address)
+        await errorPostOp.setRevertType(1)
         await errorPostOp.addStake(globalUnstakeDelaySec, { value: paymasterStake })
         await errorPostOp.deposit({ value: ONE_ETH })