diff --git a/contracts/escrow/src/lib.rs b/contracts/escrow/src/lib.rs index c789280..15d84f1 100644 --- a/contracts/escrow/src/lib.rs +++ b/contracts/escrow/src/lib.rs @@ -85,13 +85,32 @@ impl EscrowContract { .ok_or(Error::Unauthorized)?; admin.require_auth(); + let already_allowed: bool = env + .storage() + .instance() + .get(&DataKey::AllowedToken(token.clone())) + .unwrap_or(false); + env.storage() - .persistent() + .instance() .set(&DataKey::AllowedToken(token.clone()), &true); - env.storage() - .persistent() - .extend_ttl(&DataKey::AllowedToken(token.clone()), MATCH_TTL_LEDGERS, MATCH_TTL_LEDGERS); - env.storage().instance().set(&DataKey::AllowlistEnforced, &true); + + if !already_allowed { + let count: u32 = env + .storage() + .instance() + .get(&DataKey::AllowedTokenCount) + .unwrap_or(0); + let next_count = count.checked_add(1).ok_or(Error::Overflow)?; + env.storage() + .instance() + .set(&DataKey::AllowedTokenCount, &next_count); + if count == 0 { + env.storage().instance().set(&DataKey::AllowlistEnforced, &true); + } + } else { + env.storage().instance().set(&DataKey::AllowlistEnforced, &true); + } env.events().publish( (Symbol::new(&env, "admin"), symbol_short!("token_add")), @@ -109,12 +128,29 @@ impl EscrowContract { .ok_or(Error::Unauthorized)?; admin.require_auth(); - env.storage() - .persistent() - .remove(&DataKey::AllowedToken(token.clone())); + let was_allowed = env + .storage() + .instance() + .has(&DataKey::AllowedToken(token.clone())); + env.storage().instance().remove(&DataKey::AllowedToken(token.clone())); + + if was_allowed { + let count: u32 = env + .storage() + .instance() + .get(&DataKey::AllowedTokenCount) + .unwrap_or(0); + let next_count = count.saturating_sub(1); + env.storage() + .instance() + .set(&DataKey::AllowedTokenCount, &next_count); + if next_count == 0 { + env.storage().instance().set(&DataKey::AllowlistEnforced, &false); + } + } env.events().publish( - (Symbol::new(&env, "admin"), symbol_short!("token_remove")), + (Symbol::new(&env, "admin"), symbol_short!("token_removed")), token, ); Ok(()) diff --git a/contracts/escrow/src/tests/token_allowlist.rs b/contracts/escrow/src/tests/token_allowlist.rs index 1b19809..bafee82 100644 --- a/contracts/escrow/src/tests/token_allowlist.rs +++ b/contracts/escrow/src/tests/token_allowlist.rs @@ -55,6 +55,48 @@ fn test_removed_tokens_can_no_longer_be_used_for_new_matches() { ); } +#[test] +fn test_remove_last_allowed_token_disables_allowlist() { + let (env, contract_id, _oracle, player1, player2, token, _admin) = setup(); + let client = EscrowContractClient::new(&env, &contract_id); + + client.add_allowed_token(&token); + client.remove_allowed_token(&token); + + let other_token = Address::generate(&env); + let id = client.create_match( + &player1, + &player2, + &100, + &other_token, + &String::from_str(&env, "allowlist_disabled_game"), + &Platform::Lichess, + ); + assert_eq!(id, 0, "create_match should accept new token once the allowlist is disabled"); +} + +#[test] +fn test_remove_allowed_token_requires_admin() { + let (env, contract_id, _oracle, _player1, _player2, token, _admin) = setup(); + let client = EscrowContractClient::new(&env, &contract_id); + + client.add_allowed_token(&token); + + let attacker = Address::generate(&env); + env.mock_auths(&[MockAuth { + address: &attacker, + invoke: &MockAuthInvoke { + contract: &contract_id, + fn_name: "remove_allowed_token", + args: (token.clone(),).into_val(&env), + sub_invokes: &[], + }, + }]); + + let result = client.try_remove_allowed_token(&token); + assert_eq!(result, Err(Ok(Error::Unauthorized))); +} + #[test] fn test_multiple_approved_tokens_can_coexist_after_allowlist_enforcement_is_enabled() { let (env, contract_id, _oracle, player1, player2, token, _admin) = setup(); diff --git a/contracts/escrow/src/types.rs b/contracts/escrow/src/types.rs index 5d92733..04d4639 100644 --- a/contracts/escrow/src/types.rs +++ b/contracts/escrow/src/types.rs @@ -57,5 +57,6 @@ pub enum DataKey { MatchTimeout, AllowedToken(Address), AllowlistEnforced, + AllowedTokenCount, OracleRecord(u64), }