Merge pull request #156 from etherfi-protocol/vaibhav/fix-certora-rou…

…ter-issue Certora-Fix- Add Ability to Change Treasury
etherfi-protocol · Sep 3, 2024 · 4ddd751 · 4ddd751
2 parents 89b8462 + f84e8b1
commit 4ddd751
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 10 deletions.
diff --git a/src/EtherFiRewardsRouter.sol b/src/EtherFiRewardsRouter.sol
@@ -19,14 +19,16 @@ contract EtherFiRewardsRouter is OwnableUpgradeable, UUPSUpgradeable  {
 
     event EthReceived(address indexed from, uint256 value);
     event EthSent(address indexed from, address indexed to, uint256 value);
+    event UpdatedTreasury(address indexed treasury);
     event Erc20Sent(address indexed caller, address indexed token, uint256 amount);
     event Erc721Sent(address indexed caller, address indexed token, uint256 tokenId);
 
     error IncorrectRole();
 
-    constructor(address _liquidityPool, address _roleRegistry) {
+    constructor(address _liquidityPool, address _treasury, address _roleRegistry) {
         _disableInitializers();
         liquidityPool = _liquidityPool;
+        treasury = _treasury;
         roleRegistry = RoleRegistry(_roleRegistry);
     }
 

diff --git a/test/EtherFiRewardsRouter.t.sol b/test/EtherFiRewardsRouter.t.sol
@@ -1,30 +1,36 @@
 import "./TestSetup.sol";
 import "../src/EtherFiRewardsRouter.sol";
 import "../src/LiquidityPool.sol";
+import "forge-std/console2.sol";
 
 contract EtherFiRewardsRouterTest is TestSetup {
-
     address liquidityPoolAddress;
     EtherFiRewardsRouter etherfiRewardsRouterImplementation;
     UUPSProxy etherfiRewardsRouterProxy;
     EtherFiRewardsRouter etherfiRewardsRouterInstance;
 
+    function get_eeth() public {
+        vm.startPrank(address(etherfiRewardsRouterInstance));
+        vm.deal(address(etherfiRewardsRouterInstance), 10 ether);
+        liquidityPoolInstance.deposit{value: 2 ether}();
+        vm.stopPrank();
+    }
 
     function setUp() public {
         setUpTests();
         initializeRealisticFork(MAINNET_FORK);
-        vm.startPrank(owner);
-        etherfiRewardsRouterImplementation = new EtherFiRewardsRouter(address(liquidityPoolInstance), address(roleRegistry));
+        vm.startPrank(superAdmin);
+        etherfiRewardsRouterImplementation = new EtherFiRewardsRouter(address(liquidityPoolInstance), address(treasuryInstance), address(roleRegistry));
         etherfiRewardsRouterProxy = new UUPSProxy(address(etherfiRewardsRouterImplementation), "");
         etherfiRewardsRouterInstance = EtherFiRewardsRouter(payable(address(etherfiRewardsRouterProxy)));
         etherfiRewardsRouterInstance.initialize();
         vm.startPrank(superAdmin);
         roleRegistry.grantRole(etherfiRewardsRouterInstance.ETHERFI_ROUTER_ADMIN(), admin);
+        etherfiRewardsRouterInstance.transferOwnership(owner);
         vm.stopPrank();
-
         liquidityPoolAddress = address(liquidityPoolInstance);
         vm.deal(address(etherfiRewardsRouterInstance), 10 ether);
-        vm.stopPrank(); 
+        vm.stopPrank();
     }
 
     function test_withdrawToLiquidityPool() public {
@@ -38,21 +44,30 @@ contract EtherFiRewardsRouterTest is TestSetup {
 
     function test_elRouterUpgrade() public {
         vm.startPrank(owner);
-        EtherFiRewardsRouter newEtherfiRewardsRouterImplementation = new EtherFiRewardsRouter(address(liquidityPoolInstance), address(roleRegistry));
+        EtherFiRewardsRouter newEtherfiRewardsRouterImplementation = new EtherFiRewardsRouter(address(liquidityPoolInstance), address(treasuryInstance), address(roleRegistry));
         address oldImplementation = etherfiRewardsRouterInstance.getImplementation();
         etherfiRewardsRouterInstance.upgradeTo(address(newEtherfiRewardsRouterImplementation));
         address newImplementation = etherfiRewardsRouterInstance.getImplementation();
         assert(newImplementation != oldImplementation);
         assert(newImplementation == address(newEtherfiRewardsRouterImplementation));
         test_withdrawToLiquidityPool();
         vm.stopPrank();
-    } 
+    }
 
     function test_checkEventEmitted() public {
         vm.startPrank(admin);
         vm.expectEmit(true, true, false, true);
         emit EtherFiRewardsRouter.EthSent(address(etherfiRewardsRouterInstance), liquidityPoolAddress, 10 ether);
         etherfiRewardsRouterInstance.withdrawToLiquidityPool();
         vm.stopPrank();
-    }  
-}
+    }
+
+    function test_recoverERC20() public {
+        get_eeth();
+        vm.startPrank(admin);
+        uint256 balanceBefore = eETHInstance.balanceOf(address(treasuryInstance));
+        etherfiRewardsRouterInstance.recoverERC20(address(eETHInstance), 1 ether);
+        uint256 balanceAfter = eETHInstance.balanceOf(address(treasuryInstance));
+        assertApproxEqAbs(balanceAfter, balanceBefore + 1 ether, 1);
+    }
+}