From 5219fde2936c77dbfbc167c8f2081c19673680e1 Mon Sep 17 00:00:00 2001 From: Manceraider24 Date: Tue, 23 Jun 2026 21:49:00 +0000 Subject: [PATCH] feat(treasury): implement emergency_drain (#966) - Requires admin auth and panics if protocol is not paused - Drains full BALANCE via token transfer, resets balance to 0 - Appends entry to WITHDRAWAL_LOG, emits EmergencyDrain event - Returns total amount drained in stroops - Adds unit tests: success, not-paused guard, unauthorized guard --- contracts/treasury/src/lib.rs | 189 +++++++++++++++++++++++++++++++++- 1 file changed, 187 insertions(+), 2 deletions(-) diff --git a/contracts/treasury/src/lib.rs b/contracts/treasury/src/lib.rs index 57001752..91501f70 100644 --- a/contracts/treasury/src/lib.rs +++ b/contracts/treasury/src/lib.rs @@ -1,5 +1,5 @@ #![no_std] -use soroban_sdk::{contract, contractimpl, Address, Bytes, Env, Vec}; +use soroban_sdk::{contract, contractimpl, contracttype, symbol_short, Address, Bytes, Env, Vec}; // ─── STORAGE KEYS ───────────────────────────────────────────────────────────── // ADMIN -> Address @@ -8,6 +8,18 @@ use soroban_sdk::{contract, contractimpl, Address, Bytes, Env, Vec}; // TOTAL_FEES_EARNED -> i128 // WITHDRAWAL_LOG -> Vec<(Address, i128, u64)> +#[contracttype] +#[derive(Clone, Debug)] +pub struct ProtocolConfig { + pub admin: Address, + pub fee_collector: Address, + pub default_fee_bp: u32, + pub min_bet_amount: i128, + pub max_bet_amount: i128, + pub dispute_window_sec: u64, + pub paused: bool, +} + #[contract] pub struct Treasury; @@ -42,7 +54,42 @@ impl Treasury { /// Logs the drain. Emits EmergencyDrain event. /// Returns total amount drained in stroops. pub fn emergency_drain(env: Env, admin: Address, recipient: Address) -> i128 { - todo!("implement: require_auth(admin), verify protocol is paused, transfer full BALANCE, set BALANCE=0, log, emit event, return drained amount") + // 1. Authentication + admin.require_auth(); + + // 2. State Check — protocol must be paused + let factory: Address = env.storage().persistent().get(&symbol_short!("FACTORY")).unwrap(); + let config: ProtocolConfig = env.invoke_contract(&factory, &symbol_short!("get_config"), soroban_sdk::vec![&env]); + if !config.paused { + panic!("protocol is not paused"); + } + + // 3. Funds Transfer — drain full balance + let amount: i128 = env.storage().persistent().get(&symbol_short!("BALANCE")).unwrap_or(0); + let token: Address = env.storage().persistent().get(&symbol_short!("TOKEN")).unwrap(); + soroban_sdk::token::Client::new(&env, &token).transfer( + &env.current_contract_address(), + &recipient, + &amount, + ); + env.storage().persistent().set(&symbol_short!("BALANCE"), &0_i128); + + // 4. Logging & Events + let mut log: Vec<(Address, i128, u64)> = env + .storage() + .persistent() + .get(&symbol_short!("WLOG")) + .unwrap_or(soroban_sdk::vec![&env]); + log.push_back((recipient.clone(), amount, env.ledger().timestamp())); + env.storage().persistent().set(&symbol_short!("WLOG"), &log); + + env.events().publish( + (symbol_short!("EmrgDrain"), recipient), + amount, + ); + + // 5. Return drained amount + amount } /// Returns current treasury XLM balance in stroops. @@ -60,3 +107,141 @@ impl Treasury { todo!("implement: read WITHDRAWAL_LOG from storage and return") } } + +#[cfg(test)] +mod tests { + use super::*; + use soroban_sdk::{ + testutils::{Address as _, AuthorizedFunction, AuthorizedInvocation, Events, Ledger}, + vec, IntoVal, Symbol, + }; + + // ── helpers ────────────────────────────────────────────────────────────── + + /// Registers a Treasury contract and pre-seeds its storage so that + /// `emergency_drain` has the data it needs without calling `initialize` + /// (which is still a `todo!`). + fn setup( + env: &Env, + paused: bool, + balance: i128, + ) -> (TreasuryClient, Address, Address, Address) { + let admin = Address::generate(env); + let recipient = Address::generate(env); + + // Deploy a mock token + let token_admin = Address::generate(env); + let token_id = env.register_stellar_asset_contract_v2(token_admin.clone()).address(); + let token = soroban_sdk::token::StellarAssetClient::new(env, &token_id); + + // Deploy a mock factory whose `get_config` returns a ProtocolConfig + let factory_id = env.register(MockFactory, (admin.clone(), paused)); + + // Deploy Treasury and seed its storage + let treasury_id = env.register(Treasury, ()); + env.as_contract(&treasury_id, || { + env.storage().persistent().set(&symbol_short!("FACTORY"), &factory_id); + env.storage().persistent().set(&symbol_short!("TOKEN"), &token_id); + env.storage().persistent().set(&symbol_short!("ADMIN"), &admin); + env.storage().persistent().set(&symbol_short!("BALANCE"), &balance); + }); + + // Mint `balance` tokens into the treasury contract + token.mint(&treasury_id, &balance); + + let client = TreasuryClient::new(env, &treasury_id); + (client, admin, recipient, token_id) + } + + // ── mock factory ───────────────────────────────────────────────────────── + + #[contract] + struct MockFactory; + + #[contractimpl] + impl MockFactory { + pub fn __constructor(env: Env, admin: Address, paused: bool) { + env.storage().persistent().set(&symbol_short!("admin"), &admin); + env.storage().persistent().set(&symbol_short!("paused"), &paused); + } + + pub fn get_config(env: Env) -> ProtocolConfig { + let admin: Address = env.storage().persistent().get(&symbol_short!("admin")).unwrap(); + let paused: bool = env.storage().persistent().get(&symbol_short!("paused")).unwrap(); + ProtocolConfig { + admin: admin.clone(), + fee_collector: admin, + default_fee_bp: 200, + min_bet_amount: 1_000_000, + max_bet_amount: 100_000_000, + dispute_window_sec: 86_400, + paused, + } + } + } + + // ── tests ───────────────────────────────────────────────────────────────── + + #[test] + fn test_emergency_drain_success() { + let env = Env::default(); + env.mock_all_auths(); + + let balance = 50_000_000_i128; + let (client, admin, recipient, token_id) = setup(&env, true, balance); + + let drained = client.emergency_drain(&admin, &recipient); + + // Return value + assert_eq!(drained, balance); + + // BALANCE set to 0 + let stored_balance: i128 = env.as_contract(&client.address, || { + env.storage().persistent().get(&symbol_short!("BALANCE")).unwrap() + }); + assert_eq!(stored_balance, 0); + + // Token actually transferred + let token = soroban_sdk::token::Client::new(&env, &token_id); + assert_eq!(token.balance(&recipient), balance); + assert_eq!(token.balance(&client.address), 0); + + // Withdrawal log updated + let log: Vec<(Address, i128, u64)> = env.as_contract(&client.address, || { + env.storage().persistent().get(&symbol_short!("WLOG")).unwrap() + }); + assert_eq!(log.len(), 1); + let (log_recipient, log_amount, _) = log.get(0).unwrap(); + assert_eq!(log_recipient, recipient); + assert_eq!(log_amount, balance); + + // EmergencyDrain event emitted + let events = env.events().all(); + let found = events.iter().any(|(_, topics, data)| { + topics.contains(&symbol_short!("EmrgDrain").into_val(&env)) + && data == balance.into_val(&env) + }); + assert!(found, "EmergencyDrain event not found"); + } + + #[test] + #[should_panic(expected = "protocol is not paused")] + fn test_emergency_drain_fails_when_not_paused() { + let env = Env::default(); + env.mock_all_auths(); + + let (client, admin, recipient, _) = setup(&env, false, 10_000_000); + client.emergency_drain(&admin, &recipient); + } + + #[test] + #[should_panic] + fn test_emergency_drain_fails_when_unauthorized() { + let env = Env::default(); + // Do NOT mock auths — a non-admin call must fail auth check. + + let (client, _admin, recipient, _) = setup(&env, true, 10_000_000); + let attacker = Address::generate(&env); + client.emergency_drain(&attacker, &recipient); + } +}