From eeffcfe7d70e1596ea03e7f2008a38e8c61679c2 Mon Sep 17 00:00:00 2001
From: r4bbit <445106+0x-r4bbit@users.noreply.github.com>
Date: Tue, 10 Oct 2023 15:32:46 +0200
Subject: [PATCH] refactor(StakeManager): use custom error in `onlyVault`
 modifier (#28)

Also introduce tests that ensure the error is actually emitted.
---
 .gas-snapshot              | 12 ++++++++-
 contracts/StakeManager.sol |  6 ++++-
 test/StakeManager.t.sol    | 55 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/.gas-snapshot b/.gas-snapshot
index 22a3aff..383ec55 100644
--- a/.gas-snapshot
+++ b/.gas-snapshot
@@ -1 +1,11 @@
-StakeManagerTest:testDeployment() (gas: 26172)
\ No newline at end of file
+LeaveTest:testDeployment() (gas: 26172)
+LeaveTest:test_RevertWhen_SenderIsNotVault() (gas: 10562)
+LockTest:testDeployment() (gas: 26172)
+LockTest:test_RevertWhen_SenderIsNotVault() (gas: 10573)
+MigrateTest:testDeployment() (gas: 26172)
+MigrateTest:test_RevertWhen_SenderIsNotVault() (gas: 10629)
+StakeManagerTest:testDeployment() (gas: 26172)
+StakeTest:testDeployment() (gas: 26172)
+StakeTest:test_RevertWhen_SenderIsNotVault() (gas: 10638)
+UnstakeTest:testDeployment() (gas: 26172)
+UnstakeTest:test_RevertWhen_SenderIsNotVault() (gas: 10575)
\ No newline at end of file
diff --git a/contracts/StakeManager.sol b/contracts/StakeManager.sol
index e353c3b..1609857 100644
--- a/contracts/StakeManager.sol
+++ b/contracts/StakeManager.sol
@@ -7,6 +7,8 @@ import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
 import { StakeVault } from "./StakeVault.sol";
 
 contract StakeManager is Ownable {
+    error StakeManager__SenderIsNotVault();
+
     struct Account {
         uint256 lockUntil;
         uint256 balance;
@@ -40,7 +42,9 @@ contract StakeManager is Ownable {
     ERC20 public immutable stakedToken;
 
     modifier onlyVault() {
-        require(isVault[msg.sender.codehash], "Not a vault");
+        if (!isVault[msg.sender.codehash]) {
+            revert StakeManager__SenderIsNotVault();
+        }
         _;
     }
 
diff --git a/test/StakeManager.t.sol b/test/StakeManager.t.sol
index d2a866f..14e5ca3 100644
--- a/test/StakeManager.t.sol
+++ b/test/StakeManager.t.sol
@@ -30,3 +30,58 @@ contract StakeManagerTest is Test {
         assertEq(stakeManager.totalSupply(), 0);
     }
 }
+
+contract StakeTest is StakeManagerTest {
+    function setUp() public override {
+        StakeManagerTest.setUp();
+    }
+
+    function test_RevertWhen_SenderIsNotVault() public {
+        vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector);
+        stakeManager.stake(100, 1);
+    }
+}
+
+contract UnstakeTest is StakeManagerTest {
+    function setUp() public override {
+        StakeManagerTest.setUp();
+    }
+
+    function test_RevertWhen_SenderIsNotVault() public {
+        vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector);
+        stakeManager.unstake(100);
+    }
+}
+
+contract LockTest is StakeManagerTest {
+    function setUp() public override {
+        StakeManagerTest.setUp();
+    }
+
+    function test_RevertWhen_SenderIsNotVault() public {
+        vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector);
+        stakeManager.lock(100);
+    }
+}
+
+contract LeaveTest is StakeManagerTest {
+    function setUp() public override {
+        StakeManagerTest.setUp();
+    }
+
+    function test_RevertWhen_SenderIsNotVault() public {
+        vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector);
+        stakeManager.leave();
+    }
+}
+
+contract MigrateTest is StakeManagerTest {
+    function setUp() public override {
+        StakeManagerTest.setUp();
+    }
+
+    function test_RevertWhen_SenderIsNotVault() public {
+        vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector);
+        stakeManager.migrate();
+    }
+}