diff --git a/contracts/flash_loan_guard/src/lib.rs b/contracts/flash_loan_guard/src/lib.rs new file mode 100644 index 00000000..0a2d130c --- /dev/null +++ b/contracts/flash_loan_guard/src/lib.rs @@ -0,0 +1,36 @@ +use soroban_sdk::{contract, contractimpl, contracttype, Address, Env}; + +#[derive(Clone)] +#[contracttype] +pub enum DataKey { + DepositBlock(Address), +} + +@contract +pub struct FlashLoanGuardContract; + +@contractimpl +impl FlashLoanGuardContract { + pub fn deposit(env: Env, user: Address) { + user.require_auth(); + let current_block = env.ledger().sequence(); + env.storage().persistent().set(&DataKey::DepositBlock(user.clone()), ¤t_block); + } + + pub fn withdraw(env: Env, user: Address) { + user.require_auth(); + let current_block = env.ledger().sequence(); + + let deposit_block: u32 = env.storage() + .persistent() + .get(&DataKey::DepositBlock(user.clone())) + .unwrap_or(0); + + if current_block <= deposit_block { + panic!("FlashLoanReentrancy: withdrawal prohibited in the same ledger block as deposit"); + } + + // Proceed with withdrawal logic... + env.storage().persistent().remove(&DataKey::DepositBlock(user)); + } +} \ No newline at end of file diff --git a/contracts/flash_loan_guard/src/test.rs b/contracts/flash_loan_guard/src/test.rs new file mode 100644 index 00000000..d8a7ee49 --- /dev/null +++ b/contracts/flash_loan_guard/src/test.rs @@ -0,0 +1,37 @@ +#[cfg(test)] +mod test { + use super::*; + use soroban_sdk::{Env, Address}; + + #[test] + #[should_panic(expected = "FlashLoanReentrancy")] + fn test_prevents_same_block_withdrawal() { + let env = Env::default(); + env.mock_all_auths(); + let contract_id = env.register(FlashLoanGuardContract, ()); + let client = FlashLoanGuardContractClient::new(&env, &contract_id); + let user = Address::generate(&env); + + env.ledger().set_sequence_number(100); + client.deposit(&user); + + // Attempt withdrawal in the same block (100) -> should panic + client.withdraw(&user); + } + + #[test] + fn test_allows_subsequent_block_withdrawal() { + let env = Env::default(); + env.mock_all_auths(); + let contract_id = env.register(FlashLoanGuardContract, ()); + let client = FlashLoanGuardContractClient::new(&env, &contract_id); + let user = Address::generate(&env); + + env.ledger().set_sequence_number(100); + client.deposit(&user); + + // Advance ledger block + env.ledger().set_sequence_number(101); + client.withdraw(&user); // Should succeed without panicking + } +} \ No newline at end of file