diff --git a/contracts/inheritance-contract/src/lib.rs b/contracts/inheritance-contract/src/lib.rs index 0760c51c2..bc0a83cb5 100644 --- a/contracts/inheritance-contract/src/lib.rs +++ b/contracts/inheritance-contract/src/lib.rs @@ -4,8 +4,6 @@ use soroban_sdk::{contract, contracterror, contractimpl, contracttype, Address, const MAX_BENEFICIARIES: u32 = 100; const PLAN_TTL_THRESHOLD: u32 = 500; const PLAN_TTL_LEEWAY: u32 = 100; -const TEMP_TTL_THRESHOLD: u32 = 100; -const TEMP_TTL_LEEWAY: u32 = 50; #[contracterror] #[derive(Copy, Clone, Debug, Eq, PartialEq)] @@ -18,6 +16,8 @@ pub enum Error { NegativeAmount = 6, InsufficientBalance = 7, TooManyBeneficiaries = 8, + TimelockNotExpired = 9, + PayoutNotTriggered = 10, } #[contracttype] @@ -40,6 +40,7 @@ pub struct Plan { pub earn_yield: bool, pub yield_rate_bps: u32, pub is_active: bool, + pub timelock_duration: u64, } pub type InheritancePlan = Plan; @@ -66,12 +67,6 @@ impl InheritanceContract { .persistent() .extend_ttl(key, PLAN_TTL_LEEWAY, PLAN_TTL_THRESHOLD); } - - fn extend_temp_ttl(env: &Env, key: &DataKey) { - env.storage() - .temporary() - .extend_ttl(key, TEMP_TTL_LEEWAY, TEMP_TTL_THRESHOLD); - } } #[contractimpl] @@ -89,6 +84,7 @@ impl InheritanceContract { grace_period: u64, earn_yield: bool, yield_rate_bps: u32, + timelock_duration: u64, ) -> Result<(), Error> { owner.require_auth(); @@ -131,6 +127,7 @@ impl InheritanceContract { earn_yield, yield_rate_bps, is_active: true, + timelock_duration, }; env.storage().persistent().set(&key, &plan); @@ -163,11 +160,11 @@ impl InheritanceContract { /// emit payout events, and trigger anchor event emissions for fiat recipients. pub fn claim(env: Env, owner: Address) -> Result<(), Error> { let key = DataKey::Plan(owner.clone()); - if !env.storage().persistent().has(&key) { - return Err(Error::PlanNotFound); - } - - let plan: Plan = env.storage().persistent().get(&key).unwrap(); + let plan: Plan = env + .storage() + .persistent() + .get(&key) + .ok_or(Error::PlanNotFound)?; if plan.is_active { return Err(Error::InactivityPeriodNotMet); @@ -179,8 +176,38 @@ impl InheritanceContract { } let claim_key = DataKey::ClaimStatus(owner.clone()); - env.storage().temporary().set(&claim_key, &true); - Self::extend_temp_ttl(&env, &claim_key); + if env.storage().persistent().has(&claim_key) { + return Ok(()); // Already claimed + } + + env.storage().persistent().set(&claim_key, ¤t_time); + Self::extend_plan_ttl(&env, &claim_key); + + Ok(()) + } + + /// Cancel a triggered payout during the timelock window. + pub fn cancel_claim(env: Env, owner: Address) -> Result<(), Error> { + owner.require_auth(); + + let key = DataKey::Plan(owner.clone()); + let mut plan: Plan = env + .storage() + .persistent() + .get(&key) + .ok_or(Error::PlanNotFound)?; + + let claim_key = DataKey::ClaimStatus(owner.clone()); + if !env.storage().persistent().has(&claim_key) { + return Err(Error::PayoutNotTriggered); + } + + env.storage().persistent().remove(&claim_key); + + plan.is_active = true; + plan.last_ping = env.ledger().timestamp(); + env.storage().persistent().set(&key, &plan); + Self::extend_plan_ttl(&env, &key); Ok(()) } @@ -212,18 +239,22 @@ impl InheritanceContract { .get(&key) .ok_or(Error::PlanNotFound)?; - if plan.is_active { - return Err(Error::InactivityPeriodNotMet); - } + let claim_key = DataKey::ClaimStatus(owner.clone()); + let claim_time: u64 = env + .storage() + .persistent() + .get(&claim_key) + .ok_or(Error::PayoutNotTriggered)?; let current_time = env.ledger().timestamp(); - if current_time < plan.last_ping + plan.grace_period { - return Err(Error::InactivityPeriodNotMet); + if current_time < claim_time + plan.timelock_duration { + return Err(Error::TimelockNotExpired); } // Checks-effects-interactions: remove plan before transfers // to prevent double payout and guard against re-entrancy env.storage().persistent().remove(&key); + env.storage().persistent().remove(&claim_key); let token_client = soroban_sdk::token::Client::new(&env, &plan.token); let n = plan.beneficiaries.len(); @@ -265,6 +296,30 @@ impl InheritanceContract { Ok(()) } + + /// Reclaim the locked assets and delete the plan. + pub fn reclaim(env: Env, owner: Address) -> Result<(), Error> { + owner.require_auth(); + + let key = DataKey::Plan(owner.clone()); + let plan: Plan = env + .storage() + .persistent() + .get(&key) + .ok_or(Error::PlanNotFound)?; + + let claim_key = DataKey::ClaimStatus(owner.clone()); + if env.storage().persistent().has(&claim_key) { + env.storage().persistent().remove(&claim_key); + } + + env.storage().persistent().remove(&key); + + let token_client = soroban_sdk::token::Client::new(&env, &plan.token); + token_client.transfer(&env.current_contract_address(), &owner, &plan.amount); + + Ok(()) + } } #[cfg(test)] diff --git a/contracts/inheritance-contract/src/test.rs b/contracts/inheritance-contract/src/test.rs index 3765cc18b..b441860e2 100644 --- a/contracts/inheritance-contract/src/test.rs +++ b/contracts/inheritance-contract/src/test.rs @@ -1,5 +1,3 @@ -#![cfg(test)] - use super::*; use soroban_sdk::testutils::Address as _; use soroban_sdk::testutils::Ledger; @@ -46,6 +44,7 @@ fn test_create_plan_success() { &3600, &true, &500, + &86400, ); // Verify balances @@ -98,6 +97,7 @@ fn test_create_plan_insufficient_balance() { &3600, &true, &500, + &86400, ); assert_eq!(result, Err(Ok(Error::InsufficientBalance))); @@ -132,6 +132,7 @@ fn test_create_plan_negative_or_zero_amount() { &3600, &true, &500, + &86400, ); assert_eq!(result_zero, Err(Ok(Error::NegativeAmount))); @@ -144,6 +145,7 @@ fn test_create_plan_negative_or_zero_amount() { &3600, &true, &500, + &86400, ); assert_eq!(result_neg, Err(Ok(Error::NegativeAmount))); } @@ -182,6 +184,7 @@ fn test_create_plan_invalid_basis_points() { &3600, &true, &500, + &86400, ); assert_eq!(result, Err(Ok(Error::InvalidBasisPoints))); @@ -216,6 +219,7 @@ fn test_create_plan_already_exists() { &3600, &true, &500, + &86400, ); // Second creation on same owner @@ -227,6 +231,7 @@ fn test_create_plan_already_exists() { &3600, &true, &500, + &86400, ); assert_eq!(result2, Err(Ok(Error::PlanAlreadyExists))); } @@ -264,6 +269,7 @@ fn test_trigger_payout_single_beneficiary() { &3600, &true, &500, + &86400, ); // Deactivate plan @@ -273,6 +279,8 @@ fn test_trigger_payout_single_beneficiary() { env.ledger().set_timestamp(start + 4000); // Trigger payout + client.claim(&owner); + env.ledger().set_timestamp(env.ledger().timestamp() + 86400); client.trigger_payout(&owner); // Beneficiary receives full amount, contract emptied @@ -328,11 +336,14 @@ fn test_trigger_payout_multiple_beneficiaries() { &3600, &true, &500, + &86400, ); client.close_plan(&owner); env.ledger().set_timestamp(1_000_000 + 4000); + client.claim(&owner); + env.ledger().set_timestamp(env.ledger().timestamp() + 86400); client.trigger_payout(&owner); // Alice: 1000 * 5000 / 10000 = 500 @@ -382,11 +393,14 @@ fn test_trigger_payout_dust_goes_to_last_beneficiary() { &3600, &false, &0, + &86400, ); client.close_plan(&owner); env.ledger().set_timestamp(1_000_000 + 4000); + client.claim(&owner); + env.ledger().set_timestamp(env.ledger().timestamp() + 86400); client.trigger_payout(&owner); // A: 100 * 3333 / 10000 = 33 (integer truncation) @@ -428,12 +442,13 @@ fn test_trigger_payout_plan_still_active() { &3600, &false, &0, + &86400, ); // Plan is still active — close_plan was never called env.ledger().set_timestamp(1_000_000 + 4000); - let result = client.try_trigger_payout(&owner); + let result = client.try_claim(&owner); assert_eq!(result, Err(Ok(Error::InactivityPeriodNotMet))); } @@ -469,6 +484,7 @@ fn test_trigger_payout_grace_period_not_met() { &3600, &false, &0, + &86400, ); client.close_plan(&owner); @@ -476,7 +492,7 @@ fn test_trigger_payout_grace_period_not_met() { // Only 1000 seconds passed — need 3600 env.ledger().set_timestamp(1_000_000 + 1000); - let result = client.try_trigger_payout(&owner); + let result = client.try_claim(&owner); assert_eq!(result, Err(Ok(Error::InactivityPeriodNotMet))); } @@ -512,12 +528,15 @@ fn test_trigger_payout_double_payout_prevented() { &3600, &false, &0, + &86400, ); client.close_plan(&owner); env.ledger().set_timestamp(1_000_000 + 4000); // First payout succeeds + client.claim(&owner); + env.ledger().set_timestamp(env.ledger().timestamp() + 86400); client.trigger_payout(&owner); assert_eq!(token_client.balance(&beneficiary), 500); @@ -539,3 +558,100 @@ fn test_trigger_payout_no_plan() { let result = client.try_trigger_payout(&owner); assert_eq!(result, Err(Ok(Error::PlanNotFound))); } + +#[test] +fn test_cancel_claim_success() { + let env = Env::default(); + env.mock_all_auths(); + + let contract_id = env.register_contract(None, InheritanceContract); + let client = InheritanceContractClient::new(&env, &contract_id); + + let token_id = env.register_contract(None, mock_token::MockToken); + let token_client = mock_token::MockTokenClient::new(&env, &token_id); + + let owner = Address::generate(&env); + let beneficiary = Address::generate(&env); + + token_client.mint(&owner, &2000); + + let b = Beneficiary { + address: beneficiary.clone(), + allocation_bps: 10000, + fiat_anchor_info: String::from_str(&env, ""), + }; + + let start = 1_000_000; + env.ledger().set_timestamp(start); + + client.create_plan( + &owner, + &token_id, + &500, + &Vec::from_array(&env, [b]), + &3600, + &false, + &0, + &86400, + ); + + client.close_plan(&owner); + env.ledger().set_timestamp(start + 4000); + + // Trigger payout + client.claim(&owner); + + // Cancel payout + client.cancel_claim(&owner); + + // Attempting trigger_payout should now fail since the payout has been cancelled + env.ledger().set_timestamp(env.ledger().timestamp() + 86400); + let result = client.try_trigger_payout(&owner); + assert_eq!(result, Err(Ok(Error::PayoutNotTriggered))); +} + +#[test] +fn test_reclaim_success() { + let env = Env::default(); + env.mock_all_auths(); + + let contract_id = env.register_contract(None, InheritanceContract); + let client = InheritanceContractClient::new(&env, &contract_id); + + let token_id = env.register_contract(None, mock_token::MockToken); + let token_client = mock_token::MockTokenClient::new(&env, &token_id); + + let owner = Address::generate(&env); + let beneficiary = Address::generate(&env); + + token_client.mint(&owner, &2000); + + let b = Beneficiary { + address: beneficiary.clone(), + allocation_bps: 10000, + fiat_anchor_info: String::from_str(&env, ""), + }; + + let start = 1_000_000; + env.ledger().set_timestamp(start); + + client.create_plan( + &owner, + &token_id, + &500, + &Vec::from_array(&env, [b]), + &3600, + &false, + &0, + &86400, + ); + + // Owner reclaims before claim + client.reclaim(&owner); + + assert_eq!(token_client.balance(&owner), 2000); + assert_eq!(token_client.balance(&contract_id), 0); + + let result = client.try_get_plan(&owner); + assert_eq!(result, Err(Ok(Error::PlanNotFound))); +}