chore(blend): improve PoQ usage on message encapsulations (#3411)

blend/provers/src/crypto/core_and_leader/send.rs

@@ -1,14 +1,15 @@
-use core::{hash::Hash, marker::PhantomData};
+use core::{hash::Hash, marker::PhantomData, mem};
 use std::{num::NonZeroU64, sync::Arc};
 
 use lb_blend_membership::Membership;
 use lb_blend_message::{
-    Error, PaddedPayloadBody, PayloadType, crypto::proofs::PoQVerificationInputsMinusSigningKey,
-    input::EncapsulationInput,
+    Error, MAX_PAYLOAD_BODY_SIZE, 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 lb_log_targets::blend;
 use rayon::ThreadPool;
 
 use crate::{
@@ -19,6 +20,8 @@ use crate::{
     },
 };
 
+const LOG_TARGET: &str = blend::processor::core_and_leader::SEND;
+
 /// [`EpochCryptographicProcessor`] is responsible for only wrapping
 /// cover and data messages for the message indistinguishability.
 ///
@@ -27,9 +30,41 @@ pub struct EpochCryptographicProcessor<NodeId, CorePoQGenerator, ProofsGenerator
     num_blend_layers: NonZeroU64,
     membership: Membership<NodeId>,
     proofs_generator: ProofsGenerator,
+    partial_draws: PartialDraws,
     _phantom: PhantomData<CorePoQGenerator>,
 }
 
+/// Layer proofs drawn for a message that was never finished.
+///
+/// A draw that stops short — because the quota ran out, or because the caller
+/// was cancelled part-way — used to drop what it had. Every one of those proofs
+/// had already been paid for, and a core one has spent its key index for the
+/// epoch, so its nullifier can never be minted again. Keeping them here lets
+/// the next attempt carry on instead.
+///
+/// Kept per branch so resuming cannot quietly change which quota backs a
+/// message. Within a branch the proofs are interchangeable: nothing binds one
+/// to a particular payload until it is encapsulated.
+#[derive(Default)]
+struct PartialDraws {
+    leader: crate::crypto::leader::send::PartialDraws,
+    cover: Vec<BlendLayerProof>,
+}
+
+impl PartialDraws {
+    const fn for_type(&mut self, payload_type: PayloadType) -> &mut Vec<BlendLayerProof> {
+        match payload_type {
+            PayloadType::Cover => &mut self.cover,
+            PayloadType::BlockProposal => self
+                .leader
+                .for_type(crate::crypto::leader::send::PayloadType::BlockProposal),
+            PayloadType::Transaction => self
+                .leader
+                .for_type(crate::crypto::leader::send::PayloadType::Transaction),
+        }
+    }
+}
+
 impl<NodeId, CorePoQGenerator, ProofsGenerator>
     EpochCryptographicProcessor<NodeId, CorePoQGenerator, ProofsGenerator>
 {
@@ -86,6 +121,7 @@ where
                 spent_core_quota,
                 core_proof_of_quota_generator,
             ),
+            partial_draws: PartialDraws::default(),
             _phantom: PhantomData,
         }
     }
@@ -129,26 +165,67 @@ where
             .await
     }
 
-    // TODO: Think about optimizing this by, e.g., using less encapsulations if
-    // there are less than 3 proofs available, or use a proof from a different pool
-    // if needed (core proof for leadership message or leadership proof for
-    // cover message, since the protocol does not enforce that).
     async fn encapsulate_payload(
         &mut self,
         payload_type: PayloadType,
         payload: &[u8],
     ) -> Result<EncapsulatedMessageWithVerifiedPublicHeader, Error> {
-        // We validate the payload early on so we don't generate proofs unnecessarily.
-        let validated_payload = PaddedPayloadBody::try_from(payload)?;
-        let mut proofs = Vec::with_capacity(self.num_blend_layers.get() as usize);
-
-        for _ in 0..self.num_blend_layers.into() {
-            let Some(proof) = self.next_proof_for(payload_type).await else {
-                return Err(Error::ProofNotAvailable);
-            };
-            proofs.push(proof);
+        // Refuse a payload that could never fit before spending anything on it.
+        // Only the length check has to happen this early; padding it — an 18 KiB
+        // allocation with a random tail — waits until the proofs are in hand, so
+        // an attempt that comes up short or is cancelled costs nothing.
+        if payload.len() > MAX_PAYLOAD_BODY_SIZE {
+            return Err(Error::PayloadTooLarge);
         }
 
+        let Some(proofs) = self.next_proofs_for(payload_type).await else {
+            return Err(Error::ProofNotAvailable);
+        };
+
+        Ok(self.encapsulate_with(payload_type, PaddedPayloadBody::try_from(payload)?, proofs))
+    }
+
+    /// Draws a whole message's layer proofs, resuming any run a previous
+    /// attempt left unfinished.
+    ///
+    /// Proofs are accumulated on `self` rather than in a local, so a caller
+    /// that is cancelled mid-draw — a `select!` arm losing the race, say —
+    /// leaves them where the next attempt will find them.
+    ///
+    /// A branch that runs out part-way does not sink the message: the wire
+    /// format carries `ß_max` blending headers whatever happens, padding the
+    /// unused ones with random bytes, so a message can go out under fewer real
+    /// layers without telling anyone it did. Returns `None` only when not one
+    /// proof is available, which is the one case with nothing to send.
+    async fn next_proofs_for(&mut self, payload_type: PayloadType) -> Option<Vec<BlendLayerProof>> {
+        let encapsulations = self.num_blend_layers.get() as usize;
+        while self.partial_draws.for_type(payload_type).len() < encapsulations
+            && let Some(layer_proof) = self.next_proof_for(payload_type).await
+        {
+            self.partial_draws.for_type(payload_type).push(layer_proof);
+        }
+
+        let message_proofs = mem::take(self.partial_draws.for_type(payload_type));
+        if message_proofs.is_empty() {
+            return None;
+        }
+
+        if message_proofs.len() < encapsulations {
+            tracing::warn!(
+                target: LOG_TARGET,
+                "Encapsulating a {payload_type:?} message under {} of {encapsulations} layers: its quota branch is exhausted for this epoch.",
+                message_proofs.len()
+            );
+        }
+        Some(message_proofs)
+    }
+
+    fn encapsulate_with(
+        &self,
+        payload_type: PayloadType,
+        validated_payload: PaddedPayloadBody,
+        proofs: Vec<BlendLayerProof>,
+    ) -> EncapsulatedMessageWithVerifiedPublicHeader {
         let membership_size = self.membership.size();
         let proofs_and_signing_keys = proofs
             .into_iter()
@@ -188,13 +265,13 @@ where
             })
             .collect::<Vec<_>>();
 
-        Ok(EncapsulatedMessageWithVerifiedPublicHeader::try_new(
+        EncapsulatedMessageWithVerifiedPublicHeader::try_new(
             &inputs,
             payload_type,
             validated_payload,
             self.num_blend_layers.get() as usize,
         )
-        .expect("Number of encapsulation inputs is in `1..=num_blend_layers`."))
+        .expect("Number of encapsulation inputs is in `1..=num_blend_layers`.")
     }
 
     /// The `PoQ` branch each payload type draws its layer proofs from.
@@ -209,11 +286,12 @@ where
 
 #[cfg(test)]
 mod test {
+    use core::pin::pin;
     use std::{num::NonZeroU64, sync::Arc};
 
-    use futures::{StreamExt as _, stream::repeat};
+    use futures::{StreamExt as _, poll, stream::repeat};
     use lb_blend_membership::{Membership, Node};
-    use lb_blend_message::crypto::proofs::PoQVerificationInputsMinusSigningKey;
+    use lb_blend_message::{PayloadType, crypto::proofs::PoQVerificationInputsMinusSigningKey};
     use lb_blend_proofs::quota::{
         Quota,
         inputs::prove::{
@@ -224,14 +302,17 @@ mod test {
     use lb_core::crypto::ZkHash;
     use lb_cryptarchia_engine::Epoch;
     use lb_groth16::{AdditiveGroup as _, Field as _, Fr};
-    use lb_key_management_system_keys::keys::{ED25519_PUBLIC_KEY_SIZE, Ed25519PublicKey};
+    use lb_key_management_system_keys::keys::{
+        ED25519_PUBLIC_KEY_SIZE, Ed25519PublicKey, UnsecuredEd25519Key,
+    };
     use libp2p::PeerId;
     use multiaddr::Multiaddr;
     use rayon::ThreadPoolBuilder;
 
     use super::EpochCryptographicProcessor;
     use crate::crypto::test_utils::{
-        MockCorePoQGenerator, TestEpochChangeCoreAndLeaderProofsGenerator,
+        MockCorePoQGenerator, RationedCoreProofsGenerator,
+        TestEpochChangeCoreAndLeaderProofsGenerator, exhaust_core_branch, ration_core_proofs,
     };
 
     #[tokio::test]
@@ -282,4 +363,109 @@ mod test {
         let first_slot = processor.proofs_generator.0.as_mut().unwrap().next().await;
         assert!(first_slot == Some(new_private_inputs));
     }
+
+    /// A processor whose core branch is under the test's control.
+    fn rationed_processor(
+        layers: u64,
+    ) -> EpochCryptographicProcessor<PeerId, MockCorePoQGenerator, RationedCoreProofsGenerator>
+    {
+        EpochCryptographicProcessor::new(
+            NonZeroU64::new(layers).unwrap(),
+            Membership::new_without_local(&[Node {
+                address: Multiaddr::empty(),
+                id: PeerId::random(),
+                // A real key: an all-zero one decodes but has no usable shared
+                // secret, which the encapsulation would reject.
+                public_key: UnsecuredEd25519Key::from_bytes(&[7; ED25519_PUBLIC_KEY_SIZE])
+                    .public_key(),
+            }]),
+            PoQVerificationInputsMinusSigningKey {
+                core: CoreInputs {
+                    quota: Quota::ONE,
+                    zk_root: ZkHash::ZERO,
+                },
+                leader: LeaderInputs {
+                    message_quota: Quota::ONE,
+                    pol_epoch_nonce: ZkHash::ZERO,
+                    pol_ledger_aged: ZkHash::ZERO,
+                    lottery_0: Fr::ZERO,
+                    lottery_1: Fr::ZERO,
+                },
+                pow: PowInputs {
+                    pow_quota: Quota::ONE,
+                    pow_blend_difficulty: Fr::ZERO,
+                },
+            },
+            MockCorePoQGenerator,
+            Epoch::new(0),
+            Arc::new(ThreadPoolBuilder::new().build().unwrap()),
+            Quota::ZERO,
+        )
+    }
+
+    /// A branch that runs out part-way still gets its message out.
+    ///
+    /// The wire format carries `ß_max` blending headers whatever happens,
+    /// padding the unused ones with random bytes, so a message can go out
+    /// under fewer real layers without telling anyone. Failing instead
+    /// would also waste the proofs already drawn — and a core proof's key
+    /// index is spent for the epoch.
+    #[tokio::test]
+    async fn a_short_branch_sends_under_fewer_layers() {
+        let mut processor = rationed_processor(3);
+
+        ration_core_proofs(2);
+        exhaust_core_branch(true);
+
+        assert!(
+            processor.encapsulate_cover_payload(&[]).await.is_ok(),
+            "two of three layers is still a message worth sending"
+        );
+    }
+
+    /// One layer is the floor: with no proof at all there is nothing to send.
+    #[tokio::test]
+    async fn no_proofs_at_all_is_the_one_failure() {
+        let mut processor = rationed_processor(3);
+
+        ration_core_proofs(0);
+        exhaust_core_branch(true);
+
+        assert!(processor.encapsulate_cover_payload(&[]).await.is_err());
+    }
+
+    /// A caller abandoned mid-draw leaves its proofs behind for the next one.
+    ///
+    /// Each cost a proving, and a core one has spent a key index whose
+    /// nullifier can never be minted again this epoch — so a `select!` arm
+    /// losing the race must not take them with it.
+    #[tokio::test]
+    async fn proofs_drawn_before_a_cancellation_are_kept() {
+        let mut processor = rationed_processor(3);
+
+        // Two proofs, and then the branch blocks rather than ending.
+        ration_core_proofs(2);
+        exhaust_core_branch(false);
+        {
+            let draw = pin!(processor.next_proofs_for(PayloadType::Cover));
+            assert!(
+                poll!(draw).is_pending(),
+                "the draw should be waiting on a third proof"
+            );
+        }; // dropped here, as `select!` drops a losing arm
+
+        // One more proof, and then the branch is done.
+        ration_core_proofs(1);
+        exhaust_core_branch(true);
+
+        let drawn = processor
+            .next_proofs_for(PayloadType::Cover)
+            .await
+            .expect("one proof is enough to send something");
+        assert_eq!(
+            drawn.len(),
+            3,
+            "the two proofs drawn before the cancellation should have been kept"
+        );
+    }
 }

blend/provers/src/crypto/leader/send.rs

@@ -1,12 +1,13 @@
-use core::hash::Hash;
+use core::{hash::Hash, mem};
 use std::{num::NonZeroU64, sync::Arc};
 
 use lb_blend_membership::Membership;
 use lb_blend_message::{
-    Error, PaddedPayloadBody, crypto::proofs::PoQVerificationInputsMinusSigningKey,
-    input::EncapsulationInput,
+    Error, MAX_PAYLOAD_BODY_SIZE, PaddedPayloadBody,
+    crypto::proofs::PoQVerificationInputsMinusSigningKey, input::EncapsulationInput,
 };
 use lb_cryptarchia_engine::Epoch;
+use lb_log_targets::blend;
 use rayon::ThreadPool;
 
 use crate::{
@@ -17,8 +18,10 @@ use crate::{
     },
 };
 
+const LOG_TARGET: &str = blend::processor::leader::SEND;
+
 #[derive(Debug, Clone, Copy)]
-enum PayloadType {
+pub(crate) enum PayloadType {
     BlockProposal,
     Transaction,
 }
@@ -32,6 +35,23 @@ impl From<PayloadType> for lb_blend_message::PayloadType {
     }
 }
 
+/// Layer proofs drawn for a message that was never finished. See the core
+/// processor's `PartialDraws` for why they are worth keeping.
+#[derive(Default)]
+pub(crate) struct PartialDraws {
+    block_proposal: Vec<BlendLayerProof>,
+    transaction: Vec<BlendLayerProof>,
+}
+
+impl PartialDraws {
+    pub const fn for_type(&mut self, payload_type: PayloadType) -> &mut Vec<BlendLayerProof> {
+        match payload_type {
+            PayloadType::BlockProposal => &mut self.block_proposal,
+            PayloadType::Transaction => &mut self.transaction,
+        }
+    }
+}
+
 /// [`EpochCryptographicProcessor`] is responsible for only wrapping data
 /// messages (no cover messages) for the message indistinguishability.
 ///
@@ -44,6 +64,7 @@ pub struct EpochCryptographicProcessor<NodeId, ProofsGenerator> {
     membership: Membership<NodeId>,
     proofs_generator: ProofsGenerator,
     epoch: Epoch,
+    partial_draws: PartialDraws,
 }
 
 impl<NodeId, ProofsGenerator> EpochCryptographicProcessor<NodeId, ProofsGenerator> {
@@ -78,6 +99,7 @@ where
             membership,
             proofs_generator: ProofsGenerator::new(generator_settings, winning_pol_info_stream),
             epoch,
+            partial_draws: PartialDraws::default(),
         }
     }
 }
@@ -108,17 +130,61 @@ where
         payload_type: PayloadType,
         payload: &[u8],
     ) -> Result<EncapsulatedMessageWithVerifiedPublicHeader, Error> {
-        // We validate the payload early on so we don't generate proofs unnecessarily.
-        let validated_payload = PaddedPayloadBody::try_from(payload)?;
-        let mut proofs = Vec::with_capacity(self.num_blend_layers.get() as usize);
-
-        for _ in 0..self.num_blend_layers.into() {
-            let Some(proof) = self.next_proof_for(payload_type).await else {
-                return Err(Error::ProofNotAvailable);
-            };
-            proofs.push(proof);
+        // Refuse a payload that could never fit before spending anything on it.
+        // Only the length check has to happen this early; padding it — an 18 KiB
+        // allocation with a random tail — waits until the proofs are in hand, so
+        // an attempt that comes up short or is cancelled costs nothing.
+        if payload.len() > MAX_PAYLOAD_BODY_SIZE {
+            return Err(Error::PayloadTooLarge);
         }
 
+        let Some(proofs) = self.next_proofs_for(payload_type).await else {
+            return Err(Error::ProofNotAvailable);
+        };
+
+        Ok(self.encapsulate_with(payload_type, PaddedPayloadBody::try_from(payload)?, proofs))
+    }
+
+    /// Draws a whole message's layer proofs, resuming any run a previous
+    /// attempt left unfinished.
+    ///
+    /// Proofs are accumulated on `self` rather than in a local variable, so a
+    /// caller that is cancelled mid-draw leaves them where the next attempt
+    /// will find them.
+    ///
+    /// A branch that runs out part-way does not sink the message: the wire
+    /// format carries `ß_max` blending headers whatever happens, so a message
+    /// can go out under fewer real layers. Returns `None` only when not one
+    /// proof is available.
+    async fn next_proofs_for(&mut self, payload_type: PayloadType) -> Option<Vec<BlendLayerProof>> {
+        let encapsulations = self.num_blend_layers.get() as usize;
+        while self.partial_draws.for_type(payload_type).len() < encapsulations
+            && let Some(layer_proof) = self.next_proof_for(payload_type).await
+        {
+            self.partial_draws.for_type(payload_type).push(layer_proof);
+        }
+
+        let message_proofs = mem::take(self.partial_draws.for_type(payload_type));
+        if message_proofs.is_empty() {
+            return None;
+        }
+
+        if message_proofs.len() < encapsulations {
+            tracing::warn!(
+                target: LOG_TARGET,
+                "Encapsulating a {payload_type:?} message under {} of {encapsulations} layers: its quota branch is exhausted for this epoch.",
+                message_proofs.len()
+            );
+        }
+        Some(message_proofs)
+    }
+
+    fn encapsulate_with(
+        &self,
+        payload_type: PayloadType,
+        validated_payload: PaddedPayloadBody,
+        proofs: Vec<BlendLayerProof>,
+    ) -> EncapsulatedMessageWithVerifiedPublicHeader {
         let membership_size = self.membership.size();
         let proofs_and_signing_keys = proofs
             .into_iter()
@@ -160,13 +226,13 @@ where
             })
             .collect::<Vec<_>>();
 
-        Ok(EncapsulatedMessageWithVerifiedPublicHeader::try_new(
+        EncapsulatedMessageWithVerifiedPublicHeader::try_new(
             &inputs,
             payload_type.into(),
             validated_payload,
             self.num_blend_layers.get() as usize,
         )
-        .expect("Number of encapsulation inputs is in `1..=num_blend_layers`."))
+        .expect("Number of encapsulation inputs is in `1..=num_blend_layers`.")
     }
 
     /// The `PoQ` branch each payload type draws its layer proofs from.

blend/provers/src/crypto/test_utils.rs

@@ -1,9 +1,10 @@
-use core::convert::Infallible;
+use core::{cell::Cell, convert::Infallible};
 
 use async_trait::async_trait;
 use futures::future::ready;
 use lb_blend_message::{
-    crypto::proofs::PoQVerificationInputsMinusSigningKey, encap::ProofsVerifier,
+    crypto::{key_ext::Ed25519SecretKeyExt as _, proofs::PoQVerificationInputsMinusSigningKey},
+    encap::ProofsVerifier,
 };
 use lb_blend_proofs::{
     quota::{self, KeyIndex, ProofOfQuota, VerifiedProofOfQuota, inputs::prove::PublicInputs},
@@ -11,7 +12,7 @@ use lb_blend_proofs::{
 };
 use lb_core::crypto::ZkHash;
 use lb_cryptarchia_engine::Epoch;
-use lb_key_management_system_keys::keys::Ed25519PublicKey;
+use lb_key_management_system_keys::keys::{Ed25519PublicKey, UnsecuredEd25519Key};
 
 use crate::{
     CoreProofOfQuotaGenerator,
@@ -102,3 +103,69 @@ impl ProofsVerifier for TestEpochChangeProofsVerifier {
         ))
     }
 }
+
+thread_local! {
+    /// How many more core proofs [`RationedCoreProofsGenerator`] will hand out.
+    static CORE_PROOFS_AVAILABLE: Cell<usize> = const { Cell::new(0) };
+    /// Whether running out means "no more this epoch" or "not yet".
+    static CORE_BRANCH_EXHAUSTED: Cell<bool> = const { Cell::new(false) };
+}
+
+/// Lets the next `count` core proof requests succeed. Once they are used up the
+/// branch either reports itself exhausted or blocks, depending on
+/// [`exhaust_core_branch`] — which is the difference between a draw that has to
+/// settle for fewer layers and one that a caller can abandon part-way.
+///
+/// Reliable because `#[tokio::test]` runs on a current-thread runtime.
+pub fn ration_core_proofs(count: usize) {
+    CORE_PROOFS_AVAILABLE.with(|available| available.set(count));
+}
+
+/// Makes a rationed-out core branch report `None` rather than block.
+pub fn exhaust_core_branch(exhausted: bool) {
+    CORE_BRANCH_EXHAUSTED.with(|flag| flag.set(exhausted));
+}
+
+/// A generator whose core branch runs out on demand, so a test can stop a draw
+/// part-way through a message and then let it finish.
+pub struct RationedCoreProofsGenerator;
+
+#[async_trait]
+impl<CorePoQGenerator> CoreLeaderAndPowProofsGenerator<CorePoQGenerator>
+    for RationedCoreProofsGenerator
+{
+    fn new(
+        _settings: ProofsGeneratorSettings,
+        _starting_key_index: KeyIndex,
+        _proof_of_quota_generator: CorePoQGenerator,
+    ) -> Self {
+        Self
+    }
+
+    fn set_epoch_private(&mut self, _: WinningPolInfoStream, _: Epoch) {}
+
+    async fn get_next_core_proof(&mut self) -> Option<BlendLayerProof> {
+        if CORE_PROOFS_AVAILABLE.with(Cell::get) == 0 && !CORE_BRANCH_EXHAUSTED.with(Cell::get) {
+            // Not exhausted, just nothing right now: block, so a caller can be
+            // abandoned mid-draw.
+            core::future::pending::<()>().await;
+        }
+        CORE_PROOFS_AVAILABLE.with(|available| {
+            let left = available.get();
+            available.set(left.checked_sub(1)?);
+            Some(BlendLayerProof {
+                proof_of_quota: VerifiedProofOfQuota::from_bytes_unchecked([0; _]),
+                proof_of_selection: VerifiedProofOfSelection::from_bytes_unchecked([0; _]),
+                ephemeral_signing_key: UnsecuredEd25519Key::generate_with_blake_rng(),
+            })
+        })
+    }
+
+    async fn get_next_leader_proof(&mut self) -> Option<BlendLayerProof> {
+        None
+    }
+
+    async fn get_next_pow_proof(&mut self) -> Option<BlendLayerProof> {
+        None
+    }
+}

blend/provers/src/provers/core/mod.rs

@@ -6,7 +6,7 @@ use lb_blend_message::crypto::{
     key_ext::Ed25519SecretKeyExt as _, proofs::PoQVerificationInputsMinusSigningKey,
 };
 use lb_blend_proofs::{
-    quota::{KeyIndex, Quota, inputs::prove::PublicInputs},
+    quota::{KeyIndex, inputs::prove::PublicInputs},
     selection::VerifiedProofOfSelection,
 };
 use lb_groth16::fr_to_bytes;
@@ -23,7 +23,7 @@ use crate::{
 #[cfg(test)]
 mod tests;
 
-const LOG_TARGET: &str = blend::scheduling::proofs::CORE;
+const LOG_TARGET: &str = blend::prover::CORE;
 
 /// Proof generator for core `PoQ` variants.
 #[async_trait]
@@ -46,7 +46,6 @@ pub struct CoreProofsGeneratorSettings {
 }
 
 pub struct RealCoreProofsGenerator<PoQGenerator> {
-    remaining_quota: Quota,
     pub(super) settings: CoreProofsGeneratorSettings,
     proofs_stream: Pin<Box<dyn Stream<Item = BlendLayerProof> + Send + Sync>>,
     _phantom: PhantomData<PoQGenerator>,
@@ -69,11 +68,6 @@ where
                 starting_key_index,
                 buffer_size(settings.encapsulation_layers.get() as usize),
             )),
-            remaining_quota: settings
-                .public_inputs
-                .core
-                .quota
-                .saturating_sub(starting_key_index),
             settings: CoreProofsGeneratorSettings {
                 common: settings,
                 starting_core_key_index: starting_key_index,
@@ -84,13 +78,13 @@ where
 
     async fn get_next_proof(&mut self) -> Option<BlendLayerProof> {
         let start = Instant::now();
-        let Some(remaining_quota) = self.remaining_quota.checked_sub(Quota::ONE) else {
-            tracing::warn!(target: LOG_TARGET, "Core quota exhausted. No proof is generated.");
-            return None;
-        };
-        self.remaining_quota = remaining_quota;
+        // The stream is built over `quota.values_range_from(starting_key_index)`, so
+        // it runs out at exactly the point a separate counter would have: one proof
+        // per key index the epoch's quota still holds. Letting it be the only bound
+        // keeps the two from ever disagreeing, and leaves nothing to spend before
+        // the `await` that a cancelled caller could take with it.
         let Some(proof) = self.proofs_stream.next().await else {
-            tracing::warn!(target: LOG_TARGET, "No proof available from the stream.");
+            tracing::warn!(target: LOG_TARGET, "Core quota exhausted for this epoch. No proof is generated.");
             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.common.membership_size), start.elapsed().as_millis());

blend/provers/src/provers/core/tests.rs

@@ -59,3 +59,41 @@ async fn proof_generation() {
     // Next proof should be `None` since we ran out of core quota.
     assert!(core_proofs_generator.get_next_proof().await.is_none());
 }
+
+/// How many key indices the run before the restart got through.
+const SPENT: u64 = 4;
+
+/// A resumed generator hands out only what the epoch's quota has left.
+///
+/// The key-index range is the sole bound on core proofs, so it has to account
+/// for the indices a previous run already spent — otherwise resuming would
+/// either hand back proofs the quota cannot cover, or re-mint nullifiers the
+/// earlier run already put on the wire.
+#[test(tokio::test)]
+async fn resumed_generator_is_bounded_by_what_the_quota_has_left() {
+    let core_quota = Quota::new::<10>();
+    let (public_inputs, private_inputs) = valid_proof_of_quota_inputs(core_quota);
+
+    let mut core_proofs_generator = RealCoreProofsGenerator::new(
+        ProofsGeneratorSettings {
+            local_node_index: None,
+            membership_size: 1,
+            public_inputs,
+            encapsulation_layers: 1.try_into().unwrap(),
+            epoch: Epoch::new(0),
+            pow_mining_pool: Arc::new(ThreadPoolBuilder::new().build().unwrap()),
+        },
+        KeyIndex::try_new(SPENT).unwrap(),
+        CorePoQGeneratorFromPrivateCoreQuotaInputs::new(private_inputs),
+    );
+
+    let mut generated = 0u64;
+    while core_proofs_generator.get_next_proof().await.is_some() {
+        generated += 1;
+        assert!(
+            generated <= core_quota.get() - SPENT,
+            "a resumed generator must not outlive the quota it inherited"
+        );
+    }
+    assert_eq!(generated, core_quota.get() - SPENT);
+}

blend/provers/src/provers/core_and_leader/mod.rs

@@ -18,7 +18,7 @@ use crate::{
 #[cfg(test)]
 mod tests;
 
-const LOG_TARGET: &str = blend::scheduling::proofs::CORE_AND_LEADER;
+const LOG_TARGET: &str = blend::prover::CORE_AND_LEADER;
 
 /// Proof generator for core and leader `PoQ` variants.
 ///

blend/provers/src/provers/core_leader_and_pow/mod.rs

@@ -15,7 +15,7 @@ use crate::{
 #[cfg(test)]
 mod tests;
 
-const LOG_TARGET: &str = blend::scheduling::proofs::CORE_LEADER_AND_POW;
+const LOG_TARGET: &str = blend::prover::CORE_LEADER_AND_POW;
 
 /// Proof generator for all three `PoQ` variants.
 ///

blend/provers/src/provers/leader/mod.rs

@@ -27,7 +27,7 @@ use crate::{
 #[cfg(test)]
 mod tests;
 
-const LOG_TARGET: &str = blend::scheduling::proofs::LEADER;
+const LOG_TARGET: &str = blend::prover::LEADER;
 
 /// A `PoQ` generator that deals only with leadership proofs, suitable for edge
 /// nodes.

blend/provers/src/provers/leader_and_pow/mod.rs

@@ -10,7 +10,7 @@ use crate::provers::{
 #[cfg(test)]
 mod tests;
 
-const LOG_TARGET: &str = blend::scheduling::proofs::LEADER_AND_POW;
+const LOG_TARGET: &str = blend::prover::LEADER_AND_POW;
 
 /// Proof generator for the two `PoQ` variants an edge node can reach.
 ///

blend/provers/src/provers/pow/mod.rs

@@ -34,7 +34,7 @@ use crate::provers::{BlendLayerProof, ProofsGeneratorSettings};
 #[cfg(test)]
 mod tests;
 
-const LOG_TARGET: &str = blend::scheduling::proofs::POW;
+const LOG_TARGET: &str = blend::prover::POW;
 const BUFFER_SIZE: usize = 2;
 
 #[must_use]

tracing/targets/src/blend.rs

@@ -13,15 +13,21 @@ log_targets! {
         edge::BEHAVIOUR,
         handler::CORE_EDGE,
     },
+    processor::{
+        core_and_leader::SEND,
+        leader::SEND,
+    },
+    prover::{
+        CORE,
+        CORE_AND_LEADER,
+        CORE_LEADER_AND_POW,
+        LEADER,
+        LEADER_AND_POW,
+        POW
+    },
     scheduling::{
         COVER,
         DELAY,
-        proofs::CORE,
-        proofs::CORE_AND_LEADER,
-        proofs::CORE_LEADER_AND_POW,
-        proofs::LEADER,
-        proofs::LEADER_AND_POW,
-        proofs::POW,
     },
     service::{
         CORE,