Merge pull request #479 from zama-ai/fhePayment

feat: FHEPayment v0

codegen/main.ts

@@ -1,6 +1,8 @@
 import { mkdirSync, writeFileSync } from 'fs';
 
 import { ALL_OPERATORS, Network, SUPPORTED_BITS, checks, networkCodegenContext } from './common';
+import operatorsPrices from './operatorsPrices.json';
+import { generateFHEPayment } from './payments';
 import * as t from './templates';
 import * as testgen from './testgen';
 
@@ -15,6 +17,8 @@ function generateAllFiles() {
   writeFileSync('lib/Impl.sol', t.implSol(context, operators));
   writeFileSync('lib/TFHE.sol', tfheSolSource);
   writeFileSync('lib/FhevmLib.sol', t.fhevmLibSol(operators));
+  writeFileSync('lib/FHEPayment.sol', generateFHEPayment(operatorsPrices));
+  writeFileSync('payment/Payment.sol', t.paymentSol());
   mkdirSync('examples/tests', { recursive: true });
   ovShards.forEach((os) => {
     writeFileSync(`examples/tests/TFHETestSuite${os.shardNumber}.sol`, testgen.generateSmartContract(os));

codegen/operatorsPrices.json

@@ -0,0 +1,119 @@
+{
+  "fheAdd": {
+    "binary": true,
+    "scalar": { "1": 65000, "2": 94000, "3": 133000, "4": 162000, "5": 188000 },
+    "nonScalar": { "1": 65000, "2": 94000, "3": 133000, "4": 162000, "5": 188000 }
+  },
+  "fheSub": {
+    "binary": true,
+    "scalar": { "1": 65000, "2": 94000, "3": 133000, "4": 162000, "5": 188000 },
+    "nonScalar": { "1": 65000, "2": 94000, "3": 133000, "4": 162000, "5": 188000 }
+  },
+  "fheMul": {
+    "binary": true,
+    "scalar": { "1": 88000, "2": 159000, "3": 208000, "4": 264000, "5": 356000 },
+    "nonScalar": { "1": 150000, "2": 197000, "3": 262000, "4": 359000, "5": 641000 }
+  },
+  "fheDiv": { "binary": true, "scalar": { "1": 139000, "2": 238000, "3": 314000, "4": 398000, "5": 584000 } },
+  "fheRem": { "binary": true, "scalar": { "1": 286000, "2": 460000, "3": 622000, "4": 805000, "5": 1095000 } },
+  "fheBitAnd": {
+    "binary": true,
+    "nonScalar": { "0": 26000, "1": 32000, "2": 34000, "3": 34000, "4": 35000, "5": 38000 }
+  },
+  "fheBitOr": {
+    "binary": true,
+    "nonScalar": { "0": 26000, "1": 32000, "2": 34000, "3": 34000, "4": 35000, "5": 38000 }
+  },
+  "fheBitXor": {
+    "binary": true,
+    "nonScalar": { "0": 26000, "1": 32000, "2": 34000, "3": 34000, "4": 35000, "5": 38000 }
+  },
+  "fheShl": {
+    "binary": true,
+    "scalar": { "1": 35000, "2": 35000, "3": 35000, "4": 35000, "5": 38000 },
+    "nonScalar": { "1": 116000, "2": 133000, "3": 153000, "4": 183000, "5": 227000 }
+  },
+  "fheShr": {
+    "binary": true,
+    "scalar": { "1": 35000, "2": 35000, "3": 35000, "4": 35000, "5": 38000 },
+    "nonScalar": { "1": 116000, "2": 133000, "3": 153000, "4": 183000, "5": 227000 }
+  },
+  "fheRotl": {
+    "binary": true,
+    "scalar": { "1": 35000, "2": 35000, "3": 35000, "4": 35000, "5": 38000 },
+    "nonScalar": { "1": 116000, "2": 133000, "3": 153000, "4": 183000, "5": 227000 }
+  },
+  "fheRotr": {
+    "binary": true,
+    "scalar": { "1": 35000, "2": 35000, "3": 35000, "4": 35000, "5": 38000 },
+    "nonScalar": { "1": 116000, "2": 133000, "3": 153000, "4": 183000, "5": 227000 }
+  },
+  "fheEq": {
+    "binary": true,
+    "scalar": { "1": 51000, "2": 53000, "3": 54000, "4": 82000, "5": 86000, "7": 90000, "11": 300000 },
+    "nonScalar": { "1": 51000, "2": 53000, "3": 54000, "4": 82000, "5": 86000, "7": 90000, "11": 300000 }
+  },
+  "fheNe": {
+    "binary": true,
+    "scalar": { "1": 51000, "2": 53000, "3": 54000, "4": 82000, "5": 86000, "7": 90000, "11": 300000 },
+    "nonScalar": { "1": 51000, "2": 53000, "3": 54000, "4": 82000, "5": 86000, "7": 90000, "11": 300000 }
+  },
+  "fheGe": {
+    "binary": true,
+    "scalar": { "1": 70000, "2": 82000, "3": 105000, "4": 128000, "5": 156000 },
+    "nonScalar": { "1": 70000, "2": 82000, "3": 105000, "4": 128000, "5": 156000 }
+  },
+  "fheGt": {
+    "binary": true,
+    "scalar": { "1": 70000, "2": 82000, "3": 105000, "4": 128000, "5": 156000 },
+    "nonScalar": { "1": 70000, "2": 82000, "3": 105000, "4": 128000, "5": 156000 }
+  },
+  "fheLe": {
+    "binary": true,
+    "scalar": { "1": 70000, "2": 82000, "3": 105000, "4": 128000, "5": 156000 },
+    "nonScalar": { "1": 70000, "2": 82000, "3": 105000, "4": 128000, "5": 156000 }
+  },
+  "fheLt": {
+    "binary": true,
+    "scalar": { "1": 70000, "2": 82000, "3": 105000, "4": 128000, "5": 156000 },
+    "nonScalar": { "1": 70000, "2": 82000, "3": 105000, "4": 128000, "5": 156000 }
+  },
+  "fheMin": {
+    "binary": true,
+    "scalar": { "1": 121000, "2": 128000, "3": 150000, "4": 164000, "5": 192000 },
+    "nonScalar": { "1": 121000, "2": 128000, "3": 153000, "4": 183000, "5": 210000 }
+  },
+  "fheMax": {
+    "binary": true,
+    "scalar": { "1": 121000, "2": 128000, "3": 150000, "4": 164000, "5": 192000 },
+    "nonScalar": { "1": 121000, "2": 128000, "3": 153000, "4": 183000, "5": 210000 }
+  },
+  "fheNeg": {
+    "binary": false,
+    "types": { "1": 60000, "2": 95000, "3": 131000, "4": 160000, "5": 199000 }
+  },
+  "fheNot": {
+    "binary": false,
+    "types": { "0": 30000, "1": 33000, "2": 34000, "3": 35000, "4": 36000, "5": 37000 }
+  },
+  "cast": {
+    "binary": false,
+    "types": { "1": 200, "2": 200, "3": 200, "4": 200, "5": 200 }
+  },
+  "trivialEncrypt": {
+    "binary": false,
+    "types": { "0": 100, "1": 100, "2": 100, "3": 200, "4": 300, "5": 600, "7": 700 }
+  },
+  "ifThenElse": {
+    "binary": false,
+    "types": { "1": 45000, "2": 47000, "3": 47000, "4": 50000, "5": 53000, "7": 80000 }
+  },
+  "fheRand": {
+    "binary": false,
+    "types": { "2": 100000, "3": 100000, "4": 100000, "5": 100000 }
+  },
+  "fheRandBounded": {
+    "binary": false,
+    "types": { "2": 100000, "3": 100000, "4": 100000, "5": 100000 }
+  }
+}

codegen/payments.ts

@@ -0,0 +1,222 @@
+interface PriceData {
+  [key: string]: {
+    binary: boolean;
+    scalar?: { [key: string]: number };
+    nonScalar?: { [key: string]: number };
+    types?: { [key: string]: number };
+  };
+}
+
+export function generateFHEPayment(priceData: PriceData): string {
+  let output = `// SPDX-License-Identifier: BSD-3-Clause-Clear
+
+  pragma solidity ^0.8.24;
+  
+  import "./TFHEExecutorAddress.sol";
+  import "@openzeppelin/contracts/utils/Strings.sol";
+  import "@openzeppelin/contracts/access/Ownable2Step.sol";
+
+  error FHEGasBlockLimitExceeded();
+  error CallerMustBeTFHEExecutorContract();
+  error OnlyScalarOperationsAreSupported();
+  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
+      string private constant CONTRACT_NAME = "FHEPayment";
+  
+      /// @notice Version of the contract
+      uint256 private constant MAJOR_VERSION = 0;
+      uint256 private constant MINOR_VERSION = 1;
+      uint256 private constant PATCH_VERSION = 0;
+      address public immutable tfheExecutorAddress = tfheExecutorAdd;
+  
+      uint256 private constant FHE_GAS_BLOCKLIMIT = 10_000_000;
+      uint256 private constant MIN_FHE_GASPRICE = 10_000_000; // minimum of 0.01 Gwei
+      uint256 private constant FHE_GASPRICE_NATIVE_RATIO = 1000; // fhe gas price is set to 0.1% of native gas price (if above minimum)
+  
+      uint256 private lastBlock;
+      uint256 private currentBlockConsumption;
+      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) {}
+
+      function recoverBurntFunds(address receiver) external onlyOwner {
+        uint256 claimableUsedFHEGas_ = claimableUsedFHEGas;
+        claimableUsedFHEGas = 0;
+        (bool success, ) = receiver.call{value: claimableUsedFHEGas_}("");
+        if(!success) revert RecoveryFailed();
+      }
+  
+      function depositETH(address account) external payable {
+          depositsETH[account] += msg.value;
+      }
+  
+      function withdrawETH(uint256 amount, address receiver) external {
+          depositsETH[msg.sender] -= amount;
+          (bool success, ) = receiver.call{value: amount}("");
+          if(!success) revert WithdrawalFailed();
+      }
+  
+      function getAvailableDepositsETH(address account) external view returns (uint256) {
+          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) {
+              // 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) {
+              lastBlock = lastBlock_;
+              currentBlockConsumption = 0;
+          }
+      }\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 {
+        if(msg.sender != tfheExecutorAddress) revert CallerMustBeTFHEExecutorContract();
+        checkIfNewBlock();
+`;
+    } else {
+      output += `    function ${functionName}(address payer, uint8 resultType) external {
+        if(msg.sender != tfheExecutorAddress) revert CallerMustBeTFHEExecutorContract();
+`;
+    }
+
+    if (data.scalar && data.nonScalar) {
+      output += `        if (scalarByte == 0x01) {
+${generatePriceChecks(data.scalar)}
+        } else {
+${generatePriceChecks(data.nonScalar)}
+        }`;
+    } else if (data.scalar) {
+      output += `        if(scalarByte != 0x01) revert OnlyScalarOperationsAreSupported();`;
+      output += `${generatePriceChecks(data.scalar)}`;
+    } else if (data.nonScalar) {
+      output += `        if(scalarByte != 0x00) revert OnlyNonScalarOperationsAreSupported();`;
+      output += `${generatePriceChecks(data.nonScalar)}`;
+    } else {
+      if (data.types) output += `${generatePriceChecks(data.types)}`;
+    }
+
+    output += `if (currentBlockConsumption >= FHE_GAS_BLOCKLIMIT) revert FHEGasBlockLimitExceeded();
+    }\n\n`;
+  }
+
+  return (
+    output +
+    `    /// @notice Getter for the name and version of the contract
+    /// @return string representing the name and the version of the contract
+    function getVersion() external pure returns (string memory) {
+        return
+            string(
+                abi.encodePacked(
+                    CONTRACT_NAME,
+                    " v",
+                    Strings.toString(MAJOR_VERSION),
+                    ".",
+                    Strings.toString(MINOR_VERSION),
+                    ".",
+                    Strings.toString(PATCH_VERSION)
+                )
+            );
+    }
+}`
+  );
+}
+
+function generatePriceChecks(prices: { [key: string]: number }): string {
+  return Object.entries(prices)
+    .map(
+      ([resultType, price]) => `        if (resultType == ${resultType}) {
+        updateFunding(payer, ${price});
+        }`,
+    )
+    .join(' else ');
+}