@inco/lightning 0.6.1 → 0.6.5

package/src/lightning-parts/TEELifecycle.sol

@@ -1,9 +1,6 @@
 pragma solidity ^0.8.19;
 
-import {
-    BootstrapResult,
-    UpgradeResult
-} from "./TEELifecycle.types.sol";
+import {BootstrapResult, UpgradeResult, AddNodeResult} from "./TEELifecycle.types.sol";
 import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
 import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import {IQuoteVerifier} from "../interfaces/automata-interfaces/IQuoteVerifier.sol";
@@ -21,8 +18,37 @@ import {
 import {IFmspcTcbDao} from "../interfaces/automata-interfaces/IFmspcTcbDao.sol";
 import {IAutomataEnclaveIdentityDao} from "../interfaces/automata-interfaces/IAutomataEnclaveIdentityDao.sol";
 
-// todo #1031 make TEELifecycle storage upgrade compatible
+contract TEELifecycleStorage {
+    struct StorageForTEELifecycle {
+        IQuoteVerifier quoteVerifier;
+        BootstrapResult VerifiedBootstrapResult;
+        bool BootstrapComplete;
+        // @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[] ApprovedTEEVersions;
+        // @notice The Network key
+        // @todo rename to NetworkPubkey https://github.com/Inco-fhevm/inco-monorepo/issues/983
+        bytes ECIESPubkey;
+        mapping(address => bool) EOASigners;
+    }
+
+    bytes32 private constant TEELifecycleStorageLocation =
+        keccak256("inco.storage.TEELifecycle");
+
+    function getTEELifecycleStorage()
+        internal
+        pure
+        returns (StorageForTEELifecycle storage $)
+    {
+        bytes32 loc = TEELifecycleStorageLocation;
+        assembly {
+            $.slot := loc
+        }
+    }
+}
+
 abstract contract TEELifecycle is
+    TEELifecycleStorage,
     ITEELifecycle,
     OwnableUpgradeable,
     EIP712Upgradeable
@@ -37,6 +63,7 @@ abstract contract TEELifecycle is
     error TEEVersionNotFound();
     error InvalidReportMrAggregated();
     error InvalidReportDataSigner();
+    error IndexOutOfBounds();
     /// @notice The EIP712 signature recovered an incorrect address
     error InvalidEIP712Signature();
     error EOASignerAlreadyInitialized();
@@ -47,7 +74,10 @@ abstract contract TEELifecycle is
 
     // Events
     // @notice A new MR_AGGREGATED has been approved
-    event NewTEEVersionApproved(uint256 indexed version, bytes32 indexed mrAggregated);
+    event NewTEEVersionApproved(
+        uint256 indexed version,
+        bytes32 indexed mrAggregated
+    );
     event NewCovalidatorAdded(address covalidatorAddress, bytes quote);
     event BootstrapStageComplete(
         address indexed newEOASigner,
@@ -57,31 +87,26 @@ abstract contract TEELifecycle is
     // 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);
+    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)"));
+    bytes32 public constant AddNodeResultStructHash =
+        keccak256(bytes("AddNodeResult(bytes network_pubkey)"));
 
     uint16 public constant QUOTE_VERIFIER_VERSION = 4;
 
-    IQuoteVerifier public quoteVerifier;
-
-    // @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;
-
     function __TEELifecycle_init(
         IQuoteVerifier _quoteVerifier
     ) internal onlyInitializing {
-        quoteVerifier = _quoteVerifier;
-        uint16 actualQuoteVerifierVersion = quoteVerifier.quoteVersion();
+        getTEELifecycleStorage().quoteVerifier = _quoteVerifier;
+        uint16 actualQuoteVerifierVersion = _quoteVerifier.quoteVersion();
         require(
             actualQuoteVerifierVersion == QUOTE_VERIFIER_VERSION,
             InvalidQuoteVerifierVersion(
@@ -103,12 +128,14 @@ abstract contract TEELifecycle is
         require(bytes(tcbInfo.tcbInfoStr).length != 0, EmptyTcbInfo());
         require(bytes(identity.identityStr).length != 0, EmptyIdentity());
 
+        IQuoteVerifier _quoteVerifier = quoteVerifier();
+
         IFmspcTcbDao fmspcTcbDao = IFmspcTcbDao(
-            quoteVerifier.pccsRouter().fmspcTcbDaoAddr()
+            _quoteVerifier.pccsRouter().fmspcTcbDaoAddr()
         );
         fmspcTcbDao.upsertFmspcTcb(tcbInfo);
         IAutomataEnclaveIdentityDao enclaveIdDao = IAutomataEnclaveIdentityDao(
-            quoteVerifier.pccsRouter().qeIdDaoAddr()
+            _quoteVerifier.pccsRouter().qeIdDaoAddr()
         );
         (IdentityObj memory identityObj, ) = enclaveIdDao
             .EnclaveIdentityLib()
@@ -131,14 +158,16 @@ abstract contract TEELifecycle is
         bytes calldata quote,
         bytes calldata signature
     ) public onlyOwner {
+        StorageForTEELifecycle storage $ = getTEELifecycleStorage();
+
         // 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(ApprovedTEEVersions.length == 1, TEEVersionNotFound());
+        require($.ApprovedTEEVersions.length == 1, TEEVersionNotFound());
 
         bytes32 digest = bootstrapResultDigest(bootstrapResult);
         address reportDataSigner = _verifyResultForEOA(
-            ApprovedTEEVersions[0],
+            $.ApprovedTEEVersions[0],
             digest,
             quote,
             signature
@@ -148,8 +177,8 @@ abstract contract TEELifecycle is
         // 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;
+        $.ECIESPubkey = bootstrapResult.ecies_pubkey;
+        $.EOASigners[reportDataSigner] = true;
 
         emit BootstrapStageComplete(reportDataSigner, bootstrapResult);
     }
@@ -167,9 +196,14 @@ abstract contract TEELifecycle is
         bytes calldata quote,
         bytes calldata signature
     ) public {
+        StorageForTEELifecycle storage $ = getTEELifecycleStorage();
+
         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());
+        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());
@@ -182,7 +216,41 @@ abstract contract TEELifecycle is
             signature
         );
         // Make sure the new EOA signer is an existing EOA signer
-        require(EOASigners[reportDataSigner], EOASignerNotFound());
+        require($.EOASigners[reportDataSigner], EOASignerNotFound());
+    }
+
+    /**
+     * @notice Verifies the add node result, and adds the new EOA signer to the contract state.
+     * @param quote - The quote from the new covalidator that contains the current MR_AGGREGATED and the eoa address of the new party in the report data
+     */
+    function verifyAddNodeResult(
+        bytes32 newMrAggregated,
+        AddNodeResult calldata addNodeResult,
+        bytes calldata quote,
+        bytes calldata signature
+    ) public onlyOwner {
+        StorageForTEELifecycle storage $ = getTEELifecycleStorage();
+
+        require(isBootstrapComplete(), BootstrapNotComplete());
+        // Make sure the quote's network pubkey is the same as the one in the bootstrap result
+        require(
+            keccak256($.ECIESPubkey) == keccak256(addNodeResult.network_pubkey),
+            InvalidNetworkPubkey()
+        );
+        // Make sure the new MR_AGGREGATED has been pre-approved.
+        bool isApproved = _isApprovedTEEVersion(newMrAggregated);
+        require(isApproved, TEEVersionNotFound());
+
+        bytes32 digest = addNodeResultDigest(addNodeResult);
+        address reportDataSigner = _verifyResultForEOA(
+            newMrAggregated,
+            digest,
+            quote,
+            signature
+        );
+
+        // Add this new EOA signer as a new signer to the contract state
+        $.EOASigners[reportDataSigner] = true;
     }
 
     /**
@@ -192,13 +260,20 @@ abstract contract TEELifecycle is
      * @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);
+        StorageForTEELifecycle storage $ = getTEELifecycleStorage();
+        $.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) {
+    function _isApprovedTEEVersion(
+        bytes32 newMrAggregated
+    ) internal view returns (bool) {
+        StorageForTEELifecycle storage $ = getTEELifecycleStorage();
+        for (uint256 i = 0; i < $.ApprovedTEEVersions.length; i++) {
+            if ($.ApprovedTEEVersions[i] == newMrAggregated) {
                 return true;
             }
         }
@@ -230,43 +305,14 @@ abstract contract TEELifecycle is
             reportMrAggregated == newMrAggregated,
             InvalidReportMrAggregated()
         );
-        address recoveredAddress = ECDSA.recover(
-            resultEip712Digest,
-            signature
-        );
-        require(
-            recoveredAddress == reportDataSigner,
-            InvalidEIP712Signature()
-        );
+        address recoveredAddress = ECDSA.recover(resultEip712Digest, signature);
+        require(recoveredAddress == reportDataSigner, InvalidEIP712Signature());
 
         emit EOAHasUpdatedTDX(reportDataSigner, newMrAggregated);
 
         return reportDataSigner;
     }
 
-    /**
-     * @notice Adds a new covalidator to the contract state
-     * @param quote - The quote from the new covalidator that contains the current MR_AGGREGATED and the eoa address of the new party in the report data
-     */
-    function addNewCovalidator(bytes calldata quote) public onlyOwner {
-        require(isBootstrapComplete(), BootstrapNotComplete());
-
-        (bool success, bytes memory output) = _verifyAndAttestOnChain(quote);
-        require(success, string(output));
-        TD10ReportBody memory tdReport = parseTD10ReportBody(quote);
-        (address reportDataSigner, bytes32 reportMrAggregated) = parseReport(
-            tdReport
-        );
-        require(!EOASigners[reportDataSigner], EOASignerAlreadyInitialized());
-
-        bool isApproved = _isApprovedTEEVersion(reportMrAggregated);
-        require(isApproved, TEEVersionNotFound());
-
-        require(reportDataSigner != address(0), InvalidReportDataSigner());
-        emit NewCovalidatorAdded(reportDataSigner, quote);
-        EOASigners[reportDataSigner] = true;
-    }
-
     /**
      * @notice Checks if the bootstrap is complete.
      * @return true if the bootstrap is complete, false otherwise
@@ -274,7 +320,7 @@ abstract contract TEELifecycle is
     function isBootstrapComplete() public view returns (bool) {
         // 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;
+        return getTEELifecycleStorage().ECIESPubkey.length > 0;
     }
 
     /**
@@ -302,7 +348,10 @@ abstract contract TEELifecycle is
 
         // We found a supported version, begin verifying the quote
         // Note: The quote header cannot be trusted yet, it will be validated by the Verifier library
-        (success, output) = quoteVerifier.verifyQuote(header, rawQuote);
+        (success, output) = getTEELifecycleStorage().quoteVerifier.verifyQuote(
+            header,
+            rawQuote
+        );
     }
 
     /**
@@ -397,7 +446,12 @@ abstract contract TEELifecycle is
     ) public pure returns (address, bytes32) {
         return (
             address(bytes20(tdReport.reportData)),
-            computeMrAggregated(tdReport.mrTd, tdReport.rtMr0, tdReport.rtMr1, tdReport.rtMr2)
+            computeMrAggregated(
+                tdReport.mrTd,
+                tdReport.rtMr0,
+                tdReport.rtMr1,
+                tdReport.rtMr2
+            )
         );
     }
 
@@ -420,6 +474,26 @@ abstract contract TEELifecycle is
         return keccak256(message);
     }
 
+    function quoteVerifier() public view returns (IQuoteVerifier) {
+        return getTEELifecycleStorage().quoteVerifier;
+    }
+
+    function approvedTEEVersions(
+        uint256 index
+    ) external view returns (bytes32) {
+        StorageForTEELifecycle storage $ = getTEELifecycleStorage();
+        require(index < $.ApprovedTEEVersions.length, IndexOutOfBounds());
+        return $.ApprovedTEEVersions[index];
+    }
+
+    function eciesPubkey() external view returns (bytes memory) {
+        return getTEELifecycleStorage().ECIESPubkey;
+    }
+
+    function isEOASigner(address addr) external view returns (bool) {
+        return getTEELifecycleStorage().EOASigners[addr];
+    }
+
     function bootstrapResultDigest(
         BootstrapResult memory bootstrapResult
     ) public view returns (bytes32) {
@@ -447,4 +521,17 @@ abstract contract TEELifecycle is
                 )
             );
     }
+
+    function addNodeResultDigest(
+        AddNodeResult memory addNodeResult
+    ) public view returns (bytes32) {
+        return
+            _hashTypedDataV4(
+                keccak256(abi.encode(
+                    AddNodeResultStructHash,
+                    keccak256(addNodeResult.network_pubkey)
+                )
+            )
+        );
+    }
 }

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

@@ -16,3 +16,10 @@ struct BootstrapResult {
 struct UpgradeResult {
     bytes network_pubkey;
 }
+
+/**
+ * @notice A struct representing what the TDX EOA signs during the add node process.
+ */
+struct AddNodeResult {
+    bytes network_pubkey;
+}

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

@@ -1,7 +1,7 @@
 /// SPDX-License-Identifier: No License
 pragma solidity ^0.8.19;
 
-import {BootstrapResult} from "../TEELifecycle.types.sol";
+import {BootstrapResult, UpgradeResult, AddNodeResult} from "../TEELifecycle.types.sol";
 import {IQuoteVerifier} from "../../interfaces/automata-interfaces/IQuoteVerifier.sol";
 import {TD10ReportBody} from "../../interfaces/automata-interfaces/Types.sol";
 
@@ -11,8 +11,19 @@ interface ITEELifecycle {
         bytes calldata quote,
         bytes calldata signature
     ) external;
+    function verifyUpgradeResult(
+        bytes32 newMrAggregated,
+        UpgradeResult calldata upgradeResult,
+        bytes calldata quote,
+        bytes calldata signature
+    ) external;
+    function verifyAddNodeResult(
+        bytes32 newMrAggregated,
+        AddNodeResult calldata addNodeResult,
+        bytes calldata quote,
+        bytes calldata signature
+    ) external;
     function approveNewTEEVersion(bytes32 newMrAggregated) external;
-    function addNewCovalidator(bytes calldata quote) external;
     function parseTD10ReportBody(
         bytes calldata rawQuote
     ) external pure returns (TD10ReportBody memory report);

package/src/lightning-parts/primitives/HandleGeneration.sol

@@ -17,12 +17,8 @@ contract HandleGeneration is IHandleGeneration, HandleMetadata {
     ) public view returns (bytes32 generatedHandle) {
         generatedHandle = keccak256(
             abi.encodePacked(
-                EVM_HOST_CHAIN_PREFIX,
-                block.chainid, // todo cache this
-                EOps.TrivialEncrypt,
-                plaintextBytes,
-                handleType,
-                address(this) // todo cache this
+                EOps.TrivialEncrypt, // !! Note[Silas]: I reordered this to be first element for greater regularity 26/08/2025 (remove this note after 1 month) !!
+                plaintextBytes
             )
         );
         generatedHandle = embedTypeVersion(generatedHandle, handleType);
@@ -33,7 +29,24 @@ contract HandleGeneration is IHandleGeneration, HandleMetadata {
         address user,
         address contractAddress,
         ETypes inputType
-    ) public view returns (bytes32 generatedHandle) {
+    ) internal view returns (bytes32 generatedHandle) {
+        return
+            getInputHandle(
+            ciphertext,
+            address(this),
+            user,
+            contractAddress,
+            inputType
+        );
+    }
+
+    function getInputHandle(
+        bytes memory ciphertext,
+        address executorAddress,
+        address user,
+        address contractAddress,
+        ETypes inputType
+    ) internal view returns (bytes32 generatedHandle) {
         // Here we ensure that our hashing scheme is binary-compatible between IncoLightning and IncoFhevm, this helps
         // keep client-side code consistent in its ciphertext, context => handle mappings
         bytes32 ctIndexHash = keccak256(
@@ -47,7 +60,7 @@ contract HandleGeneration is IHandleGeneration, HandleMetadata {
                     prehandle,
                     EVM_HOST_CHAIN_PREFIX,
                     block.chainid, // todo cache this
-                    address(this), // todo cache this
+                    executorAddress,
                     user,
                     contractAddress
                 )
@@ -59,55 +72,8 @@ contract HandleGeneration is IHandleGeneration, HandleMetadata {
     function getOpResultHandle(
         EOps op,
         ETypes returnType,
-        bytes32 lhs,
-        bytes32 rhs
+        bytes memory packedInputs
     ) public pure returns (bytes32 generatedHandle) {
-        generatedHandle = getOpResultHandle(
-            keccak256(abi.encodePacked(op, lhs, rhs)),
-            returnType
-        );
-    }
-
-    function getOpResultHandle(
-        EOps op,
-        ETypes returnType,
-        uint256 counter,
-        bytes32 upperBound
-    ) public pure returns (bytes32 generatedHandle) {
-        generatedHandle = getOpResultHandle(
-            keccak256(abi.encodePacked(op, counter, upperBound)),
-            returnType
-        );
-    }
-
-    function getOpResultHandle(
-        EOps op,
-        ETypes returnType,
-        bytes32 value
-    ) public pure returns (bytes32 generatedHandle) {
-        generatedHandle = getOpResultHandle(
-            keccak256(abi.encodePacked(op, value)),
-            returnType
-        );
-    }
-
-    function getOpResultHandle(
-        EOps op,
-        ETypes returnType,
-        bytes32 inputA,
-        bytes32 inputB,
-        bytes32 inputC
-    ) public pure returns (bytes32 generatedHandle) {
-        generatedHandle = getOpResultHandle(
-            keccak256(abi.encodePacked(op, inputA, inputB, inputC)),
-            returnType
-        );
-    }
-
-    function getOpResultHandle(
-        bytes32 baseHandle,
-        ETypes returnType
-    ) internal pure returns (bytes32 generatedHandle) {
-        generatedHandle = embedTypeVersion(baseHandle, returnType);
+        generatedHandle = embedTypeVersion(keccak256(abi.encodePacked(op, packedInputs)), returnType);
     }
 }

package/src/lightning-parts/primitives/interfaces/IHandleGeneration.sol

@@ -4,38 +4,7 @@ pragma solidity ^0.8;
 import {ETypes, EOps} from "../../../Types.sol";
 
 interface IHandleGeneration {
-    function getTrivialEncryptHandle(
-        bytes32 plaintextBytes,
-        ETypes handleType
-    ) external view returns (bytes32 generatedHandle);
-    function getInputHandle(
-        bytes memory ciphertext,
-        address user,
-        address contractAddress,
-        ETypes inputType
-    ) external view returns (bytes32 generatedHandle);
-    function getOpResultHandle(
-        EOps op,
-        ETypes returnType,
-        bytes32 lhs,
-        bytes32 rhs
-    ) external pure returns (bytes32 generatedHandle);
-    function getOpResultHandle(
-        EOps op,
-        ETypes returnType,
-        uint256 counter,
-        bytes32 upperBound
-    ) external pure returns (bytes32 generatedHandle);
-    function getOpResultHandle(
-        EOps op,
-        ETypes returnType,
-        bytes32 value
-    ) external pure returns (bytes32 generatedHandle);
-    function getOpResultHandle(
-        EOps op,
-        ETypes returnType,
-        bytes32 inputA,
-        bytes32 inputB,
-        bytes32 inputC
-    ) external pure returns (bytes32 generatedHandle);
+    function getTrivialEncryptHandle(bytes32 plaintextBytes, ETypes handleType) external view returns (bytes32 generatedHandle);
+
+    function getOpResultHandle(EOps op, ETypes returnType, bytes memory packedInputs) external pure returns (bytes32 generatedHandle);
 }

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

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: No License
 pragma solidity ^0.8;
 
-import {TestUtils} from "@inco/shared/src/TestUtils.sol";
+import {TestUtils} from "../../shared/TestUtils.sol";
 import {HandleMetadata} from "../primitives/HandleMetadata.sol";
 import {TrivialEncryption} from "../TrivialEncryption.sol";
 import {EncryptedOperations} from "../EncryptedOperations.sol";
@@ -108,11 +108,48 @@ contract TestHandleMetadata is
     }
 
     function testEncryptedInputHandleType() public {
-        euint256 a = this.newEuint256{value: FEE}("ciphertext", address(this));
+        address self = address(this);
+        bytes32 ciphertext = keccak256(abi.encodePacked("ciphertext"));
+        euint256 a = this.newEuint256{value: FEE}(
+            getCiphertextInput(ciphertext, self, self, ETypes.Uint256),
+            self
+        );
         assert(typeOf(euint256.unwrap(a)) == ETypes.Uint256);
-        ebool b = this.newEbool{value: FEE}("ciphertext", address(this));
+        ebool b = this.newEbool{value: FEE}(
+            getCiphertextInput(ciphertext, self, self, ETypes.Bool),
+            address(this)
+        );
         assert(typeOf(ebool.unwrap(b)) == ETypes.Bool);
-        eaddress c = this.newEaddress{value: FEE}("ciphertext", address(this));
+        eaddress c = this.newEaddress{value: FEE}(
+            getCiphertextInput(
+                ciphertext,
+                self,
+                self,
+                ETypes.AddressOrUint160OrBytes20
+            ),
+            address(this)
+        );
         assert(typeOf(eaddress.unwrap(c)) == ETypes.AddressOrUint160OrBytes20);
     }
+
+    /// @notice Helper to create an input with a handle prepended to the ciphertext.
+    /// @param word A single word to be used as the ciphertext.
+    /// @param user The user address associated with the input.
+    function getCiphertextInput(
+        bytes32 word,
+        address user,
+        address contractAddress,
+        ETypes inputType
+    ) public view returns (bytes memory input) {
+        // We need a single word here to get correct encoding
+        bytes memory ciphertext = abi.encode(word);
+        bytes32 handle = getInputHandle(
+            ciphertext,
+            address(this),
+            user,
+            contractAddress,
+            inputType
+        );
+        input = abi.encode(handle, ciphertext);
+    }
 }

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

@@ -6,59 +6,52 @@ import {EncryptedInput} from "../EncryptedInput.sol";
 import {VerifierAddressGetter} from "../primitives/VerifierAddressGetter.sol";
 import {FEE, Fee} from "../Fee.sol";
 import {IncoTest} from "../../test/IncoTest.sol";
-
-contract InputsTester is EncryptedInput {
-    constructor() VerifierAddressGetter(address(0)) {}
-}
+import {console} from "forge-std/console.sol";
+import {inco} from "../../Lib.sol";
 
 contract TestInputsFee is IncoTest {
-    InputsTester inputsTester;
-
-    function setUp() public override {
-        inputsTester = new InputsTester();
-    }
 
     function testPayOnInputs() public {
         // should fail if no fee
         vm.expectRevert(Fee.FeeNotPaid.selector);
-        inputsTester.newEuint256{value: 0}(
-            fakePrepareEuint256Ciphertext(12),
+        inco.newEuint256{value: 0}(
+            fakePrepareEuint256Ciphertext(12, address(0), address(this)),
             address(0)
         );
 
         // should fail if not enough fee
         vm.expectRevert(Fee.FeeNotPaid.selector);
-        inputsTester.newEuint256{value: FEE - 1}(
-            fakePrepareEuint256Ciphertext(12),
+        inco.newEuint256{value: FEE - 1}(
+            fakePrepareEuint256Ciphertext(12, address(0), address(this)),
             address(0)
         );
 
         // should fail if too much fee
         vm.expectRevert(Fee.FeeNotPaid.selector);
-        inputsTester.newEuint256{value: FEE + 1}(
-            fakePrepareEuint256Ciphertext(12),
+        inco.newEuint256{value: FEE + 1}(
+            fakePrepareEuint256Ciphertext(12, address(0), address(this)),
             address(0)
         );
 
         // should work with exact fee
-        inputsTester.newEuint256{value: FEE}(
-            fakePrepareEuint256Ciphertext(12),
+        inco.newEuint256{value: FEE}(
+            fakePrepareEuint256Ciphertext(12, address(0), address(this)),
             address(0)
         );
     }
 
     function testPayForNewEbool() public {
         // should work with exact fee
-        inputsTester.newEbool{value: FEE}(
-            fakePrepareEboolCiphertext(true),
+        inco.newEbool{value: FEE}(
+            fakePrepareEboolCiphertext(true, address(0), address(this)),
             address(0)
         );
     }
 
     function testPayForNewEaddress() public {
         // should work with exact fee
-        inputsTester.newEaddress{value: FEE}(
-            fakePrepareEaddressCiphertext(address(this)),
+        inco.newEaddress{value: FEE}(
+            fakePrepareEaddressCiphertext(address(0), address(0), address(this)),
             address(0)
         );
     }