Simplify Gloas lookup test setup

beacon_node/beacon_chain/src/test_utils.rs

@@ -2,10 +2,7 @@ use crate::block_verification_types::{AsBlock, AvailableBlockData, LookupBlock,
 use crate::custody_context::NodeCustodyType;
 use crate::data_availability_checker::DataAvailabilityChecker;
 use crate::graffiti_calculator::GraffitiSettings;
-use crate::kzg_utils::{
-    blobs_to_data_column_sidecars_gloas, build_data_column_sidecars_fulu,
-    build_data_column_sidecars_gloas,
-};
+use crate::kzg_utils::{build_data_column_sidecars_fulu, build_data_column_sidecars_gloas};
 use crate::observed_operations::ObservationOutcome;
 pub use crate::persisted_beacon_chain::PersistedBeaconChain;
 use crate::{BeaconBlockResponseWrapper, CustodyContext, get_block_root};
@@ -1167,7 +1164,7 @@ where
     /// For pre-Gloas forks, the envelope is `None` and this behaves like `make_block`.
     pub async fn make_block_with_envelope(
         &self,
-        state: BeaconState<E>,
+        mut state: BeaconState<E>,
         slot: Slot,
     ) -> (
         SignedBlockContentsTuple<E>,
@@ -1180,6 +1177,17 @@ where
         if state.fork_name_unchecked().gloas_enabled()
             || self.spec.fork_name_at_slot::<E>(slot).gloas_enabled()
         {
+            complete_state_advance(&mut state, None, slot, &self.spec)
+                .expect("should be able to advance state to slot");
+            state.build_caches(&self.spec).expect("should build caches");
+
+            let proposer_index = state.get_beacon_proposer_index(slot, &self.spec).unwrap();
+
+            let graffiti = Graffiti::from(self.rng.lock().random::<[u8; 32]>());
+            let graffiti_settings =
+                GraffitiSettings::new(Some(graffiti), Some(GraffitiPolicy::PreserveUserGraffiti));
+            let randao_reveal = self.sign_randao_reveal(&state, proposer_index, slot);
+
             // Load the parent's payload envelope and status from the cached head.
             // TODO(gloas): we may want to pass these as arguments to support cases where we build
             // on alternate chains to the head.
@@ -1191,118 +1199,59 @@ where
                 )
             };
 
-            let (block_contents, envelope, _columns, state) = self
-                .make_gloas_block_with_status(state, slot, parent_payload_status, parent_envelope)
-                .await;
-            (block_contents, envelope, state)
+            let (block, post_block_state, _consensus_block_value) = self
+                .chain
+                .produce_block_on_state_gloas(
+                    state,
+                    None,
+                    parent_payload_status,
+                    parent_envelope,
+                    slot,
+                    randao_reveal,
+                    graffiti_settings,
+                    ProduceBlockVerification::VerifyRandao,
+                    None,
+                )
+                .await
+                .unwrap();
+
+            let signed_block = Arc::new(block.sign(
+                &self.validator_keypairs[proposer_index].sk,
+                &post_block_state.fork(),
+                post_block_state.genesis_validators_root(),
+                &self.spec,
+            ));
+
+            // Retrieve the cached envelope produced during block production and sign it.
+            let signed_envelope = self
+                .chain
+                .pending_payload_envelopes
+                .write()
+                .remove(slot)
+                .map(|envelope| {
+                    let epoch = slot.epoch(E::slots_per_epoch());
+                    let domain = self.spec.get_domain(
+                        epoch,
+                        Domain::BeaconBuilder,
+                        &post_block_state.fork(),
+                        post_block_state.genesis_validators_root(),
+                    );
+                    let message = envelope.signing_root(domain);
+                    let signature = self.validator_keypairs[proposer_index].sk.sign(message);
+                    SignedExecutionPayloadEnvelope {
+                        message: envelope,
+                        signature,
+                    }
+                });
+
+            let block_contents: SignedBlockContentsTuple<E> = (signed_block, None);
+            (block_contents, signed_envelope, post_block_state)
         } else {
             let (block_contents, state) = self.make_block(state, slot).await;
             (block_contents, None, state)
         }
     }
 
-    /// Like the Gloas branch of `make_block_with_envelope`, but takes the parent payload status and
-    /// envelope explicitly so callers can build on alternate parents (e.g. FULL vs EMPTY children).
-    pub async fn make_gloas_block_with_status(
-        &self,
-        mut state: BeaconState<E>,
-        slot: Slot,
-        parent_payload_status: proto_array::PayloadStatus,
-        parent_envelope: Option<Arc<SignedExecutionPayloadEnvelope<E>>>,
-    ) -> (
-        SignedBlockContentsTuple<E>,
-        Option<SignedExecutionPayloadEnvelope<E>>,
-        DataColumnSidecarList<E>,
-        BeaconState<E>,
-    ) {
-        complete_state_advance(&mut state, None, slot, &self.spec)
-            .expect("should be able to advance state to slot");
-        state.build_caches(&self.spec).expect("should build caches");
-
-        let proposer_index = state.get_beacon_proposer_index(slot, &self.spec).unwrap();
-
-        let graffiti = Graffiti::from(self.rng.lock().random::<[u8; 32]>());
-        let graffiti_settings =
-            GraffitiSettings::new(Some(graffiti), Some(GraffitiPolicy::PreserveUserGraffiti));
-        let randao_reveal = self.sign_randao_reveal(&state, proposer_index, slot);
-
-        let (block, post_block_state, _consensus_block_value) = self
-            .chain
-            .produce_block_on_state_gloas(
-                state,
-                None,
-                parent_payload_status,
-                parent_envelope,
-                slot,
-                randao_reveal,
-                graffiti_settings,
-                ProduceBlockVerification::VerifyRandao,
-                None,
-            )
-            .await
-            .unwrap();
-
-        let signed_block = Arc::new(block.sign(
-            &self.validator_keypairs[proposer_index].sk,
-            &post_block_state.fork(),
-            post_block_state.genesis_validators_root(),
-            &self.spec,
-        ));
-
-        let block_root = signed_block.canonical_root();
-
-        // Build the gloas data column sidecars from the blobs produced during block production.
-        // For gloas, blobs travel in the execution payload envelope, so the columns are keyed by
-        // the block root and slot rather than carried by the block body.
-        let data_columns = self
-            .chain
-            .pending_payload_envelopes
-            .write()
-            .take_blobs(slot)
-            .map(|blobs| {
-                let blob_refs: Vec<_> = blobs.iter().collect();
-                blobs_to_data_column_sidecars_gloas(
-                    &blob_refs,
-                    block_root,
-                    slot,
-                    &self.chain.kzg,
-                    &self.spec,
-                )
-                .expect("should build gloas data column sidecars")
-            })
-            .unwrap_or_default();
-
-        // Retrieve the cached envelope produced during block production and sign it.
-        let signed_envelope = self
-            .chain
-            .pending_payload_envelopes
-            .write()
-            .remove(slot)
-            .map(|envelope| {
-                let epoch = slot.epoch(E::slots_per_epoch());
-                let domain = self.spec.get_domain(
-                    epoch,
-                    Domain::BeaconBuilder,
-                    &post_block_state.fork(),
-                    post_block_state.genesis_validators_root(),
-                );
-                let message = envelope.signing_root(domain);
-                let signature = self.validator_keypairs[proposer_index].sk.sign(message);
-                SignedExecutionPayloadEnvelope {
-                    message: envelope,
-                    signature,
-                }
-            });
-
-        let block_contents: SignedBlockContentsTuple<E> = (signed_block, None);
-        (
-            block_contents,
-            signed_envelope,
-            data_columns,
-            post_block_state,
-        )
-    }
-
     /// Useful for the `per_block_processing` tests. Creates a block, and returns the state after
     /// caches are built but before the generated block is processed.
     pub async fn make_block_return_pre_state(

beacon_node/network/src/sync/tests/lookups.rs

@@ -17,7 +17,7 @@ use beacon_chain::{
     block_verification_types::{AsBlock, AvailableBlockData},
     test_utils::{
         AttestationStrategy, BeaconChainHarness, BlockStrategy, EphemeralHarnessType, NumBlobs,
-        generate_rand_block_and_blobs, test_spec,
+        generate_data_column_sidecars_from_block, generate_rand_block_and_blobs, test_spec,
     },
 };
 use beacon_processor::{BeaconProcessorChannels, DuplicateCache, Work, WorkEvent};
@@ -1057,62 +1057,76 @@ impl TestRig {
         self.network_blocks_by_slot
             .insert(genesis_block.slot(), genesis_block);
 
-        // Build imported G and A.
-        let mut parents = vec![];
-        for _ in 0..2 {
-            external_harness.advance_slot();
-            let block_root = external_harness
-                .extend_chain(
-                    1,
-                    BlockStrategy::OnCanonicalHead,
-                    AttestationStrategy::AllValidators,
-                )
-                .await;
-            let block = external_harness.get_full_block(&block_root);
-            let block_root = block.canonical_root();
-            let block_slot = block.slot();
-            let block_hash = block.as_block().payload_bid_block_hash().unwrap();
-            self.network_blocks_by_root
-                .insert(block_root, block.clone());
-            self.network_blocks_by_slot.insert(block_slot, block);
-            if let Ok(Some(envelope)) = external_harness.chain.get_payload_envelope(&block_root) {
-                self.network_envelopes_by_root
-                    .insert(block_root, Arc::new(envelope));
-            }
-            parents.push((block_root, block_slot, block_hash));
-        }
-        let [(g_root, _, g_block_hash), (a_root, a_slot, a_block_hash)] =
-            parents.try_into().unwrap();
-
-        let a_state = external_harness.get_current_state();
-        let a_envelope = self.network_envelopes_by_root.get(&a_root).cloned();
-        let g_envelope = self.network_envelopes_by_root.get(&g_root).cloned();
-
-        let child_slot = a_slot + 1;
-
-        // B: FULL child of A.
-        let (b_contents, b_envelope, b_columns, _) = external_harness
-            .make_gloas_block_with_status(
-                a_state.clone(),
-                child_slot,
-                proto_array::PayloadStatus::Full,
-                a_envelope,
+        external_harness.advance_slot();
+        let g_root = external_harness
+            .extend_chain(
+                1,
+                BlockStrategy::OnCanonicalHead,
+                AttestationStrategy::AllValidators,
             )
             .await;
-        let b_block = b_contents.0;
-        let b_root = b_block.canonical_root();
+        let g_block = external_harness.get_full_block(&g_root);
+        let g_block_hash = g_block.as_block().payload_bid_block_hash().unwrap();
+        self.network_blocks_by_root.insert(g_root, g_block.clone());
+        self.network_blocks_by_slot.insert(g_block.slot(), g_block);
+        self.network_envelopes_by_root.insert(
+            g_root,
+            Arc::new(
+                external_harness
+                    .chain
+                    .get_payload_envelope(&g_root)
+                    .unwrap()
+                    .unwrap(),
+            ),
+        );
+
+        external_harness.advance_slot();
+        let a_slot = external_harness.get_current_slot();
+        let (a_contents, a_envelope, a_state) = external_harness
+            .make_block_with_envelope(external_harness.get_current_state(), a_slot)
+            .await;
+        let a_block = a_contents.0.clone();
+        let a_root = a_block.canonical_root();
+        let a_block_hash = a_block.as_block().payload_bid_block_hash().unwrap();
+        external_harness
+            .process_block(a_slot, a_root, a_contents)
+            .await
+            .unwrap();
+
+        external_harness.advance_slot();
+        let child_slot = external_harness.get_current_slot();
 
         // C: EMPTY child of A.
-        let (c_contents, c_envelope, c_columns, _) = external_harness
-            .make_gloas_block_with_status(
-                a_state.clone(),
-                child_slot,
-                proto_array::PayloadStatus::Empty,
-                g_envelope,
-            )
+        let (c_contents, c_envelope, _) = external_harness
+            .make_block_with_envelope(a_state.clone(), child_slot)
             .await;
         let c_block = c_contents.0;
         let c_root = c_block.canonical_root();
+        let c_columns = generate_data_column_sidecars_from_block(
+            c_block.as_ref(),
+            &external_harness.chain.spec,
+        );
+
+        let a_envelope = a_envelope.expect("A should have envelope");
+        external_harness
+            .process_envelope(a_root, a_envelope.clone(), &a_state, a_block.state_root())
+            .await;
+        let a_block = external_harness.get_full_block(&a_root);
+        self.network_blocks_by_root.insert(a_root, a_block.clone());
+        self.network_blocks_by_slot.insert(a_slot, a_block);
+        self.network_envelopes_by_root
+            .insert(a_root, Arc::new(a_envelope));
+
+        // B: FULL child of A.
+        let (b_contents, b_envelope, _) = external_harness
+            .make_block_with_envelope(a_state.clone(), child_slot)
+            .await;
+        let b_block = b_contents.0;
+        let b_root = b_block.canonical_root();
+        let b_columns = generate_data_column_sidecars_from_block(
+            b_block.as_ref(),
+            &external_harness.chain.spec,
+        );
 
         assert_eq!(
             (
@@ -1216,7 +1230,16 @@ impl TestRig {
     }
 
     fn corrupt_last_column_kzg_proof(&mut self) {
-        let range_sync_block = self.get_last_block().clone();
+        let block_root = self.get_last_block().canonical_root();
+        self.corrupt_column_kzg_proof(block_root);
+    }
+
+    fn corrupt_column_kzg_proof(&mut self, block_root: Hash256) {
+        let range_sync_block = self
+            .network_blocks_by_root
+            .get(&block_root)
+            .unwrap_or_else(|| panic!("No block for root {block_root}"))
+            .clone();
         let block = range_sync_block.block_cloned();
         let blobs = range_sync_block.block_data().blobs();
         let mut columns = range_sync_block
@@ -1227,7 +1250,7 @@ impl TestRig {
         let column = Arc::make_mut(first);
         let proof = column.kzg_proofs_mut().first_mut().expect("no kzg proofs");
         *proof = kzg::KzgProof::empty();
-        self.re_insert_block(block, blobs, Some(columns));
+        self.upsert_block(block, blobs, Some(columns));
     }
 
     fn get_last_block(&self) -> &RangeSyncBlock<E> {
@@ -1247,6 +1270,15 @@ impl TestRig {
     ) {
         self.network_blocks_by_slot.clear();
         self.network_blocks_by_root.clear();
+        self.upsert_block(block, blobs, columns);
+    }
+
+    fn upsert_block(
+        &mut self,
+        block: Arc<SignedBeaconBlock<E>>,
+        blobs: Option<types::BlobSidecarList<E>>,
+        columns: Option<types::DataColumnSidecarList<E>>,
+    ) {
         let block_root = block.canonical_root();
         let block_slot = block.slot();
         let block_data = if let Some(columns) = columns {
@@ -2675,9 +2707,16 @@ async fn crypto_on_fail_with_bad_column_kzg_proof() {
     let Some(mut r) = TestRig::new_fulu_peer_test(FuluTestType::WeSupernodeThemSupernode) else {
         return;
     };
-    r.build_chain(1).await;
-    r.corrupt_last_column_kzg_proof();
-    r.trigger_with_last_block();
+    if r.is_after_gloas() {
+        r.build_chain(2).await;
+        let child = r.get_last_block().block_cloned();
+        r.corrupt_column_kzg_proof(child.parent_root());
+        r.trigger_unknown_parent_blocks_from_all_peers(&[child]);
+    } else {
+        r.build_chain(1).await;
+        r.corrupt_last_column_kzg_proof();
+        r.trigger_with_last_block();
+    }
     r.simulate(SimulateConfig::happy_path()).await;
     if cfg!(feature = "fake_crypto") {
         r.assert_successful_lookup_sync();