Audit 3.4 & 3.14 Fix ST upgrades (#694)

adamdossa · maxsam4 · commit 045ccc12a975 · 2019-06-12T13:22:18.000+05:30
* Fixes + test cases

* Small change in test

* Add fixes for 3.14
diff --git a/contracts/mocks/MockCountTransferManager.sol b/contracts/mocks/MockCountTransferManager.sol
@@ -8,6 +8,7 @@ import "../modules/TransferManager/CTM/CountTransferManager.sol";
 contract MockCountTransferManager is CountTransferManager {
 
     event Upgrader(uint256 _someData);
+    uint256 public someValue;
 
     /**
      * @notice Constructor
@@ -19,6 +20,7 @@ contract MockCountTransferManager is CountTransferManager {
 
     function initialize(uint256 _someData) public {
         require(msg.sender == address(this));
+        someValue = _someData;
         emit Upgrader(_someData);
     }
 
diff --git a/contracts/mocks/MockSecurityTokenLogic.sol b/contracts/mocks/MockSecurityTokenLogic.sol
@@ -15,6 +15,7 @@ import "../tokens/SecurityToken.sol";
 contract MockSecurityTokenLogic is SecurityToken {
 
     event UpgradeEvent(uint256 _upgrade);
+    uint256 public someValue;
 
     /**
      * @notice Initialization function
@@ -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);
     }
diff --git a/contracts/modules/UpgradableModuleFactory.sol b/contracts/modules/UpgradableModuleFactory.sol
@@ -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,
@@ -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);
     }
 
     /**
diff --git a/contracts/tokens/STFactory.sol b/contracts/tokens/STFactory.sol
@@ -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
@@ -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);
@@ -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)
@@ -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
diff --git a/test/d_count_transfer_manager.js b/test/d_count_transfer_manager.js
@@ -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;
@@ -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";
@@ -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,
@@ -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
@@ -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]);
diff --git a/test/o_security_token.js b/test/o_security_token.js
