Merge pull request #211 from zama-ai/updateCoreContracts

feat: add transient storage cleaning for InputVerifier

contracts/.npmignore

@@ -1,6 +1,6 @@
 *
 !/artifacts/contracts/ACL.sol/ACL.json
-!/artifacts/contracts/FHEPayment.sol/FHEPayment.json
+!/artifacts/contracts/FHEGasLimit.sol/FHEGasLimit.json
 !/artifacts/contracts/InputVerifier.coprocessor.sol/InputVerifier.json
 !/artifacts/contracts/InputVerifier.native.sol/InputVerifier.json
 !/artifacts/contracts/KMSVerifier.sol/KMSVerifier.json

contracts/addresses/FHEGasLimitAddress.sol

@@ -0,0 +1,5 @@
+// SPDX-License-Identifier: BSD-3-Clause-Clear
+
+pragma solidity ^0.8.24;
+
+address constant fheGasLimitAdd = 0xFb03BE574d14C256D56F09a198B586bdfc0A9de2;

contracts/codegen/inputVerifier.ts

@@ -9,10 +9,12 @@ pragma solidity ^0.8.24;
 
 import "./KMSVerifier.sol";
 import "./TFHEExecutor.sol";
-import "../addresses/KMSVerifierAddress.sol";
-import "../addresses/CoprocessorAddress.sol";
+import "../addresses/KMSVerifierAddress.sol";\n`;
+  if (isCoprocessor) {
+    output += `import "../addresses/CoprocessorAddress.sol";\n`;
+  }
 
-// Importing OpenZeppelin contracts for cryptographic signature verification and access control.
+  output += `\n// Importing OpenZeppelin contracts for cryptographic signature verification and access control.
 import "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
 import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
 import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
@@ -116,6 +118,28 @@ contract InputVerifier is UUPSUpgradeable, Ownable2StepUpgradeable, EIP712Upgrad
     function cacheProof(bytes32 proofKey) internal virtual {
         assembly {
             tstore(proofKey, 1)
+            let length := tload(0)
+            let lengthPlusOne := add(length, 1)
+            tstore(lengthPlusOne, proofKey)
+            tstore(0, lengthPlusOne)
+        }
+    }
+
+    function cleanTransientStorage() external virtual {
+        // this function removes the transient allowances, could be useful for integration with Account Abstraction when bundling several UserOps calling InputVerifier
+        assembly {
+            let length := tload(0)
+            tstore(0, 0)
+            let lengthPlusOne := add(length, 1)
+            for {
+                let i := 1
+            } lt(i, lengthPlusOne) {
+                i := add(i, 1)
+            } {
+                let handle := tload(i)
+                tstore(i, 0)
+                tstore(handle, 0)
+            }
         }
     }
 

contracts/codegen/main.ts

@@ -3,7 +3,7 @@ import { mkdirSync, writeFileSync } from 'fs';
 import { ALL_OPERATORS, SUPPORTED_BITS, checks } from './common';
 import { generateInputVerifiers } from './inputVerifier';
 import operatorsPrices from './operatorsPrices.json';
-import { generateFHEPayment } from './payments';
+import { generateFHEGasLimit } from './payments';
 import * as t from './templates';
 import * as testgen from './testgen';
 
@@ -14,10 +14,9 @@ function generateAllFiles() {
   const ovShards = testgen.splitOverloadsToShards(overloads);
   writeFileSync('lib/Impl.sol', t.implSol(operators));
   writeFileSync('lib/TFHE.sol', tfheSolSource);
-  writeFileSync('contracts/FHEPayment.sol', generateFHEPayment(operatorsPrices));
+  writeFileSync('contracts/FHEGasLimit.sol', generateFHEGasLimit(operatorsPrices));
   writeFileSync('contracts/InputVerifier.native.sol', generateInputVerifiers(false));
   writeFileSync('contracts/InputVerifier.coprocessor.sol', generateInputVerifiers(true));
-  writeFileSync('payment/Payment.sol', t.paymentSol());
   mkdirSync('contracts/tests', { recursive: true });
   ovShards.forEach((os) => {
     writeFileSync(`examples/tests/TFHETestSuite${os.shardNumber}.sol`, testgen.generateSmartContract(os));

contracts/codegen/payments.ts

@@ -7,7 +7,7 @@ interface PriceData {
   };
 }
 
-export function generateFHEPayment(priceData: PriceData): string {
+export function generateFHEGasLimit(priceData: PriceData): string {
   let output = `// SPDX-License-Identifier: BSD-3-Clause-Clear
 
   pragma solidity ^0.8.24;
@@ -21,45 +21,32 @@ export function generateFHEPayment(priceData: PriceData): string {
   error UnsupportedOperation();
   error CallerMustBeTFHEExecutorContract();
   error OnlyScalarOperationsAreSupported();
-  error RecoveryFailed();
-  error WithdrawalFailed();
-  error AccountNotEnoughFunded();
-  error AlreadyAuthorizedAllContracts();
-  error AlreadyWhitelistedContract();
-  error AllContractsNotAuthorized();
-  error ContractNotWhitelisted();
   
-  contract FHEPayment is UUPSUpgradeable, Ownable2StepUpgradeable {
+  contract FHEGasLimit is UUPSUpgradeable, Ownable2StepUpgradeable {
       /// @notice Name of the contract
-      string private constant CONTRACT_NAME = "FHEPayment";
+      string private constant CONTRACT_NAME = "FHEGasLimit";
   
       /// @notice Version of the contract
       uint256 private constant MAJOR_VERSION = 0;
       uint256 private constant MINOR_VERSION = 1;
       uint256 private constant PATCH_VERSION = 0;
+      
       address private constant tfheExecutorAddress = tfheExecutorAdd;
-  
       uint256 private constant FHE_GAS_BLOCKLIMIT = 10_000_000;
-      uint256 private constant MIN_FHE_GASPRICE = 0; // eg: 10_000_000 means a minimum of 0.01 Gwei
-      uint256 private constant FHE_GASPRICE_NATIVE_RATIO = 0; // eg: 1000 means fhe gas price is set to 0.1% of native gas price (if above minimum)
   
-      /// @custom:storage-location erc7201:fhevm.storage.FHEPayment
-      struct FHEPaymentStorage {
+      /// @custom:storage-location erc7201:fhevm.storage.FHEGasLimit
+      struct FHEGasLimitStorage {
           uint256 lastBlock;
           uint256 currentBlockConsumption;
-          uint256 claimableUsedFHEGas;
-          mapping(address payer => uint256 depositedAmount) depositsETH;
-          mapping(address user => bool allowedAllContracts) allowedAll;
-          mapping(address user => mapping(address dappContract => bool isWhitelisted)) whitelistedDapps;
       }
   
-      // keccak256(abi.encode(uint256(keccak256("fhevm.storage.FHEPayment")) - 1)) & ~bytes32(uint256(0xff))
-      bytes32 private constant FHEPaymentStorageLocation =
-          0x4c5af501c90907b9fb888b6dd79405547def38a1dc3110f42d77f5dbc3222e00;
+      // keccak256(abi.encode(uint256(keccak256("fhevm.storage.FHEGasLimit")) - 1)) & ~bytes32(uint256(0xff))
+      bytes32 private constant FHEGasLimitStorageLocation =
+        0xb5c80b3bbe0bcbcea690f6dbe62b32a45bd1ad263b78db2f25ef8414efe9bc00;
   
-      function _getFHEPaymentStorage() internal pure returns (FHEPaymentStorage storage $) {
+      function _getFHEGasLimitStorage() internal pure returns (FHEGasLimitStorage storage $) {
           assembly {
-              $.slot := FHEPaymentStorageLocation
+              $.slot := FHEGasLimitStorageLocation
           }
       }
       
@@ -68,11 +55,6 @@ export function generateFHEPayment(priceData: PriceData): string {
           return tfheExecutorAddress;
       }
   
-      function getClaimableUsedFHEGas() public view virtual returns (uint256) {
-          FHEPaymentStorage storage $ = _getFHEPaymentStorage();
-          return $.claimableUsedFHEGas;
-      }
-  
       function _authorizeUpgrade(address _newImplementation) internal virtual override onlyOwner {}
   
       /// @custom:oz-upgrades-unsafe-allow constructor
@@ -84,116 +66,14 @@ export function generateFHEPayment(priceData: PriceData): string {
       function initialize(address initialOwner) external initializer {
           __Ownable_init(initialOwner);
       }
-  
-      function recoverBurntFunds(address receiver) external virtual onlyOwner {
-          FHEPaymentStorage storage $ = _getFHEPaymentStorage();
-          uint256 claimableUsedFHEGas_ = $.claimableUsedFHEGas;
-          $.claimableUsedFHEGas = 0;
-          (bool success, ) = receiver.call{value: claimableUsedFHEGas_}("");
-          if (!success) revert RecoveryFailed();
-      }
-  
-      function depositETH(address account) external payable virtual {
-          FHEPaymentStorage storage $ = _getFHEPaymentStorage();
-          $.depositsETH[account] += msg.value;
-      }
-  
-      function withdrawETH(uint256 amount, address receiver) external virtual {
-          FHEPaymentStorage storage $ = _getFHEPaymentStorage();
-          $.depositsETH[msg.sender] -= amount;
-          (bool success, ) = receiver.call{value: amount}("");
-          if (!success) revert WithdrawalFailed();
-      }
-  
-      function getAvailableDepositsETH(address account) external view virtual returns (uint256) {
-          FHEPaymentStorage storage $ = _getFHEPaymentStorage();
-          return $.depositsETH[account];
-      }
-  
-      function didAuthorizeAllContracts(address account) external view virtual returns (bool) {
-          FHEPaymentStorage storage $ = _getFHEPaymentStorage();
-          return $.allowedAll[account];
-      }
-  
-      function didWhitelistContract(address user, address dappContract) external view virtual returns (bool) {
-          FHEPaymentStorage storage $ = _getFHEPaymentStorage();
-          return $.whitelistedDapps[user][dappContract];
-      }
-  
-      function authorizeAllContracts() external virtual {
-          FHEPaymentStorage storage $ = _getFHEPaymentStorage();
-          if ($.allowedAll[msg.sender]) revert AlreadyAuthorizedAllContracts();
-          $.allowedAll[msg.sender] = true;
-      }
-  
-      function whitelistContract(address dappContract) external virtual {
-          FHEPaymentStorage storage $ = _getFHEPaymentStorage();
-          if ($.whitelistedDapps[msg.sender][dappContract]) revert AlreadyWhitelistedContract();
-          $.whitelistedDapps[msg.sender][dappContract] = true;
-      }
-  
-      function removeAuthorizationAllContracts() external virtual {
-          FHEPaymentStorage storage $ = _getFHEPaymentStorage();
-          if (!$.allowedAll[msg.sender]) revert AllContractsNotAuthorized();
-          $.allowedAll[msg.sender] = false;
-      }
-  
-      function removeWhitelistedContract(address dappContract) external virtual {
-          FHEPaymentStorage storage $ = _getFHEPaymentStorage();
-          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 virtual {
-          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 virtual {
-          assembly {
-              tstore(0, 0)
-          }
-      }
-  
-      function updateFunding(address payer, uint256 paidAmountGas) internal virtual {
-          FHEPaymentStorage storage $ = _getFHEPaymentStorage();
-          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 updateFunding(uint256 paidAmountGas) internal virtual {
+          FHEGasLimitStorage storage $ = _getFHEGasLimitStorage();
+          $.currentBlockConsumption += paidAmountGas;
       }
   
       function checkIfNewBlock() internal virtual {
-          FHEPaymentStorage storage $ = _getFHEPaymentStorage();
+          FHEGasLimitStorage storage $ = _getFHEGasLimitStorage();
           uint256 lastBlock_ = block.number;
           if (lastBlock_ > $.lastBlock) {
               $.lastBlock = lastBlock_;
@@ -202,19 +82,19 @@ export function generateFHEPayment(priceData: PriceData): string {
       }
   
       function checkFHEGasBlockLimit() internal view virtual {
-          FHEPaymentStorage storage $ = _getFHEPaymentStorage();
+          FHEGasLimitStorage storage $ = _getFHEGasLimitStorage();
           if ($.currentBlockConsumption >= FHE_GAS_BLOCKLIMIT) revert FHEGasBlockLimitExceeded();
       }\n\n`;
 
   for (const [operation, data] of Object.entries(priceData)) {
     const functionName = `payFor${operation.charAt(0).toUpperCase() + operation.slice(1)}`;
     if (data.binary) {
-      output += `    function ${functionName}(address payer, uint8 resultType, bytes1 scalarByte) external virtual {
+      output += `    function ${functionName}(uint8 resultType, bytes1 scalarByte) external virtual {
         if(msg.sender != tfheExecutorAddress) revert CallerMustBeTFHEExecutorContract();
         checkIfNewBlock();
 `;
     } else {
-      output += `        function ${functionName}(address payer, uint8 resultType) external virtual {
+      output += `        function ${functionName}(uint8 resultType) external virtual {
         if(msg.sender != tfheExecutorAddress) revert CallerMustBeTFHEExecutorContract();
         checkIfNewBlock();
 `;
@@ -267,7 +147,7 @@ function generatePriceChecks(prices: { [key: string]: number }): string {
     Object.entries(prices)
       .map(
         ([resultType, price]) => `        if (resultType == ${resultType}) {
-        updateFunding(payer, ${price});
+        updateFunding(${price});
         }`,
       )
       .join(' else ') + 'else { revert UnsupportedOperation();}'