Ayush Suresh 4 months ago
parent
commit
2cf7ef88c8

+ 1 - 0
target_chains/stylus/contracts/pyth-receiver/src/governance_structs.rs

@@ -77,6 +77,7 @@ pub struct SetTransactionFee {
 pub struct WithdrawFee {
     pub value: u64,
     pub expo: u64,
+    pub target_address: Address,
 }
 
 #[derive(Clone, Debug, PartialEq)]

+ 28 - 2
target_chains/stylus/contracts/pyth-receiver/src/lib.rs

@@ -559,8 +559,12 @@ impl PythReceiver {
             GovernancePayload::SetValidPeriod(payload) => {
                 self.set_valid_period(payload.valid_time_period_seconds);
             }
-            GovernancePayload::SetTransactionFee(_payload) => todo!(),
-            GovernancePayload::WithdrawFee(_payload) => todo!(),
+            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(())
@@ -709,6 +713,28 @@ impl PythReceiver {
 
         Ok(())
     }
+
+    fn set_transaction_fee(&mut self, value: u64, expo: u64) {
+        let new_fee = U256::from(value) * U256::from(10).pow(U256::from(expo));
+        let _old_fee = self.transaction_fee_in_wei.get();
+
+        self.transaction_fee_in_wei.set(new_fee);
+    }
+
+    fn withdraw_fee(&mut self, value: u64, expo: u64, target_address: Address) -> Result<(), PythReceiverError> {
+        let fee_to_withdraw = U256::from(value) * 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)?;
+
+        Ok(())
+    }
 }
 
 fn verify_governance_vm(receiver: &mut PythReceiver, vm: Vaa) -> Result<(), PythReceiverError> {