From ede990f539fab8776b1fa6c3b35ed59c04d8217d Mon Sep 17 00:00:00 2001
From: Patrice Tisserand <p.tisserand@gmail.com>
Date: Wed, 27 Mar 2024 03:07:28 +0100
Subject: [PATCH] starknet: add `collection_upgrade` to upgrade collection
 owned by bridge.

---
 apps/blockchain/starknet/src/bridge.cairo     | 14 ++++-
 apps/blockchain/starknet/src/interfaces.cairo |  6 ++
 .../starknet/src/tests/bridge_t.cairo         | 60 ++++++++++++++++++-
 3 files changed, 78 insertions(+), 2 deletions(-)

diff --git a/apps/blockchain/starknet/src/bridge.cairo b/apps/blockchain/starknet/src/bridge.cairo
index fdbf6d71..a44e76d1 100644
--- a/apps/blockchain/starknet/src/bridge.cairo
+++ b/apps/blockchain/starknet/src/bridge.cairo
@@ -14,7 +14,7 @@ mod bridge {
     use starknet::contract_address::ContractAddressZeroable;
     use starknet::eth_address::EthAddressZeroable;
 
-    use starklane::interfaces::{IStarklane, IUpgradeable};
+    use starklane::interfaces::{IStarklane, IUpgradeable, IUpgradeableDispatcher, IUpgradeableDispatcherTrait, IStarklaneCollectionAdmin};
     // events
     use starklane::interfaces::{
         DepositRequestInitiated,
@@ -324,6 +324,18 @@ mod bridge {
         }
     }
 
+    #[abi(embed_v0)]
+    impl BridgeCollectionAdminImpl of IStarklaneCollectionAdmin<ContractState> {
+        fn collection_upgrade(ref self: ContractState, collection: ContractAddress, class_hash: ClassHash) {
+            assert(
+                starknet::get_caller_address() == self.bridge_admin.read(),
+                'Unauthorized replace class'
+            );
+            IUpgradeableDispatcher { contract_address: collection }
+                .upgrade(class_hash);
+        }
+    }
+
     // *** INTERNALS ***
 
     /// Ensures the caller is the bridge admin. Revert if it's not.
diff --git a/apps/blockchain/starknet/src/interfaces.cairo b/apps/blockchain/starknet/src/interfaces.cairo
index 58b6b7fd..b7051e79 100644
--- a/apps/blockchain/starknet/src/interfaces.cairo
+++ b/apps/blockchain/starknet/src/interfaces.cairo
@@ -41,6 +41,12 @@ trait IUpgradeable<T> {
     fn upgrade(ref self: T, class_hash: ClassHash);
 }
 
+#[starknet::interface]
+trait IStarklaneCollectionAdmin<T> {
+    // try to upgrade the given collection with given class_hash
+    fn collection_upgrade(ref self: T, collection: ContractAddress, class_hash: ClassHash);
+}
+
 //////////////////////////
 /// Events
 #[derive(Drop, starknet::Event)]
diff --git a/apps/blockchain/starknet/src/tests/bridge_t.cairo b/apps/blockchain/starknet/src/tests/bridge_t.cairo
index ce431a04..7ec62b6e 100644
--- a/apps/blockchain/starknet/src/tests/bridge_t.cairo
+++ b/apps/blockchain/starknet/src/tests/bridge_t.cairo
@@ -19,7 +19,10 @@ mod tests {
     use starknet::{ContractAddress, ClassHash, EthAddress};
     use starklane::{
         request::{Request, compute_request_hash},
-        interfaces::{IStarklaneDispatcher, IStarklaneDispatcherTrait, IUpgradeableDispatcher, IUpgradeableDispatcherTrait},
+        interfaces::{IStarklaneDispatcher, IStarklaneDispatcherTrait, 
+                    IUpgradeableDispatcher, IUpgradeableDispatcherTrait,
+                    IStarklaneCollectionAdminDispatcher, IStarklaneCollectionAdminDispatcherTrait,
+                    },
     };
     use starklane::token::{
         interfaces::{
@@ -594,4 +597,59 @@ mod tests {
         stop_prank(CheatTarget::One(bridge_address));        
     }
 
+    #[test]
+    fn collection_upgrade_as_admin() {
+        // Need to declare here to get the class hash before deploy anything.
+        let erc721b_contract_class = declare("erc721_bridgeable");
+
+        let BRIDGE_ADMIN = starknet::contract_address_const::<'starklane'>();
+        let BRIDGE_L1 = EthAddress { address: 'starklane_l1' };
+        let OWNER_L2 = starknet::contract_address_const::<'owner_l2'>();
+
+        let bridge_address = deploy_starklane(BRIDGE_ADMIN, BRIDGE_L1, erc721b_contract_class.class_hash);
+
+        let collection_l1: EthAddress = 0xe0c.try_into().unwrap();
+        let collection_l2: ContractAddress = 0x0.try_into().unwrap();
+        let owner_l1: EthAddress = 0xe00.try_into().unwrap();
+        let owner_l2: ContractAddress = OWNER_L2.into();
+        let ids: Span<u256> = array![0_u256, 1_u256].span();
+
+        let mut req = setup_request(
+            collection_l1,
+            collection_l2,
+            owner_l1,
+            owner_l2,
+            "name",
+            "symbol",
+            "base_uri"
+        );
+        req.ids = ids;
+        let mut buf = array![];
+        req.serialize(ref buf);
+        
+        let mut l1_handler = L1Handler {
+            contract_address: bridge_address,
+            // selector: 0x03593216f3a8b22f4cf375e5486e3d13bfde9d0f26976d20ac6f653c73f7e507,
+            function_selector: selector!("withdraw_auto_from_l1"), 
+            from_address: BRIDGE_L1.into(),
+            payload: buf.span()
+        };
+        let mut spy = spy_events(SpyOn::One(bridge_address));
+
+        l1_handler.execute().unwrap();
+        let bridge = IStarklaneDispatcher { contract_address: bridge_address };
+
+        // Deserialize the request and check some expected values.
+        let mut sp = buf.span();
+        let req = Serde::<Request>::deserialize(ref sp).unwrap();
+
+        let deployed_address = bridge.get_l2_collection_address(req.collection_l1.into());
+        assert!(!deployed_address.is_zero(), "Expected deployed erc721");
+        assert!(get_class_hash(deployed_address) == erc721b_contract_class.class_hash, "Expected class hash");
+
+        start_prank(CheatTarget::One(bridge_address), BRIDGE_ADMIN);
+        IStarklaneCollectionAdminDispatcher { contract_address: bridge_address}.collection_upgrade(deployed_address, get_class_hash(bridge_address));
+        stop_prank(CheatTarget::One(bridge_address));
+        assert!(get_class_hash(deployed_address) == get_class_hash(bridge_address), "Class hash not updated");
+    }
 }