diff --git a/Cargo.lock b/Cargo.lock index 7be22c6..9fe43e7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,16 @@ # It is not intended for manual editing. version = 3 +[[package]] +name = "aead" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d122413f284cf2d62fb1b7db97e02edb8cda96d769b16e443a4f6195e35662b0" +dependencies = [ + "crypto-common", + "generic-array", +] + [[package]] name = "ahash" version = "0.7.8" @@ -89,6 +99,41 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "chacha20" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3613f74bd2eac03dad61bd53dbe620703d4371614fe0bc3b9f04dd36fe4e818" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures", +] + +[[package]] +name = "chacha20poly1305" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10cd79432192d1c0f4e1a0fef9527696cc039165d729fb41b3f4f4f354c2dc35" +dependencies = [ + "aead", + "chacha20", + "cipher", + "poly1305", + "zeroize", +] + +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common", + "inout", + "zeroize", +] + [[package]] name = "const-oid" version = "0.9.6" @@ -168,6 +213,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" dependencies = [ "generic-array", + "rand_core 0.6.4", "typenum", ] @@ -288,6 +334,15 @@ dependencies = [ "subtle", ] +[[package]] +name = "fixed-hash" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "835c052cb0c08c1acf6ffd71c022172e18723949c8282f2b9f27efbc51e64534" +dependencies = [ + "static_assertions", +] + [[package]] name = "forward_ref" version = "1.0.0" @@ -341,6 +396,15 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "hkdf" +version = "0.12.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b5f8eb2ad728638ea2c7d47a21db23b7b58a72ed6a38256b8a1849f15fbbdf7" +dependencies = [ + "hmac", +] + [[package]] name = "hmac" version = "0.12.1" @@ -350,6 +414,15 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "inout" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0c10553d664a4d0bcff9f4215d0aac67a639cc68ef660840afe309b807bc9f5" +dependencies = [ + "generic-array", +] + [[package]] name = "itoa" version = "1.0.11" @@ -374,6 +447,36 @@ version = "0.2.155" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "97b3888a4aecf77e811145cadf6eef5901f4782c53886191b2f693f24761847c" +[[package]] +name = "minicbor" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a20020e8e2d1881d8736f64011bb5ff99f1db9947ce3089706945c8915695cb" +dependencies = [ + "minicbor-derive", +] + +[[package]] +name = "minicbor-derive" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8608fb1c805b5b6b3d5ab7bd95c40c396df622b64d77b2d621a5eae1eed050ee" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + +[[package]] +name = "minicbor-ser" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0834b86a9c56311671913d56f640d7f0b6da803df61121661cc890f0edc0eb1" +dependencies = [ + "minicbor", + "serde", +] + [[package]] name = "once_cell" version = "1.19.0" @@ -396,12 +499,33 @@ dependencies = [ "spki", ] +[[package]] +name = "poly1305" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8159bd90725d2df49889a078b54f4f79e87f1f8a8444194cdca81d38f5393abf" +dependencies = [ + "cpufeatures", + "opaque-debug", + "universal-hash", +] + [[package]] name = "ppv-lite86" version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" +[[package]] +name = "primitive-types" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b34d9fd68ae0b74a41b21c03c2f62847aa0ffea044eee893b4c140b37e244e2" +dependencies = [ + "fixed-hash", + "uint", +] + [[package]] name = "proc-macro2" version = "1.0.85" @@ -592,10 +716,10 @@ dependencies = [ [[package]] name = "secret-toolkit" version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "338c972c0a98de51ccbb859312eb7672bc64b9050b086f058748ba26a509edbb" +source = "git+https://github.com/SolarRepublic/secret-toolkit.git#1bf8c5f3f08e75276fa812cb248c6eda5cd3beae" dependencies = [ "secret-toolkit-crypto", + "secret-toolkit-notification", "secret-toolkit-permit", "secret-toolkit-serialization", "secret-toolkit-storage", @@ -606,8 +730,7 @@ dependencies = [ [[package]] name = "secret-toolkit-crypto" version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "003d7d5772c67f2240b7f298f96eb73a8a501916fe18c1d730ebfd591bf7e519" +source = "git+https://github.com/SolarRepublic/secret-toolkit.git#1bf8c5f3f08e75276fa812cb248c6eda5cd3beae" dependencies = [ "rand_chacha", "rand_core 0.6.4", @@ -616,11 +739,28 @@ dependencies = [ "sha2 0.10.8", ] +[[package]] +name = "secret-toolkit-notification" +version = "0.10.0" +source = "git+https://github.com/SolarRepublic/secret-toolkit.git#1bf8c5f3f08e75276fa812cb248c6eda5cd3beae" +dependencies = [ + "chacha20poly1305", + "generic-array", + "hkdf", + "minicbor-ser", + "primitive-types", + "ripemd", + "schemars", + "secret-cosmwasm-std", + "secret-toolkit-crypto", + "serde", + "sha2 0.10.8", +] + [[package]] name = "secret-toolkit-permit" version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4330571400b5959450fa37040609e6804a147d83f606783506bc2275f1527712" +source = "git+https://github.com/SolarRepublic/secret-toolkit.git#1bf8c5f3f08e75276fa812cb248c6eda5cd3beae" dependencies = [ "bech32", "remain", @@ -634,8 +774,7 @@ dependencies = [ [[package]] name = "secret-toolkit-serialization" version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "890adaeaa710f9f7068a807eb1553edc8c30ce9907290895c9097dd642fc613b" +source = "git+https://github.com/SolarRepublic/secret-toolkit.git#1bf8c5f3f08e75276fa812cb248c6eda5cd3beae" dependencies = [ "bincode2", "schemars", @@ -646,8 +785,7 @@ dependencies = [ [[package]] name = "secret-toolkit-storage" version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "55e8c5418af3e7ae1d1331c383b32d56c74a340dbc3b972d53555a768698f2a3" +source = "git+https://github.com/SolarRepublic/secret-toolkit.git#1bf8c5f3f08e75276fa812cb248c6eda5cd3beae" dependencies = [ "secret-cosmwasm-std", "secret-cosmwasm-storage", @@ -658,8 +796,7 @@ dependencies = [ [[package]] name = "secret-toolkit-utils" version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "83f1cba2e70fd701e3dfc6072807c02eeeb9776bee49e346a9c7745d84ff40c8" +source = "git+https://github.com/SolarRepublic/secret-toolkit.git#1bf8c5f3f08e75276fa812cb248c6eda5cd3beae" dependencies = [ "schemars", "secret-cosmwasm-std", @@ -670,8 +807,7 @@ dependencies = [ [[package]] name = "secret-toolkit-viewing-key" version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d89a0b69fa9b12735a612fa30e6e7e48130943982f1783b7ddd5c46ed09e921" +source = "git+https://github.com/SolarRepublic/secret-toolkit.git#1bf8c5f3f08e75276fa812cb248c6eda5cd3beae" dependencies = [ "base64 0.21.7", "schemars", @@ -784,6 +920,7 @@ dependencies = [ "base64 0.21.7", "constant_time_eq", "cosmwasm-schema", + "primitive-types", "rand", "schemars", "secret-cosmwasm-std", @@ -883,6 +1020,16 @@ version = "1.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" +[[package]] +name = "universal-hash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea" +dependencies = [ + "crypto-common", + "subtle", +] + [[package]] name = "version_check" version = "0.9.4" diff --git a/Cargo.toml b/Cargo.toml index fc90942..1128174 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,8 +35,10 @@ backtraces = ["cosmwasm-std/backtraces"] cosmwasm-std = { package = "secret-cosmwasm-std", version = "1.1.11" } cosmwasm-storage = { package = "secret-cosmwasm-storage", version = "1.1.11" } rand = { version = "0.8.5", default-features = false } -secret-toolkit = { version = "0.10.0", default-features = false, features = ["permit", "storage", "viewing-key"] } -secret-toolkit-crypto = { version = "0.10.0", default-features = false, features = ["hash"] } +# secret-toolkit = { version = "0.10.0", default-features = false, features = ["permit", "storage", "viewing-key"] } +secret-toolkit = { git = "https://github.com/SolarRepublic/secret-toolkit.git", default-features = false, features = ["permit", "storage", "viewing-key", "notification"] } +# secret-toolkit-crypto = { version = "0.10.0", default-features = false, features = ["hash"] } +secret-toolkit-crypto = { git = "https://github.com/SolarRepublic/secret-toolkit.git", default-features = false, features = ["hash"] } static_assertions = "1.1.0" schemars = "0.8.12" @@ -44,6 +46,7 @@ serde = { version = "1.0.158", default-features = false, features = ["derive"] } serde-big-array = "0.5.1" base64 = "0.21.0" constant_time_eq = "0.3.0" +primitive-types = { version = "0.12.2", default-features = false } [dev-dependencies] cosmwasm-schema = { version = "1.1.8" } diff --git a/src/contract.rs b/src/contract.rs index cd5f232..521d924 100644 --- a/src/contract.rs +++ b/src/contract.rs @@ -3,22 +3,24 @@ use cosmwasm_std::{ entry_point, to_binary, Addr, Api, BankMsg, Binary, BlockInfo, CanonicalAddr, Coin, CosmosMsg, Deps, DepsMut, Env, MessageInfo, Response, StdError, StdResult, Storage, Uint128 }; +use secret_toolkit::notification::hkdf_sha_256; use secret_toolkit::permit::{Permit, RevokedPermits, TokenPermissions}; use secret_toolkit::utils::{pad_handle_result, pad_query_result}; use secret_toolkit::viewing_key::{ViewingKey, ViewingKeyStore}; use secret_toolkit_crypto::{sha_256, ContractPrng}; use crate::batch; -use crate::dwb::{log_dwb, AccountTxsStore, DelayedWriteBuffer, ACCOUNT_TXS, ACCOUNT_TX_COUNT, DWB, TX_NODES}; -use crate::bucket; +use crate::dwb::{log_dwb, DelayedWriteBuffer, DWB, TX_NODES}; +//use crate::bucket; use crate::msg::{ AllowanceGivenResult, AllowanceReceivedResult, ContractStatusLevel, ExecuteAnswer, ExecuteMsg, InstantiateMsg, QueryAnswer, QueryMsg, QueryWithPermit, ResponseStatus::Success, }; use crate::receiver::Snip20ReceiveMsg; use crate::state::{ - safe_add, AllowancesStore, BalancesStore, Config, MintersStore, PrngStore, ReceiverHashStore, CONFIG, CONTRACT_STATUS, PRNG, TOTAL_SUPPLY + safe_add, AllowancesStore, Config, MintersStore, PrngStore, ReceiverHashStore, CONFIG, CONTRACT_STATUS, INTERNAL_SECRET, PRNG, TOTAL_SUPPLY }; +use crate::stored_balances::{find_start_bundle, initialize_btsb, stored_balance, stored_entry, stored_tx_count}; use crate::strings::TRANSFER_HISTORY_UNSUPPORTED_MSG; use crate::transaction_history::{ store_burn_action, store_deposit_action, store_mint_action, store_redeem_action, store_transfer_action, Tx @@ -54,7 +56,7 @@ pub fn instantiate( let admin = match msg.admin { Some(admin_addr) => deps.api.addr_validate(admin_addr.as_str())?, - None => info.sender, + None => info.sender.clone(), }; let mut total_supply: u128 = 0; @@ -62,17 +64,39 @@ pub fn instantiate( let prng_seed_hashed = sha_256(&msg.prng_seed.0); PrngStore::save(deps.storage, prng_seed_hashed)?; + // initialize the bitwise-trie of stored entries + initialize_btsb(deps.storage)?; + // initialize the delay write buffer DWB.save(deps.storage, &DelayedWriteBuffer::new()?)?; let initial_balances = msg.initial_balances.unwrap_or_default(); let raw_admin = deps.api.addr_canonicalize(admin.as_str())?; - let seed = env.block.random.as_ref().unwrap(); - let mut rng = ContractPrng::new(seed.as_slice(), &prng_seed_hashed); + let rng_seed = env.block.random.as_ref().unwrap(); + let mut rng = ContractPrng::new(rng_seed.as_slice(), &prng_seed_hashed); + + // use entropy and env.random to create an internal secret for the contract + let entropy = msg.prng_seed.0.as_slice(); + let entropy_len = 16 + info.sender.to_string().len() + entropy.len(); + let mut rng_entropy = Vec::with_capacity(entropy_len); + rng_entropy.extend_from_slice(&env.block.height.to_be_bytes()); + rng_entropy.extend_from_slice(&env.block.time.seconds().to_be_bytes()); + rng_entropy.extend_from_slice(info.sender.as_bytes()); + rng_entropy.extend_from_slice(entropy); + + // Create INTERNAL_SECRET + let salt = Some(sha_256(&rng_entropy).to_vec()); + let internal_secret = hkdf_sha_256( + &salt, + rng_seed.0.as_slice(), + "contract_internal_secret".as_bytes(), + 32, + )?; + INTERNAL_SECRET.save(deps.storage, &internal_secret)?; + for balance in initial_balances { let amount = balance.amount.u128(); let balance_address = deps.api.addr_canonicalize(balance.address.as_str())?; - perform_mint( deps.storage, &mut rng, @@ -147,7 +171,7 @@ pub fn execute(deps: DepsMut, env: Env, info: MessageInfo, msg: ExecuteMsg) -> S denom, .. } if contract_status == ContractStatusLevel::StopAllButRedeems => { - try_redeem(deps, env, info, &mut rng, amount, denom) + try_redeem(deps, env, info, amount, denom) } _ => Err(StdError::generic_err( "This contract is stopped and this action is not allowed", @@ -158,6 +182,8 @@ pub fn execute(deps: DepsMut, env: Env, info: MessageInfo, msg: ExecuteMsg) -> S ContractStatusLevel::NormalRun => {} // If it's a normal run just continue } + let secret = INTERNAL_SECRET.load(deps.storage)?; + let secret = secret.as_slice(); let response = match msg.clone() { // Native ExecuteMsg::Deposit { .. } => { @@ -167,7 +193,7 @@ pub fn execute(deps: DepsMut, env: Env, info: MessageInfo, msg: ExecuteMsg) -> S amount, denom, .. - } => try_redeem(deps, env, info, &mut rng, amount, denom), + } => try_redeem(deps, env, info, amount, denom), // Base ExecuteMsg::Transfer { @@ -212,7 +238,7 @@ pub fn execute(deps: DepsMut, env: Env, info: MessageInfo, msg: ExecuteMsg) -> S amount, memo, .. - } => try_burn(deps, env, info, &mut rng, amount, memo), + } => try_burn(deps, env, info, amount, memo), ExecuteMsg::RegisterReceive { code_hash, .. } => { try_register_receive(deps, info, code_hash) } @@ -283,13 +309,12 @@ pub fn execute(deps: DepsMut, env: Env, info: MessageInfo, msg: ExecuteMsg) -> S deps, &env, info, - &mut rng, owner, amount, memo, ), ExecuteMsg::BatchBurnFrom { actions, .. } => { - try_batch_burn_from(deps, &env, info, &mut rng, actions) + try_batch_burn_from(deps, &env, info, actions) } // Mint @@ -594,8 +619,9 @@ pub fn query_transactions( } } - let account_slice = account_raw.as_slice(); - let settled_tx_count = ACCOUNT_TX_COUNT.add_suffix(account_slice).load(deps.storage)?; + //let account_slice = account_raw.as_slice(); + let account_stored_entry = stored_entry(deps.storage, &account_raw)?; + let settled_tx_count = stored_tx_count(deps.storage, &account_stored_entry)?; let total = txs_in_dwb_count as u32 + settled_tx_count as u32; if end > total { end = total; @@ -614,24 +640,25 @@ pub fn query_transactions( //println!("OPTION 2"); txs = txs_in_dwb[start as usize..].to_vec(); // reverse chronological let mut txs_left = (end - start).saturating_sub(txs.len() as u32); - let tx_bundles_store = ACCOUNT_TXS.add_suffix(account_slice); - let tx_bundles_idx_len = tx_bundles_store.get_len(deps.storage)?; - if tx_bundles_idx_len > 0 { - let mut bundle_idx = tx_bundles_idx_len - 1; - loop { - let tx_bundle = tx_bundles_store.get_at(deps.storage, bundle_idx.clone())?; - let head_node = TX_NODES.add_suffix(&tx_bundle.head_node.to_be_bytes()).load(deps.storage)?; - let list_len = tx_bundle.list_len as u32; - if txs_left <= list_len { - txs.extend_from_slice(&head_node.to_vec(deps.storage, deps.api)?[0..txs_left as usize]); - break; - } - txs.extend(head_node.to_vec(deps.storage, deps.api)?); - txs_left = txs_left.saturating_sub(list_len); - if bundle_idx > 0 { - bundle_idx -= 1; - } else { - break; + if let Some(entry) = account_stored_entry { + let tx_bundles_idx_len = entry.history_len()?; + if tx_bundles_idx_len > 0 { + let mut bundle_idx = tx_bundles_idx_len - 1; + loop { + let tx_bundle = entry.get_tx_bundle_at(deps.storage, bundle_idx.clone())?; + let head_node = TX_NODES.add_suffix(&tx_bundle.head_node.to_be_bytes()).load(deps.storage)?; + let list_len = tx_bundle.list_len as u32; + if txs_left <= list_len { + txs.extend_from_slice(&head_node.to_vec(deps.storage, deps.api)?[0..txs_left as usize]); + break; + } + txs.extend(head_node.to_vec(deps.storage, deps.api)?); + txs_left = txs_left.saturating_sub(list_len); + if bundle_idx > 0 { + bundle_idx -= 1; + } else { + break; + } } } } @@ -644,7 +671,7 @@ pub fn query_transactions( //println!("OPTION 3"); let settled_start = settled_tx_count.saturating_sub(start - txs_in_dwb_count).saturating_sub(1); - if let Some((bundle_idx, tx_bundle, start_at)) = AccountTxsStore::find_start_bundle( + if let Some((bundle_idx, tx_bundle, start_at)) = find_start_bundle( deps.storage, &account_raw, settled_start @@ -664,21 +691,22 @@ pub fn query_transactions( if bundle_idx > 0 && txs_left > 0 { // get the next earlier bundle let mut bundle_idx = bundle_idx - 1; - let tx_bundles_store = ACCOUNT_TXS.add_suffix(account_slice); - loop { - let tx_bundle = tx_bundles_store.get_at(deps.storage, bundle_idx.clone())?; - let head_node = TX_NODES.add_suffix(&tx_bundle.head_node.to_be_bytes()).load(deps.storage)?; - let list_len = tx_bundle.list_len as u32; - if txs_left <= list_len { - txs.extend_from_slice(&head_node.to_vec(deps.storage, deps.api)?[0..txs_left as usize]); - break; - } - txs.extend(head_node.to_vec(deps.storage, deps.api)?); - txs_left = txs_left.saturating_sub(list_len); - if bundle_idx > 0 { - bundle_idx -= 1; - } else { - break; + if let Some(entry) = account_stored_entry { + loop { + let tx_bundle = entry.get_tx_bundle_at(deps.storage, bundle_idx.clone())?; + let head_node = TX_NODES.add_suffix(&tx_bundle.head_node.to_be_bytes()).load(deps.storage)?; + let list_len = tx_bundle.list_len as u32; + if txs_left <= list_len { + txs.extend_from_slice(&head_node.to_vec(deps.storage, deps.api)?[0..txs_left as usize]); + break; + } + txs.extend(head_node.to_vec(deps.storage, deps.api)?); + txs_left = txs_left.saturating_sub(list_len); + if bundle_idx > 0 { + bundle_idx -= 1; + } else { + break; + } } } } @@ -701,7 +729,7 @@ pub fn query_balance(deps: Deps, account: String) -> StdResult { let account = Addr::unchecked(account); let account = deps.api.addr_canonicalize(account.as_str())?; - let mut amount = BalancesStore::load(deps.storage, &account); + let mut amount = stored_balance(deps.storage, &account)?; let dwb = DWB.load(deps.storage)?; let dwb_index = dwb.recipient_match(&account); if dwb_index > 0 { @@ -1073,7 +1101,6 @@ fn try_redeem( deps: DepsMut, env: Env, info: MessageInfo, - rng: &mut ContractPrng, amount: Uint128, denom: Option, ) -> StdResult { @@ -1115,7 +1142,7 @@ fn try_redeem( let mut dwb = DWB.load(deps.storage)?; // settle the signer's account in buffer - dwb.settle_sender_or_owner_account(deps.storage, rng, &sender_address, tx_id, amount_raw, "redeem")?; + dwb.settle_sender_or_owner_account(deps.storage, &sender_address, tx_id, amount_raw, "redeem")?; DWB.save(deps.storage, &dwb)?; @@ -1659,7 +1686,6 @@ fn try_burn_from( deps: DepsMut, env: &Env, info: MessageInfo, - rng: &mut ContractPrng, owner: String, amount: Uint128, memo: Option, @@ -1691,9 +1717,9 @@ fn try_burn_from( let mut dwb = DWB.load(deps.storage)?; // settle the owner's account in buffer - dwb.settle_sender_or_owner_account(deps.storage, rng, &raw_owner, tx_id, raw_amount, "burn")?; + dwb.settle_sender_or_owner_account(deps.storage, &raw_owner, tx_id, raw_amount, "burn")?; if raw_burner != raw_owner { // also settle sender's account - dwb.settle_sender_or_owner_account(deps.storage, rng, &raw_burner, tx_id, 0, "burn")?; + dwb.settle_sender_or_owner_account(deps.storage, &raw_burner, tx_id, 0, "burn")?; } DWB.save(deps.storage, &dwb)?; @@ -1717,7 +1743,6 @@ fn try_batch_burn_from( deps: DepsMut, env: &Env, info: MessageInfo, - rng: &mut ContractPrng, actions: Vec, ) -> StdResult { let constants = CONFIG.load(deps.storage)?; @@ -1750,9 +1775,9 @@ fn try_batch_burn_from( let mut dwb = DWB.load(deps.storage)?; // settle the owner's account in buffer - dwb.settle_sender_or_owner_account(deps.storage, rng, &raw_owner, tx_id, amount, "burn")?; + dwb.settle_sender_or_owner_account(deps.storage, &raw_owner, tx_id, amount, "burn")?; if raw_spender != raw_owner { - dwb.settle_sender_or_owner_account(deps.storage, rng, &raw_spender, tx_id, 0, "burn")?; + dwb.settle_sender_or_owner_account(deps.storage, &raw_spender, tx_id, 0, "burn")?; } DWB.save(deps.storage, &dwb)?; @@ -1930,7 +1955,6 @@ fn try_burn( deps: DepsMut, env: Env, info: MessageInfo, - rng: &mut ContractPrng, amount: Uint128, memo: Option, ) -> StdResult { @@ -1958,7 +1982,7 @@ fn try_burn( let mut dwb = DWB.load(deps.storage)?; // settle the signer's account in buffer - dwb.settle_sender_or_owner_account(deps.storage, rng, &raw_burn_address, tx_id, raw_amount, "burn")?; + dwb.settle_sender_or_owner_account(deps.storage, &raw_burn_address, tx_id, raw_amount, "burn")?; DWB.save(deps.storage, &dwb)?; @@ -2005,10 +2029,10 @@ fn perform_transfer( let transfer_str = "transfer"; // settle the owner's account - dwb.settle_sender_or_owner_account(store, rng, from, tx_id, amount, transfer_str)?; + dwb.settle_sender_or_owner_account(store, from, tx_id, amount, transfer_str)?; // if this is a *_from action, settle the sender's account, too if sender != from { - dwb.settle_sender_or_owner_account(store, rng, sender, tx_id, 0, transfer_str)?; + dwb.settle_sender_or_owner_account(store, sender, tx_id, 0, transfer_str)?; } // TESTING @@ -2049,7 +2073,7 @@ fn perform_mint( // if minter is not recipient, settle them if minter != to { - dwb.settle_sender_or_owner_account(store, rng, minter, tx_id, 0, "mint")?; + dwb.settle_sender_or_owner_account(store, minter, tx_id, 0, "mint")?; } // add the tx info for the recipient to the buffer @@ -2438,9 +2462,9 @@ mod tests { .addr_canonicalize(Addr::unchecked("alice").as_str()) .unwrap(); - assert_eq!(5000 - 1000, BalancesStore::load(&deps.storage, &bob_addr)); + assert_eq!(5000 - 1000, stored_balance(&deps.storage, &bob_addr).unwrap()); // alice has not been settled yet - assert_ne!(1000, BalancesStore::load(&deps.storage, &alice_addr)); + assert_ne!(1000, stored_balance(&deps.storage, &alice_addr).unwrap()); let dwb = DWB.load(&deps.storage).unwrap(); println!("DWB: {dwb:?}"); @@ -2480,11 +2504,11 @@ mod tests { .addr_canonicalize(Addr::unchecked("charlie").as_str()) .unwrap(); - assert_eq!(5000 - 1000 - 100, BalancesStore::load(&deps.storage, &bob_addr)); + assert_eq!(5000 - 1000 - 100, stored_balance(&deps.storage, &bob_addr).unwrap()); // alice has not been settled yet - assert_ne!(1000, BalancesStore::load(&deps.storage, &alice_addr)); + assert_ne!(1000, stored_balance(&deps.storage, &alice_addr).unwrap()); // charlie has not been settled yet - assert_ne!(100, BalancesStore::load(&deps.storage, &charlie_addr)); + assert_ne!(100, stored_balance(&deps.storage, &charlie_addr).unwrap()); let dwb = DWB.load(&deps.storage).unwrap(); //println!("DWB: {dwb:?}"); @@ -2514,9 +2538,9 @@ mod tests { let result = handle_result.unwrap(); assert!(ensure_success(result)); - assert_eq!(5000 - 1000 - 100 - 500, BalancesStore::load(&deps.storage, &bob_addr)); + assert_eq!(5000 - 1000 - 100 - 500, stored_balance(&deps.storage, &bob_addr).unwrap()); // make sure alice has not been settled yet - assert_ne!(1500, BalancesStore::load(&deps.storage, &alice_addr)); + assert_ne!(1500, stored_balance(&deps.storage, &alice_addr).unwrap()); let dwb = DWB.load(&deps.storage).unwrap(); //println!("DWB: {dwb:?}"); @@ -2592,13 +2616,13 @@ mod tests { .addr_canonicalize(Addr::unchecked("ernie").as_str()) .unwrap(); - assert_eq!(5000 - 1000 - 100 - 500 - 200, BalancesStore::load(&deps.storage, &bob_addr)); + assert_eq!(5000 - 1000 - 100 - 500 - 200, stored_balance(&deps.storage, &bob_addr).unwrap()); // alice has not been settled yet - assert_ne!(1500, BalancesStore::load(&deps.storage, &alice_addr)); + assert_ne!(1500, stored_balance(&deps.storage, &alice_addr).unwrap()); // charlie has not been settled yet - assert_ne!(100, BalancesStore::load(&deps.storage, &charlie_addr)); + assert_ne!(100, stored_balance(&deps.storage, &charlie_addr).unwrap()); // ernie has not been settled yet - assert_ne!(200, BalancesStore::load(&deps.storage, &ernie_addr)); + assert_ne!(200, stored_balance(&deps.storage, &ernie_addr).unwrap()); let dwb = DWB.load(&deps.storage).unwrap(); //println!("DWB: {dwb:?}"); @@ -2635,9 +2659,9 @@ mod tests { .unwrap(); // alice has been settled - assert_eq!(1500 - 50, BalancesStore::load(&deps.storage, &alice_addr)); + assert_eq!(1500 - 50, stored_balance(&deps.storage, &alice_addr).unwrap()); // dora has not been settled - assert_ne!(50, BalancesStore::load(&deps.storage, &dora_addr)); + assert_ne!(50, stored_balance(&deps.storage, &dora_addr).unwrap()); let dwb = DWB.load(&deps.storage).unwrap(); //println!("DWB: {dwb:?}"); @@ -2671,7 +2695,7 @@ mod tests { let result = handle_result.unwrap(); assert!(ensure_success(result)); } - assert_eq!(5000 - 1000 - 100 - 500 - 200 - 59, BalancesStore::load(&deps.storage, &bob_addr)); + assert_eq!(5000 - 1000 - 100 - 500 - 200 - 59, stored_balance(&deps.storage, &bob_addr).unwrap()); let dwb = DWB.load(&deps.storage).unwrap(); //println!("DWB: {dwb:?}"); @@ -2695,7 +2719,7 @@ mod tests { let result = handle_result.unwrap(); assert!(ensure_success(result)); - assert_eq!(5000 - 1000 - 100 - 500 - 200 - 59 - 1, BalancesStore::load(&deps.storage, &bob_addr)); + assert_eq!(5000 - 1000 - 100 - 500 - 200 - 59 - 1, stored_balance(&deps.storage, &bob_addr).unwrap()); //let dwb = DWB.load(&deps.storage).unwrap(); //println!("DWB: {dwb:?}"); @@ -2716,7 +2740,7 @@ mod tests { let result = handle_result.unwrap(); assert!(ensure_success(result)); - assert_eq!(5000 - 1000 - 100 - 500 - 200 - 59 - 1 - 1, BalancesStore::load(&deps.storage, &bob_addr)); + assert_eq!(5000 - 1000 - 100 - 500 - 200 - 59 - 1 - 1, stored_balance(&deps.storage, &bob_addr).unwrap()); //let dwb = DWB.load(&deps.storage).unwrap(); //println!("DWB: {dwb:?}"); @@ -2740,7 +2764,7 @@ mod tests { assert!(ensure_success(result)); // alice should not settle - assert_eq!(1500 - 50, BalancesStore::load(&deps.storage, &alice_addr)); + assert_eq!(1500 - 50, stored_balance(&deps.storage, &alice_addr).unwrap()); } // alice sends 1 to dora to settle @@ -2759,7 +2783,7 @@ mod tests { let result = handle_result.unwrap(); assert!(ensure_success(result)); - assert_eq!(2724, BalancesStore::load(&deps.storage, &alice_addr)); + assert_eq!(2724, stored_balance(&deps.storage, &alice_addr).unwrap()); // now we send 50 more transactions to alice from bob for i in 1..=50 { @@ -2780,7 +2804,7 @@ mod tests { assert!(ensure_success(result)); // alice should not settle - assert_eq!(2724, BalancesStore::load(&deps.storage, &alice_addr)); + assert_eq!(2724, stored_balance(&deps.storage, &alice_addr).unwrap()); } let handle_msg = ExecuteMsg::SetViewingKey { @@ -3605,8 +3629,8 @@ mod tests { .addr_canonicalize(Addr::unchecked("alice".to_string()).as_str()) .unwrap(); - let bob_balance = BalancesStore::load(&deps.storage, &bob_canonical); - let alice_balance = BalancesStore::load(&deps.storage, &alice_canonical); + let bob_balance = stored_balance(&deps.storage, &bob_canonical).unwrap(); + let alice_balance = stored_balance(&deps.storage, &alice_canonical).unwrap(); assert_eq!(bob_balance, 5000 - 2000); assert_ne!(alice_balance, 2000); let total_supply = TOTAL_SUPPLY.load(&deps.storage).unwrap(); @@ -3748,8 +3772,8 @@ mod tests { .addr_canonicalize(Addr::unchecked("contract".to_string()).as_str()) .unwrap(); - let bob_balance = BalancesStore::load(&deps.storage, &bob_canonical); - let contract_balance = BalancesStore::load(&deps.storage, &contract_canonical); + let bob_balance = stored_balance(&deps.storage, &bob_canonical).unwrap(); + let contract_balance = stored_balance(&deps.storage, &contract_canonical).unwrap(); assert_eq!(bob_balance, 5000 - 2000); assert_ne!(contract_balance, 2000); let total_supply = TOTAL_SUPPLY.load(&deps.storage).unwrap(); @@ -3880,7 +3904,7 @@ mod tests { .addr_canonicalize(Addr::unchecked("bob".to_string()).as_str()) .unwrap(); - let bob_balance = BalancesStore::load(&deps.storage, &bob_canonical); + let bob_balance = stored_balance(&deps.storage, &bob_canonical).unwrap(); assert_eq!(bob_balance, 10000 - 2000); let total_supply = TOTAL_SUPPLY.load(&deps.storage).unwrap(); assert_eq!(total_supply, 10000 - 2000); @@ -4029,7 +4053,7 @@ mod tests { .api .addr_canonicalize(Addr::unchecked(name.to_string()).as_str()) .unwrap(); - let balance = BalancesStore::load(&deps.storage, &name_canon); + let balance = stored_balance(&deps.storage, &name_canon).unwrap(); assert_eq!(balance, 10000 - amount); } let total_supply = TOTAL_SUPPLY.load(&deps.storage).unwrap(); @@ -4063,7 +4087,7 @@ mod tests { .api .addr_canonicalize(Addr::unchecked(name.to_string()).as_str()) .unwrap(); - let balance = BalancesStore::load(&deps.storage, &name_canon); + let balance = stored_balance(&deps.storage, &name_canon).unwrap(); assert_eq!(balance, 10000 - allowance_size); } let total_supply = TOTAL_SUPPLY.load(&deps.storage).unwrap(); @@ -4413,7 +4437,7 @@ mod tests { .api .addr_canonicalize(Addr::unchecked("butler".to_string()).as_str()) .unwrap(); - assert_eq!(BalancesStore::load(&deps.storage, &canonical), 3000) + assert_eq!(stored_balance(&deps.storage, &canonical).unwrap(), 3000) } #[test] @@ -4486,7 +4510,7 @@ mod tests { .unwrap(); // stored balance not updated, still in dwb - assert_ne!(BalancesStore::load(&deps.storage, &canonical), 6000); + assert_ne!(stored_balance(&deps.storage, &canonical).unwrap(), 6000); let create_vk_msg = ExecuteMsg::CreateViewingKey { entropy: "34".to_string(), diff --git a/src/dwb.rs b/src/dwb.rs index a2a327c..8e3abff 100644 --- a/src/dwb.rs +++ b/src/dwb.rs @@ -7,9 +7,7 @@ use cosmwasm_std::{to_binary, Api, Binary, CanonicalAddr, StdError, StdResult, S use secret_toolkit::storage::{AppendStore, Item}; use crate::{ - msg::QueryAnswer, - state::{safe_add, safe_add_u64, BalancesStore,}, - transaction_history::{Tx, TRANSACTIONS}, + dwb, msg::QueryAnswer, state::{safe_add, safe_add_u64,}, stored_balances::{merge_dwb_entry, stored_balance}, transaction_history::{Tx, TRANSACTIONS} }; pub const KEY_DWB: &[u8] = b"dwb"; @@ -46,13 +44,13 @@ pub struct DelayedWriteBuffer { pub entries: [DelayedWriteBufferEntry; DWB_LEN as usize], } -#[inline] -fn random_addr(rng: &mut ContractPrng) -> CanonicalAddr { - #[cfg(test)] - return CanonicalAddr::from(&[rng.rand_bytes(), rng.rand_bytes()].concat()[0..DWB_RECIPIENT_BYTES]); // because mock canonical addr is 54 bytes - #[cfg(not(test))] - CanonicalAddr::from(&rng.rand_bytes()[0..DWB_RECIPIENT_BYTES]) // canonical addr is 20 bytes (less than 32) -} +//#[inline] +//fn random_addr(rng: &mut ContractPrng) -> CanonicalAddr { +// #[cfg(test)] +// return CanonicalAddr::from(&[rng.rand_bytes(), rng.rand_bytes()].concat()[0..DWB_RECIPIENT_BYTES]); // because mock canonical addr is 54 bytes +// #[cfg(not(test))] +// CanonicalAddr::from(&rng.rand_bytes()[0..DWB_RECIPIENT_BYTES]) // canonical addr is 20 bytes (less than 32) +//} pub fn random_in_range(rng: &mut ContractPrng, a: u32, b: u32) -> StdResult { if b <= a { @@ -82,12 +80,14 @@ impl DelayedWriteBuffer { } /// settles an entry at a given index in the buffer + #[inline] fn settle_entry( - &mut self, + &self, store: &mut dyn Storage, index: usize, ) -> StdResult<()> { - let entry = self.entries[index]; + merge_dwb_entry(store, self.entries[index], None) +/* let account = entry.recipient()?; AccountTxsStore::append_bundle( @@ -102,6 +102,7 @@ impl DelayedWriteBuffer { safe_add(&mut balance, entry.amount()? as u128); // add the amount from entry to the stored balance BalancesStore::save(store, &account, balance) +*/ } /// settles a participant's account who may or may not have an entry in the buffer @@ -109,19 +110,27 @@ impl DelayedWriteBuffer { pub fn settle_sender_or_owner_account( &mut self, store: &mut dyn Storage, - rng: &mut ContractPrng, address: &CanonicalAddr, tx_id: u64, amount_spent: u128, op_name: &str, ) -> StdResult<()> { // release the address from the buffer - let (balance, mut entry) = self.constant_time_release( + let (balance, mut dwb_entry) = self.constant_time_release( store, - rng, address )?; + if balance.checked_sub(amount_spent).is_none() { + return Err(StdError::generic_err(format!( + "insufficient funds to {op_name}: balance={balance}, required={amount_spent}", + ))); + }; + + dwb_entry.add_tx_node(store, tx_id)?; + + merge_dwb_entry(store, dwb_entry, Some(amount_spent)) +/* let head_node = entry.add_tx_node(store, tx_id)?; AccountTxsStore::append_bundle( @@ -131,16 +140,10 @@ impl DelayedWriteBuffer { entry.list_len()?, )?; - let new_balance = if let Some(balance_after_sub) = balance.checked_sub(amount_spent) { - balance_after_sub - } else { - return Err(StdError::generic_err(format!( - "insufficient funds to {op_name}: balance={balance}, required={amount_spent}", - ))); - }; BalancesStore::save(store, address, new_balance)?; Ok(()) +*/ } /// "releases" a given recipient from the buffer, removing their entry if one exists, in constant-time @@ -148,11 +151,10 @@ impl DelayedWriteBuffer { fn constant_time_release( &mut self, store: &mut dyn Storage, - rng: &mut ContractPrng, address: &CanonicalAddr ) -> StdResult<(u128, DelayedWriteBufferEntry)> { // get the address' stored balance - let mut balance = BalancesStore::load(store, address); + let mut balance = stored_balance(store, address)?; // locate the position of the entry in the buffer let matched_entry_idx = self.recipient_match(address); @@ -172,15 +174,15 @@ impl DelayedWriteBuffer { Ok((balance, entry)) } - fn unique_random_entry(&self, rng: &mut ContractPrng) -> StdResult { - // produce a new random address - let mut replacement_address = random_addr(rng); - // ensure random addr is not already in dwb (extremely unlikely!!) - while self.recipient_match(&replacement_address) > 0 { - replacement_address = random_addr(rng); - } - DelayedWriteBufferEntry::new(replacement_address) - } + //fn unique_random_entry(&self, rng: &mut ContractPrng) -> StdResult { + // // produce a new random address + // let mut replacement_address = random_addr(rng); + // // ensure random addr is not already in dwb (extremely unlikely!!) + // while self.recipient_match(&replacement_address) > 0 { + // replacement_address = random_addr(rng); + // } + // DelayedWriteBufferEntry::new(replacement_address) + //} // returns matched index for a given address pub fn recipient_match(&self, address: &CanonicalAddr) -> usize { @@ -413,9 +415,7 @@ impl DelayedWriteBufferEntry { fn add_amount(&mut self, add_tx_amount: u128) -> StdResult { // change this to safe_add if your coin needs to store amount in buffer as u128 (e.g. 18 decimals) let mut amount = self.amount()?; - let add_tx_amount_u64 = add_tx_amount - .try_into() - .or_else(|_| return Err(StdError::generic_err("dwb: deposit overflow")))?; + let add_tx_amount_u64 = amount_u64(Some(add_tx_amount))?; safe_add_u64(&mut amount, add_tx_amount_u64); self.set_amount(amount)?; @@ -423,6 +423,14 @@ impl DelayedWriteBufferEntry { } } +pub fn amount_u64(amount_spent: Option) -> StdResult { + let amount_spent = amount_spent.unwrap_or_default(); + let amount_spent_u64 = amount_spent + .try_into() + .or_else(|_| return Err(StdError::generic_err("se: spent overflow")))?; + Ok(amount_spent_u64) +} + #[derive(Serialize, Deserialize, Clone, Copy, Debug)] pub struct TxNode { /// transaction id in the TRANSACTIONS list @@ -471,70 +479,11 @@ pub struct TxBundle { /// The bundle points to a linked list of transaction nodes, which each reference /// a transaction record by its global id. /// used with add_suffix(canonical addr of account) -pub static ACCOUNT_TXS: AppendStore = AppendStore::new(KEY_ACCOUNT_TXS); +//pub static ACCOUNT_TXS: AppendStore = AppendStore::new(KEY_ACCOUNT_TXS); /// Keeps track of the total count of txs for an account (not tx bundles) /// used with add_suffix(canonical addr of account) -pub static ACCOUNT_TX_COUNT: Item = Item::new(KEY_ACCOUNT_TX_COUNT); - -pub struct AccountTxsStore {} -impl AccountTxsStore { - /// appends a new tx bundle for an account, called when non-transfer tx occurs or is settled. - pub fn append_bundle(store: &mut dyn Storage, account: &CanonicalAddr, head_node: u64, list_len: u16) -> StdResult<()> { - let account_txs_store = ACCOUNT_TXS.add_suffix(account.as_slice()); - let account_txs_len = account_txs_store.get_len(store)?; - let tx_bundle; - if account_txs_len > 0 { - // peek at the last tx bundle added - let last_tx_bundle = account_txs_store.get_at(store, account_txs_len - 1)?; - tx_bundle = TxBundle { - head_node, - list_len, - offset: last_tx_bundle.offset + u32::from(last_tx_bundle.list_len), - }; - } else { // this is the first bundle for the account - tx_bundle = TxBundle { - head_node, - list_len, - offset: 0, - }; - } - - // update the total count of txs for account - let account_tx_count_store = ACCOUNT_TX_COUNT.add_suffix(account.as_slice()); - let account_tx_count = account_tx_count_store.may_load(store)?.unwrap_or_default(); - account_tx_count_store.save(store, &(account_tx_count.saturating_add(u32::from(list_len))))?; - - account_txs_store.push(store, &tx_bundle) - } - - /// Does a binary search on the append store to find the bundle where the `start_idx` tx can be found. - /// For a paginated search `start_idx` = `page` * `page_size`. - /// Returns the bundle index, the bundle, and the index in the bundle list to start at - pub fn find_start_bundle(store: &dyn Storage, account: &CanonicalAddr, start_idx: u32) -> StdResult> { - let account_txs_store = ACCOUNT_TXS.add_suffix(account.as_slice()); - - let mut left = 0u32; - let mut right = account_txs_store.get_len(store)?; - - while left <= right { - let mid = (left + right) / 2; - let mid_bundle = account_txs_store.get_at(store, mid)?; - if start_idx >= mid_bundle.offset && start_idx < mid_bundle.offset + (mid_bundle.list_len as u32) { - // we have the correct bundle - // which index in list to start at? - let start_at = (mid_bundle.list_len as u32) - (start_idx - mid_bundle.offset) - 1; - return Ok(Some((mid, mid_bundle, start_at))); - } else if start_idx < mid_bundle.offset { - right = mid - 1; - } else { - left = mid + 1; - } - } - - Ok(None) - } -} +//pub static ACCOUNT_TX_COUNT: Item = Item::new(KEY_ACCOUNT_TX_COUNT); #[inline] fn constant_time_is_not_zero(value: i32) -> u32 { diff --git a/src/state.rs b/src/state.rs index 09abb2d..30e927e 100644 --- a/src/state.rs +++ b/src/state.rs @@ -132,47 +132,6 @@ pub fn safe_add_u64(balance: &mut u64, amount: u64) -> u64 { *balance - prev_balance } -pub static BALANCES: Item = Item::new(PREFIX_BALANCES); -pub struct BalancesStore {} -impl BalancesStore { - pub fn save(store: &mut dyn Storage, account: &CanonicalAddr, amount: u128) -> StdResult<()> { - let balances = BALANCES.add_suffix(account.as_slice()); - balances.save(store, &amount) - } - - pub fn load(store: &dyn Storage, account: &CanonicalAddr) -> u128 { - let balances = BALANCES.add_suffix(account.as_slice()); - balances.load(store).unwrap_or_default() - } - - pub fn update_balance( - store: &mut dyn Storage, - account: &CanonicalAddr, - amount_to_be_updated: u128, - should_add: bool, - operation_name: &str, - ) -> StdResult<()> { - let mut balance = Self::load(store, account); - balance = match should_add { - true => { - safe_add(&mut balance, amount_to_be_updated); - balance - } - false => { - if let Some(balance) = balance.checked_sub(amount_to_be_updated) { - balance - } else { - return Err(StdError::generic_err(format!( - "insufficient funds to {operation_name}: balance={balance}, required={amount_to_be_updated}", - ))); - } - } - }; - - Self::save(store, account, balance) - } -} - // Allowances #[derive(Serialize, Debug, Deserialize, Clone, PartialEq, Eq, Default, JsonSchema)] @@ -277,3 +236,5 @@ impl ReceiverHashStore { receiver_hash.save(store, &code_hash) } } + +pub static INTERNAL_SECRET: Item> = Item::new(b"internal-secret"); \ No newline at end of file diff --git a/src/stored_balances.rs b/src/stored_balances.rs index 6af94b6..44e078b 100644 --- a/src/stored_balances.rs +++ b/src/stored_balances.rs @@ -1,11 +1,12 @@ use constant_time_eq::constant_time_eq; -use secret_toolkit::storage::Item; +use primitive_types::U256; +use secret_toolkit::{notification::hkdf_sha_256, serialization::{Bincode2, Serde}, storage::Item}; use serde::{Serialize, Deserialize,}; use serde_big_array::BigArray; use cosmwasm_std::{CanonicalAddr, StdError, StdResult, Storage}; use crate::{ - dwb::DelayedWriteBufferEntry, state::safe_add_u64 + dwb::{amount_u64, DelayedWriteBufferEntry, TxBundle}, state::{safe_add_u64, INTERNAL_SECRET} }; // btsb = bitwise-trie of stored balances @@ -16,7 +17,6 @@ pub const KEY_BTSB_BUCKETS: &[u8] = b"btsb-buckets"; pub const KEY_BTSB_TRIE_NODES: &[u8] = b"btsb-trie-nodes"; pub const KEY_BTSB_TRIE_NODES_COUNT: &[u8] = b"btsb-trie-nodes-cnt"; - const U16_BYTES: usize = 2; const U32_BYTES: usize = 4; const U64_BYTES: usize = 8; @@ -27,24 +27,34 @@ const BTSB_BUCKET_ADDRESS_BYTES: usize = 54; #[cfg(not(test))] const BTSB_BUCKET_ADDRESS_BYTES: usize = 20; const BTSB_BUCKET_BALANCE_BYTES: usize = 8; // Max 16 (u128) -const BTSB_BUCKET_HISTORY_BYTES: usize = 5; // Max 8 (u64) +const BTSB_BUCKET_HISTORY_BYTES: usize = 4; // Max 4 (u32) const_assert!(BTSB_BUCKET_BALANCE_BYTES <= U128_BYTES); -const_assert!(BTSB_BUCKET_HISTORY_BYTES <= U64_BYTES); +const_assert!(BTSB_BUCKET_HISTORY_BYTES <= U32_BYTES); const BTSB_BUCKET_ENTRY_BYTES: usize = BTSB_BUCKET_ADDRESS_BYTES + BTSB_BUCKET_BALANCE_BYTES + BTSB_BUCKET_HISTORY_BYTES; const ZERO_ADDR: [u8; BTSB_BUCKET_ADDRESS_BYTES] = [0u8; BTSB_BUCKET_ADDRESS_BYTES]; - -#[derive(Serialize, Deserialize, Clone, Copy, Debug)] -#[cfg_attr(test, derive(Eq, PartialEq))] -pub struct StoredBalanceEntry( +/// canonical address bytes corresponding to the 33-byte null public key, in hexadecimal +#[cfg(test)] +const IMPOSSIBLE_ADDR: [u8; BTSB_BUCKET_ADDRESS_BYTES] = [0x29,0xCF,0xC6,0x37,0x62,0x55,0xA7,0x84,0x51,0xEE,0xB4,0xB1,0x29,0xED,0x8E,0xAC,0xFF,0xA2,0xFE,0xEF, + 0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00, + 0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00]; +#[cfg(not(test))] +const IMPOSSIBLE_ADDR: [u8; BTSB_BUCKET_ADDRESS_BYTES] = [0x29,0xCF,0xC6,0x37,0x62,0x55,0xA7,0x84,0x51,0xEE,0xB4,0xB1,0x29,0xED,0x8E,0xAC,0xFF,0xA2,0xFE,0xEF]; + +/// A `StoredEntry` consists of the address, balance, and tx bundle history length in a byte array representation. +/// The methods of the struct implementation also handle pushing and getting the tx bundle history in a simplified +/// append store. +#[derive(Serialize, Deserialize, Clone, Copy, Debug, PartialEq)] +#[cfg_attr(test, derive(Eq))] +pub struct StoredEntry( #[serde(with = "BigArray")] - [u8; BTSB_BUCKET_ENTRY_BYTES] + [u8; BTSB_BUCKET_ENTRY_BYTES], ); -impl StoredBalanceEntry { - pub fn new(address: CanonicalAddr) -> StdResult { +impl StoredEntry { + fn new(address: CanonicalAddr) -> StdResult { let address = address.as_slice(); if address.len() != BTSB_BUCKET_ADDRESS_BYTES { @@ -52,17 +62,35 @@ impl StoredBalanceEntry { } let mut result = [0u8; BTSB_BUCKET_ENTRY_BYTES]; - result[..BTSB_BUCKET_ENTRY_BYTES].copy_from_slice(address); + result[..BTSB_BUCKET_ADDRESS_BYTES].copy_from_slice(address); Ok(Self { - 0: result + 0: result, }) } - pub fn from(dwb_entry: DelayedWriteBufferEntry) -> StdResult { - let mut entry = StoredBalanceEntry::new(dwb_entry.recipient()?)?; - - entry.set_balace(dwb_entry.amount()?); - entry.set_history_len(1); + fn from(storage: &mut dyn Storage, dwb_entry: &DelayedWriteBufferEntry, amount_spent: Option) -> StdResult { + let mut entry = StoredEntry::new(dwb_entry.recipient()?)?; + + let amount_spent = amount_u64(amount_spent)?; + + // error should never happen because already checked in `settle_sender_or_owner_account` + let balance = if let Some(new_balance) = dwb_entry.amount()?.checked_sub(amount_spent) { + new_balance + } else { + return Err(StdError::generic_err(format!( + "insufficient funds", + ))); + }; + + entry.set_balance(balance)?; + entry.push_tx_bundle( + storage, + &TxBundle { + head_node: dwb_entry.head_node()?, + list_len: dwb_entry.list_len()?, + offset: 0, + } + )?; Ok(entry) } @@ -87,26 +115,26 @@ impl StoredBalanceEntry { Ok(u64::from_be_bytes(result)) } - fn set_balace(&mut self, val: u64) -> StdResult<()> { + fn set_balance(&mut self, val: u64) -> StdResult<()> { let start = BTSB_BUCKET_ADDRESS_BYTES; let end = start + BTSB_BUCKET_BALANCE_BYTES; self.0[start..end].copy_from_slice(&val.to_be_bytes()); Ok(()) } - pub fn history_len(&self) -> StdResult { + pub fn history_len(&self) -> StdResult { let start = BTSB_BUCKET_ADDRESS_BYTES + BTSB_BUCKET_BALANCE_BYTES; let end = start + BTSB_BUCKET_HISTORY_BYTES; let history_len_slice = &self.0[start..end]; - let mut result = [0u8; U64_BYTES]; - result[U64_BYTES - BTSB_BUCKET_HISTORY_BYTES..].copy_from_slice(history_len_slice); - Ok(u64::from_be_bytes(result)) + let mut result = [0u8; U32_BYTES]; + result[U32_BYTES - BTSB_BUCKET_HISTORY_BYTES..].copy_from_slice(history_len_slice); + Ok(u32::from_be_bytes(result)) } - fn set_history_len(&mut self, val: u64) -> StdResult<()> { + fn set_history_len(&mut self, val: u32) -> StdResult<()> { let start = BTSB_BUCKET_ADDRESS_BYTES + BTSB_BUCKET_BALANCE_BYTES; let end = start + BTSB_BUCKET_HISTORY_BYTES; - let val_bytes = &val.to_be_bytes()[U64_BYTES - BTSB_BUCKET_HISTORY_BYTES..]; + let val_bytes = &val.to_be_bytes()[U32_BYTES - BTSB_BUCKET_HISTORY_BYTES..]; if val_bytes.len() != BTSB_BUCKET_HISTORY_BYTES { return Err(StdError::generic_err("Set bucket history len error")); } @@ -114,30 +142,92 @@ impl StoredBalanceEntry { Ok(()) } - pub fn merge_dwb_entry(&mut self, entry: &DelayedWriteBufferEntry) -> StdResult<()> { - let mut balance = self.balance()?; - safe_add_u64(&mut balance, entry.amount()?); - self.set_balace(balance)?; + pub fn merge_dwb_entry( + &mut self, + storage: &mut dyn Storage, + dwb_entry: &DelayedWriteBufferEntry, + amount_spent: Option + ) -> StdResult<()> { + let history_len = self.history_len()?; + if history_len == 0 { + return Err(StdError::generic_err("use `from` to create new entry from dwb_entry")); + } - // TOOD: update history len + let mut balance = self.balance()?; + safe_add_u64(&mut balance, dwb_entry.amount()?); + + let amount_spent = amount_u64(amount_spent)?; + + // error should never happen because already checked in `settle_sender_or_owner_account` + let balance = if let Some(new_balance) = balance.checked_sub(amount_spent) { + new_balance + } else { + return Err(StdError::generic_err(format!( + "insufficient funds", + ))); + }; + + self.set_balance(balance)?; + + // peek at the last tx bundle added + let last_tx_bundle = self.get_tx_bundle_at(storage, history_len - 1)?; + let tx_bundle = TxBundle { + head_node: dwb_entry.head_node()?, + list_len: dwb_entry.list_len()?, + offset: last_tx_bundle.offset + u32::from(last_tx_bundle.list_len), + }; + self.push_tx_bundle(storage, &tx_bundle)?; Ok(()) } -} + // simplified appendstore impl for tx history + + /// gets the element at pos if within bounds + pub fn get_tx_bundle_at(&self, storage: &dyn Storage, pos: u32) -> StdResult { + let len = self.history_len()?; + if pos >= len { + return Err(StdError::generic_err("access out of bounds")); + } + self.get_tx_bundle_at_unchecked(storage, pos) + } + + /// tries to get the element at pos + fn get_tx_bundle_at_unchecked(&self, storage: &dyn Storage, pos: u32) -> StdResult { + let bundle_data = storage.get(&[KEY_BTSB_ENTRY_HISTORY, self.address_slice(), pos.to_be_bytes().as_slice()].concat()); + let bundle_data = bundle_data.ok_or_else(|| { return StdError::generic_err("tx bundle not found"); } )?; + Bincode2::deserialize( + &bundle_data + ) + } + /// Sets data at a given index + fn set_tx_bundle_at_unchecked(&self, storage: &mut dyn Storage, pos: u32, bundle: &TxBundle) -> StdResult<()> { + let bundle_data = Bincode2::serialize(bundle)?; + storage.set(&[KEY_BTSB_ENTRY_HISTORY, self.address_slice(), pos.to_be_bytes().as_slice()].concat(), &bundle_data); + Ok(()) + } + + /// Pushes a tx bundle + fn push_tx_bundle(&mut self, storage: &mut dyn Storage, bundle: &TxBundle) -> StdResult<()> { + let len = self.history_len()?; + self.set_tx_bundle_at_unchecked(storage, len, bundle)?; + self.set_history_len(len.saturating_add(1))?; + Ok(()) + } +} const BTSB_BUCKET_LEN: u16 = 128; -#[derive(Serialize, Deserialize, Clone, Copy, Debug)] -struct BtsbBucket { +#[derive(Serialize, Deserialize, Clone, Copy, Debug, PartialEq)] +pub struct BtsbBucket { pub capacity: u16, #[serde(with = "BigArray")] - pub entries: [StoredBalanceEntry; BTSB_BUCKET_LEN as usize], + pub entries: [StoredEntry; BTSB_BUCKET_LEN as usize], } -static BTSB_ENTRY_HISTORY: Item = Item::new(KEY_BTSB_ENTRY_HISTORY); +//static BTSB_ENTRY_HISTORY: Item = Item::new(KEY_BTSB_ENTRY_HISTORY); static BTSB_BUCKETS_COUNT: Item = Item::new(KEY_BTSB_BUCKETS_COUNT); static BTSB_BUCKETS: Item = Item::new(KEY_BTSB_BUCKETS); @@ -149,31 +239,28 @@ impl BtsbBucket { Ok(Self { capacity: BTSB_BUCKET_LEN, entries: [ - StoredBalanceEntry::new(CanonicalAddr::from(&ZERO_ADDR))?; BTSB_BUCKET_LEN as usize + StoredEntry::new(CanonicalAddr::from(&IMPOSSIBLE_ADDR))?; BTSB_BUCKET_LEN as usize ] }) } - pub fn add_entry(&mut self, storage: &mut dyn Storage, entry: &StoredBalanceEntry, bit_pos: u8) -> StdResult { - match self.capacity { + pub fn add_entry(&mut self, storage: &mut dyn Storage, entry: &StoredEntry) -> bool { + if self.capacity == 0 { // buffer is at capacity - 0 => Err(StdError::generic_err("")), - - // has capacity for a new entry - _ => { - // save entry to bucket - self.entries[self.entries.len() - self.capacity as usize] = entry.clone(); + return false; + } + // has capacity for a new entry + // save entry to bucket + self.entries[self.entries.len() - self.capacity as usize] = entry.clone(); - // update capacity - self.capacity -= 1; + // update capacity + self.capacity -= 1; - // done - Ok(self.capacity) - } - } + // done + true } - pub fn constant_time_find_address(&self, address: &CanonicalAddr) -> Option { + pub fn constant_time_find_address(&self, address: &CanonicalAddr) -> Option<(usize, StoredEntry)> { let address = address.as_slice(); let mut matched_index_p1: BucketEntryPosition = 0; @@ -184,11 +271,11 @@ impl BtsbBucket { match matched_index_p1 { 0 => None, - idx => Some(self.entries[idx - 1]), + idx => Some((idx - 1, self.entries[idx - 1])), } } - pub fn quick_find_entry(&self, address: &CanonicalAddr) -> Option { + pub fn quick_find_entry(&self, address: &CanonicalAddr) -> Option { let address = address.as_slice(); let mut matched_index_p1: BucketEntryPosition = 0; @@ -210,7 +297,6 @@ pub struct BitwiseTrieNode { pub bucket: u64, } - pub static BTSB_TRIE_NODES: Item = Item::new(KEY_BTSB_TRIE_NODES); pub static BTSB_TRIE_NODES_COUNT: Item = Item::new(KEY_BTSB_TRIE_NODES_COUNT); @@ -237,7 +323,7 @@ impl BitwiseTrieNode { } // loads the node's bucket from storage - pub fn bucket(self, storage: &mut dyn Storage) -> StdResult { + pub fn bucket(self, storage: &dyn Storage) -> StdResult { if self.bucket == 0 { return Err(StdError::generic_err("btsb: attempted to load bucket of branch node")); } @@ -258,9 +344,13 @@ impl BitwiseTrieNode { // locates a btsb node given an address; returns tuple of (node, bit position) -pub fn locate_btsb_node(storage: &mut dyn Storage, address: &CanonicalAddr) -> StdResult<(BitwiseTrieNode, u64, u8)> { - let hash: [u8; 32] = [0u8; 32]; - /* TODO: +pub fn locate_btsb_node(storage: &dyn Storage, address: &CanonicalAddr) -> StdResult<(BitwiseTrieNode, u64, u8)> { + //let hash: [u8; 32] = [0u8; 32]; + + let secret = INTERNAL_SECRET.load(storage)?; + let secret = secret.as_slice(); + let hash = hkdf_sha_256(&None, secret, address.as_slice(), 256)?; + /* let hash := hkdf(ikm=contractInternalSecret, info=addrress, length=256bits) */ @@ -287,123 +377,194 @@ pub fn locate_btsb_node(storage: &mut dyn Storage, address: &CanonicalAddr) -> S Ok((node, node_id, bit_pos)) } +/// Does a binary search on the append store to find the bundle where the `start_idx` tx can be found. +/// For a paginated search `start_idx` = `page` * `page_size`. +/// Returns the bundle index, the bundle, and the index in the bundle list to start at +pub fn find_start_bundle(storage: &dyn Storage, account: &CanonicalAddr, start_idx: u32) -> StdResult> { + let (node, _, _) = locate_btsb_node(storage, account)?; + let bucket = node.bucket(storage)?; + if let Some((_, entry)) = bucket.constant_time_find_address(account) { + let mut left = 0u32; + let mut right = entry.history_len()?; + + while left <= right { + let mid = (left + right) / 2; + let mid_bundle = entry.get_tx_bundle_at(storage, mid)?; + if start_idx >= mid_bundle.offset && start_idx < mid_bundle.offset + (mid_bundle.list_len as u32) { + // we have the correct bundle + // which index in list to start at? + let start_at = (mid_bundle.list_len as u32) - (start_idx - mid_bundle.offset) - 1; + return Ok(Some((mid, mid_bundle, start_at))); + } else if start_idx < mid_bundle.offset { + right = mid - 1; + } else { + left = mid + 1; + } + } + } + + Ok(None) +} + +/// gets the StoredEntry for a given account +pub fn stored_entry(storage: &dyn Storage, account: &CanonicalAddr) -> StdResult> { + let (node, _, _) = locate_btsb_node(storage, account)?; + let bucket = node.bucket(storage)?; + Ok(bucket.constant_time_find_address(account).map(|b| b.1)) +} + +/// returns the current stored balance for an entry +pub fn stored_balance(storage: &dyn Storage, address: &CanonicalAddr) -> StdResult { + if let Some(entry) = stored_entry(storage, address)? { + Ok(entry.balance()? as u128) + } else { + Ok(0_u128) + } +} + +/// Returns the total number of settled transactions for an account by peeking at last bundle +pub fn stored_tx_count(storage: &dyn Storage, entry: &Option) -> StdResult { + if let Some(entry) = entry { + // peek at last entry + let len = entry.history_len()?; + if len > 0 { + let bundle = entry.get_tx_bundle_at(storage, len - 1)?; + return Ok(bundle.offset + bundle.list_len as u32); + } + } + Ok(0) +} + // merges a dwb entry into the current node's bucket -pub fn merge_dwb_entry(storage: &mut dyn Storage, dwb_entry: DelayedWriteBufferEntry) -> StdResult<()> { +// `spent_amount` is any required subtraction due to being sender of tx +pub fn merge_dwb_entry(storage: &mut dyn Storage, dwb_entry: DelayedWriteBufferEntry, amount_spent: Option) -> StdResult<()> { // locate the node that the given entry belongs in - let (mut node, node_id, bit_pos) = locate_btsb_node(storage, &dwb_entry.recipient()?)?; + let (mut node, node_id, mut bit_pos) = locate_btsb_node(storage, &dwb_entry.recipient()?)?; // load that node's current bucket let mut bucket = node.bucket(storage)?; // search for an existing entry - match bucket.constant_time_find_address(&dwb_entry.recipient()?) { + if let Some((idx, mut found_entry)) = bucket.constant_time_find_address(&dwb_entry.recipient()?) { // found existing entry - Some(mut found_entry) => { - // merge amount and history from dwb entry - found_entry.merge_dwb_entry(&dwb_entry); - - // save updated bucket to storage - node.set_and_save_bucket(storage, bucket); - }, + // merge amount and history from dwb entry + found_entry.merge_dwb_entry(storage, &dwb_entry, amount_spent)?; + bucket.entries[idx] = found_entry; + // save updated bucket to storage + node.set_and_save_bucket(storage, bucket)?; + } else { // need to insert new entry - None => { - // create new stored balance entry - let btsb_entry = StoredBalanceEntry::from(dwb_entry)?; + // create new stored balance entry + let btsb_entry = StoredEntry::from(storage, &dwb_entry, amount_spent)?; - /* TODO: - create new storage for dwb_entry's history - */ + let secret = INTERNAL_SECRET.load(storage)?; + let secret = secret.as_slice(); + loop { // looping as many times as needed until the bucket has capacity for a new entry // try to add to the current bucket - match bucket.add_entry(storage, &btsb_entry, bit_pos) { - // bucket has capcity and it added the new entry - Ok(capacity) => { - // save bucket to storage - node.set_and_save_bucket(storage, bucket); - } - + if bucket.add_entry(storage, &btsb_entry) { + // bucket has capacity and it added the new entry + // save bucket to storage + node.set_and_save_bucket(storage, bucket)?; + // break out of the loop + break; + } else { // bucket is full; split on next bit position - Err(_) => { - // create new left and right buckets - let left_bucket = BtsbBucket::new()?; - let right_bucket = BtsbBucket::new()?; - - // each entry - for (idx, entry) in bucket.entries.iter().enumerate() { - /* TODO: - let key := hkdf(ikm=contractInternalSecret, info=canonical(addr), length=256bits) - let bit_value := (key >> (255 - bit_pos)) & 1 - if bit_value == 0: - left_bucket.add_entry(entry) - else: - right_bucket.add_entry(entry) - */ + // create new left and right buckets + let mut left_bucket = BtsbBucket::new()?; + let mut right_bucket = BtsbBucket::new()?; + + // each entry + for entry in bucket.entries { + let key = hkdf_sha_256(&None, secret, entry.address_slice(), 256)?; + let key = U256::from_big_endian(&key); + let bit_value = (key >> (255 - bit_pos)) & U256::from(1); + if bit_value == U256::from(0) { + left_bucket.add_entry(storage, &entry); + } else { + right_bucket.add_entry(storage, &entry); } + /* + let key := hkdf(ikm=contractInternalSecret, info=canonical(addr), length=256bits) + let bit_value := (key >> (255 - bit_pos)) & 1 + if bit_value == 0: + left_bucket.add_entry(entry) + else: + right_bucket.add_entry(entry) + */ + } - // save left node's bucket to storage, recycling this node's bucket ID - let left_bucket_id = node.bucket; - BTSB_BUCKETS.add_suffix(&left_bucket_id.to_be_bytes()).save(storage, &left_bucket); - - // global count of buckets - let mut buckets_count = BTSB_BUCKETS_COUNT.load(storage).unwrap_or_default(); - - // bucket ID for right node - buckets_count += 1; - let right_bucket_id = buckets_count; - BTSB_BUCKETS.add_suffix(&right_bucket_id.to_be_bytes()).save(storage, &right_bucket); - - // save updated count - BTSB_BUCKETS_COUNT.save(storage, &buckets_count)?; - - // globl count of trie nodes - let mut nodes_count = BTSB_TRIE_NODES_COUNT.load(storage).unwrap_or_default(); - - // ID for left node - nodes_count += 1; - let left_id = nodes_count; - - // ID for right node - nodes_count += 1; - let right_id = nodes_count; - - // save updated count - BTSB_TRIE_NODES_COUNT.save(storage, &nodes_count)?; - - // create left and right nodes - let left = BitwiseTrieNode { - left: 0, - right: 0, - bucket: left_bucket_id, - }; - let right = BitwiseTrieNode { - left: 0, - right: 0, - bucket: right_bucket_id, - }; - - // save left and right node to storage - BTSB_TRIE_NODES.add_suffix(&left_id.to_be_bytes()).save(storage, &left)?; - BTSB_TRIE_NODES.add_suffix(&right_id.to_be_bytes()).save(storage, &right)?; - - // convert this into a branch node - node.left = left_id; - node.right = right_id; - node.bucket = 0; - - // save node - BTSB_TRIE_NODES.add_suffix(&node_id.to_be_bytes()).save(storage, &node); - - // -- - - /* TODO - determine which child node the dwb entry belongs in, then retry insertion, - looping as many times as needed until the bucket has capacity for a new entry - */ + // save left node's bucket to storage, recycling this node's bucket ID + let left_bucket_id = node.bucket; + BTSB_BUCKETS.add_suffix(&left_bucket_id.to_be_bytes()).save(storage, &left_bucket)?; + + // global count of buckets + let mut buckets_count = BTSB_BUCKETS_COUNT.load(storage).unwrap_or_default(); + + // bucket ID for right node + buckets_count += 1; + let right_bucket_id = buckets_count; + BTSB_BUCKETS.add_suffix(&right_bucket_id.to_be_bytes()).save(storage, &right_bucket)?; + + // save updated count + BTSB_BUCKETS_COUNT.save(storage, &buckets_count)?; + + // global count of trie nodes + let mut nodes_count = BTSB_TRIE_NODES_COUNT.load(storage).unwrap_or_default(); + + // ID for left node + nodes_count += 1; + let left_id = nodes_count; + + // ID for right node + nodes_count += 1; + let right_id = nodes_count; + + // save updated count + BTSB_TRIE_NODES_COUNT.save(storage, &nodes_count)?; + + // create left and right nodes + let left = BitwiseTrieNode { + left: 0, + right: 0, + bucket: left_bucket_id, + }; + let right = BitwiseTrieNode { + left: 0, + right: 0, + bucket: right_bucket_id, + }; + + // save left and right node to storage + BTSB_TRIE_NODES.add_suffix(&left_id.to_be_bytes()).save(storage, &left)?; + BTSB_TRIE_NODES.add_suffix(&right_id.to_be_bytes()).save(storage, &right)?; + + // convert this into a branch node + node.left = left_id; + node.right = right_id; + node.bucket = 0; + + // save node + BTSB_TRIE_NODES.add_suffix(&node_id.to_be_bytes()).save(storage, &node)?; + + let key = hkdf_sha_256(&None, secret, btsb_entry.address_slice(), 256)?; + let key = U256::from_big_endian(&key); + let bit_value = (key >> (255 - bit_pos)) & U256::from(1); + + // determine which child node the dwb entry belongs in, then retry insertion + if bit_value == U256::from(0) { + node = left; + bucket = left_bucket; + } else { + node = right; + bucket = right_bucket; } + // increment bit position for next iteration of the loop + bit_pos += 1; } - }, + } } Ok(()) @@ -411,15 +572,32 @@ pub fn merge_dwb_entry(storage: &mut dyn Storage, dwb_entry: DelayedWriteBufferE // for fetching an account's stored balance during transfer executions -pub fn constant_time_get_btsb_entry(storage: &mut dyn Storage, address: CanonicalAddr) -> StdResult> { - let (mut node, node_id, bit_pos) = locate_btsb_node(storage, &address)?; +pub fn constant_time_get_btsb_entry(storage: &mut dyn Storage, address: CanonicalAddr) -> StdResult> { + let (node, _, _) = locate_btsb_node(storage, &address)?; - Ok(node.bucket(storage)?.constant_time_find_address(&address)) + Ok(node.bucket(storage)?.constant_time_find_address(&address).map(|b| b.1)) } // for fetching account's stored balance and/or history during queries -pub fn quick_get_btsb_entry(storage: &mut dyn Storage, address: CanonicalAddr) -> StdResult> { - let (mut node, node_id, bit_pos) = locate_btsb_node(storage, &address)?; +pub fn quick_get_btsb_entry(storage: &mut dyn Storage, address: CanonicalAddr) -> StdResult> { + let (node, _, _) = locate_btsb_node(storage, &address)?; Ok(node.bucket(storage)?.quick_find_entry(&address)) } + +/// initializes the btsb +pub fn initialize_btsb(storage: &mut dyn Storage) -> StdResult<()> { + let btsb_bucket = BtsbBucket::new()?; + BTSB_BUCKETS.add_suffix(&1_u64.to_be_bytes()).save(storage, &btsb_bucket)?; + BTSB_BUCKETS_COUNT.save(storage, &1)?; + BTSB_TRIE_NODES.add_suffix(&1_u64.to_be_bytes()).save( + storage, + &BitwiseTrieNode{ + left: 0, + right: 0, + bucket: 1, + } + )?; + BTSB_TRIE_NODES_COUNT.save(storage, &1)?; + Ok(()) +} \ No newline at end of file