From 86b47d37dc86a11b711de7acb46081e2f0f50624 Mon Sep 17 00:00:00 2001
From: veeso <christian.visintin@veeso.dev>
Date: Tue, 23 Jan 2024 09:58:38 +0100
Subject: [PATCH] feat: allow swap fee set only if caller is fly canister or
 owner

---
 ethereum/fly/contracts/Fly.sol | 11 +++++++++-
 ethereum/fly/test/Fly.ts       | 37 +++++++++++++++++++++++++++++++++-
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/ethereum/fly/contracts/Fly.sol b/ethereum/fly/contracts/Fly.sol
index cc2bb9b..a274cda 100644
--- a/ethereum/fly/contracts/Fly.sol
+++ b/ethereum/fly/contracts/Fly.sol
@@ -39,6 +39,15 @@ contract Fly is ERC20, Ownable {
         _;
     }
 
+    modifier isOwnerOrFlyCanister() {
+        require(
+            (msg.sender == fly_canister_address &&
+                fly_canister_address != address(0)) || msg.sender == owner(),
+            "Fly: caller is not the fly canister nor the owner"
+        );
+        _;
+    }
+
     modifier onlyTestnet() {
         require(
             block.chainid == GOERLI_CHAIN_ID ||
@@ -94,7 +103,7 @@ contract Fly is ERC20, Ownable {
      * @dev Sets the swap fee.
      * @param _swapFee The new swap fee.
      */
-    function setSwapFee(uint256 _swapFee) public onlyOwner {
+    function setSwapFee(uint256 _swapFee) public isOwnerOrFlyCanister {
         swapFee = _swapFee;
     }
 
diff --git a/ethereum/fly/test/Fly.ts b/ethereum/fly/test/Fly.ts
index 2419649..b4c0d48 100644
--- a/ethereum/fly/test/Fly.ts
+++ b/ethereum/fly/test/Fly.ts
@@ -162,12 +162,47 @@ describe("Fly", () => {
     expect(await token.swappedSupply()).to.equal(1_000);
   });
 
-  it("Should update swap fee", async () => {
+  it("Should update swap fee if owner", async () => {
     const { token } = deploy;
     await token.setSwapFee(200);
     expect(await token.swapFee()).to.equal(200);
   });
 
+  it("Should update swap fee if fly canister", async () => {
+    const { token, owner, flyCanister } = deploy;
+    // set fly canister address to owner
+    await token.setFlyCanisterAddress(owner.address);
+    // transfer ownership to fly canister
+    await token.transferOwnership(flyCanister.address);
+    // update fee
+    await token.setSwapFee(200);
+    expect(await token.swapFee()).to.equal(200);
+  });
+
+  it("Should fail to set swap fee if not owner or fly canister", async () => {
+    const { token, flyCanister } = deploy;
+    // set fly canister address to owner
+    await token.setFlyCanisterAddress(flyCanister.address);
+    // transfer ownership to fly canister
+    await token.transferOwnership(flyCanister.address);
+    // update fee
+    await expect(token.setSwapFee(200)).to.be.revertedWith(
+      "Fly: caller is not the fly canister nor the owner"
+    );
+    expect(await token.swapFee()).to.equal(INITIAL_FEE);
+  });
+
+  it("Should fail to set swap fee if not owner and fly canister is unset", async () => {
+    const { token, flyCanister } = deploy;
+    // transfer ownership to fly canister
+    await token.transferOwnership(flyCanister.address);
+    // update fee
+    await expect(token.setSwapFee(200)).to.be.revertedWith(
+      "Fly: caller is not the fly canister nor the owner"
+    );
+    expect(await token.swapFee()).to.equal(INITIAL_FEE);
+  });
+
   it("should renounce ownership", async () => {
     const { token } = deploy;
     await token.renounceOwnership();