Migrate upgrades

src/openzeppelin/lib.cairo

@@ -4,4 +4,5 @@ mod introspection;
 mod security;
 mod tests;
 mod token;
+mod upgrades;
 mod utils;

src/openzeppelin/tests.cairo

@@ -4,4 +4,5 @@ mod introspection;
 mod mocks;
 mod security;
 mod token;
+mod upgrades;
 mod utils;

src/openzeppelin/tests/mocks.cairo

@@ -17,3 +17,5 @@ mod camel721_mock;
 mod non_implementing_mock;
 mod dual721_receiver_mocks;
 mod src5_mocks;
+mod upgrades_v1;
+mod upgrades_v2;

src/openzeppelin/tests/mocks/upgrades_v1.cairo

@@ -0,0 +1,70 @@
+// This contract is a mock used to test the core functionality of the upgrade functions.
+// The functions are NOT PROTECTED.
+// DO NOT USE IN PRODUCTION.
+
+use array::ArrayTrait;
+use starknet::class_hash::ClassHash;
+
+#[starknet::interface]
+trait IUpgradesV1<TState> {
+    fn upgrade(ref self: TState, impl_hash: ClassHash);
+    fn upgrade_and_call(
+        ref self: TState, impl_hash: ClassHash, selector: felt252, calldata: Span<felt252>
+    );
+    fn set_value(ref self: TState, val: felt252);
+    fn get_value(self: @TState) -> felt252;
+    fn remove_selector(self: @TState);
+}
+
+trait UpgradesV1Trait<TState> {
+    fn set_value(ref self: TState, val: felt252);
+    fn get_value(self: @TState) -> felt252;
+    fn remove_selector(self: @TState);
+}
+
+#[starknet::contract]
+mod UpgradesV1 {
+    use array::ArrayTrait;
+    use starknet::class_hash::ClassHash;
+    use starknet::ContractAddress;
+    use openzeppelin::upgrades::upgradeable::Upgradeable;
+    use openzeppelin::upgrades::upgradeable::UpgradeableTrait;
+
+    #[storage]
+    struct Storage {
+        value: felt252
+    }
+
+    #[external(v0)]
+    impl UpgradeableImpl of UpgradeableTrait<ContractState> {
+        fn upgrade(ref self: ContractState, impl_hash: ClassHash) {
+            let mut unsafe_state = Upgradeable::unsafe_new_contract_state();
+            Upgradeable::UpgradeableImpl::upgrade(ref unsafe_state, impl_hash);
+        }
+
+        fn upgrade_and_call(
+            ref self: ContractState,
+            impl_hash: ClassHash,
+            selector: felt252,
+            calldata: Span<felt252>
+        ) {
+            let mut unsafe_state = Upgradeable::unsafe_new_contract_state();
+            Upgradeable::UpgradeableImpl::upgrade_and_call(
+                ref unsafe_state, impl_hash, selector, calldata
+            );
+        }
+    }
+
+    #[external(v0)]
+    impl UpgradesV1Impl of super::UpgradesV1Trait<ContractState> {
+        fn set_value(ref self: ContractState, val: felt252) {
+            self.value.write(val);
+        }
+
+        fn get_value(self: @ContractState) -> felt252 {
+            self.value.read()
+        }
+
+        fn remove_selector(self: @ContractState) {}
+    }
+}

src/openzeppelin/tests/mocks/upgrades_v2.cairo

@@ -0,0 +1,79 @@
+// This contract is a mock used to test the core functionality of the upgrade functions.
+// The functions are NOT PROTECTED.
+// DO NOT USE IN PRODUCTION.
+
+use array::ArrayTrait;
+use starknet::class_hash::ClassHash;
+
+#[starknet::interface]
+trait IUpgradesV2<TState> {
+    fn upgrade(ref self: TState, impl_hash: ClassHash);
+    fn upgrade_and_call(
+        ref self: TState, impl_hash: ClassHash, selector: felt252, calldata: Span<felt252>
+    );
+    fn set_value(ref self: TState, val: felt252);
+    fn set_value2(ref self: TState, val: felt252);
+    fn get_value(self: @TState) -> felt252;
+    fn get_value2(self: @TState) -> felt252;
+}
+
+trait UpgradesV2Trait<TState> {
+    fn set_value(ref self: TState, val: felt252);
+    fn set_value2(ref self: TState, val: felt252);
+    fn get_value(self: @TState) -> felt252;
+    fn get_value2(self: @TState) -> felt252;
+}
+
+#[starknet::contract]
+mod UpgradesV2 {
+    use array::ArrayTrait;
+    use starknet::class_hash::ClassHash;
+    use starknet::ContractAddress;
+    use openzeppelin::upgrades::upgradeable::Upgradeable;
+    use openzeppelin::upgrades::upgradeable::UpgradeableTrait;
+
+    #[storage]
+    struct Storage {
+        value: felt252,
+        value2: felt252,
+    }
+
+    #[external(v0)]
+    impl UpgradeableImpl of UpgradeableTrait<ContractState> {
+        fn upgrade(ref self: ContractState, impl_hash: ClassHash) {
+            let mut unsafe_state = Upgradeable::unsafe_new_contract_state();
+            Upgradeable::UpgradeableImpl::upgrade(ref unsafe_state, impl_hash);
+        }
+
+        fn upgrade_and_call(
+            ref self: ContractState,
+            impl_hash: ClassHash,
+            selector: felt252,
+            calldata: Span<felt252>
+        ) {
+            let mut unsafe_state = Upgradeable::unsafe_new_contract_state();
+            Upgradeable::UpgradeableImpl::upgrade_and_call(
+                ref unsafe_state, impl_hash, selector, calldata
+            );
+        }
+    }
+
+    #[external(v0)]
+    impl UpgradesV2Impl of super::UpgradesV2Trait<ContractState> {
+        fn set_value(ref self: ContractState, val: felt252) {
+            self.value.write(val);
+        }
+
+        fn set_value2(ref self: ContractState, val: felt252) {
+            self.value2.write(val);
+        }
+
+        fn get_value(self: @ContractState) -> felt252 {
+            self.value.read()
+        }
+
+        fn get_value2(self: @ContractState) -> felt252 {
+            self.value2.read()
+        }
+    }
+}

src/openzeppelin/tests/upgrades.cairo

@@ -0,0 +1 @@
+mod test_upgradeable;

src/openzeppelin/tests/upgrades/test_upgradeable.cairo

@@ -0,0 +1,175 @@
+use array::ArrayTrait;
+use option::OptionTrait;
+use starknet::class_hash::ClassHash;
+use starknet::class_hash::class_hash_const;
+use starknet::class_hash::Felt252TryIntoClassHash;
+use starknet::ContractAddress;
+use starknet::contract_address_const;
+use starknet::testing;
+use traits::TryInto;
+
+use openzeppelin::tests::mocks::upgrades_v1::UpgradesV1;
+use openzeppelin::tests::mocks::upgrades_v1::IUpgradesV1Dispatcher;
+use openzeppelin::tests::mocks::upgrades_v1::IUpgradesV1DispatcherTrait;
+use openzeppelin::tests::mocks::upgrades_v2::UpgradesV2;
+use openzeppelin::tests::mocks::upgrades_v2::IUpgradesV2Dispatcher;
+use openzeppelin::tests::mocks::upgrades_v2::IUpgradesV2DispatcherTrait;
+use openzeppelin::upgrades::upgradeable::Upgradeable;
+use openzeppelin::upgrades::upgradeable::Upgradeable::Upgraded;
+use openzeppelin::tests::utils;
+
+const VALUE: felt252 = 123;
+const VALUE2: felt252 = 789;
+const SET_VALUE_SELECTOR: felt252 =
+    0x3d7905601c217734671143d457f0db37f7f8883112abd34b92c4abfeafde0c3;
+const SET_VALUE2_SELECTOR: felt252 =
+    0x47c8c185eb97d3925831a1c97e43bd9077181d2b200133ede551f1c47056a3;
+const REMOVE_SELECTOR: felt252 = 0x2beeaa48bce210364c7d2f2fbb677e08136cb29e8972a6728249364dde19e6f;
+
+fn V2_CLASS_HASH() -> ClassHash {
+    UpgradesV2::TEST_CLASS_HASH.try_into().unwrap()
+}
+
+fn CLASS_HASH_ZERO() -> ClassHash {
+    class_hash_const::<0>()
+}
+
+fn ZERO() -> ContractAddress {
+    contract_address_const::<0>()
+}
+
+//
+// Setup
+//
+
+fn deploy_v1() -> IUpgradesV1Dispatcher {
+    let calldata = array![];
+    let address = utils::deploy(UpgradesV1::TEST_CLASS_HASH, calldata);
+    IUpgradesV1Dispatcher { contract_address: address }
+}
+
+//
+// upgrade
+//
+
+#[test]
+#[available_gas(2000000)]
+#[should_panic(expected: ('Class hash cannot be zero', 'ENTRYPOINT_FAILED', ))]
+fn test_upgrade_with_class_hash_zero() {
+    let v1 = deploy_v1();
+    v1.upgrade(CLASS_HASH_ZERO());
+}
+
+#[test]
+#[available_gas(2000000)]
+fn test_persisting_selector_after_upgrade() {
+    let v1 = deploy_v1();
+    v1.set_value(VALUE);
+    assert(v1.get_value() == VALUE, 'Should be VALUE');
+
+    // Upgrade and check event
+    testing::pop_log_raw(ZERO());
+    v1.upgrade(V2_CLASS_HASH());
+    let event = testing::pop_log::<Upgraded>(v1.contract_address).unwrap();
+    assert(event.implementation == V2_CLASS_HASH(), 'Invalid implementation');
+
+    let v2 = IUpgradesV2Dispatcher { contract_address: v1.contract_address };
+    v2.set_value(VALUE2);
+    assert(v2.get_value() == VALUE2, 'Selector should be callable');
+}
+
+#[test]
+#[available_gas(2000000)]
+fn test_new_selector_after_upgrade() {
+    let v1 = deploy_v1();
+
+    v1.upgrade(V2_CLASS_HASH());
+    let v2 = IUpgradesV2Dispatcher { contract_address: v1.contract_address };
+
+    v2.set_value2(VALUE);
+    assert(v2.get_value2() == VALUE, 'New selector should be callable');
+}
+
+#[test]
+#[available_gas(2000000)]
+fn test_state_persists_after_upgrade() {
+    let v1 = deploy_v1();
+    v1.set_value(VALUE);
+
+    v1.upgrade(V2_CLASS_HASH());
+    let v2 = IUpgradesV2Dispatcher { contract_address: v1.contract_address };
+
+    assert(v2.get_value() == VALUE, 'Should keep state after upgrade');
+}
+
+#[test]
+#[available_gas(2000000)]
+fn test_remove_selector_passes_in_v1() {
+    let v1 = deploy_v1();
+    v1.remove_selector();
+}
+
+#[test]
+#[available_gas(2000000)]
+#[should_panic(expected: ('ENTRYPOINT_NOT_FOUND', ))]
+fn test_remove_selector_fails_in_v2() {
+    let v1 = deploy_v1();
+    v1.upgrade(V2_CLASS_HASH());
+    // We use the v1 dispatcher because remove_selector is not in v2 interface
+    v1.remove_selector();
+}
+
+//
+// upgrade_and_call
+//
+
+#[test]
+#[available_gas(2000000)]
+#[should_panic(expected: ('Class hash cannot be zero', 'ENTRYPOINT_FAILED', ))]
+fn test_upgrade_and_call_with_class_hash_zero() {
+    let v1 = deploy_v1();
+    let calldata = array![VALUE];
+    v1.upgrade_and_call(CLASS_HASH_ZERO(), SET_VALUE_SELECTOR, calldata.span());
+}
+
+#[test]
+#[available_gas(2000000)]
+fn test_upgrade_and_call_persisting_selector() {
+    let v1 = deploy_v1();
+    assert(v1.get_value() == 0, 'Should be zero');
+
+    testing::pop_log_raw(ZERO());
+
+    // upgrade_and_call and check event
+    let calldata = array![VALUE];
+    v1.upgrade_and_call(V2_CLASS_HASH(), SET_VALUE_SELECTOR, calldata.span());
+    let event = testing::pop_log::<Upgraded>(v1.contract_address).unwrap();
+    assert(event.implementation == V2_CLASS_HASH(), 'Invalid implementation');
+
+    // Check value is set
+    let v2 = IUpgradesV2Dispatcher { contract_address: v1.contract_address };
+    assert(v2.get_value() == VALUE, 'Should set val during upgrade');
+}
+
+#[test]
+#[available_gas(2000000)]
+fn test_upgrade_and_call_new_selector() {
+    let v1 = deploy_v1();
+
+    // upgrade_and_call
+    let calldata = array![VALUE2];
+    v1.upgrade_and_call(V2_CLASS_HASH(), SET_VALUE2_SELECTOR, calldata.span());
+
+    // Check value2 is set
+    let v2 = IUpgradesV2Dispatcher { contract_address: v1.contract_address };
+    assert(v2.get_value2() == VALUE2, 'Should set val during upgrade');
+}
+
+#[test]
+#[available_gas(2000000)]
+#[should_panic(expected: ('ENTRYPOINT_NOT_FOUND', 'ENTRYPOINT_FAILED'))]
+fn test_upgrade_and_call_with_removed_selector() {
+    let v1 = deploy_v1();
+    let calldata = array![];
+    v1.upgrade_and_call(V2_CLASS_HASH(), REMOVE_SELECTOR, calldata.span());
+}

src/openzeppelin/upgrades.cairo

@@ -0,0 +1 @@
+mod upgradeable;