feat(target_chains/starknet): support SetFeeInToken instruction (#1669)

Riateche · web-flow · commit 0f7141ef2b31 · 2024-06-07T12:12:55.000+01:00
diff --git a/target_chains/starknet/contracts/src/pyth.cairo b/target_chains/starknet/contracts/src/pyth.cairo
@@ -90,6 +90,7 @@ mod pyth {
     pub struct FeeSet {
         pub old_fee: u256,
         pub new_fee: u256,
+        pub token: ContractAddress,
     }
 
     #[derive(Drop, Clone, Debug, PartialEq, Serde, starknet::Event)]
@@ -414,11 +415,10 @@ mod pyth {
             }
             match instruction.payload {
                 GovernancePayload::SetFee(payload) => {
-                    let new_fee = apply_decimal_expo(payload.value, payload.expo);
-                    let old_fee = self.single_update_fee1.read();
-                    self.single_update_fee1.write(new_fee);
-                    let event = FeeSet { old_fee, new_fee };
-                    self.emit(event);
+                    self.set_fee(payload.value, payload.expo, self.fee_token_address1.read());
+                },
+                GovernancePayload::SetFeeInToken(payload) => {
+                    self.set_fee(payload.value, payload.expo, payload.token);
                 },
                 GovernancePayload::SetDataSources(payload) => {
                     let new_data_sources = payload.sources;
@@ -716,6 +716,23 @@ mod pyth {
             };
             output_array
         }
+
+        fn set_fee(ref self: ContractState, value: u64, expo: u64, token: ContractAddress) {
+            let new_fee = apply_decimal_expo(value, expo);
+            let old_fee = if token == self.fee_token_address1.read() {
+                let old_fee = self.single_update_fee1.read();
+                self.single_update_fee1.write(new_fee);
+                old_fee
+            } else if token == self.fee_token_address2.read() {
+                let old_fee = self.single_update_fee2.read();
+                self.single_update_fee2.write(new_fee);
+                old_fee
+            } else {
+                panic_with_felt252(GovernanceActionError::InvalidGovernanceMessage.into())
+            };
+            let event = FeeSet { old_fee, new_fee, token };
+            self.emit(event);
+        }
     }
 
     fn apply_decimal_expo(value: u64, expo: u64) -> u256 {
diff --git a/target_chains/starknet/contracts/src/pyth/governance.cairo b/target_chains/starknet/contracts/src/pyth/governance.cairo
@@ -19,6 +19,7 @@ pub enum GovernanceAction {
     SetValidPeriod,
     RequestGovernanceDataSourceTransfer,
     SetWormholeAddress,
+    SetFeeInToken,
 }
 
 impl U8TryIntoGovernanceAction of TryInto<u8, GovernanceAction> {
@@ -31,6 +32,7 @@ impl U8TryIntoGovernanceAction of TryInto<u8, GovernanceAction> {
             4 => GovernanceAction::SetValidPeriod,
             5 => GovernanceAction::RequestGovernanceDataSourceTransfer,
             6 => GovernanceAction::SetWormholeAddress,
+            7 => GovernanceAction::SetFeeInToken,
             _ => { return Option::None; }
         };
         Option::Some(v)
@@ -52,6 +54,7 @@ pub enum GovernancePayload {
     // SetValidPeriod is unsupported
     RequestGovernanceDataSourceTransfer: RequestGovernanceDataSourceTransfer,
     SetWormholeAddress: SetWormholeAddress,
+    SetFeeInToken: SetFeeInToken,
 }
 
 #[derive(Drop, Clone, Debug, PartialEq, Serde)]
@@ -60,6 +63,13 @@ pub struct SetFee {
     pub expo: u64,
 }
 
+#[derive(Drop, Clone, Debug, PartialEq, Serde)]
+pub struct SetFeeInToken {
+    pub value: u64,
+    pub expo: u64,
+    pub token: ContractAddress,
+}
+
 #[derive(Drop, Clone, Debug, PartialEq, Serde)]
 pub struct SetDataSources {
     pub sources: Array<DataSource>,
@@ -155,6 +165,22 @@ pub fn parse_instruction(payload: ByteBuffer) -> GovernanceInstruction {
             let expo = reader.read_u64();
             GovernancePayload::SetFee(SetFee { value, expo })
         },
+        GovernanceAction::SetFeeInToken => {
+            let value = reader.read_u64();
+            let expo = reader.read_u64();
+            let token_len = reader.read_u8();
+            if token_len != 32 {
+                panic_with_felt252(GovernanceActionError::InvalidGovernanceMessage.into());
+            }
+            let token: felt252 = reader
+                .read_u256()
+                .try_into()
+                .expect(GovernanceActionError::InvalidGovernanceMessage.into());
+            let token = token
+                .try_into()
+                .expect(GovernanceActionError::InvalidGovernanceMessage.into());
+            GovernancePayload::SetFeeInToken(SetFeeInToken { value, expo, token })
+        },
         GovernanceAction::SetValidPeriod => { panic_with_felt252('unimplemented') },
         GovernanceAction::SetWormholeAddress => {
             let address: felt252 = reader
diff --git a/target_chains/starknet/contracts/tests/data.cairo b/target_chains/starknet/contracts/tests/data.cairo
@@ -382,6 +382,19 @@ pub fn pyth_set_fee() -> ByteBuffer {
     ByteBufferImpl::new(bytes, 23)
 }
 
+// A Pyth governance instruction to set fee signed by the test guardian #1.
+pub fn pyth_set_fee_in_token() -> ByteBuffer {
+    let bytes = array![
+        1766847064779994694408617232155063622446317599437785683244896979152308796,
+        41831183904504604246915376354509245030219494606222324288494126672855141875,
+        245200731728170526984869527586075617087934630006881191137945784647849869312,
+        49565958604199796163020368,
+        148907253456468655193350049927026865683852796092680336764850032905682767430,
+        1535109346439504966152199052711447625482878604913825938427335,
+    ];
+    ByteBufferImpl::new(bytes, 25)
+}
+
 // A Pyth governance instruction to set data sources signed by the test guardian #1.
 pub fn pyth_set_data_sources() -> ByteBuffer {
     let bytes = array![
diff --git a/target_chains/starknet/contracts/tests/pyth.cairo b/target_chains/starknet/contracts/tests/pyth.cairo
@@ -52,7 +52,9 @@ fn decode_event(mut event: Event) -> PythEvent {
         };
         PythEvent::PriceFeedUpdated(event)
     } else if key0 == event_name_hash('FeeSet') {
-        let event = FeeSet { old_fee: event.data.pop_u256(), new_fee: event.data.pop_u256(), };
+        let event = FeeSet {
+            old_fee: event.data.pop_u256(), new_fee: event.data.pop_u256(), token: event.data.pop(),
+        };
         PythEvent::FeeSet(event)
     } else if key0 == event_name_hash('DataSourcesSet') {
         let event = DataSourcesSet {
@@ -692,11 +694,16 @@ fn test_governance_set_fee_works() {
     let (from, event) = spy.events.pop_front().unwrap();
     assert!(from == pyth.contract_address);
     let event = decode_event(event);
-    let expected = FeeSet { old_fee: 1000, new_fee: 4200, };
+    let expected = FeeSet {
+        old_fee: 1000, new_fee: 4200, token: ctx.fee_contract.contract_address
+    };
     assert!(event == PythEvent::FeeSet(expected));
 
     let fee2 = pyth.get_update_fee(data::test_price_update2(), ctx.fee_contract.contract_address);
     assert!(fee2 == 4200);
+    let fee2_alt = pyth
+        .get_update_fee(data::test_price_update2(), ctx.fee_contract2.contract_address);
+    assert!(fee2_alt == 2000);
 
     start_prank(CheatTarget::One(pyth.contract_address), user);
     pyth.update_price_feeds(data::test_price_update2());
@@ -709,6 +716,62 @@ fn test_governance_set_fee_works() {
     assert!(last_price.price == 6281522520745);
 }
 
+#[test]
+fn test_governance_set_fee_in_token_works() {
+    let ctx = deploy_test();
+    let pyth = ctx.pyth;
+    let fee_contract = ctx.fee_contract;
+    let user = ctx.user;
+
+    let fee1 = pyth.get_update_fee(data::test_price_update1(), ctx.fee_contract.contract_address);
+    assert!(fee1 == 1000);
+    ctx.approve_fee(1000);
+
+    let mut balance = fee_contract.balanceOf(user);
+    start_prank(CheatTarget::One(pyth.contract_address), user);
+    pyth.update_price_feeds(data::test_price_update1());
+    stop_prank(CheatTarget::One(pyth.contract_address));
+    let new_balance = fee_contract.balanceOf(user);
+    assert!(balance - new_balance == 1000);
+    balance = new_balance;
+    let last_price = pyth
+        .get_price_unsafe(0xe62df6c8b4a85fe1a67db44dc12de5db330f7ac66b72dc658afedf0f4a415b43)
+        .unwrap_with_felt252();
+    assert!(last_price.price == 6281060000000);
+
+    let mut spy = spy_events(SpyOn::One(pyth.contract_address));
+
+    pyth.execute_governance_instruction(data::pyth_set_fee_in_token());
+
+    spy.fetch_events();
+    assert!(spy.events.len() == 1);
+    let (from, event) = spy.events.pop_front().unwrap();
+    assert!(from == pyth.contract_address);
+    let event = decode_event(event);
+    let expected = FeeSet {
+        old_fee: 2000, new_fee: 4200, token: ctx.fee_contract2.contract_address
+    };
+    assert!(event == PythEvent::FeeSet(expected));
+
+    let fee2 = pyth.get_update_fee(data::test_price_update2(), ctx.fee_contract.contract_address);
+    assert!(fee2 == 1000);
+    let fee2_alt = pyth
+        .get_update_fee(data::test_price_update2(), ctx.fee_contract2.contract_address);
+    assert!(fee2_alt == 4200);
+    ctx.approve_fee2(4200);
+
+    let balance2 = ctx.fee_contract2.balanceOf(user);
+    start_prank(CheatTarget::One(pyth.contract_address), user);
+    pyth.update_price_feeds(data::test_price_update2());
+    stop_prank(CheatTarget::One(pyth.contract_address));
+    let new_balance2 = ctx.fee_contract2.balanceOf(user);
+    assert!(balance2 - new_balance2 == 4200);
+    let last_price = pyth
+        .get_price_unsafe(0xe62df6c8b4a85fe1a67db44dc12de5db330f7ac66b72dc658afedf0f4a415b43)
+        .unwrap_with_felt252();
+    assert!(last_price.price == 6281522520745);
+}
+
 #[test]
 #[fuzzer(runs: 100, seed: 0)]
 #[should_panic]
@@ -806,8 +869,8 @@ fn test_governance_set_wormhole_works() {
 
     let user = 'user'.try_into().unwrap();
     let fee_class = declare("ERC20");
-    let fee_contract = deploy_fee_contract(fee_class, user);
-    let fee_contract2 = deploy_fee_contract(fee_class, user);
+    let fee_contract = deploy_fee_contract(fee_class, fee_address1(), user);
+    let fee_contract2 = deploy_fee_contract(fee_class, fee_address2(), user);
     let pyth = deploy_pyth_default(
         wormhole.contract_address, fee_contract.contract_address, fee_contract2.contract_address
     );
@@ -892,8 +955,8 @@ fn test_rejects_set_wormhole_without_deploying() {
 
     let user = 'user'.try_into().unwrap();
     let fee_class = declare("ERC20");
-    let fee_contract = deploy_fee_contract(fee_class, user);
-    let fee_contract2 = deploy_fee_contract(fee_class, user);
+    let fee_contract = deploy_fee_contract(fee_class, fee_address1(), user);
+    let fee_contract2 = deploy_fee_contract(fee_class, fee_address2(), user);
     let pyth = deploy_pyth_default(
         wormhole.contract_address, fee_contract.contract_address, fee_contract2.contract_address
     );
@@ -912,8 +975,8 @@ fn test_rejects_set_wormhole_with_incompatible_guardians() {
 
     let user = 'user'.try_into().unwrap();
     let fee_class = declare("ERC20");
-    let fee_contract = deploy_fee_contract(fee_class, user);
-    let fee_contract2 = deploy_fee_contract(fee_class, user);
+    let fee_contract = deploy_fee_contract(fee_class, fee_address1(), user);
+    let fee_contract2 = deploy_fee_contract(fee_class, fee_address2(), user);
     let pyth = deploy_pyth_default(
         wormhole.contract_address, fee_contract.contract_address, fee_contract2.contract_address
     );
@@ -1058,8 +1121,8 @@ fn deploy_test() -> Context {
     let user = 'user'.try_into().unwrap();
     let wormhole = super::wormhole::deploy_with_test_guardian();
     let fee_class = declare("ERC20");
-    let fee_contract = deploy_fee_contract(fee_class, user);
-    let fee_contract2 = deploy_fee_contract(fee_class, user);
+    let fee_contract = deploy_fee_contract(fee_class, fee_address1(), user);
+    let fee_contract2 = deploy_fee_contract(fee_class, fee_address2(), user);
     let pyth = deploy_pyth_default(
         wormhole.contract_address, fee_contract.contract_address, fee_contract2.contract_address
     );
@@ -1070,8 +1133,8 @@ fn deploy_mainnet() -> Context {
     let user = 'user'.try_into().unwrap();
     let wormhole = super::wormhole::deploy_with_mainnet_guardians();
     let fee_class = declare("ERC20");
-    let fee_contract = deploy_fee_contract(fee_class, user);
-    let fee_contract2 = deploy_fee_contract(fee_class, user);
+    let fee_contract = deploy_fee_contract(fee_class, fee_address1(), user);
+    let fee_contract2 = deploy_fee_contract(fee_class, fee_address2(), user);
     let pyth = deploy_pyth_default(
         wormhole.contract_address, fee_contract.contract_address, fee_contract2.contract_address
     );
@@ -1125,12 +1188,21 @@ fn deploy_pyth(
     IPythDispatcher { contract_address }
 }
 
-fn deploy_fee_contract(class: ContractClass, recipient: ContractAddress) -> IERC20CamelDispatcher {
+fn fee_address1() -> ContractAddress {
+    0x1010.try_into().unwrap()
+}
+fn fee_address2() -> ContractAddress {
+    0x049d36570d4e46f48e99674bd3fcc84644ddd6b96f7c741b1562b82f9e004dc7.try_into().unwrap()
+}
+
+fn deploy_fee_contract(
+    class: ContractClass, at: ContractAddress, recipient: ContractAddress
+) -> IERC20CamelDispatcher {
     let mut args = array![];
     let name: ByteArray = "eth";
     let symbol: ByteArray = "eth";
     (name, symbol, 100000_u256, recipient).serialize(ref args);
-    let contract_address = match class.deploy(@args) {
+    let contract_address = match class.deploy_at(@args, at) {
         Result::Ok(v) => { v },
         Result::Err(err) => { panic(err.panic_data) },
     };
diff --git a/target_chains/starknet/tools/test_vaas/src/bin/generate_test_data.rs b/target_chains/starknet/tools/test_vaas/src/bin/generate_test_data.rs
@@ -200,6 +200,29 @@ fn main() {
         "A Pyth governance instruction to set fee signed by the test guardian #1.",
     );
 
+    let pyth_set_fee_in_token_payload = vec![
+        80, 84, 71, 77, 1, 7, 234, 147, 0, 0, 0, 0, 0, 0, 0, 42, 0, 0, 0, 0, 0, 0, 0, 2, 32, 4,
+        157, 54, 87, 13, 78, 70, 244, 142, 153, 103, 75, 211, 252, 200, 70, 68, 221, 214, 185, 111,
+        124, 116, 27, 21, 98, 184, 47, 158, 0, 77, 199,
+    ];
+    let pyth_set_fee_in_token = serialize_vaa(guardians.sign_vaa(
+        &[0],
+        VaaBody {
+            timestamp: 1,
+            nonce: 2,
+            emitter_chain: 1,
+            emitter_address: u256_to_be(41.into()).into(),
+            sequence: 1.try_into().unwrap(),
+            consistency_level: 6,
+            payload: PayloadKind::Binary(pyth_set_fee_in_token_payload.clone()),
+        },
+    ));
+    print_as_cairo_fn(
+        &pyth_set_fee_in_token,
+        "pyth_set_fee_in_token",
+        "A Pyth governance instruction to set fee signed by the test guardian #1.",
+    );
+
     let pyth_set_data_sources = serialize_vaa(guardians.sign_vaa(
         &[0],
         VaaBody {
