From 4ebc156987cc315daeb78562528b5cc55aded759 Mon Sep 17 00:00:00 2001 From: Justice Date: Thu, 28 May 2026 21:08:08 +0100 Subject: [PATCH] test: add flash loan repayment coverage --- src/flash_loan/Cargo.toml | 3 + src/flash_loan/lib.rs | 59 ++-- src/flash_loan/tests.rs | 541 +++++++++++++++++++++++++-------- src/flash_loan/verification.rs | 1 - 4 files changed, 460 insertions(+), 144 deletions(-) diff --git a/src/flash_loan/Cargo.toml b/src/flash_loan/Cargo.toml index 1e0354e..eb37c22 100644 --- a/src/flash_loan/Cargo.toml +++ b/src/flash_loan/Cargo.toml @@ -14,3 +14,6 @@ soroban-sdk = "22.0.0" [dev-dependencies] soroban-sdk = { version = "22.0.0", features = ["testutils"] } + +[lints.rust] +unexpected_cfgs = { level = "warn", check-cfg = ['cfg(kani)'] } diff --git a/src/flash_loan/lib.rs b/src/flash_loan/lib.rs index 0d6a187..98b7025 100644 --- a/src/flash_loan/lib.rs +++ b/src/flash_loan/lib.rs @@ -1,6 +1,24 @@ #![no_std] -use soroban_sdk::{contract, contractclient, contractimpl, contracttype, symbol_short, token, Address, Env, Vec, Map}; +use soroban_sdk::{ + contract, contractclient, contractimpl, contracttype, symbol_short, token, Address, Env, Map, + Vec, +}; + +fn calculate_fee(amount: i128, fee_bps: u32) -> i128 { + // Fee math is shared by single and batch loans so both paths use identical rounding. + amount + .checked_mul(i128::from(fee_bps)) + .and_then(|value| value.checked_div(10_000)) + .expect("Fee calculation overflow") +} + +fn checked_repayment_amount(balance: i128, fee: i128) -> i128 { + // Repayment checks must trap on overflow instead of accepting an invalid balance target. + balance + .checked_add(fee) + .expect("Repayment calculation overflow") +} /// Storage keys for the flash loan provider #[derive(Clone)] @@ -77,13 +95,9 @@ impl FlashLoanProvider { /// * `token` - The address of the token to be lent. /// * `amount` - The amount of tokens to lend. pub fn flash_loan(env: Env, receiver: Address, token: Address, amount: i128) { - // 1. Calculate the fee (5 basis points = 0.05%) - // fee = amount * 5 / 10000 - let fee = amount.checked_mul(5).and_then(|a| a.checked_div(10000)).expect("Fee calculation overflow"); - // 1. Calculate the fee (default 5 basis points = 0.05%) let fee_bps = Self::get_fee_bps(env.clone()); - let fee = amount * fee_bps as i128 / 10000; + let fee = calculate_fee(amount, fee_bps); // 2. Initial balance check let token_client = token::Client::new(&env, &token); @@ -97,20 +111,20 @@ impl FlashLoanProvider { receiver_client.execute_loan(&token, &amount, &fee); // 5. Verify repayment - // This ensures atomic repayment enforcement. If the balance check fails, the + // This ensures atomic repayment enforcement. If the balance check fails, the // whole transaction reverts, ensuring the loan is only successful if repaid. // Soroban's call stack management and the lack of contract state in this provider // make it naturally resistant to reentrancy attacks. let balance_after = token_client.balance(&env.current_contract_address()); - - let required_repayment = balance_before.checked_add(fee).expect("Repayment calculation overflow"); + + let required_repayment = checked_repayment_amount(balance_before, fee); if balance_after < required_repayment { panic!("Flash loan not repaid with fee"); } // Topic: event name only; receiver + token (Addresses) + amounts in data. env.events() - .publish(symbol_short!("flash_ln"), (receiver, token, amount, fee)); + .publish((symbol_short!("flash_ln"),), (receiver, token, amount, fee)); } /// Executes a batch flash loan for multiple assets in a single atomic transaction. @@ -140,19 +154,23 @@ impl FlashLoanProvider { // 1. Calculate fees and check initial balances for all tokens let mut loan_details: Vec = Vec::new(&env); - let mut balance_checks: Map = Map::new(&env); + let mut required_repayments: Map = Map::new(&env); for i in 0..loans.len() { let (token, amount) = loans.get(i).unwrap(); - let fee = amount * fee_bps as i128 / 10000; + let fee = calculate_fee(amount, fee_bps); let token_client = token::Client::new(&env, &token); - let balance_before = token_client.balance(&provider_address); + let current_required = required_repayments + .get(token.clone()) + .unwrap_or_else(|| token_client.balance(&provider_address)); + // Aggregate by token so duplicate-token batches must repay every fee. + let expected_repayment = checked_repayment_amount(current_required, fee); - balance_checks.set(token.clone(), balance_before); + required_repayments.set(token.clone(), expected_repayment); loan_details.push_back(LoanDetail { token: token.clone(), - amount: *amount, + amount, fee, }); } @@ -161,7 +179,7 @@ impl FlashLoanProvider { for i in 0..loans.len() { let (token, amount) = loans.get(i).unwrap(); let token_client = token::Client::new(&env, &token); - token_client.transfer(&provider_address, &receiver, amount); + token_client.transfer(&provider_address, &receiver, &amount); } // 3. Invoke the receiver's batch execution logic @@ -173,9 +191,7 @@ impl FlashLoanProvider { let loan = loan_details.get(i).unwrap(); let token_client = token::Client::new(&env, &loan.token); let balance_after = token_client.balance(&provider_address); - let balance_before = balance_checks.get(loan.token).unwrap(); - - let expected_repayment = balance_before + loan.fee; + let expected_repayment = required_repayments.get(loan.token.clone()).unwrap(); if balance_after < expected_repayment { panic!( "Flash loan not repaid for token {:?}: expected {}, got {}", @@ -188,9 +204,12 @@ impl FlashLoanProvider { // 5. Emit batch event env.events() - .publish((symbol_short!("flash_batch"), receiver), loan_details); + .publish((symbol_short!("fl_batch"), receiver), loan_details); } } mod tests; + +#[allow(unexpected_cfgs)] +#[cfg(kani)] mod verification; diff --git a/src/flash_loan/tests.rs b/src/flash_loan/tests.rs index 5d21b28..228dac2 100644 --- a/src/flash_loan/tests.rs +++ b/src/flash_loan/tests.rs @@ -1,126 +1,329 @@ #[cfg(test)] mod tests { - use super::*; + use crate::{FlashLoanProvider, FlashLoanProviderClient, LoanDetail}; use soroban_sdk::{ - testutils::Address as _, + contract, contractimpl, symbol_short, testutils::Address as _, token::{Client as TokenClient, StellarAssetClient}, Address, Env, Vec, }; - /// A mock receiver contract for testing successful flash loans. - #[contract] - pub struct MockReceiverSuccess; + mod mock_receiver_success { + use super::*; - #[contractimpl] - impl MockReceiverSuccess { - pub fn execute_loan(env: Env, token: Address, amount: i128, fee: i128) { - let token_client = TokenClient::new(&env, &token); - let total_due = amount + fee; + /// A mock receiver contract for testing successful flash loans. + #[contract] + pub struct MockReceiverSuccess; - // Transfer back the amount + fee to the provider - token_client.transfer( - &env.current_contract_address(), - &env.storage() + #[contractimpl] + impl MockReceiverSuccess { + pub fn execute_loan(env: Env, token: Address, amount: i128, fee: i128) { + let token_client = TokenClient::new(&env, &token); + let total_due = amount + fee; + env.storage().instance().set(&symbol_short!("last_tok"), &token); + env.storage().instance().set(&symbol_short!("last_amt"), &amount); + env.storage().instance().set(&symbol_short!("last_fee"), &fee); + + // Transfer back the amount + fee to the provider + token_client.transfer( + &env.current_contract_address(), + &env.storage() + .instance() + .get::<_, Address>(&symbol_short!("provider")) + .unwrap(), + &total_due, + ); + } + + pub fn set_provider(env: Env, provider: Address) { + env.storage() .instance() - .get::<_, Address>(&symbol_short!("provider")) - .unwrap(), - &total_due, - ); - } + .set(&symbol_short!("provider"), &provider); + } - pub fn set_provider(env: Env, provider: Address) { - env.storage() - .instance() - .set(&symbol_short!("provider"), &provider); + pub fn last_token(env: Env) -> Address { + env.storage() + .instance() + .get(&symbol_short!("last_tok")) + .unwrap() + } + + pub fn last_amount(env: Env) -> i128 { + env.storage() + .instance() + .get(&symbol_short!("last_amt")) + .unwrap() + } + + pub fn last_fee(env: Env) -> i128 { + env.storage() + .instance() + .get(&symbol_short!("last_fee")) + .unwrap() + } } } + use mock_receiver_success::{MockReceiverSuccess, MockReceiverSuccessClient}; + + mod mock_receiver_failure { + use super::*; - /// A mock receiver contract for testing failed flash loans. - #[contract] - pub struct MockReceiverFailure; + /// A mock receiver contract for testing failed flash loans. + #[contract] + pub struct MockReceiverFailure; - #[contractimpl] - impl MockReceiverFailure { - pub fn execute_loan(_env: Env, _token: Address, _amount: i128, _fee: i128) { - // Do nothing, return nothing + #[contractimpl] + impl MockReceiverFailure { + pub fn execute_loan(_env: Env, _token: Address, _amount: i128, _fee: i128) { + // Do nothing, return nothing + } } } + use mock_receiver_failure::MockReceiverFailure; + + mod mock_receiver_principal_only { + use super::*; + + /// A mock receiver that repays principal but withholds the required fee. + #[contract] + pub struct MockReceiverPrincipalOnly; + + #[contractimpl] + impl MockReceiverPrincipalOnly { + pub fn execute_loan(env: Env, token: Address, amount: i128, _fee: i128) { + let token_client = TokenClient::new(&env, &token); + token_client.transfer( + &env.current_contract_address(), + &env.storage() + .instance() + .get::<_, Address>(&symbol_short!("provider")) + .unwrap(), + &amount, + ); + } - /// A mock receiver contract for testing successful batch flash loans. - #[contract] - pub struct MockBatchReceiverSuccess; + pub fn set_provider(env: Env, provider: Address) { + env.storage() + .instance() + .set(&symbol_short!("provider"), &provider); + } + } + } + use mock_receiver_principal_only::{ + MockReceiverPrincipalOnly, MockReceiverPrincipalOnlyClient, + }; - #[contractimpl] - impl MockBatchReceiverSuccess { - pub fn execute_batch_loan(env: Env, loans: Vec) { - let provider = env - .storage() - .instance() - .get::<_, Address>(&symbol_short!("provider")) - .unwrap(); + mod mock_batch_receiver_success { + use super::*; - for i in 0..loans.len() { - let loan = loans.get(i).unwrap(); - let token_client = TokenClient::new(&env, &loan.token); - let total_due = loan.amount + loan.fee; + /// A mock receiver contract for testing successful batch flash loans. + #[contract] + pub struct MockBatchReceiverSuccess; - // Transfer back the amount + fee to the provider - token_client.transfer(&env.current_contract_address(), &provider, &total_due); + #[contractimpl] + impl MockBatchReceiverSuccess { + pub fn execute_batch_loan(env: Env, loans: Vec) { + let provider = env + .storage() + .instance() + .get::<_, Address>(&symbol_short!("provider")) + .unwrap(); + let mut total_amount = 0_i128; + let mut total_fee = 0_i128; + + for i in 0..loans.len() { + let loan = loans.get(i).unwrap(); + let token_client = TokenClient::new(&env, &loan.token); + let total_due = loan.amount + loan.fee; + total_amount += loan.amount; + total_fee += loan.fee; + + // Transfer back the amount + fee to the provider + token_client.transfer(&env.current_contract_address(), &provider, &total_due); + } + + env.storage() + .instance() + .set(&symbol_short!("batchcnt"), &loans.len()); + env.storage() + .instance() + .set(&symbol_short!("sum_amt"), &total_amount); + env.storage() + .instance() + .set(&symbol_short!("sum_fee"), &total_fee); + } + + pub fn set_provider(env: Env, provider: Address) { + env.storage() + .instance() + .set(&symbol_short!("provider"), &provider); + } + + pub fn last_batch_count(env: Env) -> u32 { + env.storage() + .instance() + .get(&symbol_short!("batchcnt")) + .unwrap() + } + + pub fn last_batch_amount(env: Env) -> i128 { + env.storage() + .instance() + .get(&symbol_short!("sum_amt")) + .unwrap() + } + + pub fn last_batch_fee(env: Env) -> i128 { + env.storage() + .instance() + .get(&symbol_short!("sum_fee")) + .unwrap() } } + } + use mock_batch_receiver_success::{MockBatchReceiverSuccess, MockBatchReceiverSuccessClient}; + + mod mock_batch_receiver_failure { + use super::*; - pub fn set_provider(env: Env, provider: Address) { - env.storage() - .instance() - .set(&symbol_short!("provider"), &provider); + /// A mock receiver contract for testing failed batch flash loans. + #[contract] + pub struct MockBatchReceiverFailure; + + #[contractimpl] + impl MockBatchReceiverFailure { + pub fn execute_batch_loan(_env: Env, _loans: Vec) { + // Do nothing, return nothing + } } } + use mock_batch_receiver_failure::MockBatchReceiverFailure; + + mod mock_batch_receiver_partial_repayment { + use super::*; + + /// A mock receiver that partially repays batch loans. + #[contract] + pub struct MockBatchReceiverPartialRepayment; - /// A mock receiver contract for testing failed batch flash loans. - #[contract] - pub struct MockBatchReceiverFailure; + #[contractimpl] + impl MockBatchReceiverPartialRepayment { + pub fn execute_batch_loan(env: Env, loans: Vec) { + let provider = env + .storage() + .instance() + .get::<_, Address>(&symbol_short!("provider")) + .unwrap(); + + // Only repay the first loan + if loans.len() > 0 { + let loan = loans.get(0).unwrap(); + let token_client = TokenClient::new(&env, &loan.token); + let total_due = loan.amount + loan.fee; + token_client.transfer(&env.current_contract_address(), &provider, &total_due); + } + } - #[contractimpl] - impl MockBatchReceiverFailure { - pub fn execute_batch_loan(_env: Env, _loans: Vec) { - // Do nothing, return nothing + pub fn set_provider(env: Env, provider: Address) { + env.storage() + .instance() + .set(&symbol_short!("provider"), &provider); + } } } + use mock_batch_receiver_partial_repayment::{ + MockBatchReceiverPartialRepayment, MockBatchReceiverPartialRepaymentClient, + }; + + mod mock_batch_receiver_duplicate_underpay { + use super::*; - /// A mock receiver that partially repays batch loans. - #[contract] - pub struct MockBatchReceiverPartialRepayment; - - #[contractimpl] - impl MockBatchReceiverPartialRepayment { - pub fn execute_batch_loan(env: Env, loans: Vec) { - let provider = env - .storage() - .instance() - .get::<_, Address>(&symbol_short!("provider")) - .unwrap(); - - // Only repay the first loan - if loans.len() > 0 { - let loan = loans.get(0).unwrap(); - let token_client = TokenClient::new(&env, &loan.token); - let total_due = loan.amount + loan.fee; - token_client.transfer(&env.current_contract_address(), &provider, &total_due); + /// A duplicate-token receiver that repays all principal but only the largest fee. + #[contract] + pub struct MockBatchReceiverDuplicateUnderpay; + + #[contractimpl] + impl MockBatchReceiverDuplicateUnderpay { + pub fn execute_batch_loan(env: Env, loans: Vec) { + let provider = env + .storage() + .instance() + .get::<_, Address>(&symbol_short!("provider")) + .unwrap(); + + let mut total_amount = 0_i128; + let mut largest_fee = 0_i128; + let first_loan = loans.get(0).unwrap(); + + for i in 0..loans.len() { + let loan = loans.get(i).unwrap(); + total_amount += loan.amount; + if loan.fee > largest_fee { + largest_fee = loan.fee; + } + } + + let token_client = TokenClient::new(&env, &first_loan.token); + token_client.transfer( + &env.current_contract_address(), + &provider, + &(total_amount + largest_fee), + ); + } + + pub fn set_provider(env: Env, provider: Address) { + env.storage() + .instance() + .set(&symbol_short!("provider"), &provider); } } + } + use mock_batch_receiver_duplicate_underpay::{ + MockBatchReceiverDuplicateUnderpay, MockBatchReceiverDuplicateUnderpayClient, + }; + + fn fee_for(amount: i128, fee_bps: u32) -> i128 { + amount * fee_bps as i128 / 10_000 + } - pub fn set_provider(env: Env, provider: Address) { - env.storage() - .instance() - .set(&symbol_short!("provider"), &provider); + fn mint_to(env: &Env, token: &Address, to: &Address, amount: i128) { + if amount > 0 { + StellarAssetClient::new(env, token).mint(to, &amount); } } + fn fund_receiver_fee( + env: &Env, + token: &Address, + receiver: &Address, + amount: i128, + fee_bps: u32, + ) -> i128 { + let fee = fee_for(amount, fee_bps); + mint_to(env, token, receiver, fee); + fee + } + + fn fund_batch_receiver_fees( + env: &Env, + loans: &Vec<(Address, i128)>, + receiver: &Address, + fee_bps: u32, + ) { + for i in 0..loans.len() { + let (token, amount) = loans.get(i).unwrap(); + mint_to(env, &token, receiver, fee_for(amount, fee_bps)); + } + } + + fn balance(env: &Env, token: &Address, account: &Address) -> i128 { + TokenClient::new(env, token).balance(account) + } + fn setup(env: &Env) -> (Address, Address, Address, Address) { env.mock_all_auths(); let admin = Address::generate(env); - let provider_id = env.register_contract(None, FlashLoanProvider); + let provider_id = env.register(FlashLoanProvider, ()); let token_id = env .register_stellar_asset_contract_v2(admin.clone()) .address(); @@ -129,14 +332,14 @@ mod tests { let sac = StellarAssetClient::new(env, &token_id); sac.mint(&provider_id, &1_000_000); - (provider_id, token_id, admin, admin.clone()) + (provider_id, token_id, admin.clone(), admin) } fn setup_multiple_tokens(env: &Env, count: u32) -> (Address, Vec
, Address) { env.mock_all_auths(); let admin = Address::generate(env); - let provider_id = env.register_contract(None, FlashLoanProvider); + let provider_id = env.register(FlashLoanProvider, ()); let mut token_ids = Vec::new(env); for _ in 0..count { @@ -156,19 +359,21 @@ mod tests { let env = Env::default(); let (provider_id, token_id, _admin, _) = setup(&env); - let receiver_id = env.register_contract(None, MockReceiverSuccess); + let receiver_id = env.register(MockReceiverSuccess, ()); let receiver_client = MockReceiverSuccessClient::new(&env, &receiver_id); receiver_client.set_provider(&provider_id); let provider_client = FlashLoanProviderClient::new(&env, &provider_id); let amount = 100_000; - let fee = amount * 5 / 10000; + let fee = fund_receiver_fee(&env, &token_id, &receiver_id, amount, 5); provider_client.flash_loan(&receiver_id, &token_id, &amount); // Check provider balance: should be initial + fee - let token_client = TokenClient::new(&env, &token_id); - assert_eq!(token_client.balance(&provider_id), 1_000_000 + fee); + assert_eq!(balance(&env, &token_id, &provider_id), 1_000_000 + fee); + assert_eq!(receiver_client.last_token(), token_id); + assert_eq!(receiver_client.last_amount(), amount); + assert_eq!(receiver_client.last_fee(), fee); } #[test] @@ -177,7 +382,21 @@ mod tests { let env = Env::default(); let (provider_id, token_id, _admin, _) = setup(&env); - let receiver_id = env.register_contract(None, MockReceiverFailure); + let receiver_id = env.register(MockReceiverFailure, ()); + + let provider_client = FlashLoanProviderClient::new(&env, &provider_id); + provider_client.flash_loan(&receiver_id, &token_id, &100_000); + } + + #[test] + #[should_panic(expected = "Flash loan not repaid with fee")] + fn test_flash_loan_principal_only_repayment_fails() { + let env = Env::default(); + let (provider_id, token_id, _admin, _) = setup(&env); + + let receiver_id = env.register(MockReceiverPrincipalOnly, ()); + let receiver_client = MockReceiverPrincipalOnlyClient::new(&env, &receiver_id); + receiver_client.set_provider(&provider_id); let provider_client = FlashLoanProviderClient::new(&env, &provider_id); provider_client.flash_loan(&receiver_id, &token_id, &100_000); @@ -216,17 +435,17 @@ mod tests { let provider_client = FlashLoanProviderClient::new(&env, &provider_id); provider_client.set_fee_bps(&10); // 0.10% - let receiver_id = env.register_contract(None, MockReceiverSuccess); + let receiver_id = env.register(MockReceiverSuccess, ()); let receiver_client = MockReceiverSuccessClient::new(&env, &receiver_id); receiver_client.set_provider(&provider_id); let amount = 100_000; - let fee = amount * 10 / 10000; // 0.10% + let fee = fund_receiver_fee(&env, &token_id, &receiver_id, amount, 10); provider_client.flash_loan(&receiver_id, &token_id, &amount); - let token_client = TokenClient::new(&env, &token_id); - assert_eq!(token_client.balance(&provider_id), 1_000_000 + fee); + assert_eq!(balance(&env, &token_id, &provider_id), 1_000_000 + fee); + assert_eq!(receiver_client.last_fee(), fee); } #[test] @@ -234,7 +453,7 @@ mod tests { let env = Env::default(); let (provider_id, token_ids, _admin) = setup_multiple_tokens(&env, 3); - let receiver_id = env.register_contract(None, MockBatchReceiverSuccess); + let receiver_id = env.register(MockBatchReceiverSuccess, ()); let receiver_client = MockBatchReceiverSuccessClient::new(&env, &receiver_id); receiver_client.set_provider(&provider_id); @@ -247,20 +466,29 @@ mod tests { loans.push_back((token_ids.get(2).unwrap(), 25_000)); let fee_bps = 5; - let fee_1 = 100_000 * fee_bps / 10000; - let fee_2 = 50_000 * fee_bps / 10000; - let fee_3 = 25_000 * fee_bps / 10000; + let fee_1 = fee_for(100_000, fee_bps); + let fee_2 = fee_for(50_000, fee_bps); + let fee_3 = fee_for(25_000, fee_bps); + fund_batch_receiver_fees(&env, &loans, &receiver_id, fee_bps); provider_client.flash_loan_batch(&receiver_id, &loans); // Check provider balances: should be initial + fee for each token - let token_client_1 = TokenClient::new(&env, token_ids.get(0).unwrap()); - let token_client_2 = TokenClient::new(&env, token_ids.get(1).unwrap()); - let token_client_3 = TokenClient::new(&env, token_ids.get(2).unwrap()); - - assert_eq!(token_client_1.balance(&provider_id), 1_000_000 + fee_1); - assert_eq!(token_client_2.balance(&provider_id), 1_000_000 + fee_2); - assert_eq!(token_client_3.balance(&provider_id), 1_000_000 + fee_3); + assert_eq!( + balance(&env, &token_ids.get(0).unwrap(), &provider_id), + 1_000_000 + fee_1 + ); + assert_eq!( + balance(&env, &token_ids.get(1).unwrap(), &provider_id), + 1_000_000 + fee_2 + ); + assert_eq!( + balance(&env, &token_ids.get(2).unwrap(), &provider_id), + 1_000_000 + fee_3 + ); + assert_eq!(receiver_client.last_batch_count(), 3); + assert_eq!(receiver_client.last_batch_amount(), 175_000); + assert_eq!(receiver_client.last_batch_fee(), fee_1 + fee_2 + fee_3); } #[test] @@ -268,7 +496,7 @@ mod tests { let env = Env::default(); let (provider_id, token_ids, _admin) = setup_multiple_tokens(&env, 1); - let receiver_id = env.register_contract(None, MockBatchReceiverSuccess); + let receiver_id = env.register(MockBatchReceiverSuccess, ()); let receiver_client = MockBatchReceiverSuccessClient::new(&env, &receiver_id); receiver_client.set_provider(&provider_id); @@ -277,12 +505,15 @@ mod tests { let mut loans = Vec::new(&env); loans.push_back((token_ids.get(0).unwrap(), 100_000)); - let fee = 100_000 * 5 / 10000; + let fee = fee_for(100_000, 5); + fund_batch_receiver_fees(&env, &loans, &receiver_id, 5); provider_client.flash_loan_batch(&receiver_id, &loans); - let token_client = TokenClient::new(&env, token_ids.get(0).unwrap()); - assert_eq!(token_client.balance(&provider_id), 1_000_000 + fee); + assert_eq!( + balance(&env, &token_ids.get(0).unwrap(), &provider_id), + 1_000_000 + fee + ); } #[test] @@ -291,7 +522,7 @@ mod tests { let env = Env::default(); let (provider_id, _token_ids, _admin) = setup_multiple_tokens(&env, 1); - let receiver_id = env.register_contract(None, MockBatchReceiverSuccess); + let receiver_id = env.register(MockBatchReceiverSuccess, ()); let receiver_client = MockBatchReceiverSuccessClient::new(&env, &receiver_id); receiver_client.set_provider(&provider_id); @@ -307,13 +538,14 @@ mod tests { let env = Env::default(); let (provider_id, token_ids, _admin) = setup_multiple_tokens(&env, 2); - let receiver_id = env.register_contract(None, MockBatchReceiverFailure); + let receiver_id = env.register(MockBatchReceiverFailure, ()); let provider_client = FlashLoanProviderClient::new(&env, &provider_id); let mut loans = Vec::new(&env); loans.push_back((token_ids.get(0).unwrap(), 100_000)); loans.push_back((token_ids.get(1).unwrap(), 50_000)); + fund_batch_receiver_fees(&env, &loans, &receiver_id, 5); provider_client.flash_loan_batch(&receiver_id, &loans); } @@ -324,7 +556,7 @@ mod tests { let env = Env::default(); let (provider_id, token_ids, _admin) = setup_multiple_tokens(&env, 2); - let receiver_id = env.register_contract(None, MockBatchReceiverPartialRepayment); + let receiver_id = env.register(MockBatchReceiverPartialRepayment, ()); let receiver_client = MockBatchReceiverPartialRepaymentClient::new(&env, &receiver_id); receiver_client.set_provider(&provider_id); @@ -333,6 +565,7 @@ mod tests { let mut loans = Vec::new(&env); loans.push_back((token_ids.get(0).unwrap(), 100_000)); loans.push_back((token_ids.get(1).unwrap(), 50_000)); + fund_batch_receiver_fees(&env, &loans, &receiver_id, 5); provider_client.flash_loan_batch(&receiver_id, &loans); } @@ -345,7 +578,7 @@ mod tests { let provider_client = FlashLoanProviderClient::new(&env, &provider_id); provider_client.set_fee_bps(&15); // 0.15% - let receiver_id = env.register_contract(None, MockBatchReceiverSuccess); + let receiver_id = env.register(MockBatchReceiverSuccess, ()); let receiver_client = MockBatchReceiverSuccessClient::new(&env, &receiver_id); receiver_client.set_provider(&provider_id); @@ -354,16 +587,20 @@ mod tests { loans.push_back((token_ids.get(1).unwrap(), 50_000)); let fee_bps = 15; - let fee_1 = 100_000 * fee_bps / 10000; - let fee_2 = 50_000 * fee_bps / 10000; + let fee_1 = fee_for(100_000, fee_bps); + let fee_2 = fee_for(50_000, fee_bps); + fund_batch_receiver_fees(&env, &loans, &receiver_id, fee_bps); provider_client.flash_loan_batch(&receiver_id, &loans); - let token_client_1 = TokenClient::new(&env, token_ids.get(0).unwrap()); - let token_client_2 = TokenClient::new(&env, token_ids.get(1).unwrap()); - - assert_eq!(token_client_1.balance(&provider_id), 1_000_000 + fee_1); - assert_eq!(token_client_2.balance(&provider_id), 1_000_000 + fee_2); + assert_eq!( + balance(&env, &token_ids.get(0).unwrap(), &provider_id), + 1_000_000 + fee_1 + ); + assert_eq!( + balance(&env, &token_ids.get(1).unwrap(), &provider_id), + 1_000_000 + fee_2 + ); } #[test] @@ -371,7 +608,7 @@ mod tests { let env = Env::default(); let (provider_id, token_ids, _admin) = setup_multiple_tokens(&env, 5); - let receiver_id = env.register_contract(None, MockBatchReceiverSuccess); + let receiver_id = env.register(MockBatchReceiverSuccess, ()); let receiver_client = MockBatchReceiverSuccessClient::new(&env, &receiver_id); receiver_client.set_provider(&provider_id); @@ -381,15 +618,73 @@ mod tests { for i in 0..5 { loans.push_back((token_ids.get(i).unwrap(), (i + 1) as i128 * 10_000)); } + fund_batch_receiver_fees(&env, &loans, &receiver_id, 5); provider_client.flash_loan_batch(&receiver_id, &loans); let fee_bps = 5; for i in 0..5 { let amount = (i + 1) as i128 * 10_000; - let fee = amount * fee_bps / 10000; - let token_client = TokenClient::new(&env, token_ids.get(i).unwrap()); - assert_eq!(token_client.balance(&provider_id), 1_000_000 + fee); + let fee = fee_for(amount, fee_bps); + assert_eq!( + balance(&env, &token_ids.get(i).unwrap(), &provider_id), + 1_000_000 + fee + ); } } + + #[test] + fn test_flash_loan_batch_duplicate_token_success_accumulates_fees() { + let env = Env::default(); + let (provider_id, token_ids, _admin) = setup_multiple_tokens(&env, 1); + let token_id = token_ids.get(0).unwrap(); + + let receiver_id = env.register(MockBatchReceiverSuccess, ()); + let receiver_client = MockBatchReceiverSuccessClient::new(&env, &receiver_id); + receiver_client.set_provider(&provider_id); + + let provider_client = FlashLoanProviderClient::new(&env, &provider_id); + let mut loans = Vec::new(&env); + loans.push_back((token_id.clone(), 100_000)); + loans.push_back((token_id.clone(), 50_000)); + + let fee_1 = fee_for(100_000, 5); + let fee_2 = fee_for(50_000, 5); + fund_batch_receiver_fees(&env, &loans, &receiver_id, 5); + + provider_client.flash_loan_batch(&receiver_id, &loans); + + assert_eq!( + balance(&env, &token_id, &provider_id), + 1_000_000 + fee_1 + fee_2 + ); + assert_eq!(receiver_client.last_batch_count(), 2); + assert_eq!(receiver_client.last_batch_fee(), fee_1 + fee_2); + } + + #[test] + #[should_panic(expected = "Flash loan not repaid")] + fn test_flash_loan_batch_rejects_duplicate_token_fee_underpayment() { + let env = Env::default(); + let (provider_id, token_ids, _admin) = setup_multiple_tokens(&env, 1); + let token_id = token_ids.get(0).unwrap(); + + let receiver_id = env.register(MockBatchReceiverDuplicateUnderpay, ()); + let receiver_client = MockBatchReceiverDuplicateUnderpayClient::new(&env, &receiver_id); + receiver_client.set_provider(&provider_id); + + let provider_client = FlashLoanProviderClient::new(&env, &provider_id); + let mut loans = Vec::new(&env); + loans.push_back((token_id.clone(), 100_000)); + loans.push_back((token_id, 50_000)); + + mint_to( + &env, + &token_ids.get(0).unwrap(), + &receiver_id, + fee_for(100_000, 5), + ); + + provider_client.flash_loan_batch(&receiver_id, &loans); + } } diff --git a/src/flash_loan/verification.rs b/src/flash_loan/verification.rs index 95681fa..770b672 100644 --- a/src/flash_loan/verification.rs +++ b/src/flash_loan/verification.rs @@ -1,4 +1,3 @@ -#[cfg(kani)] mod verification { use kani;