@inco/lightning 0.5.3 → 0.6.0

package/src/lightning-parts/EncryptedInput.sol

@@ -2,17 +2,17 @@
 pragma solidity ^0.8;
 
 import {BaseAccessControlList} from "./AccessControl/BaseAccessControlList.sol";
-import {EventCounter} from "./primitives/EventCounter.sol";
 import {HandleGeneration} from "./primitives/HandleGeneration.sol";
 import {euint256, ebool, eaddress, ETypes} from "../Types.sol";
 import {IEncryptedInput} from "./interfaces/IEncryptedInput.sol";
 import {HandleAlreadyExists} from "../Errors.sol";
+import {Fee} from "./Fee.sol";
 
 abstract contract EncryptedInput is
     IEncryptedInput,
-    EventCounter,
     BaseAccessControlList,
-    HandleGeneration
+    HandleGeneration,
+    Fee
 {
     event NewInput(
         bytes32 indexed result,
@@ -26,21 +26,21 @@ abstract contract EncryptedInput is
     function newEuint256(
         bytes memory ciphertext,
         address user
-    ) external returns (euint256 newValue) {
+    ) external payable returns (euint256 newValue) {
         return euint256.wrap(newInput(ciphertext, user, ETypes.Uint256));
     }
 
     function newEbool(
         bytes memory ciphertext,
         address user
-    ) external returns (ebool newValue) {
+    ) external payable returns (ebool newValue) {
         return ebool.wrap(newInput(ciphertext, user, ETypes.Bool));
     }
 
     function newEaddress(
         bytes memory ciphertext,
         address user
-    ) external returns (eaddress newValue) {
+    ) external payable returns (eaddress newValue) {
         return
             eaddress.wrap(
                 newInput(ciphertext, user, ETypes.AddressOrUint160OrBytes20)
@@ -51,7 +51,23 @@ abstract contract EncryptedInput is
         bytes memory ciphertext,
         address user,
         ETypes inputType
+    ) internal paying returns (bytes32 newHandle) {
+        newHandle = _newInput(ciphertext, user, inputType);
+    }
+
+    function newInputNotPaying(
+        bytes memory ciphertext,
+        address user,
+        ETypes inputType
     ) internal returns (bytes32 newHandle) {
+        newHandle = _newInput(ciphertext, user, inputType);
+    }
+
+    function _newInput(
+        bytes memory ciphertext,
+        address user,
+        ETypes inputType
+    ) private returns (bytes32 newHandle) {
         newHandle = getInputHandle(ciphertext, user, msg.sender, inputType);
         // We assume that providing the same handle (which via HADU implies same plaintext, same context, and same
         // instance of encryption)

package/src/lightning-parts/EncryptedOperations.sol

@@ -3,15 +3,15 @@ pragma solidity ^0.8;
 
 import {euint256, ebool, EOps, SenderNotAllowedForHandle, ETypes, isTypeSupported, typeToBitMask} from "../Types.sol";
 import {BaseAccessControlList} from "./AccessControl/BaseAccessControlList.sol";
-import {EventCounter} from "./primitives/EventCounter.sol";
 import {HandleGeneration} from "./primitives/HandleGeneration.sol";
 import {IEncryptedOperations} from "./interfaces/IEncryptedOperations.sol";
+import {Fee} from "./Fee.sol";
 
 abstract contract EncryptedOperations is
     IEncryptedOperations,
-    EventCounter,
     BaseAccessControlList,
-    HandleGeneration
+    HandleGeneration,
+    Fee
 {
     error UnexpectedType(ETypes actual, bytes32 expectedTypes);
     error UnsupportedType(ETypes actual);
@@ -600,12 +600,12 @@ abstract contract EncryptedOperations is
 
     function eRand(
         ETypes randType
-    ) external returns (bytes32 result) {
+    ) external paying payable returns (bytes32 result) {
         require(isTypeSupported(randType), UnsupportedType(randType));
 
         result = createResultHandle(
             EOps.Rand,
-            randType,
+            randType, 
             bytes32(randCounter++),
             bytes32(uint256(randType))
         );
@@ -620,7 +620,7 @@ abstract contract EncryptedOperations is
     function eRandBounded(
         bytes32 upperBound,
         ETypes randType
-    ) external returns (bytes32 result) {
+    ) external paying payable returns (bytes32 result) {
         require(isTypeSupported(randType), UnsupportedType(randType));
         checkInput(upperBound, typeToBitMask(ETypes.Uint256));
 

package/src/lightning-parts/Fee.sol

@@ -0,0 +1,41 @@
+// SPDX-License-Identifier: No License
+pragma solidity ^0.8;
+
+// the fee should be modified through upgrades to limit gas cost
+uint256 constant FEE = 0.0001 ether;
+
+/// @notice Fee utils for lightning functions that require a fee
+/// @dev the fee may be changed through upgrades, develop your apps accordingly!
+abstract contract Fee {
+    error FeeNotPaid();
+
+    /// @notice the fee to pay through msg.value for inputs and randomness IT MAY CHANGE
+    function getFee() public pure returns (uint256) {
+        return FEE;
+    }
+
+    modifier paying() {
+        require(msg.value == FEE, FeeNotPaid());
+        _;
+    }
+
+    modifier payingMultiple(uint256 nbOfFees) {
+        require(msg.value == FEE * nbOfFees, FeeNotPaid());
+        _;
+    }
+
+    // Refund the difference between msg.value and what was actually spent
+    // assumes that all outflows and inflows to the contract are due to the user
+    // though the refund is capped at msg.value
+    modifier refundUnspent() {
+        uint256 balanceBefore = address(this).balance;
+        _;
+        uint256 balanceAfter = address(this).balance;
+        uint256 spent = balanceBefore > balanceAfter ? balanceBefore - balanceAfter : 0;
+        uint256 refund = msg.value > spent ? msg.value - spent : 0;
+        if (refund > 0) {
+            (bool success, ) = msg.sender.call{value: refund}("");
+            require(success, "Refund failed");
+        }
+    }
+}

package/src/lightning-parts/TEELifecycle.sol

@@ -2,8 +2,7 @@ pragma solidity ^0.8.19;
 
 import {
     BootstrapResult,
-    TEEVersion,
-    TEEVersionStatus
+    UpgradeResult
 } from "./TEELifecycle.types.sol";
 import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
 import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
@@ -28,37 +27,53 @@ abstract contract TEELifecycle is
     OwnableUpgradeable,
     EIP712Upgradeable
 {
+    // Errors
     error InvalidQuoteVerifierVersion(uint16 actual, uint16 expected);
     error EmptyTcbInfo();
     error EmptyIdentity();
     error BootstrapNotComplete();
     error BootstrapAlreadyCompleted();
-    /// @notice TEEVersionHistory must have exactly one version, please call approveNewTEEVersion first
-    error TEEVersionHistoryInconsistent();
-    error TEEVersionHistoryStatusIsNotPending();
+    /// @notice The TEE version was not found in the TEEVersionHistory
+    error TEEVersionNotFound();
     error InvalidReportMrAggregated();
-    error InvalidBootstrapDataSignature();
-    error EOASignerAlreadyInitialized();
     error InvalidReportDataSigner();
+    /// @notice The EIP712 signature recovered an incorrect address
+    error InvalidEIP712Signature();
+    error EOASignerAlreadyInitialized();
+    // @notice The EOA signer in the quote is not an existing EOA signer
+    error EOASignerNotFound();
+    // @notice The network pubkey signed by the TDX EOA is not the same as the one in the bootstrap result
+    error InvalidNetworkPubkey();
 
-    event QuoteVerifierUpdated(uint16 indexed version);
-    event TEEVersionUpdated(TEEVersion teeVersion);
+    // Events
+    // @notice A new MR_AGGREGATED has been approved
+    event NewTEEVersionApproved(uint256 indexed version, bytes32 indexed mrAggregated);
     event NewCovalidatorAdded(address covalidatorAddress, bytes quote);
     event BootstrapStageComplete(
         address indexed newEOASigner,
         BootstrapResult bootstrapResult
     );
+    // @notice Emitted to prove that an EOA has upgraded their TDX to a new
+    // MR_AGGREGATED. This is done by checking remote attestation of the quote
+    // and verifying the EIP712 signature of the bootstrap/upgrade result by
+    // the TDX EOA.
+    event EOAHasUpdatedTDX(address indexed eoaSigner, bytes32 indexed mrAggregated);
 
+    // Constants
     bytes32 public constant BootstrapResultStructHash =
         keccak256(bytes("BootstrapResult(bytes ecies_pubkey)"));
+    bytes32 public constant UpgradeResultStructHash =
+        keccak256(bytes("UpgradeResult(bytes network_pubkey)"));
 
     uint16 public constant QUOTE_VERIFIER_VERSION = 4;
 
     IQuoteVerifier public quoteVerifier;
-    BootstrapResult public VerifiedBootstrapResult;
-    bool public BootstrapComplete;
 
-    TEEVersion[] public TEEVersionHistory;
+    // @notice The list of approved TEE versions, each item is the MR_AGGREGATED 32 bytes
+    // @dev The index of the TEE version is the version number
+    bytes32[] public ApprovedTEEVersions;
+    // @notice The Network key
+    // @todo rename to NetworkPubkey https://github.com/Inco-fhevm/inco-monorepo/issues/983
     bytes public ECIESPubkey;
     mapping(address => bool) public EOASigners;
 
@@ -106,10 +121,10 @@ abstract contract TEELifecycle is
     }
 
     /**
-     * @notice Verifies the bootstrap data against the provided quote and signature
+     * @notice Wrapper around activateNewTEEVersion which verifies the bootstrap result and marks the new TEE version as active
      * @param bootstrapResult - The bootstrap data to verify
      * @param quote - The quote to verify against
-     * @param signature - The signature to verify against
+     * @param signature - The EIP712 signature of the bootstrap result by the TDX EOA
      */
     function verifyBootstrapResult(
         BootstrapResult calldata bootstrapResult,
@@ -119,54 +134,114 @@ abstract contract TEELifecycle is
         // Make sure the bootstrap is not already complete, and that the contract owner
         // has already submitted the pending TEE MR_AGGREGATED.
         require(!isBootstrapComplete(), BootstrapAlreadyCompleted());
-        require(TEEVersionHistory.length == 1, TEEVersionHistoryInconsistent());
-        require(
-            TEEVersionHistory[0].status == TEEVersionStatus.PENDING,
-            TEEVersionHistoryStatusIsNotPending()
+        require(ApprovedTEEVersions.length == 1, TEEVersionNotFound());
+
+        bytes32 digest = bootstrapResultDigest(bootstrapResult);
+        address reportDataSigner = _verifyResultForEOA(
+            ApprovedTEEVersions[0],
+            digest,
+            quote,
+            signature
+        );
+
+        // For bootstrap phase, we only have one EOA signer
+        // So if we arrive here in code, it means that the EOA has signed the bootstrap result
+
+        // Update contract publicly viewable state
+        ECIESPubkey = bootstrapResult.ecies_pubkey;
+        EOASigners[reportDataSigner] = true;
+
+        emit BootstrapStageComplete(reportDataSigner, bootstrapResult);
+    }
+
+    /**
+     * @notice Verifies the upgrade result for a single EOA. This function does not modify any contract state.
+     * It simply serves as an on-chain proof that each TDX has been updated to the newest MR_AGGREGATED.
+     * @param upgradeResult - The upgrade data to verify
+     * @param quote - The quote to verify against (contains the TDX EOA)
+     * @param signature - The EIP712 signature of the upgrade result by the TDX EOA
+     */
+    function verifyUpgradeResult(
+        bytes32 newMrAggregated,
+        UpgradeResult calldata upgradeResult,
+        bytes calldata quote,
+        bytes calldata signature
+    ) public {
+        require(isBootstrapComplete(), BootstrapNotComplete());
+        // Make sure the quote's network pubkey is the same as the one in the bootstrap result
+        require(keccak256(ECIESPubkey) == keccak256(upgradeResult.network_pubkey), InvalidNetworkPubkey());
+        // Make sure the new MR_AGGREGATED has been pre-approved.
+        bool isApproved = _isApprovedTEEVersion(newMrAggregated);
+        require(isApproved, TEEVersionNotFound());
+
+        bytes32 digest = upgradeResultDigest(upgradeResult);
+        address reportDataSigner = _verifyResultForEOA(
+            newMrAggregated,
+            digest,
+            quote,
+            signature
         );
+        // Make sure the new EOA signer is an existing EOA signer
+        require(EOASigners[reportDataSigner], EOASignerNotFound());
+    }
+
+    /**
+     * @notice Approves a new TEE version as PENDING and updates the TEEVersionHistory.
+     * The new TEE version can be activated by calling verifyBootstrapResult or activateNewTEEVersion.
+     * @param newMrAggregated - The MR_AGGREGATED bytes of the new TEE version
+     * @dev This function increments the version number automatically based on the current history
+     */
+    function approveNewTEEVersion(bytes32 newMrAggregated) public onlyOwner {
+        ApprovedTEEVersions.push(newMrAggregated);
+        emit NewTEEVersionApproved(ApprovedTEEVersions.length - 1, newMrAggregated);
+    }
+
+    function _isApprovedTEEVersion(bytes32 newMrAggregated) internal view returns (bool) {
+        for (uint256 i = 0; i < ApprovedTEEVersions.length; i++) {
+            if (ApprovedTEEVersions[i] == newMrAggregated) {
+                return true;
+            }
+        }
+        return false;
+    }
 
-        bytes32 _bootstrapResultDigest = bootstrapResultDigest(bootstrapResult);
+    /**
+     * @notice Verifies quote and signature from TDX for a single EOA. This is a shared function for bootstrap and upgrade.
+     * @param newMrAggregated - The MR_AGGREGATED bytes of the new TEE version
+     * @param resultEip712Digest - The EIP712 digest of the result to verify
+     * @param quote - The quote to verify against (contains the TDX EOA)
+     * @param signature - The EIP712 signature of the digest by the TDX EOA
+     * @return reportDataSigner - The address of the report data signer (i.e. the TDX EOA)
+     */
+    function _verifyResultForEOA(
+        bytes32 newMrAggregated,
+        bytes32 resultEip712Digest,
+        bytes calldata quote,
+        bytes calldata signature
+    ) internal onlyOwner returns (address) {
         (bool success, bytes memory output) = _verifyAndAttestOnChain(quote);
         require(success, string(output));
 
-        bytes32 v0MrAggregated = TEEVersionHistory[0].mrAggregated;
-
         TD10ReportBody memory tdReport = parseTD10ReportBody(quote);
         (address reportDataSigner, bytes32 reportMrAggregated) = parseReport(
             tdReport
         );
         require(
-            reportMrAggregated == v0MrAggregated,
+            reportMrAggregated == newMrAggregated,
             InvalidReportMrAggregated()
         );
         address recoveredAddress = ECDSA.recover(
-            _bootstrapResultDigest,
+            resultEip712Digest,
             signature
         );
         require(
             recoveredAddress == reportDataSigner,
-            InvalidBootstrapDataSignature()
+            InvalidEIP712Signature()
         );
 
-        VerifiedBootstrapResult = bootstrapResult;
-        TEEVersionHistory[0].status = TEEVersionStatus.ACTIVE;
-        emit BootstrapStageComplete(reportDataSigner, bootstrapResult);
-        ECIESPubkey = bootstrapResult.ecies_pubkey;
-        EOASigners[reportDataSigner] = true;
-        //TODO: update ECIES public key to ?? contract state and EOA addresses signers to the Signers contract state
-    }
-
-    /**
-     * @notice Approves a new TEE version and updates the TEEVersionHistory
-     * @param newMrAggregated - The MR_AGGREGATED bytes of the new TEE version
-     * @dev This function increments the version number automatically based on the current history
-     */
-    function approveNewTEEVersion(bytes32 newMrAggregated) public onlyOwner {
-        TEEVersionHistory.push(
-            TEEVersion({mrAggregated: newMrAggregated, status: TEEVersionStatus.PENDING})
-        );
+        emit EOAHasUpdatedTDX(reportDataSigner, newMrAggregated);
 
-        emit TEEVersionUpdated(TEEVersionHistory[TEEVersionHistory.length - 1]);
+        return reportDataSigner;
     }
 
     /**
@@ -184,25 +259,22 @@ abstract contract TEELifecycle is
         );
         require(!EOASigners[reportDataSigner], EOASignerAlreadyInitialized());
 
-        require(
-            reportMrAggregated ==
-                TEEVersionHistory[TEEVersionHistory.length - 1].mrAggregated,
-            InvalidReportMrAggregated()
-        );
+        bool isApproved = _isApprovedTEEVersion(reportMrAggregated);
+        require(isApproved, TEEVersionNotFound());
+
         require(reportDataSigner != address(0), InvalidReportDataSigner());
         emit NewCovalidatorAdded(reportDataSigner, quote);
         EOASigners[reportDataSigner] = true;
-        //TODO: Add the new covalidator signers to the Signers contract state
     }
 
     /**
-     * @notice Checks if the bootstrap is complete, meaning that there is an active TEE version.
+     * @notice Checks if the bootstrap is complete.
      * @return true if the bootstrap is complete, false otherwise
      */
     function isBootstrapComplete() public view returns (bool) {
-        return
-            TEEVersionHistory.length >= 1 &&
-            TEEVersionHistory[0].status == TEEVersionStatus.ACTIVE;
+        // The network pubkey is set once and once only, during the bootstrap process
+        // So we can use the length of the network pubkey to check if the bootstrap is complete
+        return ECIESPubkey.length > 0;
     }
 
     /**
@@ -361,4 +433,18 @@ abstract contract TEELifecycle is
                 )
             );
     }
+
+    function upgradeResultDigest(
+        UpgradeResult memory upgradeResult
+    ) public view returns (bytes32) {
+        return
+            _hashTypedDataV4(
+                keccak256(
+                    abi.encode(
+                        UpgradeResultStructHash,
+                        keccak256(upgradeResult.network_pubkey)
+                    )
+                )
+            );
+    }
 }

package/src/lightning-parts/TEELifecycle.types.sol

@@ -1,21 +1,18 @@
 // SPDX-License-Identifier: No License
 pragma solidity ^0.8.19;
 
+/**
+ * @notice A struct representing what the TDX EOA signs during the bootstrap process.
+ */
 struct BootstrapResult {
+    // TODO: rename to network_pubkey
+    // https://github.com/Inco-fhevm/inco-monorepo/issues/983
     bytes ecies_pubkey;
 }
 
-enum TEEVersionStatus {
-    PENDING,
-    ACTIVE
-}
-
 /**
- * @notice A struct representing a TEE version. The actual version number is the index in the TEEVersionHistory array.
- * @param mrAggregated - The MR_AGGREGATED of the TEE version
- * @param status - The status of the TEE version
+ * @notice A struct representing what the TDX EOA signs during the upgrade process.
  */
-struct TEEVersion {
-    bytes32 mrAggregated;
-    TEEVersionStatus status;
-}
+struct UpgradeResult {
+    bytes network_pubkey;
+}

package/src/lightning-parts/interfaces/IEncryptedInput.sol

@@ -7,13 +7,13 @@ interface IEncryptedInput {
     function newEuint256(
         bytes memory ciphertext,
         address user
-    ) external returns (euint256 newValue);
+    ) external payable returns (euint256 newValue);
     function newEbool(
         bytes memory ciphertext,
         address user
-    ) external returns (ebool newValue);
+    ) external payable returns (ebool newValue);
     function newEaddress(
         bytes memory ciphertext,
         address user
-    ) external returns (eaddress newValue);
+    ) external payable returns (eaddress newValue);
 }

package/src/lightning-parts/interfaces/IEncryptedOperations.sol

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: No License
 pragma solidity ^0.8;
 
-import {euint256} from "../../Types.sol";
+import {euint256, ebool, ETypes} from "../../Types.sol";
 
 interface IEncryptedOperations {
     function eAdd(
@@ -28,4 +28,51 @@ interface IEncryptedOperations {
         bytes32 lhs,
         bytes32 rhs
     ) external returns (bytes32 result);
+    function eBitOr(bytes32 lhs, bytes32 rhs) external returns (bytes32 result);
+    function eBitXor(
+        bytes32 lhs,
+        bytes32 rhs
+    ) external returns (bytes32 result);
+    function eShl(
+        euint256 lhs,
+        euint256 rhs
+    ) external returns (euint256 result);
+    function eShr(
+        euint256 lhs,
+        euint256 rhs
+    ) external returns (euint256 result);
+    function eRotl(
+        euint256 lhs,
+        euint256 rhs
+    ) external returns (euint256 result);
+    function eRotr(
+        euint256 lhs,
+        euint256 rhs
+    ) external returns (euint256 result);
+    function eEq(bytes32 lhs, bytes32 rhs) external returns (ebool result);
+    function eNe(bytes32 lhs, bytes32 rhs) external returns (ebool result);
+    function eGe(euint256 lhs, euint256 rhs) external returns (ebool result);
+    function eGt(euint256 lhs, euint256 rhs) external returns (ebool result);
+    function eLe(euint256 lhs, euint256 rhs) external returns (ebool result);
+    function eLt(euint256 lhs, euint256 rhs) external returns (ebool result);
+    function eMin(
+        euint256 lhs,
+        euint256 rhs
+    ) external returns (euint256 result);
+    function eMax(
+        euint256 lhs,
+        euint256 rhs
+    ) external returns (euint256 result);
+    function eNot(ebool operand) external returns (ebool result);
+    function eCast(bytes32 ct, ETypes toType) external returns (bytes32 result);
+    function eRand(ETypes randType) external payable returns (bytes32 result);
+    function eRandBounded(
+        bytes32 upperBound,
+        ETypes randType
+    ) external payable returns (bytes32 result);
+    function eIfThenElse(
+        ebool condition,
+        bytes32 ifTrue,
+        bytes32 ifFalse
+    ) external returns (bytes32 result);
 }

package/src/lightning-parts/test/Fee.t.sol

@@ -0,0 +1,101 @@
+// SPDX-License-Identifier: No License
+pragma solidity ^0.8;
+
+import {Test} from "forge-std/Test.sol";
+import {Fee, FEE} from "../Fee.sol";
+
+contract FeeTester is Fee {
+    function costOneFee() public payable paying {
+        // do nothing
+    }
+
+    function costMultipleVariableFees(
+        uint256 nbOfFees
+    ) public payable payingMultiple(nbOfFees) {
+        // do nothing
+    }
+
+    function costMultipleFixedFees() public payable payingMultiple(3) {
+        // do nothing
+    }
+}
+
+contract TestFee is Test {
+    FeeTester feeTester;
+
+    function setUp() public {
+        feeTester = new FeeTester();
+    }
+
+    function testGetFee() public view {
+        uint256 fee = feeTester.getFee();
+        assertEq(fee, FEE);
+    }
+
+    function testPaying() public {
+        // should fail if no fee
+        vm.expectRevert(Fee.FeeNotPaid.selector);
+        feeTester.costOneFee{value: 0}();
+
+        // should fail if not enough fee
+        vm.expectRevert(Fee.FeeNotPaid.selector);
+        feeTester.costOneFee{value: FEE - 1}();
+
+        // should fail if too much fee
+        vm.expectRevert(Fee.FeeNotPaid.selector);
+        feeTester.costOneFee{value: FEE + 1}();
+
+        // should work with exact fee
+        feeTester.costOneFee{value: FEE}();
+    }
+
+    function testPayingMultiple() public {
+        // should fail if no fee
+        vm.expectRevert(Fee.FeeNotPaid.selector);
+        feeTester.costMultipleVariableFees{value: 0}(3);
+
+        // should fail if not enough fee
+        vm.expectRevert(Fee.FeeNotPaid.selector);
+        feeTester.costMultipleVariableFees{value: FEE * 3 - 1}(3);
+
+        // should fail if too much fee
+        vm.expectRevert(Fee.FeeNotPaid.selector);
+        feeTester.costMultipleVariableFees{value: FEE * 3 + 1}(3);
+
+        // should work with exact fee
+        feeTester.costMultipleVariableFees{value: FEE * 3}(3);
+
+        // should work with exact fee for fixed number of fees
+        feeTester.costMultipleFixedFees{value: FEE * 3}();
+    }
+
+    function testPayingMultipleZero() public {
+        feeTester.costMultipleVariableFees{value: 0}(0);
+
+        // should fail if too much fee
+        vm.expectRevert(Fee.FeeNotPaid.selector);
+        feeTester.costMultipleVariableFees{value: FEE + 1}(0);
+
+        // should work with exact fee
+        feeTester.costMultipleVariableFees{value: 0}(0);
+    }
+
+    function testFuzzPayingMultipleVariableFees(uint8 nbOfFees) public {
+        uint256 expectedValue = FEE * nbOfFees;
+
+        // should work with exact fee
+        feeTester.costMultipleVariableFees{value: expectedValue}(nbOfFees);
+
+        // should fail if not enough fee (unless expectedValue is 0)
+        if (expectedValue > 0) {
+            vm.expectRevert(Fee.FeeNotPaid.selector);
+            feeTester.costMultipleVariableFees{value: expectedValue - 1}(
+                nbOfFees
+            );
+        }
+
+        // should fail if too much fee
+        vm.expectRevert(Fee.FeeNotPaid.selector);
+        feeTester.costMultipleVariableFees{value: expectedValue + 1}(nbOfFees);
+    }
+}

package/src/lightning-parts/test/HandleMetadata.t.sol

@@ -15,6 +15,7 @@ import {
     typeToBitMask
 } from "../../Types.sol";
 import {VerifierAddressGetter} from "../primitives/VerifierAddressGetter.sol";
+import {FEE} from "../Fee.sol";
 
 contract TestHandleMetadata is
     EIP712,
@@ -107,11 +108,11 @@ contract TestHandleMetadata is
     }
 
     function testEncryptedInputHandleType() public {
-        euint256 a = this.newEuint256("ciphertext", address(this));
+        euint256 a = this.newEuint256{value: FEE}("ciphertext", address(this));
         assert(typeOf(euint256.unwrap(a)) == ETypes.Uint256);
-        ebool b = this.newEbool("ciphertext", address(this));
+        ebool b = this.newEbool{value: FEE}("ciphertext", address(this));
         assert(typeOf(ebool.unwrap(b)) == ETypes.Bool);
-        eaddress c = this.newEaddress("ciphertext", address(this));
+        eaddress c = this.newEaddress{value: FEE}("ciphertext", address(this));
         assert(typeOf(eaddress.unwrap(c)) == ETypes.AddressOrUint160OrBytes20);
     }
 }