refactor(interchain-token): use contractstorage

Author: nbayindirli

contracts/stellar-interchain-token/src/contract.rs

@@ -5,13 +5,13 @@ use soroban_token_sdk::metadata::TokenMetadata;
 use soroban_token_sdk::TokenUtils;
 use stellar_axelar_std::events::Event;
 use stellar_axelar_std::interfaces::OwnableInterface;
-use stellar_axelar_std::ttl::{extend_instance_ttl, extend_persistent_ttl};
+use stellar_axelar_std::ttl::extend_instance_ttl;
 use stellar_axelar_std::{ensure, interfaces, Upgradable};
 
 use crate::error::ContractError;
 use crate::event::{MinterAddedEvent, MinterRemovedEvent};
 use crate::interface::InterchainTokenInterface;
-use crate::storage_types::{AllowanceDataKey, AllowanceValue, DataKey};
+use crate::storage::{self, AllowanceDataKey, AllowanceValue};
 
 #[contract]
 #[derive(Upgradable)]
@@ -30,12 +30,10 @@ impl InterchainToken {
 
         Self::write_metadata(&env, token_metadata);
 
-        env.storage().instance().set(&DataKey::TokenId, &token_id);
+        storage::set_token_id(&env, &token_id);
 
         if let Some(minter) = minter {
-            env.storage()
-                .instance()
-                .set(&DataKey::Minter(minter.clone()), &());
+            storage::set_minter_status(&env, minter.clone());
 
             MinterAddedEvent { minter }.emit(&env);
         }
@@ -96,14 +94,11 @@ impl StellarAssetInterface for InterchainToken {
 #[contractimpl]
 impl InterchainTokenInterface for InterchainToken {
     fn token_id(env: &Env) -> BytesN<32> {
-        env.storage()
-            .instance()
-            .get(&DataKey::TokenId)
-            .expect("token id not found")
+        storage::token_id(env)
     }
 
     fn is_minter(env: &Env, minter: Address) -> bool {
-        env.storage().instance().has(&DataKey::Minter(minter))
+        storage::is_minter(env, minter)
     }
 
     fn mint_from(
@@ -133,23 +128,15 @@ impl InterchainTokenInterface for InterchainToken {
     fn add_minter(env: &Env, minter: Address) {
         Self::owner(env).require_auth();
 
-        env.storage()
-            .instance()
-            .set(&DataKey::Minter(minter.clone()), &());
-
-        extend_instance_ttl(env);
+        storage::set_minter_status(env, minter.clone());
 
         MinterAddedEvent { minter }.emit(env);
     }
 
     fn remove_minter(env: &Env, minter: Address) {
         Self::owner(env).require_auth();
 
-        env.storage()
-            .instance()
-            .remove(&DataKey::Minter(minter.clone()));
-
-        extend_instance_ttl(env);
+        storage::remove_minter_status(env, minter.clone());
 
         MinterRemovedEvent { minter }.emit(env);
     }
@@ -183,8 +170,7 @@ impl token::Interface for InterchainToken {
     }
 
     fn balance(env: Env, id: Address) -> i128 {
-        extend_instance_ttl(&env);
-        Self::read_balance(&env, id)
+        storage::try_balance(&env, id).unwrap_or_default()
     }
 
     fn transfer(env: Env, from: Address, to: Address, amount: i128) {
@@ -257,26 +243,23 @@ impl InterchainToken {
     }
 
     fn read_allowance(env: &Env, from: Address, spender: Address) -> AllowanceValue {
-        let key = DataKey::Allowance(AllowanceDataKey { from, spender });
-        env.storage()
-            .temporary()
-            .get::<_, AllowanceValue>(&key)
-            .map_or(
-                AllowanceValue {
-                    amount: 0,
-                    expiration_ledger: 0,
-                },
-                |allowance| {
-                    if allowance.expiration_ledger < env.ledger().sequence() {
-                        AllowanceValue {
-                            amount: 0,
-                            expiration_ledger: allowance.expiration_ledger,
-                        }
-                    } else {
-                        allowance
+        let key = AllowanceDataKey { from, spender };
+        storage::try_allowance(env, key).map_or(
+            AllowanceValue {
+                amount: 0,
+                expiration_ledger: 0,
+            },
+            |allowance| {
+                if allowance.expiration_ledger < env.ledger().sequence() {
+                    AllowanceValue {
+                        amount: 0,
+                        expiration_ledger: allowance.expiration_ledger,
                     }
-                },
-            )
+                } else {
+                    allowance
+                }
+            },
+        )
     }
 
     fn write_allowance(
@@ -297,17 +280,15 @@ impl InterchainToken {
             ContractError::InvalidExpirationLedger
         );
 
-        let key = DataKey::Allowance(AllowanceDataKey { from, spender });
-        env.storage().temporary().set(&key, &allowance);
+        let key = AllowanceDataKey { from, spender };
+        storage::set_allowance(env, key.clone(), &allowance);
 
         if amount > 0 {
             let live_for = expiration_ledger
                 .checked_sub(env.ledger().sequence())
                 .unwrap();
 
-            env.storage()
-                .temporary()
-                .extend_ttl(&key, live_for, live_for)
+            storage::extend_allowance_ttl(env, key, live_for, live_for);
         }
     }
 
@@ -334,45 +315,21 @@ impl InterchainToken {
         }
     }
 
-    fn read_balance(env: &Env, addr: Address) -> i128 {
-        let key = DataKey::Balance(addr);
-        env.storage()
-            .persistent()
-            .get::<_, i128>(&key)
-            .inspect(|_| {
-                // Extend the TTL of the balance entry when the balance is successfully retrieved.
-                extend_persistent_ttl(env, &key);
-            })
-            .unwrap_or_default()
-    }
-
     fn receive_balance(env: &Env, addr: Address, amount: i128) {
-        let key = DataKey::Balance(addr);
+        let current_balance = storage::try_balance(env, addr.clone()).unwrap_or_default();
 
-        env.storage()
-            .persistent()
-            .update(&key, |balance: Option<i128>| {
-                balance.unwrap_or_default() + amount
-            });
+        storage::set_balance(env, addr, &(current_balance + amount));
     }
 
     fn spend_balance(env: &Env, addr: Address, amount: i128) {
-        let balance = Self::read_balance(env, addr.clone());
+        let balance = storage::try_balance(env, addr.clone()).unwrap_or_default();
 
         assert_with_error!(env, balance >= amount, ContractError::InsufficientBalance);
 
-        Self::write_balance(env, addr, balance - amount);
+        storage::set_balance(env, addr, &(balance - amount));
     }
 
     fn write_metadata(env: &Env, metadata: TokenMetadata) {
         TokenUtils::new(env).metadata().set_metadata(&metadata);
     }
-
-    fn write_balance(env: &Env, addr: Address, amount: i128) {
-        let key = DataKey::Balance(addr);
-
-        env.storage().persistent().set(&key, &amount);
-
-        extend_persistent_ttl(env, &key);
-    }
 }

contracts/stellar-interchain-token/src/lib.rs

@@ -15,7 +15,7 @@ cfg_if::cfg_if! {
         pub use interface::{InterchainTokenClient, InterchainTokenInterface};
     } else {
         pub mod event;
-        mod storage_types;
+        mod storage;
         mod contract;
 
         pub use contract::{InterchainToken, InterchainTokenClient};

contracts/stellar-interchain-token/src/storage.rs

@@ -0,0 +1,36 @@
+use soroban_sdk::{contracttype, Address, BytesN};
+use stellar_axelar_std::contractstorage;
+
+#[derive(Clone)]
+#[contracttype]
+pub struct AllowanceDataKey {
+    pub from: Address,
+    pub spender: Address,
+}
+
+#[contracttype]
+pub struct AllowanceValue {
+    pub amount: i128,
+    pub expiration_ledger: u32,
+}
+
+/// Do not use the symbol `METADATA` as a key as it is reserved for token metadata.
+#[contractstorage]
+#[derive(Clone)]
+pub enum DataKey {
+    #[temporary]
+    #[value(AllowanceValue)]
+    Allowance { key: AllowanceDataKey },
+
+    #[persistent]
+    #[value(i128)]
+    Balance { address: Address },
+
+    #[instance]
+    #[status]
+    Minter { minter: Address },
+
+    #[instance]
+    #[value(BytesN<32>)]
+    TokenId,
+}

contracts/stellar-interchain-token/src/storage_types.rs

@@ -0,0 +1,20 @@
+/// Do not use the symbol `METADATA` as a key as it is reserved for token metadata.
+#[derive(Clone)]
+pub enum DataKey {
+
+    #[temporary]
+    #[value(AllowanceValue)]
+    Allowance { key: AllowanceDataKey },
+
+    #[persistent]
+    #[value(i128)]
+    Balance { address: Address },
+
+    #[instance]
+    #[status]
+    Minter { minter: Address },
+
+    #[instance]
+    #[value(BytesN<32>)]
+    TokenId,
+}

contracts/stellar-interchain-token/src/testdata/ensure_data_key_storage_schema_is_unchanged.golden

@@ -570,3 +570,87 @@ fn clawback_fails() {
 
     token.clawback(&token.owner(), &1);
 }
+
+#[test]
+fn allowance_returns_zero_when_expired() {
+    let env = Env::default();
+
+    let user1 = Address::generate(&env);
+    let user2 = Address::generate(&env);
+    let amount = 1000;
+
+    let (token, _) = setup_token(&env);
+
+    // Set current ledger to 100
+    let current_ledger = 100;
+    env.ledger().set_sequence_number(current_ledger);
+
+    // Set allowance to expire at ledger 200
+    let expiration_ledger = 200;
+    assert_auth!(
+        user1,
+        token.approve(&user1, &user2, &amount, &expiration_ledger)
+    );
+    assert_eq!(token.allowance(&user1, &user2), amount);
+
+    // Move to ledger after expiration
+    env.ledger().set_sequence_number(expiration_ledger + 1);
+
+    // Allowance should be 0 after expiration
+    assert_eq!(token.allowance(&user1, &user2), 0);
+}
+
+#[test]
+#[should_panic(expected = "HostError: Error(Contract, #7)")] // InvalidExpirationLedger
+fn approve_fails_with_expired_ledger() {
+    let env = Env::default();
+
+    let user1 = Address::generate(&env);
+    let user2 = Address::generate(&env);
+    let amount = 1000;
+
+    let (token, _) = setup_token(&env);
+
+    // Set current ledger to 100
+    let current_ledger = 100;
+    env.ledger().set_sequence_number(current_ledger);
+
+    // Try to set allowance with already expired ledger (before current)
+    let expired_ledger = current_ledger - 1;
+
+    token
+        .mock_all_auths()
+        .approve(&user1, &user2, &amount, &expired_ledger);
+}
+
+#[test]
+#[should_panic(expected = "HostError: Error(Contract, #7)")] // InvalidExpirationLedger
+fn allowance_preserves_expiration_when_expired() {
+    let env = Env::default();
+    let user1 = Address::generate(&env);
+    let user2 = Address::generate(&env);
+    let amount = 1000;
+    let (token, _) = setup_token(&env);
+
+    // Set current ledger to 100
+    let current_ledger = 100;
+    env.ledger().set_sequence_number(current_ledger);
+
+    // Set allowance to expire at ledger 200
+    let expiration_ledger = current_ledger + 100;
+    assert_auth!(
+        user1,
+        token.approve(&user1, &user2, &amount, &expiration_ledger)
+    );
+
+    // Move past expiration
+    env.ledger().set_sequence_number(expiration_ledger + 1);
+
+    // First check returns 0
+    assert_eq!(token.allowance(&user1, &user2), 0);
+
+    // Try to set new allowance with the same expired ledger - should fail
+    token
+        .mock_all_auths()
+        .approve(&user1, &user2, &amount, &expiration_ledger);
+}