Update Upgrade and add upgrade tests

.gitignore

@@ -42,4 +42,5 @@ target
 .snfoundry_cache/
 
 .env
-dist
+dist
+.DS_STORE

lib/claim.ts

@@ -94,8 +94,8 @@ export async function signExternalClaim(signParams: {
 export async function claimExternal(args: {
   claim: Claim;
   receiver: string;
-  dustReceiver?: string;
   claimPrivateKey: string;
+  dustReceiver?: string;
   overrides?: { claimAccountAddress?: string; factoryAddress?: string; account?: Account };
   details?: UniversalDetails;
 }): Promise<TransactionReceipt> {

lib/protocol.ts

@@ -1,11 +1,11 @@
 import { Contract, byteArray, uint256 } from "starknet";
 import { deployer, manager } from ".";
 
-const cache: Record<string, Contract> = {};
+export const protocolCache: Record<string, Contract> = {};
 
 export async function deployMockERC20(): Promise<Contract> {
-  if (cache["MockERC20"]) {
-    return cache["MockERC20"];
+  if (protocolCache["MockERC20"]) {
+    return protocolCache["MockERC20"];
   }
   const mockERC20 = await manager.deployContract("MockERC20", {
     unique: true,
@@ -17,7 +17,7 @@ export async function deployMockERC20(): Promise<Contract> {
       deployer.address,
     ],
   });
-  cache["MockERC20"] = mockERC20;
+  protocolCache["MockERC20"] = mockERC20;
   return mockERC20;
 }
 
@@ -26,14 +26,14 @@ export async function setupGiftProtocol(): Promise<{
   claimAccountClassHash: string;
 }> {
   const claimAccountClassHash = await manager.declareLocalContract("ClaimAccount");
-  const cachedFactory = cache["GiftFactory"];
+  const cachedFactory = protocolCache["GiftFactory"];
   if (cachedFactory) {
     return { factory: cachedFactory, claimAccountClassHash };
   }
   const factory = await manager.deployContract("GiftFactory", {
     unique: true,
     constructorCalldata: [claimAccountClassHash, deployer.address],
   });
-  cache["GiftFactory"] = factory;
+  protocolCache["GiftFactory"] = factory;
   return { factory, claimAccountClassHash };
 }

src/contracts/gift_factory.cairo

@@ -13,7 +13,7 @@ mod GiftFactory {
     use starknet_gifting::contracts::claim_hash::{ClaimExternal, IOffChainMessageHashRev1};
     use starknet_gifting::contracts::interface::{
         IGiftAccountDispatcherTrait, IGiftFactory, ClaimData, AccountConstructorArguments, IGiftAccountDispatcher,
-        OutsideExecution, GiftStatus, StarknetSignature
+        OutsideExecution, StarknetSignature
     };
     use starknet_gifting::contracts::timelock_upgrade::{ITimelockUpgradeCallback, TimelockUpgradeComponent};
     use starknet_gifting::contracts::utils::{
@@ -248,28 +248,6 @@ mod GiftFactory {
             self.claim_class_hash.read()
         }
 
-        fn get_gift_status(self: @ContractState, claim: ClaimData) -> GiftStatus {
-            let claim_address = self.check_claim_and_get_account_address(claim);
-            let gift_balance = IERC20Dispatcher { contract_address: claim.gift_token }.balance_of(claim_address);
-            if gift_balance < claim.gift_amount {
-                return GiftStatus::ClaimedOrCancelled;
-            }
-            if (claim.gift_token == claim.fee_token) {
-                if gift_balance < claim.gift_amount + claim.fee_amount.into() {
-                    return GiftStatus::ReadyExternalOnly;
-                } else {
-                    return GiftStatus::Ready;
-                }
-            } else {
-                let fee_balance = IERC20Dispatcher { contract_address: claim.fee_token }.balance_of(claim_address);
-                if fee_balance < claim.fee_amount.into() {
-                    return GiftStatus::ReadyExternalOnly;
-                } else {
-                    return GiftStatus::Ready;
-                }
-            }
-        }
-
         fn get_claim_address(
             self: @ContractState,
             class_hash: ClassHash,
@@ -295,7 +273,7 @@ mod GiftFactory {
         }
     }
 
-
+    #[abi(embed_v0)]
     impl TimelockUpgradeCallbackImpl of ITimelockUpgradeCallback<ContractState> {
         fn perform_upgrade(ref self: ContractState, new_implementation: ClassHash, data: Span<felt252>) {
             // This should do some sanity checks 

src/contracts/interface.cairo

@@ -99,10 +99,6 @@ pub trait IGiftFactory<TContractState> {
         fee_amount: u128,
         claim_pubkey: felt252
     ) -> ContractAddress;
-
-    /// @notice Get the status of a claim
-    /// @param claim The claim data 
-    fn get_gift_status(self: @TContractState, claim: ClaimData) -> GiftStatus;
 }
 
 
@@ -184,12 +180,3 @@ pub struct AccountConstructorArguments {
     pub fee_amount: u128,
     pub claim_pubkey: felt252
 }
-
-/// @notice Enum representing the status of a gift
-/// @dev ReadyExternalOnly should only happen if there is no fee_amount or if the account reverted during claim_internal
-#[derive(Serde, Drop, Copy)]
-pub enum GiftStatus {
-    ClaimedOrCancelled,
-    Ready,
-    ReadyExternalOnly
-}

src/contracts/timelock_upgrade.cairo

@@ -7,7 +7,8 @@ pub trait ITimelockUpgrade<TContractState> {
     /// @dev After the 7-day waiting period, the upgrade can be performed within a 7-day window
     /// @dev If there is an ongoing upgrade, the previous proposition will be overwritten
     /// @param new_implementation The class hash of the new implementation
-    fn propose_upgrade(ref self: TContractState, new_implementation: ClassHash);
+    /// @param calldata The calldata to be used for the upgrade
+    fn propose_upgrade(ref self: TContractState, new_implementation: ClassHash, calldata: Array<felt252>);
 
     /// @notice Cancel the upgrade proposition
     /// @dev Will fail if there is no ongoing upgrade
@@ -23,6 +24,9 @@ pub trait ITimelockUpgrade<TContractState> {
 
     /// @notice Gets the timestamp when the upgrade is ready to be performed, 0 if no upgrade ongoing
     fn get_upgrade_ready_at(self: @TContractState) -> u64;
+
+    /// @notice Gets the hash of the calldata used for the upgrade, 0 if no upgrade ongoing
+    fn get_calldata_hash(self: @TContractState) -> felt252;
 }
 
 #[starknet::interface]
@@ -37,6 +41,7 @@ pub trait ITimelockUpgradeCallback<TContractState> {
 #[starknet::component]
 pub mod TimelockUpgradeComponent {
     use core::num::traits::Zero;
+    use core::poseidon::poseidon_hash_span;
     use openzeppelin::access::ownable::{OwnableComponent, OwnableComponent::InternalTrait};
     use starknet::{get_block_timestamp, ClassHash};
     use super::{
@@ -53,6 +58,7 @@ pub mod TimelockUpgradeComponent {
     pub struct Storage {
         pending_implementation: ClassHash,
         ready_at: u64,
+        calldata_hash: felt252,
     }
 
     #[event]
@@ -66,12 +72,13 @@ pub mod TimelockUpgradeComponent {
     #[derive(Drop, starknet::Event)]
     struct UpgradeProposed {
         new_implementation: ClassHash,
-        ready_at: u64
+        ready_at: u64,
+        calldata: Array<felt252>
     }
 
     #[derive(Drop, starknet::Event)]
     struct UpgradeCancelled {
-        new_implementation: ClassHash
+        cancelled_implementation: ClassHash
     }
 
     #[derive(Drop, starknet::Event)]
@@ -83,30 +90,34 @@ pub mod TimelockUpgradeComponent {
     impl TimelockUpgrade<
         TContractState,
         +HasComponent<TContractState>,
-        impl Ownable: OwnableComponent::HasComponent<TContractState>,
+        +OwnableComponent::HasComponent<TContractState>,
         +ITimelockUpgradeCallback<TContractState>,
     > of ITimelockUpgrade<ComponentState<TContractState>> {
-        fn propose_upgrade(ref self: ComponentState<TContractState>, new_implementation: ClassHash) {
+        fn propose_upgrade(
+            ref self: ComponentState<TContractState>, new_implementation: ClassHash, calldata: Array<felt252>
+        ) {
             self.assert_only_owner();
             assert(new_implementation.is_non_zero(), 'upgrade/new-implementation-null');
 
             let pending_implementation = self.pending_implementation.read();
             if pending_implementation.is_non_zero() {
-                self.emit(UpgradeCancelled { new_implementation: pending_implementation })
+                self.emit(UpgradeCancelled { cancelled_implementation: pending_implementation })
             }
 
             self.pending_implementation.write(new_implementation);
             let ready_at = get_block_timestamp() + MIN_SECURITY_PERIOD;
             self.ready_at.write(ready_at);
-            self.emit(UpgradeProposed { new_implementation, ready_at });
+            let calldata_hash = poseidon_hash_span(calldata.span());
+            self.calldata_hash.write(calldata_hash);
+            self.emit(UpgradeProposed { new_implementation, ready_at, calldata });
         }
 
         fn cancel_upgrade(ref self: ComponentState<TContractState>) {
             self.assert_only_owner();
-            let new_implementation = self.pending_implementation.read();
-            assert(new_implementation.is_non_zero(), 'upgrade/no-new-implementation');
+            let proposed_implementation = self.pending_implementation.read();
+            assert(proposed_implementation.is_non_zero(), 'upgrade/no-new-implementation');
             assert(self.ready_at.read() != 0, 'upgrade/not-ready');
-            self.emit(UpgradeCancelled { new_implementation });
+            self.emit(UpgradeCancelled { cancelled_implementation: proposed_implementation });
             self.reset_storage();
         }
 
@@ -115,6 +126,8 @@ pub mod TimelockUpgradeComponent {
             let new_implementation = self.pending_implementation.read();
             let ready_at = self.ready_at.read();
             let block_timestamp = get_block_timestamp();
+            let calldata_hash = poseidon_hash_span(calldata.span());
+            assert(calldata_hash == self.calldata_hash.read(), 'upgrade/invalid-calldata');
             assert(new_implementation.is_non_zero(), 'upgrade/no-pending-upgrade');
             assert(block_timestamp >= ready_at, 'upgrade/too-early');
             assert(block_timestamp < ready_at + VALID_WINDOW_PERIOD, 'upgrade/upgrade-too-late');
@@ -130,6 +143,10 @@ pub mod TimelockUpgradeComponent {
         fn get_upgrade_ready_at(self: @ComponentState<TContractState>) -> u64 {
             self.ready_at.read()
         }
+
+        fn get_calldata_hash(self: @ComponentState<TContractState>) -> felt252 {
+            self.calldata_hash.read()
+        }
     }
     #[generate_trait]
     impl PrivateImpl<
@@ -142,6 +159,7 @@ pub mod TimelockUpgradeComponent {
         fn reset_storage(ref self: ComponentState<TContractState>) {
             self.pending_implementation.write(Zero::zero());
             self.ready_at.write(0);
+            self.calldata_hash.write(0);
         }
     }
 }