From 333bf0e518cdc6ad8c598f06e4c649f0dd0b7a9f Mon Sep 17 00:00:00 2001 From: xsb <389483673@qq.com> Date: Mon, 23 Jun 2025 20:51:24 +0800 Subject: [PATCH 1/2] feat: add guard in execute tx --- ree-cookie-canister/src/canister.rs | 5 ++++- ree-cookie-canister/src/state.rs | 4 ++++ ree-cookie-canister/src/utils.rs | 26 +++++++++++++++++++++++++- 3 files changed, 33 insertions(+), 2 deletions(-) diff --git a/ree-cookie-canister/src/canister.rs b/ree-cookie-canister/src/canister.rs index 4ad32b7..deaa03f 100644 --- a/ree-cookie-canister/src/canister.rs +++ b/ree-cookie-canister/src/canister.rs @@ -14,7 +14,7 @@ use crate::{ }, state::{ExchangeState, GameStatus, PoolState}, utils::{ - calculate_premine_rune_amount, tweak_pubkey_with_empty, AddLiquidityInfo, RegisterInfo, + calculate_premine_rune_amount, tweak_pubkey_with_empty, AddLiquidityInfo, ExecuteTxGuard, RegisterInfo }, ExchangeError, Seconds, MIN_BTC_VALUE, }; @@ -291,6 +291,9 @@ pub async fn execute_tx(args: ExecuteTxArgs) -> ExecuteTxResponse { output_coins, } = intention; + let _guard = ExecuteTxGuard::new(pool_address.clone()) + .ok_or(format!("Pool {0} Executing", pool_address).to_string())?; + read_state(|s| { return s .address diff --git a/ree-cookie-canister/src/state.rs b/ree-cookie-canister/src/state.rs index b19eafb..322a14f 100644 --- a/ree-cookie-canister/src/state.rs +++ b/ree-cookie-canister/src/state.rs @@ -4,6 +4,7 @@ use ic_stable_structures::Storable; use itertools::Itertools; use ree_types::{CoinBalance, CoinId, InputCoin, OutputCoin}; use std::borrow::Cow; +use std::collections::HashSet; use crate::game::game::Game; use crate::memory::{read_state, GAMER}; @@ -26,6 +27,8 @@ pub struct ExchangeState { pub etching_key: Option, pub richswap_pool_address: String, pub game_status: GameStatus, + #[serde(default)] + pub executing_pools: HashSet } #[derive(CandidType, Deserialize, Serialize, Clone, Debug)] @@ -110,6 +113,7 @@ impl ExchangeState { etching_key: None, richswap_pool_address, game_status: GameStatus::InitKey, + executing_pools: HashSet::new(), } } diff --git a/ree-cookie-canister/src/utils.rs b/ree-cookie-canister/src/utils.rs index 9546bc4..28030dc 100644 --- a/ree-cookie-canister/src/utils.rs +++ b/ree-cookie-canister/src/utils.rs @@ -1,7 +1,31 @@ use ic_cdk::api::management_canister::bitcoin::{BitcoinNetwork, Satoshi}; use ree_types::bitcoin::key::{Secp256k1, TapTweak, TweakedPublicKey}; -use crate::{memory::read_state, *}; +use crate::{memory::{mutate_state, read_state}, *}; + +#[must_use] +pub struct ExecuteTxGuard(String); + +impl ExecuteTxGuard { + pub fn new(pool_address: String) -> Option { + mutate_state( |s| { + if s.executing_pools.contains(&pool_address) { + return None::; + } + + s.executing_pools.insert(pool_address.clone()); + Some(ExecuteTxGuard(pool_address)) + }) + } +} + +impl Drop for ExecuteTxGuard { + fn drop(&mut self) { + mutate_state(|s| { + s.executing_pools.remove(&self.0); + }); + } +} pub(crate) fn tweak_pubkey_with_empty(untweaked: Pubkey) -> TweakedPublicKey { let secp = Secp256k1::new(); From 9407dfbdb791061a4275681401602db9896220a1 Mon Sep 17 00:00:00 2001 From: xsb <389483673@qq.com> Date: Mon, 23 Jun 2025 21:33:13 +0800 Subject: [PATCH 2/2] feat: save guard state in local var instead of memory --- ree-cookie-canister/src/canister.rs | 2 +- ree-cookie-canister/src/memory.rs | 5 ++++- ree-cookie-canister/src/state.rs | 5 ----- ree-cookie-canister/src/utils.rs | 20 +++++++++++--------- 4 files changed, 16 insertions(+), 16 deletions(-) diff --git a/ree-cookie-canister/src/canister.rs b/ree-cookie-canister/src/canister.rs index deaa03f..509b23c 100644 --- a/ree-cookie-canister/src/canister.rs +++ b/ree-cookie-canister/src/canister.rs @@ -12,7 +12,7 @@ use crate::{ memory::{ mutate_state, read_state, set_state, ADDRESS_PRINCIPLE_MAP, BLOCKS, GAMER, TX_RECORDS, }, - state::{ExchangeState, GameStatus, PoolState}, + state::{ExchangeState, PoolState}, utils::{ calculate_premine_rune_amount, tweak_pubkey_with_empty, AddLiquidityInfo, ExecuteTxGuard, RegisterInfo }, diff --git a/ree-cookie-canister/src/memory.rs b/ree-cookie-canister/src/memory.rs index 09c4b08..2c318f9 100644 --- a/ree-cookie-canister/src/memory.rs +++ b/ree-cookie-canister/src/memory.rs @@ -1,4 +1,5 @@ pub(crate) use std::cell::RefCell; +use std::collections::HashSet; use candid::Principal; use ic_stable_structures::{ @@ -52,6 +53,8 @@ thread_local! { MEMORY_MANAGER.with(|m| m.borrow().get(TX_RECORDS_MEMORY_ID)), ) ); + + pub static EXECUTING_POOLS: RefCell> = RefCell::new(HashSet::new()); } @@ -94,4 +97,4 @@ where let state = get_state(); let r = f(&state); r -} +} \ No newline at end of file diff --git a/ree-cookie-canister/src/state.rs b/ree-cookie-canister/src/state.rs index 322a14f..de67818 100644 --- a/ree-cookie-canister/src/state.rs +++ b/ree-cookie-canister/src/state.rs @@ -1,10 +1,8 @@ use ic_cdk::api::management_canister::bitcoin::Satoshi; use ic_stable_structures::storable::Bound; use ic_stable_structures::Storable; -use itertools::Itertools; use ree_types::{CoinBalance, CoinId, InputCoin, OutputCoin}; use std::borrow::Cow; -use std::collections::HashSet; use crate::game::game::Game; use crate::memory::{read_state, GAMER}; @@ -27,8 +25,6 @@ pub struct ExchangeState { pub etching_key: Option, pub richswap_pool_address: String, pub game_status: GameStatus, - #[serde(default)] - pub executing_pools: HashSet } #[derive(CandidType, Deserialize, Serialize, Clone, Debug)] @@ -113,7 +109,6 @@ impl ExchangeState { etching_key: None, richswap_pool_address, game_status: GameStatus::InitKey, - executing_pools: HashSet::new(), } } diff --git a/ree-cookie-canister/src/utils.rs b/ree-cookie-canister/src/utils.rs index 28030dc..3d6dd67 100644 --- a/ree-cookie-canister/src/utils.rs +++ b/ree-cookie-canister/src/utils.rs @@ -1,28 +1,30 @@ use ic_cdk::api::management_canister::bitcoin::{BitcoinNetwork, Satoshi}; use ree_types::bitcoin::key::{Secp256k1, TapTweak, TweakedPublicKey}; -use crate::{memory::{mutate_state, read_state}, *}; +use crate::{ + memory::{read_state, EXECUTING_POOLS}, + *, +}; #[must_use] pub struct ExecuteTxGuard(String); impl ExecuteTxGuard { pub fn new(pool_address: String) -> Option { - mutate_state( |s| { - if s.executing_pools.contains(&pool_address) { - return None::; + EXECUTING_POOLS.with(|executing_pools| { + if executing_pools.borrow().contains(&pool_address) { + return None; } - - s.executing_pools.insert(pool_address.clone()); - Some(ExecuteTxGuard(pool_address)) + executing_pools.borrow_mut().insert(pool_address.clone()); + return Some(ExecuteTxGuard(pool_address)); }) } } impl Drop for ExecuteTxGuard { fn drop(&mut self) { - mutate_state(|s| { - s.executing_pools.remove(&self.0); + EXECUTING_POOLS.with_borrow_mut(|executing_pools| { + executing_pools.remove(&self.0); }); } }