Ref farming internal audit (#40)

* change mft token_id to using ':' as leading char * change user_rps stored in LookupMap * change farms into UnorderedMap Co-authored-by: marco_mac <[email protected]> Co-authored-by: Daniel Wang <[email protected]> Co-authored-by: Evgeny Kuzyakov <[email protected]>
ref-finance · Sep 1, 2021 · 2ecbe5c · 2ecbe5c
1 parent a196f4a
commit 2ecbe5c
Show file tree

Hide file tree

Showing 30 changed files with 3,084 additions and 1,134 deletions.
diff --git a/.gitignore b/.gitignore
@@ -14,3 +14,4 @@ Cargo.lock
 neardev
 
 res/ref_exchange_local.wasm
+res/ref_farming_local.wasm
diff --git a/ref-farming/Cargo.toml b/ref-farming/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ref_farming"
-version = "0.5.2"
+version = "1.0.0"
 authors = ["Marco Sun <[email protected]>"]
 edition = "2018"
 

diff --git a/ref-farming/build_local.sh b/ref-farming/build_local.sh
diff --git a/ref-farming/src/actions_of_farm.rs b/ref-farming/src/actions_of_farm.rs
@@ -2,7 +2,7 @@
 use near_sdk::{env, near_bindgen, Promise};
 use near_sdk::json_types::{U128};
 use simple_farm::{SimpleFarm, HRSimpleFarmTerms};
-use crate::utils::{gen_farm_id, MIN_SEED_DEPOSIT};
+use crate::utils::{gen_farm_id, MIN_SEED_DEPOSIT, parse_farm_id};
 use crate::errors::*;
 use crate::*;
 
@@ -72,33 +72,33 @@ impl Contract {
             farm_id.clone(),
             terms.into(),
         ));
-
-        // farm_seed.get_ref_mut().farms.push(farm);
-        farm_seed.get_ref_mut().farms.insert(farm_id.clone(), farm);
+
+        farm_seed.get_ref_mut().farms.insert(farm_id.clone());
         farm_seed.get_ref_mut().next_index += 1;
         self.data_mut().seeds.insert(&terms.seed_id, &farm_seed);
-
-        self.data_mut().farm_count += 1;
+        self.data_mut().farms.insert(&farm_id.clone(), &farm);
         farm_id
     }
 
-    /// when farm's status is Ended, and unclaimed reward is 0, 
-    /// the farm can be remove to reduce storage usage
-    pub(crate) fn internal_remove_farm(&mut self, seed_id: &SeedId) {
-
-        let mut farm_seed = self.get_seed(&seed_id);
-        let mut removable_farms: Vec<String> = vec![];
-        for farm in farm_seed.get_ref().farms.values() {
-            if farm.can_be_removed() {
-                removable_farms.push(farm.get_farm_id());
+    pub(crate) fn internal_remove_farm_by_farm_id(&mut self, farm_id: &FarmId) -> bool {
+        let (seed_id, _) = parse_farm_id(farm_id);
+        let mut removable = false;
+        if let Some(mut farm_seed) = self.get_seed_wrapped(&seed_id) {
+            let seed_amount = farm_seed.get_ref().amount;
+            if let Some(farm) = self.data().farms.get(farm_id) {
+                if farm.can_be_removed(&seed_amount) {
+                    removable = true;
+                }
+            }
+            if removable {
+                let mut farm = self.data_mut().farms.remove(farm_id).expect(ERR41_FARM_NOT_EXIST);
+                farm.move_to_clear(&seed_amount);
+                self.data_mut().outdated_farms.insert(farm_id, &farm);
+                farm_seed.get_ref_mut().farms.remove(farm_id);
+                self.data_mut().seeds.insert(&seed_id, &farm_seed);
+                return true;
             }
         }
-        for farm_id in &removable_farms {
-            farm_seed.get_ref_mut().farms.remove(farm_id);
-        }
-        if removable_farms.len() > 0 {
-            self.data_mut().seeds.insert(&seed_id, &farm_seed);
-        }
-
+        false
     }
 }
diff --git a/ref-farming/src/actions_of_reward.rs b/ref-farming/src/actions_of_reward.rs
@@ -16,18 +16,32 @@ construct_uint! {
 #[near_bindgen]
 impl Contract {
 
+    /// Clean invalid rps,
+    /// return false if the rps is still valid.
+    pub fn remove_user_rps_by_farm(&mut self, farm_id: FarmId) -> bool {
+        let sender_id = env::predecessor_account_id();
+        let mut farmer = self.get_farmer(&sender_id);
+        let (seed_id, _) = parse_farm_id(&farm_id);
+        let farm_seed = self.get_seed(&seed_id);
+        if !farm_seed.get_ref().farms.contains(&farm_id) {
+            farmer.get_ref_mut().remove_rps(&farm_id);
+            self.data_mut().farmers.insert(&sender_id, &farmer);
+            true
+        } else {
+            false
+        }
+    }
+
     pub fn claim_reward_by_farm(&mut self, farm_id: FarmId) {
         let sender_id = env::predecessor_account_id();
-        self.assert_storage_usage(&sender_id);
-        self.remove_unused_rps(&sender_id);
         self.internal_claim_user_reward_by_farm_id(&sender_id, &farm_id);
+        self.assert_storage_usage(&sender_id);
     }
 
     pub fn claim_reward_by_seed(&mut self, seed_id: SeedId) {
         let sender_id = env::predecessor_account_id();
-        self.assert_storage_usage(&sender_id);
-        self.remove_unused_rps(&sender_id);
         self.internal_claim_user_reward_by_seed_id(&sender_id, &seed_id);
+        self.assert_storage_usage(&sender_id);
     }
 
     /// Withdraws given reward token of given user.
@@ -39,9 +53,6 @@ impl Contract {
         let amount: u128 = amount.unwrap_or(U128(0)).into(); 
 
         let sender_id = env::predecessor_account_id();
-        self.assert_storage_usage(&sender_id);
-
-        self.remove_unused_rps(&sender_id);
 
         let mut farmer = self.get_farmer(&sender_id);
 
@@ -143,40 +154,22 @@ fn claim_user_reward_from_farm(
 
 impl Contract {
 
-    /// When a farm is removed, each farmer can free storage used by user_rps,
-    /// which takes a key (farm_id as max 64 bytes) and a U256 (user_rps as 32 bytes),
-    /// this clean method would be invoked in claim reward interface called by farmer.
-    pub(crate) fn remove_unused_rps(&mut self, sender_id: &AccountId) {
-        let mut farmer = self.get_farmer(sender_id);
-        let mut changed = false;
-        let farm_ids: Vec<String> = farmer.get_ref().user_rps.keys().map(|x| x.clone()).collect();
-        for farm_id in farm_ids {
-            let (seed_id, _) = parse_farm_id(&farm_id);
-            let farm_seed = self.get_seed(&seed_id);
-            if !farm_seed.get_ref().farms.contains_key(&farm_id) {
-                farmer.get_ref_mut().user_rps.remove(&farm_id);
-                changed = true;
-            }
-        }
-        if changed {
-            self.data_mut().farmers.insert(sender_id, &farmer);
-        }
-    }
-
     pub(crate) fn internal_claim_user_reward_by_seed_id(
         &mut self, 
         sender_id: &AccountId,
         seed_id: &SeedId) {
         let mut farmer = self.get_farmer(sender_id);
         if let Some(mut farm_seed) = self.get_seed_wrapped(seed_id) {
             let amount = farm_seed.get_ref().amount;
-            for farm in &mut farm_seed.get_ref_mut().farms.values_mut() {
+            for farm_id in &mut farm_seed.get_ref_mut().farms.iter() {
+                let mut farm = self.data().farms.get(farm_id).unwrap();
                 claim_user_reward_from_farm(
-                    farm, 
+                    &mut farm, 
                     farmer.get_ref_mut(),  
                     &amount,
                     true,
                 );
+                self.data_mut().farms.insert(farm_id, &farm);
             }
             self.data_mut().seeds.insert(seed_id, &farm_seed);
             self.data_mut().farmers.insert(sender_id, &farmer);
@@ -191,16 +184,16 @@ impl Contract {
 
         let (seed_id, _) = parse_farm_id(farm_id);
 
-        if let Some(mut farm_seed) = self.get_seed_wrapped(&seed_id) {
+        if let Some(farm_seed) = self.get_seed_wrapped(&seed_id) {
             let amount = farm_seed.get_ref().amount;
-            if let Some(farm) = farm_seed.get_ref_mut().farms.get_mut(farm_id) {
+            if let Some(mut farm) = self.data().farms.get(farm_id) {
                 claim_user_reward_from_farm(
-                    farm, 
+                    &mut farm, 
                     farmer.get_ref_mut(), 
                     &amount,
                     false,
                 );
-                self.data_mut().seeds.insert(&seed_id, &farm_seed);
+                self.data_mut().farms.insert(farm_id, &farm);
                 self.data_mut().farmers.insert(sender_id, &farmer);
             }
         }
@@ -221,7 +214,7 @@ impl Contract {
 
     #[inline]
     pub(crate) fn get_farmer_default(&self, from: &AccountId) -> VersionedFarmer {
-        let orig = self.data().farmers.get(from).unwrap_or(VersionedFarmer::new(0));
+        let orig = self.data().farmers.get(from).unwrap_or(VersionedFarmer::new(from.clone(), 0));
         if orig.need_upgrade() {
             orig.upgrade()
         } else {

diff --git a/ref-farming/src/actions_of_seed.rs b/ref-farming/src/actions_of_seed.rs
@@ -5,7 +5,7 @@ use near_sdk::{AccountId, Balance, PromiseResult};
 
 use crate::utils::{
     assert_one_yocto, ext_multi_fungible_token, ext_fungible_token, 
-    ext_self, parse_seed_id, GAS_FOR_FT_TRANSFER
+    ext_self, wrap_mft_token_id, parse_seed_id, GAS_FOR_FT_TRANSFER
 };
 use crate::errors::*;
 use crate::farm_seed::SeedType;
@@ -19,9 +19,6 @@ impl Contract {
     pub fn withdraw_seed(&mut self, seed_id: SeedId, amount: U128) {
         assert_one_yocto();
         let sender_id = env::predecessor_account_id();
-        self.assert_storage_usage(&sender_id);
-
-        self.remove_unused_rps(&sender_id);
 
         let amount: Balance = amount.into();
 
@@ -50,7 +47,7 @@ impl Contract {
             SeedType::MFT => {
                 let (receiver_id, token_id) = parse_seed_id(&seed_id);
                 ext_multi_fungible_token::mft_transfer(
-                    token_id,
+                    wrap_mft_token_id(&token_id),
                     sender_id.clone().try_into().unwrap(),
                     amount.into(),
                     None,
@@ -234,8 +231,8 @@ impl Contract {
 
         if farmer_seed_remain == 0 {
             // remove farmer rps of relative farm
-            for farm in &mut farm_seed.get_ref_mut().farms.values_mut() {
-                farmer.get_ref_mut().user_rps.remove(&farm.get_farm_id());
+            for farm_id in farm_seed.get_ref().farms.iter() {
+                farmer.get_ref_mut().remove_rps(farm_id);
             }
         }
         self.data_mut().farmers.insert(sender_id, &farmer);

diff --git a/ref-farming/src/errors.rs b/ref-farming/src/errors.rs
@@ -16,6 +16,7 @@ pub const ERR31_SEED_NOT_EXIST: &str = "E31: seed not exist";
 pub const ERR32_NOT_ENOUGH_SEED: &str = "E32: not enough amount of seed";
 pub const ERR33_INVALID_SEED_ID: &str = "E33: invalid seed id";
 pub const ERR34_BELOW_MIN_SEED_DEPOSITED: &str = "E34: below min_deposit of this seed";
+pub const ERR35_ILLEGAL_TOKEN_ID: &str = "E35: illegal token_id in mft_transfer_call";
 
 // farm errors //
 pub const ERR41_FARM_NOT_EXIST: &str = "E41: farm not exist";

diff --git a/ref-farming/src/farm.rs b/ref-farming/src/farm.rs
@@ -52,12 +52,6 @@ impl Farm {
         }
     }
 
-    pub fn can_be_removed(&self) -> bool {
-        match self {
-            Farm::SimpleFarm(farm) => farm.can_be_removed(),
-        }
-    }
-
     /// Returns how many reward tokens can given farmer claim.
     pub fn view_farmer_unclaimed_reward(
         &self,
@@ -85,4 +79,16 @@ impl Farm {
         }
     }
 
+    pub fn can_be_removed(&self, total_seeds: &Balance) -> bool {
+        match self {
+            Farm::SimpleFarm(farm) => farm.can_be_removed(total_seeds),
+        }
+    }
+
+    pub fn move_to_clear(&mut self, total_seeds: &Balance) -> bool {
+        match self {
+            Farm::SimpleFarm(farm) => farm.move_to_clear(total_seeds),
+        }
+    }
+
 }
diff --git a/ref-farming/src/farm_seed.rs b/ref-farming/src/farm_seed.rs
@@ -1,13 +1,13 @@
 //! FarmSeed stores information per seed about 
 //! staked seed amount and farms under it.
 
-use std::collections::HashMap;
+use std::collections::HashSet;
 use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::{Balance};
 use near_sdk::serde::{Deserialize, Serialize};
 use near_sdk::json_types::{U128};
 use crate::errors::*;
-use crate::farm::{Farm, FarmId};
+use crate::farm::FarmId;
 use crate::utils::parse_seed_id;
 
 
@@ -32,7 +32,7 @@ pub struct FarmSeed {
     pub seed_type: SeedType,
     /// all farms that accepted this seed
     /// FarmId = {seed_id}#{next_index}
-    pub farms: HashMap<FarmId, Farm>,
+    pub farms: HashSet<FarmId>,
     pub next_index: u32,
     /// total (staked) balance of this seed (Farming Token)
     pub amount: Balance,
@@ -51,7 +51,7 @@ impl FarmSeed {
         Self {
             seed_id: seed_id.clone(),
             seed_type,
-            farms: HashMap::new(),
+            farms: HashSet::new(),
             next_index: 0,
             amount: 0,
             min_deposit,
@@ -146,7 +146,7 @@ impl From<&FarmSeed> for SeedInfo {
             next_index: fs.next_index,
             amount: fs.amount.into(),
             min_deposit: fs.min_deposit.into(),
-            farms: fs.farms.keys().map(|key| key.clone()).collect(),
+            farms: fs.farms.iter().map(|key| key.clone()).collect(),
         }
     }
 }