From 33c18d3fcb8685374bf9fd52d3e0010ea848d815 Mon Sep 17 00:00:00 2001 From: Ethan Frey Date: Fri, 23 Jun 2023 10:32:06 +0200 Subject: [PATCH] Better compile-time hiding of the test commit/rollback methods --- .../provider/external-staking/src/contract.rs | 421 ++++++++++-------- .../provider/external-staking/src/msg.rs | 3 +- .../external-staking/src/multitest.rs | 40 +- 3 files changed, 248 insertions(+), 216 deletions(-) diff --git a/contracts/provider/external-staking/src/contract.rs b/contracts/provider/external-staking/src/contract.rs index accf4e86..53c99317 100644 --- a/contracts/provider/external-staking/src/contract.rs +++ b/contracts/provider/external-staking/src/contract.rs @@ -16,11 +16,11 @@ use sylvia::types::{ExecCtx, InstantiateCtx, QueryCtx}; use crate::error::ContractError; use crate::ibc::VAL_CRDT; use crate::msg::{ - AllTxsResponse, AllTxsResponseItem, AuthorizedEndpointResponse, ConfigResponse, - IbcChannelResponse, ListRemoteValidatorsResponse, PendingRewards, ReceiveVirtualStake, - StakeInfo, StakesResponse, TxResponse, + AllTxsResponse, AuthorizedEndpointResponse, ConfigResponse, IbcChannelResponse, + ListRemoteValidatorsResponse, PendingRewards, ReceiveVirtualStake, StakeInfo, StakesResponse, + TxResponse, }; -use crate::state::{Config, Distribution, PendingUnbond, Stake}; +use crate::state::{Config, Distribution, Stake}; use mesh_sync::Tx; pub const CONTRACT_NAME: &str = env!("CARGO_PKG_NAME"); @@ -100,118 +100,134 @@ impl ExternalStakingContract<'_> { /// Commits a pending stake. /// Must be called by the IBC callback handler on successful remote staking. - #[msg(exec)] // TODO: Only enable for tests / use another method (IBC message) - fn commit_stake(&self, ctx: ExecCtx, tx_id: u64) -> Result { - // Load tx - let tx = self.pending_txs.load(ctx.deps.storage, tx_id)?; - - // TODO: Verify tx comes from the right context - // Verify tx is the right type - ensure!( - matches!(tx, Tx::InFlightRemoteStaking { .. }), - ContractError::WrongTypeTx(tx_id, tx) - ); - - let (tx_amount, tx_user, tx_validator) = match tx { - Tx::InFlightRemoteStaking { - amount, - user, - validator, - .. - } => (amount, user, validator), - _ => unreachable!(), - }; - - // Load stake - let mut stake_lock = self - .stakes - .load(ctx.deps.storage, (&tx_user, &tx_validator))?; - - // Load distribution - let mut distribution_lock = self.distribution.load(ctx.deps.storage, &tx_validator)?; - - // Commit amount (need to unlock it first) - stake_lock.unlock_write()?; - let stake = stake_lock.write()?; - stake.stake += tx_amount; + #[msg(exec)] + fn test_commit_stake(&self, ctx: ExecCtx, tx_id: u64) -> Result { + #[cfg(test)] + { + // Load tx + let tx = self.pending_txs.load(ctx.deps.storage, tx_id)?; + + // TODO: Verify tx comes from the right context + // Verify tx is the right type + ensure!( + matches!(tx, Tx::InFlightRemoteStaking { .. }), + ContractError::WrongTypeTx(tx_id, tx) + ); - // Commit distribution (need to unlock it first) - distribution_lock.unlock_write()?; - let distribution = distribution_lock.write()?; - // Distribution alignment - stake - .points_alignment - .stake_increased(tx_amount, distribution.points_per_stake); - distribution.total_stake += tx_amount; + let (tx_amount, tx_user, tx_validator) = match tx { + Tx::InFlightRemoteStaking { + amount, + user, + validator, + .. + } => (amount, user, validator), + _ => unreachable!(), + }; - // Save stake - self.stakes - .save(ctx.deps.storage, (&tx_user, &tx_validator), &stake_lock)?; + // Load stake + let mut stake_lock = self + .stakes + .load(ctx.deps.storage, (&tx_user, &tx_validator))?; + + // Load distribution + let mut distribution_lock = self.distribution.load(ctx.deps.storage, &tx_validator)?; + + // Commit amount (need to unlock it first) + stake_lock.unlock_write()?; + let stake = stake_lock.write()?; + stake.stake += tx_amount; + + // Commit distribution (need to unlock it first) + distribution_lock.unlock_write()?; + let distribution = distribution_lock.write()?; + // Distribution alignment + stake + .points_alignment + .stake_increased(tx_amount, distribution.points_per_stake); + distribution.total_stake += tx_amount; + + // Save stake + self.stakes + .save(ctx.deps.storage, (&tx_user, &tx_validator), &stake_lock)?; - // Save distribution - self.distribution - .save(ctx.deps.storage, &tx_validator, &distribution_lock)?; + // Save distribution + self.distribution + .save(ctx.deps.storage, &tx_validator, &distribution_lock)?; - // Remove tx - self.pending_txs.remove(ctx.deps.storage, tx_id); + // Remove tx + self.pending_txs.remove(ctx.deps.storage, tx_id); - // TODO: Call commit hook on vault + // TODO: Call commit hook on vault - Ok(Response::new()) + Ok(Response::new()) + } + #[cfg(not(test))] + { + let _ = (ctx, tx_id); + Err(ContractError::Unauthorized {}) + } } /// Rollbacks a pending stake. /// Must be called by the IBC callback handler on failed remote staking. - #[allow(unused)] - fn rollback_stake(&self, ctx: &mut ExecCtx, tx_id: u64) -> Result<(), ContractError> { - // Load tx - let tx = self.pending_txs.load(ctx.deps.storage, tx_id)?; - - // TODO: Verify tx comes from the right context - // Verify tx is the right type - ensure!( - matches!(tx, Tx::InFlightRemoteStaking { .. }), - ContractError::WrongTypeTx(tx_id, tx) - ); + #[msg(exec)] + fn test_rollback_stake(&self, ctx: ExecCtx, tx_id: u64) -> Result { + #[cfg(test)] + { + // Load tx + let tx = self.pending_txs.load(ctx.deps.storage, tx_id)?; + + // TODO: Verify tx comes from the right context + // Verify tx is the right type + ensure!( + matches!(tx, Tx::InFlightRemoteStaking { .. }), + ContractError::WrongTypeTx(tx_id, tx) + ); - let (_tx_amount, tx_user, tx_validator) = match tx { - Tx::InFlightRemoteStaking { - amount, - user, - validator, - .. - } => (amount, user, validator), - _ => unreachable!(), - }; + let (_tx_amount, tx_user, tx_validator) = match tx { + Tx::InFlightRemoteStaking { + amount, + user, + validator, + .. + } => (amount, user, validator), + _ => unreachable!(), + }; - // Load stake - let mut stake_lock = self - .stakes - .load(ctx.deps.storage, (&tx_user, &tx_validator))?; + // Load stake + let mut stake_lock = self + .stakes + .load(ctx.deps.storage, (&tx_user, &tx_validator))?; - // Load distribution - let mut distribution_lock = self.distribution.load(ctx.deps.storage, &tx_validator)?; + // Load distribution + let mut distribution_lock = self.distribution.load(ctx.deps.storage, &tx_validator)?; - // Release stake lock - stake_lock.unlock_write()?; + // Release stake lock + stake_lock.unlock_write()?; - // Save stake - self.stakes - .save(ctx.deps.storage, (&tx_user, &tx_validator), &stake_lock)?; + // Save stake + self.stakes + .save(ctx.deps.storage, (&tx_user, &tx_validator), &stake_lock)?; - // Release distribution lock - distribution_lock.unlock_write()?; + // Release distribution lock + distribution_lock.unlock_write()?; - // Save distribution - self.distribution - .save(ctx.deps.storage, &tx_validator, &distribution_lock)?; + // Save distribution + self.distribution + .save(ctx.deps.storage, &tx_validator, &distribution_lock)?; - // Remove tx - self.pending_txs.remove(ctx.deps.storage, tx_id); + // Remove tx + self.pending_txs.remove(ctx.deps.storage, tx_id); - // TODO: Call rollback hook on vault + // TODO: Call rollback hook on vault - Ok(()) + Ok(Response::new()) + } + #[cfg(not(test))] + { + let _ = (ctx, tx_id); + Err(ContractError::Unauthorized {}) + } } /// Schedules tokens for release, adding them to the pending unbonds. After `unbonding_period` @@ -279,123 +295,140 @@ impl ExternalStakingContract<'_> { /// Commits a pending unstake. /// Must be called by the IBC callback handler on successful remote unstaking. - #[allow(unused)] - #[msg(exec)] // TODO: Only enable for tests / use another method (IBC message) - fn commit_unstake(&self, ctx: ExecCtx, tx_id: u64) -> Result { - // Load tx - let tx = self.pending_txs.load(ctx.deps.storage, tx_id)?; - - // TODO: Verify tx comes from the right context - // Verify tx is of the right type - ensure!( - matches!(tx, Tx::InFlightRemoteUnstaking { .. }), - ContractError::WrongTypeTx(tx_id, tx) - ); - - let (tx_amount, tx_user, tx_validator) = match tx { - Tx::InFlightRemoteUnstaking { - amount, - user, - validator, - .. - } => (amount, user, validator), - _ => unreachable!(), - }; - - let config = self.config.load(ctx.deps.storage)?; + #[msg(exec)] + fn test_commit_unstake(&self, ctx: ExecCtx, tx_id: u64) -> Result { + #[cfg(test)] + { + use crate::state::PendingUnbond; + + // Load tx + let tx = self.pending_txs.load(ctx.deps.storage, tx_id)?; + + // TODO: Verify tx comes from the right context + // Verify tx is of the right type + ensure!( + matches!(tx, Tx::InFlightRemoteUnstaking { .. }), + ContractError::WrongTypeTx(tx_id, tx) + ); - // Load stake - let mut stake_lock = self - .stakes - .load(ctx.deps.storage, (&tx_user, &tx_validator))?; + let (tx_amount, tx_user, tx_validator) = match tx { + Tx::InFlightRemoteUnstaking { + amount, + user, + validator, + .. + } => (amount, user, validator), + _ => unreachable!(), + }; - // Load distribution - let mut distribution_lock = self.distribution.load(ctx.deps.storage, &tx_validator)?; + let config = self.config.load(ctx.deps.storage)?; - // Commit amount (need to unlock it first) - stake_lock.unlock_write()?; - let stake = stake_lock.write()?; - stake.stake -= tx_amount; + // Load stake + let mut stake_lock = self + .stakes + .load(ctx.deps.storage, (&tx_user, &tx_validator))?; - // FIXME? Release period being computed after successful IBC tx - let release_at = ctx.env.block.time.plus_seconds(config.unbonding_period); - let unbond = PendingUnbond { - amount: tx_amount, - release_at, - }; - stake.pending_unbonds.push(unbond); + // Load distribution + let mut distribution_lock = self.distribution.load(ctx.deps.storage, &tx_validator)?; - // Commit distribution (need to unlock it first) - distribution_lock.unlock_write()?; - let distribution = distribution_lock.write()?; - // Distribution alignment - stake - .points_alignment - .stake_decreased(tx_amount, distribution.points_per_stake); - distribution.total_stake -= tx_amount; + // Commit amount (need to unlock it first) + stake_lock.unlock_write()?; + let stake = stake_lock.write()?; + stake.stake -= tx_amount; - // Save stake - self.stakes - .save(ctx.deps.storage, (&tx_user, &tx_validator), &stake_lock)?; + // FIXME? Release period being computed after successful IBC tx + let release_at = ctx.env.block.time.plus_seconds(config.unbonding_period); + let unbond = PendingUnbond { + amount: tx_amount, + release_at, + }; + stake.pending_unbonds.push(unbond); + + // Commit distribution (need to unlock it first) + distribution_lock.unlock_write()?; + let distribution = distribution_lock.write()?; + // Distribution alignment + stake + .points_alignment + .stake_decreased(tx_amount, distribution.points_per_stake); + distribution.total_stake -= tx_amount; + + // Save stake + self.stakes + .save(ctx.deps.storage, (&tx_user, &tx_validator), &stake_lock)?; - // Save distribution - self.distribution - .save(ctx.deps.storage, &tx_validator, &distribution_lock)?; + // Save distribution + self.distribution + .save(ctx.deps.storage, &tx_validator, &distribution_lock)?; - // Remove tx - self.pending_txs.remove(ctx.deps.storage, tx_id); + // Remove tx + self.pending_txs.remove(ctx.deps.storage, tx_id); - Ok(Response::new()) + Ok(Response::new()) + } + #[cfg(not(test))] + { + let _ = (ctx, tx_id); + Err(ContractError::Unauthorized {}) + } } /// Rollbacks a pending unstake. /// Must be called by the IBC callback handler on failed remote unstaking. - #[allow(unused)] - fn rollback_unstake(&self, ctx: &mut ExecCtx, tx_id: u64) -> Result<(), ContractError> { - // Load tx - let tx = self.pending_txs.load(ctx.deps.storage, tx_id)?; - - // TODO: Verify tx comes from the right context - // Verify tx is of the right type - ensure!( - matches!(tx, Tx::InFlightRemoteUnstaking { .. }), - ContractError::WrongTypeTx(tx_id, tx) - ); - let (_tx_amount, tx_user, tx_validator) = match tx { - Tx::InFlightRemoteUnstaking { - amount, - user, - validator, - .. - } => (amount, user, validator), - _ => unreachable!(), - }; + #[msg(exec)] + fn test_rollback_unstake(&self, ctx: ExecCtx, tx_id: u64) -> Result { + #[cfg(test)] + { + // Load tx + let tx = self.pending_txs.load(ctx.deps.storage, tx_id)?; + + // TODO: Verify tx comes from the right context + // Verify tx is of the right type + ensure!( + matches!(tx, Tx::InFlightRemoteUnstaking { .. }), + ContractError::WrongTypeTx(tx_id, tx) + ); + let (_tx_amount, tx_user, tx_validator) = match tx { + Tx::InFlightRemoteUnstaking { + amount, + user, + validator, + .. + } => (amount, user, validator), + _ => unreachable!(), + }; - // Load stake - let mut stake_lock = self - .stakes - .load(ctx.deps.storage, (&tx_user, &tx_validator))?; + // Load stake + let mut stake_lock = self + .stakes + .load(ctx.deps.storage, (&tx_user, &tx_validator))?; - // Load distribution - let mut distribution_lock = self.distribution.load(ctx.deps.storage, &tx_validator)?; + // Load distribution + let mut distribution_lock = self.distribution.load(ctx.deps.storage, &tx_validator)?; - // Release stake lock - stake_lock.unlock_write()?; + // Release stake lock + stake_lock.unlock_write()?; - // Save stake - self.stakes - .save(ctx.deps.storage, (&tx_user, &tx_validator), &stake_lock)?; + // Save stake + self.stakes + .save(ctx.deps.storage, (&tx_user, &tx_validator), &stake_lock)?; - // Release distribution lock - distribution_lock.unlock_write()?; + // Release distribution lock + distribution_lock.unlock_write()?; - // Save distribution - self.distribution - .save(ctx.deps.storage, &tx_validator, &distribution_lock)?; + // Save distribution + self.distribution + .save(ctx.deps.storage, &tx_validator, &distribution_lock)?; - // Remove tx - self.pending_txs.remove(ctx.deps.storage, tx_id); - Ok(()) + // Remove tx + self.pending_txs.remove(ctx.deps.storage, tx_id); + Ok(Response::new()) + } + #[cfg(not(test))] + { + let _ = (ctx, tx_id); + Err(ContractError::Unauthorized {}) + } } /// Withdraws all released tokens to the sender. @@ -656,7 +689,7 @@ impl ExternalStakingContract<'_> { .range(ctx.deps.storage, None, bound, Order::Descending) .map(|item| { let (_id, tx) = item?; - Ok::(tx) + Ok::(tx) }) .take(limit) .collect::>()?; diff --git a/contracts/provider/external-staking/src/msg.rs b/contracts/provider/external-staking/src/msg.rs index edaede58..bcb7cf42 100644 --- a/contracts/provider/external-staking/src/msg.rs +++ b/contracts/provider/external-staking/src/msg.rs @@ -96,9 +96,8 @@ pub struct PendingRewards { } pub type TxResponse = mesh_sync::Tx; -pub type AllTxsResponseItem = TxResponse; #[cw_serde] pub struct AllTxsResponse { - pub txs: Vec, + pub txs: Vec, } diff --git a/contracts/provider/external-staking/src/multitest.rs b/contracts/provider/external-staking/src/multitest.rs index 8a29dfc0..32219838 100644 --- a/contracts/provider/external-staking/src/multitest.rs +++ b/contracts/provider/external-staking/src/multitest.rs @@ -143,7 +143,7 @@ fn staking() { // This should be through `IbcPacketAckMsg` let last_external_staking_tx = get_last_external_staking_pending_tx_id(&contract).unwrap(); contract - .commit_stake(last_external_staking_tx) + .test_commit_stake(last_external_staking_tx) .call("test") .unwrap(); @@ -166,7 +166,7 @@ fn staking() { .unwrap(); contract - .commit_stake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_stake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap(); @@ -189,7 +189,7 @@ fn staking() { .unwrap(); contract - .commit_stake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_stake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap(); @@ -212,7 +212,7 @@ fn staking() { .unwrap(); contract - .commit_stake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_stake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap(); @@ -234,7 +234,7 @@ fn staking() { .unwrap(); contract - .commit_stake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_stake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap(); @@ -394,7 +394,7 @@ fn unstaking() { // This should be through `IbcPacketAckMsg` let last_external_staking_tx = get_last_external_staking_pending_tx_id(&contract).unwrap(); contract - .commit_stake(last_external_staking_tx) + .test_commit_stake(last_external_staking_tx) .call("test") .unwrap(); @@ -416,7 +416,7 @@ fn unstaking() { .unwrap(); let last_external_staking_tx = get_last_external_staking_pending_tx_id(&contract).unwrap(); contract - .commit_stake(last_external_staking_tx) + .test_commit_stake(last_external_staking_tx) .call("test") .unwrap(); @@ -438,7 +438,7 @@ fn unstaking() { .unwrap(); let last_external_staking_tx = get_last_external_staking_pending_tx_id(&contract).unwrap(); contract - .commit_stake(last_external_staking_tx) + .test_commit_stake(last_external_staking_tx) .call("test") .unwrap(); @@ -450,7 +450,7 @@ fn unstaking() { .call(users[0]) .unwrap(); contract - .commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap(); @@ -459,7 +459,7 @@ fn unstaking() { .call(users[0]) .unwrap(); contract - .commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap(); @@ -468,7 +468,7 @@ fn unstaking() { .call(users[1]) .unwrap(); contract - .commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap(); @@ -567,7 +567,7 @@ fn unstaking() { .call(users[0]) .unwrap(); contract - .commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap(); @@ -576,7 +576,7 @@ fn unstaking() { .call(users[0]) .unwrap(); contract - .commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap(); @@ -730,7 +730,7 @@ fn distribution() { // This should be through `IbcPacketAckMsg` let last_external_staking_tx = get_last_external_staking_pending_tx_id(&contract).unwrap(); contract - .commit_stake(last_external_staking_tx) + .test_commit_stake(last_external_staking_tx) .call("test") .unwrap(); @@ -751,7 +751,7 @@ fn distribution() { .call(contract.contract_addr.as_str()) .unwrap(); contract - .commit_stake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_stake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap(); @@ -772,7 +772,7 @@ fn distribution() { .call(contract.contract_addr.as_str()) .unwrap(); contract - .commit_stake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_stake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap(); @@ -975,7 +975,7 @@ fn distribution() { .call(users[1]) .unwrap(); contract - .commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap(); @@ -1002,7 +1002,7 @@ fn distribution() { .call(contract.contract_addr.as_str()) .unwrap(); contract - .commit_stake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_stake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap(); @@ -1103,7 +1103,7 @@ fn distribution() { .call(users[0]) .unwrap(); contract - .commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap(); @@ -1112,7 +1112,7 @@ fn distribution() { .call(users[1]) .unwrap(); contract - .commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) + .test_commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap()) .call("test") .unwrap();