From 738a5e0780c1b6f00c406c5be1e9c01a882deeec Mon Sep 17 00:00:00 2001 From: rajatlko13 Date: Tue, 24 Sep 2024 17:48:50 +0530 Subject: [PATCH] secret store contract splitted --- .../{SecretStorage.sol => EnclaveStore.sol} | 560 ++--------------- contracts/secret-storage/SecretStore.sol | 580 ++++++++++++++++++ 2 files changed, 646 insertions(+), 494 deletions(-) rename contracts/secret-storage/{SecretStorage.sol => EnclaveStore.sol} (50%) create mode 100644 contracts/secret-storage/SecretStore.sol diff --git a/contracts/secret-storage/SecretStorage.sol b/contracts/secret-storage/EnclaveStore.sol similarity index 50% rename from contracts/secret-storage/SecretStorage.sol rename to contracts/secret-storage/EnclaveStore.sol index 54b6e82..5e97ce5 100644 --- a/contracts/secret-storage/SecretStorage.sol +++ b/contracts/secret-storage/EnclaveStore.sol @@ -12,13 +12,14 @@ import "@openzeppelin/contracts/utils/cryptography/ECDSA.sol"; import "../AttestationAutherUpgradeable.sol"; import "../serverless-v2/tree/TreeMapUpgradeable.sol"; import "../interfaces/IAttestationVerifier.sol"; +import "./SecretStore.sol"; /** - * @title SecretStorage Contract + * @title EnclaveStore Contract * @notice Manages the registration, staking, and job assignment of execution nodes. * @dev This contract is upgradeable and uses the UUPS (Universal Upgradeable Proxy Standard) pattern. */ -contract SecretStorage is +contract EnclaveStore is Initializable, // initializer ContextUpgradeable, // _msgSender, _msgData ERC165Upgradeable, // supportsInterface @@ -31,11 +32,9 @@ contract SecretStorage is using ECDSA for bytes32; /// @notice Thrown when the provided ERC20 token address is zero. - error SecretStorageZeroAddressStakingToken(); - /// @notice Thrown when the provided ERC20 token address is zero. - error SecretStorageZeroAddressUsdcToken(); + error EnclaveStoreZeroAddressStakingToken(); /// @notice Thrown when the provided minimum stake amount is zero. - error SecretStorageZeroMinStakeAmount(); + error EnclaveStoreZeroMinStakeAmount(); /** * @dev Initializes the logic contract without any admins, safeguarding against takeover. @@ -51,43 +50,22 @@ contract SecretStorage is IAttestationVerifier attestationVerifier, uint256 maxAge, IERC20 _stakingToken, - IERC20 _usdcToken, uint256 _minStakeAmount, uint256 _slashPercentInBips, uint256 _slashMaxBips, - uint256 _noOfNodesToSelect, - uint8 _env, - uint256 _globalMaxStoreSize, - uint256 _globalMinStoreDuration, - uint256 _globalMaxStoreDuration, - uint256 _acknowledgementTimeout, - uint256 _markAliveTimeout, - uint256 _storageFeeRate, - address _stakingPaymentPool + uint8 _env ) AttestationAutherUpgradeable(attestationVerifier, maxAge) { _disableInitializers(); - if (address(_stakingToken) == address(0)) revert SecretStorageZeroAddressStakingToken(); - if (address(_usdcToken) == address(0)) revert SecretStorageZeroAddressUsdcToken(); - if (_minStakeAmount == 0) revert SecretStorageZeroMinStakeAmount(); + if (address(_stakingToken) == address(0)) revert EnclaveStoreZeroAddressStakingToken(); + if (_minStakeAmount == 0) revert EnclaveStoreZeroMinStakeAmount(); STAKING_TOKEN = _stakingToken; - USDC_TOKEN = _usdcToken; MIN_STAKE_AMOUNT = _minStakeAmount; SLASH_PERCENT_IN_BIPS = _slashPercentInBips; SLASH_MAX_BIPS = _slashMaxBips; - NO_OF_NODES_TO_SELECT = _noOfNodesToSelect; ENV = _env; - - // TODO: add checks - GLOBAL_MAX_STORE_SIZE = _globalMaxStoreSize; - GLOBAL_MIN_STORE_DURATION = _globalMinStoreDuration; - GLOBAL_MAX_STORE_DURATION = _globalMaxStoreDuration; - ACKNOWLEDGEMENT_TIMEOUT = _acknowledgementTimeout; - MARK_ALIVE_TIMEOUT = _markAliveTimeout; - STORAGE_FEE_RATE = _storageFeeRate; - STAKING_PAYMENT_POOL = _stakingPaymentPool; } //-------------------------------- Overrides start --------------------------------// @@ -107,7 +85,7 @@ contract SecretStorage is //-------------------------------- Initializer start --------------------------------// /// @notice Thrown when the provided admin address is zero. - error SecretStorageZeroAddressAdmin(); + error EnclaveStoreZeroAddressAdmin(); /** * @dev Initializes the contract with the given admin and enclave images. @@ -115,7 +93,7 @@ contract SecretStorage is * @param _images Array of enclave images to initialize. */ function initialize(address _admin, EnclaveImage[] memory _images) public initializer { - if (_admin == address(0)) revert SecretStorageZeroAddressAdmin(); + if (_admin == address(0)) revert EnclaveStoreZeroAddressAdmin(); __Context_init_unchained(); __ERC165_init_unchained(); @@ -133,9 +111,6 @@ contract SecretStorage is /// @custom:oz-upgrades-unsafe-allow state-variable-immutable IERC20 public immutable STAKING_TOKEN; - /// @custom:oz-upgrades-unsafe-allow state-variable-immutable - IERC20 public immutable USDC_TOKEN; - /// @custom:oz-upgrades-unsafe-allow state-variable-immutable uint256 public immutable MIN_STAKE_AMOUNT; @@ -150,36 +125,12 @@ contract SecretStorage is /// @custom:oz-upgrades-unsafe-allow state-variable-immutable uint8 public immutable ENV; - /// @custom:oz-upgrades-unsafe-allow state-variable-immutable - uint256 public immutable NO_OF_NODES_TO_SELECT; - - /// @custom:oz-upgrades-unsafe-allow state-variable-immutable - uint256 public immutable GLOBAL_MAX_STORE_SIZE; - - /// @custom:oz-upgrades-unsafe-allow state-variable-immutable - uint256 public immutable GLOBAL_MIN_STORE_DURATION; - - /// @custom:oz-upgrades-unsafe-allow state-variable-immutable - uint256 public immutable GLOBAL_MAX_STORE_DURATION; - - /// @custom:oz-upgrades-unsafe-allow state-variable-immutable - uint256 public immutable ACKNOWLEDGEMENT_TIMEOUT; - - /// @custom:oz-upgrades-unsafe-allow state-variable-immutable - uint256 public immutable MARK_ALIVE_TIMEOUT; - - /// @custom:oz-upgrades-unsafe-allow state-variable-immutable - uint256 public immutable STORAGE_FEE_RATE; - - /// @custom:oz-upgrades-unsafe-allow state-variable-immutable - address public immutable STAKING_PAYMENT_POOL; - /// @notice enclave stake amount will be divided by 10^18 before adding to the tree uint256 public constant STAKE_ADJUSTMENT_FACTOR = 1e18; - bytes32 public constant JOBS_ROLE = keccak256("JOBS_ROLE"); + bytes32 public constant SECRET_STORE_ROLE = keccak256("SECRET_STORE_ROLE"); - //-------------------------------- Enclave storage start --------------------------------// + //-------------------------------- EnclaveStore start --------------------------------// modifier isValidEnclaveStoreOwner(address _enclaveAddress) { _isValidEnclaveStoreOwner(_enclaveAddress); @@ -188,7 +139,7 @@ contract SecretStorage is function _isValidEnclaveStoreOwner(address _enclaveAddress) internal view { if (enclaveStorage[_enclaveAddress].owner != _msgSender()) - revert SecretStorageInvalidEnclaveOwner(); + revert EnclaveStoreInvalidEnclaveOwner(); } struct EnclaveStorage { @@ -206,7 +157,7 @@ contract SecretStorage is keccak256( abi.encode( keccak256("EIP712Domain(string name,string version)"), - keccak256("marlin.oyster.SecretStore"), + keccak256("marlin.oyster.EnclaveStore"), keccak256("1") ) ); @@ -249,21 +200,21 @@ contract SecretStorage is event EnclaveStoreStakeRemoved(address indexed enclaveAddress, uint256 removedAmount); /// @notice Thrown when the signature timestamp has expired. - error SecretStorageSignatureTooOld(); + error EnclaveStoreSignatureTooOld(); /// @notice Thrown when the signer of the registration data is invalid. - error SecretStorageInvalidSigner(); + error EnclaveStoreInvalidSigner(); /// @notice Thrown when attempting to register an enclave that already exists. - error SecretStorageEnclaveAlreadyExists(); + error EnclaveStoreEnclaveAlreadyExists(); /// @notice Thrown when attempting to drain an enclave that is already draining. - error SecretStorageEnclaveAlreadyDraining(); + error EnclaveStoreEnclaveAlreadyDraining(); /// @notice Thrown when attempting to revive an enclave that is not draining. - error SecretStorageEnclaveAlreadyRevived(); + error EnclaveStoreEnclaveAlreadyRevived(); /// @notice Thrown when attempting to deregister or remove stake from an enclave that is not draining. - error SecretStorageEnclaveNotDraining(); + error EnclaveStoreEnclaveNotDraining(); /// @notice Thrown when attempting to deregister or remove stake from an enclave that has pending jobs. - error SecretStorageEnclaveNotEmpty(); + error EnclaveStoreEnclaveNotEmpty(); /// @notice Thrown when the provided enclave owner does not match the stored owner. - error SecretStorageInvalidEnclaveOwner(); + error EnclaveStoreInvalidEnclaveOwner(); //-------------------------------- Admin methods start --------------------------------// @@ -307,7 +258,7 @@ contract SecretStorage is ) internal { address enclaveAddress = _pubKeyToAddress(_attestation.enclavePubKey); if (enclaveStorage[enclaveAddress].owner != address(0)) - revert SecretStorageEnclaveAlreadyExists(); + revert EnclaveStoreEnclaveAlreadyExists(); // attestation verification _verifyEnclaveKey(_attestationSignature, _attestation); @@ -331,18 +282,14 @@ contract SecretStorage is uint256 _signTimestamp, bytes memory _signature ) internal view { - _checkSignValidity(_signTimestamp); + if (block.timestamp > _signTimestamp + ATTESTATION_MAX_AGE) + revert EnclaveStoreSignatureTooOld(); bytes32 hashStruct = keccak256(abi.encode(REGISTER_TYPEHASH, _owner, _storageCapacity, _signTimestamp)); bytes32 digest = keccak256(abi.encodePacked("\x19\x01", DOMAIN_SEPARATOR, hashStruct)); address signer = digest.recover(_signature); - if (signer != _enclaveAddress) revert SecretStorageInvalidSigner(); - } - - function _checkSignValidity(uint256 _signTimestamp) internal view { - if (block.timestamp > _signTimestamp + ATTESTATION_MAX_AGE) - revert SecretStorageSignatureTooOld(); + if (signer != _enclaveAddress) revert EnclaveStoreInvalidSigner(); } function _register( @@ -357,7 +304,7 @@ contract SecretStorage is } function _drainEnclaveStore(address _enclaveAddress) internal { - if (enclaveStorage[_enclaveAddress].draining) revert SecretStorageEnclaveAlreadyDraining(); + if (enclaveStorage[_enclaveAddress].draining) revert EnclaveStoreEnclaveAlreadyDraining(); enclaveStorage[_enclaveAddress].draining = true; @@ -369,7 +316,7 @@ contract SecretStorage is function _reviveEnclaveStore(address _enclaveAddress) internal { EnclaveStorage memory enclaveStoreNode = enclaveStorage[_enclaveAddress]; - if (!enclaveStoreNode.draining) revert SecretStorageEnclaveAlreadyRevived(); + if (!enclaveStoreNode.draining) revert EnclaveStoreEnclaveAlreadyRevived(); enclaveStorage[_enclaveAddress].draining = false; @@ -382,8 +329,8 @@ contract SecretStorage is } function _deregisterEnclaveStore(address _enclaveAddress) internal { - if (!enclaveStorage[_enclaveAddress].draining) revert SecretStorageEnclaveNotDraining(); - if (enclaveStorage[_enclaveAddress].storageOccupied != 0) revert SecretStorageEnclaveNotEmpty(); + if (!enclaveStorage[_enclaveAddress].draining) revert EnclaveStoreEnclaveNotDraining(); + if (enclaveStorage[_enclaveAddress].storageOccupied != 0) revert EnclaveStoreEnclaveNotEmpty(); _removeStake(_enclaveAddress, enclaveStorage[_enclaveAddress].stakeAmount); @@ -410,8 +357,8 @@ contract SecretStorage is } function _removeEnclaveStoreStake(uint256 _amount, address _enclaveAddress) internal { - if (!enclaveStorage[_enclaveAddress].draining) revert SecretStorageEnclaveNotDraining(); - if (enclaveStorage[_enclaveAddress].storageOccupied != 0) revert SecretStorageEnclaveNotEmpty(); + if (!enclaveStorage[_enclaveAddress].draining) revert EnclaveStoreEnclaveNotDraining(); + if (enclaveStorage[_enclaveAddress].storageOccupied != 0) revert EnclaveStoreEnclaveNotEmpty(); _removeStake(_enclaveAddress, _amount); } @@ -517,379 +464,32 @@ contract SecretStorage is _removeEnclaveStoreStake(_amount, _enclaveAddress); } - //-------------------------------- external functions end ----------------------------------// - - //--------------------------------------- Enclave storage end -----------------------------------------// - - //--------------------------------------- User storage start -----------------------------------------// - - struct SelectedEnclave { - address enclaveAddress; - bool hasAcknowledgedStore; - uint256 lastAliveTimestamp; - } - - struct UserStorage { - address owner; - uint256 sizeLimit; - uint256 usdcDeposit; - uint256 startTimestamp; - uint256 endTimestamp; - SelectedEnclave[] selectedEnclaves; - } - - // secretId => user store data - mapping(uint256 => UserStorage) public userStorage; - - uint256 public secretId; - - event SecretStoreCreated( - uint256 indexed secretId, - address indexed owner, - uint256 sizeLimit, - uint256 endTimestamp, - uint256 usdcDeposit - ); - - event EnclaveAcknowledgedStore( - uint256 indexed secretId, - address indexed enclaveAddress - ); - - event EnclaveAcknowledgementFailed( - uint256 indexed secretId - ); - - event EnclaveStoreAlive( - uint256 indexed secretId, - address indexed enclaveAddress - ); - - event EnclaveStoreDead( - uint256 indexed secretId, - address indexed prevEnclaveAddress, - address indexed newEnclaveAddress - ); - - event SecretStoreResourceUnavailable( - uint256 indexed secretId - ); - - event SecretStoreEndTimestampUpdated( - uint256 indexed secretId, - uint256 endTimestamp - ); - - event SecretStoreTerminated( - uint256 indexed secretId - ); - - error SecretStorageInsufficientUsdcDeposit(); - error SecretStorageInvalidSizeLimit(); - error SecretStorageInvalidEndTimestamp(); - error SecretStorageUnavailableResources(); - error SecretStoreAcknowledgementTimeOver(); - error SecretStoreAcknowledgementTimeoutPending(); - error SecretStoreAcknowledgedAlready(); - error SecretStoreMarkAliveTimeoutOver(); - error SecretStoreUnacknowledged(); - error SecretStorageEnclaveNotFound(); - error SecretStorageNotUserStoreOwner(); - error SecretStorageAlreadyTerminated(); - - //-------------------------------- internal functions start ----------------------------------// - - function _createSecretStore( - uint256 _sizeLimit, - uint256 _endTimestamp, - uint256 _usdcDeposit, - address _owner - ) internal { - if(_sizeLimit == 0 || _sizeLimit > GLOBAL_MAX_STORE_SIZE) - revert SecretStorageInvalidSizeLimit(); - - if ((_endTimestamp < block.timestamp + GLOBAL_MIN_STORE_DURATION) || (_endTimestamp > block.timestamp + GLOBAL_MAX_STORE_DURATION)) - revert SecretStorageInvalidEndTimestamp(); - - // TODO: how to calculate usdcDeposit - uint256 minUsdcDeposit = (_endTimestamp - block.timestamp) * _sizeLimit * STORAGE_FEE_RATE; - _checkUsdcDeposit(_usdcDeposit, minUsdcDeposit); - - USDC_TOKEN.safeTransferFrom(_owner, address(this), _usdcDeposit); - - SelectedEnclave[] memory selectedEnclaves = _selectEnclaves(NO_OF_NODES_TO_SELECT, _sizeLimit); - if (selectedEnclaves.length < NO_OF_NODES_TO_SELECT) - revert SecretStorageUnavailableResources(); - - userStorage[++secretId] = UserStorage({ - owner: _owner, - sizeLimit: _sizeLimit, - usdcDeposit: _usdcDeposit, - startTimestamp: block.timestamp, - endTimestamp: _endTimestamp, - selectedEnclaves: selectedEnclaves - }); - - emit SecretStoreCreated(secretId, _owner, _sizeLimit, _endTimestamp, _usdcDeposit); - } - - function _checkUsdcDeposit( - uint256 _usdcDeposit, - uint256 _minUsdcDeposit - ) internal pure { - if(_usdcDeposit < _minUsdcDeposit) - revert SecretStorageInsufficientUsdcDeposit(); - } - - function _acknowledgeStore( - uint256 _secretId, - uint256 _signTimestamp, - bytes memory _signature - ) internal { - if(block.timestamp > userStorage[_secretId].startTimestamp + ACKNOWLEDGEMENT_TIMEOUT) - revert SecretStoreAcknowledgementTimeOver(); - - address enclaveAddress = _verifyAcknowledgementSign(_secretId, _signTimestamp, _signature); - - uint256 enclaveIndex = _getSelectedEnclaveIndex(_secretId, enclaveAddress); - userStorage[_secretId].selectedEnclaves[enclaveIndex].hasAcknowledgedStore = true; - userStorage[_secretId].selectedEnclaves[enclaveIndex].lastAliveTimestamp = _signTimestamp; - - emit EnclaveAcknowledgedStore(_secretId, enclaveAddress); - } - - function _verifyAcknowledgementSign( - uint256 _secretId, - uint256 _signTimestamp, - bytes memory _signature - ) internal view returns(address signer) { - _checkSignValidity(_signTimestamp); - - bytes32 hashStruct = keccak256(abi.encode(ACKNOWLEDGE_TYPEHASH, _secretId, _signTimestamp)); - bytes32 digest = keccak256(abi.encodePacked("\x19\x01", DOMAIN_SEPARATOR, hashStruct)); - signer = digest.recover(_signature); - - _allowOnlyVerified(signer); - } - - function _getSelectedEnclaveIndex( - uint256 _secretId, - address _enclaveAddress - ) internal view returns (uint256) { - uint256 len = userStorage[_secretId].selectedEnclaves.length; - for (uint256 index = 0; index < len; index++) { - if(userStorage[_secretId].selectedEnclaves[index].enclaveAddress == _enclaveAddress) - return index; - } - revert SecretStorageEnclaveNotFound(); - } - - function _acknowledgeStoreFailed( - uint256 _secretId - ) internal { - if(block.timestamp <= userStorage[_secretId].startTimestamp + ACKNOWLEDGEMENT_TIMEOUT) - revert SecretStoreAcknowledgementTimeoutPending(); - - bool ackFailed; - uint256 len = userStorage[_secretId].selectedEnclaves.length; - for (uint256 index = 0; index < len; index++) { - if(!userStorage[_secretId].selectedEnclaves[index].hasAcknowledgedStore) { - ackFailed = true; - break; - } - } - - if(!ackFailed) - revert SecretStoreAcknowledgedAlready(); - - address owner = userStorage[_secretId].owner; - uint256 usdcDeposit = userStorage[_secretId].usdcDeposit; - delete userStorage[_secretId]; - - USDC_TOKEN.safeTransfer(owner, usdcDeposit); - - emit EnclaveAcknowledgementFailed(_secretId); - } - - function _markStoreAlive( - uint256 _secretId, - uint256 _signTimestamp, - bytes memory _signature, - address _owner - ) internal { - address enclaveAddress = _verifyStoreAliveSign(_secretId, _signTimestamp, _signature); - - uint256 enclaveIndex = _getSelectedEnclaveIndex(_secretId, enclaveAddress); - if(!userStorage[_secretId].selectedEnclaves[enclaveIndex].hasAcknowledgedStore) - revert SecretStoreUnacknowledged(); - if(block.timestamp > userStorage[_secretId].selectedEnclaves[enclaveIndex].lastAliveTimestamp + MARK_ALIVE_TIMEOUT) - revert SecretStoreMarkAliveTimeoutOver(); - - uint256 endTime; - if(block.timestamp > userStorage[_secretId].endTimestamp) - endTime = userStorage[_secretId].endTimestamp; - else - endTime = block.timestamp; - - uint256 usdcPayment = ((endTime - userStorage[_secretId].selectedEnclaves[enclaveIndex].lastAliveTimestamp) * userStorage[_secretId].sizeLimit * STORAGE_FEE_RATE) / NO_OF_NODES_TO_SELECT; - userStorage[_secretId].usdcDeposit -= usdcPayment; - userStorage[_secretId].selectedEnclaves[enclaveIndex].lastAliveTimestamp = _signTimestamp; - - USDC_TOKEN.safeTransfer(_owner, usdcPayment); - - // TODO: delete from selectedNodes array and refund remaining usdc - if(block.timestamp > userStorage[_secretId].endTimestamp) { - _removeSelectedEnclave(_secretId, enclaveIndex); - if(userStorage[_secretId].selectedEnclaves.length == 0) - _refundExcessDepositAndRemoveStore(_secretId); - } - - emit EnclaveStoreAlive(_secretId, enclaveAddress); - } - - function _verifyStoreAliveSign( - uint256 _secretId, - uint256 _signTimestamp, - bytes memory _signature - ) internal view returns(address signer) { - _checkSignValidity(_signTimestamp); - - bytes32 hashStruct = keccak256(abi.encode(ALIVE_TYPEHASH, _secretId, _signTimestamp)); - bytes32 digest = keccak256(abi.encodePacked("\x19\x01", DOMAIN_SEPARATOR, hashStruct)); - signer = digest.recover(_signature); - - _allowOnlyVerified(signer); - } - - function _markStoreDead( - uint256 _secretId - ) internal { - for (uint256 index = 0; index < userStorage[_secretId].selectedEnclaves.length; ) { - bool isArrayLengthReduced = _markEnclaveDead(_secretId, userStorage[_secretId].selectedEnclaves[index].enclaveAddress); - if(!isArrayLengthReduced) - ++index; - } - - // TODO: delete data and refund remaining - if(block.timestamp > userStorage[_secretId].endTimestamp && userStorage[_secretId].selectedEnclaves.length == 0) - _refundExcessDepositAndRemoveStore(_secretId); - } - - function _markEnclaveDead( - uint256 _secretId, - address _enclaveAddress - ) internal returns (bool isArrayLengthReduced) { - uint256 enclaveIndex = _getSelectedEnclaveIndex(_secretId, _enclaveAddress); - if(!userStorage[_secretId].selectedEnclaves[enclaveIndex].hasAcknowledgedStore) - return isArrayLengthReduced; - - if(block.timestamp <= userStorage[_secretId].selectedEnclaves[enclaveIndex].lastAliveTimestamp + MARK_ALIVE_TIMEOUT) - return isArrayLengthReduced; - - if(block.timestamp > userStorage[_secretId].endTimestamp) { - isArrayLengthReduced = true; - _removeSelectedEnclave(_secretId, enclaveIndex); - } - else { - SelectedEnclave[] memory selectedEnclaves = _selectEnclaves(1, userStorage[_secretId].sizeLimit); - if (selectedEnclaves.length == 0) { - isArrayLengthReduced = true; - _removeSelectedEnclave(_secretId, enclaveIndex); - emit SecretStoreResourceUnavailable(_secretId); - } - else { - userStorage[_secretId].selectedEnclaves[enclaveIndex] = SelectedEnclave({ - enclaveAddress: selectedEnclaves[0].enclaveAddress, - hasAcknowledgedStore: false, - lastAliveTimestamp: 0 - }); - } - - emit EnclaveStoreDead( - _secretId, - _enclaveAddress, - selectedEnclaves.length != 0 ? selectedEnclaves[0].enclaveAddress : address(0) - ); - } - - // TODO: slash prev enclave(who's recipient) - _slashEnclave(_enclaveAddress, STAKING_PAYMENT_POOL); - _releaseEnclave(_enclaveAddress, userStorage[_secretId].sizeLimit); - } - - function _removeSelectedEnclave( - uint256 _secretId, - uint256 _index - ) internal { - uint256 len = userStorage[_secretId].selectedEnclaves.length; - if(_index != len - 1) - userStorage[_secretId].selectedEnclaves[_index] = userStorage[_secretId].selectedEnclaves[len - 1]; - userStorage[_secretId].selectedEnclaves.pop(); - } - - function _refundExcessDepositAndRemoveStore( - uint256 _secretId - ) internal { - address owner = userStorage[_secretId].owner; - uint256 remainingDeposit = userStorage[_secretId].usdcDeposit; - delete userStorage[_secretId]; - USDC_TOKEN.safeTransfer(owner, remainingDeposit); + /** + * @notice Allows only verified addresses to perform certain actions. + * @param _signer The address to be verified. + */ + function allowOnlyVerified(address _signer) external view { + _allowOnlyVerified(_signer); } - function _updateSecretStoreEndTimestamp( - uint256 _secretId, - uint256 _endTimestamp, - uint256 _usdcDeposit - ) internal { - if(userStorage[_secretId].owner != _msgSender()) - revert SecretStorageNotUserStoreOwner(); - - if(_endTimestamp < block.timestamp) - revert SecretStorageInvalidEndTimestamp(); - - uint256 currentEndTimestamp = userStorage[_secretId].endTimestamp; - if(block.timestamp > currentEndTimestamp) - revert SecretStorageAlreadyTerminated(); - - if(_endTimestamp > currentEndTimestamp) { - USDC_TOKEN.safeTransferFrom(_msgSender(), address(this), _usdcDeposit); - userStorage[_secretId].usdcDeposit += _usdcDeposit; - - uint256 addedDuration = _endTimestamp - currentEndTimestamp; - uint256 minUsdcDeposit = addedDuration * userStorage[_secretId].sizeLimit * STORAGE_FEE_RATE; - _checkUsdcDeposit(_usdcDeposit, minUsdcDeposit); - } - else { - uint256 removedDuration = currentEndTimestamp - _endTimestamp; - uint256 usdcRefund = removedDuration * userStorage[_secretId].sizeLimit * STORAGE_FEE_RATE; - - userStorage[_secretId].usdcDeposit -= usdcRefund; - USDC_TOKEN.safeTransfer(_msgSender(), usdcRefund); - } + //-------------------------------- external functions end ----------------------------------// - userStorage[_secretId].endTimestamp = _endTimestamp; + //--------------------------------------- EnclaveStore end -----------------------------------------// - emit SecretStoreEndTimestampUpdated(_secretId, _endTimestamp); - } + //----------------------------- SecretStoreRole functions start ---------------------------------// - function _terminateSecretStore( - uint256 _secretId - ) internal { - _updateSecretStoreEndTimestamp(_secretId, block.timestamp, 0); - - emit SecretStoreTerminated(_secretId); - } + //-------------------------------- internal functions start ----------------------------------// function _selectEnclaves( uint256 _noOfNodesToSelect, uint256 _sizeLimit - ) internal returns (SelectedEnclave[] memory selectedEnclaves) { + ) internal returns (SecretStore.SelectedEnclave[] memory selectedEnclaves) { address[] memory selectedNodes = _selectNodes(_noOfNodesToSelect); for (uint256 index = 0; index < selectedNodes.length; index++) { address enclaveAddress = selectedNodes[index]; enclaveStorage[enclaveAddress].storageOccupied += _sizeLimit; - SelectedEnclave memory selectedEnclave; + SecretStore.SelectedEnclave memory selectedEnclave; selectedEnclave.enclaveAddress = enclaveAddress; selectedEnclaves[index] = selectedEnclave; @@ -904,7 +504,10 @@ contract SecretStorage is selectedNodes = _selectN(ENV, randomizer, _noOfNodesToSelect); } - function _slashEnclave(address _enclaveAddress, address _recipient) internal returns (uint256) { + function _slashEnclave( + address _enclaveAddress, + address _recipient + ) internal returns (uint256) { uint256 totalComp = (enclaveStorage[_enclaveAddress].stakeAmount * SLASH_PERCENT_IN_BIPS) / SLASH_MAX_BIPS; enclaveStorage[_enclaveAddress].stakeAmount -= totalComp; @@ -930,61 +533,30 @@ contract SecretStorage is //-------------------------------- internal functions end ----------------------------------// - //------------------------------- external functions start ----------------------------------// - - function createSecretStore( - uint256 _sizeLimit, - uint256 _endTimestamp, - uint256 _usdcDeposit - ) external { - _createSecretStore(_sizeLimit, _endTimestamp, _usdcDeposit, _msgSender()); - } - - function acknowledgeStore( - uint256 _secretId, - uint256 _signTimestamp, - bytes memory _signature - ) external { - _acknowledgeStore(_secretId, _signTimestamp, _signature); - } - - function acknowledgeStoreFailed( - uint256 _secretId - ) external { - _acknowledgeStoreFailed(_secretId); - } - - function markStoreAlive( - uint256 _secretId, - uint256 _signTimestamp, - bytes memory _signature - ) external { - _markStoreAlive(_secretId, _signTimestamp, _signature, _msgSender()); - } + //-------------------------------- external functions start ----------------------------------// - function markStoreDead( - uint256 _secretId - ) external { - _markStoreDead(_secretId); + function selectEnclaves( + uint256 _noOfNodesToSelect, + uint256 _sizeLimit + ) external onlyRole(SECRET_STORE_ROLE) returns (SecretStore.SelectedEnclave[] memory selectedEnclaves) { + return _selectEnclaves(_noOfNodesToSelect, _sizeLimit); } - function updateSecretStoreEndTimestamp( - uint256 _secretId, - uint256 _endTimestamp, - uint256 _usdcDeposit - ) external { - _updateSecretStoreEndTimestamp(_secretId, _endTimestamp, _usdcDeposit); + function slashEnclave( + address _enclaveAddress, + address _recipient + ) external onlyRole(SECRET_STORE_ROLE) returns (uint256) { + return _slashEnclave(_enclaveAddress, _recipient); } - function terminateSecretStore( - uint256 _secretId - ) external { - _terminateSecretStore(_secretId); + function releaseEnclave( + address _enclaveAddress, + uint256 _sizeLimit + ) external onlyRole(SECRET_STORE_ROLE) { + _releaseEnclave(_enclaveAddress, _sizeLimit); } - //-------------------------------- external functions end ----------------------------------// - - //----------------------------------- User storage end --------------------------------------// + //---------------------------------- external functions end ----------------------------------// - //-------------------------------- SecretStorage functions end --------------------------------// + //-------------------------------- SecretStoreRole functions end --------------------------------// } diff --git a/contracts/secret-storage/SecretStore.sol b/contracts/secret-storage/SecretStore.sol new file mode 100644 index 0000000..03de2f1 --- /dev/null +++ b/contracts/secret-storage/SecretStore.sol @@ -0,0 +1,580 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.0; + +import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol"; +import "@openzeppelin/contracts-upgradeable/utils/ContextUpgradeable.sol"; +import "@openzeppelin/contracts-upgradeable/utils/introspection/ERC165Upgradeable.sol"; +import "@openzeppelin/contracts-upgradeable/access/AccessControlUpgradeable.sol"; +import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol"; +import "@openzeppelin/contracts/token/ERC20/IERC20.sol"; +import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; +import "@openzeppelin/contracts/utils/cryptography/ECDSA.sol"; +import "./EnclaveStore.sol"; + +/** + * @title SecretStore Contract + * @notice Manages the registration, staking, and job assignment of execution nodes. + * @dev This contract is upgradeable and uses the UUPS (Universal Upgradeable Proxy Standard) pattern. + */ +contract SecretStore is + Initializable, // initializer + ContextUpgradeable, // _msgSender, _msgData + ERC165Upgradeable, // supportsInterface + AccessControlUpgradeable, + UUPSUpgradeable // public upgrade +{ + using SafeERC20 for IERC20; + using ECDSA for bytes32; + + /// @notice Thrown when the provided ERC20 token address is zero. + error SecretStoreZeroAddressUsdcToken(); + + /** + * @dev Initializes the logic contract without any admins, safeguarding against takeover. + * @param attestationVerifier The attestation verifier contract. + * @param maxAge Maximum age for attestations. + * @param _token The ERC20 token used for staking. + * @param _minStakeAmount Minimum stake amount required. + * @param _slashPercentInBips Slashing percentage in basis points. + * @param _slashMaxBips Maximum basis points for slashing. + */ + /// @custom:oz-upgrades-unsafe-allow constructor + constructor( + IERC20 _usdcToken, + uint256 _noOfNodesToSelect, + uint256 _globalMaxStoreSize, + uint256 _globalMinStoreDuration, + uint256 _globalMaxStoreDuration, + uint256 _acknowledgementTimeout, + uint256 _markAliveTimeout, + uint256 _storageFeeRate, + address _stakingPaymentPool, + address _storageEnclaves + ) { + _disableInitializers(); + + if (address(_usdcToken) == address(0)) revert SecretStoreZeroAddressUsdcToken(); + + USDC_TOKEN = _usdcToken; + NO_OF_NODES_TO_SELECT = _noOfNodesToSelect; + + // TODO: add checks + GLOBAL_MAX_STORE_SIZE = _globalMaxStoreSize; + GLOBAL_MIN_STORE_DURATION = _globalMinStoreDuration; + GLOBAL_MAX_STORE_DURATION = _globalMaxStoreDuration; + ACKNOWLEDGEMENT_TIMEOUT = _acknowledgementTimeout; + MARK_ALIVE_TIMEOUT = _markAliveTimeout; + STORAGE_FEE_RATE = _storageFeeRate; + STAKING_PAYMENT_POOL = _stakingPaymentPool; + ENCLAVE_STORE = EnclaveStore(_storageEnclaves); + } + + //-------------------------------- Overrides start --------------------------------// + + /// @inheritdoc ERC165Upgradeable + function supportsInterface( + bytes4 interfaceId + ) public view virtual override(ERC165Upgradeable, AccessControlUpgradeable) returns (bool) { + return super.supportsInterface(interfaceId); + } + + /// @inheritdoc UUPSUpgradeable + function _authorizeUpgrade(address /*account*/) internal view override onlyRole(DEFAULT_ADMIN_ROLE) {} + + //-------------------------------- Overrides end --------------------------------// + + //-------------------------------- Initializer start --------------------------------// + + /// @notice Thrown when the provided admin address is zero. + error SecretStoreZeroAddressAdmin(); + + /** + * @dev Initializes the contract with the given admin. + * @param _admin The address of the admin. + */ + function initialize(address _admin) public initializer { + if (_admin == address(0)) revert SecretStoreZeroAddressAdmin(); + + __Context_init_unchained(); + __ERC165_init_unchained(); + __AccessControl_init_unchained(); + __UUPSUpgradeable_init_unchained(); + + _grantRole(DEFAULT_ADMIN_ROLE, _admin); + } + + //-------------------------------- Initializer end --------------------------------// + + /// @custom:oz-upgrades-unsafe-allow state-variable-immutable + IERC20 public immutable USDC_TOKEN; + + /// @custom:oz-upgrades-unsafe-allow state-variable-immutable + uint256 public immutable NO_OF_NODES_TO_SELECT; + + /// @custom:oz-upgrades-unsafe-allow state-variable-immutable + uint256 public immutable GLOBAL_MAX_STORE_SIZE; + + /// @custom:oz-upgrades-unsafe-allow state-variable-immutable + uint256 public immutable GLOBAL_MIN_STORE_DURATION; + + /// @custom:oz-upgrades-unsafe-allow state-variable-immutable + uint256 public immutable GLOBAL_MAX_STORE_DURATION; + + /// @custom:oz-upgrades-unsafe-allow state-variable-immutable + uint256 public immutable ACKNOWLEDGEMENT_TIMEOUT; + + /// @custom:oz-upgrades-unsafe-allow state-variable-immutable + uint256 public immutable MARK_ALIVE_TIMEOUT; + + /// @custom:oz-upgrades-unsafe-allow state-variable-immutable + uint256 public immutable STORAGE_FEE_RATE; + + /// @custom:oz-upgrades-unsafe-allow state-variable-immutable + address public immutable STAKING_PAYMENT_POOL; + + /// @custom:oz-upgrades-unsafe-allow state-variable-immutable + EnclaveStore public immutable ENCLAVE_STORE; + + //------------------------------------ SecretStore start -----------------------------------// + + bytes32 private constant DOMAIN_SEPARATOR = + keccak256( + abi.encode( + keccak256("EIP712Domain(string name,string version)"), + keccak256("marlin.oyster.SecretStore"), + keccak256("1") + ) + ); + + bytes32 private constant ACKNOWLEDGE_TYPEHASH = + keccak256("Acknowledge(uint256 secretId,uint256 signTimestamp)"); + + bytes32 private constant ALIVE_TYPEHASH = + keccak256("Alive(uint256 secretId,uint256 signTimestamp)"); + + /// @notice Thrown when the signature timestamp has expired. + error SecretStoreSignatureTooOld(); + + struct SelectedEnclave { + address enclaveAddress; + bool hasAcknowledgedStore; + uint256 lastAliveTimestamp; + } + + struct UserStorage { + address owner; + uint256 sizeLimit; + uint256 usdcDeposit; + uint256 startTimestamp; + uint256 endTimestamp; + SelectedEnclave[] selectedEnclaves; + } + + // secretId => user store data + mapping(uint256 => UserStorage) public userStorage; + + uint256 public secretId; + + event SecretStoreCreated( + uint256 indexed secretId, + address indexed owner, + uint256 sizeLimit, + uint256 endTimestamp, + uint256 usdcDeposit + ); + + event EnclaveAcknowledgedStore( + uint256 indexed secretId, + address indexed enclaveAddress + ); + + event EnclaveAcknowledgementFailed( + uint256 indexed secretId + ); + + event EnclaveStoreAlive( + uint256 indexed secretId, + address indexed enclaveAddress + ); + + event EnclaveStoreDead( + uint256 indexed secretId, + address indexed prevEnclaveAddress, + address indexed newEnclaveAddress + ); + + event SecretStoreResourceUnavailable( + uint256 indexed secretId + ); + + event SecretStoreEndTimestampUpdated( + uint256 indexed secretId, + uint256 endTimestamp + ); + + event SecretStoreTerminated( + uint256 indexed secretId + ); + + error SecretStoreInsufficientUsdcDeposit(); + error SecretStoreInvalidSizeLimit(); + error SecretStoreInvalidEndTimestamp(); + error SecretStoreUnavailableResources(); + error SecretStoreAcknowledgementTimeOver(); + error SecretStoreAcknowledgementTimeoutPending(); + error SecretStoreAcknowledgedAlready(); + error SecretStoreMarkAliveTimeoutOver(); + error SecretStoreUnacknowledged(); + error SecretStoreEnclaveNotFound(); + error SecretStoreNotUserStoreOwner(); + error SecretStoreAlreadyTerminated(); + + //-------------------------------- internal functions start ----------------------------------// + + function _createSecretStore( + uint256 _sizeLimit, + uint256 _endTimestamp, + uint256 _usdcDeposit, + address _owner + ) internal { + if(_sizeLimit == 0 || _sizeLimit > GLOBAL_MAX_STORE_SIZE) + revert SecretStoreInvalidSizeLimit(); + + if ((_endTimestamp < block.timestamp + GLOBAL_MIN_STORE_DURATION) || (_endTimestamp > block.timestamp + GLOBAL_MAX_STORE_DURATION)) + revert SecretStoreInvalidEndTimestamp(); + + // TODO: how to calculate usdcDeposit + uint256 minUsdcDeposit = (_endTimestamp - block.timestamp) * _sizeLimit * STORAGE_FEE_RATE; + _checkUsdcDeposit(_usdcDeposit, minUsdcDeposit); + + USDC_TOKEN.safeTransferFrom(_owner, address(this), _usdcDeposit); + + SelectedEnclave[] memory selectedEnclaves = ENCLAVE_STORE.selectEnclaves(NO_OF_NODES_TO_SELECT, _sizeLimit); + if (selectedEnclaves.length < NO_OF_NODES_TO_SELECT) + revert SecretStoreUnavailableResources(); + + uint256 id = ++secretId; + userStorage[id].owner = _owner; + userStorage[id].sizeLimit = _sizeLimit; + userStorage[id].usdcDeposit = _usdcDeposit; + userStorage[id].startTimestamp = block.timestamp; + userStorage[id].endTimestamp = _endTimestamp; + + // cannot allocate memory array directly to storage var + for (uint256 index = 0; index < selectedEnclaves.length; index++) { + userStorage[id].selectedEnclaves.push(selectedEnclaves[index]); + } + + emit SecretStoreCreated(secretId, _owner, _sizeLimit, _endTimestamp, _usdcDeposit); + } + + function _checkUsdcDeposit( + uint256 _usdcDeposit, + uint256 _minUsdcDeposit + ) internal pure { + if(_usdcDeposit < _minUsdcDeposit) + revert SecretStoreInsufficientUsdcDeposit(); + } + + function _acknowledgeStore( + uint256 _secretId, + uint256 _signTimestamp, + bytes memory _signature + ) internal { + if(block.timestamp > userStorage[_secretId].startTimestamp + ACKNOWLEDGEMENT_TIMEOUT) + revert SecretStoreAcknowledgementTimeOver(); + + address enclaveAddress = _verifyAcknowledgementSign(_secretId, _signTimestamp, _signature); + + uint256 enclaveIndex = _getSelectedEnclaveIndex(_secretId, enclaveAddress); + userStorage[_secretId].selectedEnclaves[enclaveIndex].hasAcknowledgedStore = true; + userStorage[_secretId].selectedEnclaves[enclaveIndex].lastAliveTimestamp = _signTimestamp; + + emit EnclaveAcknowledgedStore(_secretId, enclaveAddress); + } + + function _verifyAcknowledgementSign( + uint256 _secretId, + uint256 _signTimestamp, + bytes memory _signature + ) internal view returns(address signer) { + _checkSignValidity(_signTimestamp); + + bytes32 hashStruct = keccak256(abi.encode(ACKNOWLEDGE_TYPEHASH, _secretId, _signTimestamp)); + bytes32 digest = keccak256(abi.encodePacked("\x19\x01", DOMAIN_SEPARATOR, hashStruct)); + signer = digest.recover(_signature); + + ENCLAVE_STORE.allowOnlyVerified(signer); + } + + function _checkSignValidity(uint256 _signTimestamp) internal view { + if (block.timestamp > _signTimestamp + ENCLAVE_STORE.ATTESTATION_MAX_AGE()) + revert SecretStoreSignatureTooOld(); + } + + function _getSelectedEnclaveIndex( + uint256 _secretId, + address _enclaveAddress + ) internal view returns (uint256) { + uint256 len = userStorage[_secretId].selectedEnclaves.length; + for (uint256 index = 0; index < len; index++) { + if(userStorage[_secretId].selectedEnclaves[index].enclaveAddress == _enclaveAddress) + return index; + } + revert SecretStoreEnclaveNotFound(); + } + + function _acknowledgeStoreFailed( + uint256 _secretId + ) internal { + if(block.timestamp <= userStorage[_secretId].startTimestamp + ACKNOWLEDGEMENT_TIMEOUT) + revert SecretStoreAcknowledgementTimeoutPending(); + + bool ackFailed; + uint256 len = userStorage[_secretId].selectedEnclaves.length; + for (uint256 index = 0; index < len; index++) { + if(!userStorage[_secretId].selectedEnclaves[index].hasAcknowledgedStore) { + ackFailed = true; + break; + } + } + + if(!ackFailed) + revert SecretStoreAcknowledgedAlready(); + + address owner = userStorage[_secretId].owner; + uint256 usdcDeposit = userStorage[_secretId].usdcDeposit; + delete userStorage[_secretId]; + + USDC_TOKEN.safeTransfer(owner, usdcDeposit); + + emit EnclaveAcknowledgementFailed(_secretId); + } + + function _markStoreAlive( + uint256 _secretId, + uint256 _signTimestamp, + bytes memory _signature, + address _owner + ) internal { + address enclaveAddress = _verifyStoreAliveSign(_secretId, _signTimestamp, _signature); + + uint256 enclaveIndex = _getSelectedEnclaveIndex(_secretId, enclaveAddress); + if(!userStorage[_secretId].selectedEnclaves[enclaveIndex].hasAcknowledgedStore) + revert SecretStoreUnacknowledged(); + if(block.timestamp > userStorage[_secretId].selectedEnclaves[enclaveIndex].lastAliveTimestamp + MARK_ALIVE_TIMEOUT) + revert SecretStoreMarkAliveTimeoutOver(); + + uint256 endTime; + if(block.timestamp > userStorage[_secretId].endTimestamp) + endTime = userStorage[_secretId].endTimestamp; + else + endTime = block.timestamp; + + uint256 usdcPayment = ((endTime - userStorage[_secretId].selectedEnclaves[enclaveIndex].lastAliveTimestamp) * userStorage[_secretId].sizeLimit * STORAGE_FEE_RATE) / NO_OF_NODES_TO_SELECT; + userStorage[_secretId].usdcDeposit -= usdcPayment; + userStorage[_secretId].selectedEnclaves[enclaveIndex].lastAliveTimestamp = _signTimestamp; + + USDC_TOKEN.safeTransfer(_owner, usdcPayment); + + // TODO: delete from selectedNodes array and refund remaining usdc + if(block.timestamp > userStorage[_secretId].endTimestamp) { + _removeSelectedEnclave(_secretId, enclaveIndex); + if(userStorage[_secretId].selectedEnclaves.length == 0) + _refundExcessDepositAndRemoveStore(_secretId); + } + + emit EnclaveStoreAlive(_secretId, enclaveAddress); + } + + function _verifyStoreAliveSign( + uint256 _secretId, + uint256 _signTimestamp, + bytes memory _signature + ) internal view returns(address signer) { + _checkSignValidity(_signTimestamp); + + bytes32 hashStruct = keccak256(abi.encode(ALIVE_TYPEHASH, _secretId, _signTimestamp)); + bytes32 digest = keccak256(abi.encodePacked("\x19\x01", DOMAIN_SEPARATOR, hashStruct)); + signer = digest.recover(_signature); + + ENCLAVE_STORE.allowOnlyVerified(signer); + } + + function _markStoreDead( + uint256 _secretId + ) internal { + for (uint256 index = 0; index < userStorage[_secretId].selectedEnclaves.length; ) { + bool isArrayLengthReduced = _markEnclaveDead(_secretId, userStorage[_secretId].selectedEnclaves[index].enclaveAddress); + if(!isArrayLengthReduced) + ++index; + } + + // TODO: delete data and refund remaining + if(block.timestamp > userStorage[_secretId].endTimestamp && userStorage[_secretId].selectedEnclaves.length == 0) + _refundExcessDepositAndRemoveStore(_secretId); + } + + function _markEnclaveDead( + uint256 _secretId, + address _enclaveAddress + ) internal returns (bool isArrayLengthReduced) { + uint256 enclaveIndex = _getSelectedEnclaveIndex(_secretId, _enclaveAddress); + if(!userStorage[_secretId].selectedEnclaves[enclaveIndex].hasAcknowledgedStore) + return isArrayLengthReduced; + + if(block.timestamp <= userStorage[_secretId].selectedEnclaves[enclaveIndex].lastAliveTimestamp + MARK_ALIVE_TIMEOUT) + return isArrayLengthReduced; + + if(block.timestamp > userStorage[_secretId].endTimestamp) { + isArrayLengthReduced = true; + _removeSelectedEnclave(_secretId, enclaveIndex); + } + else { + SelectedEnclave[] memory selectedEnclaves = ENCLAVE_STORE.selectEnclaves(1, userStorage[_secretId].sizeLimit); + if (selectedEnclaves.length == 0) { + isArrayLengthReduced = true; + _removeSelectedEnclave(_secretId, enclaveIndex); + emit SecretStoreResourceUnavailable(_secretId); + } + else { + userStorage[_secretId].selectedEnclaves[enclaveIndex] = SelectedEnclave({ + enclaveAddress: selectedEnclaves[0].enclaveAddress, + hasAcknowledgedStore: false, + lastAliveTimestamp: 0 + }); + } + + emit EnclaveStoreDead( + _secretId, + _enclaveAddress, + selectedEnclaves.length != 0 ? selectedEnclaves[0].enclaveAddress : address(0) + ); + } + + // TODO: slash prev enclave(who's recipient) + ENCLAVE_STORE.slashEnclave(_enclaveAddress, STAKING_PAYMENT_POOL); + ENCLAVE_STORE.releaseEnclave(_enclaveAddress, userStorage[_secretId].sizeLimit); + } + + function _removeSelectedEnclave( + uint256 _secretId, + uint256 _index + ) internal { + uint256 len = userStorage[_secretId].selectedEnclaves.length; + if(_index != len - 1) + userStorage[_secretId].selectedEnclaves[_index] = userStorage[_secretId].selectedEnclaves[len - 1]; + userStorage[_secretId].selectedEnclaves.pop(); + } + + function _refundExcessDepositAndRemoveStore( + uint256 _secretId + ) internal { + address owner = userStorage[_secretId].owner; + uint256 remainingDeposit = userStorage[_secretId].usdcDeposit; + delete userStorage[_secretId]; + USDC_TOKEN.safeTransfer(owner, remainingDeposit); + } + + function _updateSecretStoreEndTimestamp( + uint256 _secretId, + uint256 _endTimestamp, + uint256 _usdcDeposit + ) internal { + if(userStorage[_secretId].owner != _msgSender()) + revert SecretStoreNotUserStoreOwner(); + + if(_endTimestamp < block.timestamp) + revert SecretStoreInvalidEndTimestamp(); + + uint256 currentEndTimestamp = userStorage[_secretId].endTimestamp; + if(block.timestamp > currentEndTimestamp) + revert SecretStoreAlreadyTerminated(); + + if(_endTimestamp > currentEndTimestamp) { + USDC_TOKEN.safeTransferFrom(_msgSender(), address(this), _usdcDeposit); + userStorage[_secretId].usdcDeposit += _usdcDeposit; + + uint256 addedDuration = _endTimestamp - currentEndTimestamp; + uint256 minUsdcDeposit = addedDuration * userStorage[_secretId].sizeLimit * STORAGE_FEE_RATE; + _checkUsdcDeposit(_usdcDeposit, minUsdcDeposit); + } + else { + uint256 removedDuration = currentEndTimestamp - _endTimestamp; + uint256 usdcRefund = removedDuration * userStorage[_secretId].sizeLimit * STORAGE_FEE_RATE; + + userStorage[_secretId].usdcDeposit -= usdcRefund; + USDC_TOKEN.safeTransfer(_msgSender(), usdcRefund); + } + + userStorage[_secretId].endTimestamp = _endTimestamp; + + emit SecretStoreEndTimestampUpdated(_secretId, _endTimestamp); + } + + function _terminateSecretStore( + uint256 _secretId + ) internal { + _updateSecretStoreEndTimestamp(_secretId, block.timestamp, 0); + + emit SecretStoreTerminated(_secretId); + } + + //-------------------------------- internal functions end ----------------------------------// + + //------------------------------- external functions start ----------------------------------// + + function createSecretStore( + uint256 _sizeLimit, + uint256 _endTimestamp, + uint256 _usdcDeposit + ) external { + _createSecretStore(_sizeLimit, _endTimestamp, _usdcDeposit, _msgSender()); + } + + function acknowledgeStore( + uint256 _secretId, + uint256 _signTimestamp, + bytes memory _signature + ) external { + _acknowledgeStore(_secretId, _signTimestamp, _signature); + } + + function acknowledgeStoreFailed( + uint256 _secretId + ) external { + _acknowledgeStoreFailed(_secretId); + } + + function markStoreAlive( + uint256 _secretId, + uint256 _signTimestamp, + bytes memory _signature + ) external { + _markStoreAlive(_secretId, _signTimestamp, _signature, _msgSender()); + } + + function markStoreDead( + uint256 _secretId + ) external { + _markStoreDead(_secretId); + } + + function updateSecretStoreEndTimestamp( + uint256 _secretId, + uint256 _endTimestamp, + uint256 _usdcDeposit + ) external { + _updateSecretStoreEndTimestamp(_secretId, _endTimestamp, _usdcDeposit); + } + + function terminateSecretStore( + uint256 _secretId + ) external { + _terminateSecretStore(_secretId); + } + + //-------------------------------- external functions end ----------------------------------// + + //-------------------------------- SecretStore functions end --------------------------------// +}