@cofhe/mock-contracts 0.0.0-alpha-20260409113701

package/contracts/MockCoFHE.sol

@@ -0,0 +1,444 @@
+// SPDX-License-Identifier: BSD-3-Clause-Clear
+pragma solidity >=0.8.25 <0.9.0;
+
+import { Strings } from '@openzeppelin/contracts/utils/Strings.sol';
+import { FHE } from '@fhenixprotocol/cofhe-contracts/FHE.sol';
+import { FunctionId, Utils } from '@fhenixprotocol/cofhe-contracts/ICofhe.sol';
+import { console } from 'hardhat/console.sol';
+
+address constant ZK_VERIFIER_SIGNER_ADDRESS = 0x6E12D8C87503D4287c294f2Fdef96ACd9DFf6bd2;
+uint256 constant ZK_VERIFIER_SIGNER_PRIVATE_KEY = 49099792800763675079532137679706322989817545357788440619111868498148356080914;
+
+address constant DECRYPT_RESULT_SIGNER_ADDRESS = 0x70997970C51812dc3A010C7d01b50e0d17dc79C8;
+uint256 constant DECRYPT_RESULT_SIGNER_PRIVATE_KEY = 0x59c6995e998f97a5a0044966f0945389dc9e86dae88c7a8412f4603b6b78690d;
+
+/**
+ * @dev Mock implementation of the CoFHE coprocessor, used to test FHE ops in isolation.
+ * Is inherited by MockTaskManager.
+ *
+ * It is responsible for storing a map of ctHash -> value
+ * and for performing the operations on the values.
+ *
+ * It is intended as a 1:1 drop-in replacement for the real CoFHE coprocessor, with the following differences:
+ * - AsyncCallbacks are called synchronously (with a mock 1-10 second delay).
+ * - Unencrypted values are available onchain via the `mockStorage` map.
+ *
+ * NOTE: This is not used in production
+ */
+abstract contract MockCoFHE {
+  // Pulled from TMCommon
+  uint256 constant uintTypeMask = (type(uint8).max >> 1); // 0x7f - 7 bits reserved for uint type in the one before last byte
+  uint256 constant triviallyEncryptedMask = type(uint8).max - uintTypeMask; //0x80  1 bit reserved for isTriviallyEncrypted
+  uint256 constant shiftedTypeMask = uintTypeMask << 8; // 0x7f007 bits reserved for uint type in the one before last byte
+
+  bool public logOps = true;
+
+  mapping(uint256 => uint256) public mockStorage;
+  mapping(uint256 => bool) public inMockStorage;
+
+  error InputNotInMockStorage(uint256 ctHash);
+
+  // Used internally to check if we missed any operations in the mocks
+  error InvalidUnaryOperation(string operation);
+  error InvalidTwoInputOperation(string operation);
+  error InvalidThreeInputOperation(string operation);
+
+  // OPTIONS
+
+  function setLogOps(bool _logOps) public {
+    logOps = _logOps;
+  }
+
+  // Utils
+
+  function getUintTypeFromHash(uint256 hash) internal pure returns (uint8) {
+    return uint8((hash & shiftedTypeMask) >> 8);
+  }
+
+  function getUtypeStringFromHash(uint256 hash) internal pure returns (string memory) {
+    uint8 inputType = getUintTypeFromHash(hash);
+    if (inputType == Utils.EBOOL_TFHE) return 'ebool';
+    if (inputType == Utils.EUINT8_TFHE) return 'euint8';
+    if (inputType == Utils.EUINT16_TFHE) return 'euint16';
+    if (inputType == Utils.EUINT32_TFHE) return 'euint32';
+    if (inputType == Utils.EUINT64_TFHE) return 'euint64';
+    if (inputType == Utils.EUINT128_TFHE) return 'euint128';
+    if (inputType == Utils.EUINT256_TFHE) return 'euint256';
+    if (inputType == Utils.EADDRESS_TFHE) return 'eaddress';
+    return 'unknown';
+  }
+
+  function getUtypeBits(uint256 hash) internal pure returns (uint256) {
+    uint8 inputType = getUintTypeFromHash(hash);
+    if (inputType == Utils.EBOOL_TFHE) return 8;
+    if (inputType == Utils.EUINT8_TFHE) return 8;
+    if (inputType == Utils.EUINT16_TFHE) return 16;
+    if (inputType == Utils.EUINT32_TFHE) return 32;
+    if (inputType == Utils.EUINT64_TFHE) return 64;
+    if (inputType == Utils.EUINT128_TFHE) return 128;
+    if (inputType == Utils.EUINT256_TFHE) return 256;
+    if (inputType == Utils.EADDRESS_TFHE) return 160;
+    return 0;
+  }
+
+  function getUtypeMask(uint256 hash) internal pure returns (uint256) {
+    uint256 bits = getUtypeBits(hash);
+    unchecked {
+      return (1 << bits) - 1;
+    }
+  }
+
+  function removeFirstLetter(string memory str) public pure returns (string memory) {
+    bytes memory strBytes = bytes(str);
+    if (strBytes.length == 0) return '';
+    bytes memory result = new bytes(strBytes.length - 1);
+    for (uint i = 1; i < strBytes.length; i++) {
+      result[i - 1] = strBytes[i];
+    }
+    return string(result);
+  }
+
+  function getIsBoolTypeFromHash(uint256 hash) internal pure returns (bool) {
+    uint8 inputType = getUintTypeFromHash(hash);
+    return (inputType ^ Utils.EBOOL_TFHE) == 0;
+  }
+
+  function strEq(string memory _a, string memory _b) internal pure returns (bool) {
+    return keccak256(abi.encodePacked(_a)) == keccak256(abi.encodePacked(_b));
+  }
+
+  function opIs(string memory op, FunctionId fid) internal pure returns (bool) {
+    return strEq(op, Utils.functionIdToString(fid));
+  }
+
+  function sliceString(string memory str, uint start, uint length) public pure returns (string memory) {
+    bytes memory strBytes = bytes(str);
+    require(start + length <= strBytes.length, 'Out of bounds');
+
+    bytes memory result = new bytes(length);
+    for (uint i = 0; i < length; i++) {
+      result[i] = strBytes[start + i];
+    }
+
+    return string(result);
+  }
+
+  function logCtHash(uint256 ctHash) internal view returns (string memory) {
+    string memory hashStr = Strings.toString(ctHash);
+    uint256 length = bytes(hashStr).length;
+    if (length <= 6) {
+      return hashStr;
+    }
+
+    bool stored = inMockStorage[ctHash];
+    uint256 value = mockStorage[ctHash];
+    bool isBool = getIsBoolTypeFromHash(ctHash);
+
+    string memory valueString = isBool ? (value == 1 ? 'true' : 'false') : Strings.toString(value);
+
+    string memory truncated = string.concat(
+      getUtypeStringFromHash(ctHash),
+      '(',
+      sliceString(hashStr, 0, 4),
+      '..',
+      sliceString(hashStr, length - 4, 4),
+      ')[',
+      stored ? valueString : 'EMPTY',
+      ']'
+    );
+
+    return truncated;
+  }
+
+  string constant LOG_PREFIX = unicode'├ ';
+  string constant LOG_DIVIDER = unicode' | ';
+
+  function padRight(string memory input, uint256 length, bytes1 padChar) internal pure returns (string memory) {
+    bytes memory inputBytes = bytes(input);
+    if (inputBytes.length >= length) return input;
+
+    bytes memory padded = new bytes(length);
+    uint256 i = 0;
+    for (; i < inputBytes.length; i++) {
+      padded[i] = inputBytes[i];
+    }
+    for (; i < length; i++) {
+      padded[i] = padChar;
+    }
+    return string(padded);
+  }
+
+  function logOperation(string memory operation, string memory inputs, string memory output) internal view {
+    if (logOps)
+      console.log(string.concat(LOG_PREFIX, padRight(operation, 16, ' '), LOG_DIVIDER, inputs, '  =>  ', output));
+  }
+
+  function logAllow(string memory operation, uint256 ctHash, address account) internal view {
+    if (logOps)
+      console.log(
+        string.concat(
+          LOG_PREFIX,
+          padRight(operation, 16, ' '),
+          LOG_DIVIDER,
+          logCtHash(ctHash),
+          ' -> ',
+          Strings.toHexString(account)
+        )
+      );
+  }
+
+  // Storage functions
+
+  function _set(uint256 ctHash, uint256 value, bool log) internal {
+    mockStorage[ctHash] = value;
+    inMockStorage[ctHash] = true;
+
+    if (log) logOperation('set', '', logCtHash(ctHash));
+  }
+
+  function _set(uint256 ctHash, uint256 value) internal {
+    uint256 mask = getUtypeMask(ctHash);
+    _set(ctHash, value & mask, false);
+  }
+
+  function _set(uint256 ctHash, bool value) internal {
+    _set(ctHash, value ? 1 : 0);
+  }
+
+  function _get(uint256 ctHash) internal view returns (uint256) {
+    if (!inMockStorage[ctHash]) revert InputNotInMockStorage(ctHash);
+
+    uint256 mask = getUtypeMask(ctHash);
+    return mockStorage[ctHash] & mask;
+  }
+
+  // Public functions
+
+  function MOCK_setInEuintKey(uint256 ctHash, uint256 value) public {
+    _set(ctHash, value);
+  }
+
+  // Mock Log
+
+  function MOCK_logAllow(string memory operation, uint256 ctHash, address account) public view {
+    logAllow(operation, ctHash, account);
+  }
+
+  // Mock functions
+
+  function MOCK_verifyKeyInStorage(uint256 ctHash) internal view {
+    if (!inMockStorage[ctHash]) revert InputNotInMockStorage(ctHash);
+  }
+
+  function MOCK_unaryOperation(uint256 ctHash, string memory operation, uint256 input) internal {
+    if (opIs(operation, FunctionId.random)) {
+      _set(ctHash, uint256(blockhash(block.number - 1)));
+      logOperation('FHE.random', '', logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.cast)) {
+      _set(ctHash, _get(input));
+      logOperation('FHE.cast', logCtHash(input), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.not)) {
+      bool inputIsTruthy = _get(input) == 1;
+      _set(ctHash, !inputIsTruthy);
+      logOperation('FHE.not', logCtHash(input), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.square)) {
+      unchecked {
+        _set(ctHash, _get(input) * _get(input));
+      }
+      logOperation('FHE.square', string.concat(logCtHash(input), ' * ', logCtHash(input)), logCtHash(ctHash));
+      return;
+    }
+    revert InvalidUnaryOperation(operation);
+  }
+
+  function MOCK_twoInputOperation(uint256 ctHash, string memory operation, uint256 input1, uint256 input2) internal {
+    if (opIs(operation, FunctionId.sub)) {
+      unchecked {
+        _set(ctHash, _get(input1) - _get(input2));
+      }
+      logOperation('FHE.sub', string.concat(logCtHash(input1), ' - ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.add)) {
+      unchecked {
+        _set(ctHash, _get(input1) + _get(input2));
+      }
+      logOperation('FHE.add', string.concat(logCtHash(input1), ' + ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.xor)) {
+      unchecked {
+        _set(ctHash, _get(input1) ^ _get(input2));
+      }
+      logOperation('FHE.xor', string.concat(logCtHash(input1), ' ^ ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.and)) {
+      unchecked {
+        _set(ctHash, _get(input1) & _get(input2));
+      }
+      logOperation('FHE.and', string.concat(logCtHash(input1), ' & ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.or)) {
+      unchecked {
+        _set(ctHash, _get(input1) | _get(input2));
+      }
+      logOperation('FHE.or', string.concat(logCtHash(input1), ' | ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.div)) {
+      uint256 cleartext2 = _get(input2);
+      if (cleartext2 == 0) {
+        _set(ctHash, type(uint256).max);
+      } else {
+        unchecked {
+          _set(ctHash, _get(input1) / cleartext2);
+        }
+      }
+      logOperation('FHE.div', string.concat(logCtHash(input1), ' / ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.rem)) {
+      unchecked {
+        _set(ctHash, _get(input1) % _get(input2));
+      }
+      logOperation('FHE.rem', string.concat(logCtHash(input1), ' % ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.mul)) {
+      unchecked {
+        _set(ctHash, _get(input1) * _get(input2));
+      }
+      logOperation('FHE.mul', string.concat(logCtHash(input1), ' * ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.shl)) {
+      unchecked {
+        _set(ctHash, _get(input1) << _get(input2));
+      }
+      logOperation('FHE.shl', string.concat(logCtHash(input1), ' << ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.shr)) {
+      unchecked {
+        _set(ctHash, _get(input1) >> _get(input2));
+      }
+      logOperation('FHE.shr', string.concat(logCtHash(input1), ' >> ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.gte)) {
+      _set(ctHash, _get(input1) >= _get(input2));
+      logOperation('FHE.gte', string.concat(logCtHash(input1), ' >= ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.lte)) {
+      _set(ctHash, _get(input1) <= _get(input2));
+      logOperation('FHE.lte', string.concat(logCtHash(input1), ' <= ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.lt)) {
+      _set(ctHash, _get(input1) < _get(input2));
+      logOperation('FHE.lt', string.concat(logCtHash(input1), ' < ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.gt)) {
+      _set(ctHash, _get(input1) > _get(input2));
+      logOperation('FHE.gt', string.concat(logCtHash(input1), ' > ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.min)) {
+      uint256 min;
+      unchecked {
+        min = _get(input1) < _get(input2) ? _get(input1) : _get(input2);
+      }
+      _set(ctHash, min);
+
+      logOperation(
+        'FHE.min',
+        string.concat('min(', logCtHash(input1), ', ', logCtHash(input2), ')'),
+        logCtHash(ctHash)
+      );
+      return;
+    }
+    if (opIs(operation, FunctionId.max)) {
+      uint256 max;
+      unchecked {
+        max = _get(input1) > _get(input2) ? _get(input1) : _get(input2);
+      }
+      _set(ctHash, max);
+
+      logOperation(
+        'FHE.max',
+        string.concat('max(', logCtHash(input1), ', ', logCtHash(input2), ')'),
+        logCtHash(ctHash)
+      );
+      return;
+    }
+    if (opIs(operation, FunctionId.eq)) {
+      _set(ctHash, _get(input1) == _get(input2));
+
+      logOperation('FHE.eq', string.concat(logCtHash(input1), ' == ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.ne)) {
+      _set(ctHash, _get(input1) != _get(input2));
+
+      logOperation('FHE.ne', string.concat(logCtHash(input1), ' != ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.rol)) {
+      unchecked {
+        _set(ctHash, _get(input1) << _get(input2));
+      }
+
+      logOperation('FHE.rol', string.concat(logCtHash(input1), ' << ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    if (opIs(operation, FunctionId.ror)) {
+      unchecked {
+        _set(ctHash, _get(input1) >> _get(input2));
+      }
+
+      logOperation('FHE.ror', string.concat(logCtHash(input1), ' >> ', logCtHash(input2)), logCtHash(ctHash));
+      return;
+    }
+    revert InvalidTwoInputOperation(operation);
+  }
+
+  function MOCK_threeInputOperation(
+    uint256 ctHash,
+    string memory operation,
+    uint256 input1,
+    uint256 input2,
+    uint256 input3
+  ) internal {
+    if (opIs(operation, FunctionId.trivialEncrypt)) {
+      _set(ctHash, input1);
+
+      logOperation(
+        string.concat('FHE.asE', removeFirstLetter(getUtypeStringFromHash(ctHash))),
+        string.concat(removeFirstLetter(getUtypeStringFromHash(ctHash)), '(', Strings.toString(input1), ')'),
+        logCtHash(ctHash)
+      );
+      return;
+    }
+    if (opIs(operation, FunctionId.select)) {
+      _set(ctHash, _get(input1) == 1 ? _get(input2) : _get(input3));
+
+      logOperation(
+        'FHE.select',
+        string.concat(logCtHash(input1), ' ? ', logCtHash(input2), ' : ', logCtHash(input3)),
+        logCtHash(ctHash)
+      );
+      return;
+    }
+    revert InvalidThreeInputOperation(operation);
+  }
+}