Rabbithole balance functions

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "lib/forge-std"]
 	path = lib/forge-std
 	url = https://github.com/foundry-rs/forge-std
+[submodule "lib/solady"]
+	path = lib/solady
+	url = https://github.com/vectorized/solady

foundry.toml

@@ -2,7 +2,7 @@
 auto_detect_solc = true
 fs_permissions = [{access = "read", path = "./addresses"}, {access = "read", path = "./package.json"}]
 fuzz_runs = 500
-libs = ['lib', '_imagine']
+libs = ['lib']
 optimizer = true
 optimizer_runs = 500000
 out = 'out'

remappings.txt

@@ -1 +1,2 @@
-forge-std/=lib/forge-std/src/
+forge-std/=lib/forge-std/src/
+solady/=lib/solady/src/

src/ProtocolRewards.sol

@@ -1,13 +1,13 @@
 // SPDX-License-Identifier: MIT
 pragma solidity 0.8.17;
 
-import {Enjoy} from "../_imagine/Enjoy.sol";
-import {EIP712} from "./lib/EIP712.sol";
+import {EIP712} from "solady/utils/EIP712.sol";
+import {Ownable} from "solady/auth/Ownable.sol";
 import {IProtocolRewards} from "./interfaces/IProtocolRewards.sol";
 
 /// @title ProtocolRewards
 /// @notice Manager of deposits & withdrawals for protocol rewards
-contract ProtocolRewards is Enjoy, IProtocolRewards, EIP712 {
+contract ProtocolRewards is IProtocolRewards, EIP712, Ownable {
     /// @notice The EIP-712 typehash for gasless withdraws
     bytes32 public constant WITHDRAW_TYPEHASH = keccak256("Withdraw(address from,address to,uint256 amount,uint256 nonce,uint256 deadline)");
 
@@ -17,7 +17,21 @@ contract ProtocolRewards is Enjoy, IProtocolRewards, EIP712 {
     /// @notice An account's nonce for gasless withdraws
     mapping(address => uint256) public nonces;
 
-    constructor() payable EIP712("ProtocolRewards", "1") {}
+    /// @notice Total Balance across all accounts
+    uint256 public totalBalance;
+
+    constructor() payable EIP712() {}
+
+    function _domainNameAndVersion()
+        internal
+        pure
+        virtual
+        override
+        returns (string memory name, string memory version)
+    {
+        name = "ProtocolRewards";
+        version = "1";
+    }
 
     /// @notice The total amount of ETH held in the contract
     function totalSupply() external view returns (uint256) {
@@ -33,11 +47,32 @@ contract ProtocolRewards is Enjoy, IProtocolRewards, EIP712 {
             revert ADDRESS_ZERO();
         }
 
-        balanceOf[to] += msg.value;
+        _increaseBalance(to, msg.value);
 
         emit Deposit(msg.sender, to, reason, msg.value, comment);
     }
 
+    /// @notice Allow admin to increase balance of an address only up the amount of excess ETH in the contract
+    /// @dev TODO we will probably need a batch version of this function
+    /// @param to The address to increase the balance of
+    /// @param amount The amount to increase the balance by
+    function increaseBalance(address to, uint256 amount) external onlyOwner {
+        if (amount > address(this).balance - totalBalance) {
+            revert INVALID_AMOUNT();
+        }
+        _increaseBalance(to, amount);
+    }
+
+    function _increaseBalance(address to, uint256 amount) internal {
+        balanceOf[to] += amount;
+        totalBalance += amount;
+    }
+
+    function _decreaseBalance(address to, uint256 amount) internal {
+        balanceOf[to] -= amount;
+        totalBalance -= amount;
+    }
+
     /// @notice Generic function to deposit ETH for multiple recipients, with an optional comment
     /// @param recipients recipients to send the amount to, array aligns with amounts
     /// @param amounts amounts to send to each recipient, array aligns with recipients
@@ -75,7 +110,7 @@ contract ProtocolRewards is Enjoy, IProtocolRewards, EIP712 {
                 revert ADDRESS_ZERO();
             }
 
-            balanceOf[currentRecipient] += currentAmount;
+            _increaseBalance(currentRecipient, currentAmount);
 
             emit Deposit(msg.sender, currentRecipient, reasons[i], currentAmount, comment);
 
@@ -85,66 +120,6 @@ contract ProtocolRewards is Enjoy, IProtocolRewards, EIP712 {
         }
     }
 
-    /// @notice Used by Zora ERC-721 & ERC-1155 contracts to deposit protocol rewards
-    /// @param creator Creator for NFT rewards
-    /// @param creatorReward Creator reward amount
-    /// @param createReferral Creator referral
-    /// @param createReferralReward Creator referral reward
-    /// @param mintReferral Mint referral user
-    /// @param mintReferralReward Mint referral amount
-    /// @param firstMinter First minter reward
-    /// @param firstMinterReward First minter reward amount
-    /// @param zora ZORA recipient
-    /// @param zoraReward ZORA amount
-    function depositRewards(
-        address creator,
-        uint256 creatorReward,
-        address createReferral,
-        uint256 createReferralReward,
-        address mintReferral,
-        uint256 mintReferralReward,
-        address firstMinter,
-        uint256 firstMinterReward,
-        address zora,
-        uint256 zoraReward
-    ) external payable {
-        if (msg.value != (creatorReward + createReferralReward + mintReferralReward + firstMinterReward + zoraReward)) {
-            revert INVALID_DEPOSIT();
-        }
-
-        unchecked {
-            if (creator != address(0)) {
-                balanceOf[creator] += creatorReward;
-            }
-            if (createReferral != address(0)) {
-                balanceOf[createReferral] += createReferralReward;
-            }
-            if (mintReferral != address(0)) {
-                balanceOf[mintReferral] += mintReferralReward;
-            }
-            if (firstMinter != address(0)) {
-                balanceOf[firstMinter] += firstMinterReward;
-            }
-            if (zora != address(0)) {
-                balanceOf[zora] += zoraReward;
-            }
-        }
-
-        emit RewardsDeposit(
-            creator,
-            createReferral,
-            mintReferral,
-            firstMinter,
-            zora,
-            msg.sender,
-            creatorReward,
-            createReferralReward,
-            mintReferralReward,
-            firstMinterReward,
-            zoraReward
-        );
-    }
-
     /// @notice Withdraw protocol rewards
     /// @param to Withdraws from msg.sender to this address
     /// @param amount Amount to withdraw (0 for total balance)
@@ -163,7 +138,7 @@ contract ProtocolRewards is Enjoy, IProtocolRewards, EIP712 {
             amount = balanceOf[owner];
         }
 
-        balanceOf[owner] -= amount;
+        _decreaseBalance(owner, amount);
 
         emit Withdraw(owner, to, amount);
 
@@ -220,7 +195,7 @@ contract ProtocolRewards is Enjoy, IProtocolRewards, EIP712 {
             withdrawHash = keccak256(abi.encode(WITHDRAW_TYPEHASH, from, to, amount, nonces[from]++, deadline));
         }
 
-        bytes32 digest = _hashTypedDataV4(withdrawHash);
+        bytes32 digest = _hashTypedData(withdrawHash);
 
         address recoveredAddress = ecrecover(digest, v, r, s);
 
@@ -240,7 +215,7 @@ contract ProtocolRewards is Enjoy, IProtocolRewards, EIP712 {
             amount = balanceOf[from];
         }
 
-        balanceOf[from] -= amount;
+        _decreaseBalance(from, amount);
 
         emit Withdraw(from, to, amount);