Fix proportional emergency withdrawals

contracts/sorosave/src/admin.rs

@@ -1,4 +1,4 @@
-use soroban_sdk::{Address, Env, String};
+use soroban_sdk::{Address, Env, Map, String};
 
 use crate::errors::ContractError;
 use crate::storage;
@@ -127,16 +127,66 @@ pub fn emergency_withdraw(env: &Env, admin: Address, group_id: u64) -> Result<()
         return Err(ContractError::GroupCompleted);
     }
 
-    // Calculate remaining balance and distribute equally
     let token_client = soroban_sdk::token::Client::new(env, &group.token);
     let contract_addr = env.current_contract_address();
     let balance = token_client.balance(&contract_addr);
 
     if balance > 0 {
-        let per_member = balance / group.members.len() as i128;
-        if per_member > 0 {
+        let mut contribution_amounts: Map<Address, i128> = Map::new(env);
+        let mut total_contributed = 0i128;
+
+        if group.current_round > 0 {
+            if let Some(round) = storage::get_round(env, group_id, group.current_round) {
+                for member in group.members.iter() {
+                    if round.contributions.contains_key(member.clone()) {
+                        contribution_amounts.set(member.clone(), group.contribution_amount);
+                        total_contributed += group.contribution_amount;
+                    }
+                }
+            }
+        }
+
+        if total_contributed > 0 {
+            let mut distributed = 0i128;
+            let mut remainder_recipient = group.admin.clone();
+            let mut largest_contribution = 0i128;
+
             for member in group.members.iter() {
-                token_client.transfer(&contract_addr, &member, &per_member);
+                let contributed = contribution_amounts.get(member.clone()).unwrap_or(0);
+                if contributed <= 0 {
+                    continue;
+                }
+
+                if contributed > largest_contribution {
+                    largest_contribution = contributed;
+                    remainder_recipient = member.clone();
+                }
+
+                let share = (balance * contributed) / total_contributed;
+                if share > 0 {
+                    token_client.transfer(&contract_addr, &member, &share);
+                    distributed += share;
+                }
+            }
+
+            let remainder = balance - distributed;
+            if remainder > 0 {
+                token_client.transfer(&contract_addr, &remainder_recipient, &remainder);
+            }
+        } else {
+            let per_member = balance / group.members.len() as i128;
+            let mut distributed = 0i128;
+
+            if per_member > 0 {
+                for member in group.members.iter() {
+                    token_client.transfer(&contract_addr, &member, &per_member);
+                    distributed += per_member;
+                }
+            }
+
+            let remainder = balance - distributed;
+            if remainder > 0 {
+                token_client.transfer(&contract_addr, &group.admin, &remainder);
             }
         }
     }

contracts/sorosave/src/test.rs

@@ -1,4 +1,8 @@
-use soroban_sdk::{testutils::Address as _, token::StellarAssetClient, Address, Env, String};
+use soroban_sdk::{
+    testutils::Address as _,
+    token::{Client as TokenClient, StellarAssetClient},
+    Address, Env, String,
+};
 
 use crate::types::GroupStatus;
 use crate::{SoroSaveContract, SoroSaveContractClient};
@@ -222,3 +226,92 @@ fn test_set_group_admin() {
     let group = client.get_group(&group_id);
     assert_eq!(group.admin, new_admin);
 }
+
+#[test]
+fn test_emergency_withdraw_mid_round_returns_actual_contributions() {
+    let (env, admin, client, _token) = setup_env();
+
+    let token_admin = Address::generate(&env);
+    let token_id = env.register_stellar_asset_contract_v2(token_admin);
+    let token = token_id.address();
+    let asset_client = StellarAssetClient::new(&env, &token);
+    let token_client = TokenClient::new(&env, &token);
+
+    let member1 = Address::generate(&env);
+    let member2 = Address::generate(&env);
+    let member3 = Address::generate(&env);
+    let member4 = Address::generate(&env);
+
+    for member in [
+        admin.clone(),
+        member1.clone(),
+        member2.clone(),
+        member3.clone(),
+        member4.clone(),
+    ] {
+        asset_client.mint(&member, &20_000_000);
+    }
+
+    let group_id = client.create_group(
+        &admin,
+        &String::from_str(&env, "Emergency Withdraw Unequal"),
+        &token,
+        &1_000_000,
+        &86400,
+        &5,
+    );
+    client.join_group(&member1, &group_id);
+    client.join_group(&member2, &group_id);
+    client.join_group(&member3, &group_id);
+    client.join_group(&member4, &group_id);
+    client.start_group(&admin, &group_id);
+
+    // Complete and pay out two full rounds so their funds are no longer in escrow.
+    for contributor in [
+        admin.clone(),
+        member1.clone(),
+        member2.clone(),
+        member3.clone(),
+        member4.clone(),
+    ] {
+        client.contribute(&contributor, &group_id);
+    }
+    client.distribute_payout(&group_id);
+
+    for contributor in [
+        admin.clone(),
+        member1.clone(),
+        member2.clone(),
+        member3.clone(),
+        member4.clone(),
+    ] {
+        client.contribute(&contributor, &group_id);
+    }
+    client.distribute_payout(&group_id);
+
+    // Trigger the emergency path mid-round 3 after only two members contributed.
+    client.contribute(&admin, &group_id);
+    client.contribute(&member1, &group_id);
+
+    let contract_addr = client.address.clone();
+    let balance_before = token_client.balance(&contract_addr);
+    assert_eq!(balance_before, 2_000_000);
+
+    let admin_before = token_client.balance(&admin);
+    let member1_before = token_client.balance(&member1);
+    let member2_before = token_client.balance(&member2);
+    let member3_before = token_client.balance(&member3);
+    let member4_before = token_client.balance(&member4);
+
+    client.emergency_withdraw(&admin, &group_id);
+
+    assert_eq!(token_client.balance(&admin) - admin_before, 1_000_000);
+    assert_eq!(token_client.balance(&member1) - member1_before, 1_000_000);
+    assert_eq!(token_client.balance(&member2) - member2_before, 0);
+    assert_eq!(token_client.balance(&member3) - member3_before, 0);
+    assert_eq!(token_client.balance(&member4) - member4_before, 0);
+    assert_eq!(token_client.balance(&contract_addr), 0);
+
+    let group = client.get_group(&group_id);
+    assert_eq!(group.status, GroupStatus::Completed);
+}