@inco/lightning 0.7.10 → 0.8.0-devnet-1

package/src/libs/incoLightning_testnet_v2_889158349.sol

@@ -16,6 +16,11 @@ function typeOf(bytes32 handle) pure returns (ETypes) {
 }
 
 library e {
+    error CallFailedAfterFeeRefresh();
+
+    /// @dev slot to store the fee for inco operations
+    bytes32 private constant FEE_SLOT = keccak256("inco.fee");
+
     function sanitize(euint256 a) internal returns (euint256) {
         if (euint256.unwrap(a) == bytes32(0)) {
             return asEuint256(0);
@@ -355,17 +360,22 @@ library e {
 
     /// @dev costs the inco fee
     function rand() internal returns (euint256) {
-        return euint256.wrap(inco.eRand{value: inco.getFee()}(ETypes.Uint256));
+        bytes32 result = _callWithFeeRetry(abi.encodeWithSelector(inco.eRand.selector, ETypes.Uint256));
+        return euint256.wrap(result);
     }
 
     /// @dev costs the inco fee
     function randBounded(uint256 upperBound) internal returns (euint256) {
-        return euint256.wrap(inco.eRandBounded{value: inco.getFee()}(euint256.unwrap(asEuint256(upperBound)), ETypes.Uint256));
+        bytes32 boundHandle = euint256.unwrap(asEuint256(upperBound));
+        bytes32 result = _callWithFeeRetry(abi.encodeWithSelector(inco.eRandBounded.selector, boundHandle, ETypes.Uint256));
+        return euint256.wrap(result);
     }
 
     /// @dev costs the inco fee
     function randBounded(euint256 upperBound) internal returns (euint256) {
-        return euint256.wrap(inco.eRandBounded{value: inco.getFee()}(euint256.unwrap(s(upperBound)), ETypes.Uint256));
+        bytes32 boundHandle = euint256.unwrap(s(upperBound));
+        bytes32 result = _callWithFeeRetry(abi.encodeWithSelector(inco.eRandBounded.selector, boundHandle, ETypes.Uint256));
+        return euint256.wrap(result);
     }
 
     function asEuint256(uint256 a) internal returns (euint256) {
@@ -397,7 +407,8 @@ library e {
     /// @notice Creates a new encrypted uint256 for the given user.
     ///  @dev costs the inco fee
     function newEuint256(bytes memory ciphertext, address user) internal returns (euint256) {
-        return inco.newEuint256{value: inco.getFee()}(ciphertext, user);
+        bytes32 result = _callWithFeeRetry(abi.encodeWithSelector(inco.newEuint256.selector, ciphertext, user));
+        return euint256.wrap(result);
     }
 
     /// @notice Creates a new encrypted bool assuming msg.sender is the user
@@ -409,7 +420,8 @@ library e {
     /// @notice Creates a new encrypted bool for the given user.
     ///  @dev costs the inco fee
     function newEbool(bytes memory ciphertext, address user) internal returns (ebool) {
-        return inco.newEbool{value: inco.getFee()}(ciphertext, user);
+        bytes32 result = _callWithFeeRetry(abi.encodeWithSelector(inco.newEbool.selector, ciphertext, user));
+        return ebool.wrap(result);
     }
 
     /// @notice Creates a new encrypted address assuming msg.sender is the user
@@ -421,7 +433,8 @@ library e {
     /// @notice Creates a new encrypted address for the given user.
     ///  @dev costs the inco fee
     function newEaddress(bytes memory ciphertext, address user) internal returns (eaddress) {
-        return inco.newEaddress{value: inco.getFee()}(ciphertext, user);
+        bytes32 result = _callWithFeeRetry(abi.encodeWithSelector(inco.newEaddress.selector, ciphertext, user));
+        return eaddress.wrap(result);
     }
 
     function allow(euint256 a, address to) internal {
@@ -475,4 +488,47 @@ library e {
     function select(ebool control, eaddress ifTrue, eaddress ifFalse) internal returns (eaddress) {
         return eaddress.wrap(inco.eIfThenElse(s(control), eaddress.unwrap(s(ifTrue)), eaddress.unwrap(s(ifFalse))));
     }
+
+    /// @dev Store fee in the custom slot
+    ///  @param _fee The fee to store
+    function _setFee(uint256 _fee) private {
+        bytes32 slot = FEE_SLOT;
+        assembly {
+            sstore(slot, _fee)
+        }
+    }
+
+    /// @dev Retrieve fee from the custom slot
+    ///  @return fee The stored fee
+    function _getFee() private view returns (uint256 fee) {
+        bytes32 slot = FEE_SLOT;
+        assembly {
+            fee := sload(slot)
+        }
+    }
+
+    /// @dev Get current fee with fallback to inco.getFee() if not cached
+    function getCurrentFee() private returns (uint256) {
+        uint256 cachedFee = _getFee();
+        if (cachedFee == 0) {
+            cachedFee = inco.getFee();
+            _setFee(cachedFee);
+        }
+        return cachedFee;
+    }
+
+    /// @dev Execute a call to inco with fee, retrying with fresh fee if it fails
+    ///  @param callData The encoded function call (use abi.encodeWithSelector)
+    ///  @return result The bytes32 result from the call
+    function _callWithFeeRetry(bytes memory callData) private returns (bytes32) {
+        uint256 fee = getCurrentFee();
+        (bool success, bytes memory result) = address(inco).call{value: fee}(callData);
+        if (!success) {
+            fee = inco.getFee();
+            _setFee(fee);
+            (success, result) = address(inco).call{value: fee}(callData);
+            require(success, CallFailedAfterFeeRefresh());
+        }
+        return abi.decode(result, (bytes32));
+    }
 }

package/src/lightning-parts/Fee.sol

@@ -2,13 +2,15 @@
 pragma solidity ^0.8;
 
 // the fee should be modified through upgrades to limit gas cost
-uint256 constant FEE = 0.0001 ether;
+uint256 constant FEE = 0.000001 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();
+    error FeeWithdrawalFailed();
+    error NoFeesToWithdraw();
 
     /// @notice the fee to pay through msg.value for inputs and randomness IT MAY CHANGE
     function getFee() public pure returns (uint256) {
@@ -25,19 +27,13 @@ abstract contract Fee {
         _;
     }
 
-    // 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");
-        }
+    /// @notice Internal helper to withdraw accumulated fees to a recipient
+    /// @param recipient The address to send the fees to
+    function _withdrawFeesTo(address recipient) internal {
+        uint256 fees = address(this).balance;
+        require(fees > 0, NoFeesToWithdraw());
+        (bool success,) = payable(recipient).call{value: fees}("");
+        require(success, FeeWithdrawalFailed());
     }
 
 }

package/src/lightning-parts/TEELifecycle.sol

@@ -29,8 +29,7 @@ contract TEELifecycleStorage {
         // @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;
+        bytes networkPubkey;
     }
 
     bytes32 private constant TEE_LIFECYCLE_STORAGE_LOCATION = keccak256("inco.storage.TEELifecycle");
@@ -80,9 +79,9 @@ abstract contract TEELifecycle is
     event SignerHasUpdatedTDX(address indexed eoaSigner, bytes32 indexed mrAggregated);
 
     // Constants
-    bytes32 public constant BOOTSTRAP_RESULT_STRUCT_HASH = keccak256(bytes("BootstrapResult(bytes ecies_pubkey)"));
-    bytes32 public constant UPGRADE_RESULT_STRUCT_HASH = keccak256(bytes("UpgradeResult(bytes network_pubkey)"));
-    bytes32 public constant ADD_NODE_RESULT_STRUCT_HASH = keccak256(bytes("AddNodeResult(bytes network_pubkey)"));
+    bytes32 public constant BOOTSTRAP_RESULT_STRUCT_HASH = keccak256(bytes("BootstrapResult(bytes networkPubkey)"));
+    bytes32 public constant UPGRADE_RESULT_STRUCT_HASH = keccak256(bytes("UpgradeResult(bytes networkPubkey)"));
+    bytes32 public constant ADD_NODE_RESULT_STRUCT_HASH = keccak256(bytes("AddNodeResult(bytes networkPubkey)"));
 
     uint16 public constant QUOTE_VERIFIER_VERSION = 4;
 
@@ -141,7 +140,7 @@ abstract contract TEELifecycle is
 
         // Update contract publicly viewable state
         addSigner(reportDataSigner);
-        $.eciesPubkey = bootstrapResult.ecies_pubkey;
+        $.networkPubkey = bootstrapResult.networkPubkey;
 
         emit BootstrapStageComplete(reportDataSigner, bootstrapResult);
     }
@@ -163,7 +162,7 @@ abstract contract TEELifecycle is
 
         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.networkPubkey), InvalidNetworkPubkey());
+        require(keccak256($.networkPubkey) == keccak256(upgradeResult.networkPubkey), InvalidNetworkPubkey());
         // Make sure the new MR_AGGREGATED has been pre-approved.
         bool isApproved = _isApprovedTeeVersion(newMrAggregated);
         require(isApproved, TEEVersionNotFound());
@@ -188,7 +187,7 @@ abstract contract TEELifecycle is
 
         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.networkPubkey), InvalidNetworkPubkey());
+        require(keccak256($.networkPubkey) == keccak256(addNodeResult.networkPubkey), InvalidNetworkPubkey());
         // Make sure the new MR_AGGREGATED has been pre-approved.
         bool isApproved = _isApprovedTeeVersion(newMrAggregated);
         require(isApproved, TEEVersionNotFound());
@@ -257,7 +256,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 getTeeLifecycleStorage().eciesPubkey.length > 0;
+        return getTeeLifecycleStorage().networkPubkey.length > 0;
     }
 
     /**
@@ -384,14 +383,14 @@ abstract contract TEELifecycle is
         return $.approvedTeeVersions[index];
     }
 
-    function eciesPubkey() external view returns (bytes memory) {
-        return getTeeLifecycleStorage().eciesPubkey;
+    function networkPubkey() external view returns (bytes memory) {
+        return getTeeLifecycleStorage().networkPubkey;
     }
 
     function bootstrapResultDigest(BootstrapResult memory bootstrapResult) public view returns (bytes32) {
         return
             _hashTypedDataV4(
-                keccak256(abi.encode(BOOTSTRAP_RESULT_STRUCT_HASH, keccak256(bootstrapResult.ecies_pubkey)))
+                keccak256(abi.encode(BOOTSTRAP_RESULT_STRUCT_HASH, keccak256(bootstrapResult.networkPubkey)))
             );
     }
 

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

@@ -5,10 +5,7 @@ pragma solidity ^0.8.19;
  * @notice A struct representing what the TDX EOA signs during the bootstrap process.
  */
 struct BootstrapResult {
-    // TODO: rename to networkPubkey
-    // https://github.com/Inco-fhevm/inco-monorepo/issues/983
-    // forge-lint: disable-next-line(mixed-case-variable)
-    bytes ecies_pubkey;
+    bytes networkPubkey;
 }
 
 /**

package/src/periphery/IncoUtils.sol

@@ -0,0 +1,42 @@
+// SPDX-License-Identifier: No License
+pragma solidity ^0.8;
+
+import {StorageSlot} from "@openzeppelin/contracts/utils/StorageSlot.sol";
+
+// Re-export FEE constant for convenience - consumers can import both IncoUtils and FEE from this file
+import {FEE} from "../lightning-parts/Fee.sol";
+
+contract IncoUtils {
+
+    error RefundFailed();
+    error ReentrantCall();
+
+    uint256 private constant NOT_ENTERED = 1;
+    uint256 private constant ENTERED = 2;
+
+    bytes32 private constant REENTRANCY_GUARD_SLOT = keccak256("inco.storage.IncoUtils.reentrancyGuard");
+
+    /// @notice Refund the difference between msg.value and what was actually spent
+    /// @dev Assumes all outflows and inflows to the contract are due to the user; refund is capped at msg.value
+    /// @dev Includes built-in reentrancy protection using OZ StorageSlot pattern (modifiers cannot use other modifiers)
+    modifier refundUnspent() {
+        StorageSlot.Uint256Slot storage status = StorageSlot.getUint256Slot(REENTRANCY_GUARD_SLOT);
+
+        require(status.value != ENTERED, ReentrantCall());
+        status.value = ENTERED;
+
+        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, RefundFailed());
+        }
+
+        status.value = NOT_ENTERED;
+    }
+
+}

package/src/test/AddTwo.sol

@@ -3,12 +3,12 @@ pragma solidity ^0.8;
 
 import {euint256, ebool} from "../Types.sol";
 import {IncoLightning} from "../IncoLightning.sol";
-import {Fee} from "../lightning-parts/Fee.sol";
+import {IncoUtils, FEE} from "../periphery/IncoUtils.sol";
 import {DecryptionAttestation} from "../lightning-parts/DecryptionAttester.types.sol";
 
 // To implement such a contract, we would normally import e form Lib.sol. For test purposes, we take inco as
 // a constructor argument instead, so we can test it from other deployment addresses.
-contract AddTwo is Fee {
+contract AddTwo is IncoUtils {
 
     IncoLightning inco;
 
@@ -36,7 +36,7 @@ contract AddTwo is Fee {
         refundUnspent
         returns (euint256 result, euint256 resultRevealed)
     {
-        euint256 value = inco.newEuint256{value: getFee()}(uint256EInput, msg.sender);
+        euint256 value = inco.newEuint256{value: FEE}(uint256EInput, msg.sender);
         result = addTwo(value);
 
         inco.allow(euint256.unwrap(result), address(this));

package/src/test/FakeIncoInfra/MockRemoteAttestation.sol

@@ -43,7 +43,7 @@ contract MockRemoteAttestation is TestUtils {
     // Helper function to create a successful bootstrap result
     function successfulBootstrapResult(
         ITEELifecycle teeLifecycle,
-        bytes memory eciesPublicKey,
+        bytes memory networkPubkey,
         address signer,
         uint256 signerPrivKey
     )
@@ -55,14 +55,13 @@ contract MockRemoteAttestation is TestUtils {
             bytes32 mrAggregated
         )
     {
-        bytes memory eciesPubkey = eciesPublicKey;
         // See DEFAULT_MRTD in attestation/src/remote_attestation.rs
         bytes memory mrtd =
             hex"010101010101010101010101010101010101010101010101010101010101010101010101010101010101010101010101";
         // See DEFAULT_MR_AGGREGATED in attestation/src/remote_attestation.rs to
         // see the calculation of the default value.
         mrAggregated = hex"c3a67bac251d4946d7b17481d39631676042fe3afab06e70c22105ad8383c19f";
-        bootstrapResult = BootstrapResult({ecies_pubkey: eciesPubkey});
+        bootstrapResult = BootstrapResult({networkPubkey: networkPubkey});
 
         quote = createQuote(mrtd, signer);
         signature = signBootstrapResult(bootstrapResult, signerPrivKey, teeLifecycle);

package/src/test/IncoTest.sol

@@ -39,12 +39,12 @@ contract IncoTest is MockOpHandler, DeployUtils, FakeDecryptionAttester, MockRem
     function setUp() public virtual {
         deployCreateX();
         vm.startPrank(testDeployer);
-        vm.setEnv("USE_TDX_HW", "false"); // results in the test deployment using the FakeQuoteVerifier
+        vm.setEnv("SHOULD_SETUP_TEE_SIGNER", "true"); // results in a network pubkey and decrypt signer being populated in the TEE Lifecycle
         (IIncoLightning proxy,) = deployIncoLightningUsingConfig({
             deployer: testDeployer,
             // The highest precedent deployment
             // We select the pepper that will be used that will be generated in the lib.sol (usually "testnet"), but currently "alphanet" has higher major version
-            pepper: "testnet",
+            pepper: "devnet",
             quoteVerifier: new FakeQuoteVerifier()
         });
         IOwnable(address(proxy)).transferOwnership(owner);

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

@@ -138,7 +138,7 @@ contract TEELifecycleMockTest is MockRemoteAttestation, TEELifecycle {
     {
         (bootstrapPartyPrivkey, bootstrapPartyAddress) = getLabeledKeyPair("bootstrapParty");
         mrAggregated = testMrAggregated;
-        bootstrapResult = BootstrapResult({ecies_pubkey: testNetworkPubkey});
+        bootstrapResult = BootstrapResult({networkPubkey: testNetworkPubkey});
 
         quote = createQuote(testMrtd, bootstrapPartyAddress);
         signature = signBootstrapResult(bootstrapResult, bootstrapPartyPrivkey);

package/src/test/TestFakeInfra.t.sol

@@ -14,6 +14,7 @@ contract TakesEInput is IncoTest {
 
     euint256 public a;
     ebool public b;
+    eaddress public c;
     uint256 public decryptedA;
 
     function setA(bytes memory uint256EInput) external {
@@ -25,6 +26,10 @@ contract TakesEInput is IncoTest {
         b = boolEInput.newEbool(msg.sender);
     }
 
+    function setC(bytes memory addressEInput) external {
+        c = addressEInput.newEaddress(msg.sender);
+    }
+
 }
 
 // its meta: this is testing correct behavior of our testing infrastructure
@@ -238,6 +243,264 @@ contract TestFakeInfra is IncoTest {
         assertEq(getUint256Value(a), 1);
     }
 
+    function testERandBoundedWithEuint256() public {
+        euint256 bound = e.asEuint256(100);
+        processAllOperations();
+        euint256 a = e.randBounded(bound);
+        processAllOperations();
+        assertEq(getUint256Value(a), 1);
+    }
+
+    // ============ FEE CACHING TESTS ============
+
+    // The fee slot used by the library
+    bytes32 constant FEE_SLOT = keccak256("inco.fee");
+
+    function testFeeCachingOnRand() public {
+        // Verify the slot is empty before any operation
+        bytes32 cachedFeeBefore = vm.load(address(this), FEE_SLOT);
+        assertEq(uint256(cachedFeeBefore), 0, "Fee should not be cached initially");
+
+        // Call rand which should cache the fee
+        euint256 a = e.rand();
+        processAllOperations();
+        assertEq(getUint256Value(a), 1);
+
+        // Verify the fee is now cached
+        bytes32 cachedFeeAfter = vm.load(address(this), FEE_SLOT);
+        assertEq(uint256(cachedFeeAfter), inco.getFee(), "Cached fee should match inco.getFee()");
+    }
+
+    function testFeeCachingPersistsAcrossOperations() public {
+        // Verify the slot is empty before any operation
+        bytes32 cachedFeeBefore = vm.load(address(this), FEE_SLOT);
+        assertEq(uint256(cachedFeeBefore), 0, "Fee should not be cached initially");
+
+        // First operation caches the fee
+        e.rand();
+        processAllOperations();
+
+        // Get the cached fee
+        bytes32 cachedFeeFirst = vm.load(address(this), FEE_SLOT);
+        assertEq(uint256(cachedFeeFirst), inco.getFee(), "Fee should be cached after first operation");
+
+        // Second operation should use cached fee
+        e.rand();
+        processAllOperations();
+
+        // Verify fee is still the same (wasn't re-fetched)
+        bytes32 cachedFeeSecond = vm.load(address(this), FEE_SLOT);
+        assertEq(uint256(cachedFeeFirst), uint256(cachedFeeSecond), "Cached fee should persist across operations");
+    }
+
+    function testFeeCachingOnRandBounded() public {
+        // Verify the slot is empty before any operation
+        bytes32 cachedFeeBefore = vm.load(address(this), FEE_SLOT);
+        assertEq(uint256(cachedFeeBefore), 0, "Fee should not be cached initially");
+
+        // Call randBounded which should cache the fee
+        euint256 a = e.randBounded(100);
+        processAllOperations();
+        assertEq(getUint256Value(a), 1);
+
+        // Verify the fee is now cached
+        bytes32 cachedFeeAfter = vm.load(address(this), FEE_SLOT);
+        assertEq(uint256(cachedFeeAfter), inco.getFee(), "Cached fee should match inco.getFee()");
+    }
+
+    function testFeeCachingOnNewEuint256() public {
+        // Create a contract that will use e.newEuint256
+        TakesEInput inputContract = new TakesEInput();
+        vm.deal(address(inputContract), 1 ether);
+
+        // Verify the slot is empty in the input contract before operation
+        bytes32 cachedFeeBefore = vm.load(address(inputContract), FEE_SLOT);
+        assertEq(uint256(cachedFeeBefore), 0, "Fee should not be cached initially");
+
+        // Call setA which uses e.newEuint256 internally
+        address self = address(this);
+        bytes memory ciphertext = fakePrepareEuint256Ciphertext(42, self, address(inputContract));
+        inputContract.setA(ciphertext);
+        processAllOperations();
+
+        // Verify the fee is now cached in the input contract
+        bytes32 cachedFeeAfter = vm.load(address(inputContract), FEE_SLOT);
+        assertEq(uint256(cachedFeeAfter), inco.getFee(), "Cached fee should match inco.getFee()");
+    }
+
+    function testFeeCachingOnNewEbool() public {
+        // Create a contract that will use e.newEbool
+        TakesEInput inputContract = new TakesEInput();
+        vm.deal(address(inputContract), 1 ether);
+
+        // Verify the slot is empty in the input contract before operation
+        bytes32 cachedFeeBefore = vm.load(address(inputContract), FEE_SLOT);
+        assertEq(uint256(cachedFeeBefore), 0, "Fee should not be cached initially");
+
+        // Call setB which uses e.newEbool internally
+        address self = address(this);
+        bytes memory ciphertext = fakePrepareEboolCiphertext(true, self, address(inputContract));
+        inputContract.setB(ciphertext);
+        processAllOperations();
+
+        // Verify the fee is now cached in the input contract
+        bytes32 cachedFeeAfter = vm.load(address(inputContract), FEE_SLOT);
+        assertEq(uint256(cachedFeeAfter), inco.getFee(), "Cached fee should match inco.getFee()");
+    }
+
+    function testFeeCachingOnEAddress() public {
+        // Create a contract that will use e.newEaddress
+        TakesEInput inputContract = new TakesEInput();
+        vm.deal(address(inputContract), 1 ether);
+
+        // Verify the slot is empty in the input contract before operation
+        bytes32 cachedFeeBefore = vm.load(address(inputContract), FEE_SLOT);
+        assertEq(uint256(cachedFeeBefore), 0, "Fee should not be cached initially");
+
+        // Call setC which uses e.newEaddress internally
+        address self = address(this);
+        bytes memory ciphertext = fakePrepareEaddressCiphertext(alice, self, address(inputContract));
+        inputContract.setC(ciphertext);
+        processAllOperations();
+
+        // Verify the fee is now cached in the input contract
+        bytes32 cachedFeeAfter = vm.load(address(inputContract), FEE_SLOT);
+        assertEq(uint256(cachedFeeAfter), inco.getFee(), "Cached fee should match inco.getFee()");
+    }
+
+    function testFeeRetryWithStaleCachedFee() public {
+        // First, cache a valid fee by calling rand
+        e.rand();
+        processAllOperations();
+
+        // Now manually set the cached fee to a wrong value
+        // This simulates a fee increase after caching
+        vm.store(address(this), FEE_SLOT, bytes32(uint256(1)));
+
+        // Verify the stale fee is set
+        bytes32 staleFee = vm.load(address(this), FEE_SLOT);
+        assertEq(uint256(staleFee), 1, "Stale fee should be 1 wei");
+
+        // Call rand again - the first call with stale fee should fail,
+        // but retry with fresh fee should succeed (no revert = success)
+        euint256 a = e.rand();
+        processAllOperations();
+
+        // Verify we got a valid handle (non-zero)
+        assertTrue(euint256.unwrap(a) != bytes32(0), "Should return a valid handle");
+
+        // Verify the fee cache was updated to the correct value
+        bytes32 updatedFee = vm.load(address(this), FEE_SLOT);
+        assertEq(uint256(updatedFee), inco.getFee(), "Fee should be updated after retry");
+    }
+
+    function testFeeRetryOnRandBounded() public {
+        // Cache a valid fee first
+        e.rand();
+        processAllOperations();
+
+        // Set stale fee
+        vm.store(address(this), FEE_SLOT, bytes32(uint256(1)));
+
+        // randBounded should also retry successfully (no revert = success)
+        euint256 a = e.randBounded(100);
+        processAllOperations();
+
+        // Verify we got a valid handle (non-zero)
+        assertTrue(euint256.unwrap(a) != bytes32(0), "Should return a valid handle");
+
+        // Verify fee was updated
+        bytes32 updatedFee = vm.load(address(this), FEE_SLOT);
+        assertEq(uint256(updatedFee), inco.getFee(), "Fee should be updated after retry");
+    }
+
+    function testFeeRetryOnNewEuint256() public {
+        TakesEInput inputContract = new TakesEInput();
+        vm.deal(address(inputContract), 1 ether);
+
+        // First call to cache the fee
+        address self = address(this);
+        bytes memory ciphertext1 = fakePrepareEuint256Ciphertext(42, self, address(inputContract));
+        inputContract.setA(ciphertext1);
+        processAllOperations();
+
+        // Verify first value was stored correctly
+        assertEq(getUint256Value(inputContract.a()), 42, "First value should be 42");
+
+        // Set stale fee in the input contract
+        vm.store(address(inputContract), FEE_SLOT, bytes32(uint256(1)));
+
+        // Second call should retry and succeed
+        bytes memory ciphertext2 = fakePrepareEuint256Ciphertext(99, self, address(inputContract));
+        inputContract.setA(ciphertext2);
+        processAllOperations();
+
+        // Verify the retried value was stored correctly
+        assertEq(getUint256Value(inputContract.a()), 99, "Retried value should be 99");
+
+        // Verify fee was updated
+        bytes32 updatedFee = vm.load(address(inputContract), FEE_SLOT);
+        assertEq(uint256(updatedFee), inco.getFee(), "Fee should be updated after retry");
+    }
+
+    function testFeeRetryOnNewEbool() public {
+        TakesEInput inputContract = new TakesEInput();
+        vm.deal(address(inputContract), 1 ether);
+
+        // First call to cache the fee
+        address self = address(this);
+        bytes memory ciphertext1 = fakePrepareEboolCiphertext(true, self, address(inputContract));
+        inputContract.setB(ciphertext1);
+        processAllOperations();
+
+        // Verify first value was stored correctly
+        assertEq(getBoolValue(inputContract.b()), true, "First value should be true");
+
+        // Set stale fee
+        vm.store(address(inputContract), FEE_SLOT, bytes32(uint256(1)));
+
+        // Second call should retry and succeed
+        bytes memory ciphertext2 = fakePrepareEboolCiphertext(false, self, address(inputContract));
+        inputContract.setB(ciphertext2);
+        processAllOperations();
+
+        // Verify the retried value was stored correctly
+        assertEq(getBoolValue(inputContract.b()), false, "Retried value should be false");
+
+        // Verify fee was updated
+        bytes32 updatedFee = vm.load(address(inputContract), FEE_SLOT);
+        assertEq(uint256(updatedFee), inco.getFee(), "Fee should be updated after retry");
+    }
+
+    function testFeeRetryOnNewEaddress() public {
+        TakesEInput inputContract = new TakesEInput();
+        vm.deal(address(inputContract), 1 ether);
+
+        // First call to cache the fee
+        address self = address(this);
+        bytes memory ciphertext1 = fakePrepareEaddressCiphertext(alice, self, address(inputContract));
+        inputContract.setC(ciphertext1);
+        processAllOperations();
+
+        // Verify first value was stored correctly
+        assertEq(getAddressValue(inputContract.c()), alice, "First value should be alice");
+
+        // Set stale fee
+        vm.store(address(inputContract), FEE_SLOT, bytes32(uint256(1)));
+
+        // Second call should retry and succeed
+        bytes memory ciphertext2 = fakePrepareEaddressCiphertext(bob, self, address(inputContract));
+        inputContract.setC(ciphertext2);
+        processAllOperations();
+
+        // Verify the retried value was stored correctly
+        assertEq(getAddressValue(inputContract.c()), bob, "Retried value should be bob");
+
+        // Verify fee was updated
+        bytes32 updatedFee = vm.load(address(inputContract), FEE_SLOT);
+        assertEq(uint256(updatedFee), inco.getFee(), "Fee should be updated after retry");
+    }
+
     function testEIfThenElse() public {
         ebool controlA = e.asEbool(true);
         ebool controlB = e.asEbool(false);