Require payments for storage increase and refund difference

examples/fungible-token/Cargo.toml

@@ -21,6 +21,7 @@ opt-level = "z"
 lto = true
 debug = false
 panic = "abort"
+overflow-checks = true
 
 [workspace]
 members = []

examples/fungible-token/src/lib.rs

@@ -14,11 +14,14 @@
 use borsh::{BorshDeserialize, BorshSerialize};
 use near_sdk::collections::UnorderedMap;
 use near_sdk::json_types::U128;
-use near_sdk::{env, near_bindgen, AccountId, Balance};
+use near_sdk::{env, near_bindgen, AccountId, Balance, Promise, StorageUsage};
 
 #[global_allocator]
 static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;
 
+/// Price per 1 byte of storage from mainnet genesis config.
+const STORAGE_PRICE_PER_BYTE: Balance = 100000000000000000000;
+
 /// Contains balance and allowances information for one account.
 #[derive(BorshDeserialize, BorshSerialize)]
 pub struct Account {
@@ -87,20 +90,26 @@ impl FungibleToken {
 
     /// Sets the `allowance` for `escrow_account_id` on the account of the caller of this contract
     /// (`predecessor_id`) who is the balance owner.
+    /// Requirements:
+    /// * Caller of the method has to attach deposit enough to cover storage difference at the
+    ///   fixed storage price defined in the contract.
+    #[payable]
     pub fn set_allowance(&mut self, escrow_account_id: AccountId, allowance: U128) {
+        let initial_storage = env::storage_usage();
         assert!(
             env::is_valid_account_id(escrow_account_id.as_bytes()),
             "Escrow account ID is invalid"
         );
         let allowance = allowance.into();
         let owner_id = env::predecessor_account_id();
         if escrow_account_id == owner_id {
-            env::panic(b"Can't set allowance for yourself");
+            env::panic(b"Can not set allowance for yourself");
         }
         let mut account = self.get_account(&owner_id);
 
         account.set_allowance(&escrow_account_id, allowance);
         self.set_account(&owner_id, &account);
+        self.refund_storage(initial_storage);
     }
 
     /// Transfers the `amount` of tokens from `owner_id` to the `new_owner_id`.
@@ -110,7 +119,11 @@ impl FungibleToken {
     /// * If this function is called by an escrow account (`owner_id != predecessor_account_id`),
     ///   then the allowance of the caller of the function (`predecessor_account_id`) on
     ///   the account of `owner_id` should be greater or equal than the transfer `amount`.
+    /// * Caller of the method has to attach deposit enough to cover storage difference at the
+    ///   fixed storage price defined in the contract.
+    #[payable]
     pub fn transfer_from(&mut self, owner_id: AccountId, new_owner_id: AccountId, amount: U128) {
+        let initial_storage = env::storage_usage();
         assert!(env::is_valid_account_id(owner_id.as_bytes()), "Owner's account ID is invalid");
         assert!(
             env::is_valid_account_id(new_owner_id.as_bytes()),
@@ -146,14 +159,20 @@ impl FungibleToken {
         let mut new_account = self.get_account(&new_owner_id);
         new_account.balance += amount;
         self.set_account(&new_owner_id, &new_account);
+        self.refund_storage(initial_storage);
     }
 
     /// Transfer `amount` of tokens from the caller of the contract (`predecessor_id`) to
     /// `new_owner_id`.
     /// Act the same was as `transfer_from` with `owner_id` equal to the caller of the contract
     /// (`predecessor_id`).
+    /// Requirements:
+    /// * Caller of the method has to attach deposit enough to cover storage difference at the
+    ///   fixed storage price defined in the contract.
+    #[payable]
     pub fn transfer(&mut self, new_owner_id: AccountId, amount: U128) {
-        // NOTE: New owner's Account ID checked in transfer_from
+        // NOTE: New owner's Account ID checked in transfer_from.
+        // Storage fees are also refunded in transfer_from.
         self.transfer_from(env::predecessor_account_id(), new_owner_id, amount);
     }
 
@@ -194,6 +213,29 @@ impl FungibleToken {
         let account_hash = env::sha256(owner_id.as_bytes());
         self.accounts.insert(&account_hash, &account);
     }
+
+    fn refund_storage(&self, initial_storage: StorageUsage) {
+        let current_storage = env::storage_usage();
+        let attached_deposit = env::attached_deposit();
+        let refund_amount = if current_storage > initial_storage {
+            let required_deposit =
+                Balance::from(current_storage - initial_storage) * STORAGE_PRICE_PER_BYTE;
+            assert!(
+                required_deposit <= attached_deposit,
+                "The required attached deposit is {}, but the given attached deposit is is {}",
+                required_deposit,
+                attached_deposit,
+            );
+            attached_deposit - required_deposit
+        } else {
+            attached_deposit
+                + Balance::from(initial_storage - current_storage) * STORAGE_PRICE_PER_BYTE
+        };
+        if refund_amount > 0 {
+            env::log(format!("Refunding {} tokens for storage", refund_amount).as_bytes());
+            Promise::new(env::predecessor_account_id()).transfer(refund_amount);
+        }
+    }
 }
 
 #[cfg(not(target_arch = "wasm32"))]
@@ -214,16 +256,6 @@ mod tests {
         "carol.near".to_string()
     }
 
-    fn catch_unwind_silent<F: FnOnce() -> R + std::panic::UnwindSafe, R>(
-        f: F,
-    ) -> std::thread::Result<R> {
-        let prev_hook = std::panic::take_hook();
-        std::panic::set_hook(Box::new(|_| {}));
-        let result = std::panic::catch_unwind(f);
-        std::panic::set_hook(prev_hook);
-        result
-    }
-
     fn get_context(predecessor_account_id: AccountId) -> VMContext {
         VMContext {
             current_account_id: alice(),
@@ -233,7 +265,7 @@ mod tests {
             input: vec![],
             block_index: 0,
             block_timestamp: 0,
-            account_balance: 0,
+            account_balance: 1_000_000_000_000_000_000_000_000_000u128,
             account_locked_balance: 0,
             storage_usage: 10u64.pow(6),
             attached_deposit: 0,
@@ -256,55 +288,104 @@ mod tests {
     }
 
     #[test]
+    #[should_panic]
     fn test_new_twice_fails() {
         let context = get_context(carol());
         testing_env!(context);
         let total_supply = 1_000_000_000_000_000u128;
-        let _contract = FungibleToken::new(bob(), total_supply.into());
-        catch_unwind_silent(|| {
-            FungibleToken::new(bob(), total_supply.into());
-        })
-        .unwrap_err();
+        {
+            let _contract = FungibleToken::new(bob(), total_supply.into());
+        }
+        FungibleToken::new(bob(), total_supply.into());
     }
 
     #[test]
     fn test_transfer() {
-        let context = get_context(carol());
-        testing_env!(context);
+        let mut context = get_context(carol());
+        testing_env!(context.clone());
         let total_supply = 1_000_000_000_000_000u128;
         let mut contract = FungibleToken::new(carol(), total_supply.into());
+        context.storage_usage = env::storage_usage();
+
+        context.attached_deposit = 1000 * STORAGE_PRICE_PER_BYTE;
+        testing_env!(context.clone());
         let transfer_amount = total_supply / 3;
         contract.transfer(bob(), transfer_amount.into());
+        context.storage_usage = env::storage_usage();
+        context.account_balance = env::account_balance();
+
+        context.is_view = true;
+        context.attached_deposit = 0;
+        testing_env!(context.clone());
         assert_eq!(contract.get_balance(carol()).0, (total_supply - transfer_amount));
         assert_eq!(contract.get_balance(bob()).0, transfer_amount);
     }
 
     #[test]
+    #[should_panic(expected = "Can not set allowance for yourself")]
     fn test_self_allowance_fail() {
-        let context = get_context(carol());
-        testing_env!(context);
+        let mut context = get_context(carol());
+        testing_env!(context.clone());
+        let total_supply = 1_000_000_000_000_000u128;
+        let mut contract = FungibleToken::new(carol(), total_supply.into());
+        context.attached_deposit = STORAGE_PRICE_PER_BYTE * 1000;
+        testing_env!(context.clone());
+        contract.set_allowance(carol(), (total_supply / 2).into());
+    }
+
+    #[test]
+    #[should_panic(
+        expected = "The required attached deposit is 33100000000000000000000, but the given attached deposit is is 0"
+    )]
+    fn test_self_allowance_fail_no_deposit() {
+        let mut context = get_context(carol());
+        testing_env!(context.clone());
         let total_supply = 1_000_000_000_000_000u128;
         let mut contract = FungibleToken::new(carol(), total_supply.into());
-        catch_unwind_silent(move || {
-            contract.set_allowance(carol(), (total_supply / 2).into());
-        })
-        .unwrap_err();
+        context.attached_deposit = 0;
+        testing_env!(context.clone());
+        contract.set_allowance(bob(), (total_supply / 2).into());
     }
 
     #[test]
     fn test_carol_escrows_to_bob_transfers_to_alice() {
         // Acting as carol
-        testing_env!(get_context(carol()));
+        let mut context = get_context(carol());
+        testing_env!(context.clone());
         let total_supply = 1_000_000_000_000_000u128;
         let mut contract = FungibleToken::new(carol(), total_supply.into());
+        context.storage_usage = env::storage_usage();
+
+        context.is_view = true;
+        testing_env!(context.clone());
         assert_eq!(contract.get_total_supply().0, total_supply);
+
         let allowance = total_supply / 3;
         let transfer_amount = allowance / 3;
+        context.is_view = false;
+        context.attached_deposit = STORAGE_PRICE_PER_BYTE * 1000;
+        testing_env!(context.clone());
         contract.set_allowance(bob(), allowance.into());
+        context.storage_usage = env::storage_usage();
+        context.account_balance = env::account_balance();
+
+        context.is_view = true;
+        context.attached_deposit = 0;
+        testing_env!(context.clone());
         assert_eq!(contract.get_allowance(carol(), bob()).0, allowance);
+
         // Acting as bob now
-        testing_env!(get_context(bob()));
+        context.is_view = false;
+        context.attached_deposit = STORAGE_PRICE_PER_BYTE * 1000;
+        context.predecessor_account_id = bob();
+        testing_env!(context.clone());
         contract.transfer_from(carol(), alice(), transfer_amount.into());
+        context.storage_usage = env::storage_usage();
+        context.account_balance = env::account_balance();
+
+        context.is_view = true;
+        context.attached_deposit = 0;
+        testing_env!(context.clone());
         assert_eq!(contract.get_balance(carol()).0, total_supply - transfer_amount);
         assert_eq!(contract.get_balance(alice()).0, transfer_amount);
         assert_eq!(contract.get_allowance(carol(), bob()).0, allowance - transfer_amount);
@@ -313,20 +394,83 @@ mod tests {
     #[test]
     fn test_carol_escrows_to_bob_locks_and_transfers_to_alice() {
         // Acting as carol
-        testing_env!(get_context(carol()));
+        let mut context = get_context(carol());
+        testing_env!(context.clone());
         let total_supply = 1_000_000_000_000_000u128;
         let mut contract = FungibleToken::new(carol(), total_supply.into());
+        context.storage_usage = env::storage_usage();
+
+        context.is_view = true;
+        testing_env!(context.clone());
         assert_eq!(contract.get_total_supply().0, total_supply);
+
         let allowance = total_supply / 3;
         let transfer_amount = allowance / 3;
+        context.is_view = false;
+        context.attached_deposit = STORAGE_PRICE_PER_BYTE * 1000;
+        testing_env!(context.clone());
         contract.set_allowance(bob(), allowance.into());
+        context.storage_usage = env::storage_usage();
+        context.account_balance = env::account_balance();
+
+        context.is_view = true;
+        context.attached_deposit = 0;
+        testing_env!(context.clone());
         assert_eq!(contract.get_allowance(carol(), bob()).0, allowance);
-        // Acting as bob now
-        testing_env!(get_context(bob()));
         assert_eq!(contract.get_balance(carol()).0, total_supply);
+
+        // Acting as bob now
+        context.is_view = false;
+        context.attached_deposit = STORAGE_PRICE_PER_BYTE * 1000;
+        context.predecessor_account_id = bob();
+        testing_env!(context.clone());
         contract.transfer_from(carol(), alice(), transfer_amount.into());
+        context.storage_usage = env::storage_usage();
+        context.account_balance = env::account_balance();
+
+        context.is_view = true;
+        context.attached_deposit = 0;
+        testing_env!(context.clone());
         assert_eq!(contract.get_balance(carol()).0, (total_supply - transfer_amount));
         assert_eq!(contract.get_balance(alice()).0, transfer_amount);
         assert_eq!(contract.get_allowance(carol(), bob()).0, allowance - transfer_amount);
     }
+
+    #[test]
+    fn test_self_allowance_set_for_refund() {
+        let mut context = get_context(carol());
+        testing_env!(context.clone());
+        let total_supply = 1_000_000_000_000_000u128;
+        let mut contract = FungibleToken::new(carol(), total_supply.into());
+        context.storage_usage = env::storage_usage();
+
+        let initial_balance = context.account_balance;
+        let initial_storage = context.storage_usage;
+        context.attached_deposit = STORAGE_PRICE_PER_BYTE * 1000;
+        testing_env!(context.clone());
+        contract.set_allowance(bob(), (total_supply / 2).into());
+        context.storage_usage = env::storage_usage();
+        context.account_balance = env::account_balance();
+        assert_eq!(
+            context.account_balance,
+            initial_balance
+                + Balance::from(context.storage_usage - initial_storage) * STORAGE_PRICE_PER_BYTE
+        );
+
+        let initial_balance = context.account_balance;
+        let initial_storage = context.storage_usage;
+        testing_env!(context.clone());
+        context.attached_deposit = 0;
+        testing_env!(context.clone());
+        contract.set_allowance(bob(), 0.into());
+        context.storage_usage = env::storage_usage();
+        context.account_balance = env::account_balance();
+        assert!(context.storage_usage < initial_storage);
+        assert!(context.account_balance < initial_balance);
+        assert_eq!(
+            context.account_balance,
+            initial_balance
+                - Balance::from(initial_storage - context.storage_usage) * STORAGE_PRICE_PER_BYTE
+        );
+    }
 }