fix: token paymaster decimal bug

contracts/samples/TokenPaymaster.sol

@@ -139,7 +139,12 @@ contract TokenPaymaster is BasePaymaster, UniswapHelper, OracleHelper {
                     cachedPriceWithMarkup = clientSuppliedPrice;
                 }
             }
-            uint256 tokenAmount = weiToToken(preChargeNative, cachedPriceWithMarkup);
+            uint8 tokenDecimals = token.decimals();
+            uint256 tokenAmount = weiToToken(
+                preChargeNative, 
+                tokenDecimals,
+                cachedPriceWithMarkup
+            );
             SafeERC20.safeTransferFrom(token, userOp.sender, address(this), tokenAmount);
             context = abi.encode(tokenAmount, userOp.sender);
             validationResult = _packValidationData(
@@ -169,7 +174,12 @@ contract TokenPaymaster is BasePaymaster, UniswapHelper, OracleHelper {
             uint256 cachedPriceWithMarkup = _cachedPrice * PRICE_DENOMINATOR / priceMarkup;
         // Refund tokens based on actual gas cost
             uint256 actualChargeNative = actualGasCost + tokenPaymasterConfig.refundPostopCost * actualUserOpFeePerGas;
-            uint256 actualTokenNeeded = weiToToken(actualChargeNative, cachedPriceWithMarkup);
+            uint8 tokenDecimals = token.decimals();
+            uint256 actualTokenNeeded = weiToToken(
+                actualChargeNative,
+                tokenDecimals,
+                cachedPriceWithMarkup
+            );
             if (preCharge > actualTokenNeeded) {
                 // If the initially provided token amount is greater than the actual amount needed, refund the difference
                 SafeERC20.safeTransfer(

contracts/samples/utils/UniswapHelper.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.23;
 /* solhint-disable not-rely-on-time */
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";
 
 import "@uniswap/v3-periphery/contracts/interfaces/ISwapRouter.sol";
 import "@uniswap/v3-periphery/contracts/interfaces/IPeripheryPayments.sol";
@@ -26,15 +27,15 @@ abstract contract UniswapHelper {
     ISwapRouter public immutable uniswap;
 
     /// @notice The ERC20 token used for transaction fee payments
-    IERC20 public immutable token;
+    IERC20Metadata public immutable token;
 
     /// @notice The ERC-20 token that wraps the native asset for current chain
     IERC20 public immutable wrappedNative;
 
     UniswapHelperConfig private uniswapHelperConfig;
 
     constructor(
-        IERC20 _token,
+        IERC20Metadata _token,
         IERC20 _wrappedNative,
         ISwapRouter _uniswap,
         UniswapHelperConfig memory _uniswapHelperConfig
@@ -50,9 +51,10 @@ abstract contract UniswapHelper {
         uniswapHelperConfig = _uniswapHelperConfig;
     }
 
-    function _maybeSwapTokenToWeth(IERC20 tokenIn, uint256 quote) internal returns (uint256) {
+    function _maybeSwapTokenToWeth(IERC20Metadata tokenIn, uint256 quote) internal returns (uint256) {
         uint256 tokenBalance = tokenIn.balanceOf(address(this));
-        uint256 amountOutMin = addSlippage(tokenToWei(tokenBalance, quote), uniswapHelperConfig.slippage);
+        uint8 tokenDecimals = tokenIn.decimals();
+        uint256 amountOutMin = addSlippage(tokenToWei(tokenBalance, tokenDecimals, quote), uniswapHelperConfig.slippage);
         if (amountOutMin < uniswapHelperConfig.minSwapAmount) {
             return 0;
         }
@@ -71,12 +73,20 @@ abstract contract UniswapHelper {
     }
 
 
-    function tokenToWei(uint256 amount, uint256 price) public pure returns (uint256) {
-        return amount * price / PRICE_DENOMINATOR;
+    function tokenToWei(uint256 amount, uint8 tokenDecimals, uint256 price) public pure returns (uint256) {
+        if (tokenDecimals >= 18) {
+            return amount * price / (PRICE_DENOMINATOR * 10 ** (tokenDecimals - 18));
+        } else {
+            return amount * price * 10 ** (18 - tokenDecimals) / PRICE_DENOMINATOR;
+        }
     }
 
-    function weiToToken(uint256 amount, uint256 price) public pure returns (uint256) {
-        return amount * PRICE_DENOMINATOR / price;
+    function weiToToken(uint256 amount, uint8 tokenDecimals, uint256 price) public pure returns (uint256) {
+        if (tokenDecimals >= 18) {
+            return amount * PRICE_DENOMINATOR * 10 ** (tokenDecimals - 18) / price;
+        } else {
+            return amount * PRICE_DENOMINATOR / (price * 10 ** (18 - tokenDecimals));
+        }
     }
 
     function unwrapWeth(uint256 amount) internal {

test/samples/OracleHelper.test.ts

@@ -63,7 +63,7 @@ describe('OracleHelper', function () {
     it('should figure out the correct price', async function () {
       await testEnv.paymaster.updateCachedPrice(true)
       const cachedPrice = await testEnv.paymaster.cachedPrice()
-      const tokensPerEtherCalculated = await testEnv.paymaster.weiToToken(parseEther('1'), cachedPrice)
+      const tokensPerEtherCalculated = await testEnv.paymaster.weiToToken(parseEther('1'), 18, cachedPrice)
       assert.equal(cachedPrice.toString(), testEnv.expectedPrice.toString(), 'price not right')
       assert.equal(tokensPerEtherCalculated.toString(), testEnv.expectedTokensPerEtherCalculated.toString(), 'tokens amount not right')
     })

test/samples/TokenPaymaster.test.ts

@@ -1,5 +1,5 @@
 import { ContractReceipt, ContractTransaction, Wallet, utils, BigNumber } from 'ethers'
-import { hexlify, hexZeroPad, Interface, parseEther } from 'ethers/lib/utils'
+import { hexlify, hexZeroPad, Interface, parseEther, parseUnits } from 'ethers/lib/utils'
 import { assert, expect } from 'chai'
 import { ethers } from 'hardhat'
 
@@ -243,7 +243,7 @@ describe('TokenPaymaster', function () {
     const expectedTokenPriceWithMarkup = priceDenominator
       .mul(initialPriceToken).div(initialPriceEther) // expectedTokenPrice of 0.2 as BigNumber
       .mul(10).div(15) // added 150% priceMarkup
-    const expectedTokenCharge = actualGasCostPaymaster.add(addedPostOpCost).mul(priceDenominator).div(expectedTokenPriceWithMarkup)
+    const expectedTokenCharge = actualGasCostPaymaster.add(addedPostOpCost).mul(priceDenominator).div(expectedTokenPriceWithMarkup).div(BigNumber.from(10).pow(18 - 6)) // token decimals is 6
     const postOpGasCost = actualGasCostEntryPoint.sub(actualGasCostPaymaster)
     assert.equal(decodedLogs.length, 5)
     assert.equal(decodedLogs[4].args.success, true)
@@ -330,10 +330,14 @@ describe('TokenPaymaster', function () {
     const preChargeTokens = decodedLogs[0].args.value
     const requiredGas = BigNumber.from(op.callGasLimit).add(BigNumber.from(op.verificationGasLimit).add(BigNumber.from(op.paymasterVerificationGasLimit))).add(BigNumber.from(op.paymasterPostOpGasLimit)).add(op.preVerificationGas).add(40000 /*  REFUND_POSTOP_COST */)
     const requiredPrefund = requiredGas.mul(op.maxFeePerGas)
-    const preChargeTokenPrice = requiredPrefund.mul(priceDenominator).div(preChargeTokens)
+    const preChargeTokenPrice = requiredPrefund.mul(priceDenominator).div(preChargeTokens).div(BigNumber.from(10).pow(18 - 6)) // token decimals is 6
 
-    // TODO: div 1e10 to hide rounding errors. look into it - 1e10 is too much.
-    assert.equal(preChargeTokenPrice.div(1e10).toString(), overrideTokenPrice.div(1e10).toString())
+    // assert we didn't charge more than user allowed
+    assert.isTrue(overrideTokenPrice.lte(preChargeTokenPrice))
+
+    // assert the token amount is the most optimal that can be charged, which means if charging more token the price will be lower than the one supplied by the client
+    const tokenPriceIfCharingMore = requiredPrefund.mul(priceDenominator).div(preChargeTokens.add(1)).div(BigNumber.from(10).pow(18 - 6)) // token decimals is 6
+    assert.isTrue(tokenPriceIfCharingMore.lte(overrideTokenPrice))
     await ethers.provider.send('evm_revert', [snapshot])
   })
 
@@ -375,9 +379,16 @@ describe('TokenPaymaster', function () {
     const preChargeTokens = decodedLogs[0].args.value
     const requiredGas = BigNumber.from(op.callGasLimit).add(BigNumber.from(op.verificationGasLimit).add(BigNumber.from(op.paymasterVerificationGasLimit))).add(BigNumber.from(op.paymasterPostOpGasLimit)).add(op.preVerificationGas).add(40000 /*  REFUND_POSTOP_COST */)
     const requiredPrefund = requiredGas.mul(op.maxFeePerGas)
-    const preChargeTokenPrice = requiredPrefund.mul(priceDenominator).div(preChargeTokens)
+    const preChargeTokenPrice = requiredPrefund.mul(priceDenominator).div(preChargeTokens).div(BigNumber.from(10).pow(18 - 6)) // token decimals is 6
+    const expectedPrice = currentCachedPrice.mul(10).div(15)
+
+    // assert we didn't charge more than the amount calculated by the cached price
+    assert.isTrue(expectedPrice.lte(preChargeTokenPrice))
+
+    // assert the token amount is the most optimal that can be charged, which means if charging more token the price will be lower than the oracle price
+    const tokenPriceIfCharingMore = requiredPrefund.mul(priceDenominator).div(preChargeTokens.add(1)).div(BigNumber.from(10).pow(18 - 6)) // token decimals is 6
+    assert.isTrue(tokenPriceIfCharingMore.lte(expectedPrice))
 
-    assert.equal(preChargeTokenPrice.toString(), currentCachedPrice.mul(10).div(15).toString())
     await ethers.provider.send('evm_revert', [snapshot])
   })
 
@@ -427,7 +438,7 @@ describe('TokenPaymaster', function () {
     const snapshot = await ethers.provider.send('evm_snapshot', [])
 
     // Make sure account has small amount of tokens
-    await token.transfer(account.address, parseEther('0.01'))
+    await token.transfer(account.address, parseUnits('0.01', 6))
     await token.sudoApprove(account.address, paymaster.address, ethers.constants.MaxUint256)
 
     // Ether price increased 100 times!
@@ -437,7 +448,7 @@ describe('TokenPaymaster', function () {
     await ethers.provider.send('evm_increaseTime', [200])
 
     // Withdraw most of the tokens the account hs inside the inner transaction
-    const withdrawTokensCall = await token.populateTransaction.transfer(token.address, parseEther('0.009')).then(tx => tx.data!)
+    const withdrawTokensCall = await token.populateTransaction.transfer(token.address, parseUnits('0.009', 6)).then(tx => tx.data!)
     const callData = await account.populateTransaction.execute(token.address, 0, withdrawTokensCall).then(tx => tx.data!)
 
     let op = await fillUserOp({
@@ -496,7 +507,7 @@ describe('TokenPaymaster', function () {
     assert.equal(decodedLogs[4].name, 'StubUniswapExchangeEvent')
     assert.equal(decodedLogs[8].name, 'Received')
     assert.equal(decodedLogs[9].name, 'Deposited')
-    const deFactoExchangeRate = decodedLogs[4].args.amountOut.toString() / decodedLogs[4].args.amountIn.toString()
+    const deFactoExchangeRate = decodedLogs[4].args.amountOut.toString() / decodedLogs[4].args.amountIn.toString() / 1e12
     const expectedPrice = initialPriceToken / initialPriceEther
     assert.closeTo(deFactoExchangeRate, expectedPrice, 0.001)
   })