Skip to content

Fix LP token storage handling #27

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
referencedev merged 3 commits into from
May 17, 2021
Merged

Fix LP token storage handling #27

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
54b2eb1
Fix charging for storage for LP tokens
referencedev May 14, 2021
08d070c
Handle storage during mft_transfer
referencedev May 15, 2021
c6e4776
Update release wasm
referencedev May 17, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 26 additions & 28 deletions ref-exchange/src/account_deposit.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,17 +16,17 @@ const MIN_ACCOUNT_DEPOSIT_LENGTH: u128 = MAX_ACCOUNT_LENGTH + 16 + 4;
/// Account deposits information and storage cost.
#[derive(BorshSerialize, BorshDeserialize, Default)]
#[cfg_attr(feature = "test", derive(Clone))]
pub struct AccountDeposit {
/// Native amount sent to the exchange.
/// Used for storage now, but in future can be used for trading as well.
pub amount: Balance,
/// Amounts of various tokens in this account.
pub struct Account {
/// Native NEAR amount sent to the exchange.
/// Used for storage right now, but in future can be used for trading as well.
pub near_amount: Balance,
/// Amounts of various tokens deposited to this account.
pub tokens: HashMap<AccountId, Balance>,
}

impl AccountDeposit {
/// Adds amount to the balance of given token while checking that storage is covered.
pub(crate) fn add(&mut self, token: &AccountId, amount: Balance) {
impl Account {
/// Deposit amount to the balance of given token while checking that storage is covered.
pub(crate) fn deposit(&mut self, token: &AccountId, amount: Balance) {
if let Some(x) = self.tokens.get_mut(token) {
*x = *x + amount;
} else {
Expand All @@ -35,9 +35,9 @@ impl AccountDeposit {
}
}

/// Subtract from `token` balance.
/// Withdraw amount of `token` from the internal balance.
/// Panics if `amount` is bigger than the current balance.
pub(crate) fn sub(&mut self, token: &AccountId, amount: Balance) {
pub(crate) fn withdraw(&mut self, token: &AccountId, amount: Balance) {
let value = *self.tokens.get(token).expect(ERR21_TOKEN_NOT_REG);
assert!(value >= amount, "{}", ERR22_NOT_ENOUGH_TOKENS);
self.tokens.insert(token.clone(), value - amount);
Expand All @@ -51,13 +51,13 @@ impl AccountDeposit {

/// Returns how much NEAR is available for storage.
pub fn storage_available(&self) -> Balance {
self.amount - self.storage_usage()
self.near_amount - self.storage_usage()
}

/// Asserts there is sufficient amount of $NEAR to cover storage usage.
pub fn assert_storage_usage(&self) {
assert!(
self.storage_usage() <= self.amount,
self.storage_usage() <= self.near_amount,
"{}",
ERR11_INSUFFICIENT_STORAGE
);
Expand Down Expand Up @@ -96,7 +96,7 @@ impl Contract {
let sender_id = env::predecessor_account_id();
let mut deposits = self.get_account_deposits(&sender_id);
deposits.register(&token_ids);
self.deposited_amounts.insert(&sender_id, &deposits);
self.accounts.insert(&sender_id, &deposits);
}

/// Unregister given token from user's account deposit.
Expand All @@ -107,7 +107,7 @@ impl Contract {
for token_id in token_ids {
deposits.unregister(token_id.as_ref());
}
self.deposited_amounts.insert(&sender_id, &deposits);
self.accounts.insert(&sender_id, &deposits);
}

/// Withdraws given token from the deposits of given user.
Expand All @@ -121,11 +121,11 @@ impl Contract {
let sender_id = env::predecessor_account_id();
let mut deposits = self.get_account_deposits(&sender_id);
// Note: subtraction and deregistration will be reverted if the promise fails.
deposits.sub(&token_id, amount);
deposits.withdraw(&token_id, amount);
if unregister == Some(true) {
deposits.unregister(&token_id);
}
self.deposited_amounts.insert(&sender_id, &deposits);
self.accounts.insert(&sender_id, &deposits);
ext_fungible_token::ft_transfer(
sender_id.clone().try_into().unwrap(),
amount.into(),
Expand Down Expand Up @@ -163,8 +163,8 @@ impl Contract {
PromiseResult::Failed => {
// This reverts the changes from withdraw function.
let mut deposits = self.get_account_deposits(&sender_id);
deposits.add(&token_id, amount.0);
self.deposited_amounts.insert(&sender_id, &deposits);
deposits.deposit(&token_id, amount.0);
self.accounts.insert(&sender_id, &deposits);
}
};
}
Expand All @@ -175,9 +175,9 @@ impl Contract {
/// If account already exists, adds amount to it.
/// This should be used when it's known that storage is prepaid.
pub(crate) fn internal_register_account(&mut self, account_id: &AccountId, amount: Balance) {
let mut deposit_amount = self.deposited_amounts.get(&account_id).unwrap_or_default();
deposit_amount.amount += amount;
self.deposited_amounts.insert(&account_id, &deposit_amount);
let mut deposit_amount = self.accounts.get(&account_id).unwrap_or_default();
deposit_amount.near_amount += amount;
self.accounts.insert(&account_id, &deposit_amount);
}

/// Record deposit of some number of tokens to this contract.
Expand All @@ -195,16 +195,14 @@ impl Contract {
"{}",
ERR12_TOKEN_NOT_WHITELISTED
);
account_deposit.add(token_id, amount);
self.deposited_amounts.insert(sender_id, &account_deposit);
account_deposit.deposit(token_id, amount);
self.accounts.insert(sender_id, &account_deposit);
}

// Returns `from` AccountDeposit.
#[inline]
pub(crate) fn get_account_deposits(&self, from: &AccountId) -> AccountDeposit {
self.deposited_amounts
.get(from)
.expect(ERR10_ACC_NOT_REGISTERED)
pub(crate) fn get_account_deposits(&self, from: &AccountId) -> Account {
self.accounts.get(from).expect(ERR10_ACC_NOT_REGISTERED)
}

/// Returns current balance of given token for given user. If there is nothing recorded, returns 0.
Expand All @@ -213,7 +211,7 @@ impl Contract {
sender_id: &AccountId,
token_id: &AccountId,
) -> Balance {
self.deposited_amounts
self.accounts
.get(sender_id)
.and_then(|d| d.tokens.get(token_id).cloned())
.unwrap_or_default()
Expand Down
4 changes: 3 additions & 1 deletion ref-exchange/src/errors.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,10 @@
pub const ERR10_ACC_NOT_REGISTERED: &str = "E10: account not registered";
pub const ERR11_INSUFFICIENT_STORAGE: &str = "E11: insufficient $NEAR storage deposit";
pub const ERR12_TOKEN_NOT_WHITELISTED: &str = "E12: token not whitelisted";
pub const ERR13_LP_NOT_REGISTERED: &str = "E13: LP not registered";
pub const ERR14_LP_ALREADY_REGISTERED: &str = "E14: LP already registered";

// Account Deposits //
// Accounts //

pub const ERR21_TOKEN_NOT_REG: &str = "E21: token not registered";
pub const ERR22_NOT_ENOUGH_TOKENS: &str = "E22: not enough tokens in deposit";
Expand Down
85 changes: 55 additions & 30 deletions ref-exchange/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,11 @@ use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
use near_sdk::collections::{LookupMap, UnorderedSet, Vector};
use near_sdk::json_types::{ValidAccountId, U128};
use near_sdk::{
assert_one_yocto, env, log, near_bindgen, AccountId, PanicOnDefault, Promise, PromiseResult,
assert_one_yocto, env, log, near_bindgen, AccountId, Balance, PanicOnDefault, Promise,
PromiseResult, StorageUsage,
};

use crate::account_deposit::AccountDeposit;
use crate::account_deposit::Account;
pub use crate::action::*;
use crate::errors::*;
use crate::pool::Pool;
Expand Down Expand Up @@ -44,8 +45,8 @@ pub struct Contract {
referral_fee: u32,
/// List of all the pools.
pools: Vector<Pool>,
/// Balances of deposited tokens for each account.
deposited_amounts: LookupMap<AccountId, AccountDeposit>,
/// Accounts registered, keeping track all the amounts deposited, storage and more.
accounts: LookupMap<AccountId, Account>,
/// Set of whitelisted tokens by "owner".
whitelisted_tokens: UnorderedSet<AccountId>,
}
Expand All @@ -59,7 +60,7 @@ impl Contract {
exchange_fee,
referral_fee,
pools: Vector::new(b"p".to_vec()),
deposited_amounts: LookupMap::new(b"d".to_vec()),
accounts: LookupMap::new(b"d".to_vec()),
whitelisted_tokens: UnorderedSet::new(b"w".to_vec()),
}
}
Expand Down Expand Up @@ -112,7 +113,11 @@ impl Contract {
amounts: Vec<U128>,
min_amounts: Option<Vec<U128>>,
) {
assert_one_yocto();
assert!(
env::attached_deposit() > 0,
"Requires attached deposit of at least 1 yoctoNEAR"
);
let prev_storage = env::storage_usage();
let sender_id = env::predecessor_account_id();
let mut amounts: Vec<u128> = amounts.into_iter().map(|amount| amount.into()).collect();
let mut pool = self.pools.get(pool_id).expect("ERR_NO_POOL");
Expand All @@ -124,20 +129,22 @@ impl Contract {
assert!(amount >= &min_amount.0, "ERR_MIN_AMOUNT");
}
}
let mut deposits = self.deposited_amounts.get(&sender_id).unwrap_or_default();
let mut deposits = self.accounts.get(&sender_id).unwrap_or_default();
let tokens = pool.tokens();
// Subtract updated amounts from deposits. This will fail if there is not enough funds for any of the tokens.
for i in 0..tokens.len() {
deposits.sub(&tokens[i], amounts[i]);
deposits.withdraw(&tokens[i], amounts[i]);
}
self.deposited_amounts.insert(&sender_id, &deposits);
self.accounts.insert(&sender_id, &deposits);
self.pools.replace(pool_id, &pool);
self.internal_check_storage(prev_storage);
}

/// Remove liquidity from the pool into general pool of liquidity.
#[payable]
pub fn remove_liquidity(&mut self, pool_id: u64, shares: U128, min_amounts: Vec<U128>) {
assert_one_yocto();
let prev_storage = env::storage_usage();
let sender_id = env::predecessor_account_id();
let mut pool = self.pools.get(pool_id).expect("ERR_NO_POOL");
let amounts = pool.remove_liquidity(
Expand All @@ -150,34 +157,43 @@ impl Contract {
);
self.pools.replace(pool_id, &pool);
let tokens = pool.tokens();
let mut deposits = self.deposited_amounts.get(&sender_id).unwrap_or_default();
let mut deposits = self.accounts.get(&sender_id).unwrap_or_default();
for i in 0..tokens.len() {
deposits.add(&tokens[i], amounts[i]);
deposits.deposit(&tokens[i], amounts[i]);
}
// Freed up storage balance from LP tokens will be returned to near_balance.
if prev_storage > env::storage_usage() {
deposits.near_amount +=
(prev_storage - env::storage_usage()) as Balance * env::storage_byte_cost();
}
self.deposited_amounts.insert(&sender_id, &deposits);
self.accounts.insert(&sender_id, &deposits);
}
}

/// Internal methods implementation.
impl Contract {
/// Check how much storage taken costs and refund the left over back.
fn internal_check_storage(&self, prev_storage: StorageUsage) {
let storage_cost = env::storage_usage()
.checked_sub(prev_storage)
.unwrap_or_default() as Balance
* env::storage_byte_cost();
let refund = env::attached_deposit()
.checked_sub(storage_cost)
.expect("ERR_STORAGE_DEPOSIT");
if refund > 0 {
Promise::new(env::predecessor_account_id()).transfer(refund);
}
}

/// Adds given pool to the list and returns it's id.
/// If there is not enough attached balance to cover storage, fails.
/// If too much attached - refunds it back.
fn internal_add_pool(&mut self, pool: Pool) -> u64 {
let prev_storage = env::storage_usage();
let id = self.pools.len() as u64;
self.pools.push(&pool);

// Check how much storage cost and refund the left over back.
let storage_cost = (env::storage_usage() - prev_storage) as u128 * env::storage_byte_cost();
assert!(
storage_cost <= env::attached_deposit(),
"ERR_STORAGE_DEPOSIT"
);
let refund = env::attached_deposit() - storage_cost;
if refund > 0 {
Promise::new(env::predecessor_account_id()).transfer(refund);
}
self.internal_check_storage(prev_storage);
id
}

Expand All @@ -193,8 +209,8 @@ impl Contract {
min_amount_out: u128,
referral_id: &Option<AccountId>,
) -> u128 {
let mut deposits = self.deposited_amounts.get(&sender_id).unwrap_or_default();
deposits.sub(token_in, amount_in);
let mut deposits = self.accounts.get(&sender_id).unwrap_or_default();
deposits.withdraw(token_in, amount_in);
let mut pool = self.pools.get(pool_id).expect("ERR_NO_POOL");
let amount_out = pool.swap(
token_in,
Expand All @@ -204,8 +220,8 @@ impl Contract {
&self.owner_id,
referral_id,
);
deposits.add(token_out, amount_out);
self.deposited_amounts.insert(&sender_id, &deposits);
deposits.deposit(token_out, amount_out);
self.accounts.insert(&sender_id, &deposits);
self.pools.replace(pool_id, &pool);
amount_out
}
Expand Down Expand Up @@ -286,7 +302,7 @@ mod tests {
deposit_tokens(context, contract, accounts(3), token_amounts.clone());
testing_env!(context
.predecessor_account_id(account_id.clone())
.attached_deposit(1)
.attached_deposit(to_yocto("0.00067"))
.build());
contract.add_liquidity(
pool_id,
Expand Down Expand Up @@ -363,7 +379,15 @@ mod tests {
);
assert_eq!(contract.get_deposit(accounts(1), accounts(2)).0, one_near);

testing_env!(context.predecessor_account_id(accounts(3)).build());
testing_env!(context
.predecessor_account_id(accounts(3))
.attached_deposit(to_yocto("0.0067"))
.build());
contract.mft_register("0".to_string(), accounts(1));
testing_env!(context
.predecessor_account_id(accounts(3))
.attached_deposit(1)
.build());
// transfer 1m shares in pool 0 to acc 1.
contract.mft_transfer("0".to_string(), accounts(1), U128(1_000_000), None);

Expand Down Expand Up @@ -405,9 +429,10 @@ mod tests {
.attached_deposit(to_yocto("1"))
.build());
let id = contract.add_simple_pool(vec![accounts(1), accounts(2)], 25);
testing_env!(context.attached_deposit(1).build());
testing_env!(context.attached_deposit(to_yocto("0.00067")).build());
contract.add_liquidity(id, vec![U128(to_yocto("50")), U128(to_yocto("10"))], None);
contract.add_liquidity(id, vec![U128(to_yocto("50")), U128(to_yocto("50"))], None);
testing_env!(context.attached_deposit(1).build());
contract.remove_liquidity(id, U128(to_yocto("1")), vec![U128(1), U128(1)]);

// Check that amounts add up to deposits.
Expand Down
31 changes: 23 additions & 8 deletions ref-exchange/src/multi_fungible_token.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,18 +65,17 @@ impl Contract {
}
TokenOrPool::Token(token_id) => {
let mut sender_account = self
.deposited_amounts
.accounts
.get(&sender_id)
.expect(ERR10_ACC_NOT_REGISTERED);
let mut receiver_account = self
.deposited_amounts
.accounts
.get(receiver_id)
.expect(ERR10_ACC_NOT_REGISTERED);
sender_account.sub(&token_id, amount);
receiver_account.add(&token_id, amount);
self.deposited_amounts.insert(&sender_id, &sender_account);
self.deposited_amounts
.insert(&receiver_id, &receiver_account);
sender_account.withdraw(&token_id, amount);
receiver_account.deposit(&token_id, amount);
self.accounts.insert(&sender_id, &sender_account);
self.accounts.insert(&receiver_id, &receiver_account);
log!(
"Transfer {}: {} from {} to {}",
token_id,
Expand Down Expand Up @@ -119,6 +118,22 @@ impl Contract {
}
}

/// Register LP token of given pool for given account.
/// Fails if token_id is not a pool.
#[payable]
pub fn mft_register(&mut self, token_id: String, account_id: ValidAccountId) {
let prev_storage = env::storage_usage();
match parse_token_id(token_id) {
TokenOrPool::Token(_) => env::panic(b"ERR_INVALID_REGISTER"),
TokenOrPool::Pool(pool_id) => {
let mut pool = self.pools.get(pool_id).expect("ERR_NO_POOL");
pool.share_register(account_id.as_ref());
self.pools.replace(pool_id, &pool);
self.internal_check_storage(prev_storage);
}
}
}

/// Transfer one of internal tokens: LP or balances.
/// `token_id` can either by account of the token or pool number.
#[payable]
Expand Down Expand Up @@ -206,7 +221,7 @@ impl Contract {
let refund_amount = std::cmp::min(receiver_balance, unused_amount);
// If sender's account was deleted, we assume that they have also withdrew all the liquidity from pools.
// Funds are sent to the owner account.
let refund_to = if self.deposited_amounts.get(&sender_id).is_some() {
let refund_to = if self.accounts.get(&sender_id).is_some() {
sender_id
} else {
self.owner_id.clone()
Expand Down
Loading