Fix `ShardTransition.shard_data_roots` and add test

1. Fix `ShardTransition.shard_data_roots`: use `get_block_data_merkle_root` helper to calculate it.
2. Rework `get_valid_custody_chunk_response` testing helper: accept
`block_length_or_custody_data`
3. Add `test_with_custody_challenge_and_response` test

setup.py

@@ -221,6 +221,13 @@ get_attesting_indices = cache_this(
 
 
 PHASE1_SUNDRY_FUNCTIONS = '''
+
+def get_block_data_merkle_root(data: ByteList) -> Root:
+    # To get the Merkle root of the block data, we need the Merkle root without the length Mixing
+    # The below implements this in the Remerkleable framework
+    return data.get_backing().get_left().merkle_root()
+
+
 _get_start_shard = get_start_shard
 get_start_shard = cache_this(
     lambda state, slot: (state.validators.hash_tree_root(), slot),

specs/phase1/validator.md

@@ -275,6 +275,8 @@ Set `attestation_data.shard_head_root = hash_tree_root(shard_head_block)`.
 
 Set `shard_transition` to the value returned by `get_shard_transition(head_state, shard, shard_blocks)`.
 
+`get_block_data_merkle_root(data: ByteList) -> Root` is the function that returns the Merkle root of the block data without the length mixing.
+
 ```python
 def get_shard_transition_fields(
     beacon_state: BeaconState,
@@ -294,7 +296,7 @@ def get_shard_transition_fields(
     for slot in offset_slots:
         if slot in shard_block_slots:
             shard_block = shard_blocks[shard_block_slots.index(slot)]
-            shard_data_roots.append(hash_tree_root(shard_block.message.body))
+            shard_data_roots.append(get_block_data_merkle_root(shard_block.message.body))
         else:
             shard_block = SignedShardBlock(message=ShardBlock(slot=slot, shard=shard))
             shard_data_roots.append(Root())

tests/core/pyspec/eth2spec/test/helpers/custody.py

@@ -136,17 +136,16 @@ def build_proof(anchor, leaf_index):
     return list(reversed(proof))
 
 
-def get_block_data_merkle_root(data_as_bytelist):
-    # To get the Merkle root of the block data, we need the Merkle root without the length Mixing
-    # The below implements this in the Remerkleable framework
-    return data_as_bytelist.get_backing().get_left().merkle_root()
-
-
-def get_valid_custody_chunk_response(spec, state, chunk_challenge, block_length, challenge_index,
+def get_valid_custody_chunk_response(spec, state, chunk_challenge, challenge_index,
+                                     block_length_or_custody_data,
                                      invalid_chunk_data=False):
-    custody_data = get_custody_test_vector(block_length)
+    if isinstance(block_length_or_custody_data, int):
+        custody_data = get_custody_test_vector(block_length_or_custody_data)
+    else:
+        custody_data = block_length_or_custody_data
+
     custody_data_block = ByteList[spec.MAX_SHARD_BLOCK_SIZE](custody_data)
-    chunks = custody_chunkify(spec, custody_data)
+    chunks = custody_chunkify(spec, custody_data_block)
 
     chunk_index = chunk_challenge.chunk_index
 
@@ -166,7 +165,7 @@ def get_custody_test_vector(bytelength, offset=0):
 
 
 def get_sample_shard_transition(spec, start_slot, block_lengths):
-    b = [get_block_data_merkle_root(ByteList[spec.MAX_SHARD_BLOCK_SIZE](get_custody_test_vector(x)))
+    b = [spec.get_block_data_merkle_root(ByteList[spec.MAX_SHARD_BLOCK_SIZE](get_custody_test_vector(x)))
          for x in block_lengths]
     shard_transition = spec.ShardTransition(
         start_slot=start_slot,
@@ -201,5 +200,5 @@ def get_custody_slashable_shard_transition(spec, start_slot, block_lengths, cust
     slashable_test_vector = get_custody_slashable_test_vector(spec, custody_secret,
                                                               block_lengths[0], slashable=slashable)
     block_data = ByteList[spec.MAX_SHARD_BLOCK_SIZE](slashable_test_vector)
-    shard_transition.shard_data_roots[0] = get_block_data_merkle_root(block_data)
+    shard_transition.shard_data_roots[0] = spec.get_block_data_merkle_root(block_data)
     return shard_transition, slashable_test_vector

tests/core/pyspec/eth2spec/test/phase1/block_processing/test_process_chunk_challenge.py

@@ -242,7 +242,8 @@ def test_custody_response(spec, state):
 
     chunk_challenge_index = state.custody_chunk_challenge_index - 1
 
-    custody_response = get_valid_custody_chunk_response(spec, state, challenge, 2**15 // 3, chunk_challenge_index)
+    custody_response = get_valid_custody_chunk_response(
+        spec, state, challenge, chunk_challenge_index, block_length_or_custody_data=2**15 // 3)
 
     yield from run_custody_chunk_response_processing(spec, state, custody_response)
 
@@ -270,7 +271,8 @@ def test_custody_response_multiple_epochs(spec, state):
 
     chunk_challenge_index = state.custody_chunk_challenge_index - 1
 
-    custody_response = get_valid_custody_chunk_response(spec, state, challenge, 2**15 // 3, chunk_challenge_index)
+    custody_response = get_valid_custody_chunk_response(
+        spec, state, challenge, chunk_challenge_index, block_length_or_custody_data=2**15 // 3)
 
     yield from run_custody_chunk_response_processing(spec, state, custody_response)
 
@@ -298,6 +300,7 @@ def test_custody_response_many_epochs(spec, state):
 
     chunk_challenge_index = state.custody_chunk_challenge_index - 1
 
-    custody_response = get_valid_custody_chunk_response(spec, state, challenge, 2**15 // 3, chunk_challenge_index)
+    custody_response = get_valid_custody_chunk_response(
+        spec, state, challenge, chunk_challenge_index, block_length_or_custody_data=2**15 // 3)
 
     yield from run_custody_chunk_response_processing(spec, state, custody_response)

tests/core/pyspec/eth2spec/test/phase1/epoch_processing/test_process_custody_final_updates.py

@@ -159,7 +159,8 @@ def test_validator_withdrawal_resume_after_chunk_challenge_response(spec, state)
     assert state.validators[validator_index].withdrawable_epoch == spec.FAR_FUTURE_EPOCH
 
     chunk_challenge_index = state.custody_chunk_challenge_index - 1
-    custody_response = get_valid_custody_chunk_response(spec, state, challenge, 2**15 // 3, chunk_challenge_index)
+    custody_response = get_valid_custody_chunk_response(
+        spec, state, challenge, chunk_challenge_index, block_length_or_custody_data=2**15 // 3)
 
     _, _, _ = run_custody_chunk_response_processing(spec, state, custody_response)
 

tests/core/pyspec/eth2spec/test/phase1/sanity/test_blocks.py

@@ -7,8 +7,13 @@ from eth2spec.test.context import (
 )
 from eth2spec.test.helpers.attestations import get_valid_on_time_attestation
 from eth2spec.test.helpers.block import build_empty_block
+from eth2spec.test.helpers.custody import (
+    get_valid_chunk_challenge,
+    get_valid_custody_chunk_response,
+)
 from eth2spec.test.helpers.shard_block import (
     build_shard_block,
+    get_committee_index_of_shard,
     get_sample_shard_block_body,
     get_shard_transitions,
 )
@@ -16,6 +21,25 @@ from eth2spec.test.helpers.shard_transitions import is_full_crosslink
 from eth2spec.test.helpers.state import state_transition_and_sign_block, transition_to_valid_shard_slot, transition_to
 
 
+def run_beacon_block(spec, state, block, valid=True):
+    yield 'pre', state.copy()
+
+    if not valid:
+        state_transition_and_sign_block(spec, state, block, expect_fail=True)
+        yield 'block', block
+        yield 'post', None
+        return
+
+    state_transition_and_sign_block(spec, state, block)
+    yield 'block', block
+    yield 'post', state
+
+
+#
+# Beacon block with non-empty shard transitions
+#
+
+
 def run_beacon_block_with_shard_blocks(spec, state, target_len_offset_slot, committee_index, shard, valid=True):
     transition_to(spec, state, state.slot + target_len_offset_slot)
 
@@ -102,3 +126,45 @@ def test_process_beacon_block_with_empty_proposal_transition(spec, state):
     assert state.shard_states[shard].slot == state.slot - 1
 
     yield from run_beacon_block_with_shard_blocks(spec, state, target_len_offset_slot, committee_index, shard)
+
+
+#
+# Beacon block with custody operations
+#
+
+
+@with_all_phases_except([PHASE0])
+@spec_state_test
+def test_with_custody_challenge_and_response(spec, state):
+    # NOTE: this test is only for full crosslink (minimal config), not for mainnet
+    if not is_full_crosslink(spec, state):
+        # skip
+        return
+
+    state = transition_to_valid_shard_slot(spec, state)
+
+    # build shard block
+    shard = 0
+    committee_index = get_committee_index_of_shard(spec, state, state.slot, shard)
+    body = get_sample_shard_block_body(spec)
+    shard_block = build_shard_block(spec, state, shard, body=body, slot=state.slot, signed=True)
+    shard_block_dict: Dict[spec.Shard, Sequence[spec.SignedShardBlock]] = {shard: [shard_block]}
+    shard_transitions = get_shard_transitions(spec, state, shard_block_dict)
+    attestation = get_valid_on_time_attestation(
+        spec, state, index=committee_index,
+        shard_transition=shard_transitions[shard], signed=True,
+    )
+
+    block = build_empty_block(spec, state, slot=state.slot + 1)
+    block.body.attestations = [attestation]
+    block.body.shard_transitions = shard_transitions
+
+    # Custody operations
+    challenge = get_valid_chunk_challenge(spec, state, attestation, shard_transitions[shard])
+    block.body.chunk_challenges = [challenge]
+    chunk_challenge_index = state.custody_chunk_challenge_index
+    custody_response = get_valid_custody_chunk_response(
+        spec, state, challenge, chunk_challenge_index, block_length_or_custody_data=body)
+    block.body.chunk_challenge_responses = [custody_response]
+
+    yield from run_beacon_block(spec, state, block)