From 9a920013d27b7ba8ef995bde01ddea8bf485c867 Mon Sep 17 00:00:00 2001 From: Sergio Garcia Date: Tue, 9 Jul 2024 15:00:15 +0100 Subject: [PATCH 1/3] no message --- src/contracts/gift_factory.cairo | 9 ++++---- src/contracts/timelock_upgrade.cairo | 32 +++++++++++++++++++++++++--- src/mocks/future_factory.cairo | 6 +++++- 3 files changed, 39 insertions(+), 8 deletions(-) diff --git a/src/contracts/gift_factory.cairo b/src/contracts/gift_factory.cairo index 2e5806e..4f4ab4e 100644 --- a/src/contracts/gift_factory.cairo +++ b/src/contracts/gift_factory.cairo @@ -75,6 +75,7 @@ mod GiftFactory { #[abi(embed_v0)] impl OwnableImpl = OwnableComponent::OwnableImpl; impl InternalImpl = OwnableComponent::InternalImpl; + impl TimelockUpgradeInternalImpl = TimelockUpgradeComponent::TimelockUpgradeInternalImpl; // Pausable component!(path: PausableComponent, storage: pausable, event: PausableEvent); @@ -227,10 +228,10 @@ mod GiftFactory { #[abi(embed_v0)] impl TimelockUpgradeCallbackImpl of ITimelockUpgradeCallback { - fn perform_upgrade(ref self: ContractState, new_implementation: ClassHash, data: Span) { - // This should do some sanity checks - // We should check that the new implementation is a valid implementation - // Execute the upgrade using replace_class_syscall(...) + fn perform_upgrade(ref self: ContractState, new_implementation: ClassHash, data: Array) { + self.timelock_upgrade.assert_and_reset_lock(); + // This should do some sanity checks and ensure that the new implementation is a valid implementation, + // then it will replace_class_syscall panic_with_felt252('gift-fac/downgrade-not-allowed'); } } diff --git a/src/contracts/timelock_upgrade.cairo b/src/contracts/timelock_upgrade.cairo index 445d4aa..3c4e4fd 100644 --- a/src/contracts/timelock_upgrade.cairo +++ b/src/contracts/timelock_upgrade.cairo @@ -40,7 +40,7 @@ pub trait ITimelockUpgradeCallback { /// @dev Currently empty as the upgrade logic will be handled in the contract we upgrade to /// @param new_implementation The class hash of the new implementation /// @param data The data to be used for the upgrade - fn perform_upgrade(ref self: TContractState, new_implementation: ClassHash, data: Span); + fn perform_upgrade(ref self: TContractState, new_implementation: ClassHash, data: Array); } #[starknet::component] @@ -62,7 +62,9 @@ pub mod TimelockUpgradeComponent { #[storage] pub struct Storage { - pending_upgrade: PendingUpgrade + pending_upgrade: PendingUpgrade, + /// true only during the upgrade call + upgrade_lock: bool, } #[event] @@ -141,14 +143,37 @@ pub mod TimelockUpgradeComponent { assert(current_timestamp < ready_at + VALID_WINDOW_PERIOD, 'upgrade/upgrade-too-late'); self.pending_upgrade.write(Default::default()); + + self.upgrade_lock.write(true); ITimelockUpgradeCallbackLibraryDispatcher { class_hash: implementation } - .perform_upgrade(implementation, calldata.span()); + .perform_upgrade(implementation, calldata); + self.upgrade_lock.write(false); } + fn get_pending_upgrade(self: @ComponentState) -> PendingUpgrade { self.pending_upgrade.read() } } + + #[generate_trait] + // #[embeddable_as(TimelockUpgradeInternalImpl)] + pub impl TimelockUpgradeInternalImpl< + TContractState, impl Ownable: OwnableComponent::HasComponent, +HasComponent + > of ITimelockUpgradeInternal { + /// @notice Should be called but the `perform_upgrade` method to make sure this method can only by called when upgrading + fn assert_and_reset_lock(ref self: ComponentState) { + assert(self.upgrade_lock.read(), 'upgrade/only-during-upgrade'); + self.upgrade_lock.write(false); + } + fn emit_upgrade_executed( + ref self: ComponentState, new_implementation: ClassHash, calldata: Array + ) { + self.emit(UpgradeExecuted { new_implementation, calldata }); + } + } + + #[generate_trait] impl PrivateImpl< TContractState, impl Ownable: OwnableComponent::HasComponent, +HasComponent @@ -159,6 +184,7 @@ pub mod TimelockUpgradeComponent { } } + impl DefaultClassHash of Default { fn default() -> ClassHash { Zero::zero() diff --git a/src/mocks/future_factory.cairo b/src/mocks/future_factory.cairo index 5ba9497..43e143c 100644 --- a/src/mocks/future_factory.cairo +++ b/src/mocks/future_factory.cairo @@ -19,6 +19,7 @@ mod FutureFactory { component!(path: TimelockUpgradeComponent, storage: timelock_upgrade, event: TimelockUpgradeEvent); #[abi(embed_v0)] impl TimelockUpgradeImpl = TimelockUpgradeComponent::TimelockUpgradeImpl; + impl TimelockUpgradeInternalImpl = TimelockUpgradeComponent::TimelockUpgradeInternalImpl; #[storage] struct Storage { @@ -48,8 +49,11 @@ mod FutureFactory { #[abi(embed_v0)] impl TimelockUpgradeCallbackImpl of ITimelockUpgradeCallback { - fn perform_upgrade(ref self: ContractState, new_implementation: ClassHash, data: Span) { + fn perform_upgrade(ref self: ContractState, new_implementation: ClassHash, data: Array) { + self.timelock_upgrade.assert_and_reset_lock(); starknet::syscalls::replace_class_syscall(new_implementation).unwrap(); + self.timelock_upgrade.emit_upgrade_executed(new_implementation, data); } } } + From 52b346fcdaeec81672dcfa29b5591360ba038a57 Mon Sep 17 00:00:00 2001 From: Sergio Garcia Date: Tue, 9 Jul 2024 15:19:46 +0100 Subject: [PATCH 2/3] improve upgrade mechanism --- src/contracts/gift_factory.cairo | 6 ++++-- src/contracts/timelock_upgrade.cairo | 5 ++--- tests-integration/upgrade.test.ts | 13 +++++++++++++ 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/src/contracts/gift_factory.cairo b/src/contracts/gift_factory.cairo index 4f4ab4e..5b6e302 100644 --- a/src/contracts/gift_factory.cairo +++ b/src/contracts/gift_factory.cairo @@ -231,8 +231,10 @@ mod GiftFactory { fn perform_upgrade(ref self: ContractState, new_implementation: ClassHash, data: Array) { self.timelock_upgrade.assert_and_reset_lock(); // This should do some sanity checks and ensure that the new implementation is a valid implementation, - // then it will replace_class_syscall - panic_with_felt252('gift-fac/downgrade-not-allowed'); + // then it can call replace_class_syscall and emit the UpgradeExecuted event + panic_with_felt252( + 'gift-fac/downgrade-not-allowed' + ); // since this is the first version nobody should be calling this method } } diff --git a/src/contracts/timelock_upgrade.cairo b/src/contracts/timelock_upgrade.cairo index 3c4e4fd..762a68b 100644 --- a/src/contracts/timelock_upgrade.cairo +++ b/src/contracts/timelock_upgrade.cairo @@ -157,11 +157,10 @@ pub mod TimelockUpgradeComponent { } #[generate_trait] - // #[embeddable_as(TimelockUpgradeInternalImpl)] pub impl TimelockUpgradeInternalImpl< - TContractState, impl Ownable: OwnableComponent::HasComponent, +HasComponent + TContractState, +HasComponent > of ITimelockUpgradeInternal { - /// @notice Should be called but the `perform_upgrade` method to make sure this method can only by called when upgrading + /// @notice Should be called by the `perform_upgrade` method to make sure this method can only by called when upgrading fn assert_and_reset_lock(ref self: ComponentState) { assert(self.upgrade_lock.read(), 'upgrade/only-during-upgrade'); self.upgrade_lock.write(false); diff --git a/tests-integration/upgrade.test.ts b/tests-integration/upgrade.test.ts index 2035528..250a03e 100644 --- a/tests-integration/upgrade.test.ts +++ b/tests-integration/upgrade.test.ts @@ -49,6 +49,11 @@ describe("Test Factory Upgrade", function () { newFactory.connect(deployer); await newFactory.get_num().should.eventually.equal(1n); + // we can't call the perform_upgrade method directly + await expectRevertWithErrorMessage("upgrade/only-during-upgrade", () => + factory.perform_upgrade(newFactoryClassHash, []), + ); + // clear deployment cache delete protocolCache["GiftFactory"]; }); @@ -79,6 +84,14 @@ describe("Test Factory Upgrade", function () { await expectRevertWithErrorMessage("Caller is not the owner", () => factory.upgrade([])); }); + it("no calls to perform_upgrade", async function () { + const { factory } = await setupGiftProtocol(); + const newFactoryClassHash = "0x1"; + await expectRevertWithErrorMessage("upgrade/only-during-upgrade", () => + factory.perform_upgrade(newFactoryClassHash, []), + ); + }); + it("Invalid Calldata", async function () { const { factory } = await setupGiftProtocol(); const newFactoryClassHash = "0x1"; From 7a8466d2305f09d7e485c848a19c068c6bea28bb Mon Sep 17 00:00:00 2001 From: Sergio Garcia Date: Fri, 12 Jul 2024 09:53:33 +0100 Subject: [PATCH 3/3] assert lock was reset --- src/contracts/timelock_upgrade.cairo | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/contracts/timelock_upgrade.cairo b/src/contracts/timelock_upgrade.cairo index 762a68b..0688cfb 100644 --- a/src/contracts/timelock_upgrade.cairo +++ b/src/contracts/timelock_upgrade.cairo @@ -147,7 +147,7 @@ pub mod TimelockUpgradeComponent { self.upgrade_lock.write(true); ITimelockUpgradeCallbackLibraryDispatcher { class_hash: implementation } .perform_upgrade(implementation, calldata); - self.upgrade_lock.write(false); + assert(!self.upgrade_lock.read(), 'upgrade/lock-not-reset') }