Fix burning issues

staking-factory/src/lib.rs

@@ -4,8 +4,8 @@ use near_sdk::json_types::{Base58CryptoHash, U128};
 use near_sdk::serde::{Deserialize, Serialize};
 use near_sdk::serde_json::json;
 use near_sdk::{
-    assert_self, env, ext_contract, is_promise_success, log, near_bindgen, sys, AccountId, Balance,
-    CryptoHash, PanicOnDefault, Promise, PromiseOrValue, PublicKey,
+    env, ext_contract, is_promise_success, log, near_bindgen, sys, AccountId, Balance, CryptoHash,
+    PanicOnDefault, Promise, PromiseOrValue, PublicKey,
 };
 
 /// The 30 NEAR tokens required for the storage of the staking pool.
@@ -200,14 +200,13 @@ impl StakingPoolFactory {
     /// Callback after a staking pool was created.
     /// Returns the promise to whitelist the staking pool contract if the pool creation succeeded.
     /// Otherwise refunds the attached deposit and returns `false`.
+    #[private]
     pub fn on_staking_pool_create(
         &mut self,
         staking_pool_account_id: AccountId,
         attached_deposit: U128,
         predecessor_account_id: AccountId,
     ) -> PromiseOrValue<bool> {
-        assert_self();
-
         let staking_pool_created = is_promise_success();
 
         if staking_pool_created {

staking-factory/tests/spec.rs

@@ -100,11 +100,12 @@ fn create_staking_pool(
             code_hash,
             user.account_id(),
             STAKING_KEY.parse().unwrap(),
-            fee,
+            fee
         ),
         deposit = to_yocto(POOL_DEPOSIT)
     )
 }
+
 pub fn should_fail(r: ExecutionResult) {
     match r.status() {
         ExecutionStatus::Failure(_) => {}

staking-farm/src/farm.rs

@@ -1,7 +1,8 @@
 use near_contract_standards::fungible_token::core_impl::ext_fungible_token;
 use near_sdk::json_types::U64;
-use near_sdk::Timestamp;
+use near_sdk::{promise_result_as_success, Timestamp};
 
+use crate::internal::ZERO_ADDRESS;
 use crate::stake::ext_self;
 use crate::*;
 
@@ -145,9 +146,10 @@ impl StakingContract {
             .get(&farm_id)
             .cloned()
             .unwrap_or(U256::zero());
-        if let Some(distribution) =
-            self.internal_calculate_distribution(&farm, self.total_stake_shares)
-        {
+        if let Some(distribution) = self.internal_calculate_distribution(
+            &farm,
+            self.total_stake_shares - self.total_burn_shares,
+        ) {
             (
                 farm.last_distribution.rps,
                 (U256::from(account.stake_shares) * (distribution.rps - user_rps) / DENOMINATOR)
@@ -159,9 +161,10 @@ impl StakingContract {
     }
 
     fn internal_distribute(&mut self, farm: &mut Farm) {
-        if let Some(distribution) =
-            self.internal_calculate_distribution(&farm, self.total_stake_shares)
-        {
+        if let Some(distribution) = self.internal_calculate_distribution(
+            &farm,
+            self.total_stake_shares - self.total_burn_shares,
+        ) {
             if distribution.reward_round != farm.last_distribution.reward_round {
                 farm.last_distribution = distribution;
             }
@@ -236,43 +239,35 @@ impl StakingContract {
 
 #[near_bindgen]
 impl StakingContract {
+    /// Callback after checking owner for the delegated claim.
+    #[private]
     pub fn callback_post_get_owner(
         &mut self,
         token_id: AccountId,
         delegator_id: AccountId,
         account_id: AccountId,
     ) -> Promise {
-        let owner_id = match env::promise_result(0) {
-            PromiseResult::Successful(result) => AccountId::new_unchecked(
-                near_sdk::serde_json::from_slice(&result).expect("Failed to parse"),
-            ),
-            _ => panic!("get_owner MUST HAVE RESULT"),
-        };
+        let owner_id: AccountId = near_sdk::serde_json::from_slice(
+            &promise_result_as_success().expect("get_owner must have result"),
+        )
+        .expect("Failed to parse");
         assert_eq!(owner_id, account_id, "Caller is not an owner");
         self.internal_claim(&token_id, &delegator_id, &account_id)
     }
 
     /// Callback from depositing funds to the user's account.
     /// If it failed, return funds to the user's account.
+    #[private]
     pub fn callback_post_withdraw_reward(
         &mut self,
         token_id: AccountId,
         sender_id: AccountId,
         amount: U128,
     ) {
-        assert_eq!(
-            env::promise_results_count(),
-            1,
-            "ERR_CALLBACK_POST_WITHDRAW_INVALID",
-        );
-        match env::promise_result(0) {
-            PromiseResult::NotReady => unreachable!(),
-            PromiseResult::Successful(_) => {}
-            PromiseResult::Failed => {
-                // This reverts the changes from the claim function.
-                self.internal_user_token_deposit(&sender_id, &token_id, amount.0);
-            }
-        };
+        if promise_result_as_success().is_none() {
+            // This reverts the changes from the claim function.
+            self.internal_user_token_deposit(&sender_id, &token_id, amount.0);
+        }
     }
 
     /// Claim given tokens for given account.
@@ -307,6 +302,9 @@ impl StakingContract {
     }
 
     pub fn get_unclaimed_reward(&self, account_id: AccountId, farm_id: u64) -> U128 {
+        if account_id == AccountId::new_unchecked(ZERO_ADDRESS.to_string()) {
+            return U128(0);
+        }
         let account = self.accounts.get(&account_id).expect("ERR_NO_ACCOUNT");
         let farm = self.farms.get(farm_id).expect("ERR_NO_FARM");
         let (_rps, reward) = self.internal_unclaimed_balance(&account, farm_id, &farm);

staking-farm/src/internal.rs

@@ -4,10 +4,10 @@ use near_sdk::log;
 
 /// Zero address is implicit address that doesn't have a key for it.
 /// Used for burning tokens.
-const ZERO_ADDRESS: &str = "0000000000000000000000000000000000000000000000000000000000000000";
+pub const ZERO_ADDRESS: &str = "0000000000000000000000000000000000000000000000000000000000000000";
 
 /// Minimum amount that will be sent to burn. This is to ensure there is enough storage on the other side.
-const MIN_BURN_AMOUNT: Balance = 1694457700619870000000;
+pub const MIN_BURN_AMOUNT: Balance = 1694457700619870000000;
 
 impl StakingContract {
     /********************/
@@ -47,10 +47,9 @@ impl StakingContract {
         amount
     }
 
-    pub(crate) fn internal_withdraw(&mut self, amount: Balance) {
+    pub(crate) fn internal_withdraw(&mut self, account_id: &AccountId, amount: Balance) {
         assert!(amount > 0, "Withdrawal amount should be positive");
 
-        let account_id = env::predecessor_account_id();
         let mut account = self.internal_get_account(&account_id);
         assert!(
             account.unstaked >= amount,
@@ -70,7 +69,7 @@ impl StakingContract {
             account.unstaked
         );
 
-        Promise::new(account_id).transfer(amount);
+        Promise::new(account_id.clone()).transfer(amount);
         self.last_total_balance -= amount;
     }
 
@@ -131,10 +130,9 @@ impl StakingContract {
         );
     }
 
-    pub(crate) fn inner_unstake(&mut self, amount: u128) {
+    pub(crate) fn inner_unstake(&mut self, account_id: &AccountId, amount: u128) {
         assert!(amount > 0, "Unstaking amount should be positive");
 
-        let account_id = env::predecessor_account_id();
         let mut account = self.internal_get_account(&account_id);
 
         // Distribute rewards from all the farms for the given user.
@@ -176,6 +174,9 @@ impl StakingContract {
 
         self.total_staked_balance -= unstake_amount;
         self.total_stake_shares -= num_shares;
+        if account_id == &AccountId::new_unchecked(ZERO_ADDRESS.to_string()) {
+            self.total_burn_shares -= num_shares;
+        }
 
         log!(
             "@{} unstaking {}. Spent {} staking shares. Total {} unstaked balance and {} \
@@ -193,6 +194,28 @@ impl StakingContract {
         );
     }
 
+    pub(crate) fn internal_unstake_all(&mut self, account_id: &AccountId) {
+        // Unstake action always restakes
+        self.internal_ping();
+
+        let account = self.internal_get_account(&account_id);
+        let amount = self.staked_amount_from_num_shares_rounded_down(account.stake_shares);
+        self.inner_unstake(account_id, amount);
+
+        self.internal_restake();
+    }
+
+    /// Add given number of staked shares to the given account.
+    fn internal_add_shares(&mut self, account_id: &AccountId, num_shares: NumStakeShares) {
+        if num_shares > 0 {
+            let mut account = self.internal_get_account(&account_id);
+            account.stake_shares += num_shares;
+            self.internal_save_account(&account_id, &account);
+            // Increasing the total amount of "stake" shares.
+            self.total_stake_shares += num_shares;
+        }
+    }
+
     /// Distributes rewards after the new epoch. It's automatically called before every action.
     /// Returns true if the current epoch height is different from the last epoch height.
     pub(crate) fn internal_ping(&mut self) -> bool {
@@ -206,7 +229,7 @@ impl StakingContract {
         // NOTE: We need to subtract `attached_deposit` in case `ping` called from `deposit` call
         // since the attached deposit gets included in the `account_balance`, and we have not
         // accounted it yet.
-        let mut total_balance =
+        let total_balance =
             env::account_locked_balance() + env::account_balance() - env::attached_deposit();
 
         assert!(
@@ -215,58 +238,48 @@ impl StakingContract {
             total_balance,
             self.last_total_balance
         );
-        let mut total_reward = total_balance - self.last_total_balance;
+        let total_reward = total_balance - self.last_total_balance;
         if total_reward > 0 {
-            // Burn fee gets computed first.
-            let mut burn_fee = self.burn_fee_fraction.multiply(total_reward);
-
-            if burn_fee > 0 {
-                // TODO: replace with burn host function when available.
-                if burn_fee < MIN_BURN_AMOUNT {
-                    burn_fee = 0
-                } else {
-                    Promise::new(AccountId::new_unchecked(ZERO_ADDRESS.to_string()))
-                        .transfer(burn_fee);
-                }
-            }
-
-            // All subsequent computations are done without part that is going to be burnt.
-            total_reward -= burn_fee;
-            total_balance -= burn_fee;
+            // The validation fee that will be burnt.
+            let burn_fee = self.burn_fee_fraction.multiply(total_reward);
 
             // The validation fee that the contract owner takes.
-            let owners_fee = self.reward_fee_fraction.multiply(total_reward);
+            let owners_fee = self.reward_fee_fraction.multiply(total_reward - burn_fee);
 
             // Distributing the remaining reward to the delegators first.
-            let remaining_reward = total_reward - owners_fee;
+            let remaining_reward = total_reward - owners_fee - burn_fee;
             self.total_staked_balance += remaining_reward;
 
+            // Now buying "stake" shares for the burn.
+            let num_burn_shares = self.num_shares_from_staked_amount_rounded_down(burn_fee);
+            self.internal_add_shares(
+                &AccountId::new_unchecked(ZERO_ADDRESS.to_string()),
+                num_burn_shares,
+            );
+            self.total_burn_shares += num_burn_shares;
+
             // Now buying "stake" shares for the contract owner at the new share price.
-            let num_shares = self.num_shares_from_staked_amount_rounded_down(owners_fee);
-            if num_shares > 0 {
-                // Updating owner's inner account
-                let owner_id = StakingContract::get_owner_id();
-                let mut account = self.internal_get_account(&owner_id);
-                account.stake_shares += num_shares;
-                self.internal_save_account(&owner_id, &account);
-                // Increasing the total amount of "stake" shares.
-                self.total_stake_shares += num_shares;
-            }
+            let num_owner_shares = self.num_shares_from_staked_amount_rounded_down(owners_fee);
+            self.internal_add_shares(&StakingContract::get_owner_id(), num_owner_shares);
+
             // Increasing the total staked balance by the owners fee, no matter whether the owner
             // received any shares or not.
-            self.total_staked_balance += owners_fee;
+            self.total_staked_balance += owners_fee + burn_fee;
 
             log!(
-                "Epoch {}: Contract received total rewards of {} tokens and {} burnt. \
+                "Epoch {}: Contract received total rewards of {} tokens. \
                  New total staked balance is {}. Total number of shares {}",
                 epoch_height,
                 total_reward,
-                burn_fee,
                 self.total_staked_balance,
                 self.total_stake_shares,
             );
-            if num_shares > 0 {
-                log!("Total rewards fee is {} stake shares.", num_shares);
+            if num_owner_shares > 0 || num_burn_shares > 0 {
+                log!(
+                    "Total rewards fee is {} and burn is {} stake shares.",
+                    num_owner_shares,
+                    num_burn_shares
+                );
             }
         }
 

staking-farm/src/lib.rs

@@ -45,6 +45,17 @@ construct_uint! {
 /// updated in the previous epoch. It will not unlock the funds for 4 epochs.
 const NUM_EPOCHS_TO_UNLOCK: EpochHeight = 4;
 
+/// Tracking balance for burning.
+#[derive(BorshDeserialize, BorshSerialize)]
+pub struct BurnInfo {
+    /// The unstaked balance that can be burnt.
+    pub unstaked: Balance,
+    /// Number of "stake" shares that must be burnt.
+    pub stake_shares: Balance,
+    /// The minimum epoch height when the burn is allowed.
+    pub unstaked_available_epoch_height: EpochHeight,
+}
+
 #[near_bindgen]
 #[derive(BorshDeserialize, BorshSerialize)]
 pub struct StakingContract {
@@ -60,6 +71,9 @@ pub struct StakingContract {
     pub total_stake_shares: NumStakeShares,
     /// The total staked balance.
     pub total_staked_balance: Balance,
+    /// The total burn share balance, that will not be accounted in the farming.
+    pub total_burn_shares: NumStakeShares,
+    /// The total amount to burn that will be available
     /// The fraction of the reward that goes to the owner of the staking pool for running the
     /// validator node.
     pub reward_fee_fraction: Ratio,
@@ -147,6 +161,7 @@ impl StakingContract {
             last_total_balance: account_balance,
             total_staked_balance,
             total_stake_shares: NumStakeShares::from(total_staked_balance),
+            total_burn_shares: 0,
             reward_fee_fraction,
             burn_fee_fraction,
             accounts: UnorderedMap::new(b"u".to_vec()),