Upgradable tokens

contracts/ModuleRegistry.sol

@@ -119,6 +119,7 @@ contract ModuleRegistry is IModuleRegistry, EternalStorage {
      * @dev Any module can be added during token creation without being registered if it is defined in the token proxy deployment contract
      * @dev The feature switch for custom modules is labelled "customModulesAllowed"
      * @param _moduleFactory is the address of the relevant module factory
+     * @param _isUpgrade whether or not the function is being called as a result of an upgrade
      */
     function useModule(address _moduleFactory, bool _isUpgrade) external {
         if (IFeatureRegistry(getAddressValue(Encoder.getKey("featureRegistry"))).getFeatureStatus("customModulesAllowed")) {
@@ -129,12 +130,16 @@ contract ModuleRegistry is IModuleRegistry, EternalStorage {
         } else {
             require(getBoolValue(Encoder.getKey("verified", _moduleFactory)), "ModuleFactory must be verified");
         }
-        require(isCompatibleModule(_moduleFactory, msg.sender), "Incompatible versions");
         // This if statement is required to be able to add modules from the STFactory contract during deployment
         // before the token has been registered to the STR.
-        if ((!_isUpgrade) && ISecurityTokenRegistry(getAddressValue(Encoder.getKey("securityTokenRegistry"))).isSecurityToken(msg.sender)) {
-            pushArray(Encoder.getKey("reputation", _moduleFactory), msg.sender);
-            emit ModuleUsed(_moduleFactory, msg.sender);
+        if (ISecurityTokenRegistry(getAddressValue(Encoder.getKey("securityTokenRegistry"))).isSecurityToken(msg.sender)
+            || (ISecurityTokenRegistry(getAddressValue(Encoder.getKey("securityTokenRegistry"))).getSTFactoryAddress() == msg.sender)
+        ) {
+            require(isCompatibleModule(_moduleFactory, msg.sender), "Incompatible versions");
+            if (!_isUpgrade) {
+                pushArray(Encoder.getKey("reputation", _moduleFactory), msg.sender);
+                emit ModuleUsed(_moduleFactory, msg.sender);
+            }
         }
     }
 

contracts/mocks/STFactoryMock.sol

@@ -29,7 +29,7 @@ contract STFactoryMock is ISTFactory {
         string calldata _tokenDetails,
         address _issuer,
         bool _divisible,
-        address _treasuryWallet,
+        address /* _treasuryWallet */,
         address _polymathRegistry
     )
         external

contracts/mocks/SecurityTokenMock.sol

@@ -24,8 +24,8 @@ contract SecurityTokenMock is SecurityToken {
         string memory _name,
         string memory _symbol,
         uint8 _decimals,
-        uint256 _granularity,
-        string memory _tokenDetails,
+        uint256 /* _granularity */,
+        string memory /* _tokenDetails */,
         address _polymathRegistry,
         address _delegate
     )

contracts/tokens/STFactory.sol

@@ -4,6 +4,7 @@ import "./SecurityTokenProxy.sol";
 import "../proxy/OwnedUpgradeabilityProxy.sol";
 import "../interfaces/ISTFactory.sol";
 import "../interfaces/ISecurityToken.sol";
+import "../interfaces/IPolymathRegistry.sol";
 import "../interfaces/IOwnable.sol";
 import "openzeppelin-solidity/contracts/ownership/Ownable.sol";
 import "../interfaces/IModuleRegistry.sol";
@@ -17,12 +18,11 @@ contract STFactory is ISTFactory, Ownable {
 
     address public transferManagerFactory;
     DataStoreFactory public dataStoreFactory;
-    // Mapping from Security Token address to token version
-    mapping (address => uint256) tokenVersion;
-    // Mapping from Security Token address to registry
-    // We do this so that we can use the registry to check module compatibility on upgrades
-    // And as a secondary benefit that a token was deployed initially through this factory
-    mapping (address => address) tokenToRegistry;
+    IPolymathRegistry public polymathRegistry;
+
+    // Mapping from Security Token address to token upgrade version.
+    // A mapping to 0 means a token has not yet been deployed
+    mapping (address => uint256) tokenUpgrade;
 
     struct LogicContract {
         string version;
@@ -43,12 +43,17 @@ contract STFactory is ISTFactory, Ownable {
     event DefaultTransferManagerUpdated(address indexed _oldTransferManagerFactory, address indexed _newTransferManagerFactory);
     event DefaultDataStoreUpdated(address indexed _oldDataStoreFactory, address indexed _newDataStoreFactory);
 
-    constructor(address _transferManagerFactory, address _dataStoreFactory, string memory _version, address _logicContract, bytes memory _initializationData) public {
+    constructor(address _polymathRegistry, address _transferManagerFactory, address _dataStoreFactory, string memory _version, address _logicContract, bytes memory _initializationData) 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");
         transferManagerFactory = _transferManagerFactory;
         dataStoreFactory = DataStoreFactory(_dataStoreFactory);
+        polymathRegistry = IPolymathRegistry(_polymathRegistry);
+
+        // Start at 1 so that we can distinguish deployed tokens in tokenUpgrade
+        latestUpgrade = 1;
         logicContracts[latestUpgrade].logicContract = _logicContract;
         logicContracts[latestUpgrade].initializationData = _initializationData;
         logicContracts[latestUpgrade].version = _version;
@@ -66,7 +71,7 @@ contract STFactory is ISTFactory, Ownable {
         address _issuer,
         bool _divisible,
         address _treasuryWallet,
-        address _polymathRegistry
+        address /* _polymathRegistry */
     )
         external
         returns(address)
@@ -76,8 +81,7 @@ contract STFactory is ISTFactory, Ownable {
             _symbol,
             _decimals,
             _tokenDetails,
-            _divisible,
-            _polymathRegistry
+            _divisible
         );
         //NB When dataStore is generated, the security token address is automatically set via the constructor in DataStoreProxy.
         if (address(dataStoreFactory) != address(0)) {
@@ -96,8 +100,7 @@ contract STFactory is ISTFactory, Ownable {
         string memory _symbol,
         uint8 _decimals,
         string memory _tokenDetails,
-        bool _divisible,
-        address _polymathRegistry
+        bool _divisible
     ) internal returns(address) {
         // Creates proxy contract and sets some initial storage
         SecurityTokenProxy proxy = new SecurityTokenProxy(
@@ -106,15 +109,15 @@ contract STFactory is ISTFactory, Ownable {
             _decimals,
             _divisible ? 1 : uint256(10) ** _decimals,
             _tokenDetails,
-            _polymathRegistry
+            address(polymathRegistry)
         );
         // Sets logic contract
         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)
         (bool success, ) = address(proxy).call(logicContracts[latestUpgrade].initializationData);
         require(success, "Unsuccessful initialization");
-        tokenToRegistry[address(proxy)] = _polymathRegistry;
+        tokenUpgrade[address(proxy)] = latestUpgrade;
         return address(proxy);
     }
 
@@ -141,13 +144,13 @@ contract STFactory is ISTFactory, Ownable {
      */
     function upgradeToken(uint8 _maxModuleType) external {
         // Check the token was created by this factory
-        require(tokenToRegistry[msg.sender] != address(0), "Invalid token");
-        uint256 newVersion = tokenVersion[msg.sender] + 1;
+        require(tokenUpgrade[msg.sender] != 0, "Invalid token");
+        uint256 newVersion = tokenUpgrade[msg.sender] + 1;
         require(newVersion <= latestUpgrade, "Incorrect version");
         OwnedUpgradeabilityProxy(address(uint160(msg.sender))).upgradeToAndCall(logicContracts[newVersion].version, logicContracts[newVersion].logicContract, logicContracts[newVersion].upgradeData);
-        tokenVersion[msg.sender] = newVersion;
+        tokenUpgrade[msg.sender] = newVersion;
         // Check that all modules remain valid
-        IModuleRegistry moduleRegistry = IModuleRegistry(IPolymathRegistry(tokenToRegistry[msg.sender]).getAddress("ModuleRegistry"));
+        IModuleRegistry moduleRegistry = IModuleRegistry(polymathRegistry.getAddress("ModuleRegistry"));
         address moduleFactory;
         bool isArchived;
         for (uint8 i = 1; i < _maxModuleType; i++) {

migrations/2_deploy_contracts.js

@@ -372,7 +372,7 @@ module.exports = function(deployer, network, accounts) {
         .then(() => {
             // H) Deploy the STVersionProxy001 Contract which contains the logic of deployment of securityToken.
             let tokenInitBytesCall = web3.eth.abi.encodeFunctionCall(tokenInitBytes, [STGetter.address]);
-            return deployer.deploy(STFactory, GeneralTransferManagerFactory.address, DataStoreFactory.address, "3.0.0", SecurityTokenLogic.address, tokenInitBytesCall, { from: PolymathAccount });
+            return deployer.deploy(STFactory, polymathRegistry.address, GeneralTransferManagerFactory.address, DataStoreFactory.address, "3.0.0", SecurityTokenLogic.address, tokenInitBytesCall, { from: PolymathAccount });
         })
         .then(() => {
             // K) Deploy the FeatureRegistry contract to control feature switches
@@ -407,6 +407,18 @@ module.exports = function(deployer, network, accounts) {
                 from: PolymathAccount
             });
         })
+        .then(() => {
+            return SecurityTokenRegistry.at(SecurityTokenRegistryProxy.address);
+        })
+        .then((securityTokenRegistry) => {
+            return securityTokenRegistry.setProtocolFactory(STFactory.address, 3, 0, 0);
+        })
+        .then(() => {
+            return SecurityTokenRegistry.at(SecurityTokenRegistryProxy.address);
+        })
+        .then((securityTokenRegistry) => {
+            return securityTokenRegistry.setLatestVersion(3, 0, 0);
+        })
         .then(() => {
             // Assign the address into the SecurityTokenRegistry key
             return polymathRegistry.changeAddress("SecurityTokenRegistry", SecurityTokenRegistryProxy.address, { from: PolymathAccount });

test/b_capped_sto.js

@@ -6,6 +6,7 @@ import { setUpPolymathNetwork, deployGPMAndVerifyed, deployCappedSTOAndVerifyed,
 import { catchRevert } from "./helpers/exceptions";
 
 const CappedSTOFactory = artifacts.require("./CappedSTOFactory.sol");
+const STFactory = artifacts.require("./STFactory.sol");
 const CappedSTO = artifacts.require("./CappedSTO.sol");
 const DummySTO = artifacts.require("./DummySTO.sol");
 const SecurityToken = artifacts.require("./SecurityToken.sol");
@@ -177,6 +178,10 @@ contract("CappedSTO", async (accounts) => {
 
         it("Should generate the new security token with the same symbol as registered above", async () => {
             await I_PolyToken.approve(I_STRProxied.address, initRegFee, { from: token_owner });
+            let t = await I_STRGetter.getSTFactoryAddress.call();
+            console.log(t);
+            let foo = await STFactory.at(t);
+            console.log(await foo.polymathRegistry.call());
 
             let tx = await I_STRProxied.generateSecurityToken(name, symbol, tokenDetails, false, treasury_wallet, 0, { from: token_owner });
 
@@ -849,7 +854,7 @@ contract("CappedSTO", async (accounts) => {
                 let endTime = startTime + duration.days(30);
                 const cap = web3.utils.toWei("10000");
                 const dummyBytesSig = encodeModuleCall(DummySTOParameters, [startTime, endTime, cap, "Hello"]);
-                const tx = await I_SecurityToken_ETH.addModule(I_DummySTOFactory.address, dummyBytesSig, maxCost, new BN(0), { from: token_owner });
+                const tx = await I_SecurityToken_ETH.addModule(I_DummySTOFactory.address, dummyBytesSig, maxCost, new BN(0), false, { from: token_owner });
                 console.log(tx.logs[2]);
                 assert.equal(tx.logs[2].args._types[0], stoKey, `Wrong module type added`);
                 assert.equal(

test/helpers/createInstances.js

@@ -253,7 +253,7 @@ async function deploySTFactory(account_polymath) {
         ]
     };
     let tokenInitBytesCall = web3.eth.abi.encodeFunctionCall(tokenInitBytes, [I_STGetter.address]);
-    I_STFactory = await STFactory.new(I_GeneralTransferManagerFactory.address, I_DataStoreFactory.address, "3.0.0", I_SecurityToken.address, tokenInitBytesCall, { from: account_polymath });
+    I_STFactory = await STFactory.new(I_PolymathRegistry.address, I_GeneralTransferManagerFactory.address, I_DataStoreFactory.address, "3.0.0", I_SecurityToken.address, tokenInitBytesCall, { from: account_polymath });
 
     assert.notEqual(I_STFactory.address.valueOf(), "0x0000000000000000000000000000000000000000", "STFactory contract was not deployed");
 

test/u_module_registry_proxy.js

@@ -177,7 +177,7 @@ contract("ModuleRegistryProxy", async (accounts) => {
                 ]
             };
             let tokenInitBytesCall = web3.eth.abi.encodeFunctionCall(tokenInitBytes, [I_STGetter.address]);
-            I_STFactory = await STFactory.new(I_GeneralTransferManagerFactory.address, I_DataStoreFactory.address, "3.0.0", I_SecurityTokenLogic.address, tokenInitBytesCall, { from: account_polymath });
+            I_STFactory = await STFactory.new(I_PolymathRegistry.address, I_GeneralTransferManagerFactory.address, I_DataStoreFactory.address, "3.0.0", I_SecurityTokenLogic.address, tokenInitBytesCall, { from: account_polymath });
 
             assert.notEqual(I_STFactory.address.valueOf(), address_zero, "STFactory contract was not deployed");
         });