@inco/lightning 0.3.2 → 0.5.0

package/src/test/FakeIncoInfra/FakeDecryptionAttester.sol

@@ -0,0 +1,198 @@
+// SPDX-License-Identifier: No License
+pragma solidity ^0.8;
+
+import {FakeIncoInfraBase} from "./FakeIncoInfraBase.sol";
+import {inco} from "../../Lib.sol";
+import {DecryptionAttestation} from "../../lightning-parts/DecryptionAttester.types.sol";
+import {AllowanceProof} from "../../lightning-parts/AccessControl/AdvancedAccessControl.sol";
+import {SenderNotAllowedForHandle, EOps, ETypes} from "../../Types.sol";
+import {FakeComputeServer} from "./FakeComputeServer.sol";
+
+/// @notice simulates the TEE / covalidator offchain API for decryption attestations requests
+/// The signing and checking of the decryption request is out of scope of this simulation, it is assumed that when
+/// called in tests, it emulates the act of a requester emitting a decryption request
+/// One can request decryption attestations on one handle, or on a handle not existing onchain resulting from an
+/// operation on existing handles, or on a mix of existing handles and plaintexts values (which get trivial encrypted)
+/// as long as the handles are allowed for the requester
+contract FakeDecryptionAttester is FakeIncoInfraBase, FakeComputeServer {
+    struct HandleWithProof {
+        bytes32 handle;
+        AllowanceProof proof; // sharer == address(0) means no proof
+    }
+
+    // The following 4 functions mimic expected behavior of the TEE / covalidator over offchain API
+
+    /// @notice request decryption attestation for an existing handle
+    function getDecryptionAttestation(
+        address requester,
+        HandleWithProof memory handle
+    )
+        internal
+        returns (
+            DecryptionAttestation memory decryption,
+            bytes memory signature
+        )
+    {
+        checkAccessControl(requester, handle);
+        (decryption, signature) = _getDecryptionAttestation(
+            handle.handle,
+            get(handle.handle)
+        );
+    }
+
+    // The following 3 functions are for requesting decryption over a handle resulting of an operation
+    // only binary ops are supported for now (2 params)
+
+    /// @notice request decryption attestation for a binary operation on 2 existing handles
+    function getDecryptionAttestation(
+        address requester,
+        HandleWithProof memory lhs,
+        HandleWithProof memory rhs,
+        EOps op
+    )
+        internal
+        returns (
+            DecryptionAttestation memory decryption,
+            bytes memory signature
+        )
+    {
+        checkAccessControl(requester, lhs);
+        checkAccessControl(requester, rhs);
+        bytes32 result = computeBinaryOp(get(lhs.handle), get(rhs.handle), op);
+        (decryption, signature) = _getDecryptionAttestation(
+            lhs.handle,
+            rhs.handle,
+            result,
+            op
+        );
+    }
+
+    /// @notice request decryption attestation for a binary operation on a handle and a plaintext value (lhs)
+    function getDecryptionAttestation(
+        address requester,
+        uint256 lhs,
+        HandleWithProof memory rhs,
+        EOps op
+    )
+        internal
+        returns (
+            DecryptionAttestation memory decryption,
+            bytes memory signature
+        )
+    {
+        checkAccessControl(requester, rhs);
+        bytes32 lhsHandle = inco.getTrivialEncryptHandle(
+            bytes32(lhs),
+            ETypes.Uint256
+        );
+        bytes32 result = computeBinaryOp(bytes32(lhs), get(rhs.handle), op);
+        (decryption, signature) = _getDecryptionAttestation(
+            lhsHandle,
+            rhs.handle,
+            result,
+            op
+        );
+    }
+
+    /// @notice request decryption attestation for a binary operation on a handle and a plaintext value (rhs)
+    function getDecryptionAttestation(
+        address requester,
+        HandleWithProof memory lhs,
+        uint256 rhs,
+        EOps op
+    )
+        internal
+        returns (
+            DecryptionAttestation memory decryption,
+            bytes memory signature
+        )
+    {
+        checkAccessControl(requester, lhs);
+        bytes32 rhsHandle = inco.getTrivialEncryptHandle(
+            bytes32(rhs),
+            ETypes.Uint256
+        );
+
+        bytes32 result = computeBinaryOp(get(lhs.handle), bytes32(rhs), op);
+        (decryption, signature) = _getDecryptionAttestation(
+            lhs.handle,
+            rhsHandle,
+            result,
+            op
+        );
+    }
+
+    /// Private methods ///
+
+    function checkAccessControl(
+        address requester,
+        HandleWithProof memory handle
+    ) private {
+        if (handle.proof.sharer == address(0)) {
+            require(
+                inco.isAllowed(handle.handle, requester),
+                SenderNotAllowedForHandle(handle.handle, requester)
+            );
+        } else {
+            require(
+                inco.incoVerifier().isAllowedWithProof(
+                    handle.handle,
+                    requester,
+                    handle.proof
+                ),
+                SenderNotAllowedForHandle(handle.handle, requester)
+            );
+        }
+    }
+
+    function _getDecryptionAttestation(
+        bytes32 handle,
+        bytes32 value
+    )
+        private
+        view
+        returns (
+            DecryptionAttestation memory decryption,
+            bytes memory signature
+        )
+    {
+        decryption = DecryptionAttestation({handle: handle, value: value});
+        signature = signDecryption(decryption);
+    }
+
+    function signDecryption(
+        DecryptionAttestation memory decryption
+    ) private view returns (bytes memory signature) {
+        // todo change this to don't call inco and call verifier directly
+        bytes32 digest = inco.incoVerifier().decryptionAttestationDigest(
+            decryption
+        );
+        signature = getSignatureForDigest(digest, teePrivKey);
+    }
+
+    function _getDecryptionAttestation(
+        bytes32 lhsHandle,
+        bytes32 rhsHandle,
+        bytes32 encodedResult,
+        EOps op
+    )
+        private
+        view
+        returns (
+            DecryptionAttestation memory decryption,
+            bytes memory signature
+        )
+    {
+        ETypes resultType = opToResultType(op);
+        decryption = DecryptionAttestation({
+            handle: inco.getOpResultHandle(
+                op,
+                resultType,
+                lhsHandle,
+                rhsHandle
+            ),
+            value: encodedResult
+        });
+        signature = signDecryption(decryption);
+    }
+}

package/src/test/FakeIncoInfra/FakeIncoInfraBase.sol

@@ -3,9 +3,7 @@ pragma solidity ^0.8;
 
 import {ebool, euint256, ETypes} from "../../Types.sol";
 import {inco} from "../../Lib.sol";
-import {Vm} from "forge-std/Test.sol";
 import {TestUtils} from "@inco/shared/src/TestUtils.sol";
-import {DecryptionResult} from "../../lightning-parts/DecryptionHandler.sol";
 import {KVStore} from "./KVStore.sol";
 
 /// @notice simulates what inco does offchain but over plaintexts
@@ -19,19 +17,6 @@ contract FakeIncoInfraBase is TestUtils, KVStore {
         (teePrivKey, teePubkeyAddress) = getLabeledKeyPair("tee");
     }
 
-    function handleDecryptionRequest(
-        uint256 requestId,
-        bytes32 handle
-    ) internal {
-        DecryptionResult memory result = DecryptionResult({
-            abiEncodedResult: get(handle),
-            requestId: requestId,
-            handle: handle
-        });
-        vm.prank(teePubkeyAddress);
-        inco.fulfillRequest(result, "");
-    }
-
     function fakePrepareEuint256Ciphertext(
         uint256 value
     ) internal pure returns (bytes memory ciphertext) {

package/src/test/FakeIncoInfra/FakeQuoteVerifier.sol

@@ -1,9 +1,9 @@
 // SPDX-License-Identifier: No License
 pragma solidity ^0.8;
 
-import {IPCCSRouter} from "@automata-network/dcap-attestation/interfaces/IPCCSRouter.sol";
-import {Header} from "@automata-network/dcap-attestation/types/CommonStruct.sol";
-import {IQuoteVerifier} from "automata-dcap-attestation/interfaces/IQuoteVerifier.sol";
+import {IPCCSRouter} from "../../interfaces/automata-interfaces/IPCCSRouter.sol";
+import {Header} from "../../interfaces/automata-interfaces/Types.sol";
+import {IQuoteVerifier} from "../../interfaces/automata-interfaces/IQuoteVerifier.sol";
 
 // This contract is used to test the IncoLightning contract. It is a simple implementation of the QuoteVerifier interface.
 // It is used to test the IncoLightning contract without relying on the real QuoteVerifier contract.
@@ -18,12 +18,16 @@ contract FakeQuoteVerifier is IQuoteVerifier {
         return 4;
     }
 
-    function verifyQuote(Header calldata, bytes calldata quote) external pure returns (bool, bytes memory) {
+    function verifyQuote(
+        Header calldata,
+        bytes calldata quote
+    ) external pure returns (bool, bytes memory) {
         return (true, quote);
     }
 
-    function verifyZkOutput(bytes calldata quote) external pure returns (bool, bytes memory) {
+    function verifyZkOutput(
+        bytes calldata quote
+    ) external pure returns (bool, bytes memory) {
         return (true, quote);
     }
-
 }

package/src/test/FakeIncoInfra/MockOpHandler.sol

@@ -4,7 +4,6 @@ pragma solidity ^0.8;
 import {Vm} from "forge-std/Test.sol";
 import {inco} from "../../Lib.sol";
 import {ebool, euint256, ETypes, EOps} from "../../Types.sol";
-import {DecryptionRequested} from "../../lightning-parts/DecryptionHandler.sol";
 import {FakeComputeServer} from "./FakeComputeServer.sol";
 import {FakeIncoInfraBase} from "./FakeIncoInfraBase.sol";
 import {asBytes32} from "@inco/shared/src/TypeUtils.sol";
@@ -71,12 +70,6 @@ contract MockOpHandler is FakeIncoInfraBase, FakeComputeServer {
                 (ETypes, bytes, uint256)
             );
             handleEInput(result, inputType, ciphertext);
-        } else if (op == EOps.DecryptionRequested) {
-            uint256 requestId = uint256(log.topics[1]);
-            bytes32 handle = log.topics[2];
-            // other fields ignored for now
-            // todo cheats to trigger decryption later
-            handleDecryptionRequest(requestId, handle);
         }
     }
 

package/src/test/FakeIncoInfra/MockRemoteAttestation.sol

@@ -2,10 +2,11 @@
 pragma solidity ^0.8;
 
 import {TestUtils} from "@inco/shared/src/TestUtils.sol";
-import {TEELifecycle} from "@inco/lightning/src/lightning-parts/TEELifecycle.sol";
-
-import {HEADER_LENGTH, MINIMUM_QUOTE_LENGTH, TDX_TEE} from "@automata-network/dcap-attestation/types/Constants.sol";
-import {BootstrapResult} from "@inco/lightning/src/lightning-parts/TEELifecycle.types.sol";
+import {
+    HEADER_LENGTH,
+    MINIMUM_QUOTE_LENGTH,
+    TDX_TEE
+} from "../../interfaces/automata-interfaces/Types.sol";
 
 contract MockRemoteAttestation is TestUtils {
     function createQuote(
@@ -31,7 +32,18 @@ contract MockRemoteAttestation is TestUtils {
         bytes memory prefix = new bytes(HEADER_LENGTH + 136 - 8);
         bytes memory middle = new bytes(520 - 184);
         bytes memory reportDataSuffix = new bytes(44);
-        bytes memory suffix = new bytes(MINIMUM_QUOTE_LENGTH - HEADER_LENGTH - 584);
-        quote = abi.encodePacked(version, tdxTEEType, prefix, mrtd, middle, abi.encodePacked(signer), reportDataSuffix, suffix);
+        bytes memory suffix = new bytes(
+            MINIMUM_QUOTE_LENGTH - HEADER_LENGTH - 584
+        );
+        quote = abi.encodePacked(
+            version,
+            tdxTEEType,
+            prefix,
+            mrtd,
+            middle,
+            abi.encodePacked(signer),
+            reportDataSuffix,
+            suffix
+        );
     }
-}
+}

package/src/test/FakeIncoInfra/getOpForSelector.sol

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: No License
 pragma solidity ^0.8;
 
-import {DecryptionRequested} from "../../lightning-parts/DecryptionHandler.sol";
 import {EncryptedOperations} from "../../lightning-parts/EncryptedOperations.sol";
 import {TrivialEncryption} from "../../lightning-parts/TrivialEncryption.sol";
 import {EncryptedInput} from "../../lightning-parts/EncryptedInput.sol";
@@ -63,8 +62,6 @@ function getOpForSelector(bytes32 opEventSelector) pure returns (EOps) {
         return EOps.Rand;
     } else if (opEventSelector == EncryptedOperations.ERandBounded.selector) {
         return EOps.RandBounded;
-    } else if (opEventSelector == DecryptionRequested.selector) {
-        return EOps.DecryptionRequested;
     } else {
         revert("getOpForSelector: Unsupported selector");
     }

package/src/test/IncoTest.sol

@@ -2,14 +2,16 @@
 pragma solidity ^0.8;
 
 import {MockOpHandler} from "./FakeIncoInfra/MockOpHandler.sol";
-import {IncoLightning} from "../IncoLightning.sol";
+import {IIncoLightning} from "../interfaces/IIncoLightning.sol";
 import {inco} from "../Lib.sol";
 import {DeployUtils} from "../DeployUtils.sol";
 import {deployedBy} from "../Lib.sol";
+import {FakeDecryptionAttester} from "./FakeIncoInfra/FakeDecryptionAttester.sol";
 import {console} from "forge-std/console.sol";
 import {FakeQuoteVerifier} from "./FakeIncoInfra/FakeQuoteVerifier.sol";
+import {IOwnable} from "@inco/shared/src/IOwnable.sol";
 
-contract IncoTest is MockOpHandler, DeployUtils {
+contract IncoTest is MockOpHandler, DeployUtils, FakeDecryptionAttester {
     address immutable owner;
     address immutable testDeployer;
 
@@ -27,16 +29,14 @@ contract IncoTest is MockOpHandler, DeployUtils {
         deployCreateX();
         vm.startPrank(testDeployer);
         vm.setEnv("USE_TDX_HW", "false"); // results in the test deployment using the FakeQuoteVerifier
-        IncoLightning proxy = deployIncoLightningUsingConfig({
+        (IIncoLightning proxy, ) = deployIncoLightningUsingConfig({
             deployer: testDeployer,
-        // The highest precedent deployment
+            // The highest precedent deployment
             pepper: "testnet",
-            minorVersionForSalt: 1,
-            patchVersionForSalt: 29,
-            includePreviewFeatures: false,
-            teeLifecycleAddress: address(0)
+            quoteVerifier: new FakeQuoteVerifier()
         });
-        proxy.transferOwnership(owner);
+        IOwnable(address(proxy)).transferOwnership(owner);
+        IOwnable(address(inco.incoVerifier())).transferOwnership(owner);
         vm.stopPrank();
         console.log(
             "Deployed %s (proxy) to: %s",
@@ -48,8 +48,9 @@ contract IncoTest is MockOpHandler, DeployUtils {
             address(proxy) == address(inco),
             "generated inco address in Lib.sol does not match address of inco deployed by IncoTest"
         );
-        vm.prank(owner);
-        inco.addSigner(teePubkeyAddress);
+        vm.startPrank(owner);
+        inco.incoVerifier().addSigner(teePubkeyAddress);
+        vm.stopPrank();
         vm.recordLogs();
     }
 }

package/src/test/TEELifecycle/TEELifecycleMockTest.t.sol

@@ -1,145 +1,166 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.0;
 
-import "@inco/lightning/src/lightning-parts/TEELifecycle.sol";
-import "@inco/lightning/src/lightning-parts/TEELifecycle.types.sol";
+import {TEELifecycle} from "../../lightning-parts/TEELifecycle.sol";
+import {BootstrapResult} from "../../lightning-parts/TEELifecycle.types.sol";
 import {MockRemoteAttestation} from "../FakeIncoInfra/MockRemoteAttestation.sol";
 import {FakeQuoteVerifier} from "../FakeIncoInfra/FakeQuoteVerifier.sol";
-
-import "forge-std/Vm.sol";
-import "forge-std/Test.sol";
+import {Test} from "forge-std/Test.sol";
 
 contract TEELifecycleMockTest is Test, MockRemoteAttestation, TEELifecycle {
-
     function setUp() public {
         quoteVerifier = new FakeQuoteVerifier();
     }
 
-    function test_successfulBootstrap() public {
-        (BootstrapResult memory bootstrapResult, , , bytes memory quote, bytes memory signature, bytes memory mrtd) = successfulBootstrapResult();
+    function testSuccessfulBootstrap() public {
+        (
+            BootstrapResult memory bootstrapResult,
+            ,
+            ,
+            bytes memory quote,
+            bytes memory signature,
+            bytes memory mrtd
+        ) = successfulBootstrapResult();
         vm.startPrank(this.owner());
         this.approveNewTEEVersion(mrtd);
-        this.verifyBootstrapResult(
-            bootstrapResult,
-            quote,
-            signature
-        );
+        this.verifyBootstrapResult(bootstrapResult, quote, signature);
         assertTrue(this.isBootstrapComplete(), "Bootstrap should be complete");
         vm.stopPrank();
     }
-    
-    function test_invalidMrtd() public {
-        bytes memory badMrtd = hex"1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef";
 
-        (BootstrapResult memory bootstrapResult, , address bootstrapPartyAddress, bytes memory quote, bytes memory signature, bytes memory mrtd) = successfulBootstrapResult();
+    function testInvalidMrtd() public {
+        bytes
+            memory badMrtd = hex"1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef";
+
+        (
+            BootstrapResult memory bootstrapResult,
+            ,
+            address bootstrapPartyAddress,
+            bytes memory quote,
+            bytes memory signature,
+            bytes memory mrtd
+        ) = successfulBootstrapResult();
 
         quote = createQuote(badMrtd, bootstrapPartyAddress); // Replace with bad MRTD
         vm.startPrank(this.owner());
         this.approveNewTEEVersion(mrtd);
-        vm.expectRevert(bytes("Invalid report MRTD"));
-        this.verifyBootstrapResult(
-            bootstrapResult,
-            quote,
-            signature
-        );
+        vm.expectRevert(TEELifecycle.InvalidReportMRTD.selector);
+        this.verifyBootstrapResult(bootstrapResult, quote, signature);
         vm.stopPrank();
     }
 
-    function test_invalidSignature() public {
-        (BootstrapResult memory bootstrapResult, , , bytes memory quote, , bytes memory mrtd) = successfulBootstrapResult();
-        (uint256 bootstrapPartyFakePrivkey, ) = getLabeledKeyPair("bootstrapPartyFake");
-        bytes memory signatureInvalid = signBootstrapResult(bootstrapResult, bootstrapPartyFakePrivkey);
-        vm.startPrank(this.owner());
-        this.approveNewTEEVersion(mrtd);
-        vm.expectRevert(bytes("Invalid signature for bootstrap data"));
-        this.verifyBootstrapResult(
+    function testInvalidSignature() public {
+        (
+            BootstrapResult memory bootstrapResult,
+            ,
+            ,
+            bytes memory quote,
+            ,
+            bytes memory mrtd
+        ) = successfulBootstrapResult();
+        (uint256 bootstrapPartyFakePrivkey, ) = getLabeledKeyPair(
+            "bootstrapPartyFake"
+        );
+        bytes memory signatureInvalid = signBootstrapResult(
             bootstrapResult,
-            quote,
-            signatureInvalid
+            bootstrapPartyFakePrivkey
         );
+        vm.startPrank(this.owner());
+        this.approveNewTEEVersion(mrtd);
+        vm.expectRevert(TEELifecycle.InvalidBootstrapDataSignature.selector);
+        this.verifyBootstrapResult(bootstrapResult, quote, signatureInvalid);
         vm.stopPrank();
     }
 
-    function test_bootstrapAlreadyComplete() public {
-        (BootstrapResult memory bootstrapResult, , , bytes memory quote, bytes memory signature, bytes memory mrtd) = successfulBootstrapResult();
+    function testBootstrapAlreadyComplete() public {
+        (
+            BootstrapResult memory bootstrapResult,
+            ,
+            ,
+            bytes memory quote,
+            bytes memory signature,
+            bytes memory mrtd
+        ) = successfulBootstrapResult();
         vm.startPrank(this.owner());
         this.approveNewTEEVersion(mrtd);
-        this.verifyBootstrapResult(
-            bootstrapResult,
-            quote,
-            signature
-        );
-        vm.expectRevert(bytes("Bootstrap already completed"));
-        this.verifyBootstrapResult(
-            bootstrapResult,
-            quote,
-            signature
-        );
+        this.verifyBootstrapResult(bootstrapResult, quote, signature);
+        vm.expectRevert(TEELifecycle.BootstrapAlreadyCompleted.selector);
+        this.verifyBootstrapResult(bootstrapResult, quote, signature);
         vm.stopPrank();
     }
 
-    function test_approveNewTEEInvalidMrtd() public {
+    function testApproveNewTEEInvalidMrtd() public {
         bytes memory mrtd = hex"deadbeef";
         vm.startPrank(this.owner());
-        vm.expectRevert(bytes("MRTD must be 48 bytes"));
+        vm.expectRevert(TEELifecycle.MrtdInvalidLength.selector);
         this.approveNewTEEVersion(mrtd);
         vm.stopPrank();
     }
 
-    function test_bootstrapNotCompleteNewCoval() public {
-        bytes memory mrtd = hex"2a90c8fa38672cafd791d994beb6836b99383b2563736858632284f0f760a6446efd1e7ec457cf08b629ea630f7b4525";
+    function testBootstrapNotCompleteNewCoval() public {
+        bytes
+            memory mrtd = hex"2a90c8fa38672cafd791d994beb6836b99383b2563736858632284f0f760a6446efd1e7ec457cf08b629ea630f7b4525";
         (, address newCoval) = getLabeledKeyPair("newCoval");
         bytes memory quote = createQuote(mrtd, newCoval);
         vm.startPrank(this.owner());
-        vm.expectRevert(bytes("Bootstrap not complete"));
+        vm.expectRevert(TEELifecycle.BootstrapNotComplete.selector);
         this.addNewCovalidator(quote);
         vm.stopPrank();
     }
 
-    function test_invalidMrtdNewCoval() public {
-        (BootstrapResult memory bootstrapResult, , , bytes memory quote, bytes memory signature, bytes memory mrtd) = successfulBootstrapResult();
+    function testInvalidMrtdNewCoval() public {
+        (
+            BootstrapResult memory bootstrapResult,
+            ,
+            ,
+            bytes memory quote,
+            bytes memory signature,
+            bytes memory mrtd
+        ) = successfulBootstrapResult();
         vm.startPrank(this.owner());
         this.approveNewTEEVersion(mrtd);
-        this.verifyBootstrapResult(
-            bootstrapResult,
-            quote,
-            signature
-        );
-        bytes memory badMrtd = hex"1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef";
+        this.verifyBootstrapResult(bootstrapResult, quote, signature);
+        bytes
+            memory badMrtd = hex"1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef";
         (, address newCoval) = getLabeledKeyPair("newCoval");
         bytes memory quoteNew = createQuote(badMrtd, newCoval);
-        
-        vm.expectRevert(bytes("Invalid report MRTD"));
+
+        vm.expectRevert(TEELifecycle.InvalidMrtdReport.selector);
         this.addNewCovalidator(quoteNew);
         vm.stopPrank();
     }
 
     // Helper function to create a successful bootstrap result
-    function successfulBootstrapResult() internal returns (BootstrapResult memory bootstrapResult, uint256 bootstrapPartyPrivkey, address bootstrapPartyAddress, bytes memory quote, bytes memory signature, bytes memory mrtd) {
-        (bootstrapPartyPrivkey, bootstrapPartyAddress) = getLabeledKeyPair("bootstrapParty");
-        bytes memory eciesPubkey = hex"04ff5c6dd72ad7583288b84ee2598e081fe0bc6ef543c342e925a5dfcff9afb2444d25454d7d5dcfadc9ed99477c245efa93caf58d7f58143300d81cc948e7bdf5";
+    function successfulBootstrapResult()
+        internal
+        returns (
+            BootstrapResult memory bootstrapResult,
+            uint256 bootstrapPartyPrivkey,
+            address bootstrapPartyAddress,
+            bytes memory quote,
+            bytes memory signature,
+            bytes memory mrtd
+        )
+    {
+        (bootstrapPartyPrivkey, bootstrapPartyAddress) = getLabeledKeyPair(
+            "bootstrapParty"
+        );
+        bytes
+            memory eciesPubkey = hex"04ff5c6dd72ad7583288b84ee2598e081fe0bc6ef543c342e925a5dfcff9afb2444d25454d7d5dcfadc9ed99477c245efa93caf58d7f58143300d81cc948e7bdf5";
         mrtd = hex"2a90c8fa38672cafd791d994beb6836b99383b2563736858632284f0f760a6446efd1e7ec457cf08b629ea630f7b4525";
 
-        bootstrapResult = BootstrapResult({
-            ecies_pubkey: eciesPubkey
-        });
+        bootstrapResult = BootstrapResult({ecies_pubkey: eciesPubkey});
 
-        quote = createQuote(
-            mrtd,
-            bootstrapPartyAddress
-        );
+        quote = createQuote(mrtd, bootstrapPartyAddress);
         signature = signBootstrapResult(bootstrapResult, bootstrapPartyPrivkey);
     }
-    
+
     // Helper function to sign the bootstrap result
     function signBootstrapResult(
         BootstrapResult memory bootstrapResult,
         uint256 privateKey
     ) internal view returns (bytes memory) {
-        bytes32 bootstrapResultDigest = bootstrapResultDigest(
-            bootstrapResult
-        );
+        bytes32 bootstrapResultDigest = bootstrapResultDigest(bootstrapResult);
         return getSignatureForDigest(bootstrapResultDigest, privateKey);
     }
-
-}
+}

package/src/test/TestAddTwo.t.sol

@@ -7,9 +7,15 @@ import {IncoTest} from "./IncoTest.sol";
 import {AddTwo} from "./AddTwo.sol";
 
 contract TestAddTwo is IncoTest {
-    function testAddTwo() public {
-        AddTwo addTwo = new AddTwo(inco);
+    AddTwo addTwo;
+
+    function setUp() public override {
+        super.setUp();
+        addTwo = new AddTwo(inco);
         vm.label(address(addTwo), "addTwo");
+    }
+
+    function testAddTwo() public {
         euint256 a = inco.asEuint256(3);
         inco.allow(euint256.unwrap(a), address(addTwo));
         euint256 b = addTwo.addTwo(a);
@@ -17,13 +23,14 @@ contract TestAddTwo is IncoTest {
         assertEq(getUint256Value(b), 5);
     }
 
-    function testAddTwoScalar() public {
-        AddTwo addTwo = new AddTwo(inco);
-        vm.label(address(addTwo), "addTwoScalar");
-        euint256 a = inco.asEuint256(3);
-        inco.allow(euint256.unwrap(a), address(addTwo));
-        euint256 b = addTwo.addTwoScalar(a);
+    function testAddTwoEoaAndPublicReveal() public {
+        (euint256 result, euint256 revealedResult) = addTwo.addTwoEOA(
+            fakePrepareEuint256Ciphertext(3)
+        );
         processAllOperations();
-        assertEq(getUint256Value(b), 5);
+        assertEq(getUint256Value(result), 5);
+        assertEq(getUint256Value(revealedResult), 5);
+        assertFalse(inco.isAllowed(euint256.unwrap(result), bob));
+        assertTrue(inco.isAllowed(euint256.unwrap(revealedResult), bob));
     }
 }