@cofhe/mock-contracts 0.8.0 → 1.0.0

package/contracts/MockTaskManager.sol

@@ -16,6 +16,7 @@ error InvalidInputsAmount(string operation, uint256 got, uint256 expected);
 error InvalidOperationInputs(string operation);
 error TooManyInputs(string operation, uint256 got, uint256 maxAllowed);
 error InvalidBytesLength(uint256 got, uint256 expected);
+error LengthMismatch();
 // Type and security validation errors
 error InvalidTypeOrSecurityZone(string operation);
 error InvalidInputType(uint8 actual, uint8 expected);
@@ -31,7 +32,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
@@ -159,10 +159,6 @@ library TMCommon {
     return uint256((SHIFTED_TYPE_MASK | SECURITY_ZONE_MASK) & hash);
   }
 
-  function isTriviallyEncryptedFromHash(uint256 hash) internal pure returns (bool) {
-    return (hash & TRIVIALLY_ENCRYPTED_MASK) == TRIVIALLY_ENCRYPTED_MASK;
-  }
-
   function computeDecryptResultHash(
     uint256 ctHash,
     uint256 result,
@@ -211,8 +207,17 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
   }
 
   function setSecurityZones(int32 minSZ, int32 maxSZ) external onlyOwner {
-    securityZoneMin = minSZ;
-    securityZoneMax = maxSZ;
+    _setSecurityZones(minSZ, maxSZ);
+  }
+
+  /// @dev Same check as the real task manager's `_setSecurityZones`, the one write path for the bounds.
+  ///      Without it `setSecurityZones(2, 1)` would be stored here while the real task manager reverts.
+  function _setSecurityZones(int32 newMin, int32 newMax) private {
+    if (newMin > newMax) {
+      revert InvalidSecurityZone(newMin, newMin, newMax);
+    }
+    securityZoneMin = newMin;
+    securityZoneMax = newMax;
   }
 
   function isInitialized() public view returns (bool) {
@@ -249,6 +254,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;
 
@@ -370,9 +378,7 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
   }
 
   function checkAllowed(uint256 ctHash) internal view {
-    if (!TMCommon.isTriviallyEncryptedFromHash(ctHash)) {
-      if (!acl.isAllowed(ctHash, msg.sender)) revert ACLNotAllowed(ctHash, msg.sender);
-    }
+    if (!acl.isAllowed(ctHash, msg.sender)) revert ACLNotAllowed(ctHash, msg.sender);
   }
 
   function isUnaryOperation(FunctionId funcId) internal pure returns (bool) {
@@ -481,6 +487,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);
@@ -511,6 +543,9 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
     }
 
     validateInputs(encryptedHashes, funcId, operation);
+    if (funcId == FunctionId.trivialEncrypt) {
+      validateTrivialEncryptInputs(extraInputs);
+    }
     uint256[] memory inputs = TMCommon.combineInputs(encryptedHashes, extraInputs);
 
     int32 securityZone = getSecurityZone(funcId, encryptedHashes, extraInputs);
@@ -544,12 +579,7 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
   /// @notice Publish a signed decrypt result to the chain
   /// @dev Anyone with a valid signature from the decrypt network can call this
   function publishDecryptResult(uint256 ctHash, uint256 result, bytes calldata signature) external {
-    _verifyDecryptResult(ctHash, result, signature, true);
-    _decryptResultReady[ctHash] = true;
-    _decryptResult[ctHash] = result;
-
-    _decryptResultReadyTimestamp[ctHash] = uint64(block.timestamp);
-    emit DecryptionResult(ctHash, result, msg.sender);
+    _publishDecryptResult(ctHash, result, signature);
   }
 
   function publishDecryptResultBatch(
@@ -557,14 +587,25 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
     uint256[] calldata results,
     bytes[] calldata signatures
   ) external {
-    // Mock implementation: process each result individually
-    // Note: publishDecryptResult is defined later in the contract
-    for (uint256 i = 0; i < ctHashes.length; i++) {
-      // Call via external interface to avoid forward reference issue
-      this.publishDecryptResult(ctHashes[i], results[i], signatures[i]);
+    uint256 length = ctHashes.length;
+    if (results.length != length || signatures.length != length) revert LengthMismatch();
+
+    for (uint256 i = 0; i < length; i++) {
+      _publishDecryptResult(ctHashes[i], results[i], signatures[i]);
     }
   }
 
+  /// @dev Shared by both publish functions, so a batch keeps its caller as the `DecryptionResult` requestor,
+  ///      like the real task manager. An external self-call would make the requestor this contract.
+  function _publishDecryptResult(uint256 ctHash, uint256 result, bytes calldata signature) private {
+    _verifyDecryptResult(ctHash, result, signature, true);
+    _decryptResultReady[ctHash] = true;
+    _decryptResult[ctHash] = result;
+
+    _decryptResultReadyTimestamp[ctHash] = uint64(block.timestamp);
+    emit DecryptionResult(ctHash, result, msg.sender);
+  }
+
   function verifyDecryptResult(uint256 ctHash, uint256 result, bytes calldata signature) external view returns (bool) {
     return _verifyDecryptResult(ctHash, result, signature, true);
   }
@@ -657,47 +698,36 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
   }
 
   function allow(uint256 ctHash, address account) external {
-    if (!TMCommon.isTriviallyEncryptedFromHash(ctHash)) {
-      acl.allow(ctHash, account, msg.sender);
+    acl.allow(ctHash, account, msg.sender);
 
-      // NOTE: MOCK
-      MOCK_trackedLogAllow(account == msg.sender ? 'FHE.allowThis' : 'FHE.allow', ctHash, account);
-    }
+    // NOTE: MOCK
+    MOCK_trackedLogAllow(account == msg.sender ? 'FHE.allowThis' : 'FHE.allow', ctHash, account);
   }
 
   function allowGlobal(uint256 ctHash) external {
-    if (!TMCommon.isTriviallyEncryptedFromHash(ctHash)) {
-      acl.allowGlobal(ctHash, msg.sender);
+    acl.allowGlobal(ctHash, msg.sender);
 
-      // NOTE: MOCK
-      MOCK_trackedLogAllow('FHE.allowGlobal', ctHash, msg.sender);
-    }
+    // NOTE: MOCK
+    MOCK_trackedLogAllow('FHE.allowGlobal', ctHash, msg.sender);
   }
 
   function allowTransient(uint256 ctHash, address account) external {
-    if (!TMCommon.isTriviallyEncryptedFromHash(ctHash)) {
-      acl.allowTransient(ctHash, account, msg.sender);
+    acl.allowTransient(ctHash, account, msg.sender);
 
-      // NOTE: MOCK
-      MOCK_trackedLogAllow('FHE.allowTransient', ctHash, account);
-    }
+    // NOTE: MOCK
+    MOCK_trackedLogAllow('FHE.allowTransient', ctHash, account);
   }
 
   function allowForDecryption(uint256 ctHash) external {
-    if (!TMCommon.isTriviallyEncryptedFromHash(ctHash)) {
-      uint256[] memory hashes = new uint256[](1);
-      hashes[0] = ctHash;
-      acl.allowForDecryption(hashes, msg.sender);
+    uint256[] memory hashes = new uint256[](1);
+    hashes[0] = ctHash;
+    acl.allowForDecryption(hashes, msg.sender);
 
-      // NOTE: MOCK
-      MOCK_trackedLogAllow('FHE.allowForDecryption', ctHash, msg.sender);
-    }
+    // NOTE: MOCK
+    MOCK_trackedLogAllow('FHE.allowForDecryption', ctHash, msg.sender);
   }
 
   function isAllowed(uint256 ctHash, address account) external view returns (bool) {
-    if (TMCommon.isTriviallyEncryptedFromHash(ctHash)) {
-      return true;
-    }
     return acl.isAllowed(ctHash, account);
   }
 
@@ -769,17 +799,11 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
   }
 
   function setSecurityZoneMax(int32 securityZone) external onlyOwner {
-    if (securityZone < securityZoneMin) {
-      revert InvalidSecurityZone(securityZone, securityZoneMin, securityZoneMax);
-    }
-    securityZoneMax = securityZone;
+    _setSecurityZones(securityZoneMin, securityZone);
   }
 
   function setSecurityZoneMin(int32 securityZone) external onlyOwner {
-    if (securityZone > securityZoneMax) {
-      revert InvalidSecurityZone(securityZone, securityZoneMin, securityZoneMax);
-    }
-    securityZoneMin = securityZone;
+    _setSecurityZones(securityZone, securityZoneMax);
   }
 
   function setACLContract(address _aclAddress) external onlyOwner {
@@ -813,12 +837,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(
@@ -826,7 +884,10 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
     uint256[] calldata results,
     bytes[] calldata signatures
   ) external view returns (bool) {
-    for (uint256 i = 0; i < ctHashes.length; i++) {
+    uint256 length = ctHashes.length;
+    if (results.length != length || signatures.length != length) revert LengthMismatch();
+
+    for (uint256 i = 0; i < length; i++) {
       if (!_verifyDecryptResult(ctHashes[i], results[i], signatures[i], true)) {
         return false;
       }
@@ -839,8 +900,11 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
     uint256[] calldata results,
     bytes[] calldata signatures
   ) external view returns (bool[] memory) {
-    bool[] memory out = new bool[](ctHashes.length);
-    for (uint256 i = 0; i < ctHashes.length; i++) {
+    uint256 length = ctHashes.length;
+    if (results.length != length || signatures.length != length) revert LengthMismatch();
+
+    bool[] memory out = new bool[](length);
+    for (uint256 i = 0; i < length; i++) {
       out[i] = _verifyDecryptResult(ctHashes[i], results[i], signatures[i], false);
     }
     return out;

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