From 40f467024d2c7d00a9ca6783c0a1a64c1d307623 Mon Sep 17 00:00:00 2001
From: Christian Kamm <mail@ckamm.de>
Date: Tue, 23 Aug 2022 14:01:10 +0200
Subject: [PATCH] FlashLoan: Add pre-health check

---
 .../mango-v4/src/instructions/flash_loan.rs     | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/programs/mango-v4/src/instructions/flash_loan.rs b/programs/mango-v4/src/instructions/flash_loan.rs
index 2616ff5a0..b56b43afa 100644
--- a/programs/mango-v4/src/instructions/flash_loan.rs
+++ b/programs/mango-v4/src/instructions/flash_loan.rs
@@ -316,7 +316,15 @@ pub fn flash_loan_end<'key, 'accounts, 'remaining, 'info>(
 
     // Check health before balance adjustments
     let retriever = new_fixed_order_account_retriever(health_ais, &account.borrow())?;
-    let _pre_health = compute_health(&account.borrow(), HealthType::Init, &retriever)?;
+    let pre_health = compute_health(&account.borrow(), HealthType::Init, &retriever)?;
+    msg!("pre_health {:?}", pre_health);
+    account
+        .fixed
+        .maybe_recover_from_being_liquidated(pre_health);
+    require!(
+        !account.fixed.being_liquidated(),
+        MangoError::BeingLiquidated
+    );
 
     // Prices for logging
     let mut prices = vec![];
@@ -390,8 +398,11 @@ pub fn flash_loan_end<'key, 'accounts, 'remaining, 'info>(
     // Check health after account position changes
     let post_health =
         compute_health_from_fixed_accounts(&account.borrow(), HealthType::Init, health_ais)?;
-    msg!("post_cpi_health {:?}", post_health);
-    require!(post_health >= 0, MangoError::HealthMustBePositive);
+    msg!("post_health {:?}", post_health);
+    require!(
+        post_health >= 0 || post_health > pre_health,
+        MangoError::HealthMustBePositive
+    );
     account
         .fixed
         .maybe_recover_from_being_liquidated(post_health);