From 3a466b3b1b741ae7ecee37a5ad7e735f225f9d26 Mon Sep 17 00:00:00 2001
From: Denis Fadeev <denis@fadeev.org>
Date: Fri, 22 Sep 2023 10:11:33 +0400
Subject: [PATCH] custom errors

---
 omnichain/staking/contracts/Staking.sol | 22 ++++++++++------------
 1 file changed, 10 insertions(+), 12 deletions(-)

diff --git a/omnichain/staking/contracts/Staking.sol b/omnichain/staking/contracts/Staking.sol
index 966731d2..2dbfaed6 100644
--- a/omnichain/staking/contracts/Staking.sol
+++ b/omnichain/staking/contracts/Staking.sol
@@ -10,6 +10,11 @@ contract Staking is ERC20, zContract {
     error SenderNotSystemContract();
     error WrongChain(uint256 chainID);
     error UnknownAction(uint8 action);
+    error Overflow();
+    error Underflow();
+    error WrongAmount();
+    error NotAuthorized();
+    error NoRewardsToClaim();
 
     SystemContract public immutable systemContract;
     uint256 constant BITCOIN = 18332;
@@ -58,10 +63,6 @@ contract Staking is ERC20, zContract {
         uint256 amount,
         bytes calldata message
     ) external override onlySystem {
-        if (msg.sender != address(systemContract)) {
-            revert SenderNotSystemContract();
-        }
-
         if (chainID != context.chainID) {
             revert WrongChain(context.chainID);
         }
@@ -87,7 +88,7 @@ contract Staking is ERC20, zContract {
 
     function stakeZRC(address staker, uint256 amount) internal {
         stake[staker] += amount;
-        require(stake[staker] >= amount, "Overflow detected");
+        if (stake[staker] < amount) revert Overflow();
 
         lastStakeTime[staker] = block.timestamp;
         updateRewards(staker);
@@ -108,7 +109,7 @@ contract Staking is ERC20, zContract {
         address zrc20 = systemContract.gasCoinZRC20ByChainId(chainID);
         (, uint256 gasFee) = IZRC20(zrc20).withdrawGasFee();
 
-        require(amount >= gasFee, "Amount should be greater than the gas fee");
+        if (amount < gasFee) revert WrongAmount();
 
         bytes memory recipient = withdraw[staker];
 
@@ -116,7 +117,7 @@ contract Staking is ERC20, zContract {
         IZRC20(zrc20).withdraw(recipient, amount - gasFee);
 
         stake[staker] = 0;
-        require(stake[staker] <= amount, "Underflow detected");
+        if (stake[staker] > amount) revert Underflow();
 
         lastStakeTime[staker] = block.timestamp;
     }
@@ -152,12 +153,9 @@ contract Staking is ERC20, zContract {
     }
 
     function claimRewards(address staker) public {
-        require(
-            beneficiary[staker] == msg.sender,
-            "Not authorized to claim rewards"
-        );
+        if (beneficiary[staker] != msg.sender) revert NotAuthorized();
         uint256 rewardAmount = queryRewards(staker);
-        require(rewardAmount > 0, "No rewards to claim");
+        if (rewardAmount <= 0) revert NoRewardsToClaim();
         updateRewards(staker);
     }
 }