update drip to be onlyOwner and allow caller to provide rate limiting lookup key

Signed-off-by: Aaron Sutula <[email protected]>

Author: asutula

src/Faucet.sol

@@ -23,8 +23,8 @@ event Funding(address indexed from, uint256 value);
 /// @title Faucet Contract
 /// @dev A simple faucet contract for distributing tokens
 contract Faucet is Ownable {
-    /// @dev Mapping to store the next allowed request time for each address
-    mapping(address => uint256) internal _nextRequestAt;
+    /// @dev Mapping to store the next allowed request time for each key
+    mapping(string => uint256) internal _nextRequestAt;
 
     /// @dev Amount of tokens to drip per request
     uint256 internal _dripAmount = 100;
@@ -58,16 +58,18 @@ contract Faucet is Ownable {
     /// @dev Distributes tokens to the specified recipient
     /// @dev Reverts if the recipient has requested tokens too recently
     /// @param recipient The address to receive the tokens
-    function drip(address payable recipient) external {
-        if (_nextRequestAt[recipient] > block.timestamp) {
+    /// @param key The key to identify the recipient used in the _nextRequestAt mapping
+    function drip(address payable recipient, string calldata key) external onlyOwner {
+        uint256 nextRequestAt = _nextRequestAt[key];
+        if (nextRequestAt > block.timestamp) {
             revert TryLater();
         }
-        if (address(this).balance < _dripAmount) {
+        uint256 amount = _dripAmount;
+        if (address(this).balance < amount) {
             revert FaucetEmpty();
         }
-
-        _nextRequestAt[recipient] = block.timestamp + (5 minutes);
-        recipient.transfer(_dripAmount);
+        _nextRequestAt[key] = block.timestamp + (5 minutes);
+        recipient.transfer(amount);
     }
 
     /// @dev Sets the amount of tokens to distribute per request

test/Faucet.t.sol

@@ -14,7 +14,9 @@ contract FaucetTest is Test {
     Faucet internal faucet;
     Vm.Wallet internal wallet;
     Vm.Wallet internal chain;
+    address owner;
     address constant tester = address(0x1804c8AB1F12E6bbf3894d4083f33e07309d1f38);
+    string constant key = "test";
     uint256 constant mintAmount = 1000 * 10 ** 18;
 
     function setUp() public virtual {
@@ -23,31 +25,37 @@ contract FaucetTest is Test {
         wallet = vm.createWallet("user");
         FaucetDeployer faucetDeployer = new FaucetDeployer();
         faucet = faucetDeployer.run("local", mintAmount / 2);
+        owner = faucet.owner();
         assertEq(faucet.supply(), mintAmount / 2);
     }
 
     function test_DripTransfer() public {
         assertEq(wallet.addr.balance, 0);
 
-        faucet.drip(payable(wallet.addr));
+        vm.prank(owner);
+        faucet.drip(payable(wallet.addr), key);
 
         assertEq(faucet.supply(), mintAmount / 2 - faucet.dripAmount());
         assertEq(wallet.addr.balance, faucet.dripAmount());
     }
 
     function test_DripTransferNoDelayFail() public {
-        faucet.drip(payable(wallet.addr));
+        vm.prank(owner);
+        faucet.drip(payable(wallet.addr), key);
 
         vm.expectRevert(TryLater.selector);
-        faucet.drip(payable(wallet.addr));
+        vm.prank(owner);
+        faucet.drip(payable(wallet.addr), key);
     }
 
     function test_DripTransferDelay() public {
-        faucet.drip(payable(wallet.addr));
+        vm.startPrank(owner);
+        faucet.drip(payable(wallet.addr), key);
 
         vm.warp(block.timestamp + (5 minutes));
-
-        faucet.drip(payable(wallet.addr));
+        
+        faucet.drip(payable(wallet.addr), key);
+        vm.stopPrank();
 
         assertEq(wallet.addr.balance, 2 * faucet.dripAmount());
     }