From b48701d8bdc7813b75e2bea037b8e80da7420413 Mon Sep 17 00:00:00 2001 From: twygod Date: Wed, 10 Nov 2021 20:48:13 +0800 Subject: [PATCH 1/3] fix: Unbounded iteration in deleteMarket. --- contracts/market/MarketRegistry.sol | 63 +++++++++-------------------- test/unit/testMarketRegistry.js | 2 + 2 files changed, 22 insertions(+), 43 deletions(-) diff --git a/contracts/market/MarketRegistry.sol b/contracts/market/MarketRegistry.sol index 8359417..271947f 100644 --- a/contracts/market/MarketRegistry.sol +++ b/contracts/market/MarketRegistry.sol @@ -1,6 +1,8 @@ //SPDX-License-Identifier: UNLICENSED pragma solidity ^0.8.4; +import "@openzeppelin/contracts/utils/structs/EnumerableSet.sol"; + import "../Controller.sol"; /** @@ -8,13 +10,15 @@ import "../Controller.sol"; * @dev Registering and managing all the lending markets. */ contract MarketRegistry is Controller { + using EnumerableSet for EnumerableSet.AddressSet; + struct Market { address uToken; address userManager; } - address[] public uTokenList; - address[] public userManagerList; + EnumerableSet.AddressSet private uTokenList; + EnumerableSet.AddressSet private userManagerList; mapping(address => Market) public tokens; event LogAddUToken(address indexed tokenAddress, address contractAddress); @@ -43,60 +47,33 @@ contract MarketRegistry is Controller { * @return Stored uToken address */ function getUTokens() public view returns (address[] memory) { - return uTokenList; + return uTokenList.values(); } function getUserManagers() public view returns (address[] memory) { - return userManagerList; + return userManagerList.values(); } function addUToken(address token, address uToken) public newToken(token) onlyAdmin { - uTokenList.push(uToken); - tokens[token].uToken = uToken; - emit LogAddUToken(token, uToken); + if (tokens[token].uToken == address(0)) { + uTokenList.add(uToken); + tokens[token].uToken = uToken; + emit LogAddUToken(token, uToken); + } } function addUserManager(address token, address userManager) public newUserManager(token) onlyAdmin { - userManagerList.push(userManager); - tokens[token].userManager = userManager; - emit LogAddUserManager(token, userManager); + if (tokens[token].userManager == address(0)) { + userManagerList.add(userManager); + tokens[token].userManager = userManager; + emit LogAddUserManager(token, userManager); + } } function deleteMarket(address token) public onlyAdmin { - address oldUToken = tokens[token].uToken; - bool uTokenExist = false; - uint256 uTokenIndex = 0; - - for (uint256 i = 0; i < uTokenList.length; i++) { - if (oldUToken == uTokenList[i]) { - uTokenExist = true; - uTokenIndex = i; - } - } - - if (uTokenExist) { - uTokenList[uTokenIndex] = uTokenList[uTokenList.length - 1]; - uTokenList.pop(); - } - + uTokenList.remove(tokens[token].uToken); + userManagerList.remove(tokens[token].userManager); delete tokens[token].uToken; - - address oldUserManager = tokens[token].userManager; - bool userManagerExist = false; - uint256 userManagerIndex = 0; - - for (uint256 i = 0; i < userManagerList.length; i++) { - if (oldUserManager == userManagerList[i]) { - userManagerExist = true; - userManagerIndex = i; - } - } - - if (userManagerExist) { - userManagerList[userManagerIndex] = userManagerList[userManagerList.length - 1]; - userManagerList.pop(); - } - delete tokens[token].userManager; } } diff --git a/test/unit/testMarketRegistry.js b/test/unit/testMarketRegistry.js index 2c131a5..0986251 100644 --- a/test/unit/testMarketRegistry.js +++ b/test/unit/testMarketRegistry.js @@ -28,6 +28,7 @@ describe("MarketRegistry Contract", () => { res.uToken.should.eq(uToken.address); utokens = await marketRegistry.getUTokens(); + console.log("utokens: ", utokens); utokens.length.should.eq(1); await marketRegistry.deleteMarket(ETH); @@ -49,6 +50,7 @@ describe("MarketRegistry Contract", () => { res.userManager.should.eq(userManager.address); markets = await marketRegistry.getUserManagers(); + console.log("userManager: ", markets); markets.length.should.eq(1); await marketRegistry.deleteMarket(ETH); From 86938814ca170cfe104de2d2797f9d131e8d7d87 Mon Sep 17 00:00:00 2001 From: twygod Date: Mon, 15 Nov 2021 12:59:10 +0800 Subject: [PATCH 2/3] chg: revert when the address already exists use addUToken and addUserManager. --- contracts/market/MarketRegistry.sol | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/contracts/market/MarketRegistry.sol b/contracts/market/MarketRegistry.sol index 271947f..5a41b17 100644 --- a/contracts/market/MarketRegistry.sol +++ b/contracts/market/MarketRegistry.sol @@ -55,19 +55,17 @@ contract MarketRegistry is Controller { } function addUToken(address token, address uToken) public newToken(token) onlyAdmin { - if (tokens[token].uToken == address(0)) { - uTokenList.add(uToken); - tokens[token].uToken = uToken; - emit LogAddUToken(token, uToken); - } + require(tokens[token].uToken == address(0), "MarketRegistry: uToken has set"); + uTokenList.add(uToken); + tokens[token].uToken = uToken; + emit LogAddUToken(token, uToken); } function addUserManager(address token, address userManager) public newUserManager(token) onlyAdmin { - if (tokens[token].userManager == address(0)) { - userManagerList.add(userManager); - tokens[token].userManager = userManager; - emit LogAddUserManager(token, userManager); - } + require(tokens[token].userManager == address(0), "MarketRegistry: userManager has set"); + userManagerList.add(userManager); + tokens[token].userManager = userManager; + emit LogAddUserManager(token, userManager); } function deleteMarket(address token) public onlyAdmin { From aad076c0a976ba6a8581434d5761c442e024e7e5 Mon Sep 17 00:00:00 2001 From: twygod Date: Mon, 15 Nov 2021 13:03:54 +0800 Subject: [PATCH 3/3] chg: remove repeated require. --- contracts/market/MarketRegistry.sol | 2 -- 1 file changed, 2 deletions(-) diff --git a/contracts/market/MarketRegistry.sol b/contracts/market/MarketRegistry.sol index 5a41b17..3a462ab 100644 --- a/contracts/market/MarketRegistry.sol +++ b/contracts/market/MarketRegistry.sol @@ -55,14 +55,12 @@ contract MarketRegistry is Controller { } function addUToken(address token, address uToken) public newToken(token) onlyAdmin { - require(tokens[token].uToken == address(0), "MarketRegistry: uToken has set"); uTokenList.add(uToken); tokens[token].uToken = uToken; emit LogAddUToken(token, uToken); } function addUserManager(address token, address userManager) public newUserManager(token) onlyAdmin { - require(tokens[token].userManager == address(0), "MarketRegistry: userManager has set"); userManagerList.add(userManager); tokens[token].userManager = userManager; emit LogAddUserManager(token, userManager);