refactor: split governance instructions from lib.rs into governance.rs

- Extract governance-related functionality into separate governance.rs module
- Keep all pyth operations in lib.rs as requested
- Maintain idiomatic Rust module organization
- All code compiles successfully with only minor warnings

Co-Authored-By: ayush.suresh@dourolabs.xyz <byteSlayer31037@gmail.com>

target_chains/stylus/contracts/pyth-receiver/src/governance.rs

@@ -0,0 +1,318 @@
+use alloc::vec::Vec;
+use stylus_sdk::{
+    alloy_primitives::{Address, FixedBytes, U16, U256, U32, U64},
+    call::Call,
+    prelude::*,
+};
+
+use crate::{
+    error::PythReceiverError,
+    governance_structs::{self, *},
+    structs::DataSource,
+    IWormholeContract, PythReceiver,
+};
+use wormhole_vaas::{Readable, Vaa, Writeable};
+
+impl PythReceiver {
+    pub fn execute_governance_instruction_internal(
+        &mut self,
+        data: Vec<u8>,
+    ) -> Result<(), PythReceiverError> {
+        let wormhole: IWormholeContract = IWormholeContract::new(self.wormhole.get());
+        let config = Call::new();
+
+        wormhole
+            .parse_and_verify_vm(config, Vec::from(data.clone()))
+            .map_err(|_| PythReceiverError::InvalidWormholeMessage)?;
+
+        let vm = Vaa::read(&mut Vec::from(data.clone()).as_slice())
+            .map_err(|_| PythReceiverError::InvalidVaa)?;
+
+        verify_governance_vm(self, vm.clone())?;
+
+        let instruction = governance_structs::parse_instruction(vm.body.payload.to_vec())
+            .map_err(|_| PythReceiverError::InvalidGovernanceMessage)?;
+
+        let chain_id_config = Call::new();
+
+        let wormhole_id = wormhole
+            .chain_id(chain_id_config)
+            .map_err(|_| PythReceiverError::WormholeUninitialized)?;
+
+        if instruction.target_chain_id != 0 && instruction.target_chain_id != wormhole_id {
+            return Err(PythReceiverError::InvalidGovernanceTarget);
+        }
+
+        match instruction.payload {
+            GovernancePayload::SetFee(payload) => {
+                self.set_fee(payload.value, payload.expo);
+            }
+            GovernancePayload::SetDataSources(payload) => {
+                set_data_sources(self, payload.sources);
+            }
+            GovernancePayload::SetWormholeAddress(payload) => {
+                self.set_wormhole_address(payload.address, data.clone())?;
+            }
+            GovernancePayload::RequestGovernanceDataSourceTransfer(_) => {
+                return Err(PythReceiverError::InvalidGovernanceMessage);
+            }
+            GovernancePayload::AuthorizeGovernanceDataSourceTransfer(payload) => {
+                self.authorize_governance_transfer(payload.claim_vaa)?;
+            }
+            GovernancePayload::UpgradeContract(_payload) => {}
+            GovernancePayload::SetValidPeriod(payload) => {
+                self.set_valid_period(payload.valid_time_period_seconds);
+            }
+            GovernancePayload::SetTransactionFee(payload) => {
+                self.set_transaction_fee(payload.value, payload.expo);
+            }
+            GovernancePayload::WithdrawFee(payload) => {
+                self.withdraw_fee(payload.value, payload.expo, payload.target_address)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn upgrade_contract(&self, _new_implementation: FixedBytes<32>) {
+        unimplemented!("Upgrade contract not yet implemented");
+    }
+
+    fn set_fee(&mut self, value: u64, expo: u64) {
+        let new_fee = U256::from(value).saturating_mul(U256::from(10).pow(U256::from(expo)));
+        let old_fee = self.single_update_fee_in_wei.get();
+
+        self.single_update_fee_in_wei.set(new_fee);
+
+        log(self.vm(), crate::FeeSet { old_fee, new_fee });
+    }
+
+    fn set_valid_period(&mut self, valid_time_period_seconds: u64) {
+        let old_valid_period = self.valid_time_period_seconds.get();
+        let new_valid_period = U256::from(valid_time_period_seconds);
+        self.valid_time_period_seconds.set(new_valid_period);
+
+        log(
+            self.vm(),
+            crate::ValidPeriodSet {
+                old_valid_period,
+                new_valid_period,
+            },
+        );
+    }
+
+    fn set_wormhole_address(
+        &mut self,
+        address: Address,
+        data: Vec<u8>,
+    ) -> Result<(), PythReceiverError> {
+        let wormhole: IWormholeContract = IWormholeContract::new(address);
+        let config = Call::new();
+
+        wormhole
+            .parse_and_verify_vm(config, data.clone())
+            .map_err(|_| PythReceiverError::InvalidVaa)?;
+
+        let vm = Vaa::read(&mut data.as_slice())
+            .map_err(|_| PythReceiverError::VaaVerificationFailed)?;
+
+        if vm.body.emitter_chain != self.governance_data_source_chain_id.get().to::<u16>() {
+            return Err(PythReceiverError::InvalidGovernanceMessage);
+        }
+
+        if vm.body.emitter_address.as_slice()
+            != self.governance_data_source_emitter_address.get().as_slice()
+        {
+            return Err(PythReceiverError::InvalidGovernanceMessage);
+        }
+
+        if vm.body.sequence.to::<u64>() <= self.last_executed_governance_sequence.get().to::<u64>()
+        {
+            return Err(PythReceiverError::InvalidWormholeAddressToSet);
+        }
+
+        let data = governance_structs::parse_instruction(vm.body.payload.to_vec())
+            .map_err(|_| PythReceiverError::InvalidGovernanceMessage)?;
+
+        match data.payload {
+            GovernancePayload::SetWormholeAddress(payload) => {
+                if payload.address != address {
+                    return Err(PythReceiverError::InvalidWormholeAddressToSet);
+                }
+            }
+            _ => return Err(PythReceiverError::InvalidGovernanceMessage),
+        }
+
+        self.wormhole.set(address);
+        Ok(())
+    }
+
+    fn authorize_governance_transfer(
+        &mut self,
+        claim_vaa: Vec<u8>,
+    ) -> Result<(), PythReceiverError> {
+        let wormhole: IWormholeContract = IWormholeContract::new(self.wormhole.get());
+        let config = Call::new();
+        wormhole
+            .parse_and_verify_vm(config, claim_vaa.clone())
+            .map_err(|_| PythReceiverError::InvalidWormholeMessage)?;
+
+        let claim_vm = Vaa::read(&mut Vec::from(claim_vaa).as_slice())
+            .map_err(|_| PythReceiverError::VaaVerificationFailed)?;
+
+        let instruction = governance_structs::parse_instruction(claim_vm.body.payload.to_vec())
+            .map_err(|_| PythReceiverError::InvalidGovernanceMessage)?;
+
+        let config2 = Call::new();
+        if instruction.target_chain_id != 0
+            && instruction.target_chain_id != wormhole.chain_id(config2).unwrap_or(0)
+        {
+            return Err(PythReceiverError::InvalidGovernanceTarget);
+        }
+
+        let request_payload = match instruction.payload {
+            GovernancePayload::RequestGovernanceDataSourceTransfer(payload) => payload,
+            _ => return Err(PythReceiverError::InvalidGovernanceMessage),
+        };
+
+        let current_index = self.governance_data_source_index.get().to::<u32>();
+        let new_index = request_payload.governance_data_source_index;
+
+        if current_index >= new_index {
+            return Err(PythReceiverError::OldGovernanceMessage);
+        }
+
+        self.governance_data_source_index.set(U32::from(new_index));
+        let old_data_source_emitter_address = self.governance_data_source_emitter_address.get();
+
+        self.governance_data_source_chain_id
+            .set(U16::from(claim_vm.body.emitter_chain));
+        let emitter_bytes: [u8; 32] = claim_vm
+            .body
+            .emitter_address
+            .as_slice()
+            .try_into()
+            .map_err(|_| PythReceiverError::InvalidEmitterAddress)?;
+        self.governance_data_source_emitter_address
+            .set(FixedBytes::from(emitter_bytes));
+
+        let last_executed_governance_sequence = claim_vm.body.sequence.to::<u64>();
+        self.last_executed_governance_sequence
+            .set(U64::from(last_executed_governance_sequence));
+
+        log(
+            self.vm(),
+            crate::GovernanceDataSourceSet {
+                old_chain_id: current_index as u16,
+                old_emitter_address: old_data_source_emitter_address,
+                new_chain_id: claim_vm.body.emitter_chain,
+                new_emitter_address: FixedBytes::from(emitter_bytes),
+                initial_sequence: last_executed_governance_sequence,
+            },
+        );
+
+        Ok(())
+    }
+
+    fn set_transaction_fee(&mut self, value: u64, expo: u64) {
+        let new_fee = U256::from(value).saturating_mul(U256::from(10).pow(U256::from(expo)));
+        let old_fee = self.transaction_fee_in_wei.get();
+
+        self.transaction_fee_in_wei.set(new_fee);
+
+        log(self.vm(), crate::TransactionFeeSet { old_fee, new_fee });
+    }
+
+    fn withdraw_fee(
+        &mut self,
+        value: u64,
+        expo: u64,
+        target_address: Address,
+    ) -> Result<(), PythReceiverError> {
+        let fee_to_withdraw =
+            U256::from(value).saturating_mul(U256::from(10).pow(U256::from(expo)));
+        let current_balance = self.vm().balance(self.vm().contract_address());
+
+        if current_balance < fee_to_withdraw {
+            return Err(PythReceiverError::InsufficientFee);
+        }
+
+        self.vm()
+            .transfer_eth(target_address, fee_to_withdraw)
+            .map_err(|_| PythReceiverError::InsufficientFee)?;
+
+        log(
+            self.vm(),
+            crate::FeeWithdrawn {
+                target_address,
+                fee_amount: fee_to_withdraw,
+            },
+        );
+
+        Ok(())
+    }
+}
+
+pub fn verify_governance_vm(receiver: &mut PythReceiver, vm: Vaa) -> Result<(), PythReceiverError> {
+    if vm.body.emitter_chain != receiver.governance_data_source_chain_id.get().to::<u16>() {
+        return Err(PythReceiverError::InvalidGovernanceMessage);
+    }
+
+    if vm.body.emitter_address.as_slice()
+        != receiver
+            .governance_data_source_emitter_address
+            .get()
+            .as_slice()
+    {
+        return Err(PythReceiverError::InvalidGovernanceMessage);
+    }
+
+    let current_sequence = vm.body.sequence.to::<u64>();
+    let last_executed_sequence = receiver.last_executed_governance_sequence.get().to::<u64>();
+
+    if current_sequence <= last_executed_sequence {
+        return Err(PythReceiverError::GovernanceMessageAlreadyExecuted);
+    }
+
+    receiver
+        .last_executed_governance_sequence
+        .set(U64::from(current_sequence));
+
+    Ok(())
+}
+
+pub fn set_data_sources(receiver: &mut PythReceiver, data_sources: Vec<DataSource>) {
+    let mut old_data_sources = Vec::new();
+    for i in 0..receiver.valid_data_sources.len() {
+        if let Some(storage_data_source) = receiver.valid_data_sources.get(i) {
+            let data_source = DataSource {
+                chain_id: storage_data_source.chain_id.get(),
+                emitter_address: storage_data_source.emitter_address.get(),
+            };
+            old_data_sources.push(data_source.emitter_address);
+            receiver.is_valid_data_source.setter(data_source).set(false);
+        }
+    }
+
+    receiver.valid_data_sources.erase();
+
+    let mut new_data_sources = Vec::new();
+    for data_source in data_sources {
+        let mut storage_data_source = receiver.valid_data_sources.grow();
+        storage_data_source.chain_id.set(data_source.chain_id);
+        storage_data_source
+            .emitter_address
+            .set(data_source.emitter_address);
+
+        new_data_sources.push(data_source.emitter_address);
+        receiver.is_valid_data_source.setter(data_source).set(true);
+    }
+
+    log(
+        receiver.vm(),
+        crate::DataSourcesSet {
+            old_data_sources,
+            new_data_sources,
+        },
+    );
+}

target_chains/stylus/contracts/pyth-receiver/src/lib.rs

@@ -6,6 +6,7 @@
 extern crate alloc;
 
 mod error;
+mod governance;
 mod governance_structs;
 #[cfg(test)]
 mod integration_tests;
@@ -391,7 +392,6 @@ impl PythReceiver {
         check_uniqueness: bool,
     ) -> Result<Vec<PriceFeedReturn>, PythReceiverError> {
         let update_data_array: &[u8] = &update_data;
-        // Check the first 4 bytes of the update_data_array for the magic header
         if update_data_array.len() < 4 {
             return Err(PythReceiverError::InvalidUpdateData);
         }
@@ -535,64 +535,7 @@ impl PythReceiver {
         &mut self,
         data: Vec<u8>,
     ) -> Result<(), PythReceiverError> {
-        let wormhole: IWormholeContract = IWormholeContract::new(self.wormhole.get());
-        let config = Call::new();
-
-        wormhole
-            .parse_and_verify_vm(config, Vec::from(data.clone()))
-            .map_err(|_| PythReceiverError::InvalidWormholeMessage)?;
-
-        let vm = Vaa::read(&mut Vec::from(data.clone()).as_slice())
-            .map_err(|_| PythReceiverError::InvalidVaa)?;
-
-        verify_governance_vm(self, vm.clone())?;
-
-        let instruction = governance_structs::parse_instruction(vm.body.payload.to_vec())
-            .map_err(|_| PythReceiverError::InvalidGovernanceMessage)?;
-
-        let chain_id_config = Call::new();
-
-        let wormhole_id = wormhole
-            .chain_id(chain_id_config)
-            .map_err(|_| PythReceiverError::WormholeUninitialized)?;
-
-        if instruction.target_chain_id != 0 && instruction.target_chain_id != wormhole_id {
-            return Err(PythReceiverError::InvalidGovernanceTarget);
-        }
-
-        match instruction.payload {
-            GovernancePayload::SetFee(payload) => {
-                self.set_fee(payload.value, payload.expo);
-            }
-            GovernancePayload::SetDataSources(payload) => {
-                set_data_sources(self, payload.sources);
-            }
-            GovernancePayload::SetWormholeAddress(payload) => {
-                self.set_wormhole_address(payload.address, data.clone())?;
-            }
-            GovernancePayload::RequestGovernanceDataSourceTransfer(_) => {
-                return Err(PythReceiverError::InvalidGovernanceMessage);
-            }
-            GovernancePayload::AuthorizeGovernanceDataSourceTransfer(payload) => {
-                self.authorize_governance_transfer(payload.claim_vaa)?;
-            }
-            GovernancePayload::UpgradeContract(_payload) => {}
-            GovernancePayload::SetValidPeriod(payload) => {
-                self.set_valid_period(payload.valid_time_period_seconds);
-            }
-            GovernancePayload::SetTransactionFee(payload) => {
-                self.set_transaction_fee(payload.value, payload.expo);
-            }
-            GovernancePayload::WithdrawFee(payload) => {
-                self.withdraw_fee(payload.value, payload.expo, payload.target_address)?;
-            }
-        }
-
-        Ok(())
-    }
-
-    fn upgrade_contract(&self, _new_implementation: FixedBytes<32>) {
-        unimplemented!("Upgrade contract not yet implemented");
+        self.execute_governance_instruction_internal(data)
     }
 
     fn is_no_older_than(&self, publish_time: U64, max_age: u64) -> bool {
@@ -601,8 +544,6 @@ impl PythReceiver {
             <= max_age
     }
 
-    // Stylus doesn't provide a way to mock up the testing timestamp
-    // so at the moment I'm using the testing trait to let me test old timestamps
     fn get_current_timestamp(&self) -> u64 {
         #[cfg(test)]
         {
@@ -613,209 +554,6 @@ impl PythReceiver {
             self.vm().block_timestamp()
         }
     }
-
-    fn set_fee(&mut self, value: u64, expo: u64) {
-        let new_fee = U256::from(value).saturating_mul(U256::from(10).pow(U256::from(expo)));
-        let old_fee = self.single_update_fee_in_wei.get();
-
-        self.single_update_fee_in_wei.set(new_fee);
-
-        log(self.vm(), FeeSet { old_fee, new_fee });
-    }
-
-    fn set_valid_period(&mut self, valid_time_period_seconds: u64) {
-        let old_valid_period = self.valid_time_period_seconds.get();
-        let new_valid_period = U256::from(valid_time_period_seconds);
-        self.valid_time_period_seconds.set(new_valid_period);
-
-        log(
-            self.vm(),
-            ValidPeriodSet {
-                old_valid_period,
-                new_valid_period,
-            },
-        );
-    }
-
-    fn set_wormhole_address(
-        &mut self,
-        address: Address,
-        data: Vec<u8>,
-    ) -> Result<(), PythReceiverError> {
-        let wormhole: IWormholeContract = IWormholeContract::new(address);
-        let config = Call::new();
-
-        // Make sure the contract at the new Wormhole address is initialized and functional by testing it on the SetWormholeAddress instruction VAA
-        wormhole
-            .parse_and_verify_vm(config, data.clone())
-            .map_err(|_| PythReceiverError::InvalidVaa)?;
-
-        let vm = Vaa::read(&mut data.as_slice())
-            .map_err(|_| PythReceiverError::VaaVerificationFailed)?;
-
-        if vm.body.emitter_chain != self.governance_data_source_chain_id.get().to::<u16>() {
-            return Err(PythReceiverError::InvalidGovernanceMessage);
-        }
-
-        if vm.body.emitter_address.as_slice()
-            != self.governance_data_source_emitter_address.get().as_slice()
-        {
-            return Err(PythReceiverError::InvalidGovernanceMessage);
-        }
-
-        if vm.body.sequence.to::<u64>() <= self.last_executed_governance_sequence.get().to::<u64>()
-        {
-            return Err(PythReceiverError::InvalidWormholeAddressToSet);
-        }
-
-        let data = governance_structs::parse_instruction(vm.body.payload.to_vec())
-            .map_err(|_| PythReceiverError::InvalidGovernanceMessage)?;
-
-        match data.payload {
-            GovernancePayload::SetWormholeAddress(payload) => {
-                if payload.address != address {
-                    return Err(PythReceiverError::InvalidWormholeAddressToSet);
-                }
-            }
-            _ => return Err(PythReceiverError::InvalidGovernanceMessage),
-        }
-
-        self.wormhole.set(address);
-        Ok(())
-    }
-
-    fn authorize_governance_transfer(
-        &mut self,
-        claim_vaa: Vec<u8>,
-    ) -> Result<(), PythReceiverError> {
-        let wormhole: IWormholeContract = IWormholeContract::new(self.wormhole.get());
-        let config = Call::new();
-        wormhole
-            .parse_and_verify_vm(config, claim_vaa.clone())
-            .map_err(|_| PythReceiverError::InvalidWormholeMessage)?;
-
-        let claim_vm = Vaa::read(&mut Vec::from(claim_vaa).as_slice())
-            .map_err(|_| PythReceiverError::VaaVerificationFailed)?;
-
-        let instruction = governance_structs::parse_instruction(claim_vm.body.payload.to_vec())
-            .map_err(|_| PythReceiverError::InvalidGovernanceMessage)?;
-
-        let config2 = Call::new();
-        if instruction.target_chain_id != 0
-            && instruction.target_chain_id != wormhole.chain_id(config2).unwrap_or(0)
-        {
-            return Err(PythReceiverError::InvalidGovernanceTarget);
-        }
-
-        let request_payload = match instruction.payload {
-            GovernancePayload::RequestGovernanceDataSourceTransfer(payload) => payload,
-            _ => return Err(PythReceiverError::InvalidGovernanceMessage),
-        };
-
-        let current_index = self.governance_data_source_index.get().to::<u32>();
-        let new_index = request_payload.governance_data_source_index;
-
-        if current_index >= new_index {
-            return Err(PythReceiverError::OldGovernanceMessage);
-        }
-
-        self.governance_data_source_index.set(U32::from(new_index));
-        let old_data_source_emitter_address = self.governance_data_source_emitter_address.get();
-
-        self.governance_data_source_chain_id
-            .set(U16::from(claim_vm.body.emitter_chain));
-        let emitter_bytes: [u8; 32] = claim_vm
-            .body
-            .emitter_address
-            .as_slice()
-            .try_into()
-            .map_err(|_| PythReceiverError::InvalidEmitterAddress)?;
-        self.governance_data_source_emitter_address
-            .set(FixedBytes::from(emitter_bytes));
-
-        let last_executed_governance_sequence = claim_vm.body.sequence.to::<u64>();
-        self.last_executed_governance_sequence
-            .set(U64::from(last_executed_governance_sequence));
-
-        log(
-            self.vm(),
-            GovernanceDataSourceSet {
-                old_chain_id: current_index as u16,
-                old_emitter_address: old_data_source_emitter_address,
-                new_chain_id: claim_vm.body.emitter_chain,
-                new_emitter_address: FixedBytes::from(emitter_bytes),
-                initial_sequence: last_executed_governance_sequence,
-            },
-        );
-
-        Ok(())
-    }
-
-    fn set_transaction_fee(&mut self, value: u64, expo: u64) {
-        let new_fee = U256::from(value).saturating_mul(U256::from(10).pow(U256::from(expo)));
-        let old_fee = self.transaction_fee_in_wei.get();
-
-        self.transaction_fee_in_wei.set(new_fee);
-
-        log(self.vm(), TransactionFeeSet { old_fee, new_fee });
-    }
-
-    fn withdraw_fee(
-        &mut self,
-        value: u64,
-        expo: u64,
-        target_address: Address,
-    ) -> Result<(), PythReceiverError> {
-        let fee_to_withdraw =
-            U256::from(value).saturating_mul(U256::from(10).pow(U256::from(expo)));
-        let current_balance = self.vm().balance(self.vm().contract_address());
-
-        if current_balance < fee_to_withdraw {
-            return Err(PythReceiverError::InsufficientFee);
-        }
-
-        self.vm()
-            .transfer_eth(target_address, fee_to_withdraw)
-            .map_err(|_| PythReceiverError::InsufficientFee)?;
-
-        log(
-            self.vm(),
-            FeeWithdrawn {
-                target_address,
-                fee_amount: fee_to_withdraw,
-            },
-        );
-
-        Ok(())
-    }
-}
-
-fn verify_governance_vm(receiver: &mut PythReceiver, vm: Vaa) -> Result<(), PythReceiverError> {
-    if vm.body.emitter_chain != receiver.governance_data_source_chain_id.get().to::<u16>() {
-        return Err(PythReceiverError::InvalidGovernanceMessage);
-    }
-
-    if vm.body.emitter_address.as_slice()
-        != receiver
-            .governance_data_source_emitter_address
-            .get()
-            .as_slice()
-    {
-        return Err(PythReceiverError::InvalidGovernanceMessage);
-    }
-
-    let current_sequence = vm.body.sequence.to::<u64>();
-    let last_executed_sequence = receiver.last_executed_governance_sequence.get().to::<u64>();
-
-    if current_sequence <= last_executed_sequence {
-        return Err(PythReceiverError::GovernanceMessageAlreadyExecuted);
-    }
-
-    receiver
-        .last_executed_governance_sequence
-        .set(U64::from(current_sequence));
-
-    Ok(())
 }
 
 fn parse_wormhole_proof(vaa: Vaa) -> Result<MerkleRoot<Keccak160>, PythReceiverError> {
@@ -826,39 +564,3 @@ fn parse_wormhole_proof(vaa: Vaa) -> Result<MerkleRoot<Keccak160>, PythReceiverE
     });
     Ok(root)
 }
-
-fn set_data_sources(receiver: &mut PythReceiver, data_sources: Vec<DataSource>) {
-    let mut old_data_sources = Vec::new();
-    for i in 0..receiver.valid_data_sources.len() {
-        if let Some(storage_data_source) = receiver.valid_data_sources.get(i) {
-            let data_source = DataSource {
-                chain_id: storage_data_source.chain_id.get(),
-                emitter_address: storage_data_source.emitter_address.get(),
-            };
-            old_data_sources.push(data_source.emitter_address);
-            receiver.is_valid_data_source.setter(data_source).set(false);
-        }
-    }
-
-    receiver.valid_data_sources.erase();
-
-    let mut new_data_sources = Vec::new();
-    for data_source in data_sources {
-        let mut storage_data_source = receiver.valid_data_sources.grow();
-        storage_data_source.chain_id.set(data_source.chain_id);
-        storage_data_source
-            .emitter_address
-            .set(data_source.emitter_address);
-
-        new_data_sources.push(data_source.emitter_address);
-        receiver.is_valid_data_source.setter(data_source).set(true);
-    }
-
-    log(
-        receiver.vm(),
-        DataSourcesSet {
-            old_data_sources,
-            new_data_sources,
-        },
-    );
-}