diff --git a/Cargo.toml b/Cargo.toml index 15d7f021f..a0503442f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -127,7 +127,7 @@ utxo = { path = "utxo" } [workspace.package] edition = "2024" -rust-version = "1.91" +rust-version = "1.92" version = "1.3.0" license = "MIT" diff --git a/mempool/src/pool/tx_pool/collect_txs.rs b/mempool/src/pool/tx_pool/collect_txs.rs index e938cc4c0..e60627f32 100644 --- a/mempool/src/pool/tx_pool/collect_txs.rs +++ b/mempool/src/pool/tx_pool/collect_txs.rs @@ -359,6 +359,23 @@ pub fn get_best_tx_ids_by_score_and_ancestry( mod get_best_tx_ids_by_score_and_ancestry_impl { use super::*; + // Stack item for collect_nearest_selected_ancestors. + struct StackItem<'a, I: Iterator>> { + tx_entry: &'a TxMempoolEntry, + tx_ancestors: BTreeSet>, + parents_iter: I, + } + + fn new_stack_item<'a>( + tx_entry: &'a TxMempoolEntry, + ) -> StackItem<'a, impl Iterator>> { + StackItem { + tx_entry, + tx_ancestors: BTreeSet::new(), + parents_iter: tx_entry.parents(), + } + } + // For each ancestor path starting from the specified tx, take the nearest ancestor that is in `selected_tx_ids` // and put it into a set. E.g. in // /--S1--**--S2 @@ -379,30 +396,41 @@ mod get_best_tx_ids_by_score_and_ancestry_impl { selected_tx_ids: &BTreeSet>, ancestors_map: &mut BTreeMap, BTreeSet>>, ) -> Result<(), TxCollectionError> { - let tx_id = tx_entry.tx_id(); + if ancestors_map.contains_key(tx_entry.tx_id()) { + return Ok(()); + } - if !ancestors_map.contains_key(tx_id) { - let mut tx_ancestors = BTreeSet::new(); + let mut stack = vec![new_stack_item(tx_entry)]; - for parent_id in tx_entry.parents() { + 'outer: loop { + // We start with a non-empty stack and then check if it's empty at the end of the outer loop. + let item = stack.last_mut().expect("known to be present"); + + for parent_id in item.parents_iter.by_ref() { if selected_tx_ids.contains(parent_id) { // The nearest selected ancestor has been found, no need to go deeper. - tx_ancestors.insert(*parent_id); + item.tx_ancestors.insert(*parent_id); + } else if let Some(existing_ancestors) = ancestors_map.get(parent_id) { + item.tx_ancestors.extend(existing_ancestors); } else { let parent_tx_entry = mempool.store.get_existing_entry(parent_id)?; - collect_nearest_selected_ancestors( - mempool, - parent_tx_entry, - selected_tx_ids, - ancestors_map, - )?; - - let parent_ancestors = ancestors_map.get(parent_id).expect("must be present"); - tx_ancestors.extend(parent_ancestors); + stack.push(new_stack_item(parent_tx_entry)); + continue 'outer; } } - ancestors_map.insert(*tx_id, tx_ancestors); + // We've completely processed the current item, so put the collected `tx_ancestors` + // into `ancestors_map`. If there is another item on top of the stack, it's this item's + // child, whose `tx_ancestors` must be extended with this item's `tx_ancestors`. + // Otherwise we're done. + let item = stack.pop().expect("known to be present"); + let entry = ancestors_map.entry(*item.tx_entry.tx_id()).insert_entry(item.tx_ancestors); + + if let Some(child_item) = stack.last_mut() { + child_item.tx_ancestors.extend(entry.get()); + } else { + break; + } } Ok(()) diff --git a/mempool/src/pool/tx_pool/tests/tx_ids_by_score_and_ancestry.rs b/mempool/src/pool/tx_pool/tests/tx_ids_by_score_and_ancestry.rs index eaf5f24fa..c96a9d47c 100644 --- a/mempool/src/pool/tx_pool/tests/tx_ids_by_score_and_ancestry.rs +++ b/mempool/src/pool/tx_pool/tests/tx_ids_by_score_and_ancestry.rs @@ -174,3 +174,87 @@ async fn tx_ids_by_score_and_ancestry(#[case] seed: Seed) { ] ); } + +// While using a small stack and unlimited `max_cluster_tx_count`, create a long chain of +// dependent txs and pass the first and the last of them to `get_best_tx_ids_by_score_and_ancestry` +// (so that it has to check all the txs in between). The expected result is that the test doesn't +// crash. This basically checks that `get_best_tx_ids_by_score_and_ancestry` doesn't use recursion. +#[rstest] +#[case(Seed::from_entropy())] +#[trace] +fn stack_overflow_on_misconfigured_cluster_size_limit(#[case] seed: Seed) { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + // Use a small stack, same logic as in the `stack_overflow_on_transaction_addition` test. + .thread_stack_size(256 * 1024) + .enable_all() + .build() + .unwrap(); + + let test_body = move || { + let mut rng = make_seedable_rng(seed); + let mut tf = TestFramework::builder(&mut rng).build(); + let genesis_id = tf.genesis().get_id(); + + // This was enough to crash the original code using the given stack size both in + // debug and release builds (actually, 1000 was also enough, but we use a bigger + // value "just in case"). + let chain_len = 1500; + + let mut amount = 1_000_000; + let funding_tx = make_tx( + &mut rng, + &[(OutPointSourceId::BlockReward(genesis_id.into()), 0)], + &[amount], + ); + let funding_tx_id = funding_tx.transaction().get_id(); + tf.make_block_builder() + .add_transaction(funding_tx) + .build_and_process(&mut rng) + .unwrap(); + + let mut mempool = setup_with_chainstate_generic( + tf.chainstate(), + MempoolConfig { + min_tx_relay_fee_rate: FeeRate::from_amount_per_kb(Amount::ZERO).into(), + max_cluster_size_bytes: usize::MAX.into(), + max_cluster_tx_count: usize::MAX.into(), + }, + Default::default(), + ); + + // Disable heavy checks because the test is already relatively slow. + mempool.disable_heavy_validity_checks(); + + let mut prev_source: OutPointSourceId = funding_tx_id.into(); + let mut tx_ids = Vec::new(); + + for _ in 0..chain_len { + let tx = make_tx(&mut rng, &[(prev_source, 0)], &[amount - 1]); + amount -= 1; + let tx_id = tx.transaction().get_id(); + prev_source = tx_id.into(); + + let result = mempool.add_transaction_test(tx); + + assert_eq!(result, Ok(TxStatus::InMempool)); + tx_ids.push(tx_id); + } + + let selected_tx_ids = [*tx_ids.first().unwrap(), *tx_ids.last().unwrap()]; + let result = mempool + .get_best_tx_ids_by_score_and_ancestry(&BTreeSet::from(selected_tx_ids), tx_ids.len()) + .unwrap(); + assert_eq!(result, selected_tx_ids); + }; + + runtime.block_on(async move { + // Same as in `stack_overflow_on_transaction_addition`, need to use `spawn` or `spawn_blocking` + // here. + tokio::task::spawn_blocking(move || { + test_body(); + }) + .await + .unwrap(); + }); +}