init

CanzaApps · Apr 3, 2024 · 2f863ac · 2f863ac
1 parent b31933e
commit 2f863ac
Showing 1 changed file with 54 additions and 25 deletions.
diff --git a/contracts/Treasury.sol b/contracts/Treasury.sol
@@ -7,7 +7,7 @@ import "@openzeppelin/contracts/security/ReentrancyGuard.sol";
 import "./ITreasury.sol";
 
 contract Treasury is ITreasury, Ownable, ReentrancyGuard {
-    bool paused;
+    bool public paused;
     address[] multisig;
     mapping(address => bool) public multisign;
     address baki_vault;
@@ -31,7 +31,7 @@ contract Treasury is ITreasury, Ownable, ReentrancyGuard {
         address _multisig3,
         string memory _asset,
         address _assetAddress
-    ) {
+    ) Ownable(msg.sender) {
         multisig.push(_multisig1);
         multisig.push(_multisig2);
         multisig.push(_multisig3);
@@ -49,9 +49,11 @@ contract Treasury is ITreasury, Ownable, ReentrancyGuard {
         address _to,
         uint256 _amount,
         string memory _asset
-    ) external isPaused nonReentrant onlyOwner {
+    ) external isPaused nonReentrant {
+        bool auth = authorized();
+        if (auth == false) return;
         bool result = checkMutisig();
-        if (result == false) revert();
+        require(result == true, "This action requires multisig");
         require(
             IERC20(currencies[_asset]).balanceOf(address(this)) >= _amount,
             "Insufficient balance"
@@ -74,7 +76,7 @@ contract Treasury is ITreasury, Ownable, ReentrancyGuard {
         string memory _asset
     ) external isPaused nonReentrant {
         if (baki_vault != msg.sender) {
-            paused = true;
+            _pause();
             revert();
         }
         require(
@@ -93,35 +95,44 @@ contract Treasury is ITreasury, Ownable, ReentrancyGuard {
     }
 
     // Add Baki vault address
-    function addBakiVault(
-        address _vault
-    ) external isPaused onlyOwner nonReentrant {
+    function addBakiVault(address _vault) external isPaused nonReentrant {
+        bool auth = authorized();
+        if (auth == false) return;
         require(_vault != address(0), "address cannot be a zero address");
         baki_vault = _vault;
     }
 
-    function addAsset(
-        address _address,
-        string memory _asset
-    ) external onlyOwner isPaused nonReentrant {
+    function addAsset(address _address, string memory _asset)
+        external
+        isPaused
+        nonReentrant
+    {
+        bool auth = authorized();
+        if (auth == false) return;
         require(_address != address(0), "address cannot be a zero address");
         currencies[_asset] = _address;
     }
 
-    function viewBalance(
-        string memory _asset
-    ) external view onlyOwner isPaused returns (uint256) {
+    function viewBalance(string memory _asset)
+        external
+        view
+        onlyOwner
+        isPaused
+        returns (uint256)
+    {
         return IERC20(currencies[_asset]).balanceOf(address(this));
     }
 
     // Control functuons
     function pauseVault() external onlyOwner nonReentrant {
-        paused = true;
+        _pause();
     }
 
-    function unPauseVault() external onlyOwner nonReentrant {
+    function unPauseVault() external nonReentrant {
+        bool auth = authorized();
+        if (auth == false) return;
         bool result = checkMutisig();
-        if (result == false) revert();
+        require(result == true, "This action requires multisig");
         paused = false;
         rollMultisig();
     }
@@ -130,21 +141,26 @@ contract Treasury is ITreasury, Ownable, ReentrancyGuard {
         return paused;
     }
 
-    function addMultisig(address _signer) external onlyOwner nonReentrant {
+    function addMultisig(address _signer) external isPaused nonReentrant {
+        bool auth = authorized();
+        if (auth == false) return;
         bool result = checkMutisig();
-        if (result == false) revert();
+        require(result == true, "This action requires multisig");
         require(_signer != address(0), "address cannot be a zero address");
         multisig.push(_signer);
         multisign[_signer] = false;
         rollMultisig();
     }
 
-    function updateMultisig(
-        address _signer,
-        uint256 _index
-    ) external onlyOwner nonReentrant {
+    function updateMultisig(address _signer, uint256 _index)
+        external
+        isPaused
+        nonReentrant
+    {
+        bool auth = authorized();
+        if (auth == false) return;
         bool result = checkMutisig();
-        if (result == false) revert();
+        require(result == true, "This action requires multisig");
         require(_signer != address(0), "address cannot be a zero address");
         multisig[_index] = _signer;
         rollMultisig();
@@ -176,4 +192,17 @@ contract Treasury is ITreasury, Ownable, ReentrancyGuard {
         }
         return true;
     }
+
+    function authorized() internal returns (bool) {
+        // Add your custom logic here
+        if (msg.sender == owner()) {
+            return true;
+        }
+        _pause();
+        return false;
+    }
+
+    function _pause() internal {
+        paused = true;
+    }
 }