feat: add remove-allowed-token

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(())

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();

contracts/escrow/src/types.rs

@@ -57,5 +57,6 @@ pub enum DataKey {
     MatchTimeout,
     AllowedToken(Address),
     AllowlistEnforced,
+    AllowedTokenCount,
     OracleRecord(u64),
 }