@cofhe/mock-contracts 0.7.1 → 1.0.0-alpha.0

package/contracts/MockCoFHEAddressBook.sol

@@ -0,0 +1,64 @@
+// SPDX-License-Identifier: BSD-3-Clause-Clear
+pragma solidity >=0.8.19 <0.9.0;
+
+import { ICoFHEAddressBook } from '@fhenixprotocol/cofhe-contracts/ICofhe.sol';
+
+/**
+ * @title  MockCoFHEAddressBook
+ * @notice Mock of the CoFHEAddressBook: maps a TaskManager id to the TaskManager that FHE.sol
+ *         releases pinned to that id talk to. Deployed to `COFHE_ADDRESS_BOOK` (from FHE.sol)
+ *         with `hardhat_setCode` / `deployCodeTo`, so no constructor or initializer runs.
+ * @dev    The production book is an Ownable2Step UUPS proxy; setters here are unrestricted.
+ */
+contract MockCoFHEAddressBook is ICoFHEAddressBook {
+  mapping(uint256 id => address tm) private taskManagers;
+
+  /// @notice          Emitted when an id is set or unset (`current == address(0)`).
+  /// @param id        TaskManager id.
+  /// @param previous  Address the id resolved to before, zero if it was unset.
+  /// @param current   Address the id resolves to now, zero if it was unset.
+  event TaskManagerSet(uint256 indexed id, address indexed previous, address indexed current);
+
+  /// @notice     Returned when an id resolves to nothing.
+  /// @param id   TaskManager id.
+  error TaskManagerNotSet(uint256 id);
+
+  /// @notice Returned when setting an id to the zero address.
+  error InvalidAddress();
+
+  function exists() public pure returns (bool) {
+    return true;
+  }
+
+  /// @notice     Resolves an id to its TaskManager.
+  /// @param id   TaskManager id.
+  /// @return tm  The TaskManager address; reverts with {TaskManagerNotSet} when unset.
+  function getTm(uint256 id) external view returns (address tm) {
+    tm = taskManagers[id];
+    if (tm == address(0)) {
+      revert TaskManagerNotSet(id);
+    }
+  }
+
+  /// @notice     Points an id at a TaskManager, replacing any previous target.
+  /// @param id   TaskManager id.
+  /// @param tm   TaskManager address.
+  function setTm(uint256 id, address tm) external {
+    if (tm == address(0)) {
+      revert InvalidAddress();
+    }
+    emit TaskManagerSet(id, taskManagers[id], tm);
+    taskManagers[id] = tm;
+  }
+
+  /// @notice     Retires an id: every FHE.sol build pinned to it reverts from then on.
+  /// @param id   TaskManager id.
+  function unsetTm(uint256 id) external {
+    address previous = taskManagers[id];
+    if (previous == address(0)) {
+      revert TaskManagerNotSet(id);
+    }
+    emit TaskManagerSet(id, previous, address(0));
+    delete taskManagers[id];
+  }
+}

package/contracts/MockERC1967Proxy.sol

@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: MIT
+pragma solidity >=0.8.25 <0.9.0;
+
+import { ERC1967Proxy } from '@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol';
+
+/**
+ * @notice OpenZeppelin's ERC-1967 proxy under a name of its own (no clash with a project's own
+ * `ERC1967Proxy` artifact). The ACP share registry is the production contract, which refuses to
+ * be initialized directly, so the mocks deploy it behind this proxy as production does.
+ */
+contract MockERC1967Proxy is ERC1967Proxy {
+  constructor(address implementation, bytes memory data) payable ERC1967Proxy(implementation, data) {}
+}

package/contracts/MockPermissioned.sol

@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: MIT
+pragma solidity >=0.8.25 <0.9.0;
+
+import { PermissionedUpgradeable } from './Permissioned.sol';
+
+/**
+ * @notice The ACL's ACP checks, taken unchanged from cofhe-contracts (`PermissionedUpgradeable`
+ * in Permissioned.sol), on a contract deployed directly rather than behind a proxy. The EIP-712
+ * domain ("ACL", "2") is defined by code there, so it needs no initialization.
+ */
+contract MockPermissioned is PermissionedUpgradeable {
+  /// @dev Exposed for tests that build typed-data digests.
+  function hashTypedDataV4(bytes32 structHash) public view virtual returns (bytes32) {
+    return _hashTypedDataV4(structHash);
+  }
+}

package/contracts/MockTaskManager.sol

@@ -31,7 +31,6 @@ error OnlyAggregatorAllowed(address caller);
 
 // Operation-specific errors
 error RandomFunctionNotSupported();
-error NotImplemented();
 
 library TMCommon {
   uint256 private constant HASH_MASK_FOR_METADATA = type(uint256).max - type(uint16).max; // 2 bytes reserved for metadata
@@ -249,6 +248,9 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
   // Signer address for decrypt result verification (threshold network's signing key)
   address public decryptResultSigner;
 
+  // Counter mixed into generated random seeds, as in the real task manager
+  uint256 private randomCounter;
+
   // Storage contract for plaintext results of decrypt operations
   // PlaintextsStorage public plaintextsStorage;
 
@@ -263,22 +265,63 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
     return true;
   }
 
-  function sendEventCreated(uint256 ctHash, string memory operation, uint256[] memory inputs) private {
-    if (inputs.length == 1 || opIs(operation, FunctionId.cast)) {
-      emit TaskCreated(ctHash, operation, inputs[0], 0, 0);
+  error OnlySelf();
 
-      // NOTE: MOCK
-      MOCK_unaryOperation(ctHash, operation, inputs[0]);
-    } else if (inputs.length == 2) {
-      emit TaskCreated(ctHash, operation, inputs[0], inputs[1], 0);
+  function sendEventCreated(
+    uint256 ctHash,
+    FunctionId funcId,
+    string memory operation,
+    uint256[] memory inputs
+  ) private {
+    uint8 arity = (inputs.length == 1 || funcId == FunctionId.cast) ? 1 : (inputs.length == 2 ? 2 : 3);
 
-      // NOTE: MOCK
-      MOCK_twoInputOperation(ctHash, operation, inputs[0], inputs[1]);
+    if (arity == 1) {
+      emit TaskCreated(ctHash, operation, inputs[0], 0, 0);
+    } else if (arity == 2) {
+      emit TaskCreated(ctHash, operation, inputs[0], inputs[1], 0);
     } else {
       emit TaskCreated(ctHash, operation, inputs[0], inputs[1], inputs[2]);
+    }
 
-      // NOTE: MOCK
-      MOCK_threeInputOperation(ctHash, operation, inputs[0], inputs[1], inputs[2]);
+    // NOTE: MOCK - plaintext replication is mock-only work that doesn't exist in the real
+    // task manager, so (under forge, when enabled) it runs with gas metering paused; when
+    // metered (hardhat) its cost is reported via MockGasConsumed instead.
+    (MockGasMode mode, uint256 startGas) = _mockGasTrackStart();
+    if (mode != MockGasMode.PAUSED_BY_MOCK) {
+      // MEASURE, or the caller's test already paused metering (nothing of ours to leak
+      // on revert in either case).
+      _mockDispatch(ctHash, funcId, inputs, arity);
+      _mockGasTrackEnd(mode, startGas);
+      return;
+    }
+
+    // External self-call so a revert inside the mock op can be caught and gas metering
+    // resumed before bubbling the original error - otherwise the paused state leaks into
+    // the rest of the test and every later call reports ~0 gas. The extra call itself
+    // runs unmetered.
+    try this.MOCK_dispatchOperation(ctHash, funcId, inputs, arity) {
+      _mockGasTrackEnd(mode, startGas);
+    } catch (bytes memory err) {
+      _mockGasTrackEnd(mode, startGas);
+      assembly ('memory-safe') {
+        revert(add(err, 0x20), mload(err))
+      }
+    }
+  }
+
+  /// @dev Trampoline for `sendEventCreated`'s try/catch; not part of the mocked interface.
+  function MOCK_dispatchOperation(uint256 ctHash, FunctionId funcId, uint256[] calldata inputs, uint8 arity) external {
+    if (msg.sender != address(this)) revert OnlySelf();
+    _mockDispatch(ctHash, funcId, inputs, arity);
+  }
+
+  function _mockDispatch(uint256 ctHash, FunctionId funcId, uint256[] memory inputs, uint8 arity) private {
+    if (arity == 1) {
+      MOCK_unaryOperation(ctHash, funcId, inputs[0]);
+    } else if (arity == 2) {
+      MOCK_twoInputOperation(ctHash, funcId, inputs[0], inputs[1]);
+    } else {
+      MOCK_threeInputOperation(ctHash, funcId, inputs[0], inputs[1], inputs[2]);
     }
   }
 
@@ -299,12 +342,18 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
     //     );
     // }
 
-    // NOTE: MOCK
+    // NOTE: MOCK - the real task manager only emits a task event here and the result lands
+    // in a later transaction; storing it synchronously is mock-only work. Only the cheap
+    // existence check runs before the tracked block (so an InputNotInMockStorage revert
+    // can't leak paused gas metering); the read itself is mock-only and stays inside.
+    MOCK_verifyKeyInStorage(ctHash);
+    (MockGasMode mode, uint256 startGas) = _mockGasTrackStart();
     _decryptResultReady[ctHash] = true;
     _decryptResult[ctHash] = _get(ctHash);
 
     uint64 asyncOffset = uint64((block.timestamp % 10) + 1);
     _decryptResultReadyTimestamp[ctHash] = uint64(block.timestamp) + asyncOffset;
+    _mockGasTrackEnd(mode, startGas);
   }
 
   function getDecryptResult(uint256 ctHash) public view returns (uint256) {
@@ -399,9 +448,11 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
     }
   }
 
-  function validateInputs(uint256[] memory encryptedHashes, FunctionId funcId) internal view {
-    string memory functionName = Utils.functionIdToString(funcId);
-
+  function validateInputs(
+    uint256[] memory encryptedHashes,
+    FunctionId funcId,
+    string memory functionName
+  ) internal view {
     if (encryptedHashes.length == 0) {
       if (!isPlaintextOperation(funcId)) {
         revert InvalidOperationInputs(functionName);
@@ -432,6 +483,32 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
     validateFunctionInputTypes(funcId, functionName, encryptedHashes);
   }
 
+  /// @dev Same checks as the real task manager: the plaintext must fit the target type and the
+  ///      security zone must be in range. Without them `FHE.asEuint8(256)` would be stored masked
+  ///      to 0 here while the real task manager reverts.
+  function validateTrivialEncryptInputs(uint256[] memory extraInputs) internal view {
+    uint256 securityZone = extraInputs[2];
+    if (securityZone > uint256(int256(type(int32).max)) || !isValidSecurityZone(int32(int256(securityZone)))) {
+      revert InvalidSecurityZone(int32(int256(securityZone)), securityZoneMin, securityZoneMax);
+    }
+
+    uint256 value = extraInputs[0];
+    uint256 toType = extraInputs[1];
+    uint256 max;
+    if (toType == Utils.EUINT8_TFHE) max = type(uint8).max;
+    else if (toType == Utils.EUINT16_TFHE) max = type(uint16).max;
+    else if (toType == Utils.EUINT32_TFHE) max = type(uint32).max;
+    else if (toType == Utils.EUINT64_TFHE) max = type(uint64).max;
+    else if (toType == Utils.EUINT128_TFHE) max = type(uint128).max;
+    else if (toType == Utils.EADDRESS_TFHE) max = type(uint160).max;
+    else if (toType == Utils.EBOOL_TFHE) max = 1;
+    else return;
+
+    if (value > max) {
+      revert InvalidInputForFunction('trivialEncrypt', uint8(toType));
+    }
+  }
+
   function validateSelectInputs(uint256[] memory encryptedHashes) internal pure {
     if (encryptedHashes.length != 3) {
       revert InvalidInputsAmount('select', encryptedHashes.length, 3);
@@ -455,19 +532,23 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
     if (funcId == FunctionId.random) {
       revert RandomFunctionNotSupported();
     }
+    string memory operation = Utils.functionIdToString(funcId);
     uint256 inputsLength = encryptedHashes.length + extraInputs.length;
     if (inputsLength > 3) {
-      revert TooManyInputs(Utils.functionIdToString(funcId), inputsLength, 3);
+      revert TooManyInputs(operation, inputsLength, 3);
     }
 
-    validateInputs(encryptedHashes, funcId);
+    validateInputs(encryptedHashes, funcId, operation);
+    if (funcId == FunctionId.trivialEncrypt) {
+      validateTrivialEncryptInputs(extraInputs);
+    }
     uint256[] memory inputs = TMCommon.combineInputs(encryptedHashes, extraInputs);
 
     int32 securityZone = getSecurityZone(funcId, encryptedHashes, extraInputs);
     uint256 ctHash = TMCommon.calcPlaceholderKey(returnType, securityZone, inputs, funcId);
 
     acl.allowTransient(ctHash, msg.sender, address(this));
-    sendEventCreated(ctHash, Utils.functionIdToString(funcId), inputs);
+    sendEventCreated(ctHash, funcId, operation, inputs);
 
     return ctHash;
   }
@@ -611,7 +692,7 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
       acl.allow(ctHash, account, msg.sender);
 
       // NOTE: MOCK
-      MOCK_logAllow(account == msg.sender ? 'FHE.allowThis' : 'FHE.allow', ctHash, account);
+      MOCK_trackedLogAllow(account == msg.sender ? 'FHE.allowThis' : 'FHE.allow', ctHash, account);
     }
   }
 
@@ -620,7 +701,7 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
       acl.allowGlobal(ctHash, msg.sender);
 
       // NOTE: MOCK
-      MOCK_logAllow('FHE.allowGlobal', ctHash, msg.sender);
+      MOCK_trackedLogAllow('FHE.allowGlobal', ctHash, msg.sender);
     }
   }
 
@@ -629,7 +710,7 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
       acl.allowTransient(ctHash, account, msg.sender);
 
       // NOTE: MOCK
-      MOCK_logAllow('FHE.allowTransient', ctHash, account);
+      MOCK_trackedLogAllow('FHE.allowTransient', ctHash, account);
     }
   }
 
@@ -640,7 +721,7 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
       acl.allowForDecryption(hashes, msg.sender);
 
       // NOTE: MOCK
-      MOCK_logAllow('FHE.allowForDecryption', ctHash, msg.sender);
+      MOCK_trackedLogAllow('FHE.allowForDecryption', ctHash, msg.sender);
     }
   }
 
@@ -655,14 +736,14 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
     acl.shareCtHash(ctHash, msg.sender, receiver);
 
     // NOTE: MOCK
-    MOCK_logAllow('FHE.share', ctHash, receiver);
+    MOCK_trackedLogAllow('FHE.share', ctHash, receiver);
   }
 
   function receiveCtHash(uint256 ctHash, address expectedSharer) external {
     acl.receiveCtHash(ctHash, expectedSharer, msg.sender);
 
     // NOTE: MOCK
-    MOCK_logAllow('FHE.receive', ctHash, expectedSharer);
+    MOCK_trackedLogAllow('FHE.receive', ctHash, expectedSharer);
   }
 
   /// @dev Per-input message hash used by batch verification:
@@ -763,12 +844,46 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
   // Stub implementations for new ITaskManager interface methods (inc PR #48)
 
   function createRandomTask(uint8 returnType, uint256 seed, int32 securityZone) external returns (uint256) {
-    // Mock implementation: just return a pseudo-random hash based on seed
-    return uint256(keccak256(abi.encode(returnType, seed, securityZone, block.timestamp)));
+    if (!isValidSecurityZone(securityZone)) {
+      revert InvalidSecurityZone(securityZone, securityZoneMin, securityZoneMax);
+    }
+
+    if (seed == 0) {
+      seed = uint256(
+        keccak256(
+          abi.encodePacked(blockhash(block.number - 1), block.timestamp, randomCounter, block.chainid, securityZone)
+        )
+      );
+      unchecked {
+        randomCounter++;
+      }
+    }
+
+    // Same handle derivation as the real task manager: msg.sender is part of the preimage.
+    uint256[] memory inputs = new uint256[](2);
+    inputs[0] = seed;
+    inputs[1] = uint256(uint160(msg.sender));
+
+    uint256 ctHash = TMCommon.calcPlaceholderKey(returnType, securityZone, inputs, FunctionId.random);
+    acl.allowTransient(ctHash, msg.sender, address(this));
+    emit TaskCreated(
+      ctHash,
+      Utils.functionIdToString(FunctionId.random),
+      seed,
+      uint256(uint32(securityZone)),
+      inputs[1]
+    );
+
+    // NOTE: MOCK - the real task manager leaves the value to the coprocessor; store one here.
+    (MockGasMode mode, uint256 startGas) = _mockGasTrackStart();
+    _set(ctHash, uint256(keccak256(abi.encode(ctHash))));
+    _mockGasTrackEnd(mode, startGas);
+
+    return ctHash;
   }
 
   function isPubliclyAllowed(uint256 ctHash) external view returns (bool) {
-    revert NotImplemented();
+    return acl.globalAllowed(ctHash);
   }
 
   function verifyDecryptResultBatch(

package/contracts/MockThresholdNetwork.sol

@@ -9,7 +9,7 @@ pragma solidity >=0.8.19 <0.9.0;
 
 import { MockACL } from './MockACL.sol';
 import { MockTaskManager } from './MockTaskManager.sol';
-import { MockPermissioned, ACP } from './Permissioned.sol';
+import { ACP, PermissionedUpgradeable } from './Permissioned.sol';
 
 contract MockThresholdNetwork {
   MockTaskManager public mockTaskManager;
@@ -181,16 +181,16 @@ contract MockThresholdNetwork {
 
   function decodeLowLevelReversion(bytes memory data) public pure returns (string memory error) {
     bytes4 selector = bytes4(data);
-    if (selector == MockPermissioned.PermissionInvalid_Expired.selector) {
+    if (selector == PermissionedUpgradeable.PermissionInvalid_Expired.selector) {
       return 'PermissionInvalid_Expired';
     }
-    if (selector == MockPermissioned.PermissionInvalid_IssuerSignature.selector) {
+    if (selector == PermissionedUpgradeable.PermissionInvalid_IssuerSignature.selector) {
       return 'PermissionInvalid_IssuerSignature';
     }
-    if (selector == MockPermissioned.PermissionInvalid_RecipientSignature.selector) {
+    if (selector == PermissionedUpgradeable.PermissionInvalid_RecipientSignature.selector) {
       return 'PermissionInvalid_RecipientSignature';
     }
-    if (selector == MockPermissioned.PermissionInvalid_Disabled.selector) {
+    if (selector == PermissionedUpgradeable.PermissionInvalid_Disabled.selector) {
       return 'PermissionInvalid_Disabled';
     }
     // Handle other errors

package/contracts/MockZkVerifier.sol

@@ -3,7 +3,7 @@
 
 pragma solidity >=0.8.19 <0.9.0;
 
-import { TASK_MANAGER_ADDRESS } from '@fhenixprotocol/cofhe-contracts/FHE.sol';
+import { Common } from '@fhenixprotocol/cofhe-contracts/FHE.sol';
 import { EncryptedInput } from '@fhenixprotocol/cofhe-contracts/ICofhe.sol';
 import { MockTaskManager } from './MockTaskManager.sol';
 
@@ -104,7 +104,7 @@ contract MockZkVerifier {
   }
 
   function insertCtHash(uint256 ctHash, uint256 value) public {
-    MockTaskManager(TASK_MANAGER_ADDRESS).MOCK_setInEuintKey(ctHash, value);
+    MockTaskManager(address(Common.tm())).MOCK_setInEuintKey(ctHash, value);
     salt += 1;
   }