FlashLoan: Don't deduce vault_len, add more checks (#542)

It looks like it wasn't possible to exploit the program by re-initializing the user-owned token accounts used in flash loan because the later use of health_ais with the health computation would error out if any token account was included there. However, the check and a few others were too indirect. In this patch: - We pass the number of loans into FlashLoanEnd explicitly (verified from FlashLoanBegin) - Add explicit checks for token mints, so it's no longer possible to use token accounts for foreign mints in Begin when the loan amount is zero, and it's clearer to see that the bookkeeping in End won't break if the user reinited the account for a different mint. - Also add a few other extra comments and checks. The updated FlashLoanEnd instruction is called FlashLoanEndV2
2023-04-14 15:18:02 +02:00 · 2023-04-14 15:18:02 +02:00 · 98a79505e1
parent 350d558ee3
commit 98a79505e1
7 changed files with 214 additions and 14 deletions
--- a/lib/client/src/client.rs
+++ b/lib/client/src/client.rs
@ -1293,6 +1293,7 @@ impl MangoClient {
            },
            0u64,
        ];
+        let num_loans: u8 = loan_amounts.len().try_into().unwrap();

        // This relies on the fact that health account banks will be identical to the first_bank above!
        let health_ams = self
@ -1364,7 +1365,8 @@ impl MangoClient {
                ams.push(to_readonly_account_meta(self.group()));
                ams
            },
-            data: anchor_lang::InstructionData::data(&mango_v4::instruction::FlashLoanEnd {
+            data: anchor_lang::InstructionData::data(&mango_v4::instruction::FlashLoanEndV2 {
+                num_loans,
                flash_loan_type: mango_v4::accounts_ix::FlashLoanType::Swap,
            }),
        });
--- a/mango_v4.json
+++ b/mango_v4.json
@ -1543,6 +1543,38 @@
        }
      ]
    },
+    {
+      "name": "flashLoanEndV2",
+      "accounts": [
+        {
+          "name": "account",
+          "isMut": true,
+          "isSigner": false
+        },
+        {
+          "name": "owner",
+          "isMut": false,
+          "isSigner": true
+        },
+        {
+          "name": "tokenProgram",
+          "isMut": false,
+          "isSigner": false
+        }
+      ],
+      "args": [
+        {
+          "name": "numLoans",
+          "type": "u8"
+        },
+        {
+          "name": "flashLoanType",
+          "type": {
+            "defined": "FlashLoanType"
+          }
+        }
+      ]
+    },
    {
      "name": "healthRegionBegin",
      "accounts": [
@ -7701,6 +7733,11 @@
          "type": "i128",
          "index": false
        },
+        {
+          "name": "oracleSlot",
+          "type": "u64",
+          "index": false
+        },
        {
          "name": "stablePrice",
          "type": "i128",
@ -8504,6 +8541,26 @@
          "index": false
        }
      ]
+    },
+    {
+      "name": "FilledPerpOrderLog",
+      "fields": [
+        {
+          "name": "mangoGroup",
+          "type": "publicKey",
+          "index": false
+        },
+        {
+          "name": "perpMarketIndex",
+          "type": "u16",
+          "index": false
+        },
+        {
+          "name": "seqNum",
+          "type": "u64",
+          "index": false
+        }
+      ]
    }
  ],
  "errors": [
--- a/programs/mango-v4/src/instructions/flash_loan.rs
+++ b/programs/mango-v4/src/instructions/flash_loan.rs
@ -19,6 +19,12 @@ pub fn flash_loan_begin<'key, 'accounts, 'remaining, 'info>(
    ctx: Context<'key, 'accounts, 'remaining, 'info, FlashLoanBegin<'info>>,
    loan_amounts: Vec<u64>,
 ) -> Result<()> {
+    let num_loans = loan_amounts.len();
+    require_gt!(num_loans, 0);
+
+    // Loans of 0 are acceptable and common: Users often want to loan some of token A,
+    // nothing of token B, swap A to B and then deposit the gains.
+
    let account = ctx.accounts.account.load_full_mut()?;

    // account constraint #1
@ -27,7 +33,6 @@ pub fn flash_loan_begin<'key, 'accounts, 'remaining, 'info>(
        MangoError::SomeError
    );

-    let num_loans = loan_amounts.len();
    require_eq!(ctx.remaining_accounts.len(), 3 * num_loans + 1);
    let banks = &ctx.remaining_accounts[..num_loans];
    let vaults = &ctx.remaining_accounts[num_loans..2 * num_loans];
@ -70,8 +75,13 @@ pub fn flash_loan_begin<'key, 'accounts, 'remaining, 'info>(
        let vault = Account::<TokenAccount>::try_from(vault_ai)?;
        let token_account = Account::<TokenAccount>::try_from(token_account_ai)?;

+        require_keys_eq!(token_account.mint, bank.mint);
+
+        // This check is likely unnecessary
        require_keys_neq!(token_account.owner, group_ai.key());

+        require_eq!(bank.flash_loan_approved_amount, 0);
+        require_eq!(bank.flash_loan_token_account_initial, u64::MAX);
        bank.flash_loan_approved_amount = *amount;
        bank.flash_loan_token_account_initial = token_account.amount;

@ -144,9 +154,11 @@ pub fn flash_loan_begin<'key, 'accounts, 'remaining, 'info>(

                // must be the FlashLoanEnd instruction
                require!(
-                    ix.data[0..8] == crate::instruction::FlashLoanEnd::discriminator(),
+                    ix.data[0..8] == crate::instruction::FlashLoanEndV2::discriminator(),
                    MangoError::SomeError
                );
+                // the correct number of loans is passed to the End instruction
+                require_eq!(ix.data[8] as usize, num_loans);

                require_msg!(
                    ctx.accounts.account.key() == ix.accounts[0].pubkey,
@ -186,8 +198,16 @@ struct TokenVaultChange {

 pub fn flash_loan_end<'key, 'accounts, 'remaining, 'info>(
    ctx: Context<'key, 'accounts, 'remaining, 'info, FlashLoanEnd<'info>>,
+    num_loans: u8,
    flash_loan_type: FlashLoanType,
 ) -> Result<()> {
+    require_gt!(num_loans, 0);
+
+    // FlashLoanEnd can only be called in the same tx as a FlashLoanBegin because:
+    // - FlashLoanBegin checks for a matching FlashLoanEnd in the same tx
+    // - FlashLoanBegin sets flash_loan_token_account_initial on a bank, which is
+    //   validated below. (and there must be at least one bank-vault-token account triple)
+
    let mut account = ctx.accounts.account.load_full_mut()?;

    // account constraint #1
@ -203,14 +223,7 @@ pub fn flash_loan_end<'key, 'accounts, 'remaining, 'info>(
    require_keys_eq!(group, group_ai.key());

    // Find index at which vaults start
-    let vaults_len = ctx.remaining_accounts[..remaining_len - 1]
-        .iter()
-        .rev()
-        .map_while(|ai| Account::<TokenAccount>::try_from(ai).ok())
-        .position(|token_account| token_account.owner == group)
-        .ok_or_else(|| {
-            error_msg!("expected at least one group-owned vault token account to be passed")
-        })?;
+    let vaults_len: usize = num_loans.into();
    let vaults_index = remaining_len - 2 * vaults_len - 1;

    let health_ais = &ctx.remaining_accounts[..vaults_index];
@ -246,10 +259,14 @@ pub fn flash_loan_end<'key, 'accounts, 'remaining, 'info>(
        let token_account_ai = &token_accounts[vault_index];
        let token_account = Account::<TokenAccount>::try_from(token_account_ai)?;

+        // The token account could have been re-initialized for a different mint
+        require_keys_eq!(token_account.mint, bank.mint);
+
        // Ensure this bank/vault combination was mentioned in the Begin instruction:
        // The Begin instruction only checks that End ends with the same vault accounts -
        // but there could be an extra vault account in End, or a different bank could be
        // used for the same vault.
+        // This check guarantees that FlashLoanBegin was called on this bank.
        require_neq!(bank.flash_loan_token_account_initial, u64::MAX);

        // Create the token position now, so we can compute the pre-health with fixed order health accounts
--- a/programs/mango-v4/src/lib.rs
+++ b/programs/mango-v4/src/lib.rs
@ -39,6 +39,7 @@ declare_id!("4MangoMjqJ2firMokCjjGgoK8d4MXcrgL7XJaL3w6fVg");
 #[program]
 pub mod mango_v4 {
    use super::*;
+    use error::*;

    pub fn group_create(
        ctx: Context<GroupCreate>,
@ -355,9 +356,17 @@ pub mod mango_v4 {
    pub fn flash_loan_end<'key, 'accounts, 'remaining, 'info>(
        ctx: Context<'key, 'accounts, 'remaining, 'info, FlashLoanEnd<'info>>,
        flash_loan_type: FlashLoanType,
+    ) -> Result<()> {
+        Err(error_msg!("FlashLoanEnd was replaced by FlashLoanEndV2"))
+    }
+
+    pub fn flash_loan_end_v2<'key, 'accounts, 'remaining, 'info>(
+        ctx: Context<'key, 'accounts, 'remaining, 'info, FlashLoanEnd<'info>>,
+        num_loans: u8,
+        flash_loan_type: FlashLoanType,
    ) -> Result<()> {
        #[cfg(feature = "enable-gpl")]
-        instructions::flash_loan_end(ctx, flash_loan_type)?;
+        instructions::flash_loan_end(ctx, num_loans, flash_loan_type)?;
        Ok(())
    }

--- a/programs/mango-v4/tests/program_test/mango_client.rs
+++ b/programs/mango-v4/tests/program_test/mango_client.rs
@ -499,13 +499,14 @@ pub struct FlashLoanEndInstruction {
 #[async_trait::async_trait(?Send)]
 impl ClientInstruction for FlashLoanEndInstruction {
    type Accounts = mango_v4::accounts::FlashLoanEnd;
-    type Instruction = mango_v4::instruction::FlashLoanEnd;
+    type Instruction = mango_v4::instruction::FlashLoanEndV2;
    async fn to_instruction(
        &self,
        account_loader: impl ClientAccountLoader + 'async_trait,
    ) -> (Self::Accounts, instruction::Instruction) {
        let program_id = mango_v4::id();
        let instruction = Self::Instruction {
+            num_loans: 1,
            flash_loan_type: self.flash_loan_type,
        };

--- a/ts/client/src/client.ts
+++ b/ts/client/src/client.ts
@ -2680,7 +2680,7 @@ export class MangoClient {
    };

    const flashLoanEndIx = await this.program.methods
-      .flashLoanEnd(flashLoanType)
+      .flashLoanEndV2(2, flashLoanType)
      .accounts({
        account: mangoAccount.publicKey,
      })
--- a/ts/client/src/mango_v4.ts
+++ b/ts/client/src/mango_v4.ts
@ -1543,6 +1543,38 @@ export type MangoV4 = {
        }
      ]
    },
+    {
+      "name": "flashLoanEndV2",
+      "accounts": [
+        {
+          "name": "account",
+          "isMut": true,
+          "isSigner": false
+        },
+        {
+          "name": "owner",
+          "isMut": false,
+          "isSigner": true
+        },
+        {
+          "name": "tokenProgram",
+          "isMut": false,
+          "isSigner": false
+        }
+      ],
+      "args": [
+        {
+          "name": "numLoans",
+          "type": "u8"
+        },
+        {
+          "name": "flashLoanType",
+          "type": {
+            "defined": "FlashLoanType"
+          }
+        }
+      ]
+    },
    {
      "name": "healthRegionBegin",
      "accounts": [
@ -7701,6 +7733,11 @@ export type MangoV4 = {
          "type": "i128",
          "index": false
        },
+        {
+          "name": "oracleSlot",
+          "type": "u64",
+          "index": false
+        },
        {
          "name": "stablePrice",
          "type": "i128",
@ -8504,6 +8541,26 @@ export type MangoV4 = {
          "index": false
        }
      ]
+    },
+    {
+      "name": "FilledPerpOrderLog",
+      "fields": [
+        {
+          "name": "mangoGroup",
+          "type": "publicKey",
+          "index": false
+        },
+        {
+          "name": "perpMarketIndex",
+          "type": "u16",
+          "index": false
+        },
+        {
+          "name": "seqNum",
+          "type": "u64",
+          "index": false
+        }
+      ]
    }
  ],
  "errors": [
@ -10290,6 +10347,38 @@ export const IDL: MangoV4 = {
        }
      ]
    },
+    {
+      "name": "flashLoanEndV2",
+      "accounts": [
+        {
+          "name": "account",
+          "isMut": true,
+          "isSigner": false
+        },
+        {
+          "name": "owner",
+          "isMut": false,
+          "isSigner": true
+        },
+        {
+          "name": "tokenProgram",
+          "isMut": false,
+          "isSigner": false
+        }
+      ],
+      "args": [
+        {
+          "name": "numLoans",
+          "type": "u8"
+        },
+        {
+          "name": "flashLoanType",
+          "type": {
+            "defined": "FlashLoanType"
+          }
+        }
+      ]
+    },
    {
      "name": "healthRegionBegin",
      "accounts": [
@ -16448,6 +16537,11 @@ export const IDL: MangoV4 = {
          "type": "i128",
          "index": false
        },
+        {
+          "name": "oracleSlot",
+          "type": "u64",
+          "index": false
+        },
        {
          "name": "stablePrice",
          "type": "i128",
@ -17251,6 +17345,26 @@ export const IDL: MangoV4 = {
          "index": false
        }
      ]
+    },
+    {
+      "name": "FilledPerpOrderLog",
+      "fields": [
+        {
+          "name": "mangoGroup",
+          "type": "publicKey",
+          "index": false
+        },
+        {
+          "name": "perpMarketIndex",
+          "type": "u16",
+          "index": false
+        },
+        {
+          "name": "seqNum",
+          "type": "u64",
+          "index": false
+        }
+      ]
    }
  ],
  "errors": [