Contracts review

smart-contracts/.env.example

@@ -1,5 +1,7 @@
 COINMARKETCAP_API_KEY= # for coverage report to estimate deployment and calls price
 ETH_NODE_ADDRESS=https://arb-sepolia.g.alchemy.com/v2/SOME_API_KEY
 ETHERSCAN_API_KEY= # for coverage report to estimate deployment and calls price
+ARBITRUM_API_KEY= # for coverage report to estimate deployment and calls price
 MOR_TOKEN_ADDRESS= # MOR token address
-OWNER_PRIVATE_KEY= # contract owner private key
+OWNER_PRIVATE_KEY= # contract owner private key
+PRIVATE_KEY= # deployer private key

smart-contracts/.eslintrc.json

@@ -0,0 +1,22 @@
+{
+  "parser": "@typescript-eslint/parser",
+  "parserOptions": {
+    "ecmaVersion": 12,
+    "sourceType": "module",
+    "project": "./tsconfig.json"
+  },
+  "plugins": ["@typescript-eslint"],
+  "extends": [
+    "eslint:recommended",
+    "plugin:@typescript-eslint/recommended",
+    "plugin:prettier/recommended"
+  ],
+  "rules": {
+    "@typescript-eslint/no-unused-vars": "error",
+    "@typescript-eslint/no-floating-promises": "error"
+  },
+  "env": {
+    "browser": true,
+    "es2021": true
+  }
+}

smart-contracts/contracts/diamond/facets/Marketplace.sol

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.24;
 
+import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
 import {SafeERC20, IERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 
 import {OwnableDiamondStorage} from "../presets/OwnableDiamondStorage.sol";
@@ -21,102 +22,96 @@ contract Marketplace is
     BidStorage
 {
     using SafeERC20 for IERC20;
+    using EnumerableSet for EnumerableSet.Bytes32Set;
 
-    function __Marketplace_init(
-        address token_
-    ) external initializer(MARKETPLACE_STORAGE_SLOT) initializer(BID_STORAGE_SLOT) {
-        _getBidStorage().token = IERC20(token_);
+    function __Marketplace_init(address token_) external initializer(BIDS_STORAGE_SLOT) {
+        BidsStorage storage bidsStorage = getBidsStorage();
+        bidsStorage.token = token_;
     }
 
-    /// @notice sets a bid fee
-    function setBidFee(uint256 bidFee_) external onlyOwner {
-        _getMarketplaceStorage().bidFee = bidFee_;
-        emit FeeUpdated(bidFee_);
+    function setMarketplaceBidFee(uint256 bidFee_) external onlyOwner {
+        MarketStorage storage marketStorage = getMarketStorage();
+        marketStorage.bidFee = bidFee_;
+
+        emit MaretplaceFeeUpdated(bidFee_);
     }
 
-    /// @notice posts a new bid for a model
-    function postModelBid(
-        address provider_,
-        bytes32 modelId_,
-        uint256 pricePerSecond_
-    ) external returns (bytes32 bidId) {
-        if (!_ownerOrProvider(provider_)) {
-            revert NotOwnerOrProvider();
-        }
-        if (!isProviderActive(provider_)) {
-            revert ProviderNotFound();
+    function postModelBid(bytes32 modelId_, uint256 pricePerSecond_) external returns (bytes32 bidId) {
+        address provider_ = _msgSender();
+
+        if (!getIsProviderActive(provider_)) {
+            revert MarketplaceProviderNotFound();
         }
-        if (!isModelActive(modelId_)) {
-            revert ModelNotFound();
+        if (!getIsModelActive(modelId_)) {
+            revert MarketplaceModelNotFound();
         }
 
-        return _postModelBid(provider_, modelId_, pricePerSecond_);
-    }
+        BidsStorage storage bidsStorage = getBidsStorage();
+        MarketStorage storage marketStorage = getMarketStorage();
 
-    /// @notice deletes a bid
-    function deleteModelBid(bytes32 bidId_) external {
-        if (!_isBidActive(bidId_)) {
-            revert ActiveBidNotFound();
-        }
-        if (!_ownerOrProvider(getBid(bidId_).provider)) {
-            revert NotOwnerOrProvider();
-        }
+        IERC20(bidsStorage.token).safeTransferFrom(_msgSender(), address(this), marketStorage.bidFee);
+        marketStorage.feeBalance += marketStorage.bidFee;
 
-        _deleteBid(bidId_);
-    }
+        bytes32 providerModelId_ = getProviderModelId(provider_, modelId_);
+        uint256 providerModelNonce_ = bidsStorage.providerModelNonce[providerModelId_]++;
+        bytes32 bidId_ = getBidId(provider_, modelId_, providerModelNonce_);
 
-    /// @notice withdraws the fee balance
-    function withdraw(address recipient_, uint256 amount_) external onlyOwner {
-        if (amount_ > getFeeBalance()) {
-            revert NotEnoughBalance();
+        if (providerModelNonce_ != 0) {
+            bytes32 oldBidId_ = getBidId(provider_, modelId_, providerModelNonce_ - 1);
+            if (isBidActive(oldBidId_)) {
+                _deleteBid(oldBidId_);
+            }
         }
 
-        decreaseFeeBalance(amount_);
-        getToken().safeTransfer(recipient_, amount_);
-    }
+        Bid storage bid = bidsStorage.bids[bidId_];
+        bid.provider = provider_;
+        bid.modelId = modelId_;
+        bid.pricePerSecond = pricePerSecond_;
+        bid.nonce = providerModelNonce_;
+        bid.createdAt = uint128(block.timestamp);
+
+        bidsStorage.providerBids[provider_].add(bidId_);
+        bidsStorage.providerActiveBids[provider_].add(bidId_);
+        bidsStorage.modelBids[modelId_].add(bidId_);
+        bidsStorage.modelActiveBids[modelId_].add(bidId_);
+
+        emit MarketplaceBidPosted(provider_, modelId_, providerModelNonce_);
 
-    function _incrementBidNonce(address provider_, bytes32 modelId_) private returns (uint256) {
-        return _incrementBidNonce(getProviderModelId(provider_, modelId_));
+        return bidId_;
     }
 
-    function _postModelBid(address provider_, bytes32 modelId_, uint256 pricePerSecond_) private returns (bytes32) {
-        uint256 fee_ = getBidFee();
-        getToken().safeTransferFrom(_msgSender(), address(this), fee_);
-        increaseFeeBalance(fee_);
+    function deleteModelBid(bytes32 bidId_) external {
+        BidsStorage storage bidsStorage = getBidsStorage();
+        _onlyAccount(bidsStorage.bids[bidId_].provider);
 
-        // TEST IT if it increments nonce correctly
-        uint256 nonce_ = _incrementBidNonce(provider_, modelId_);
-        if (nonce_ != 0) {
-            bytes32 oldBidId_ = getBidId(provider_, modelId_, nonce_ - 1);
-            if (_isBidActive(oldBidId_)) {
-                _deleteBid(oldBidId_);
-            }
+        if (!isBidActive(bidId_)) {
+            revert MarketplaceActiveBidNotFound();
         }
 
-        bytes32 bidId_ = getBidId(provider_, modelId_, nonce_);
-
-        setBid(bidId_, Bid(provider_, modelId_, pricePerSecond_, nonce_, uint128(block.timestamp), 0));
+        _deleteBid(bidId_);
+    }
 
-        addProviderBid(provider_, bidId_);
-        addModelBid(modelId_, bidId_);
+    function withdraw(address recipient_, uint256 amount_) external onlyOwner {
+        BidsStorage storage bidsStorage = getBidsStorage();
+        MarketStorage storage marketStorage = getMarketStorage();
 
-        addProviderActiveBids(provider_, bidId_);
-        addModelActiveBids(modelId_, bidId_);
+        amount_ = amount_ > marketStorage.feeBalance ? marketStorage.feeBalance : amount_;
 
-        emit BidPosted(provider_, modelId_, nonce_);
+        marketStorage.feeBalance -= amount_;
 
-        return bidId_;
+        IERC20(bidsStorage.token).safeTransfer(recipient_, amount_);
     }
 
     function _deleteBid(bytes32 bidId_) private {
-        Bid storage bid = getBid(bidId_);
+        BidsStorage storage bidsStorage = getBidsStorage();
+        Bid storage bid = bidsStorage.bids[bidId_];
 
         bid.deletedAt = uint128(block.timestamp);
 
-        removeProviderActiveBids(bid.provider, bidId_);
-        removeModelActiveBids(bid.modelId, bidId_);
+        bidsStorage.providerActiveBids[bid.provider].remove(bidId_);
+        bidsStorage.modelActiveBids[bid.modelId].remove(bidId_);
 
-        emit BidDeleted(bid.provider, bid.modelId, bid.nonce);
+        emit MarketplaceBidDeleted(bid.provider, bid.modelId, bid.nonce);
     }
 
     function getBidId(address provider_, bytes32 modelId_, uint256 nonce_) public pure returns (bytes32) {
@@ -126,14 +121,4 @@ contract Marketplace is
     function getProviderModelId(address provider_, bytes32 modelId_) public pure returns (bytes32) {
         return keccak256(abi.encodePacked(provider_, modelId_));
     }
-
-    function _ownerOrProvider(address provider_) private view returns (bool) {
-        return _msgSender() == owner() || _msgSender() == provider_;
-    }
-
-    function _isBidActive(bytes32 bidId_) private view returns (bool) {
-        Bid memory bid_ = getBid(bidId_);
-
-        return bid_.createdAt != 0 && bid_.deletedAt == 0;
-    }
 }

smart-contracts/contracts/diamond/facets/ModelRegistry.sol

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.24;
 
+import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
 import {SafeERC20, IERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 
 import {OwnableDiamondStorage} from "../presets/OwnableDiamondStorage.sol";
@@ -11,91 +12,83 @@ import {ModelStorage} from "../storages/ModelStorage.sol";
 import {IModelRegistry} from "../../interfaces/facets/IModelRegistry.sol";
 
 contract ModelRegistry is IModelRegistry, OwnableDiamondStorage, ModelStorage, BidStorage {
+    using EnumerableSet for EnumerableSet.Bytes32Set;
     using SafeERC20 for IERC20;
 
-    function __ModelRegistry_init() external initializer(MODEL_STORAGE_SLOT) {}
+    function __ModelRegistry_init() external initializer(MODELS_STORAGE_SLOT) {}
 
-    function setModelMinimumStake(uint256 modelMinimumStake_) external onlyOwner {
-        _setModelMinimumStake(modelMinimumStake_);
-        emit ModelMinimumStakeSet(modelMinimumStake_);
+    function modelSetMinStake(uint256 modelMinimumStake_) external onlyOwner {
+        ModelsStorage storage modelsStorage = getModelsStorage();
+        modelsStorage.modelMinimumStake = modelMinimumStake_;
+
+        emit ModelMinimumStakeUpdated(modelMinimumStake_);
     }
 
-    /// @notice Registers or updates existing model
     function modelRegister(
-        // TODO: it is not secure (frontrunning) to take the modelId as key
         bytes32 modelId_,
         bytes32 ipfsCID_,
         uint256 fee_,
-        uint256 addStake_,
-        address owner_,
+        uint256 amount_,
         string calldata name_,
-        string[] calldata tags_
+        string[] memory tags_
     ) external {
-        if (!_isOwnerOrModelOwner(owner_)) {
-            revert NotOwnerOrModelOwner();
-        }
+        ModelsStorage storage modelsStorage = getModelsStorage();
+        Model storage model = modelsStorage.models[modelId_];
 
-        Model memory model_ = models(modelId_);
-        // TODO: there is no way to decrease the stake
-        uint256 newStake_ = model_.stake + addStake_;
-        if (newStake_ < modelMinimumStake()) {
-            revert StakeTooLow();
+        uint256 newStake_ = model.stake + amount_;
+        uint256 minStake_ = modelsStorage.modelMinimumStake;
+        if (newStake_ < minStake_) {
+            revert ModelStakeTooLow(newStake_, minStake_);
         }
 
-        if (addStake_ > 0) {
-            getToken().safeTransferFrom(_msgSender(), address(this), addStake_);
+        if (amount_ > 0) {
+            BidsStorage storage bidsStorage = getBidsStorage();
+            IERC20(bidsStorage.token).safeTransferFrom(_msgSender(), address(this), amount_);
         }
 
-        uint128 createdAt_ = model_.createdAt;
-        if (createdAt_ == 0) {
-            // model never existed
-            addModel(modelId_);
-            setModelActive(modelId_, true);
-            createdAt_ = uint128(block.timestamp);
+        if (model.createdAt == 0) {
+            modelsStorage.modelIds.add(modelId_);
+
+            model.createdAt = uint128(block.timestamp);
+            model.owner = _msgSender();
         } else {
-            if (!_isOwnerOrModelOwner(model_.owner)) {
-                revert NotOwnerOrModelOwner();
-            }
-            if (model_.isDeleted) {
-                setModelActive(modelId_, true);
-            }
+            _onlyAccount(model.owner);
         }
 
-        setModel(modelId_, Model(ipfsCID_, fee_, newStake_, owner_, name_, tags_, createdAt_, false));
+        model.stake = newStake_;
+        model.ipfsCID = ipfsCID_;
+        model.fee = fee_; // TODO: validate fee and get usage places
+        model.name = name_;
+        model.tags = tags_;
+        model.isDeleted = false;
 
-        emit ModelRegisteredUpdated(owner_, modelId_);
+        modelsStorage.activeModels.add(modelId_);
+
+        emit ModelRegisteredUpdated(_msgSender(), modelId_);
     }
 
     function modelDeregister(bytes32 modelId_) external {
-        Model storage model = models(modelId_);
+        ModelsStorage storage modelsStorage = getModelsStorage();
+        Model storage model = modelsStorage.models[modelId_];
 
-        if (!isModelExists(modelId_)) {
-            revert ModelNotFound();
-        }
-        if (!_isOwnerOrModelOwner(model.owner)) {
-            revert NotOwnerOrModelOwner();
-        }
+        _onlyAccount(model.owner);
         if (!isModelActiveBidsEmpty(modelId_)) {
             revert ModelHasActiveBids();
         }
+        if (model.isDeleted) {
+            revert ModelHasAlreadyDeregistered();
+        }
 
-        uint256 stake_ = model.stake;
+        uint256 withdrawAmount_ = model.stake;
 
         model.stake = 0;
         model.isDeleted = true;
 
-        setModelActive(modelId_, false);
+        modelsStorage.activeModels.remove(modelId_);
 
-        getToken().safeTransfer(model.owner, stake_);
+        BidsStorage storage bidsStorage = getBidsStorage();
+        IERC20(bidsStorage.token).safeTransfer(model.owner, withdrawAmount_);
 
         emit ModelDeregistered(model.owner, modelId_);
     }
-
-    function isModelExists(bytes32 modelId_) public view returns (bool) {
-        return models(modelId_).createdAt != 0;
-    }
-
-    function _isOwnerOrModelOwner(address modelOwner_) internal view returns (bool) {
-        return _msgSender() == owner() || _msgSender() == modelOwner_;
-    }
 }