Skip to content

Commit

Permalink
Merge pull request #14 from zkLinkProtocol/secure3_fix
Browse files Browse the repository at this point in the history
Secure3 fix
  • Loading branch information
zkbenny committed Apr 10, 2024
2 parents e97f756 + 75ecc87 commit 2e50cb5
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 27 deletions.
2 changes: 1 addition & 1 deletion contracts/interfaces/IMergeTokenPortal.sol
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ interface IMergeTokenPortal {

event SourceTokenRemoved(address indexed sourceToken);

event SourceTokenLocked(address indexed sourceToken, bool isLocked);
event SourceTokenStatusUpdated(address indexed sourceToken, bool isLocked);

event DepositLimitUpdated(address indexed sourceToken, uint256 depositLimit);

Expand Down
46 changes: 36 additions & 10 deletions contracts/merge/MergeTokenPortal.sol
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import {IERC20Upgradeable} from "@openzeppelin/contracts-upgradeable/token/ERC20/IERC20Upgradeable.sol";
import {IERC20MetadataUpgradeable} from "@openzeppelin/contracts-upgradeable/token/ERC20/extensions/IERC20MetadataUpgradeable.sol";
import {SafeERC20Upgradeable} from "@openzeppelin/contracts-upgradeable/token/ERC20/utils/SafeERC20Upgradeable.sol";
import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
Expand All @@ -10,20 +10,23 @@ import {IMergeTokenPortal} from "../interfaces/IMergeTokenPortal.sol";
import {IERC20MergeToken} from "../interfaces/IERC20MergeToken.sol";

contract MergeTokenPortal is IMergeTokenPortal, UUPSUpgradeable, OwnableUpgradeable, ReentrancyGuardUpgradeable {
using SafeERC20Upgradeable for IERC20Upgradeable;
using SafeERC20Upgradeable for IERC20MetadataUpgradeable;

/// @dev A mapping source token address => source token status.
mapping(address sourceToken => SourceTokenInfo) public sourceTokenInfoMap;

// @dev Security Council address
address public securityCouncil;

/// @dev A mapping merge token address => is merge token supported.
mapping(address mergeToken => mapping(address sourceToken => bool)) public isMergeTokenSupported;

/**
* @dev This empty reserved space is put in place to allow future versions to add new
* variables without shifting down storage in the inheritance chain.
* See https://docs.openzeppelin.com/contracts/4.x/upgradeable#storage_gaps
*/
uint256[48] private __gap;
uint256[47] private __gap;

modifier onlyOwnerOrSecurityCouncil() {
require(
Expand Down Expand Up @@ -59,6 +62,8 @@ contract MergeTokenPortal is IMergeTokenPortal, UUPSUpgradeable, OwnableUpgradea

/// @notice Deposit source token to mint merge token
function deposit(address _sourceToken, uint256 _amount, address _receiver) external override nonReentrant {
require(_sourceToken != address(0), "Invalid source token address");
require(_receiver != address(0), "Invalid receiver address");
require(_amount > 0, "Deposit amount is zero");
SourceTokenInfo storage tokenInfo = sourceTokenInfoMap[_sourceToken];
require(tokenInfo.isSupported, "Source token is not supported");
Expand All @@ -68,7 +73,10 @@ contract MergeTokenPortal is IMergeTokenPortal, UUPSUpgradeable, OwnableUpgradea
require(afterBalance <= tokenInfo.depositLimit, "Source token deposit limit exceeded");
tokenInfo.balance = afterBalance;

IERC20Upgradeable(_sourceToken).safeTransferFrom(msg.sender, address(this), _amount);
uint256 _sourceTokenBeforeBalance = IERC20MetadataUpgradeable(_sourceToken).balanceOf(address(this));
IERC20MetadataUpgradeable(_sourceToken).safeTransferFrom(msg.sender, address(this), _amount);
uint256 _sourceTokenAfterBalance = IERC20MetadataUpgradeable(_sourceToken).balanceOf(address(this));
require(_sourceTokenAfterBalance - _sourceTokenBeforeBalance == _amount, "Not support deflationary token");

address mergeToken = tokenInfo.mergeToken;
IERC20MergeToken(mergeToken).mint(_receiver, _amount);
Expand All @@ -78,6 +86,8 @@ contract MergeTokenPortal is IMergeTokenPortal, UUPSUpgradeable, OwnableUpgradea

/// @notice Burn merge token and get source token back
function withdraw(address _sourceToken, uint256 _amount, address _receiver) external override nonReentrant {
require(_sourceToken != address(0), "Invalid source token address");
require(_receiver != address(0), "Invalid receiver address");
require(_amount > 0, "Withdraw amount is zero");
SourceTokenInfo storage tokenInfo = sourceTokenInfoMap[_sourceToken];
require(tokenInfo.isSupported, "Source token is not supported");
Expand All @@ -89,32 +99,47 @@ contract MergeTokenPortal is IMergeTokenPortal, UUPSUpgradeable, OwnableUpgradea

address mergeToken = tokenInfo.mergeToken;
IERC20MergeToken(mergeToken).burn(msg.sender, _amount);
uint256 _sourceTokenBeforeBalance = IERC20MetadataUpgradeable(_sourceToken).balanceOf(address(this));
IERC20MetadataUpgradeable(_sourceToken).safeTransfer(_receiver, _amount);
uint256 _sourceTokenAfterBalance = IERC20MetadataUpgradeable(_sourceToken).balanceOf(address(this));

IERC20Upgradeable(_sourceToken).safeTransfer(_receiver, _amount);
require(_sourceTokenAfterBalance + _amount == _sourceTokenBeforeBalance, "Not support deflationary token");

emit WithdrawFromMerge(_sourceToken, mergeToken, msg.sender, _amount, _receiver);
}

/// @notice Add source token
function addSourceToken(address _sourceToken, address _mergeToken, uint256 _depositLimit) external onlyOwner {
SourceTokenInfo storage tokenInfo = sourceTokenInfoMap[_sourceToken];
require(!tokenInfo.isSupported, "Source token is already supported");
bool isSupported = sourceTokenInfoMap[_sourceToken].isSupported;
require(!isSupported, "Source token is already supported");
require(!isMergeTokenSupported[_mergeToken][_sourceToken], "Merge token is already supported");
require(_sourceToken != address(0) && _mergeToken != address(0), "Invalid token address");
require(_sourceToken != _mergeToken, "Should not Match");
uint8 _sourceTokenDecimals = IERC20MetadataUpgradeable(_sourceToken).decimals();
uint8 _mergeTokenDecimals = IERC20MetadataUpgradeable(_mergeToken).decimals();
require(_sourceTokenDecimals == _mergeTokenDecimals, "Token decimals are not the same");

sourceTokenInfoMap[_sourceToken] = SourceTokenInfo({
isSupported: true,
isLocked: false,
mergeToken: _mergeToken,
balance: 0,
depositLimit: _depositLimit
});
isMergeTokenSupported[_mergeToken][_sourceToken] = true;

emit SourceTokenAdded(_sourceToken, _mergeToken, _depositLimit);
}

/// @notice Remove source token
function removeSourceToken(address _sourceToken) external onlyOwnerOrSecurityCouncil {
SourceTokenInfo storage tokenInfo = sourceTokenInfoMap[_sourceToken];
require(tokenInfo.balance == 0, "Source Token balance is not zero");
bool isSupported = sourceTokenInfoMap[_sourceToken].isSupported;
require(isSupported, "Source token is already removed");
uint256 balance = sourceTokenInfoMap[_sourceToken].balance;
require(balance == 0, "Source Token balance is not zero");

address mergeToken = sourceTokenInfoMap[_sourceToken].mergeToken;
delete isMergeTokenSupported[mergeToken][_sourceToken];
delete sourceTokenInfoMap[_sourceToken];

emit SourceTokenRemoved(_sourceToken);
Expand All @@ -127,13 +152,14 @@ contract MergeTokenPortal is IMergeTokenPortal, UUPSUpgradeable, OwnableUpgradea

tokenInfo.isLocked = _isLocked;

emit SourceTokenLocked(_sourceToken, _isLocked);
emit SourceTokenStatusUpdated(_sourceToken, _isLocked);
}

/// @notice Set deposit limit
function setDepositLimit(address _sourceToken, uint256 _limit) external onlyOwner {
SourceTokenInfo storage tokenInfo = sourceTokenInfoMap[_sourceToken];
require(tokenInfo.isSupported, "Source token is not supported");
require(_limit >= tokenInfo.balance, "Invalid Specification");

tokenInfo.depositLimit = _limit;

Expand Down
44 changes: 28 additions & 16 deletions test/MergeTokenPortal.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -86,23 +86,25 @@ describe('MergeToeknPortal', function () {
erc20MergeAddr2Token = await ERC20TokenFactory.deploy(mergeTokenPortal.target, 'user2', 'ADD2TK', 18);
});
it('Should add source token correctly', async function () {
await mergeTokenPortal.connect(owner).addSourceToken(ownerAddr, recipientAddr, 100);
const sourceTokenInfo = await mergeTokenPortal.connect(owner).getSourceTokenInfos(ownerAddr);
await mergeTokenPortal.connect(owner).addSourceToken(erc20MergeAddr1Token, erc20MergeAddr2Token, 100);
const sourceTokenInfo = await mergeTokenPortal.connect(owner).getSourceTokenInfos(erc20MergeAddr1Token);
expect(sourceTokenInfo.isSupported).to.equal(true);
expect(sourceTokenInfo.isLocked).to.equal(false);
expect(sourceTokenInfo.mergeToken).to.equal(recipientAddr);
expect(sourceTokenInfo.mergeToken).to.equal(erc20MergeAddr2Token);
expect(sourceTokenInfo.balance).to.equal(0n);
expect(sourceTokenInfo.depositLimit).to.equal(100n);
});

it('Should add source token correctly2', async function () {
await mergeTokenPortal.connect(owner).addSourceToken(ownerAddr, recipientAddr, 100);
await mergeTokenPortal.removeSourceToken(ownerAddr);
await mergeTokenPortal.connect(owner).addSourceToken(ownerAddr, recipientAddr, 1000);
const sourceTokenInfo = await mergeTokenPortal.connect(owner).getSourceTokenInfos(ownerAddr);
await mergeTokenPortal.connect(owner).addSourceToken(erc20MergeAddr1Token, erc20MergeAddr2Token, 100);
await mergeTokenPortal.connect(owner).removeSourceToken(erc20MergeAddr1Token);
expect(await mergeTokenPortal.isMergeTokenSupported(erc20MergeAddr2Token, erc20MergeAddr1Token)).to.equal(false);

await mergeTokenPortal.connect(owner).addSourceToken(erc20MergeAddr1Token, erc20MergeAddr2Token, 1000);
const sourceTokenInfo = await mergeTokenPortal.connect(owner).getSourceTokenInfos(erc20MergeAddr1Token);
expect(sourceTokenInfo.isSupported).to.equal(true);
expect(sourceTokenInfo.isLocked).to.equal(false);
expect(sourceTokenInfo.mergeToken).to.equal(recipientAddr);
expect(sourceTokenInfo.mergeToken).to.equal(erc20MergeAddr2Token);
expect(sourceTokenInfo.balance).to.equal(0n);
expect(sourceTokenInfo.depositLimit).to.equal(1000n);
});
Expand All @@ -116,29 +118,39 @@ describe('MergeToeknPortal', function () {
});

it('Should allow owner remove a source token', async function () {
await mergeTokenPortal.connect(owner).addSourceToken(ownerAddr, recipientAddr, 100);
await mergeTokenPortal.connect(owner).removeSourceToken(ownerAddr);
const sourceTokenInfo = await mergeTokenPortal.getSourceTokenInfos(ownerAddr);
await mergeTokenPortal.connect(owner).addSourceToken(erc20MergeAddr1Token, erc20MergeAddr2Token, 100);
let sourceTokenInfo = await mergeTokenPortal.getSourceTokenInfos(erc20MergeAddr1Token);
expect(sourceTokenInfo.isSupported).to.equal(true);
expect(sourceTokenInfo.isLocked).to.equal(false);
expect(sourceTokenInfo.balance).to.equal(0n);
expect(sourceTokenInfo.depositLimit).to.equal(100n);
expect(sourceTokenInfo.mergeToken).to.equal(erc20MergeAddr2Token);
expect(await mergeTokenPortal.isMergeTokenSupported(erc20MergeAddr2Token, erc20MergeAddr1Token)).to.equal(true);

await mergeTokenPortal.connect(owner).removeSourceToken(erc20MergeAddr1Token);
sourceTokenInfo = await mergeTokenPortal.getSourceTokenInfos(erc20MergeAddr1Token);
expect(sourceTokenInfo.isSupported).to.equal(false);
expect(sourceTokenInfo.isLocked).to.equal(false);
expect(sourceTokenInfo.balance).to.equal(0n);
expect(sourceTokenInfo.depositLimit).to.equal(0n);
expect(sourceTokenInfo.mergeToken).to.equal(ZERO_ADDRESS);
expect(await mergeTokenPortal.isMergeTokenSupported(erc20MergeAddr2Token, erc20MergeAddr1Token)).to.equal(false);
});

it('Should allow commitee remove a source token2', async function () {
await mergeTokenPortal.connect(owner).addSourceToken(ownerAddr, recipientAddr, 100);
await mergeTokenPortal.connect(commitee).removeSourceToken(ownerAddr);
await mergeTokenPortal.connect(owner).addSourceToken(erc20MergeAddr1Token, erc20MergeAddr2Token, 100);
await mergeTokenPortal.connect(commitee).removeSourceToken(erc20MergeAddr1Token);

const sourceTokenInfo = await mergeTokenPortal.getSourceTokenInfos(ownerAddr);
const sourceTokenInfo = await mergeTokenPortal.getSourceTokenInfos(erc20MergeAddr1Token);
expect(sourceTokenInfo.isSupported).to.equal(false);
expect(sourceTokenInfo.isLocked).to.equal(false);
expect(sourceTokenInfo.balance).to.equal(0n);
expect(sourceTokenInfo.depositLimit).to.equal(0n);
});

it('Should not allow non-owner or non-commitee remove a source token', async function () {
await mergeTokenPortal.connect(owner).addSourceToken(ownerAddr, recipientAddr, 100);
await expect(mergeTokenPortal.connect(user1).removeSourceToken(ownerAddr)).to.be.revertedWith(
await mergeTokenPortal.connect(owner).addSourceToken(erc20MergeAddr1Token, erc20MergeAddr2Token, 100);
await expect(mergeTokenPortal.connect(user1).removeSourceToken(erc20MergeAddr1Token)).to.be.revertedWith(
'Only owner or commitee can call this function',
);
});
Expand Down

0 comments on commit 2e50cb5

Please sign in to comment.