Merge pull request #22 from ayanworks/feat/audit-changes

refactor: security changes based on audit

contracts/PolygonDidRegistry.sol

@@ -1,5 +1,5 @@
 //SPDX-License-Identifier: Apache 2.0
-pragma solidity ^0.8.16;
+pragma solidity 0.8.16;
 
 /**
  *@title PolygonDidRegistry
@@ -15,13 +15,6 @@ contract PolygonDidRegistry {
         string didDoc;
     }
 
-    modifier onlyController(address _id) {
-        require(
-            polyDIDs[_id].controller == msg.sender,
-            "message sender is not the controller of the DID Doc"
-        );
-        _;
-    }
     mapping(address => PolyDID) polyDIDs;
     mapping(uint256 => address) activeDIDs;
     mapping(address => uint256) activeAddress;
@@ -30,43 +23,63 @@ contract PolygonDidRegistry {
     event DIDCreated(address id, string doc);
     event DIDUpdated(address id, string doc);
     event TransferOwnership(address newOwner);
-    event ResourceAdded(address _id, string _resourceId, string _resourcePayload);
+    event ResourceAdded(
+        address _id,
+        string _resourceId,
+        string _resourcePayload
+    );
 
     bool private initialized;
 
     /**
-     *@dev initializes the ownership of contract
+     *@dev modifiers for access control
      **/
 
-    function initialize() public {
-        require(!initialized, "Contract instance has already been initialized");
+    modifier onlyController(address _id) {
+        require(
+            polyDIDs[_id].controller == msg.sender,
+            'message sender is not the controller of the DID Doc'
+        );
+        _;
+    }
+
+    modifier nonReentrant() {
+        require(!initialized, 'Contract instance has already been initialized');
         initialized = true;
-        owner = msg.sender;
         totalDIDs = 0;
+        _;
     }
 
     modifier onlyOwner() {
-        require(msg.sender == owner, "message sender is not the owner");
+        require(msg.sender == owner, 'message sender is not the owner');
         _;
     }
 
+    /**
+     *@dev initializes the ownership of contract
+     **/
+
+    function initialize() external nonReentrant {
+        require(msg.sender != address(0), 'Invalid owner address');
+        owner = msg.sender;
+    }
+
     /**
      *@dev transfer the ownership of contract
      *@param _newOwner - Address of the new owner to whom the ownership needs to be passed
      **/
 
-    function transferOwnership(address _newOwner)
-        public
-        onlyOwner
-        returns (string memory)
-    {
-        if (owner != _newOwner) {
-            owner = _newOwner;
-            emit TransferOwnership(owner);
-            return ("Ownership transferred successfully");
-        } else {
-            return ("Ownership cannot be transferred to the same account");
-        }
+    function transferOwnership(
+        address _newOwner
+    ) external onlyOwner returns (string memory) {
+        require(_newOwner != address(0), 'Invalid owner address');
+        require(
+            owner != _newOwner,
+            'Ownership cannot be transferred to the same account'
+        );
+        owner = _newOwner;
+        emit TransferOwnership(owner);
+        return ('Ownership transferred successfully');
     }
 
     /**
@@ -83,15 +96,19 @@ contract PolygonDidRegistry {
      *@param _doc - A string object that holds the DID Doc
      */
 
-    function createDID(address _id, string memory _doc)
-        public
+    function createDID(
+        address _id,
+        string memory _doc
+    )
+        external
         returns (
             address controller,
             uint256 created,
             uint256 updated,
             string memory didDoc
         )
     {
+        require(_id != address(0), 'Invalid address provided');
         polyDIDs[_id].controller = msg.sender;
         polyDIDs[_id].created = block.timestamp;
         polyDIDs[_id].updated = block.timestamp;
@@ -113,7 +130,9 @@ contract PolygonDidRegistry {
      *@param _id - Address that refers to the DID doc position
      */
 
-    function getDIDDoc(address _id) public view returns (string memory, string[] memory) {
+    function getDIDDoc(
+        address _id
+    ) public view returns (string memory, string[] memory) {
         string[] memory result = new string[](keysById[_id].length);
 
         for (uint256 i = 0; i < keysById[_id].length; i++) {
@@ -126,26 +145,19 @@ contract PolygonDidRegistry {
      *@dev Reads total number of DIDs and total number of active DIDs from Chain
      */
 
-    function getTotalNumberOfDIDs()
-        public
-        view
-        returns (uint256 _totalDIDs)
-    {
+    function getTotalNumberOfDIDs() public view returns (uint256 _totalDIDs) {
         return (totalDIDs);
     }
 
-
     /**
      *@dev Reads one DID at a time from Chain based on index
      *@param _index - Uint256 type variable that refers to the DID position
      *@return _did - returns the DID Doc assciated with the index. Returns null if the DID Doc is deleted.
      */
 
-    function getDIDDocByIndex(uint256 _index)
-        public
-        view
-        returns (string memory)
-    {
+    function getDIDDocByIndex(
+        uint256 _index
+    ) public view returns (string memory) {
         return polyDIDs[activeDIDs[_index]].didDoc;
     }
 
@@ -155,8 +167,11 @@ contract PolygonDidRegistry {
      *@param _doc - A String that holds the DID doc
      */
 
-    function updateDIDDoc(address _id, string memory _doc)
-        public
+    function updateDIDDoc(
+        address _id,
+        string memory _doc
+    )
+        external
         onlyController(_id)
         returns (
             address controller,
@@ -181,20 +196,32 @@ contract PolygonDidRegistry {
      *@param _id - Address that refers to the DID doc
      *@param _resourceId - Id that refers to the resource
      */
-    function addResource(address _id, string memory _resourceId, string memory _resourcePayload) public onlyController(_id) returns (address, string memory, string memory) {
-    resourceData[_id][_resourceId] = _resourcePayload;
-    keysById[_id].push(_resourceId);
-    emit ResourceAdded(_id, _resourceId, _resourcePayload);
-    return (_id, _resourceId, _resourcePayload);
-}
-
-     /**
+    function addResource(
+        address _id,
+        string memory _resourceId,
+        string memory _resourcePayload
+    )
+        external
+        onlyController(_id)
+        returns (address, string memory, string memory)
+    {
+        require(_id != address(0), 'Invalid address provided');
+        resourceData[_id][_resourceId] = _resourcePayload;
+        keysById[_id].push(_resourceId);
+        emit ResourceAdded(_id, _resourceId, _resourcePayload);
+        return (_id, _resourceId, _resourcePayload);
+    }
+
+    /**
      *@dev Reads DID linked resource from Chain
      *@param _id - Address that refers to the DID doc
      *@param _resourceId - Id that refers to a specific resource
      */
 
-    function getResource(address _id, string memory _resourceId) public view returns (string memory) {
+    function getResource(
+        address _id,
+        string memory _resourceId
+    ) public view returns (string memory) {
         return resourceData[_id][_resourceId];
     }
 
@@ -203,7 +230,9 @@ contract PolygonDidRegistry {
      *@param _id - Address that refers to the DID doc
      */
 
-    function getAllResources(address _id) public view returns (string[] memory) {
+    function getAllResources(
+        address _id
+    ) public view returns (string[] memory) {
         string[] memory result = new string[](keysById[_id].length);
 
         for (uint256 i = 0; i < keysById[_id].length; i++) {