Merge pull request #7 from sqrlfirst/feat/add-upgradeable-component

adding upgradeable feature to contracts
OilerNetwork · Jun 6, 2024 · a3a0ef9 · a3a0ef9
2 parents 3cdfcfe + af75d91
commit a3a0ef9
Show file tree

Hide file tree

Showing 5 changed files with 98 additions and 8 deletions.
diff --git a/starknet/src/L1_headers_store/contract.cairo b/starknet/src/L1_headers_store/contract.cairo
@@ -24,22 +24,30 @@ pub mod L1HeaderStore {
     // *************************************************************************
     // Core lib imports 
     use openzeppelin::access::ownable::ownable::OwnableComponent::InternalTrait;
-    use starknet::{ContractAddress, EthAddress, get_caller_address};
+    use openzeppelin::upgrades::UpgradeableComponent;
+    use openzeppelin::upgrades::interface::IUpgradeable;
+    use starknet::{ContractAddress, EthAddress, get_caller_address, ClassHash};
 
     component!(path: OwnableComponent, storage: ownable, event: OwnableEvent);
+    component!(path: UpgradeableComponent, storage: upgradeable, event: upgradeableEvent);
 
     // Ownable Mixin
     #[abi(embed_v0)]
     impl OwnableMixinImpl = OwnableComponent::OwnableMixinImpl<ContractState>;
     impl OwnableInternalImpl = OwnableComponent::InternalImpl<ContractState>;
 
+    // Upgradeable 
+    impl UpgradeableInternalImpl = UpgradeableComponent::InternalImpl<ContractState>;
+
     // *************************************************************************
     //                              STORAGE
     // *************************************************************************
     #[storage]
     struct Storage {
         #[substorage(v0)]
         ownable: OwnableComponent::Storage,
+        #[substorage(v0)]
+        upgradeable: UpgradeableComponent::Storage,
         initialized: bool,
         l1_messages_origin: ContractAddress,
         latest_l1_block: u64,
@@ -62,7 +70,9 @@ pub mod L1HeaderStore {
     #[derive(Drop, starknet::Event)]
     enum Event {
         #[flat]
-        OwnableEvent: OwnableComponent::Event
+        OwnableEvent: OwnableComponent::Event,
+        #[flat]
+        upgradeableEvent: UpgradeableComponent::Event
     }
 
     // *************************************************************************
@@ -180,6 +190,16 @@ pub mod L1HeaderStore {
         }
     }
 
+    #[abi(embed_v0)]
+    impl UpgradeableImpl of IUpgradeable<ContractState> {
+        /// Upgrades the contract class hash to `new_class_hash`.
+        /// This may only be called by the contract owner.
+        fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
+            self.ownable.assert_only_owner();
+            self.upgradeable._upgrade(new_class_hash);
+        }
+    }
+
     // *************************************************************************
     //                          INTERNAL FUNCTIONS
     // *************************************************************************

diff --git a/starknet/src/L1_messages_proxy/contract.cairo b/starknet/src/L1_messages_proxy/contract.cairo
@@ -9,19 +9,25 @@ pub mod L1MessagesProxy {
     use fossil::L1_messages_proxy::interface::IL1MessagesProxy;
     use fossil::library::words64_utils::words64_to_u256;
     use openzeppelin::access::ownable::OwnableComponent;
+    use openzeppelin::upgrades::UpgradeableComponent;
+    use openzeppelin::upgrades::interface::IUpgradeable;
     // *************************************************************************
     //                               IMPORTS
     // *************************************************************************
     // Core lib imports   
-    use starknet::{ContractAddress, EthAddress};
+    use starknet::{ContractAddress, EthAddress, ClassHash};
 
     component!(path: OwnableComponent, storage: ownable, event: OwnableEvent);
+    component!(path: UpgradeableComponent, storage: upgradeable, event: upgradeableEvent);
 
     // Ownable Mixin
     #[abi(embed_v0)]
     impl OwnableMixinImpl = OwnableComponent::OwnableMixinImpl<ContractState>;
     impl OwnableInternalImpl = OwnableComponent::InternalImpl<ContractState>;
 
+    // Upgradeable 
+    impl UpgradeableInternalImpl = UpgradeableComponent::InternalImpl<ContractState>;
+
     /// Struct to receive message (parentHash_, blockNumber_) from L1.
     #[derive(Drop, Serde)]
     struct L1Payload {
@@ -39,6 +45,8 @@ pub mod L1MessagesProxy {
     struct Storage {
         #[substorage(v0)]
         ownable: OwnableComponent::Storage,
+        #[substorage(v0)]
+        upgradeable: UpgradeableComponent::Storage,
         initialized: bool,
         l1_messages_sender: EthAddress,
         l1_headers_store: IL1HeadersStoreDispatcher,
@@ -51,7 +59,9 @@ pub mod L1MessagesProxy {
     #[derive(Drop, starknet::Event)]
     enum Event {
         #[flat]
-        OwnableEvent: OwnableComponent::Event
+        OwnableEvent: OwnableComponent::Event,
+        #[flat]
+        upgradeableEvent: UpgradeableComponent::Event
     }
 
     // *************************************************************************
@@ -149,4 +159,14 @@ pub mod L1MessagesProxy {
             self.l1_headers_store.read().contract_address
         }
     }
+
+    #[abi(embed_v0)]
+    impl UpgradeableImpl of IUpgradeable<ContractState> {
+        /// Upgrades the contract class hash to `new_class_hash`.
+        /// This may only be called by the contract owner.
+        fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
+            self.ownable.assert_only_owner();
+            self.upgradeable._upgrade(new_class_hash);
+        }
+    }
 }
diff --git a/starknet/src/fact_registry/contract.cairo b/starknet/src/fact_registry/contract.cairo
@@ -17,13 +17,32 @@ pub mod FactRegistry {
         rlp_utils::{to_rlp_array, extract_data, extract_element}, keccak_utils::keccak_words64
     };
     use fossil::types::{OptionsSet, Words64Sequence, RLPItem};
+    use openzeppelin::access::ownable::OwnableComponent;
+    use openzeppelin::upgrades::UpgradeableComponent;
+    use openzeppelin::upgrades::interface::IUpgradeable;
     // Core lib imports 
-    use starknet::{ContractAddress, EthAddress, contract_address_const};
+    use starknet::{ContractAddress, EthAddress, contract_address_const, ClassHash};
+
+    component!(path: OwnableComponent, storage: ownable, event: OwnableEvent);
+    component!(path: UpgradeableComponent, storage: upgradeable, event: upgradeableEvent);
+
+    // Ownable Mixin
+    #[abi(embed_v0)]
+    impl OwnableMixinImpl = OwnableComponent::OwnableMixinImpl<ContractState>;
+    impl OwnableInternalImpl = OwnableComponent::InternalImpl<ContractState>;
+
+    // Upgradeable 
+    impl UpgradeableInternalImpl = UpgradeableComponent::InternalImpl<ContractState>;
+
     // *************************************************************************
     //                              STORAGE
     // *************************************************************************
     #[storage]
     struct Storage {
+        #[substorage(v0)]
+        ownable: OwnableComponent::Storage,
+        #[substorage(v0)]
+        upgradeable: UpgradeableComponent::Storage,
         initialized: bool,
         l1_headers_store: IL1HeadersStoreDispatcher,
         verified_account_storage_hash: LegacyMap::<(EthAddress, u64), u256>,
@@ -33,6 +52,18 @@ pub mod FactRegistry {
         verified_storage: LegacyMap::<(EthAddress, u64, u256), (bool, u256)>,
     }
 
+    // *************************************************************************
+    //                             EVENTS
+    // *************************************************************************
+    #[event]
+    #[derive(Drop, starknet::Event)]
+    enum Event {
+        #[flat]
+        OwnableEvent: OwnableComponent::Event,
+        #[flat]
+        upgradeableEvent: UpgradeableComponent::Event
+    }
+
     // *************************************************************************
     //                          EXTERNAL FUNCTIONS
     // *************************************************************************
@@ -43,12 +74,17 @@ pub mod FactRegistry {
         /// 
         /// # Arguments
         /// * `l1_headers_store_addr` - The address of L1 Header Store contract.
-        fn initialize(ref self: ContractState, l1_headers_store_addr: starknet::ContractAddress) {
+        fn initialize(
+            ref self: ContractState,
+            l1_headers_store_addr: starknet::ContractAddress,
+            admin: starknet::ContractAddress
+        ) {
             assert!(self.initialized.read() == false, "FactRegistry: already initialized");
             self.initialized.write(true);
             self
                 .l1_headers_store
                 .write(IL1HeadersStoreDispatcher { contract_address: l1_headers_store_addr });
+            self.ownable.initializer(admin);
         }
 
         /// Verifies the account information for a given Ethereum address  at a given block using a provided state root proof and stores the verified value.
@@ -274,6 +310,16 @@ pub mod FactRegistry {
         }
     }
 
+    #[abi(embed_v0)]
+    impl UpgradeableImpl of IUpgradeable<ContractState> {
+        /// Upgrades the contract class hash to `new_class_hash`.
+        /// This may only be called by the contract owner.
+        fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
+            self.ownable.assert_only_owner();
+            self.upgradeable._upgrade(new_class_hash);
+        }
+    }
+
     // *************************************************************************
     //                          INTERNAL FUNCTIONS
     // *************************************************************************

diff --git a/starknet/src/fact_registry/interface.cairo b/starknet/src/fact_registry/interface.cairo
@@ -2,7 +2,11 @@ use fossil::types::{OptionsSet, Words64Sequence};
 
 #[starknet::interface]
 pub trait IFactRegistry<TState> {
-    fn initialize(ref self: TState, l1_headers_store_addr: starknet::ContractAddress);
+    fn initialize(
+        ref self: TState,
+        l1_headers_store_addr: starknet::ContractAddress,
+        admin: starknet::ContractAddress
+    );
     fn prove_account(
         ref self: TState,
         option: OptionsSet,

diff --git a/starknet/tests/test_utils.cairo b/starknet/tests/test_utils.cairo
@@ -38,7 +38,7 @@ pub fn setup() -> Dispatchers {
     let (contract_address, _) = contract.deploy(@array![]).unwrap();
     let proxy = IL1MessagesProxyDispatcher { contract_address };
 
-    registry.initialize(store.contract_address);
+    registry.initialize(store.contract_address, OWNER());
     store.initialize(proxy.contract_address, OWNER());
     proxy.initialize(L1_ORIGIN(), store.contract_address, OWNER());