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

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/lightning-parts/TEELifecycle.sol

@@ -70,6 +70,8 @@ abstract contract TEELifecycle is
     // Events
     // @notice A new MR_AGGREGATED has been approved
     event NewTEEVersionApproved(uint256 indexed version, bytes32 indexed mrAggregated);
+    // @notice A previously approved TEE version has been removed
+    event TEEVersionRemoved(bytes32 indexed mrAggregated);
     event NewCovalidatorAdded(address covalidatorAddress, bytes quote);
     event BootstrapStageComplete(address indexed newEoaSigner, BootstrapResult bootstrapResult);
     // @notice Emitted to prove that an EOA has upgraded their TDX to a new
@@ -211,6 +213,28 @@ abstract contract TEELifecycle is
         emit NewTEEVersionApproved($.approvedTeeVersions.length - 1, newMrAggregated);
     }
 
+    /**
+     * @notice Removes a previously approved TEE version from the contract state
+     * @param mrAggregated - The MR_AGGREGATED bytes of the TEE version to remove
+     */
+    function removeApprovedTeeVersion(bytes32 mrAggregated) public onlyOwner {
+        StorageForTeeLifecycle storage $ = getTeeLifecycleStorage();
+        bool found = false;
+        for (uint256 i = 0; i < $.approvedTeeVersions.length; i++) {
+            if ($.approvedTeeVersions[i] == mrAggregated) {
+                // Shift all elements after index i to the left to preserve insertion order
+                for (uint256 j = i; j < $.approvedTeeVersions.length - 1; j++) {
+                    $.approvedTeeVersions[j] = $.approvedTeeVersions[j + 1];
+                }
+                $.approvedTeeVersions.pop();
+                found = true;
+                break;
+            }
+        }
+        require(found, TEEVersionNotFound());
+        emit TEEVersionRemoved(mrAggregated);
+    }
+
     function _isApprovedTeeVersion(bytes32 newMrAggregated) internal view returns (bool) {
         StorageForTeeLifecycle storage $ = getTeeLifecycleStorage();
         for (uint256 i = 0; i < $.approvedTeeVersions.length; i++) {

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/TEELifecycle/TEELifecycleMockTest.t.sol

@@ -124,6 +124,61 @@ contract TEELifecycleMockTest is MockRemoteAttestation, TEELifecycle {
         vm.stopPrank();
     }
 
+    function testRemoveApprovedTeeVersionPreservesOrder() public {
+        bytes32 mrAggregated1 = hex"1111111111111111111111111111111111111111111111111111111111111111";
+        bytes32 mrAggregated2 = hex"2222222222222222222222222222222222222222222222222222222222222222";
+        bytes32 mrAggregated3 = hex"3333333333333333333333333333333333333333333333333333333333333333";
+
+        vm.startPrank(this.owner());
+
+        // Add three versions
+        this.approveNewTeeVersion(mrAggregated1);
+        this.approveNewTeeVersion(mrAggregated2);
+        this.approveNewTeeVersion(mrAggregated3);
+
+        // Verify all exist in order
+        assertEq(this.approvedTeeVersions(0), mrAggregated1);
+        assertEq(this.approvedTeeVersions(1), mrAggregated2);
+        assertEq(this.approvedTeeVersions(2), mrAggregated3);
+
+        // Remove the middle one (mrAggregated2)
+        this.removeApprovedTeeVersion(mrAggregated2);
+
+        // Verify insertion order is preserved: mrAggregated1 stays at 0, mrAggregated3 shifts to 1
+        assertEq(this.approvedTeeVersions(0), mrAggregated1);
+        assertEq(this.approvedTeeVersions(1), mrAggregated3);
+
+        // Verify index 2 is now out of bounds
+        vm.expectRevert(TEELifecycle.IndexOutOfBounds.selector);
+        this.approvedTeeVersions(2);
+
+        vm.stopPrank();
+    }
+
+    function testRemoveApprovedTeeVersionNotFound() public {
+        bytes32 nonExistentMrAggregated = hex"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
+
+        vm.startPrank(this.owner());
+        vm.expectRevert(TEELifecycle.TEEVersionNotFound.selector);
+        this.removeApprovedTeeVersion(nonExistentMrAggregated);
+        vm.stopPrank();
+    }
+
+    function testRemoveApprovedTeeVersionOnlyOwner() public {
+        bytes32 mrAggregated = hex"1111111111111111111111111111111111111111111111111111111111111111";
+
+        vm.startPrank(this.owner());
+        this.approveNewTeeVersion(mrAggregated);
+        vm.stopPrank();
+
+        // Try to remove as non-owner
+        address nonOwner = address(0x1234);
+        vm.startPrank(nonOwner);
+        vm.expectRevert();
+        this.removeApprovedTeeVersion(mrAggregated);
+        vm.stopPrank();
+    }
+
     // Helper function to create a successful bootstrap result
     function successfulBootstrapResult()
         internal