feat: added fallback on origin, AA compatibility, tests

zama-ai · Aug 29, 2024 · 625970f · 625970f
1 parent 0049e9d
commit 625970f
Show file tree

Hide file tree

Showing 8 changed files with 484 additions and 20 deletions.
diff --git a/codegen/payments.ts b/codegen/payments.ts
@@ -23,6 +23,10 @@ export function generateFHEPayment(priceData: PriceData): string {
   error RecoveryFailed();
   error WithdrawalFailed();
   error AccountNotEnoughFunded();
+  error AlreadyAuthorizedAllContracts();
+  error AlreadyWhitelistedContract();
+  error AllContractsNotAuthorized();
+  error ContractNotWhitelisted();
   
   contract FHEPayment is Ownable2Step {
       /// @notice Name of the contract
@@ -43,6 +47,8 @@ export function generateFHEPayment(priceData: PriceData): string {
       uint256 public claimableUsedFHEGas;
   
       mapping(address payer => uint256 depositedAmount) private depositsETH;
+      mapping(address user => bool allowedAllContracts) private allowedAll;
+      mapping(address user => mapping(address dappContract => bool isWhitelisted)) private whitelistedDapps;
 
       constructor() Ownable(msg.sender) {}
 
@@ -67,19 +73,81 @@ export function generateFHEPayment(priceData: PriceData): string {
           return depositsETH[account];
       }
 
-      function updateFunding(address payer, uint256 paidAmountGas) private {
-        uint256 ratio_gas = (tx.gasprice*FHE_GASPRICE_NATIVE_RATIO)/1_000_000;
-        uint256 effective_fhe_gasPrice = ratio_gas > MIN_FHE_GASPRICE ? ratio_gas : MIN_FHE_GASPRICE;
-        uint256 paidAmountWei = effective_fhe_gasPrice*paidAmountGas;
-        uint256 depositedAmount = depositsETH[payer];
-        if(paidAmountWei>depositedAmount) revert AccountNotEnoughFunded();
-        unchecked{
-          depositsETH[payer] = depositedAmount - paidAmountWei;
+      function didAuthorizeAllContracts(address account) external view returns (bool) {
+          return allowedAll[account];
+      }
+
+      function didWhitelistContract(address user, address dappContract) external view returns (bool) {
+          return whitelistedDapps[user][dappContract];
+      }
+
+      function authorizeAllContracts() external {
+          if (allowedAll[msg.sender]) revert AlreadyAuthorizedAllContracts();
+          allowedAll[msg.sender] = true;
+      }
+
+      function whitelistContract(address dappContract) external {
+          if (whitelistedDapps[msg.sender][dappContract]) revert AlreadyWhitelistedContract();
+          whitelistedDapps[msg.sender][dappContract] = true;
+      }
+
+      function removeAuthorizationAllContracts() external {
+          if (!allowedAll[msg.sender]) revert AllContractsNotAuthorized();
+          allowedAll[msg.sender] = false;
+      }
+
+      function removeWhitelistedContract(address dappContract) external {
+          if (!whitelistedDapps[msg.sender][dappContract]) revert ContractNotWhitelisted();
+          whitelistedDapps[msg.sender][dappContract] = false;
+      }
+
+      // @notice: to be used in the context of account abstraction, before an FHE tx, to make the contract address replace tx.origin as a spender
+      function becomeTransientSpender() external {
+        assembly {
+            tstore(0, caller())
         }
-        currentBlockConsumption += paidAmountGas;
-        claimableUsedFHEGas += paidAmountWei;
       }
-  
+
+      // @notice: to be used in the context of account abstraction, after an FHE tx, to avoid issues if batched with other userOps
+      function stopBeingTransientSpender() external {
+        assembly {
+            tstore(0, 0)
+        }
+      }
+
+      function updateFunding(address payer, uint256 paidAmountGas) private {
+          uint256 ratio_gas = (tx.gasprice * FHE_GASPRICE_NATIVE_RATIO) / 1_000_000;
+          uint256 effective_fhe_gasPrice = ratio_gas > MIN_FHE_GASPRICE ? ratio_gas : MIN_FHE_GASPRICE;
+          uint256 paidAmountWei = effective_fhe_gasPrice * paidAmountGas;
+          uint256 depositedAmount = depositsETH[payer];
+          if (paidAmountWei > depositedAmount) {
+              // if dApp is not enough funded, fallbacks to user (tx.origin by default, in case of an EOA, 
+              // otherwise a smart contract account should call \`becomeTransientSpender\` before, in the same tx
+              address spender;
+              assembly {
+                spender := tload(0)
+              }
+              spender = spender == address(0) ? tx.origin : spender;
+              if (allowedAll[spender] || whitelistedDapps[spender][payer]) {
+                  uint256 depositedAmountUser = depositsETH[spender];
+                  if (paidAmountWei > depositedAmountUser) revert AccountNotEnoughFunded();
+                  unchecked {
+                      depositsETH[spender] = depositedAmountUser - paidAmountWei;
+                  }
+                  currentBlockConsumption += paidAmountGas;
+                  claimableUsedFHEGas += paidAmountWei;
+              } else {
+                  revert AccountNotEnoughFunded();
+              }
+          } else {
+              unchecked {
+                  depositsETH[payer] = depositedAmount - paidAmountWei;
+              }
+              currentBlockConsumption += paidAmountGas;
+              claimableUsedFHEGas += paidAmountWei;
+          }
+      }
+
       function checkIfNewBlock() private {
           uint256 lastBlock_ = block.number;
           if (block.number > lastBlock) {

diff --git a/examples/PaymentLimit.sol b/examples/PaymentLimit.sol
@@ -0,0 +1,40 @@
+// SPDX-License-Identifier: BSD-3-Clause-Clear
+
+pragma solidity ^0.8.24;
+
+import "../lib/TFHE.sol";
+import "@openzeppelin/contracts/access/Ownable2Step.sol";
+import "../payment/Payment.sol";
+
+contract PaymentLimit {
+    constructor() payable {
+        Payment.depositForThis(msg.value);
+    }
+
+    function wayunderBlockFHEGasLimit() external {
+        // should pass if only tx in block
+        euint64 x = TFHE.asEuint64(2);
+        euint64 result;
+        for (uint256 i; i < 3; i++) {
+            result = TFHE.mul(result, x);
+        }
+    }
+
+    function underBlockFHEGasLimit() external {
+        // should pass if only tx in block
+        euint64 x = TFHE.asEuint64(2);
+        euint64 result;
+        for (uint256 i; i < 15; i++) {
+            result = TFHE.mul(result, x);
+        }
+    }
+
+    function aboveBlockFHEGasLimit() external {
+        // should revert due to exceeding block fheGas limit
+        euint64 x = TFHE.asEuint64(2);
+        euint64 result;
+        for (uint256 i; i < 16; i++) {
+            result = TFHE.mul(result, x);
+        }
+    }
+}
diff --git a/examples/Regression1.sol b/examples/Regression1.sol
@@ -1,5 +1,5 @@
-/* SPDX-License-Identifier: MIT */
-pragma solidity ^0.8.20;
+// SPDX-License-Identifier: BSD-3-Clause-Clear
+pragma solidity ^0.8.24;
 
 contract Regression1 {
     error IndexOutOfBound();

diff --git a/examples/SmartAccount.sol b/examples/SmartAccount.sol
@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: BSD-3-Clause-Clear
+pragma solidity ^0.8.24;
+
+import "@openzeppelin/contracts/access/Ownable2Step.sol";
+
+contract SmartAccount is Ownable2Step {
+    struct Transaction {
+        address target;
+        uint256 value;
+        bytes data;
+    }
+
+    constructor() Ownable(msg.sender) {}
+
+    function executeBatch(Transaction[] memory transactions) public payable onlyOwner {
+        for (uint i = 0; i < transactions.length; i++) {
+            Transaction memory transaction = transactions[i];
+            (bool success, ) = transaction.target.call{value: transaction.value}(transaction.data);
+            require(success, "Transaction failed");
+        }
+    }
+}
diff --git a/examples/TracingSubCalls.sol b/examples/TracingSubCalls.sol
@@ -1,5 +1,6 @@
-// SPDX-License-Identifier: MIT
-pragma solidity ^0.8.0;
+// SPDX-License-Identifier: BSD-3-Clause-Clear
+
+pragma solidity ^0.8.24;
 import "../lib/TFHE.sol";
 
 contract TracingSubCalls {

diff --git a/lib/FHEPayment.sol b/lib/FHEPayment.sol
@@ -13,6 +13,10 @@ error OnlyNonScalarOperationsAreSupported();
 error RecoveryFailed();
 error WithdrawalFailed();
 error AccountNotEnoughFunded();
+error AlreadyAuthorizedAllContracts();
+error AlreadyWhitelistedContract();
+error AllContractsNotAuthorized();
+error ContractNotWhitelisted();
 
 contract FHEPayment is Ownable2Step {
     /// @notice Name of the contract
@@ -33,6 +37,8 @@ contract FHEPayment is Ownable2Step {
     uint256 public claimableUsedFHEGas;
 
     mapping(address payer => uint256 depositedAmount) private depositsETH;
+    mapping(address user => bool allowedAllContracts) private allowedAll;
+    mapping(address user => mapping(address dappContract => bool isWhitelisted)) private whitelistedDapps;
 
     constructor() Ownable(msg.sender) {}
 
@@ -57,17 +63,79 @@ contract FHEPayment is Ownable2Step {
         return depositsETH[account];
     }
 
+    function didAuthorizeAllContracts(address account) external view returns (bool) {
+        return allowedAll[account];
+    }
+
+    function didWhitelistContract(address user, address dappContract) external view returns (bool) {
+        return whitelistedDapps[user][dappContract];
+    }
+
+    function authorizeAllContracts() external {
+        if (allowedAll[msg.sender]) revert AlreadyAuthorizedAllContracts();
+        allowedAll[msg.sender] = true;
+    }
+
+    function whitelistContract(address dappContract) external {
+        if (whitelistedDapps[msg.sender][dappContract]) revert AlreadyWhitelistedContract();
+        whitelistedDapps[msg.sender][dappContract] = true;
+    }
+
+    function removeAuthorizationAllContracts() external {
+        if (!allowedAll[msg.sender]) revert AllContractsNotAuthorized();
+        allowedAll[msg.sender] = false;
+    }
+
+    function removeWhitelistedContract(address dappContract) external {
+        if (!whitelistedDapps[msg.sender][dappContract]) revert ContractNotWhitelisted();
+        whitelistedDapps[msg.sender][dappContract] = false;
+    }
+
+    // @notice: to be used in the context of account abstraction, before an FHE tx, to make the contract address replace tx.origin as a spender
+    function becomeTransientSpender() external {
+        assembly {
+            tstore(0, caller())
+        }
+    }
+
+    // @notice: to be used in the context of account abstraction, after an FHE tx, to avoid issues if batched with other userOps
+    function stopBeingTransientSpender() external {
+        assembly {
+            tstore(0, 0)
+        }
+    }
+
     function updateFunding(address payer, uint256 paidAmountGas) private {
         uint256 ratio_gas = (tx.gasprice * FHE_GASPRICE_NATIVE_RATIO) / 1_000_000;
         uint256 effective_fhe_gasPrice = ratio_gas > MIN_FHE_GASPRICE ? ratio_gas : MIN_FHE_GASPRICE;
         uint256 paidAmountWei = effective_fhe_gasPrice * paidAmountGas;
         uint256 depositedAmount = depositsETH[payer];
-        if (paidAmountWei > depositedAmount) revert AccountNotEnoughFunded();
-        unchecked {
-            depositsETH[payer] = depositedAmount - paidAmountWei;
+        if (paidAmountWei > depositedAmount) {
+            // if dApp is not enough funded, fallbacks to user (tx.origin by default, in case of an EOA,
+            // otherwise a smart contract account should call `becomeTransientSpender` before, in the same tx
+            address spender;
+            assembly {
+                spender := tload(0)
+            }
+            spender = spender == address(0) ? tx.origin : spender;
+            if (allowedAll[spender] || whitelistedDapps[spender][payer]) {
+                uint256 depositedAmountUser = depositsETH[spender];
+                if (paidAmountWei > depositedAmountUser) revert AccountNotEnoughFunded();
+                unchecked {
+                    depositsETH[spender] = depositedAmountUser - paidAmountWei;
+                }
+                currentBlockConsumption += paidAmountGas;
+                claimableUsedFHEGas += paidAmountWei;
+            } else {
+                revert AccountNotEnoughFunded();
+            }
+        } else {
+            unchecked {
+                depositsETH[payer] = depositedAmount - paidAmountWei;
+            }
+            currentBlockConsumption += paidAmountGas;
+            claimableUsedFHEGas += paidAmountWei;
         }
-        currentBlockConsumption += paidAmountGas;
-        claimableUsedFHEGas += paidAmountWei;
     }
 
     function checkIfNewBlock() private {

diff --git a/test/paymentUtils.ts b/test/paymentUtils.ts
@@ -0,0 +1,13 @@
+import dotenv from 'dotenv';
+import fs from 'fs';
+import { ethers } from 'hardhat';
+
+export async function initializeFHEPayment() {
+  const fhePaymentFactory = await ethers.getContractFactory('FHEPayment');
+  const parsedFHEPayment = dotenv.parse(fs.readFileSync('lib/.env.fhepayment'));
+  const fhePayment = fhePaymentFactory.attach(parsedFHEPayment.FHE_PAYMENT_CONTRACT_ADDRESS);
+  return fhePayment;
+}
+
+export const FHE_GASPRICE_NATIVE_RATIO = 1000n;
+export const MIN_FHE_GASPRICE = 10_000_000n;