Skip to content

Commit

Permalink
Fix delegation support in the wallet
Browse files Browse the repository at this point in the history 
- previously all delegation staking was added as own in the cache
- fix delegation rollback when a reorg happens
- fix delegation nonce rollback when abandoning inactive transactions
  with descendants
  • Loading branch information
@OBorce
OBorce committed Sep 20, 2023
1 parent a1e451c commit 3cf71fd
Show file tree
Hide file tree
Showing 6 changed files with 406 additions and 126 deletions.
31 changes: 27 additions & 4 deletions test/functional/wallet_delegations.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -324,9 +324,9 @@ async def async_test(self):

assert "Success" in await wallet.create_new_account()
assert "Success" in await wallet.select_account(1)
address = await wallet.new_address()
acc1_address = await wallet.new_address()
assert "Success" in await wallet.select_account(DEFAULT_ACCOUNT_INDEX)
assert "The transaction was submitted successfully" in await wallet.send_to_address(address, 100)
assert "The transaction was submitted successfully" in await wallet.send_to_address(acc1_address, 100)
assert "Success" in await wallet.select_account(1)
transactions = node.mempool_transactions()

Expand All @@ -345,7 +345,7 @@ async def async_test(self):
assert pools[0].balance == 40000

assert "Success" in await wallet.select_account(1)
delegation_id = await wallet.create_delegation(address, pools[0].pool_id)
delegation_id = await wallet.create_delegation(acc1_address, pools[0].pool_id)
assert delegation_id is not None
transactions = node.mempool_transactions()

Expand Down Expand Up @@ -373,7 +373,7 @@ async def async_test(self):
last_delegation_balance = delegations[0].balance
for _ in range(4, 10):
tip_id = node.chainstate_best_block_id()
assert "The transaction was submitted successfully" in await wallet.send_to_address(address, 1)
assert "The transaction was submitted successfully" in await wallet.send_to_address(acc1_address, 1)
transactions = node.mempool_transactions()
self.wait_until(lambda: node.chainstate_best_block_id() != tip_id, timeout = 5)

Expand All @@ -382,6 +382,29 @@ async def async_test(self):
assert delegations[0].balance > last_delegation_balance
last_delegation_balance = delegations[0].balance

# stake to acc1 delegation from acc 0
assert "Success" in await wallet.select_account(DEFAULT_ACCOUNT_INDEX)
assert "Success" in await wallet.stake_delegation(10, delegation_id)
self.wait_until(lambda: node.chainstate_best_block_id() != tip_id, timeout = 5)

# check that we still don't have any delagations for this account
delegations = await wallet.list_delegation_ids()
assert len(delegations) == 0

# create a delegation from acc 0 but with destination address for acc1
delegation_id = await wallet.create_delegation(acc1_address, pools[0].pool_id)
tip_id = node.chainstate_best_block_id()
self.wait_until(lambda: node.chainstate_best_block_id() != tip_id, timeout = 5)

# check that we still don't have any delagations for this account
delegations = await wallet.list_delegation_ids()
assert len(delegations) == 0

assert "Success" in await wallet.select_account(1)
delegations = await wallet.list_delegation_ids()
assert len(delegations) == 2
assert delegation_id in [delegation.delegation_id for delegation in delegations]


if __name__ == '__main__':
WalletSubmitTransaction().main()
Expand Down
41 changes: 30 additions & 11 deletions wallet/src/account/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ use crypto::key::hdkd::u31::U31;
use crypto::key::PublicKey;
use crypto::vrf::{VRFPrivateKey, VRFPublicKey};
use itertools::Itertools;
use std::cmp::Reverse;
use std::collections::BTreeMap;
use std::ops::Add;
use std::sync::Arc;
Expand Down Expand Up @@ -450,7 +451,7 @@ impl Account {
amount,
current_block_height,
)?;
let delegation_data = self.output_cache.delegation_data(delegation_id)?;
let delegation_data = self.find_delegation(&delegation_id)?;
let nonce = delegation_data
.last_nonce
.map_or(Some(AccountNonce::new(0)), |nonce| nonce.increment())
Expand Down Expand Up @@ -527,11 +528,16 @@ impl Account {
}

pub fn get_delegations(&self) -> impl Iterator<Item = (&DelegationId, &DelegationData)> {
self.output_cache.delegation_ids()
self.output_cache
.delegation_ids()
.filter(|(_, data)| self.is_mine_or_watched_destination(&data.destination))
}

pub fn find_delegation(&self, delegation_id: DelegationId) -> WalletResult<&DelegationData> {
self.output_cache.delegation_data(delegation_id)
pub fn find_delegation(&self, delegation_id: &DelegationId) -> WalletResult<&DelegationData> {
self.output_cache
.delegation_data(delegation_id)
.filter(|data| self.is_mine_or_watched_destination(&data.destination))
.ok_or(WalletError::DelegationNotFound(*delegation_id))
}

pub fn create_stake_pool_tx(
Expand Down Expand Up @@ -765,12 +771,19 @@ impl Account {
/// Return true if this transaction output is can be spent by this account or if it is being
/// watched.
fn is_mine_or_watched(&self, txo: &TxOutput) -> bool {
Self::get_tx_output_destination(txo).map_or(false, |d| match d {
Self::get_tx_output_destination(txo)
.map_or(false, |d| self.is_mine_or_watched_destination(d))
}

/// Return true if this transaction output is can be spent by this account or if it is being
/// watched.
fn is_mine_or_watched_destination(&self, destination: &Destination) -> bool {
match destination {
Destination::Address(pkh) => self.key_chain.is_public_key_hash_mine(pkh),
Destination::PublicKey(pk) => self.key_chain.is_public_key_mine(pk),
Destination::AnyoneCanSpend => false,
Destination::ScriptHash(_) | Destination::ClassicMultisig(_) => false,
})
}
}

fn mark_outputs_as_seen(
Expand Down Expand Up @@ -862,14 +875,17 @@ impl Account {
wallet_events: &impl WalletEvents,
common_block_height: BlockHeight,
) -> WalletResult<()> {
let revoked_txs = self
let mut revoked_txs = self
.output_cache
.txs_with_unconfirmed()
.iter()
.filter_map(|(id, tx)| match tx.state() {
TxState::Confirmed(height, _, _) => {
TxState::Confirmed(height, _, idx) => {
if height > common_block_height {
Some(AccountWalletTxId::new(self.get_account_id(), id.clone()))
Some((
AccountWalletTxId::new(self.get_account_id(), id.clone()),
(height, idx),
))
} else {
None
}
Expand All @@ -881,7 +897,10 @@ impl Account {
})
.collect::<Vec<_>>();

for tx_id in revoked_txs {
// sort from latest tx down to remove them in order
revoked_txs.sort_by_key(|&(_, height_idx)| Reverse(height_idx));

for (tx_id, _) in revoked_txs {
db_tx.del_transaction(&tx_id)?;
wallet_events.del_transaction(&tx_id);
self.output_cache.remove_tx(&tx_id.into_item_id())?;
Expand All @@ -905,7 +924,7 @@ impl Account {
.map_or(false, |txo| self.is_mine_or_watched(txo)),
TxInput::Account(outpoint) => match outpoint.account() {
AccountSpending::Delegation(delegation_id, _) => {
self.output_cache.owns_delegation(delegation_id)
self.find_delegation(delegation_id).is_ok()
}
},
});
Expand Down
167 changes: 93 additions & 74 deletions wallet/src/account/output_cache/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,8 @@ pub struct DelegationData {
pub pool_id: PoolId,
pub destination: Destination,
pub last_nonce: Option<AccountNonce>,
/// last parent transaction if the parent is unconfirmed
pub last_parent: Option<OutPointSourceId>,
pub not_staked_yet: bool,
}
impl DelegationData {
Expand All @@ -48,6 +50,7 @@ impl DelegationData {
pool_id,
destination,
last_nonce: None,
last_parent: None,
not_staked_yet: true,
}
}
Expand Down Expand Up @@ -162,14 +165,8 @@ impl OutputCache {
self.delegations.iter()
}

pub fn delegation_data(&self, delegation_id: DelegationId) -> WalletResult<&DelegationData> {
self.delegations
.get(&delegation_id)
.ok_or(WalletError::DelegationNotFound(delegation_id))
}

pub fn owns_delegation(&self, delegation_id: &DelegationId) -> bool {
self.delegations.contains_key(delegation_id)
pub fn delegation_data(&self, delegation_id: &DelegationId) -> Option<&DelegationData> {
self.delegations.get(delegation_id)
}

pub fn add_tx(&mut self, tx_id: OutPointSourceId, tx: WalletTx) -> WalletResult<()> {
Expand Down Expand Up @@ -220,11 +217,10 @@ impl OutputCache {
});
}
TxOutput::DelegateStaking(_, delegation_id) => {
let delegation_data = self
.delegations
.get_mut(delegation_id)
.ok_or(WalletError::InconsistentDelegationAddition(*delegation_id))?;
delegation_data.not_staked_yet = false;
if let Some(delegation_data) = self.delegations.get_mut(delegation_id) {
delegation_data.not_staked_yet = false;
}
// Else it is not ours
}
TxOutput::CreateDelegationId(destination, pool_id) => {
let input0_outpoint = tx
Expand Down Expand Up @@ -273,30 +269,14 @@ impl OutputCache {
if !already_present {
match outpoint.account() {
Delegation(delegation_id, _) => {
match self.delegations.get_mut(delegation_id) {
Some(data) => {
let next_nonce = data
.last_nonce
.map_or(Some(AccountNonce::new(0)), |nonce| {
nonce.increment()
})
.ok_or(WalletError::DelegationNonceOverflow(
*delegation_id,
))?;
ensure!(
outpoint.nonce() >= next_nonce,
WalletError::InconsistentDelegationDuplicateNonce(
*delegation_id,
outpoint.nonce()
)
);
data.last_nonce = Some(outpoint.nonce());
}
None => {
return Err(WalletError::InconsistentDelegationRemoval(
*delegation_id,
));
}
if let Some(data) = self.delegations.get_mut(delegation_id) {
Self::update_delegation_state(
&mut self.unconfirmed_descendants,
data,
delegation_id,
outpoint,
tx_id,
)?;
}
}
}
Expand All @@ -307,6 +287,37 @@ impl OutputCache {
Ok(())
}

/// Update delegation state with new tx input
fn update_delegation_state(
unconfirmed_descendants: &mut BTreeMap<OutPointSourceId, BTreeSet<OutPointSourceId>>,
data: &mut DelegationData,
delegation_id: &DelegationId,
outpoint: &common::chain::AccountOutPoint,
tx_id: &OutPointSourceId,
) -> Result<(), WalletError> {
let next_nonce = data
.last_nonce
.map_or(Some(AccountNonce::new(0)), |nonce| nonce.increment())
.ok_or(WalletError::DelegationNonceOverflow(*delegation_id))?;

ensure!(
outpoint.nonce() == next_nonce,
WalletError::InconsistentDelegationDuplicateNonce(*delegation_id, outpoint.nonce())
);

data.last_nonce = Some(outpoint.nonce());
// update unconfirmed descendants
if let Some(descendants) = data
.last_parent
.as_ref()
.and_then(|parent_tx_id| unconfirmed_descendants.get_mut(parent_tx_id))
{
descendants.insert(tx_id.clone());
}
data.last_parent = Some(tx_id.clone());
Ok(())
}

pub fn remove_tx(&mut self, tx_id: &OutPointSourceId) -> WalletResult<()> {
let tx_opt = self.txs.remove(tx_id);
if let Some(tx) = tx_opt {
Expand All @@ -318,13 +329,10 @@ impl OutputCache {
}
TxInput::Account(outpoint) => match outpoint.account() {
Delegation(delegation_id, _) => {
match self.delegations.get_mut(delegation_id) {
Some(data) => {
data.last_nonce = outpoint.nonce().decrement();
}
None => {
return Err(WalletError::NoUtxos);
}
if let Some(data) = self.delegations.get_mut(delegation_id) {
data.last_nonce = outpoint.nonce().decrement();
data.last_parent =
find_parent(&self.unconfirmed_descendants, tx_id.clone());
}
}
},
Expand Down Expand Up @@ -474,40 +482,40 @@ impl OutputCache {
if let Some(descendants) = self.unconfirmed_descendants.remove(&outpoint_source_id) {
to_abandon.extend(descendants.into_iter())
}
}

match self.txs.entry(outpoint_source_id) {
Entry::Occupied(mut entry) => {
match entry.get_mut() {
WalletTx::Block(_) => Err(WalletError::CannotFindTransactionWithId(tx_id)),
WalletTx::Tx(tx) => match tx.state() {
TxState::Inactive(_) => {
tx.set_state(TxState::Abandoned);
for input in tx.get_transaction().inputs() {
match input {
TxInput::Utxo(outpoint) => {
self.consumed.insert(outpoint.clone(), *tx.state());
}
TxInput::Account(outpoint) => match outpoint.account() {
Delegation(delegation_id, _) => {
match self.delegations.get_mut(delegation_id) {
Some(data) => {
data.last_nonce =
outpoint.nonce().decrement();
}
None => {
return Err(WalletError::InconsistentDelegationRemoval(*delegation_id));
}
}
}
},
for tx_id in all_abandoned.iter().rev().copied() {
match self.txs.entry(tx_id.into()) {
Entry::Occupied(mut entry) => match entry.get_mut() {
WalletTx::Block(_) => Err(WalletError::CannotFindTransactionWithId(tx_id)),
WalletTx::Tx(tx) => match tx.state() {
TxState::Inactive(_) => {
tx.set_state(TxState::Abandoned);
for input in tx.get_transaction().inputs() {
match input {
TxInput::Utxo(outpoint) => {
self.consumed.insert(outpoint.clone(), *tx.state());
}
TxInput::Account(outpoint) => match outpoint.account() {
Delegation(delegation_id, _) => {
if let Some(data) =
self.delegations.get_mut(delegation_id)
{
data.last_nonce = outpoint.nonce().decrement();
data.last_parent = find_parent(
&self.unconfirmed_descendants,
tx_id.into(),
);
}
}
},
}
Ok(())
}
state => Err(WalletError::CannotAbandonTransaction(*state)),
},
}
}
Ok(())
}
state => Err(WalletError::CannotAbandonTransaction(*state)),
},
},
Entry::Vacant(_) => Err(WalletError::CannotFindTransactionWithId(tx_id)),
}?;
}
Expand Down Expand Up @@ -576,3 +584,14 @@ fn valid_timelock(
fn is_in_state(tx: &WalletTx, utxo_states: UtxoStates) -> bool {
utxo_states.contains(get_utxo_state(&tx.state()))
}

/// Find the parent tx if it is in the unconfirmed transactions
fn find_parent(
unconfirmed_descendants: &BTreeMap<OutPointSourceId, BTreeSet<OutPointSourceId>>,
tx_id: OutPointSourceId,
) -> Option<OutPointSourceId> {
unconfirmed_descendants
.iter()
.find_map(|(parent_id, descendants)| descendants.contains(&tx_id).then_some(parent_id))
.cloned()
}
Loading

0 comments on commit 3cf71fd

Please sign in to comment.