From 471536571d3af6866a2bd447c542cdae428f73c1 Mon Sep 17 00:00:00 2001 From: Nikolaus Heger Date: Wed, 2 Sep 2026 18:22:26 +0800 Subject: [PATCH 1/2] fix(wormhole): prove collect-rewards against the head, not the finalized block QPoW finality trails the head by a long way (~100 blocks on staging mainnet), so any leaf minted inside that window is absent from the ZK tree at the finalized block. collect_rewards proved against chain_getFinalizedHead, which made recent rewards unsweepable and failed with "Leaf index N not found in ZK tree at block H" part-way through proof generation. Prove against the head instead, matching every other read path in the CLI. A proof invalidated by a reorg is re-run, which is far cheaper than waiting for finality. --at-block still pins a specific block for callers who want one. Also check settlement before generating proofs rather than discovering it on the Nth one: leaves settle in index order, so one probe of the highest selected leaf covers the common case and a binary search finds the boundary otherwise. Leaves that are still too new are now reported and skipped instead of aborting the whole sweep, and the all-unsettled case gets an error that says what to do about it. Splits get_zk_merkle_proof into try_get_zk_merkle_proof, which returns Ok(None) for an unsettled leaf so callers can tell that apart from an RPC failure. --- src/cli/wormhole.rs | 27 ++++++++--- src/collect_rewards_lib.rs | 93 +++++++++++++++++++++++++++++++++----- 2 files changed, 101 insertions(+), 19 deletions(-) diff --git a/src/cli/wormhole.rs b/src/cli/wormhole.rs index 17e08bd..a23a6f3 100644 --- a/src/cli/wormhole.rs +++ b/src/cli/wormhole.rs @@ -204,6 +204,22 @@ pub async fn get_zk_merkle_proof( leaf_index: u64, at_block: subxt::utils::H256, ) -> crate::error::Result { + try_get_zk_merkle_proof(quantus_client, leaf_index, at_block).await?.ok_or_else(|| { + crate::error::QuantusError::Generic(format!( + "Leaf index {} not found in ZK tree at block {:?}", + leaf_index, at_block + )) + }) +} + +/// As [`get_zk_merkle_proof`], but a leaf that is not yet settled into the tree at +/// `at_block` is `Ok(None)` rather than an error. RPC and consistency failures still +/// fail, so callers can tell "too new to prove yet" apart from "the node is unwell". +pub async fn try_get_zk_merkle_proof( + quantus_client: &QuantusClient, + leaf_index: u64, + at_block: subxt::utils::H256, +) -> crate::error::Result> { let proof_params = rpc_params![leaf_index, at_block]; let proof: Option = quantus_client .rpc_client() @@ -216,12 +232,9 @@ pub async fn get_zk_merkle_proof( )) })?; - let proof = proof.ok_or_else(|| { - crate::error::QuantusError::Generic(format!( - "Leaf index {} not found in ZK tree at block {:?}", - leaf_index, at_block - )) - })?; + let Some(proof) = proof else { + return Ok(None); + }; if proof.leaf_index != leaf_index { return Err(crate::error::QuantusError::Generic(format!( @@ -233,7 +246,7 @@ pub async fn get_zk_merkle_proof( // enforces depth == siblings.len(). debug_assert_eq!(proof.depth as usize, proof.siblings.len()); - Ok(proof) + Ok(Some(proof)) } /// Compute sorted siblings and position hints from unsorted siblings. diff --git a/src/collect_rewards_lib.rs b/src/collect_rewards_lib.rs index d1d09ff..5494124 100644 --- a/src/collect_rewards_lib.rs +++ b/src/collect_rewards_lib.rs @@ -18,6 +18,7 @@ use crate::{ }, cli::wormhole::{ compute_merkle_positions, get_zk_merkle_proof, parse_secret_hex as parse_secret_hex_str, + try_get_zk_merkle_proof, }, subsquid::{ compute_address_hash, get_hash_prefix, SubsquidClient, Transfer, TransferQueryParams, @@ -39,6 +40,34 @@ use subxt::{ext::codec::Encode, tx::TxStatus}; /// Result type for collect rewards operations pub type Result = std::result::Result; +/// How many of `ascending` (sorted leaf indices) are settled into the ZK tree at +/// `at_block`. Leaves settle in index order, so provability is monotonic: probing the +/// highest index answers the common case in one call, and a binary search finds the +/// boundary when some are still too new. +async fn count_settled_leaves( + quantus_client: &QuantusClient, + ascending: &[u64], + at_block: subxt::utils::H256, +) -> Result { + let Some(&highest) = ascending.last() else { + return Ok(0); + }; + if try_get_zk_merkle_proof(quantus_client, highest, at_block).await?.is_some() { + return Ok(ascending.len()); + } + + let (mut settled, mut unsettled) = (0usize, ascending.len()); + while settled < unsettled { + let probe = settled + (unsettled - settled) / 2; + if try_get_zk_merkle_proof(quantus_client, ascending[probe], at_block).await?.is_some() { + settled = probe + 1; + } else { + unsettled = probe; + } + } + Ok(settled) +} + const MAX_PRE_SUBMISSION_NULLIFIER_QUERIES: usize = 7; /// Error type for collect rewards operations @@ -362,22 +391,62 @@ pub async fn collect_rewards( .await .map_err(|e| CollectRewardsError::from(format!("Failed to get block: {}", e)))? } else { - // Prove against the latest finalized block. Best-block proofs can be - // invalidated by reorgs before finality (same class as recursive flows). - use subxt::ext::jsonrpsee::{core::client::ClientT, rpc_params}; - let finalized_hash: subxt::utils::H256 = quantus_client - .rpc_client() - .request("chain_getFinalizedHead", rpc_params![]) + // Prove against the head. QPoW finality trails the head by a long way, and a + // leaf minted inside that window is not yet in the finalized tree — proving + // against the finalized block makes recent rewards permanently unsweepable. + // A proof reorged out is re-run; that is much cheaper than waiting for finality. + let latest_hash = quantus_client.get_latest_block().await?; + quantus_client + .client() + .blocks() + .at(latest_hash) .await - .map_err(|e| { - CollectRewardsError::from(format!("Failed to get finalized block hash: {}", e)) - })?; - quantus_client.client().blocks().at(finalized_hash).await.map_err(|e| { - CollectRewardsError::from(format!("Failed to get finalized block: {}", e)) - })? + .map_err(|e| CollectRewardsError::from(format!("Failed to get latest block: {}", e)))? }; let proof_block_hash = proof_block.hash(); + // Leaves settle into the tree in index order, so at any block one threshold + // separates provable from not-yet-settled. Checking it here costs a single RPC + // call on the common path and turns a failure that used to surface part-way + // through proof generation into an immediate one. + let mut leaf_indices = Vec::with_capacity(selected_transfers.len()); + for transfer in &selected_transfers { + leaf_indices.push(transfer.leaf_index.parse::().map_err(|_| { + CollectRewardsError::from(format!("Invalid leaf_index: {}", transfer.leaf_index)) + })?); + } + let mut ascending = leaf_indices.clone(); + ascending.sort_unstable(); + let settled = count_settled_leaves(&quantus_client, &ascending, proof_block_hash).await?; + + if settled < ascending.len() { + let cutoff = ascending[settled]; + progress.on_step( + "warning", + &format!( + "skipping {} of {} transfer(s): leaf index >= {} is not yet settled in the ZK tree at block {:?}. Re-run to sweep them once the chain has folded them in.", + ascending.len() - settled, + ascending.len(), + cutoff, + proof_block_hash + ), + ); + selected_transfers = std::mem::take(&mut selected_transfers) + .into_iter() + .zip(&leaf_indices) + .filter(|(_, leaf_index)| **leaf_index < cutoff) + .map(|(transfer, _)| transfer) + .collect(); + } + + if selected_transfers.is_empty() { + return Err(CollectRewardsError::from(format!( + "None of the {} selected transfer(s) are provable at block {:?}: their leaves are not yet settled in the ZK tree. Re-run once the chain has folded them in, or pass --at-block with a block that includes them.", + ascending.len(), + proof_block_hash + ))); + } + // Step 4: Generate proofs progress.on_step("proofs", &format!("Generating {} proofs", selected_transfers.len())); From 2932cd1c47289711a6cbd049ca8419be422c133f Mon Sep 17 00:00:00 2001 From: Nikolaus Heger Date: Wed, 2 Sep 2026 20:10:15 +0800 Subject: [PATCH 2/2] fix(wormhole): pick the withdrawal set from settled leaves only The settlement filter ran after the --amount selection, so a large not-yet-settled leaf could be selected, then dropped, leaving the sweep short of the requested amount or failing outright while an older settled transfer could have covered it. Establish the settlement boundary before selecting: resolve the proof block, probe the boundary over all unspent transfers, then run the largest-first selection over the settled ones only. The requested amount is validated against the provable total, and the shortfall error says how many transfers were held back as too new. Also format with the pinned nightly rustfmt. --- src/cli/wormhole.rs | 14 +- src/collect_rewards_lib.rs | 286 +++++++++++++++++++++++-------------- 2 files changed, 190 insertions(+), 110 deletions(-) diff --git a/src/cli/wormhole.rs b/src/cli/wormhole.rs index a23a6f3..49d9149 100644 --- a/src/cli/wormhole.rs +++ b/src/cli/wormhole.rs @@ -204,12 +204,14 @@ pub async fn get_zk_merkle_proof( leaf_index: u64, at_block: subxt::utils::H256, ) -> crate::error::Result { - try_get_zk_merkle_proof(quantus_client, leaf_index, at_block).await?.ok_or_else(|| { - crate::error::QuantusError::Generic(format!( - "Leaf index {} not found in ZK tree at block {:?}", - leaf_index, at_block - )) - }) + try_get_zk_merkle_proof(quantus_client, leaf_index, at_block) + .await? + .ok_or_else(|| { + crate::error::QuantusError::Generic(format!( + "Leaf index {} not found in ZK tree at block {:?}", + leaf_index, at_block + )) + }) } /// As [`get_zk_merkle_proof`], but a leaf that is not yet settled into the tree at diff --git a/src/collect_rewards_lib.rs b/src/collect_rewards_lib.rs index 5494124..f0d0596 100644 --- a/src/collect_rewards_lib.rs +++ b/src/collect_rewards_lib.rs @@ -59,7 +59,10 @@ async fn count_settled_leaves( let (mut settled, mut unsettled) = (0usize, ascending.len()); while settled < unsettled { let probe = settled + (unsettled - settled) / 2; - if try_get_zk_merkle_proof(quantus_client, ascending[probe], at_block).await?.is_some() { + if try_get_zk_merkle_proof(quantus_client, ascending[probe], at_block) + .await? + .is_some() + { settled = probe + 1; } else { unsettled = probe; @@ -68,6 +71,65 @@ async fn count_settled_leaves( Ok(settled) } +/// Pick the transfers to withdraw, largest first, from the ones already settled into +/// the ZK tree. `settled_cutoff` is the lowest leaf index not yet in the tree, or +/// `None` when every leaf is settled; transfers at or above it are held back so they +/// cannot crowd a provable transfer out of the selection. Returns the selection and +/// how many transfers were held back. +fn select_provable_transfers( + unspent_transfers: Vec, + settled_cutoff: Option, + requested: Option, +) -> Result<(Vec, usize)> { + // Parse every unspent transfer (fail early if the indexer returned garbage), then + // keep only the provable ones. + let mut provable: Vec<(Transfer, u128)> = Vec::new(); + let mut skipped = 0usize; + for t in unspent_transfers { + let context = format!("transfer {}", t.id); + let amount = parse_transfer_amount(&t.amount, &context)?; + let leaf_index = parse_leaf_index(&t.leaf_index, &context)?; + if settled_cutoff.is_some_and(|cutoff| leaf_index >= cutoff) { + skipped += 1; + continue; + } + provable.push((t, amount)); + } + + let mut total_available: u128 = 0; + for (_, amount) in &provable { + total_available = + checked_add_amount(total_available, *amount, "total available transfers")?; + } + + let withdraw_amount = requested.unwrap_or(total_available); + if withdraw_amount > total_available { + let unsettled_note = if skipped > 0 { + format!(" and {} transfer(s) whose leaves are not yet settled in the ZK tree", skipped) + } else { + String::new() + }; + return Err(CollectRewardsError::from(format!( + "Requested {} but only {} available (after filtering spent nullifiers{})", + withdraw_amount, total_available, unsettled_note + ))); + } + + // Sort by amount descending (largest first) + provable.sort_by_key(|k| std::cmp::Reverse(k.1)); + + let mut selected_transfers = Vec::new(); + let mut selected_total: u128 = 0; + for (t, amt) in provable { + if selected_total >= withdraw_amount { + break; + } + selected_transfers.push(t); + selected_total = checked_add_amount(selected_total, amt, "selected transfers")?; + } + Ok((selected_transfers, skipped)) +} + const MAX_PRE_SUBMISSION_NULLIFIER_QUERIES: usize = 7; /// Error type for collect rewards operations @@ -322,131 +384,89 @@ pub async fn collect_rewards( }); } - // Calculate total available (only unspent) - let mut total_available: u128 = 0; - for t in &unspent_transfers { - let amount = parse_transfer_amount(&t.amount, &format!("transfer {}", t.id))?; - total_available = checked_add_amount(total_available, amount, "total available transfers")?; - } - - // Determine amount to withdraw - let withdraw_amount = config.amount.unwrap_or(total_available); - if withdraw_amount > total_available { - return Err(CollectRewardsError::from(format!( - "Requested {} but only {} available (after filtering spent nullifiers)", - withdraw_amount, total_available - ))); - } - - // Parse amounts for sorting (fail early if any are invalid) - let mut transfers_with_amounts: Vec<(Transfer, u128)> = Vec::new(); - for t in unspent_transfers { - let amt = parse_transfer_amount(&t.amount, &format!("transfer {}", t.id))?; - transfers_with_amounts.push((t, amt)); - } - - // Sort by amount descending (largest first) - transfers_with_amounts.sort_by_key(|k| std::cmp::Reverse(k.1)); - - let mut selected_transfers = Vec::new(); - let mut selected_total: u128 = 0; - for (t, amt) in transfers_with_amounts { - if selected_total >= withdraw_amount { - break; - } - selected_transfers.push(t); - selected_total = checked_add_amount(selected_total, amt, "selected transfers")?; - } - - if config.dry_run { - return Ok(CollectRewardsResult { - wormhole_address, - destination_address: config.destination_address, - total_withdrawn: 0, - batches: vec![], - transfers_processed: selected_transfers.len(), - }); - } - // Get block for proofs - either specific block or latest - let proof_block = if let Some(block_num) = config.at_block { - // Fetch block hash for the specified block number - use subxt::ext::jsonrpsee::{core::client::ClientT, rpc_params}; - let block_hash: Option = quantus_client - .rpc_client() - .request("chain_getBlockHash", rpc_params![block_num]) - .await - .map_err(|e| { - CollectRewardsError::from(format!( - "Failed to get block hash for block {}: {}", - block_num, e - )) + let proof_block = + if let Some(block_num) = config.at_block { + // Fetch block hash for the specified block number + use subxt::ext::jsonrpsee::{core::client::ClientT, rpc_params}; + let block_hash: Option = quantus_client + .rpc_client() + .request("chain_getBlockHash", rpc_params![block_num]) + .await + .map_err(|e| { + CollectRewardsError::from(format!( + "Failed to get block hash for block {}: {}", + block_num, e + )) + })?; + let block_hash = block_hash.ok_or_else(|| { + CollectRewardsError::from(format!("Block {} not found", block_num)) })?; - let block_hash = block_hash - .ok_or_else(|| CollectRewardsError::from(format!("Block {} not found", block_num)))?; - quantus_client - .client() - .blocks() - .at(block_hash) - .await - .map_err(|e| CollectRewardsError::from(format!("Failed to get block: {}", e)))? - } else { - // Prove against the head. QPoW finality trails the head by a long way, and a - // leaf minted inside that window is not yet in the finalized tree — proving - // against the finalized block makes recent rewards permanently unsweepable. - // A proof reorged out is re-run; that is much cheaper than waiting for finality. - let latest_hash = quantus_client.get_latest_block().await?; - quantus_client - .client() - .blocks() - .at(latest_hash) - .await - .map_err(|e| CollectRewardsError::from(format!("Failed to get latest block: {}", e)))? - }; + quantus_client + .client() + .blocks() + .at(block_hash) + .await + .map_err(|e| CollectRewardsError::from(format!("Failed to get block: {}", e)))? + } else { + // Prove against the head. QPoW finality trails the head by a long way, and a + // leaf minted inside that window is not yet in the finalized tree — proving + // against the finalized block makes recent rewards permanently unsweepable. + // A proof reorged out is re-run; that is much cheaper than waiting for finality. + let latest_hash = quantus_client.get_latest_block().await?; + quantus_client.client().blocks().at(latest_hash).await.map_err(|e| { + CollectRewardsError::from(format!("Failed to get latest block: {}", e)) + })? + }; let proof_block_hash = proof_block.hash(); // Leaves settle into the tree in index order, so at any block one threshold - // separates provable from not-yet-settled. Checking it here costs a single RPC - // call on the common path and turns a failure that used to surface part-way - // through proof generation into an immediate one. - let mut leaf_indices = Vec::with_capacity(selected_transfers.len()); - for transfer in &selected_transfers { - leaf_indices.push(transfer.leaf_index.parse::().map_err(|_| { - CollectRewardsError::from(format!("Invalid leaf_index: {}", transfer.leaf_index)) - })?); + // separates provable from not-yet-settled. It has to be established before the + // amount selection, or a not-yet-settled leaf crowds out a settled one that could + // have covered the request. Costs a single RPC call on the common path. + let mut ascending = Vec::with_capacity(unspent_transfers.len()); + for transfer in &unspent_transfers { + ascending + .push(parse_leaf_index(&transfer.leaf_index, &format!("transfer {}", transfer.id))?); } - let mut ascending = leaf_indices.clone(); ascending.sort_unstable(); let settled = count_settled_leaves(&quantus_client, &ascending, proof_block_hash).await?; + let settled_cutoff = ascending.get(settled).copied(); + + let unspent_count = unspent_transfers.len(); + let (selected_transfers, skipped) = + select_provable_transfers(unspent_transfers, settled_cutoff, config.amount)?; - if settled < ascending.len() { - let cutoff = ascending[settled]; + if skipped > 0 { progress.on_step( "warning", &format!( - "skipping {} of {} transfer(s): leaf index >= {} is not yet settled in the ZK tree at block {:?}. Re-run to sweep them once the chain has folded them in.", - ascending.len() - settled, - ascending.len(), - cutoff, + "Skipping {} of {} unspent transfer(s): leaf index >= {} is not yet settled in the ZK tree at block {:?}. Re-run to sweep them once the chain has folded them in.", + skipped, + unspent_count, + settled_cutoff.unwrap_or_default(), proof_block_hash ), ); - selected_transfers = std::mem::take(&mut selected_transfers) - .into_iter() - .zip(&leaf_indices) - .filter(|(_, leaf_index)| **leaf_index < cutoff) - .map(|(transfer, _)| transfer) - .collect(); } if selected_transfers.is_empty() { return Err(CollectRewardsError::from(format!( - "None of the {} selected transfer(s) are provable at block {:?}: their leaves are not yet settled in the ZK tree. Re-run once the chain has folded them in, or pass --at-block with a block that includes them.", - ascending.len(), - proof_block_hash + "None of the {} unspent transfer(s) are provable at block {:?}: their leaves are not yet settled in the ZK tree. Re-run once the chain has folded them in, or pass --at-block with a block that includes them.", + unspent_count, proof_block_hash ))); } + if config.dry_run { + return Ok(CollectRewardsResult { + wormhole_address, + destination_address: config.destination_address, + total_withdrawn: 0, + batches: vec![], + transfers_processed: selected_transfers.len(), + }); + } + // Step 4: Generate proofs progress.on_step("proofs", &format!("Generating {} proofs", selected_transfers.len())); @@ -1214,6 +1234,64 @@ mod tests { assert_eq!(result, 100); } + fn unspent_transfer(id: &str, leaf_index: u64, amount: u128) -> Transfer { + Transfer { + id: id.to_string(), + block_id: "b1".to_string(), + block_height: 1, + timestamp: "2024-01-01T00:00:00.000Z".to_string(), + extrinsic_hash: None, + from_id: "from".to_string(), + to_id: "to".to_string(), + amount: amount.to_string(), + fee: "0".to_string(), + from_hash: "aa".to_string(), + to_hash: "bb".to_string(), + leaf_index: leaf_index.to_string(), + transfer_count: "0".to_string(), + } + } + + #[test] + fn select_provable_transfers_backfills_past_an_unsettled_leaf() { + // The largest transfer is the newest leaf and is not yet in the ZK tree. The + // request must be covered from the settled transfer instead of failing. + let unspent = vec![unspent_transfer("new", 9, 100), unspent_transfer("old", 1, 60)]; + let (selected, skipped) = select_provable_transfers(unspent, Some(9), Some(50)).unwrap(); + + assert_eq!(skipped, 1); + assert_eq!(selected.iter().map(|t| t.id.as_str()).collect::>(), vec!["old"]); + } + + #[test] + fn select_provable_transfers_rejects_a_request_the_settled_leaves_cannot_cover() { + let unspent = vec![unspent_transfer("new", 9, 100), unspent_transfer("old", 1, 60)]; + let err = select_provable_transfers(unspent, Some(9), Some(80)) + .expect_err("80 is not coverable from the 60 that is settled"); + + assert!(err.message.contains("only 60 available"), "unexpected error: {}", err.message); + assert!(err.message.contains("not yet settled"), "unexpected error: {}", err.message); + } + + #[test] + fn select_provable_transfers_sweeps_everything_when_all_leaves_are_settled() { + let unspent = vec![unspent_transfer("new", 9, 100), unspent_transfer("old", 1, 60)]; + let (selected, skipped) = select_provable_transfers(unspent, None, None).unwrap(); + + assert_eq!(skipped, 0); + // Largest first. + assert_eq!(selected.iter().map(|t| t.id.as_str()).collect::>(), vec!["new", "old"]); + } + + #[test] + fn select_provable_transfers_selects_nothing_when_no_leaf_is_settled() { + let unspent = vec![unspent_transfer("new", 9, 100), unspent_transfer("older", 1, 60)]; + let (selected, skipped) = select_provable_transfers(unspent, Some(1), None).unwrap(); + + assert!(selected.is_empty()); + assert_eq!(skipped, 2); + } + const TEST_SECRET_HEX: &str = "0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20"; const TEST_MNEMONIC: &str = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon art";