From d76e1532df513072d951cb1b7132b9f29d888366 Mon Sep 17 00:00:00 2001
From: Dmitry <dymitr.paremski@thesis.co>
Date: Wed, 24 Jan 2024 11:35:38 +0100
Subject: [PATCH] Checking for treasury zero address and adding tests

Added a check with custom error in case a new treasury address is about
to be set to zero.
Added tests to verify treasury address during the contract deployment
and during the update by the governance.
---
 core/contracts/Acre.sol      |  6 +++++
 core/test/Acre.test.ts       | 46 ++++++++++++++++++++++++++++++++++++
 core/test/Deployment.test.ts | 18 +++++++++++---
 3 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/core/contracts/Acre.sol b/core/contracts/Acre.sol
index 956a5304f..90a4aad17 100644
--- a/core/contracts/Acre.sol
+++ b/core/contracts/Acre.sol
@@ -75,6 +75,9 @@ contract Acre is ERC4626Fees, Ownable {
         IERC20 _tbtc,
         address _treasury
     ) ERC4626(_tbtc) ERC20("Acre Staked Bitcoin", "stBTC") Ownable(msg.sender) {
+        if (address(_treasury) == address(0)) {
+            revert ZeroAddress();
+        }
         treasury = _treasury;
         // TODO: Revisit the exact values closer to the launch.
         minimumDepositAmount = 0.001 * 1e18; // 0.001 tBTC
@@ -86,6 +89,9 @@ contract Acre is ERC4626Fees, Ownable {
     /// @param newTreasury New treasury wallet address.
     function updateTreasury(address newTreasury) external onlyOwner {
         // TODO: Introduce a parameters update process.
+        if (address(newTreasury) == address(0)) {
+            revert ZeroAddress();
+        }
         treasury = newTreasury;
 
         emit TreasuryUpdated(newTreasury);
diff --git a/core/test/Acre.test.ts b/core/test/Acre.test.ts
index 5e8381d1f..0013271e1 100644
--- a/core/test/Acre.test.ts
+++ b/core/test/Acre.test.ts
@@ -1141,6 +1141,52 @@ describe("Acre", () => {
     })
   })
 
+  describe("updateTreasury", () => {
+    let snapshot: SnapshotRestorer
+
+    before(async () => {
+      snapshot = await takeSnapshot()
+    })
+
+    after(async () => {
+      await snapshot.restore()
+    })
+
+    context("when caller is not governance", () => {
+      it("should revert", async () => {
+        await expect(acre.connect(thirdParty).updateTreasury(ZeroAddress))
+          .to.be.revertedWithCustomError(acre, "OwnableUnauthorizedAccount")
+          .withArgs(thirdParty.address)
+      })
+    })
+
+    context("when caller is governance", () => {
+      context("when a new treasury is zero address", () => {
+        it("should revert", async () => {
+          await expect(
+            acre.connect(governance).updateTreasury(ZeroAddress),
+          ).to.be.revertedWithCustomError(acre, "ZeroAddress")
+        })
+      })
+
+      context("when a new treasury is non-zero address", () => {
+        let newTreasury: string
+
+        before(async () => {
+          // Treasury is set by the deployment scripts. See deployment tests
+          // where initial parameters are checked.
+          newTreasury = await ethers.Wallet.createRandom().getAddress()
+
+          await acre.connect(governance).updateTreasury(newTreasury)
+        })
+
+        it("should update the treasury", async () => {
+          expect(await acre.treasury()).to.be.equal(newTreasury)
+        })
+      })
+    })
+  })
+
   describe("maxMint", () => {
     beforeAfterEachSnapshotWrapper()
 
diff --git a/core/test/Deployment.test.ts b/core/test/Deployment.test.ts
index 362e2f43c..5fc01978b 100644
--- a/core/test/Deployment.test.ts
+++ b/core/test/Deployment.test.ts
@@ -10,9 +10,9 @@ import type { Acre, Dispatcher, TestERC20 } from "../typechain"
 
 async function fixture() {
   const { tbtc, acre, dispatcher } = await deployment()
-  const { governance, maintainer } = await getNamedSigner()
+  const { governance, maintainer, treasury } = await getNamedSigner()
 
-  return { acre, dispatcher, tbtc, governance, maintainer }
+  return { acre, dispatcher, tbtc, governance, maintainer, treasury }
 }
 
 describe("Deployment", () => {
@@ -20,12 +20,24 @@ describe("Deployment", () => {
   let dispatcher: Dispatcher
   let tbtc: TestERC20
   let maintainer: HardhatEthersSigner
+  let treasury: HardhatEthersSigner
 
   before(async () => {
-    ;({ acre, dispatcher, tbtc, maintainer } = await loadFixture(fixture))
+    ;({ acre, dispatcher, tbtc, maintainer, treasury } =
+      await loadFixture(fixture))
   })
 
   describe("Acre", () => {
+    describe("constructor", () => {
+      context("when treasury has been set", () => {
+        it("should be set to a treasury address", async () => {
+          const actualTreasury = await acre.treasury()
+
+          expect(actualTreasury).to.be.equal(await treasury.getAddress())
+        })
+      })
+    })
+
     describe("updateDispatcher", () => {
       context("when a dispatcher has been set", () => {
         it("should be set to a dispatcher address by the deployment script", async () => {