diff --git a/src/cli/wormhole.rs b/src/cli/wormhole.rs index 17e08bd..49d9149 100644 --- a/src/cli/wormhole.rs +++ b/src/cli/wormhole.rs @@ -204,6 +204,24 @@ 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 +234,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 +248,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..f0d0596 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,96 @@ 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) +} + +/// 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 @@ -293,40 +384,77 @@ 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")?; - } + // 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 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)) + })? + }; + let proof_block_hash = proof_block.hash(); - // 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 - ))); + // Leaves settle into the tree in index order, so at any block one threshold + // 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))?); } - - // 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)); + 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 skipped > 0 { + progress.on_step( + "warning", + &format!( + "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 + ), + ); } - // 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 selected_transfers.is_empty() { + return Err(CollectRewardsError::from(format!( + "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 { @@ -339,45 +467,6 @@ pub async fn collect_rewards( }); } - // 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 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 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![]) - .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)) - })? - }; - let proof_block_hash = proof_block.hash(); - // Step 4: Generate proofs progress.on_step("proofs", &format!("Generating {} proofs", selected_transfers.len())); @@ -1145,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";