Charge user in postOp

contracts/common/BiconomyTokenPaymasterErrors.sol

@@ -81,8 +81,14 @@ contract BiconomyTokenPaymasterErrors {
      */
     error WithdrawalFailed();
 
+
+    /**
+     * @notice Throws when PM was not able to charge user
+     */
+    error FailedToChargeTokens(address account, address token, uint256 amount, bytes32 userOpHash);
+
     /**
-     * @notice Emitted when ETH is withdrawn from the paymaster
+     * Throws when account has insufficient token balance to pay for gas
      */
-    event EthWithdrawn(address indexed recipient, uint256 indexed amount);
+    error InsufficientTokenBalance(address account, address token, uint256 amount, bytes32 userOpHash);
 }

contracts/interfaces/IBiconomyTokenPaymaster.sol

@@ -8,7 +8,6 @@ interface IBiconomyTokenPaymaster {
     enum PaymasterMode {
         EXTERNAL, // Price provided by external service. Authenticated using signature from verifyingSigner
         INDEPENDENT // Price queried from oracle. No signature needed from external service.
-
     }
 
     // Struct for storing information about the token
@@ -23,18 +22,25 @@ interface IBiconomyTokenPaymaster {
     event UpdatedVerifyingSigner(address indexed oldSigner, address indexed newSigner, address indexed actor);
     event UpdatedFeeCollector(address indexed oldFeeCollector, address indexed newFeeCollector, address indexed actor);
     event UpdatedPriceExpiryDuration(uint256 indexed oldValue, uint256 indexed newValue);
-    event TokensRefunded(
-        address indexed userOpSender, address indexed token, uint256 refundAmount, bytes32 indexed userOpHash
-    );
-    event PaidGasInTokens(
+
+    event PaidGasInTokensIndependent(
         address indexed userOpSender,
         address indexed token,
-        uint256 nativeCharge,
+        uint256 gasCostBeforePostOpAndPenalty,
         uint256 tokenCharge,
         uint32 priceMarkup,
         uint256 tokenPrice,
-        bytes32 indexed userOpHash
+        bytes32 userOpHash
     );
+    event PaidGasInTokensExternal(
+        address indexed userOpSender,
+        address indexed token,
+        uint256 tokenAmount,
+        bytes32 userOpHash
+    );
+
+    event EthWithdrawn(address indexed recipient, uint256 indexed amount);
+
     event Received(address indexed sender, uint256 value);
     event TokensWithdrawn(address indexed token, address indexed to, uint256 indexed amount, address actor);
     event AddedToTokenDirectory(address indexed tokenAddress, IOracle indexed oracle, uint8 decimals);

contracts/libraries/TokenPaymasterParserLib.sol

@@ -31,17 +31,15 @@ library TokenPaymasterParserLib {
             uint48 validUntil,
             uint48 validAfter,
             address tokenAddress,
-            uint256 tokenPrice, 
-            uint32 externalPriceMarkup,
+            uint256 estimatedTokenAmount, 
             bytes calldata signature
         )
     {
         validUntil = uint48(bytes6(modeSpecificData[:6]));
         validAfter = uint48(bytes6(modeSpecificData[6:12]));
         tokenAddress = address(bytes20(modeSpecificData[12:32]));
-        tokenPrice = uint256(bytes32(modeSpecificData[32:64]));
-        externalPriceMarkup = uint32(bytes4(modeSpecificData[64:68]));
-        signature = modeSpecificData[68:];
+        estimatedTokenAmount = uint256(bytes32(modeSpecificData[32:64]));
+        signature = modeSpecificData[64:];
     }
 
     function parseIndependentModeSpecificData(

contracts/token/BiconomyTokenPaymaster.sol

@@ -391,8 +391,7 @@ contract BiconomyTokenPaymaster is
         uint48 validUntil,
         uint48 validAfter,
         address tokenAddress,
-        uint256 tokenPrice,
-        uint32 externalPriceMarkup
+        uint256 estimatedTokenAmount
     )
         public
         view
@@ -415,8 +414,7 @@ contract BiconomyTokenPaymaster is
                 validUntil,
                 validAfter,
                 tokenAddress,
-                tokenPrice,
-                externalPriceMarkup
+                estimatedTokenAmount
             )
         );
     }
@@ -479,14 +477,6 @@ contract BiconomyTokenPaymaster is
             revert InvalidPaymasterMode();
         }
 
-        // callGasLimit + paymasterPostOpGas
-        uint256 maxPenalty = (
-            (
-                uint128(uint256(userOp.accountGasLimits))
-                    + uint128(bytes16(userOp.paymasterAndData[_PAYMASTER_POSTOP_GAS_OFFSET:_PAYMASTER_DATA_OFFSET]))
-            ) * 10 * userOp.unpackMaxFeePerGas()
-        ) / 100;
-
         if (mode == PaymasterMode.EXTERNAL) {
             // Use the price and other params specified in modeSpecificData by the verifyingSigner
             // Useful for supporting tokens which don't have oracle support
@@ -495,8 +485,7 @@ contract BiconomyTokenPaymaster is
                 uint48 validUntil,
                 uint48 validAfter,
                 address tokenAddress,
-                uint256 tokenPrice, // Note: what backend should pass is nativeTokenPriceInUsd/tokenPriceInUsd * 10^token decimals
-                uint32 externalPriceMarkup,
+                uint256 estimatedTokenAmount,
                 bytes memory signature
             ) = modeSpecificData.parseExternalModeSpecificData();
 
@@ -506,7 +495,7 @@ contract BiconomyTokenPaymaster is
 
             bool validSig = verifyingSigner.isValidSignatureNow(
                 ECDSA_solady.toEthSignedMessageHash(
-                    getHash(userOp, validUntil, validAfter, tokenAddress, tokenPrice, externalPriceMarkup)
+                    getHash(userOp, validUntil, validAfter, tokenAddress, estimatedTokenAmount)
                 ),
                 signature
             );
@@ -516,38 +505,35 @@ contract BiconomyTokenPaymaster is
                 return ("", _packValidationData(true, validUntil, validAfter));
             }
 
-            if (externalPriceMarkup > _MAX_PRICE_MARKUP || externalPriceMarkup < _PRICE_DENOMINATOR) {
-                revert InvalidPriceMarkup();
-            }
-
-
-            uint256 tokenAmount;
-            {
-                uint256 maxFeePerGas = UserOperationLib.unpackMaxFeePerGas(userOp);
-                tokenAmount = ((maxCost + maxPenalty + (unaccountedGas * maxFeePerGas)) * externalPriceMarkup * tokenPrice)
-                    / (_NATIVE_TOKEN_DECIMALS * _PRICE_DENOMINATOR);
+            if(IERC20(tokenAddress).balanceOf(userOp.sender) < estimatedTokenAmount) {
+                revert InsufficientTokenBalance(userOp.sender, tokenAddress, estimatedTokenAmount, userOpHash);
             }
 
-            // Transfer full amount to this address. Unused amount will be refunded in postOP
-            SafeTransferLib.safeTransferFrom(tokenAddress, userOp.sender, address(this), tokenAmount);
-
-            // deduct max penalty from the token amount we pass to the postOp
-            // so we don't refund it at postOp
-            context = abi.encode(
-                userOp.sender,
-                tokenAddress,
-                tokenAmount-((maxPenalty*tokenPrice*externalPriceMarkup)/(_NATIVE_TOKEN_DECIMALS*_PRICE_DENOMINATOR)),
-                tokenPrice,
-                externalPriceMarkup,
-                userOpHash
+            context = abi.encodePacked(
+                PaymasterMode.EXTERNAL,
+                abi.encode(
+                    userOp.sender,
+                    tokenAddress,
+                    estimatedTokenAmount,
+                    userOpHash
+                )
             );
             validationData = _packValidationData(false, validUntil, validAfter);
+
+        /// INDEPENDENT MODE
         } else if (mode == PaymasterMode.INDEPENDENT) {
             // Use only oracles for the token specified in modeSpecificData
             if (modeSpecificData.length != 20) {
                 revert InvalidTokenAddress();
             }
 
+            uint256 maxPenalty = (
+            (
+                uint128(uint256(userOp.accountGasLimits))
+                    + uint128(bytes16(userOp.paymasterAndData[_PAYMASTER_POSTOP_GAS_OFFSET:_PAYMASTER_DATA_OFFSET]))
+                ) * 10 //* userOp.unpackMaxFeePerGas()
+            ) / 100;
+
             // Get address for token used to pay
             address tokenAddress = modeSpecificData.parseIndependentModeSpecificData();
             uint256 tokenPrice = _getPrice(tokenAddress);
@@ -566,17 +552,22 @@ contract BiconomyTokenPaymaster is
             }
 
             // Transfer full amount to this address. Unused amount will be refunded in postOP
-            SafeTransferLib.safeTransferFrom(tokenAddress, userOp.sender, address(this), tokenAmount);
+            // SafeTransferLib.safeTransferFrom(tokenAddress, userOp.sender, address(this), tokenAmount);
+            if(IERC20(tokenAddress).balanceOf(userOp.sender) < tokenAmount) {
+                revert InsufficientTokenBalance(userOp.sender, tokenAddress, tokenAmount, userOpHash);
+            }
 
-            context =
+            context = abi.encodePacked(
+                PaymasterMode.INDEPENDENT,
                 abi.encode(
                     userOp.sender,
                     tokenAddress,
-                    tokenAmount-((maxPenalty*tokenPrice*priceMarkup)/(_NATIVE_TOKEN_DECIMALS*_PRICE_DENOMINATOR)),
+                    maxPenalty,
                     tokenPrice,
                     priceMarkup,
                     userOpHash
-                );
+                )
+            );
             validationData = 0; // Validation success and price is valid indefinetly
         }
     }
@@ -596,31 +587,46 @@ contract BiconomyTokenPaymaster is
     )
         internal
         override
-    {
-        // Decode context data
-        (
-            address userOpSender,
-            address tokenAddress,
-            uint256 prechargedAmount,
-            uint256 tokenPrice,
-            uint32 appliedPriceMarkup,
-            bytes32 userOpHash
-        ) = abi.decode(context, (address, address, uint256, uint256, uint32, bytes32));
-
-        // Calculate the actual cost in tokens based on the actual gas cost and the token price
-        uint256 actualTokenAmount = (
-            (actualGasCost + (unaccountedGas * actualUserOpFeePerGas)) * appliedPriceMarkup * tokenPrice
-        ) / (_NATIVE_TOKEN_DECIMALS * _PRICE_DENOMINATOR);
-        if (prechargedAmount > actualTokenAmount) {
-            // If the user was overcharged, refund the excess tokens
-            uint256 refundAmount = prechargedAmount - actualTokenAmount;
-            SafeTransferLib.safeTransfer(tokenAddress, userOpSender, refundAmount);
-            emit TokensRefunded(userOpSender, tokenAddress, refundAmount, userOpHash);
-        }
+    {   
+        PaymasterMode pmMode = PaymasterMode(uint8(context[0]));
+        if (pmMode == PaymasterMode.EXTERNAL) {
+            // Decode context data
+            (
+                address userOpSender,
+                address tokenAddress,
+                uint256 estimatedTokenAmount,
+                bytes32 userOpHash
+            ) = abi.decode(context[1:], (address, address, uint256, bytes32));
+
+            if (SafeTransferLib.trySafeTransferFrom(tokenAddress, userOpSender, address(this), estimatedTokenAmount)) {
+                emit PaidGasInTokensExternal(userOpSender, tokenAddress, estimatedTokenAmount, userOpHash);
+            } else {
+                revert FailedToChargeTokens(userOpSender, tokenAddress, estimatedTokenAmount, userOpHash);
+            }
 
-        emit PaidGasInTokens(
-            userOpSender, tokenAddress, actualGasCost, actualTokenAmount, appliedPriceMarkup, tokenPrice, userOpHash
-        );
+        } else if (pmMode == PaymasterMode.INDEPENDENT) {
+            (
+                address userOpSender,
+                address tokenAddress,
+                uint256 maxPenalty,
+                uint256 tokenPrice,
+                uint32 appliedPriceMarkup,
+                bytes32 userOpHash
+            ) = abi.decode(context[1:], (address, address, uint256, uint256, uint32, bytes32));
+            // Calculate the amount to charge. unaccountedGas and maxPenalty are used, as we do not know the exact gas spent for postop and actual penalty at this point
+            // this is obviously overcharge, however, the excess amount can be refunded by backend, when we know the exact gas spent (emitted by EP after executing UserOp)
+            uint256 tokenAmount = (
+                (actualGasCost + ((unaccountedGas + maxPenalty)) * actualUserOpFeePerGas)) * appliedPriceMarkup * tokenPrice
+            / (_NATIVE_TOKEN_DECIMALS * _PRICE_DENOMINATOR);
+
+            if (SafeTransferLib.trySafeTransferFrom(tokenAddress, userOpSender, address(this), tokenAmount)) {
+                emit PaidGasInTokensIndependent(
+                    userOpSender, tokenAddress, actualGasCost, tokenAmount, appliedPriceMarkup, tokenPrice, userOpHash
+                );
+            } else {
+                revert FailedToChargeTokens(userOpSender, tokenAddress, tokenAmount, userOpHash);
+            }
+        }
     }
 
     function _validateTokenInfo(TokenInfo memory tokenInfo) internal view {

test/base/TestBase.sol

@@ -14,6 +14,9 @@ import { Exec } from "account-abstraction/utils/Exec.sol";
 import { IPaymaster } from "account-abstraction/interfaces/IPaymaster.sol";
 
 import { Nexus } from "@nexus/contracts/Nexus.sol";
+import { INexus } from "@nexus/contracts/interfaces/INexus.sol";
+import { IERC7579Account } from "@nexus/contracts/interfaces/IERC7579Account.sol";
+import { IExecutionHelper } from "@nexus/contracts/interfaces/base/IExecutionHelper.sol";
 import { CheatCodes } from "@nexus/test/foundry/utils/CheatCodes.sol";
 import { BaseEventsAndErrors } from "./BaseEventsAndErrors.sol";
 import { MockToken } from "@nexus/contracts/mocks/MockToken.sol";
@@ -63,8 +66,7 @@ abstract contract TestBase is CheatCodes, TestHelper, BaseEventsAndErrors {
         uint48 validUntil;
         uint48 validAfter;
         address tokenAddress;
-        uint256 tokenPrice;
-        uint32 externalPriceMarkup;
+        uint256 estimatedTokenAmount;
     }
 
     // Used to buffer user op gas limits
@@ -342,8 +344,7 @@ abstract contract TestBase is CheatCodes, TestHelper, BaseEventsAndErrors {
             pmData.validUntil,
             pmData.validAfter,
             pmData.tokenAddress,
-            pmData.tokenPrice,
-            pmData.externalPriceMarkup,
+            pmData.estimatedTokenAmount,
             new bytes(65) // Zero signature
         );
 
@@ -352,7 +353,7 @@ abstract contract TestBase is CheatCodes, TestHelper, BaseEventsAndErrors {
 
         // Generate hash to be signed
         bytes32 paymasterHash =
-            paymaster.getHash(userOp, pmData.validUntil, pmData.validAfter, pmData.tokenAddress, pmData.tokenPrice, pmData.externalPriceMarkup);
+            paymaster.getHash(userOp, pmData.validUntil, pmData.validAfter, pmData.tokenAddress, pmData.estimatedTokenAmount);
 
         // Sign the hash
         signature = signMessage(signer, paymasterHash);
@@ -367,8 +368,7 @@ abstract contract TestBase is CheatCodes, TestHelper, BaseEventsAndErrors {
             pmData.validUntil,
             pmData.validAfter,
             pmData.tokenAddress,
-            pmData.tokenPrice,
-            pmData.externalPriceMarkup,
+            pmData.estimatedTokenAmount,
             signature
         );
     }
@@ -482,10 +482,10 @@ abstract contract TestBase is CheatCodes, TestHelper, BaseEventsAndErrors {
         // Assert we never undercharge
         assertGe(gasPaidBySAInNativeTokens, BUNDLER.addr.balance - initialBundlerBalance);
 
-        // Ensure that max 2% difference between what is should have been charged and what was charged
+        // Ensure that max 3% difference between what should have been charged and what was charged
         // this difference comes from difference of postop gas and estimated postop gas (paymaster.unaccountedGas)
         // and from estimation of real penalty which is not emitted by EP :(
-        assertApproxEqRel((totalGasFeePaid + maxPenalty - realPenalty) * priceMarkup / _PRICE_MARKUP_DENOMINATOR, gasPaidBySAInNativeTokens, 0.02e18, "If this fails, check the test case inline comments");
+        assertApproxEqRel((totalGasFeePaid + maxPenalty - realPenalty) * priceMarkup / _PRICE_MARKUP_DENOMINATOR, gasPaidBySAInNativeTokens, 0.03e18, "If this fails, check the test case inline comments");
     }
 
     function _toSingletonArray(address addr) internal pure returns (address[] memory) {

test/mocks/PaymasterParserLibExposed.sol

@@ -15,8 +15,7 @@ library PaymasterParserLibExposed {
             uint48 validUntil,
             uint48 validAfter,
             address tokenAddress,
-            uint256 tokenPrice, 
-            uint32 externalPriceMarkup,
+            uint256 estimatedTokenAmount, 
             bytes calldata signature
         ) {
         return modeSpecificData.parseExternalModeSpecificData();

test/mocks/PaymasterParserLibWrapper.sol

@@ -15,11 +15,10 @@ contract PaymasterParserLibWrapper {
             uint48 validUntil,
             uint48 validAfter,
             address tokenAddress,
-            uint256 tokenPrice, 
-            uint32 externalPriceMarkup,
+            uint256 estimatedTokenAmount, 
             bytes memory signature
         ) {
-        (validUntil, validAfter, tokenAddress, tokenPrice, externalPriceMarkup, signature) = modeSpecificData.parseExternalModeSpecificData();
+        (validUntil, validAfter, tokenAddress, estimatedTokenAmount, signature) = modeSpecificData.parseExternalModeSpecificData();
     }
 
     function parseIndependentModeSpecificData(bytes calldata modeSpecificData) external pure returns (address tokenAddress) {