aaiello/zeta-335-add-supply-cap-to-smart-contracts (#13)

* aaiello/zeta-335-add-supply-cap-to-smart-contracts * comments addressed * renames * default value to maxSupply * update test * Add describe, enforce max supply on revert * Improve test message * Move setMaxSupply function Co-authored-by: Lucas Janon <[email protected]>
zeta-chain · Jun 23, 2022 · 50aefdc · 50aefdc
1 parent d651c44
commit 50aefdc
Show file tree

Hide file tree

Showing 3 changed files with 89 additions and 0 deletions.
diff --git a/packages/protocol-contracts/contracts/ZetaConnector.non-eth.sol b/packages/protocol-contracts/contracts/ZetaConnector.non-eth.sol
@@ -17,6 +17,8 @@ interface ZetaToken is IERC20 {
 }
 
 contract ZetaConnectorNonEth is ZetaConnectorBase {
+    uint256 public maxSupply = 2**256 - 1;
+
     constructor(
         address zetaTokenAddress_,
         address tssAddress_,
@@ -27,6 +29,10 @@ contract ZetaConnectorNonEth is ZetaConnectorBase {
         return ZetaToken(zetaToken).balanceOf(address(this));
     }
 
+    function setMaxSupply(uint256 maxSupply_) external onlyTssAddress {
+        maxSupply = maxSupply_;
+    }
+
     function send(ZetaInterfaces.SendInput calldata input) external override whenNotPaused {
         ZetaToken(zetaToken).burnFrom(msg.sender, input.zetaAmount);
 
@@ -49,6 +55,7 @@ contract ZetaConnectorNonEth is ZetaConnectorBase {
         bytes calldata message,
         bytes32 internalSendHash
     ) external override whenNotPaused onlyTssAddress {
+        if (zetaAmount + ZetaToken(zetaToken).totalSupply() > maxSupply) revert ExceedsMaxSupply(maxSupply);
         ZetaToken(zetaToken).mint(destinationAddress, zetaAmount, internalSendHash);
 
         if (message.length > 0) {
@@ -76,6 +83,7 @@ contract ZetaConnectorNonEth is ZetaConnectorBase {
         bytes calldata message,
         bytes32 internalSendHash
     ) external override whenNotPaused onlyTssAddress {
+        if (zetaAmount + ZetaToken(zetaToken).totalSupply() > maxSupply) revert ExceedsMaxSupply(maxSupply);
         ZetaToken(zetaToken).mint(originSenderAddress, zetaAmount, internalSendHash);
 
         if (message.length > 0) {

diff --git a/packages/protocol-contracts/contracts/interfaces/ConnectorErrors.sol b/packages/protocol-contracts/contracts/interfaces/ConnectorErrors.sol
@@ -9,4 +9,6 @@ interface ConnectorErrors {
     error InvalidAddress();
 
     error ZetaTransferError();
+
+    error ExceedsMaxSupply(uint256 maxSupply);
 }
diff --git a/packages/protocol-contracts/test/ZetaConnector.spec.ts b/packages/protocol-contracts/test/ZetaConnector.spec.ts
@@ -671,5 +671,84 @@ describe("ZetaConnector tests", () => {
         expect(e2.length).to.equal(1);
       });
     });
+
+    describe("MaxSupply", () => {
+      describe("setMaxSupply", () => {
+        it("Should revert if the caller is not the TSS address", async () => {
+          await expect(zetaConnectorNonEthContract.connect(randomSigner).setMaxSupply(0)).to.revertedWith(
+            `CallerIsNotTss("${randomSigner.address}")`
+          );
+        });
+
+        it("Should revert if want to mint more than MaxSupply", async () => {
+          await zetaConnectorNonEthContract.connect(tssSigner).setMaxSupply(999);
+          await expect(
+            zetaConnectorNonEthContract
+              .connect(tssSigner)
+              .onReceive(
+                randomSigner.address,
+                1,
+                zetaReceiverMockContract.address,
+                1000,
+                new ethers.utils.AbiCoder().encode(["string"], ["hello"]),
+                "0x0000000000000000000000000000000000000000000000000000000000000000"
+              )
+          ).to.revertedWith(`ExceedsMaxSupply(999)`);
+        });
+      });
+
+      describe("onReceive, onRevert (mint)", () => {
+        it("Should mint if total supply + supply to add < max supply", async () => {
+          const supplyToAdd = 1000;
+          const initialSupply = await zetaTokenNonEthContract.totalSupply();
+
+          await zetaConnectorNonEthContract.connect(tssSigner).setMaxSupply(initialSupply.add(supplyToAdd));
+
+          await expect(
+            zetaConnectorNonEthContract
+              .connect(tssSigner)
+              .onReceive(
+                randomSigner.address,
+                1,
+                zetaReceiverMockContract.address,
+                supplyToAdd,
+                new ethers.utils.AbiCoder().encode(["string"], ["hello"]),
+                "0x0000000000000000000000000000000000000000000000000000000000000000"
+              )
+          ).to.be.not.reverted;
+
+          const finalSupply = await zetaTokenNonEthContract.totalSupply();
+
+          expect(finalSupply).to.eq(initialSupply.add(supplyToAdd));
+
+          await expect(
+            zetaConnectorNonEthContract
+              .connect(tssSigner)
+              .onReceive(
+                randomSigner.address,
+                1,
+                zetaReceiverMockContract.address,
+                1,
+                new ethers.utils.AbiCoder().encode(["string"], ["hello"]),
+                "0x0000000000000000000000000000000000000000000000000000000000000000"
+              )
+          ).to.revertedWith(`ExceedsMaxSupply(${initialSupply.add(supplyToAdd)})`);
+
+          await expect(
+            zetaConnectorNonEthContract
+              .connect(tssSigner)
+              .onRevert(
+                randomSigner.address,
+                1,
+                randomSigner.address,
+                2,
+                1000,
+                new ethers.utils.AbiCoder().encode(["string"], ["hello"]),
+                "0x0000000000000000000000000000000000000000000000000000000000000000"
+              )
+          ).to.revertedWith(`ExceedsMaxSupply(${initialSupply.add(supplyToAdd)})`);
+        });
+      });
+    });
   });
 });