Add upgrade_and_call

CHANGELOG.md

@@ -9,6 +9,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- `IUpgradeAndCall` interface (#1148)
+- `upgrade_and_call` function in UpgradeableComponent's InternalImpl (#1148)
+
 ## 0.17.0 (2024-09-23)
 
 ### Added

docs/modules/ROOT/pages/api/upgrades.adoc

@@ -36,6 +36,35 @@ Upgrades the contract code by updating its {class_hash}.
 
 NOTE: This function is usually protected by an xref:access.adoc[Access Control] mechanism.
 
+[.contract]
+[[IUpgradeAndCall]]
+=== `++IUpgradeAndCall++` link:https://github.com/OpenZeppelin/cairo-contracts/blob/release-v0.16.0/packages/upgrades/src/interface.cairo[{github-icon},role=heading-link]
+
+:Upgraded: xref:UpgradeableComponent-Upgraded[Upgraded]
+
+```cairo
+use openzeppelin::upgrades::interface::IUpgradeAndCall;
+```
+
+Interface for an upgradeable contract that couples an upgrade with a function call in the upgraded context.
+
+[.contract-index]
+.Functions
+--
+* xref:#IUpgradeAndCall-upgrade_and_call[`++upgrade_and_call(new_class_hash, selector, calldata)++`]
+--
+
+[#IUpgradeAndCall-Functions]
+==== Functions
+
+[.contract-item]
+[[IUpgradeAndCall-upgrade_and_call]]
+==== `[.contract-item-name]#++upgrade_and_call++#++(new_class_hash: ClassHash, selector: felt252, calldata: Span<felt252>) → Span<felt252>++` [.item-kind]#external#
+
+Upgrades the contract code by updating its {class_hash} and calls `selector` with the upgraded context.
+
+NOTE: This function is usually protected by an xref:access.adoc[Access Control] mechanism.
+
 [.contract]
 [[UpgradeableComponent]]
 === `++UpgradeableComponent++` link:https://github.com/OpenZeppelin/cairo-contracts/blob/release-v0.17.0/packages/upgrades/src/upgradeable.cairo[{github-icon},role=heading-link]
@@ -52,6 +81,7 @@ Upgradeable component.
 .InternalImpl
 
 * xref:#UpgradeableComponent-upgrade[`++upgrade(self, new_class_hash)++`]
+* xref:#UpgradeableComponent-upgrade_and_call[`++upgrade_and_call(self, new_class_hash, selector, calldata)++`]
 --
 
 [.contract-index]
@@ -63,6 +93,8 @@ Upgradeable component.
 [#UpgradeableComponent-Internal-Functions]
 ==== Internal Functions
 
+:snip6: https://github.com/starknet-io/SNIPs/blob/main/SNIPS/snip-6.md[SNIP-6]
+
 [.contract-item]
 [[UpgradeableComponent-upgrade]]
 ==== `[.contract-item-name]#++upgrade++#++(ref self: ContractState, new_class_hash: ClassHash)++` [.item-kind]#internal#
@@ -75,6 +107,26 @@ Requirements:
 
 Emits an {Upgraded} event.
 
+[.contract-item]
+[[UpgradeableComponent-upgrade_and_call]]
+==== `[.contract-item-name]#++upgrade_and_call++#++(ref self: ContractState, new_class_hash: ClassHash, selector: felt252, calldata: Span<felt252>) → Span<felt252>++` [.item-kind]#internal#
+
+Replaces the contract's class hash with `new_class_hash` and then calls `selector`
+from the upgraded context.
+This function returns the unwrapped `call_contract_syscall` return value(s), if available, of the `selector` call.
+
+Requirements:
+
+- `new_class_hash` must be different from zero.
+
+NOTE: The function call comes from the upgraded contract itself and not the account.
+
+NOTE: A similar behavior to `upgrade_and_call` can also be achieved with a list of calls from an account since the {snip6} account standard supports multicall.
+An account can execute a list of calls with xref:IUpgradeable-upgrade[upgrade] being the first element in the list and the extra function call as the second.
+With this approach, the calls will execute from the account's context and can't be front-ran.
+
+Emits an {Upgraded} event.
+
 [#UpgradeableComponent-Events]
 ==== Events
 

packages/upgrades/README.md

@@ -7,6 +7,7 @@ This crate provides an interface and component used for upgradeability.
 ### Interfaces
 
 - `IUpgradeable`
+- `IUpgradeAndCall`
 
 ### Components
 

packages/upgrades/src/interface.cairo

@@ -7,3 +7,10 @@ use starknet::ClassHash;
 pub trait IUpgradeable<TState> {
     fn upgrade(ref self: TState, new_class_hash: ClassHash);
 }
+
+#[starknet::interface]
+pub trait IUpgradeAndCall<TState> {
+    fn upgrade_and_call(
+        ref self: TState, new_class_hash: ClassHash, selector: felt252, calldata: Span<felt252>
+    ) -> Span<felt252>;
+}

packages/upgrades/src/tests/mocks/upgrades_mocks.cairo

@@ -7,6 +7,9 @@ use starknet::ClassHash;
 #[starknet::interface]
 pub(crate) trait IUpgradesV1<TState> {
     fn upgrade(ref self: TState, new_class_hash: ClassHash);
+    fn upgrade_and_call(
+        ref self: TState, new_class_hash: ClassHash, selector: felt252, calldata: Span<felt252>
+    ) -> Span<felt252>;
     fn set_value(ref self: TState, val: felt252);
     fn get_value(self: @TState) -> felt252;
     fn remove_selector(self: @TState);
@@ -42,6 +45,15 @@ pub(crate) mod UpgradesV1 {
             self.upgradeable.upgrade(new_class_hash);
         }
 
+        fn upgrade_and_call(
+            ref self: ContractState,
+            new_class_hash: ClassHash,
+            selector: felt252,
+            calldata: Span<felt252>
+        ) -> Span<felt252> {
+            self.upgradeable.upgrade_and_call(new_class_hash, selector, calldata)
+        }
+
         fn set_value(ref self: ContractState, val: felt252) {
             self.value.write(val);
         }
@@ -57,6 +69,9 @@ pub(crate) mod UpgradesV1 {
 #[starknet::interface]
 pub(crate) trait IUpgradesV2<TState> {
     fn upgrade(ref self: TState, new_class_hash: ClassHash);
+    fn upgrade_and_call(
+        ref self: TState, new_class_hash: ClassHash, selector: felt252, calldata: Span<felt252>
+    ) -> Span<felt252>;
     fn set_value(ref self: TState, val: felt252);
     fn set_value2(ref self: TState, val: felt252);
     fn get_value(self: @TState) -> felt252;
@@ -94,6 +109,15 @@ pub(crate) mod UpgradesV2 {
             self.upgradeable.upgrade(new_class_hash);
         }
 
+        fn upgrade_and_call(
+            ref self: ContractState,
+            new_class_hash: ClassHash,
+            selector: felt252,
+            calldata: Span<felt252>
+        ) -> Span<felt252> {
+            self.upgradeable.upgrade_and_call(new_class_hash, selector, calldata)
+        }
+
         fn set_value(ref self: ContractState, val: felt252) {
             self.value.write(val);
         }

packages/upgrades/src/tests/test_upgradeable.cairo

@@ -1,6 +1,10 @@
+use crate::tests::mocks::upgrades_mocks::{
+    IUpgradesV1SafeDispatcher, IUpgradesV1SafeDispatcherTrait
+};
 use crate::tests::mocks::upgrades_mocks::{IUpgradesV1Dispatcher, IUpgradesV1DispatcherTrait};
 use crate::tests::mocks::upgrades_mocks::{IUpgradesV2Dispatcher, IUpgradesV2DispatcherTrait};
 use openzeppelin_test_common::upgrades::UpgradeableSpyHelpers;
+use openzeppelin_testing as utils;
 use openzeppelin_testing::constants::{CLASS_HASH_ZERO, FELT_VALUE as VALUE};
 use openzeppelin_testing::{declare_class, deploy};
 use snforge_std::{spy_events, ContractClass};
@@ -69,12 +73,84 @@ fn test_remove_selector_passes_in_v1() {
 }
 
 #[test]
-#[ignore] // REASON: should_panic attribute not fit for complex panic messages.
-#[should_panic(expected: ('ENTRYPOINT_NOT_FOUND',))]
+#[feature("safe_dispatcher")]
 fn test_remove_selector_fails_in_v2() {
     let (v1, v2_class) = setup_test();
 
     v1.upgrade(v2_class.class_hash);
-    // We use the v1 dispatcher because remove_selector is not in v2 interface
-    v1.remove_selector();
+
+    // We use the v1 dispatcher because `remove_selector` is not in v2 interface
+    let safe_dispatcher = IUpgradesV1SafeDispatcher { contract_address: v1.contract_address };
+    let mut result = safe_dispatcher.remove_selector();
+    utils::assert_entrypoint_not_found_error(
+        result, selector!("remove_selector"), v1.contract_address
+    );
+}
+
+
+//
+// upgrade_and_call
+//
+
+#[test]
+#[should_panic(expected: ('Class hash cannot be zero',))]
+fn test_upgrade_and_call_with_class_hash_zero() {
+    let (v1, _) = setup_test();
+    let calldata = array![VALUE];
+    v1.upgrade_and_call(CLASS_HASH_ZERO(), selector!("set_value2"), calldata.span());
+}
+
+#[test]
+fn test_upgrade_and_call_with_new_selector() {
+    let (v1, v2_class) = setup_test();
+    let mut spy = spy_events();
+
+    let calldata = array![VALUE];
+    let new_selector = selector!("set_value2");
+    v1.upgrade_and_call(v2_class.class_hash, new_selector, calldata.span());
+    spy.assert_only_event_upgraded(v1.contract_address, v2_class.class_hash);
+
+    let v2 = IUpgradesV2Dispatcher { contract_address: v1.contract_address };
+    assert_eq!(v2.get_value2(), VALUE);
+}
+
+#[test]
+fn test_upgrade_and_call_with_return_value() {
+    let (v1, v2_class) = setup_test();
+
+    // Set value to get with upgrade_and_call
+    v1.set_value(VALUE);
+
+    let calldata = array![];
+    let selector = selector!("get_value");
+    let call_res = v1.upgrade_and_call(v2_class.class_hash, selector, calldata.span());
+
+    assert_eq!(call_res.len(), 1, "Return span should include one value");
+    assert_eq!(*call_res.at(0), VALUE);
+}
+
+#[test]
+fn test_upgrade_and_call_with_no_return_value() {
+    let (v1, v2_class) = setup_test();
+
+    let calldata = array![VALUE];
+    let selector = selector!("set_value");
+    let call_res = v1.upgrade_and_call(v2_class.class_hash, selector, calldata.span());
+
+    assert_eq!(call_res.len(), 0);
+}
+
+#[test]
+#[ignore] // REASON:
+#[feature("safe_dispatcher")]
+fn test_upgrade_and_call_with_removed_selector() {
+    let (v1, v2_class) = setup_test();
+    let removed_selector = selector!("remove_selector");
+    let calldata = array![];
+
+    // We use the v1 dispatcher because `remove_selector` is not in v2 interface
+    let safe_dispatcher = IUpgradesV1SafeDispatcher { contract_address: v1.contract_address };
+    let mut result = safe_dispatcher
+        .upgrade_and_call(v2_class.class_hash, removed_selector, calldata.span());
+    utils::assert_entrypoint_not_found_error(result, removed_selector, v1.contract_address);
 }

packages/upgrades/src/upgradeable.cairo

@@ -9,6 +9,7 @@ pub mod UpgradeableComponent {
     use core::num::traits::Zero;
     use starknet::ClassHash;
     use starknet::SyscallResultTrait;
+    use starknet::syscalls::{call_contract_syscall, replace_class_syscall};
 
     #[storage]
     pub struct Storage {}
@@ -42,8 +43,32 @@ pub mod UpgradeableComponent {
         /// Emits an `Upgraded` event.
         fn upgrade(ref self: ComponentState<TContractState>, new_class_hash: ClassHash) {
             assert(!new_class_hash.is_zero(), Errors::INVALID_CLASS);
-            starknet::syscalls::replace_class_syscall(new_class_hash).unwrap_syscall();
+            replace_class_syscall(new_class_hash).unwrap_syscall();
             self.emit(Upgraded { class_hash: new_class_hash });
         }
+
+        /// Replaces the contract's class hash with `new_class_hash` and then calls `selector`
+        /// from the upgraded context.
+        /// This function returns the unwrapped `call_contract_syscall` return value(s), if
+        /// available, of the `selector` call.
+        ///
+        /// Requirements:
+        ///
+        /// - `new_class_hash` is not zero.
+        ///
+        /// Emits an `Upgraded` event.
+        fn upgrade_and_call(
+            ref self: ComponentState<TContractState>,
+            new_class_hash: ClassHash,
+            selector: felt252,
+            calldata: Span<felt252>
+        ) -> Span<felt252> {
+            self.upgrade(new_class_hash);
+            let this = starknet::get_contract_address();
+            // `call_contract_syscall` is used in order to call `selector` from the new class.
+            // See:
+            // https://docs.starknet.io/documentation/architecture_and_concepts/Contracts/system-calls-cairo1/#replace_class
+            call_contract_syscall(this, selector, calldata).unwrap_syscall()
+        }
     }
 }