Merge pull request #157 from Brushfam/oz/l-01

contracts/src/token/psp22/extensions/flashmint.rs

@@ -30,7 +30,10 @@ pub use crate::{
 pub use flashmint::Internal as _;
 use ink::{
     env::CallFlags,
-    prelude::vec::Vec,
+    prelude::{
+        vec,
+        vec::Vec,
+    },
 };
 use openbrush::traits::{
     AccountId,
@@ -76,7 +79,18 @@ pub trait FlashLenderImpl: Storage<psp22::Data> + psp22::Internal + PSP22 + Inte
             return Err(FlashLenderError::AllowanceDoesNotAllowRefund)
         }
         psp22::Internal::_approve_from_to(self, receiver_account, this, current_allowance - amount - fee)?;
-        psp22::Internal::_burn_from(self, receiver_account, amount + fee)?;
+
+        let flash_fee_receiver = self._flash_fee_receiver();
+
+        if let Some(fee_receiver) = flash_fee_receiver {
+            if fee == 0 {
+                psp22::Internal::_burn_from(self, receiver_account, amount + fee)?;
+            } else {
+                psp22::Internal::_burn_from(self, receiver_account, amount)?;
+                psp22::Internal::_transfer_from_to(self, receiver_account, fee_receiver, fee, vec![])?;
+            }
+        }
+
         Ok(())
     }
 }
@@ -92,6 +106,8 @@ pub trait Internal {
         amount: Balance,
         data: Vec<u8>,
     ) -> Result<(), FlashLenderError>;
+
+    fn _flash_fee_receiver(&self) -> Option<AccountId>;
 }
 
 pub trait InternalImpl: Storage<psp22::Data> + Internal {
@@ -128,4 +144,8 @@ pub trait InternalImpl: Storage<psp22::Data> + Internal {
 
         result
     }
+
+    fn _flash_fee_receiver(&self) -> Option<AccountId> {
+        None
+    }
 }

lang/codegen/src/implementations.rs

@@ -464,6 +464,10 @@ pub(crate) fn impl_flashmint(impl_args: &mut ImplArgs) {
             ) -> Result<(), FlashLenderError> {
                 flashmint::InternalImpl::_on_flashloan(self, receiver_account, token, fee, amount, data)
             }
+
+            fn _flash_fee_receiver(&self) -> Option<AccountId> {
+                flashmint::InternalImpl::_flash_fee_receiver(self)
+            }
         }
     ))
     .expect("Should parse");