From 98a79505e123aa27a6f2d178462e8d394b60cef5 Mon Sep 17 00:00:00 2001 From: Christian Kamm Date: Fri, 14 Apr 2023 15:18:02 +0200 Subject: [PATCH] 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 --- lib/client/src/client.rs | 4 +- mango_v4.json | 57 +++++++++ .../mango-v4/src/instructions/flash_loan.rs | 37 ++++-- programs/mango-v4/src/lib.rs | 11 +- .../tests/program_test/mango_client.rs | 3 +- ts/client/src/client.ts | 2 +- ts/client/src/mango_v4.ts | 114 ++++++++++++++++++ 7 files changed, 214 insertions(+), 14 deletions(-) diff --git a/lib/client/src/client.rs b/lib/client/src/client.rs index 3bbac7684..2a15e5f38 100644 --- 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, }), }); diff --git a/mango_v4.json b/mango_v4.json index eee1598b1..63be09512 100644 --- 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": [ diff --git a/programs/mango-v4/src/instructions/flash_loan.rs b/programs/mango-v4/src/instructions/flash_loan.rs index 7085935ba..1edb77ca1 100644 --- 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, ) -> 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::::try_from(vault_ai)?; let token_account = Account::::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::::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::::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 diff --git a/programs/mango-v4/src/lib.rs b/programs/mango-v4/src/lib.rs index 67fe67a6c..45d6dd0f6 100644 --- 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, @@ -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(()) } diff --git a/programs/mango-v4/tests/program_test/mango_client.rs b/programs/mango-v4/tests/program_test/mango_client.rs index ac7c22299..6af48de4c 100644 --- 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, }; diff --git a/ts/client/src/client.ts b/ts/client/src/client.ts index 1c1cbd3e8..8aa083516 100644 --- 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, }) diff --git a/ts/client/src/mango_v4.ts b/ts/client/src/mango_v4.ts index 4350ccf37..b4f03f25f 100644 --- 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": [