@cofhe/mock-contracts 0.7.1 → 0.8.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/MockTaskManager.sol

@@ -263,22 +263,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 +340,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 +446,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);
@@ -455,19 +504,20 @@ 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);
     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 +661,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 +670,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 +679,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 +690,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 +705,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:

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;
   }