Merge pull request #60 from primevprotocol/use-ownable

Use openzeppelin's Ownable.sol for whitelist
primev · Dec 19, 2023 · 33473c1 · 33473c1
2 parents af88356 + c33a145
commit 33473c1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 14 deletions.
diff --git a/contracts/Whitelist.sol b/contracts/Whitelist.sol
@@ -1,32 +1,28 @@
 // SPDX-License-Identifier: BSL 1.1
 pragma solidity ^0.8.15;
 
+import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+
 // Contract that allows an admin to add/remove addresses from the whitelist,
 // and allows whitelisted addresses to mint/burn native tokens.
-contract Whitelist {
-    address public admin;
+contract Whitelist is Ownable {
+
     mapping(address => bool) public whitelistedAddresses;
 
     // Mint/burn precompile addresses.
     // See: https://github.com/primevprotocol/go-ethereum/blob/03ae168c6ac15dda8c5a3f123e2b9f3350aad613/core/vm/contracts.go
     address constant MINT = address(0x89);
     address constant BURN = address(0x90);
 
-    constructor(address _admin) {
-        require(_admin != address(0), "Admin address cannot be zero");
-        admin = _admin;
-    }
-
-    modifier onlyAdmin() {
-        require(msg.sender == admin, "Only admin can call this function");
-        _;
+    constructor(address _owner) Ownable() {
+        _transferOwnership(_owner);
     }
 
-    function addToWhitelist(address _address) external onlyAdmin {
+    function addToWhitelist(address _address) external onlyOwner {
         whitelistedAddresses[_address] = true;
     }
 
-    function removeFromWhitelist(address _address) external onlyAdmin {
+    function removeFromWhitelist(address _address) external onlyOwner {
         whitelistedAddresses[_address] = false;
     }
 

diff --git a/test/WhitelistTest.sol b/test/WhitelistTest.sol
@@ -44,7 +44,7 @@ contract WhitelistTest is Test {
 
     function test_RevertNormalUserAddToWhitelist() public {
         vm.prank(normalUser);
-        vm.expectRevert("Only admin can call this function");
+        vm.expectRevert("Ownable: caller is not the owner");
         whitelist.addToWhitelist(addressInstance);
     }
 
@@ -53,7 +53,7 @@ contract WhitelistTest is Test {
         whitelist.addToWhitelist(addressInstance);
         assertTrue(whitelist.isWhitelisted(addressInstance));
         vm.prank(normalUser);
-        vm.expectRevert("Only admin can call this function");
+        vm.expectRevert("Ownable: caller is not the owner");
         whitelist.removeFromWhitelist(addressInstance);
     }
 }