From 7ebe8ec15eaa1a355b5b434d2a62814126d233f1 Mon Sep 17 00:00:00 2001 From: Antonio Date: Wed, 26 Aug 2026 06:52:56 +0000 Subject: [PATCH] fix(blend): improved recovery state logic (#3394) --- .../crypto/core_and_leader/send.rs | 13 +- .../core_and_leader/send_and_receive.rs | 3 + .../src/message_blend/crypto/mod.rs | 6 + .../src/message_blend/crypto/test_utils.rs | 1 + .../src/message_blend/provers/core/mod.rs | 35 ++- .../src/message_blend/provers/core/tests.rs | 6 +- .../provers/core_and_leader/mod.rs | 54 ++-- .../provers/core_and_leader/tests.rs | 60 +++-- .../provers/core_leader_and_pow/mod.rs | 4 + .../provers/core_leader_and_pow/tests.rs | 7 +- services/blend/src/core/mod.rs | 113 +++++---- services/blend/src/core/processor.rs | 1 + services/blend/src/core/tests/mod.rs | 233 +++++++++++++++++- services/blend/src/core/tests/utils.rs | 3 +- services/blend/src/test_utils/crypto.rs | 26 +- 15 files changed, 462 insertions(+), 103 deletions(-) diff --git a/blend/scheduling/src/message_blend/crypto/core_and_leader/send.rs b/blend/scheduling/src/message_blend/crypto/core_and_leader/send.rs index 14bca7cfc..d627e3050 100644 --- a/blend/scheduling/src/message_blend/crypto/core_and_leader/send.rs +++ b/blend/scheduling/src/message_blend/crypto/core_and_leader/send.rs @@ -5,6 +5,7 @@ use lb_blend_message::{ Error, PaddedPayloadBody, PayloadType, crypto::proofs::PoQVerificationInputsMinusSigningKey, input::EncapsulationInput, }; +use lb_blend_proofs::quota::Quota; use lb_cryptarchia_engine::Epoch; use lb_groth16::fr_to_bytes; use rayon::ThreadPool; @@ -34,6 +35,12 @@ pub struct EpochCryptographicProcessor EpochCryptographicProcessor { + /// `ß_max`: how many layer proofs one encapsulation draws from the + /// generator, and therefore how much quota it spends. + pub const fn num_blend_layers(&self) -> NonZeroU64 { + self.num_blend_layers + } + #[cfg(test)] pub const fn proofs_generator(&self) -> &ProofsGenerator { &self.proofs_generator @@ -58,9 +65,10 @@ where core_proof_of_quota_generator: CorePoQGenerator, epoch: Epoch, pow_mining_pool: Arc, + spent_core_quota: Quota, ) -> Self { tracing::trace!( - "Creating epoch cryptographic processor with public info {public_info:?} and epoch {epoch:?}" + "Creating epoch cryptographic processor with public info {public_info:?} and epoch {epoch:?}, resuming core key indices from {spent_core_quota}" ); let generator_settings = ProofsGeneratorSettings { @@ -76,6 +84,8 @@ where membership, proofs_generator: ProofsGenerator::new( generator_settings, + // Spent core quota == starting key index + spent_core_quota, core_proof_of_quota_generator, ), _phantom: PhantomData, @@ -257,6 +267,7 @@ mod test { MockCorePoQGenerator, Epoch::new(0), Arc::new(ThreadPoolBuilder::new().build().unwrap()), + Quota::ZERO, ); let new_private_inputs = ProofOfLeadershipQuotaInputs { diff --git a/blend/scheduling/src/message_blend/crypto/core_and_leader/send_and_receive.rs b/blend/scheduling/src/message_blend/crypto/core_and_leader/send_and_receive.rs index ad35c7d82..01c7e11fc 100644 --- a/blend/scheduling/src/message_blend/crypto/core_and_leader/send_and_receive.rs +++ b/blend/scheduling/src/message_blend/crypto/core_and_leader/send_and_receive.rs @@ -50,6 +50,7 @@ where non_ephemeral_encryption_key, num_blend_layers, pow_mining_pool, + spent_core_quota, } = settings; Self { receiver_processor: ReceiverEpochCryptographicProcessor::new( @@ -65,6 +66,7 @@ where core_proof_of_quota_generator, epoch, pow_mining_pool, + spent_core_quota, ), } } @@ -179,6 +181,7 @@ mod test { non_ephemeral_encryption_key: [0; _].into(), num_blend_layers: NonZeroU64::new(1).unwrap(), pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, }, Membership::new_without_local(&[Node { address: Multiaddr::empty(), diff --git a/blend/scheduling/src/message_blend/crypto/mod.rs b/blend/scheduling/src/message_blend/crypto/mod.rs index b2febf82a..54f3ea03d 100644 --- a/blend/scheduling/src/message_blend/crypto/mod.rs +++ b/blend/scheduling/src/message_blend/crypto/mod.rs @@ -10,6 +10,7 @@ use lb_blend_message::{ }, }, }; +use lb_blend_proofs::quota::Quota; use lb_codec::{BinaryDecode as _, BinaryEncode as _}; use lb_key_management_system_keys::keys::X25519PrivateKey; use rayon::ThreadPool; @@ -38,6 +39,11 @@ pub struct EpochCryptographicProcessorSettings { /// The dedicated thread pool the `PoW` puzzle search runs on, built once /// for the service and shared by every epoch's processor. pub pow_mining_pool: Arc, + /// How much of this epoch's core quota has already been spent, counted in + /// key indices. Zero for an epoch entered fresh; recovered from the + /// persisted state when a restart lands part-way through one, so the + /// generator resumes rather than replaying key nullifiers. + pub spent_core_quota: Quota, } #[must_use] diff --git a/blend/scheduling/src/message_blend/crypto/test_utils.rs b/blend/scheduling/src/message_blend/crypto/test_utils.rs index 2d4474315..2da778e08 100644 --- a/blend/scheduling/src/message_blend/crypto/test_utils.rs +++ b/blend/scheduling/src/message_blend/crypto/test_utils.rs @@ -47,6 +47,7 @@ impl CoreLeaderAndPowProofsGenerator { fn new( _settings: ProofsGeneratorSettings, + _starting_key_index: KeyIndex, _proof_of_quota_generator: CorePoQGenerator, ) -> Self { Self(None) diff --git a/blend/scheduling/src/message_blend/provers/core/mod.rs b/blend/scheduling/src/message_blend/provers/core/mod.rs index 82727aafb..e598798c4 100644 --- a/blend/scheduling/src/message_blend/provers/core/mod.rs +++ b/blend/scheduling/src/message_blend/provers/core/mod.rs @@ -29,15 +29,25 @@ const LOG_TARGET: &str = blend::scheduling::proofs::CORE; #[async_trait] pub trait CoreProofsGenerator: Sized { /// Instantiate a new generator for the duration of an epoch. - fn new(settings: ProofsGeneratorSettings, proof_of_quota_generator: PoQGenerator) -> Self; + fn new( + settings: ProofsGeneratorSettings, + starting_key_index: KeyIndex, + proof_of_quota_generator: PoQGenerator, + ) -> Self; /// Request a new core proof from the prover. It returns `None` if the /// maximum core quota has already been reached for this epoch. async fn get_next_proof(&mut self) -> Option; } +#[derive(Debug, Clone)] +pub struct CoreProofsGeneratorSettings { + pub common: ProofsGeneratorSettings, + pub starting_core_key_index: KeyIndex, +} + pub struct RealCoreProofsGenerator { remaining_quota: Quota, - pub(super) settings: ProofsGeneratorSettings, + pub(super) settings: CoreProofsGeneratorSettings, proofs_stream: Pin + Send + Sync>>, _phantom: PhantomData, } @@ -47,16 +57,27 @@ impl CoreProofsGenerator for RealCoreProofsGenerator where PoQGenerator: CoreProofOfQuotaGenerator + Clone + Send + Sync + 'static, { - fn new(settings: ProofsGeneratorSettings, proof_of_quota_generator: PoQGenerator) -> Self { + fn new( + settings: ProofsGeneratorSettings, + starting_key_index: KeyIndex, + proof_of_quota_generator: PoQGenerator, + ) -> Self { Self { proofs_stream: Box::pin(create_proof_stream( settings.public_inputs, proof_of_quota_generator, - KeyIndex::ZERO, + starting_key_index, buffer_size(settings.encapsulation_layers.get() as usize), )), - remaining_quota: settings.public_inputs.core.quota, - settings, + remaining_quota: settings + .public_inputs + .core + .quota + .saturating_sub(starting_key_index), + settings: CoreProofsGeneratorSettings { + common: settings, + starting_core_key_index: starting_key_index, + }, _phantom: PhantomData, } } @@ -72,7 +93,7 @@ where tracing::warn!(target: LOG_TARGET, "No proof available from the stream."); return None; }; - tracing::trace!(target: LOG_TARGET, "Generated core Blend layer proof with key nullifier {:?} addressed to node at index {:?} in {:?} ms.", hex::encode(fr_to_bytes(&proof.proof_of_quota.key_nullifier())), proof.proof_of_selection.expected_index(self.settings.membership_size), start.elapsed().as_millis()); + tracing::trace!(target: LOG_TARGET, "Generated core Blend layer proof with key nullifier {:?} addressed to node at index {:?} in {:?} ms.", hex::encode(fr_to_bytes(&proof.proof_of_quota.key_nullifier())), proof.proof_of_selection.expected_index(self.settings.common.membership_size), start.elapsed().as_millis()); Some(proof) } } diff --git a/blend/scheduling/src/message_blend/provers/core/tests.rs b/blend/scheduling/src/message_blend/provers/core/tests.rs index a1402abd7..eb8d02fbd 100644 --- a/blend/scheduling/src/message_blend/provers/core/tests.rs +++ b/blend/scheduling/src/message_blend/provers/core/tests.rs @@ -1,6 +1,9 @@ use std::sync::Arc; -use lb_blend_proofs::{quota::Quota, selection::inputs::VerifyInputs}; +use lb_blend_proofs::{ + quota::{KeyIndex, Quota}, + selection::inputs::VerifyInputs, +}; use lb_cryptarchia_engine::Epoch; use rayon::ThreadPoolBuilder; use test_log::test; @@ -28,6 +31,7 @@ async fn proof_generation() { epoch: Epoch::new(0), pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), }, + KeyIndex::ZERO, CorePoQGeneratorFromPrivateCoreQuotaInputs::new(private_inputs.clone()), ); diff --git a/blend/scheduling/src/message_blend/provers/core_and_leader/mod.rs b/blend/scheduling/src/message_blend/provers/core_and_leader/mod.rs index 71c0bb918..971036c78 100644 --- a/blend/scheduling/src/message_blend/provers/core_and_leader/mod.rs +++ b/blend/scheduling/src/message_blend/provers/core_and_leader/mod.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use async_trait::async_trait; use lb_blend_message::crypto::proofs::PoQVerificationInputsMinusSigningKey; +use lb_blend_proofs::quota::KeyIndex; use lb_cryptarchia_engine::Epoch; use lb_log_targets::blend; @@ -31,6 +32,7 @@ pub trait CoreAndLeaderProofsGenerator: Sized { /// Instantiate a new generator for the duration of an epoch. fn new( settings: ProofsGeneratorSettings, + starting_key_index: KeyIndex, core_proof_of_quota_generator: CorePoQGenerator, ) -> Self; /// Notify the proof generator about the stream of winning `PoL` slots for @@ -59,10 +61,13 @@ pub struct RealCoreAndLeaderProofsGenerator { impl RealCoreAndLeaderProofsGenerator { #[cfg(test)] - pub fn override_settings(&mut self, new_settings: ProofsGeneratorSettings) { + pub fn override_settings( + &mut self, + new_settings: crate::message_blend::provers::core::CoreProofsGeneratorSettings, + ) { self.core_proofs_generator.settings = new_settings.clone(); if let Some(leader_proofs_generator) = &mut self.leader_proofs_generator { - leader_proofs_generator.settings = new_settings; + leader_proofs_generator.settings = new_settings.common; } } } @@ -75,11 +80,13 @@ where { fn new( settings: ProofsGeneratorSettings, + starting_key_index: KeyIndex, core_proof_of_quota_generator: CorePoQGenerator, ) -> Self { Self { core_proofs_generator: RealCoreProofsGenerator::new( settings, + starting_key_index, core_proof_of_quota_generator, ), leader_proofs_generator: None, @@ -95,18 +102,29 @@ where ) { // TODO: Change trait API to avoid runtime panics. let (current_generator_epoch, current_leader_inputs) = ( - self.core_proofs_generator.settings.epoch, - self.core_proofs_generator.settings.public_inputs.leader, + self.core_proofs_generator.settings.common.epoch, + self.core_proofs_generator + .settings + .common + .public_inputs + .leader, ); assert!( current_generator_epoch == reference_epoch, "set_epoch_private should be called with a reference epoch matching the current core proofs generator's epoch." ); - let current_epoch_local_node_index = self.core_proofs_generator.settings.local_node_index; - let current_epoch_membership_size = self.core_proofs_generator.settings.membership_size; - let current_epoch_core_public_inputs = - self.core_proofs_generator.settings.public_inputs.core; - let current_epoch_pow_public_inputs = self.core_proofs_generator.settings.public_inputs.pow; + let current_epoch_local_node_index = + self.core_proofs_generator.settings.common.local_node_index; + let current_epoch_membership_size = + self.core_proofs_generator.settings.common.membership_size; + let current_epoch_core_public_inputs = self + .core_proofs_generator + .settings + .common + .public_inputs + .core; + let current_epoch_pow_public_inputs = + self.core_proofs_generator.settings.common.public_inputs.pow; self.leader_proofs_generator = Some(RealLeaderProofsGenerator::new( ProofsGeneratorSettings { @@ -118,8 +136,14 @@ where leader: current_leader_inputs, pow: current_epoch_pow_public_inputs, }, - encapsulation_layers: self.core_proofs_generator.settings.encapsulation_layers, - pow_mining_pool: Arc::clone(&self.core_proofs_generator.settings.pow_mining_pool), + encapsulation_layers: self + .core_proofs_generator + .settings + .common + .encapsulation_layers, + pow_mining_pool: Arc::clone( + &self.core_proofs_generator.settings.common.pow_mining_pool, + ), }, winning_pol_info_stream, )); @@ -129,10 +153,10 @@ where let proof = self.core_proofs_generator.get_next_proof().await?; tracing::trace!( target: LOG_TARGET, - epoch = ?self.core_proofs_generator.settings.epoch, - quota = %self.core_proofs_generator.settings.public_inputs.core.quota, - membership_size = self.core_proofs_generator.settings.membership_size, - local_node_index = ?self.core_proofs_generator.settings.local_node_index, + epoch = ?self.core_proofs_generator.settings.common.epoch, + quota = %self.core_proofs_generator.settings.common.public_inputs.core.quota, + membership_size = self.core_proofs_generator.settings.common.membership_size, + local_node_index = ?self.core_proofs_generator.settings.common.local_node_index, key_nullifier = ?proof.proof_of_quota.key_nullifier(), signing_key = ?proof.ephemeral_signing_key.public_key(), "generated core PoQ" diff --git a/blend/scheduling/src/message_blend/provers/core_and_leader/tests.rs b/blend/scheduling/src/message_blend/provers/core_and_leader/tests.rs index 5cdbde326..db17086ac 100644 --- a/blend/scheduling/src/message_blend/provers/core_and_leader/tests.rs +++ b/blend/scheduling/src/message_blend/provers/core_and_leader/tests.rs @@ -1,13 +1,17 @@ use std::sync::Arc; use futures::stream::repeat; -use lb_blend_proofs::{quota::Quota, selection::inputs::VerifyInputs}; +use lb_blend_proofs::{ + quota::{KeyIndex, Quota}, + selection::inputs::VerifyInputs, +}; use lb_cryptarchia_engine::Epoch; use rayon::ThreadPoolBuilder; use test_log::test; use crate::message_blend::provers::{ ProofsGeneratorSettings, + core::CoreProofsGeneratorSettings, core_and_leader::{CoreAndLeaderProofsGenerator as _, RealCoreAndLeaderProofsGenerator}, test_utils::{ CorePoQGeneratorFromPrivateCoreQuotaInputs, @@ -30,6 +34,7 @@ async fn proof_generation() { epoch: Epoch::new(0), pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), }, + KeyIndex::ZERO, CorePoQGeneratorFromPrivateCoreQuotaInputs::new(core_private_inputs), ); @@ -71,13 +76,16 @@ async fn proof_generation() { // We override all the settings since we fixtures for core and leadership proofs // use a different set of public inputs. - core_and_leader_proofs_generator.override_settings(ProofsGeneratorSettings { - local_node_index: None, - membership_size: 1, - public_inputs: leadership_public_inputs, - encapsulation_layers: 1.try_into().unwrap(), - epoch: Epoch::new(0), - pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + core_and_leader_proofs_generator.override_settings(CoreProofsGeneratorSettings { + common: ProofsGeneratorSettings { + local_node_index: None, + membership_size: 1, + public_inputs: leadership_public_inputs, + encapsulation_layers: 1.try_into().unwrap(), + epoch: Epoch::new(0), + pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + }, + starting_core_key_index: KeyIndex::ZERO, }); core_and_leader_proofs_generator .set_epoch_private(Box::pin(repeat(leadership_private_inputs)), Epoch::new(0)); @@ -109,6 +117,7 @@ async fn proof_generation() { } } +#[expect(clippy::too_many_lines, reason = "Test function.")] #[test(tokio::test)] async fn epoch_private_info() { let core_quota = Quota::new::<10>(); @@ -126,18 +135,22 @@ async fn epoch_private_info() { epoch: Epoch::new(0), pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), }, + KeyIndex::ZERO, CorePoQGeneratorFromPrivateCoreQuotaInputs::new(core_private_inputs.clone()), ); // Switch to leadership inputs before wiring leader private epoch info, because // we use fixtures that yield different public inputs. - core_and_leader_proofs_generator.override_settings(ProofsGeneratorSettings { - local_node_index: None, - membership_size: 1, - public_inputs: leadership_public_inputs, - encapsulation_layers: 1.try_into().unwrap(), - epoch: Epoch::new(0), - pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + core_and_leader_proofs_generator.override_settings(CoreProofsGeneratorSettings { + common: ProofsGeneratorSettings { + local_node_index: None, + membership_size: 1, + public_inputs: leadership_public_inputs, + encapsulation_layers: 1.try_into().unwrap(), + epoch: Epoch::new(0), + pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + }, + starting_core_key_index: KeyIndex::ZERO, }); core_and_leader_proofs_generator @@ -196,13 +209,16 @@ async fn epoch_private_info() { // We override all the settings since we fixtures for core and leadership proofs // use a different set of public inputs. - core_and_leader_proofs_generator.override_settings(ProofsGeneratorSettings { - local_node_index: None, - membership_size: 1, - public_inputs: core_public_inputs, - encapsulation_layers: 1.try_into().unwrap(), - epoch: Epoch::new(0), - pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + core_and_leader_proofs_generator.override_settings(CoreProofsGeneratorSettings { + common: ProofsGeneratorSettings { + local_node_index: None, + membership_size: 1, + public_inputs: core_public_inputs, + encapsulation_layers: 1.try_into().unwrap(), + epoch: Epoch::new(0), + pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + }, + starting_core_key_index: KeyIndex::ZERO, }); // We test that core proof generation still works fine diff --git a/blend/scheduling/src/message_blend/provers/core_leader_and_pow/mod.rs b/blend/scheduling/src/message_blend/provers/core_leader_and_pow/mod.rs index 7f63ce86b..16545599a 100644 --- a/blend/scheduling/src/message_blend/provers/core_leader_and_pow/mod.rs +++ b/blend/scheduling/src/message_blend/provers/core_leader_and_pow/mod.rs @@ -1,4 +1,5 @@ use async_trait::async_trait; +use lb_blend_proofs::quota::KeyIndex; use lb_cryptarchia_engine::Epoch; use lb_log_targets::blend; @@ -28,6 +29,7 @@ pub trait CoreLeaderAndPowProofsGenerator: Sized { /// Instantiate a new generator for the duration of an epoch. fn new( settings: ProofsGeneratorSettings, + starting_key_index: KeyIndex, core_proof_of_quota_generator: CorePoQGenerator, ) -> Self; /// Notify the proof generator about the stream of winning `PoL` slots for @@ -65,11 +67,13 @@ where { fn new( settings: ProofsGeneratorSettings, + starting_key_index: KeyIndex, core_proof_of_quota_generator: CorePoQGenerator, ) -> Self { Self { core_and_leader_proofs_generator: RealCoreAndLeaderProofsGenerator::new( settings.clone(), + starting_key_index, core_proof_of_quota_generator, ), // The `PoW` branch depends only on public epoch information, so diff --git a/blend/scheduling/src/message_blend/provers/core_leader_and_pow/tests.rs b/blend/scheduling/src/message_blend/provers/core_leader_and_pow/tests.rs index c12f43989..cea35ed5e 100644 --- a/blend/scheduling/src/message_blend/provers/core_leader_and_pow/tests.rs +++ b/blend/scheduling/src/message_blend/provers/core_leader_and_pow/tests.rs @@ -1,6 +1,9 @@ use std::sync::Arc; -use lb_blend_proofs::{quota::Quota, selection::inputs::VerifyInputs}; +use lb_blend_proofs::{ + quota::{KeyIndex, Quota}, + selection::inputs::VerifyInputs, +}; use lb_cryptarchia_engine::Epoch; use rayon::ThreadPoolBuilder; use test_log::test; @@ -42,6 +45,7 @@ async fn pow_proof_generation() { epoch: Epoch::new(0), pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), }, + KeyIndex::ZERO, CorePoQGeneratorFromPrivateCoreQuotaInputs::new(core_private_inputs), ); @@ -86,6 +90,7 @@ async fn core_proofs_are_delegated() { epoch: Epoch::new(0), pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), }, + KeyIndex::ZERO, CorePoQGeneratorFromPrivateCoreQuotaInputs::new(core_private_inputs), ); diff --git a/services/blend/src/core/mod.rs b/services/blend/src/core/mod.rs index 46b24fc5a..e099e8312 100644 --- a/services/blend/src/core/mod.rs +++ b/services/blend/src/core/mod.rs @@ -617,26 +617,6 @@ where pow: current_epoch_public_info.poq_pow_public_inputs, }; - let crypto_processor = CurrentEpochCryptographicProcessor::< - _, - KmsAdapter::CorePoQGenerator, - ProofsGenerator, - ProofsVerifier, - >::try_new_with_core_condition_check( - current_epoch_public_info.membership.clone(), - blend_config.minimum_network_size, - EpochCryptographicProcessorSettings { - non_ephemeral_encryption_key: blend_config.non_ephemeral_signing_key.derive_x25519(), - num_blend_layers: blend_config.num_blend_layers, - pow_mining_pool: Arc::clone(&blend_config.pow_mining_pool), - }, - current_epoch_poq_verification_inputs, - current_epoch_core_poq_generator - .expect("Core PoQ generator must be present at startup: the proxy service only launches CoreMode when the node is part of the core membership."), - current_epoch_public_info.epoch, - ) - .expect("The initial membership should satisfy the core node condition"); - // Initialize the current epoch state. If the epoch matches the stored one, // retrieves the tracked consumed core quota. Else, fallback to `0`. let current_recovery_checkpoint = match last_saved_state.take() { @@ -655,26 +635,47 @@ where current_epoch_public_info.epoch ); - // Everything else in a stale state belongs to the epoch it was - // saved under, but a transaction still waiting for a `PoW` solution - // has not been encapsulated and so belongs to none: it outlives the - // state that carried it, the same way it outlives an epoch rotation. - let pending_transactions = - maybe_stale_state.map_or_else(VecDeque::new, |state| state.into_components().4); - - ServiceState::with_epoch( - current_epoch_public_info.epoch, - pending_transactions, - EpochBlendingTokenCollector::new( - &reward::EpochInfo::new( + let current_epoch_reward_info = reward::EpochInfo::new( current_epoch_public_info.epoch, ¤t_epoch_public_info.poq_leadership_public_inputs.pol_epoch_nonce, current_epoch_public_info.membership.size() as u64, current_epoch_public_info.poq_core_public_inputs.quota, blend_config.activity_threshold_sensitivity, - ).expect("Reward epoch info must be created successfully. Panicking since the service cannot continue with this epoch") - ), - None, + ).expect("Reward epoch info must be created successfully. Panicking since the service cannot continue with this epoch"); + + // Everything else in a stale state belongs to the epoch it was + // saved under, but a transaction still waiting for a `PoW` solution + // has not been encapsulated and so belongs to none: it outlives the + // state that carried it, the same way it outlives an epoch rotation. + // + // The tokens that state collected are the exception. A state saved + // under the immediately preceding epoch holds a full epoch's worth + // of them, and they are still worth an activity proof: rotating that + // collector here is the same move the running service makes at an + // epoch boundary, and it hands the proof to the submission below. + // A gap of two or more epochs is past submitting for, so it is + // dropped. + let (pending_transactions, recovered_old_epoch_token_collector) = maybe_stale_state + .map_or_else( + || (VecDeque::new(), None), + |state| { + let is_previous_epoch = state.last_seen_epoch().strict_add(1.into()) + == current_epoch_public_info.epoch; + let (_, _, _, _, pending_transactions, token_collector, ..) = + state.into_components(); + let old_epoch_token_collector = is_previous_epoch.then(|| { + tracing::debug!(target: LOG_TARGET, "Recovered a token collector for the immediately preceding epoch. Rotating it so its activity proof is not lost."); + token_collector.rotate_epoch(¤t_epoch_reward_info).1 + }); + (pending_transactions, old_epoch_token_collector) + }, + ); + + ServiceState::with_epoch( + current_epoch_public_info.epoch, + pending_transactions, + EpochBlendingTokenCollector::new(¤t_epoch_reward_info), + recovered_old_epoch_token_collector, state_updater, ) .expect("service state should be created successfully") @@ -692,11 +693,34 @@ where } let current_recovery_checkpoint = state_updater.commit_changes(); + let epoch_core_quota = + blend_config.epoch_core_quota(current_epoch_public_info.membership.size()); + let spent_core_quota = current_recovery_checkpoint.spent_quota(); + + let crypto_processor = CurrentEpochCryptographicProcessor::< + _, + KmsAdapter::CorePoQGenerator, + ProofsGenerator, + ProofsVerifier, + >::try_new_with_core_condition_check( + current_epoch_public_info.membership.clone(), + blend_config.minimum_network_size, + EpochCryptographicProcessorSettings { + non_ephemeral_encryption_key: blend_config.non_ephemeral_signing_key.derive_x25519(), + num_blend_layers: blend_config.num_blend_layers, + pow_mining_pool: Arc::clone(&blend_config.pow_mining_pool), + spent_core_quota, + }, + current_epoch_poq_verification_inputs, + current_epoch_core_poq_generator + .expect("Core PoQ generator must be present at startup: the proxy service only launches CoreMode when the node is part of the core membership."), + current_epoch_public_info.epoch, + ) + .expect("The initial membership should satisfy the core node condition"); + let message_scheduler = SchedulerWrapper::new_with_initial_messages( SchedulerEpochInfo { - core_quota: blend_config - .epoch_core_quota(current_epoch_public_info.membership.size()) - .saturating_sub(current_recovery_checkpoint.spent_quota()), + core_quota: epoch_core_quota.saturating_sub(spent_core_quota), epoch: current_epoch_public_info.epoch, }, BlakeRng::from_entropy(), @@ -1226,6 +1250,7 @@ where .derive_x25519(), num_blend_layers: settings.num_blend_layers, pow_mining_pool: Arc::clone(&settings.pow_mining_pool), + spent_core_quota: Quota::ZERO, }, new_poq_verification_inputs, core_poq_generator, @@ -1912,8 +1937,6 @@ where if state_updater.remove_sent_data_message(data_message_to_blend).is_err() { tracing::warn!(target: LOG_TARGET, "Recovered data message should be present in the recovery state but was not found."); } - // Each data message that is sent is one less cover message that should be generated, hence we consume one core quota per data message here. - state_updater.consume_core_quota(Quota::ONE); }).map( |data_message_to_blend| -> BoxFuture<'_, ()> { backend.publish(data_message_to_blend, current_epoch).boxed() @@ -2129,14 +2152,18 @@ where .encapsulate_cover_payload(&random_sized_bytes::<{ size_of::() }>()) .await .expect("Should not fail to generate new cover message"); + // Each message consumes `num_blend_layers` indices. + state_updater.consume_core_quota( + Quota::try_new(cryptographic_processor.num_blend_layers().get()) + .expect("Number of blend layers must fit within the `PoQ` quota width."), + ); let self_decapsulation_output = cryptographic_processor .receiver() .decapsulate_message_recursive(encapsulated_cover_message.clone()); let Ok(multi_layer_decapsulation_output) = self_decapsulation_output else { - // First layer not addressed to ourselves. Publish as regular cover message, - // hence we consume a core quota. + // First layer not addressed to ourselves, so it goes out fully encapsulated. + // The quota it spent was already recorded above. tracing::trace!(target: LOG_TARGET, "Locally generated cover message does not have its outermost layer addressed to us. Sending it out fully encapsulated..."); - state_updater.consume_core_quota(Quota::ONE); return Some(encapsulated_cover_message.into()); }; let (blending_tokens, message_type) = multi_layer_decapsulation_output.into_components(); diff --git a/services/blend/src/core/processor.rs b/services/blend/src/core/processor.rs index 1ef4187b6..187504e7a 100644 --- a/services/blend/src/core/processor.rs +++ b/services/blend/src/core/processor.rs @@ -401,6 +401,7 @@ mod tests { non_ephemeral_encryption_key: key(local_id).0.derive_x25519(), num_blend_layers: NonZeroU64::new(1).unwrap(), pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, } } diff --git a/services/blend/src/core/tests/mod.rs b/services/blend/src/core/tests/mod.rs index 0de27f04c..8ee531008 100644 --- a/services/blend/src/core/tests/mod.rs +++ b/services/blend/src/core/tests/mod.rs @@ -1,4 +1,4 @@ -use core::{num::NonZeroU64, time::Duration}; +use core::{iter::once, num::NonZeroU64, time::Duration}; use std::{collections::VecDeque, sync::Arc}; use futures::{StreamExt as _, stream::repeat}; @@ -40,7 +40,10 @@ use crate::{ membership::{MembershipInfo, ZkInfo, chain::BlendEpochState}, message::{BlendPayload, ServiceMessage}, test_utils::{ - crypto::{GatedPowProofsGenerator, MockCoreAndLeaderProofsGenerator, PowGate}, + crypto::{ + GatedPowProofsGenerator, MockCoreAndLeaderProofsGenerator, PowGate, + recorded_starting_core_key_indices, reset_starting_core_key_indices, + }, epoch::OncePolStreamProvider, }, }; @@ -88,6 +91,7 @@ async fn test_handle_incoming_blend_message() { non_ephemeral_encryption_key: settings.non_ephemeral_signing_key.derive_x25519(), num_blend_layers: settings.num_blend_layers, pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, }, &public_info, (), @@ -141,6 +145,7 @@ async fn test_handle_incoming_blend_message() { non_ephemeral_encryption_key: settings.non_ephemeral_signing_key.derive_x25519(), num_blend_layers: settings.num_blend_layers, pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, }, &public_info, (), @@ -150,7 +155,8 @@ async fn test_handle_incoming_blend_message() { let (_, _, _, _, _, current_token_collector, _, state_updater) = recovery_checkpoint.into_components(); let (new_token_collector, old_token_collector) = - current_token_collector.rotate_epoch(&reward_epoch_info(&public_info)); + EpochBlendingTokenCollector::clone(¤t_token_collector) + .rotate_epoch(&reward_epoch_info(&public_info)); // Check that decapsulating the same message fails with the new processor // but succeeds with the old one. Also, it should be scheduled in the old @@ -240,6 +246,7 @@ async fn test_handle_incoming_blend_message() { non_ephemeral_encryption_key: settings.non_ephemeral_signing_key.derive_x25519(), num_blend_layers: settings.num_blend_layers, pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, }, &new_epoch_info(epoch, membership, &settings), (), @@ -325,6 +332,7 @@ async fn test_duplicate_decapsulated_replica_handled_gracefully() { non_ephemeral_encryption_key: settings.non_ephemeral_signing_key.derive_x25519(), num_blend_layers: settings.num_blend_layers, pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, }, &public_info, (), @@ -419,6 +427,7 @@ async fn test_handle_incoming_blend_message_with_invalid_poq() { non_ephemeral_encryption_key: settings.non_ephemeral_signing_key.derive_x25519(), num_blend_layers: settings.num_blend_layers, pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, }, &public_info_0, (), @@ -439,6 +448,7 @@ async fn test_handle_incoming_blend_message_with_invalid_poq() { non_ephemeral_encryption_key: settings.non_ephemeral_signing_key.derive_x25519(), num_blend_layers: settings.num_blend_layers, pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, }, &public_info_1, (), @@ -574,6 +584,7 @@ async fn test_handle_epoch_event() { non_ephemeral_encryption_key: settings.non_ephemeral_signing_key.derive_x25519(), num_blend_layers: settings.num_blend_layers, pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, }, &public_info, (), @@ -748,6 +759,7 @@ async fn test_handle_epoch_event_membership_change_rewires_backend_and_generator non_ephemeral_encryption_key: settings.non_ephemeral_signing_key.derive_x25519(), num_blend_layers: settings.num_blend_layers, pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, }, &public_info, (), @@ -851,6 +863,7 @@ async fn transition_to_new_epoch_with_secret(secret_epoch: Epoch) -> Vec non_ephemeral_encryption_key: settings.non_ephemeral_signing_key.derive_x25519(), num_blend_layers: settings.num_blend_layers, pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, }, &public_info, (), @@ -953,6 +966,7 @@ async fn test_handle_epoch_event_empty_epoch_retires() { non_ephemeral_encryption_key: settings.non_ephemeral_signing_key.derive_x25519(), num_blend_layers: settings.num_blend_layers, pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, }, &public_info, (), @@ -1027,6 +1041,7 @@ async fn test_handle_epoch_event_non_empty_without_local_core_path_retires() { non_ephemeral_encryption_key: settings.non_ephemeral_signing_key.derive_x25519(), num_blend_layers: settings.num_blend_layers, pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, }, &public_info, (), @@ -1538,6 +1553,7 @@ async fn test_proof_generator_epoch_binding() { non_ephemeral_encryption_key: settings.non_ephemeral_signing_key.derive_x25519(), num_blend_layers: settings.num_blend_layers, pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, }, &public_info_0, (), @@ -1548,6 +1564,7 @@ async fn test_proof_generator_epoch_binding() { non_ephemeral_encryption_key: settings.non_ephemeral_signing_key.derive_x25519(), num_blend_layers: settings.num_blend_layers, pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()), + spent_core_quota: Quota::ZERO, }, &public_info_1, (), @@ -1670,12 +1687,18 @@ async fn test_proof_generator_epoch_binding() { async fn test_initialize_recovers_matching_saved_state() { let minimal_network_size = 2; let (membership, local_private_key) = new_membership(minimal_network_size); - let settings = settings( - local_private_key.clone(), - u64::from(minimal_network_size).try_into().unwrap(), - (), - 0, - ); + let settings = { + let mut settings = settings( + local_private_key.clone(), + u64::from(minimal_network_size).try_into().unwrap(), + (), + 0, + ); + // More than one layer, so the emission slots the recovery state counts and the + // key indices the generator resumes from cannot be confused for each other. + settings.num_blend_layers = 3.try_into().unwrap(); + settings + }; let initial_epoch = 0.into(); @@ -1712,10 +1735,13 @@ async fn test_initialize_recovers_matching_saved_state() { ) .unwrap(); let mut updater = saved_state.start_updating(); - updater.consume_core_quota(Quota::new::<5>()); + // Five cover messages' worth, at three layers each. + updater.consume_core_quota(Quota::new::<15>()); updater.queue_unencapsulated_transaction(b"transaction".to_vec()); let saved_state = updater.commit_changes(); + reset_starting_core_key_indices(); + let ( _remaining_epoch_stream, _current_public_info, @@ -1746,7 +1772,7 @@ async fn test_initialize_recovers_matching_saved_state() { assert_eq!( recovered_checkpoint.spent_quota(), - Quota::new::<5>(), + Quota::new::<15>(), "Matching epoch: spent_quota should be restored from saved state" ); assert_eq!(recovered_checkpoint.last_seen_epoch(), initial_epoch); @@ -1757,6 +1783,15 @@ async fn test_initialize_recovers_matching_saved_state() { Some(&b"transaction".to_vec()), "Matching epoch: a queued transaction should be restored from saved state" ); + // The restored quota is also what tells the core proof generator where to pick + // up — it is counted in proofs, and the generator hands out one per key index. + // Re-proving an index re-derives a key nullifier this node already put on the + // wire, and peers drop that as a duplicate. + assert_eq!( + recorded_starting_core_key_indices(), + vec![Quota::new::<15>()], + "Matching epoch: the generator should resume from the spent quota" + ); // Mismatched epoch: fresh state should be created @@ -1849,6 +1884,182 @@ async fn test_initialize_recovers_matching_saved_state() { ); } +/// The tokens collected during an epoch are worth an activity proof, and a +/// restart into the *next* epoch must not throw them away. +#[test_log::test(tokio::test)] +async fn test_initialize_submits_activity_proof_for_the_previous_epoch() { + let minimal_network_size = 2; + let (membership, local_private_key) = new_membership(minimal_network_size); + let mut settings = settings( + local_private_key.clone(), + u64::from(minimal_network_size).try_into().unwrap(), + (), + 0, + ); + // A long epoch makes the core quota, and with it the activity threshold, high + // enough for a single token to clear it. + settings.time.rounds_per_epoch = 648_000.try_into().unwrap(); + + // Saved under epoch 0; the node comes back up in epoch 1. + let saved_epoch = 0.into(); + let current_epoch = Epoch::new(1); + + let membership_info = MembershipInfo { + membership: membership.clone(), + zk: Some(ZkInfo { + root: ZkHash::ZERO, + core_and_path_selectors: Some([(ZkHash::ZERO, false); CORE_MERKLE_TREE_HEIGHT]), + }), + }; + let (membership_stream, membership_sender) = new_stream(); + membership_sender + .send(test_blend_epoch_state(1, membership_info)) + .await + .unwrap(); + + let (overwatch_handle, _overwatch_cmd_receiver, state_updater, _state_receiver) = + dummy_overwatch_resources(); + let (sdp_relay, mut sdp_relay_receiver) = sdp_relay(); + + let saved_public_info = new_epoch_info(saved_epoch, membership.clone(), &settings); + let saved_state = ServiceState::with_epoch( + saved_epoch, + VecDeque::new(), + EpochBlendingTokenCollector::new(&reward_epoch_info(&saved_public_info)), + None, + state_updater.clone(), + ) + .unwrap(); + let mut updater = saved_state.start_updating(); + updater.collect_current_epoch_tokens(once(BlendingToken::new( + Ed25519Key::from_bytes(&[0; _]).public_key(), + VerifiedProofOfQuota::from_bytes_unchecked([0; _]), + VerifiedProofOfSelection::from_bytes_unchecked([0; _]), + ))); + let saved_state = updater.commit_changes(); + + let ( + _remaining_epoch_stream, + _current_public_info, + _crypto_processor, + recovered_checkpoint, + _pending_transactions, + _message_scheduler, + _backend, + _rng, + ) = initialize::< + NodeId, + TestBlendBackend, + TestPayloadDispatcher, + MockCoreAndLeaderProofsGenerator, + MockProofsVerifier, + MockKmsAdapter, + RuntimeServiceId, + >( + settings.clone(), + membership_stream, + overwatch_handle, + MockKmsAdapter, + &sdp_relay, + Some(saved_state), + state_updater, + ) + .await; + + assert_eq!( + recovered_checkpoint.last_seen_epoch(), + current_epoch, + "the recovered state should track the epoch the node came back up in" + ); + assert!( + matches!( + sdp_relay_receiver.try_recv(), + Ok(lb_sdp_service::SdpMessage::PostActivity { + metadata: ActivityMetadata::Blend(_), + }) + ), + "the previous epoch's tokens should be submitted as an activity proof, not dropped" + ); +} + +/// A state older than the immediately preceding epoch is past submitting for. +/// +/// The counterpart to +/// [`test_initialize_submits_activity_proof_for_the_previous_epoch`]. +#[test_log::test(tokio::test)] +async fn test_initialize_drops_activity_proof_older_than_one_epoch() { + let minimal_network_size = 2; + let (membership, local_private_key) = new_membership(minimal_network_size); + let mut settings = settings( + local_private_key.clone(), + u64::from(minimal_network_size).try_into().unwrap(), + (), + 0, + ); + settings.time.rounds_per_epoch = 648_000.try_into().unwrap(); + + let membership_info = MembershipInfo { + membership: membership.clone(), + zk: Some(ZkInfo { + root: ZkHash::ZERO, + core_and_path_selectors: Some([(ZkHash::ZERO, false); CORE_MERKLE_TREE_HEIGHT]), + }), + }; + // The node comes back up in epoch 3, so a state saved under epoch 1 is + // genuinely two epochs behind rather than merely different. + let (membership_stream, membership_sender) = new_stream(); + membership_sender + .send(test_blend_epoch_state(3, membership_info)) + .await + .unwrap(); + + let (overwatch_handle, _overwatch_cmd_receiver, state_updater, _state_receiver) = + dummy_overwatch_resources(); + let (sdp_relay, mut sdp_relay_receiver) = sdp_relay(); + + let stale_epoch = 1.into(); + let stale_public_info = new_epoch_info(stale_epoch, membership.clone(), &settings); + let stale_state = ServiceState::with_epoch( + stale_epoch, + VecDeque::new(), + EpochBlendingTokenCollector::new(&reward_epoch_info(&stale_public_info)), + None, + state_updater.clone(), + ) + .unwrap(); + let mut updater = stale_state.start_updating(); + updater.collect_current_epoch_tokens(once(BlendingToken::new( + Ed25519Key::from_bytes(&[0; _]).public_key(), + VerifiedProofOfQuota::from_bytes_unchecked([0; _]), + VerifiedProofOfSelection::from_bytes_unchecked([0; _]), + ))); + let stale_state = updater.commit_changes(); + + let (.., _rng) = initialize::< + NodeId, + TestBlendBackend, + TestPayloadDispatcher, + MockCoreAndLeaderProofsGenerator, + MockProofsVerifier, + MockKmsAdapter, + RuntimeServiceId, + >( + settings.clone(), + membership_stream, + overwatch_handle, + MockKmsAdapter, + &sdp_relay, + Some(stale_state), + state_updater, + ) + .await; + + assert!( + sdp_relay_receiver.try_recv().is_err(), + "a collector more than one epoch old should be dropped, not submitted" + ); +} + /// A transaction waits for a `PoW` solution without holding up anything else. /// /// The puzzle search behind a transaction's layer proofs takes long enough that diff --git a/services/blend/src/core/tests/utils.rs b/services/blend/src/core/tests/utils.rs index c4e5ccd86..7a9e3df41 100644 --- a/services/blend/src/core/tests/utils.rs +++ b/services/blend/src/core/tests/utils.rs @@ -36,7 +36,7 @@ use lb_core::crypto::ZkHash; use lb_groth16::{AdditiveGroup as _, Fr, fr_from_bytes_unchecked, fr_to_bytes}; use lb_key_management_system_service::keys::{Ed25519PublicKey, UnsecuredEd25519Key}; use lb_network_service::{NetworkService, backends::NetworkBackend}; -use lb_poq::CorePathAndSelectors; +use lb_poq::{CorePathAndSelectors, KeyIndex}; use lb_sdp_service::SdpMessage; use overwatch::{ overwatch::{OverwatchHandle, commands::OverwatchCommand}, @@ -475,6 +475,7 @@ impl CoreLeaderAndPowProofsGenerator { fn new( settings: ProofsGeneratorSettings, + _starting_key_index: KeyIndex, _core_proof_of_quota_generator: CorePoQGenerator, ) -> Self { Self(settings.public_inputs.leader.pol_epoch_nonce) diff --git a/services/blend/src/test_utils/crypto.rs b/services/blend/src/test_utils/crypto.rs index ba39dc6ae..ceed4ce2c 100644 --- a/services/blend/src/test_utils/crypto.rs +++ b/services/blend/src/test_utils/crypto.rs @@ -10,7 +10,7 @@ use lb_blend::{ encap::ProofsVerifier, }, proofs::{ - quota::{ProofOfQuota, VerifiedProofOfQuota}, + quota::{KeyIndex, ProofOfQuota, VerifiedProofOfQuota}, selection::{ProofOfSelection, VerifiedProofOfSelection, inputs::VerifyInputs}, }, scheduling::message_blend::provers::{ @@ -22,6 +22,27 @@ use lb_chain_service::Epoch; use lb_key_management_system_service::keys::{Ed25519PublicKey, UnsecuredEd25519Key}; use tokio::sync::watch; +thread_local! { + /// Records the core key index each [`MockCoreAndLeaderProofsGenerator`] was + /// built to start from, so tests can assert that a recovered quota reaches + /// the generator rather than it silently restarting at zero. Reliable + /// because `#[tokio::test]` uses a single-threaded runtime, so the value is + /// test-isolated. + static STARTING_CORE_KEY_INDICES: RefCell> = const { RefCell::new(Vec::new()) }; +} + +/// Clears the record of generator starting key indices. Call before the code +/// under test to isolate the constructions of interest. +pub fn reset_starting_core_key_indices() { + STARTING_CORE_KEY_INDICES.with(|indices| indices.borrow_mut().clear()); +} + +/// Returns the starting key index of every generator built since the last +/// reset, in construction order. +pub fn recorded_starting_core_key_indices() -> Vec { + STARTING_CORE_KEY_INDICES.with(|indices| indices.borrow().clone()) +} + pub struct MockCoreAndLeaderProofsGenerator; #[async_trait] @@ -30,8 +51,10 @@ impl CoreLeaderAndPowProofsGenerator { fn new( _settings: ProofsGeneratorSettings, + starting_key_index: KeyIndex, _core_proof_of_quota_generator: CorePoQGenerator, ) -> Self { + STARTING_CORE_KEY_INDICES.with(|indices| indices.borrow_mut().push(starting_key_index)); Self } @@ -184,6 +207,7 @@ impl CoreLeaderAndPowProofsGenerator { fn new( _settings: ProofsGeneratorSettings, + _starting_key_index: KeyIndex, _core_proof_of_quota_generator: CorePoQGenerator, ) -> Self { Self