Skip to content

Commit

Permalink
Audit 3.4 & 3.14 Fix ST upgrades (#694)
Browse files Browse the repository at this point in the history
* Fixes + test cases

* Small change in test

* Add fixes for 3.14
  • Loading branch information
adamdossa authored and maxsam4 committed Jun 12, 2019
1 parent 004194e commit 045ccc1
Show file tree
Hide file tree
Showing 6 changed files with 244 additions and 17 deletions.
2 changes: 2 additions & 0 deletions contracts/mocks/MockCountTransferManager.sol
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import "../modules/TransferManager/CTM/CountTransferManager.sol";
contract MockCountTransferManager is CountTransferManager {

event Upgrader(uint256 _someData);
uint256 public someValue;

/**
* @notice Constructor
Expand All @@ -19,6 +20,7 @@ contract MockCountTransferManager is CountTransferManager {

function initialize(uint256 _someData) public {
require(msg.sender == address(this));
someValue = _someData;
emit Upgrader(_someData);
}

Expand Down
18 changes: 18 additions & 0 deletions contracts/mocks/MockSecurityTokenLogic.sol
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import "../tokens/SecurityToken.sol";
contract MockSecurityTokenLogic is SecurityToken {

event UpgradeEvent(uint256 _upgrade);
uint256 public someValue;

/**
* @notice Initialization function
Expand All @@ -23,10 +24,27 @@ contract MockSecurityTokenLogic is SecurityToken {
*/
function upgrade(address _getterDelegate, uint256 _upgrade) external {
getterDelegate = _getterDelegate;
someValue = _upgrade;
//securityTokenVersion = SemanticVersion(3, 1, 0);
emit UpgradeEvent(_upgrade);
}

/**
* @notice Initialization function
* @dev Expected to be called atomically with the proxy being created, by the owner of the token
* @dev Can only be called once
*/
function initialize(address _getterDelegate, uint256 _someValue) public {
//Expected to be called atomically with the proxy being created
require(!initialized, "Already initialized");
getterDelegate = _getterDelegate;
securityTokenVersion = SemanticVersion(3, 0, 0);
updateFromRegistry();
tokenFactory = msg.sender;
initialized = true;
someValue = _someValue;
}

function newFunction(uint256 _upgrade) external {
emit UpgradeEvent(_upgrade);
}
Expand Down
33 changes: 28 additions & 5 deletions contracts/modules/UpgradableModuleFactory.sol
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ import "../proxy/OwnedUpgradeabilityProxy.sol";
*/
contract UpgradableModuleFactory is ModuleFactory {

event LogicContractSet(string _version, address _logicContract, bytes _upgradeData);
event LogicContractSet(string _version, uint256 _upgrade, address _logicContract, bytes _upgradeData);

event ModuleUpgraded(
address indexed _module,
Expand Down Expand Up @@ -68,12 +68,35 @@ contract UpgradableModuleFactory is ModuleFactory {
require(_logicContract != logicContracts[latestUpgrade].logicContract, "Same version");
require(_logicContract != address(0), "Invalid address");
latestUpgrade++;
logicContracts[latestUpgrade].version = _version;
logicContracts[latestUpgrade].logicContract = _logicContract;
logicContracts[latestUpgrade].upgradeData = _upgradeData;
_modifyLogicContract(latestUpgrade, _version, _logicContract, _upgradeData);
}

/**
* @notice Used to update an existing token logic contract
* @param _upgrade logic contract to upgrade
* @param _version Version of upgraded module
* @param _logicContract Address of deployed module logic contract referenced from proxy
* @param _upgradeData Data to be passed in call to upgradeToAndCall when a token upgrades its module
*/
function updateLogicContract(uint256 _upgrade, string calldata _version, address _logicContract, bytes calldata _upgradeData) external onlyOwner {
require(_upgrade <= latestUpgrade, "Invalid upgrade");
// version & contract must differ from previous version, otherwise upgrade proxy will fail
if (_upgrade > 0) {
require(keccak256(abi.encodePacked(_version)) != keccak256(abi.encodePacked(logicContracts[_upgrade - 1].version)), "Same version");
require(_logicContract != logicContracts[_upgrade - 1].logicContract, "Same version");
}
require(_logicContract != address(0), "Invalid address");
require(_upgradeData.length > 4, "Invalid Upgrade");
_modifyLogicContract(_upgrade, _version, _logicContract, _upgradeData);
}

function _modifyLogicContract(uint256 _upgrade, string memory _version, address _logicContract, bytes memory _upgradeData) internal {
logicContracts[_upgrade].version = _version;
logicContracts[_upgrade].logicContract = _logicContract;
logicContracts[_upgrade].upgradeData = _upgradeData;
IModuleRegistry moduleRegistry = IModuleRegistry(IPolymathRegistry(polymathRegistry).getAddress("ModuleRegistry"));
moduleRegistry.unverifyModule(address(this));
emit LogicContractSet(_version, _logicContract, _upgradeData);
emit LogicContractSet(_version, _upgrade, _logicContract, _upgradeData);
}

/**
Expand Down
47 changes: 39 additions & 8 deletions contracts/tokens/STFactory.sol
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ contract STFactory is ISTFactory, Ownable {

uint256 public latestUpgrade;

event LogicContractSet(string _version, address _logicContract, bytes _upgradeData);
event LogicContractSet(string _version, uint256 _upgrade, address _logicContract, bytes _initializationData, bytes _upgradeData);
event TokenUpgraded(
address indexed _securityToken,
uint256 indexed _version
Expand All @@ -50,13 +50,14 @@ contract STFactory is ISTFactory, Ownable {
string memory _version,
address _logicContract,
bytes memory _initializationData
)
public
)
public
{
require(_logicContract != address(0), "Invalid Address");
require(_transferManagerFactory != address(0), "Invalid Address");
require(_dataStoreFactory != address(0), "Invalid Address");
require(_polymathRegistry != address(0), "Invalid Address");
require(_initializationData.length > 4, "Invalid Initialization");
transferManagerFactory = _transferManagerFactory;
dataStoreFactory = DataStoreFactory(_dataStoreFactory);
polymathRegistry = IPolymathRegistry(_polymathRegistry);
Expand Down Expand Up @@ -121,6 +122,7 @@ contract STFactory is ISTFactory, Ownable {
address(polymathRegistry)
);
// Sets logic contract
emit Log(latestUpgrade);
proxy.upgradeTo(logicContracts[latestUpgrade].version, logicContracts[latestUpgrade].logicContract);
// Initialises security token contract - needed for functions that can only be called by the
// owner of the contract, or are specific to this particular logic contract (e.g. setting version)
Expand All @@ -130,23 +132,52 @@ contract STFactory is ISTFactory, Ownable {
return address(proxy);
}

event Log(uint256 _upgrade);

/**
* @notice Used to set a new token logic contract
* @param _version Version of upgraded module
* @param _logicContract Address of deployed module logic contract referenced from proxy
* @param _upgradeData Data to be passed in call to upgradeToAndCall when a token upgrades its module
*/
function setLogicContract(string calldata _version, address _logicContract, bytes calldata _upgradeData) external onlyOwner {
function setLogicContract(string calldata _version, address _logicContract, bytes calldata _initializationData, bytes calldata _upgradeData) external onlyOwner {
require(keccak256(abi.encodePacked(_version)) != keccak256(abi.encodePacked(logicContracts[latestUpgrade].version)), "Same version");
require(_logicContract != logicContracts[latestUpgrade].logicContract, "Same version");
require(_logicContract != address(0), "Invalid address");
require(_initializationData.length > 4, "Invalid Initialization");
require(_upgradeData.length > 4, "Invalid Upgrade");
latestUpgrade++;
logicContracts[latestUpgrade].version = _version;
logicContracts[latestUpgrade].logicContract = _logicContract;
logicContracts[latestUpgrade].upgradeData = _upgradeData;
emit LogicContractSet(_version, _logicContract, _upgradeData);
_modifyLogicContract(latestUpgrade, _version, _logicContract, _initializationData, _upgradeData);
}

/**
* @notice Used to update an existing token logic contract
* @param _upgrade logic contract to upgrade
* @param _version Version of upgraded module
* @param _logicContract Address of deployed module logic contract referenced from proxy
* @param _upgradeData Data to be passed in call to upgradeToAndCall when a token upgrades its module
*/
function updateLogicContract(uint256 _upgrade, string calldata _version, address _logicContract, bytes calldata _initializationData, bytes calldata _upgradeData) external onlyOwner {
require(_upgrade <= latestUpgrade, "Invalid upgrade");
require(_upgrade > 0, "Invalid upgrade");
// version & contract must differ from previous version, otherwise upgrade proxy will fail
if (_upgrade > 1) {
require(keccak256(abi.encodePacked(_version)) != keccak256(abi.encodePacked(logicContracts[_upgrade - 1].version)), "Same version");
require(_logicContract != logicContracts[_upgrade - 1].logicContract, "Same version");
}
require(_logicContract != address(0), "Invalid address");
require(_initializationData.length > 4, "Invalid Initialization");
require(_upgradeData.length > 4, "Invalid Upgrade");
_modifyLogicContract(_upgrade, _version, _logicContract, _initializationData, _upgradeData);
}

function _modifyLogicContract(uint256 _upgrade, string memory _version, address _logicContract, bytes memory _initializationData, bytes memory _upgradeData) internal {
logicContracts[_upgrade].version = _version;
logicContracts[_upgrade].logicContract = _logicContract;
logicContracts[_upgrade].upgradeData = _upgradeData;
logicContracts[_upgrade].initializationData = _initializationData;
emit LogicContractSet(_version, _upgrade, _logicContract, _initializationData, _upgradeData);
}
/**
* @notice Used to upgrade a token
* @param _maxModuleType maximum module type enumeration
Expand Down
66 changes: 65 additions & 1 deletion test/d_count_transfer_manager.js
Original file line number Diff line number Diff line change
Expand Up @@ -37,8 +37,10 @@ contract("CountTransferManager", async (accounts) => {
let I_GeneralPermissionManager;
let I_CountTransferManager;
let I_CountTransferManager2;
let I_CountTransferManager3;
let I_GeneralTransferManager;
let I_GeneralTransferManager2;
let I_GeneralTransferManager3;
let I_ExchangeTransferManager;
let I_ModuleRegistry;
let I_ModuleRegistryProxy;
Expand All @@ -49,17 +51,21 @@ contract("CountTransferManager", async (accounts) => {
let I_STFactory;
let I_SecurityToken;
let I_SecurityToken2;
let I_SecurityToken3;
let I_PolyToken;
let I_PolymathRegistry;
let I_STRGetter;
let I_STGetter;
let I_MockCountTransferManagerLogic;
let stGetter;
let stGetter2;
let stGetter3;

// SecurityToken Details
const name = "Team";
const symbol = "sap";
const symbol2 = "sapp";
const symbol3 = "sapp3"
const tokenDetails = "This is equity type of issuance";
const decimals = 18;
const contact = "team@polymath.network";
Expand Down Expand Up @@ -373,6 +379,21 @@ contract("CountTransferManager", async (accounts) => {
I_GeneralTransferManager2 = await GeneralTransferManager.at(moduleData);
});

it("deploy another new token & auto attach modules", async () => {
//register ticker and deploy token
await I_PolyToken.approve(I_STRProxied.address, initRegFee, { from: token_owner });
let tx = await I_STRProxied.registerTicker(token_owner, symbol3, contact, { from: token_owner });

await I_PolyToken.approve(I_STRProxied.address, initRegFee, { from: token_owner });

let tx2 = await I_STRProxied.generateNewSecurityToken(name, symbol3, tokenDetails, false, token_owner, 0, { from: token_owner });

I_SecurityToken3 = await SecurityToken.at(tx2.logs[1].args._securityTokenAddress);
stGetter3 = await STGetter.at(I_SecurityToken3.address);
let moduleData = (await stGetter3.getModulesByType(2))[0];
I_GeneralTransferManager3 = await GeneralTransferManager.at(moduleData);
});

it("add 3 holders to the token", async () => {
await I_GeneralTransferManager2.modifyKYCData(
account_investor1,
Expand Down Expand Up @@ -455,8 +476,21 @@ contract("CountTransferManager", async (accounts) => {
console.log("current max holder number is " + (await I_CountTransferManager2.maxHolderCount({ from: token_owner })));
});

it("Should successfully attach the CountTransferManager with the third security token and set max holder to 2", async () => {
const tx = await I_SecurityToken3.addModule(I_CountTransferManagerFactory.address, bytesSTO, new BN(0), new BN(0), false, { from: token_owner });
assert.equal(tx.logs[2].args._types[0].toNumber(), transferManagerKey, "CountTransferManager doesn't get deployed");
assert.equal(
web3.utils.toAscii(tx.logs[2].args._name).replace(/\u0000/g, ""),
"CountTransferManager",
"CountTransferManager module was not added"
);
I_CountTransferManager3 = await CountTransferManager.at(tx.logs[2].args._module);
await I_CountTransferManager3.changeHolderCount(2, { from: token_owner });
console.log("current max holder number is " + (await I_CountTransferManager3.maxHolderCount({ from: token_owner })));
});

it("Should upgrade the CTM", async () => {
let I_MockCountTransferManagerLogic = await MockCountTransferManager.new("0x0000000000000000000000000000000000000000", "0x0000000000000000000000000000000000000000", { from: account_polymath });
I_MockCountTransferManagerLogic = await MockCountTransferManager.new("0x0000000000000000000000000000000000000000", "0x0000000000000000000000000000000000000000", { from: account_polymath });
let bytesCM = encodeProxyCall(["uint256"], [11]);
await catchRevert(
// Fails as no upgrade available
Expand All @@ -482,9 +516,39 @@ contract("CountTransferManager", async (accounts) => {
let tx = await I_MRProxied.verifyModule(I_CountTransferManagerFactory.address, { from: account_polymath });
await I_SecurityToken2.upgradeModule(I_CountTransferManager2.address, { from: token_owner });
let I_MockCountTransferManager = await MockCountTransferManager.at(I_CountTransferManager2.address);
let newValue = await I_MockCountTransferManager.someValue.call();
assert(newValue.toNumber(), 11);
await I_MockCountTransferManager.newFunction();
});

it("Should modify the upgrade data and upgrade", async () => {
let bytesCM = encodeProxyCall(["uint256"], [12]);
await catchRevert(
// Fails due to the same version being used
I_CountTransferManagerFactory.updateLogicContract(1, "3.0.0", I_MockCountTransferManagerLogic.address, bytesCM, { from: account_polymath })
);
await catchRevert(
// Fails due to the wrong contract being used
I_CountTransferManagerFactory.updateLogicContract(1, "4.0.0", "0x0000000000000000000000000000000000000000", bytesCM, { from: account_polymath })
);
await catchRevert(
// Fails due to the wrong owner being used
I_CountTransferManagerFactory.updateLogicContract(1, "4.0.0", "0x0000000000000000000000000000000000000000", bytesCM, { from: token_owner })
);
await I_CountTransferManagerFactory.updateLogicContract(1, "4.0.0", I_MockCountTransferManagerLogic.address, bytesCM, { from: account_polymath });
await catchRevert(
// Fails as upgraded module has been unverified
I_SecurityToken3.upgradeModule(I_CountTransferManager3.address, { from: token_owner })
);
let tx = await I_MRProxied.verifyModule(I_CountTransferManagerFactory.address, { from: account_polymath });
await I_SecurityToken3.upgradeModule(I_CountTransferManager3.address, { from: token_owner });
let I_MockCountTransferManager = await MockCountTransferManager.at(I_CountTransferManager3.address);
let newValue = await I_MockCountTransferManager.someValue.call();
assert(newValue.toNumber(), 12);
await I_MockCountTransferManager.newFunction();

});

it("Should upgrade the CTM again", async () => {
let I_MockCountTransferManagerLogic = await MockCountTransferManager.new("0x0000000000000000000000000000000000000000", "0x0000000000000000000000000000000000000000", { from: account_polymath });
let bytesCM = encodeProxyCall(["uint256"], [11]);
Expand Down
Loading

0 comments on commit 045ccc1

Please sign in to comment.