@inco/lightning 0.8.0-devnet → 0.8.0-devnet-1

package/src/libs/incoLightning_testnet_v0_183408998.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/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

@@ -9,6 +9,8 @@ uint256 constant FEE = 0.000001 ether;
 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/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/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);

package/src/test/TestFeeWithdrawal.t.sol

@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: No License
+pragma solidity ^0.8;
+
+import {inco} from "../Lib.sol";
+import {IncoTest} from "./IncoTest.sol";
+import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import {Fee} from "../lightning-parts/Fee.sol";
+
+contract RejectingContract {
+    // Contract that rejects all Ether transfers
+
+    }
+
+contract TestFeeWithdrawal is IncoTest {
+
+    function setUp() public override {
+        super.setUp();
+    }
+
+    function testWithdrawFees_Succeeds() public {
+        vm.deal(address(inco), 1 ether);
+        vm.prank(owner);
+        inco.withdrawFees();
+        assertEq(address(inco).balance, 0);
+        assertEq(address(owner).balance, 1 ether);
+    }
+
+    function testWithdrawFees_Reverts_NoFeesToWithdraw() public {
+        vm.expectRevert(Fee.NoFeesToWithdraw.selector);
+        vm.prank(owner);
+        inco.withdrawFees();
+    }
+
+    function testWithdrawFees_Reverts_NotOwner() public {
+        vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, alice));
+        vm.prank(alice);
+        inco.withdrawFees();
+    }
+
+    function testWithdrawFees_Reverts_FeeWithdrawalFailed() public {
+        // Deploy a contract that rejects Ether
+        RejectingContract rejectingOwner = new RejectingContract();
+
+        // Change owner to the rejecting contract
+        vm.prank(owner);
+        inco.transferOwnership(address(rejectingOwner));
+
+        // Add fees to withdraw
+        vm.deal(address(inco), 1 ether);
+
+        // Expect the FeeWithdrawalFailed error
+        vm.expectRevert(Fee.FeeWithdrawalFailed.selector);
+
+        // Try to withdraw (will fail because rejectingOwner can't receive Ether)
+        vm.prank(address(rejectingOwner));
+        inco.withdrawFees();
+    }
+
+}
+