Skip to content

Commit

Permalink
feat: made all fhevm contracts UUPS
Browse files Browse the repository at this point in the history 
chore: fixed prettier

fix: typo in address starting with 0

fix: get correct value of counter when doing snapshot
  • Loading branch information
@jatZama
jatZama committed Sep 12, 2024
1 parent fa5cb0e commit a2786b4
Show file tree
Hide file tree
Showing 30 changed files with 946 additions and 387 deletions.
8 changes: 2 additions & 6 deletions CustomProvider.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,13 +28,9 @@ class CustomProvider extends ProviderWrapper {
this.lastBlockSnapshotForDecrypt = parseInt(blockNumberHex);
const parsedEnvCoprocessor = dotenv.parse(fs.readFileSync('lib/.env.exec'));
const coprocAdd = parsedEnvCoprocessor.TFHE_EXECUTOR_CONTRACT_ADDRESS;
const callData = {
to: coprocAdd,
data: '0x1f20d85c',
};
this.lastCounterRand = await this._wrappedProvider.request({
method: 'eth_call',
params: [callData, 'latest'],
method: 'eth_getStorageAt',
params: [coprocAdd, '0xa436a06f0efce5ea38c956a21e24202a59b3b746d48a23fb52b4a5bc33fe3e00', 'latest'],
});
return result;
}
Expand Down
207 changes: 128 additions & 79 deletions codegen/payments.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,9 @@ export function generateFHEPayment(priceData: PriceData): string {
import "./TFHEExecutorAddress.sol";
import "@openzeppelin/contracts/utils/Strings.sol";
import "@openzeppelin/contracts/access/Ownable2Step.sol";
import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
error FHEGasBlockLimitExceeded();
error CallerMustBeTFHEExecutorContract();
error OnlyScalarOperationsAreSupported();
Expand All @@ -28,143 +29,191 @@ export function generateFHEPayment(priceData: PriceData): string {
error AllContractsNotAuthorized();
error ContractNotWhitelisted();
contract FHEPayment is Ownable2Step {
contract FHEPayment is UUPSUpgradeable, Ownable2StepUpgradeable {
/// @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;
address public constant 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;
/// @custom:storage-location erc7201:fhevm.storage.FHEPayment
struct FHEPaymentStorage {
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;
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 _getFHEPaymentStorage() internal pure returns (FHEPaymentStorage storage $) {
assembly {
$.slot := FHEPaymentStorageLocation
}
}
function getTFHEExecutorAddress() public view virtual returns (address) {
return tfheExecutorAddress;
}
function depositETH(address account) external payable {
depositsETH[account] += msg.value;
function getClaimableUsedFHEGas() public view virtual returns (uint256) {
FHEPaymentStorage storage $ = _getFHEPaymentStorage();
return $.claimableUsedFHEGas;
}
function withdrawETH(uint256 amount, address receiver) external {
depositsETH[msg.sender] -= amount;
function _authorizeUpgrade(address _newImplementation) internal virtual override onlyOwner {}
/// @custom:oz-upgrades-unsafe-allow constructor
constructor() {
_disableInitializers();
}
/// @notice Initializes the contract setting \`initialOwner\` as the initial owner
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();
if (!success) revert WithdrawalFailed();
}
function getAvailableDepositsETH(address account) external view returns (uint256) {
return depositsETH[account];
function getAvailableDepositsETH(address account) external view virtual returns (uint256) {
FHEPaymentStorage storage $ = _getFHEPaymentStorage();
return $.depositsETH[account];
}
function didAuthorizeAllContracts(address account) external view returns (bool) {
return allowedAll[account];
function didAuthorizeAllContracts(address account) external view virtual returns (bool) {
FHEPaymentStorage storage $ = _getFHEPaymentStorage();
return $.allowedAll[account];
}
function didWhitelistContract(address user, address dappContract) external view returns (bool) {
return whitelistedDapps[user][dappContract];
function didWhitelistContract(address user, address dappContract) external view virtual returns (bool) {
FHEPaymentStorage storage $ = _getFHEPaymentStorage();
return $.whitelistedDapps[user][dappContract];
}
function authorizeAllContracts() external {
if (allowedAll[msg.sender]) revert AlreadyAuthorizedAllContracts();
allowedAll[msg.sender] = true;
function authorizeAllContracts() external virtual {
FHEPaymentStorage storage $ = _getFHEPaymentStorage();
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 whitelistContract(address dappContract) external virtual {
FHEPaymentStorage storage $ = _getFHEPaymentStorage();
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 removeAuthorizationAllContracts() external virtual {
FHEPaymentStorage storage $ = _getFHEPaymentStorage();
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;
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 {
assembly {
tstore(0, caller())
}
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 {
assembly {
tstore(0, 0)
}
function stopBeingTransientSpender() external virtual {
assembly {
tstore(0, 0)
}
}
function updateFunding(address payer, uint256 paidAmountGas) private {
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];
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,
// 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 := tload(0)
}
spender = spender == address(0) ? tx.origin : spender;
if (allowedAll[spender] || whitelistedDapps[spender][payer]) {
uint256 depositedAmountUser = depositsETH[spender];
if ($.allowedAll[spender] || $.whitelistedDapps[spender][payer]) {
uint256 depositedAmountUser = $.depositsETH[spender];
if (paidAmountWei > depositedAmountUser) revert AccountNotEnoughFunded();
unchecked {
depositsETH[spender] = depositedAmountUser - paidAmountWei;
$.depositsETH[spender] = depositedAmountUser - paidAmountWei;
}
currentBlockConsumption += paidAmountGas;
claimableUsedFHEGas += paidAmountWei;
$.currentBlockConsumption += paidAmountGas;
$.claimableUsedFHEGas += paidAmountWei;
} else {
revert AccountNotEnoughFunded();
}
} else {
unchecked {
depositsETH[payer] = depositedAmount - paidAmountWei;
$.depositsETH[payer] = depositedAmount - paidAmountWei;
}
currentBlockConsumption += paidAmountGas;
claimableUsedFHEGas += paidAmountWei;
$.currentBlockConsumption += paidAmountGas;
$.claimableUsedFHEGas += paidAmountWei;
}
}
function checkIfNewBlock() private {
function checkIfNewBlock() internal virtual {
FHEPaymentStorage storage $ = _getFHEPaymentStorage();
uint256 lastBlock_ = block.number;
if (block.number > lastBlock) {
lastBlock = lastBlock_;
currentBlockConsumption = 0;
if (block.number > $.lastBlock) {
$.lastBlock = lastBlock_;
$.currentBlockConsumption = 0;
}
}
function checkFHEGasBlockLimit() internal view virtual {
FHEPaymentStorage storage $ = _getFHEPaymentStorage();
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 {
output += ` function ${functionName}(address payer, uint8 resultType, bytes1 scalarByte) external virtual {
if(msg.sender != tfheExecutorAddress) revert CallerMustBeTFHEExecutorContract();
checkIfNewBlock();
`;
} else {
output += ` function ${functionName}(address payer, uint8 resultType) external {
output += ` function ${functionName}(address payer, uint8 resultType) external virtual {
if(msg.sender != tfheExecutorAddress) revert CallerMustBeTFHEExecutorContract();
`;
}
Expand All @@ -185,15 +234,15 @@ ${generatePriceChecks(data.nonScalar)}
if (data.types) output += `${generatePriceChecks(data.types)}`;
}

output += `if (currentBlockConsumption >= FHE_GAS_BLOCKLIMIT) revert FHEGasBlockLimitExceeded();
output += `checkFHEGasBlockLimit();
}\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) {
function getVersion() external pure virtual returns (string memory) {
return
string(
abi.encodePacked(
Expand Down
32 changes: 32 additions & 0 deletions examples/ACLUpgradedExample.sol
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
// SPDX-License-Identifier: BSD-3-Clause-Clear

pragma solidity ^0.8.24;

import "../lib/ACL.sol";

contract ACLUpgradedExample is ACL {
/// @notice Name of the contract
string private constant CONTRACT_NAME = "ACL";

/// @notice Version of the contract
uint256 private constant MAJOR_VERSION = 0;
uint256 private constant MINOR_VERSION = 2;
uint256 private constant PATCH_VERSION = 0;

/// @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 virtual override returns (string memory) {
return
string(
abi.encodePacked(
CONTRACT_NAME,
" v",
Strings.toString(MAJOR_VERSION),
".",
Strings.toString(MINOR_VERSION),
".",
Strings.toString(PATCH_VERSION)
)
);
}
}
32 changes: 32 additions & 0 deletions examples/ACLUpgradedExample2.sol
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
// SPDX-License-Identifier: BSD-3-Clause-Clear

pragma solidity ^0.8.24;

import "../lib/ACL.sol";

contract ACLUpgradedExample2 is ACL {
/// @notice Name of the contract
string private constant CONTRACT_NAME = "ACL";

/// @notice Version of the contract
uint256 private constant MAJOR_VERSION = 0;
uint256 private constant MINOR_VERSION = 3;
uint256 private constant PATCH_VERSION = 0;

/// @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 virtual override returns (string memory) {
return
string(
abi.encodePacked(
CONTRACT_NAME,
" v",
Strings.toString(MAJOR_VERSION),
".",
Strings.toString(MINOR_VERSION),
".",
Strings.toString(PATCH_VERSION)
)
);
}
}
Loading

0 comments on commit a2786b4

Please sign in to comment.