fix: write admin_address to 0 and update is_admin checks

lfglabs-dev · Jun 5, 2024 · a72c078 · a72c078
1 parent 309e39d
commit a72c078
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 12 deletions.
diff --git a/src/naming/main.cairo b/src/naming/main.cairo
@@ -156,6 +156,7 @@ mod Naming {
         self._pricing_contract.write(pricing);
         self._referral_contract.write(referral);
         self._admin_address.write(admin);
+        self.ownable.initializer(admin);
     }
 
     component!(path: storage_read_component, storage: storage_read, event: StorageReadEvent);
@@ -711,12 +712,13 @@ mod Naming {
         fn update_admin(ref self: ContractState, new_admin: ContractAddress) {
             assert(get_caller_address() == self._admin_address.read(), 'you are not admin');
             self.ownable.initializer(new_admin);
+            self._admin_address.write(Zeroable::zero());
         }
 
         fn set_expiry(
             ref self: ContractState, root_domain: felt252, expiry: u64
         ) {
-            assert(get_caller_address() == self._admin_address.read(), 'you are not admin');
+            self.ownable.assert_only_owner();
             let hashed_domain = self.hash_domain(array![root_domain].span());
             let domain_data = self._domain_data.read(hashed_domain);
             let data = DomainData {
@@ -735,7 +737,7 @@ mod Naming {
         }
 
         fn claim_balance(ref self: ContractState, erc20: ContractAddress) {
-            assert(get_caller_address() == self._admin_address.read(), 'you are not admin');
+            self.ownable.assert_only_owner();
             let balance = IERC20CamelDispatcher { contract_address: erc20 }
                 .balanceOf(get_contract_address());
             let has_claimed = IERC20CamelDispatcher { contract_address: erc20 }
@@ -744,45 +746,45 @@ mod Naming {
         }
 
         fn set_discount(ref self: ContractState, discount_id: felt252, discount: Discount) {
-            assert(get_caller_address() == self._admin_address.read(), 'you are not admin');
+            self.ownable.assert_only_owner();
             self.discounts.write(discount_id, discount);
         }
 
         fn set_pricing_contract(ref self: ContractState, pricing_contract: ContractAddress) {
-            assert(get_caller_address() == self._admin_address.read(), 'you are not admin');
+            self.ownable.assert_only_owner();
             self._pricing_contract.write(pricing_contract);
         }
 
         fn set_referral_contract(ref self: ContractState, referral_contract: ContractAddress) {
-            assert(get_caller_address() == self._admin_address.read(), 'you are not admin');
+            self.ownable.assert_only_owner();
             self._referral_contract.write(referral_contract);
         }
 
         fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
-            assert(get_caller_address() == self._admin_address.read(), 'you are not admin');
+            self.ownable.assert_only_owner();
             // todo: use components
             assert(!new_class_hash.is_zero(), 'Class hash cannot be zero');
             starknet::replace_class_syscall(new_class_hash).unwrap();
         }
 
         fn set_server_pub_key(ref self: ContractState, new_key: felt252) {
-            assert(get_caller_address() == self._admin_address.read(), 'you are not admin');
+            self.ownable.assert_only_owner();
             self._server_pub_key.write(new_key);
         }
 
         fn whitelist_renewal_contract(ref self: ContractState, contract: ContractAddress) {
-            assert(get_caller_address() == self._admin_address.read(), 'you are not admin');
+            self.ownable.assert_only_owner();
             self._whitelisted_renewal_contracts.write(contract, true);
         }
 
         fn blacklist_renewal_contract(ref self: ContractState, contract: ContractAddress) {
-            assert(get_caller_address() == self._admin_address.read(), 'you are not admin');
+            self.ownable.assert_only_owner();
             self._whitelisted_renewal_contracts.write(contract, false);
         }
 
 
         fn toggle_ar_discount_renew(ref self: ContractState) {
-            assert(get_caller_address() == self._admin_address.read(), 'you are not admin');
+            self.ownable.assert_only_owner();
             self._ar_discount_renew_enabled.write(!self._ar_discount_renew_enabled.read());
         }
     }

diff --git a/src/tests/naming/test_abuses.cairo b/src/tests/naming/test_abuses.cairo
@@ -230,7 +230,7 @@ fn test_non_admin_cannot_set_admin() {
 
 #[test]
 #[available_gas(2000000000)]
-#[should_panic(expected: ('you are not admin', 'ENTRYPOINT_FAILED'))]
+#[should_panic(expected: ('Caller is not the owner', 'ENTRYPOINT_FAILED'))]
 fn test_non_admin_cannot_claim_balance() {
     // setup
     let (eth, _, _, naming) = deploy();

diff --git a/src/tests/naming/test_admin_update.cairo b/src/tests/naming/test_admin_update.cairo
@@ -22,7 +22,6 @@ fn test_update_admin() {
     let new_admin = contract_address_const::<0x456>();
 
     let ownable2Step = IOwnableTwoStepDispatcher { contract_address: naming.contract_address };
-    assert(ownable2Step.owner() == contract_address_const::<0>(), 'admin should be 0');
 
     // we call the update_admin function with the new admin
     set_contract_address(admin);