Provide fixes after testing

dappforce · Apr 2, 2024 · b3d4a77 · b3d4a77
1 parent 5ab9be0
commit b3d4a77
Showing 1 changed file with 24 additions and 6 deletions.
diff --git a/pallets/energy/src/lib.rs b/pallets/energy/src/lib.rs
@@ -338,12 +338,18 @@ pub mod pallet {
         }
     }
 
+    /// Keeps track of whether transaction was paid using proxy's real account energy.
+    pub enum IsProxy<AccountId> {
+        Yes(AccountId),
+        No,
+    }
+
     /// Keeps track of how the user paid for the transaction.
     pub enum LiquidityInfo<T: Config> {
         /// Nothing have been paid.
         Nothing,
         /// Transaction have been paid using energy.
-        Energy(BalanceOf<T>),
+        Energy(BalanceOf<T>, IsProxy<T::AccountId>),
         /// Transaction have been paid using the native method.
         Native(<T::NativeOnChargeTransaction as OnChargeTransaction<T>>::LiquidityInfo),
     }
@@ -373,11 +379,12 @@ pub mod pallet {
             let energy_fee = Self::native_token_to_energy(fee_without_tip);
 
             let maybe_proxy_call = <T as Config>::RuntimeCall::from_ref(call).is_sub_type();
-
+
+            let mut is_who_a_proxy = false;
             let energy_provider = match maybe_proxy_call {
                 Some(pallet_proxy::Call::proxy { real, .. }) => {
                     let real_account = T::Lookup::lookup(real.clone())?;
-                    let is_who_a_proxy = pallet_proxy::Pallet::<T>::find_proxy(&real_account, who, None).is_ok();
+                    is_who_a_proxy = pallet_proxy::Pallet::<T>::find_proxy(&real_account, who, None).is_ok();
 
                     if is_who_a_proxy {
                         real_account
@@ -414,7 +421,14 @@ pub mod pallet {
             match Self::ensure_can_consume_energy(&energy_provider, energy_fee) {
                 Ok(()) => {
                     Self::consume_energy(&energy_provider, energy_fee);
-                    Ok(LiquidityInfo::Energy(energy_fee))
+                    Ok(LiquidityInfo::Energy(
+                        energy_fee,
+                        if is_who_a_proxy {
+                            IsProxy::Yes(energy_provider)
+                        } else {
+                            IsProxy::No
+                        },
+                    ))
                 },
                 Err(_) => Err(InvalidTransaction::Payment.into()),
             }
@@ -439,14 +453,18 @@ pub mod pallet {
                         tip,
                         fallback_info,
                     ),
-                LiquidityInfo::Energy(paid) => {
+                LiquidityInfo::Energy(paid, proxy) => {
                     let corrected_fee_without_tip = corrected_fee.saturating_sub(tip);
                     let corrected_energy_fee =
                         Self::native_token_to_energy(corrected_fee_without_tip);
 
                     let refund_amount = paid.saturating_sub(corrected_energy_fee);
+                    let refund_destination = match proxy {
+                        IsProxy::Yes(ref account) => account,
+                        IsProxy::No => who,
+                    };
 
-                    Self::capture_energy(who, refund_amount);
+                    Self::capture_energy(refund_destination, refund_amount);
 
                     Ok(())
                 },