add test coverage for incomplete storage ranges with and without proof

Signed-off-by: garyschulte <[email protected]>

ethereum/eth/src/test/java/org/hyperledger/besu/ethereum/eth/sync/snapsync/PersistDataStepTest.java

@@ -16,7 +16,11 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import org.hyperledger.besu.datatypes.Hash;
@@ -31,18 +35,20 @@
 import org.hyperledger.besu.services.tasks.Task;
 
 import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import org.apache.tuweni.bytes.Bytes;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
 
 public class PersistDataStepTest {
 
   private final WorldStateStorage worldStateStorage =
       new InMemoryKeyValueStorageProvider().createWorldStateStorage(DataStorageFormat.BONSAI);
   private final SnapSyncProcessState snapSyncState = mock(SnapSyncProcessState.class);
   private final SnapWorldDownloadState downloadState = mock(SnapWorldDownloadState.class);
-
   private final SnapSyncConfiguration snapSyncConfiguration = mock(SnapSyncConfiguration.class);
 
   private final PersistDataStep persistDataStep =
@@ -73,6 +79,59 @@ public void shouldSkipPersistDataWhenNoData() {
         .isEmpty();
   }
 
+  @Test
+  public void shouldHealWhenIncompleteStorageDataWithoutProof() {
+    // in normal operation this case would likely be triggered by a repivot.
+    // for simplicity, we limit the storage slots returned and return the partial
+    // range without a proof, implying the slots should be the complete storage range
+
+    final List<Task<SnapDataRequest>> tasks =
+        TaskGenerator.createAccountRequest(1, 1, 1, true, downloadState, false);
+    final List<Task<SnapDataRequest>> result = persistDataStep.persist(tasks);
+
+    assertThat(result).isSameAs(tasks);
+    // assert incomplete storage data without proof causes the account to be added to the heal list
+    verify(downloadState, times(1)).addAccountToHealingList(any(Bytes.class));
+    verify(downloadState, times(1)).enqueueRequest(any(AccountRangeDataRequest.class));
+    assertThat(worldStateStorage.getNodeData(Bytes.EMPTY, tasks.get(0).getData().getRootHash()))
+        .isEmpty();
+  }
+
+  @Test
+  @SuppressWarnings({"unchecked", "rawtypes"})
+  public void shouldEnqueueChildRequestWhenIncompleteStorageDataWithProof() {
+
+    final List<Task<SnapDataRequest>> tasks =
+        TaskGenerator.createAccountRequest(1, 1, 1, true, downloadState, true);
+    final List<Task<SnapDataRequest>> result = persistDataStep.persist(tasks);
+
+    assertThat(result).isSameAs(tasks);
+    // assert the current heal behavior of account marked, without new account request
+    verify(downloadState, times(1)).addAccountToHealingList(any(Bytes.class));
+    verify(downloadState, times(0)).enqueueRequest(any(AccountRangeDataRequest.class));
+
+    // assert that the incomplete storage data with proof enqueues a child storage request.
+    // verification is messy due to the stream generic parameter
+    ArgumentCaptor<Stream> rawArgumentCaptor = ArgumentCaptor.forClass(Stream.class);
+    verify(downloadState, atLeast(1)).enqueueRequests(rawArgumentCaptor.capture());
+    var enqueuedChildRequests =
+        rawArgumentCaptor.getAllValues().stream()
+            .map(stream -> stream.collect(Collectors.toList()))
+            .map(List.class::cast)
+            .filter(list -> !list.isEmpty())
+            .filter(list -> list.get(0) instanceof StorageRangeDataRequest)
+            .findFirst();
+    assertThat(enqueuedChildRequests).isPresent();
+    assertThat(enqueuedChildRequests.get()).isNotEmpty();
+
+    // assert the parent storage request is not persisted, child requests should complete first
+    assertThat(worldStateStorage.getNodeData(Bytes.EMPTY, tasks.get(0).getData().getRootHash()))
+        .isEmpty();
+  }
+
+  @Test
+  public void shouldHealWhenProofInvalid() {}
+
   private void assertDataPersisted(final List<Task<SnapDataRequest>> tasks) {
     tasks.forEach(
         task -> {

ethereum/eth/src/test/java/org/hyperledger/besu/ethereum/eth/sync/snapsync/TaskGenerator.java

@@ -42,17 +42,26 @@
 public class TaskGenerator {
 
   public static List<Task<SnapDataRequest>> createAccountRequest(final boolean withData) {
+    return createAccountRequest(1, 1, 100, withData, null, false);
+  }
 
+  static List<Task<SnapDataRequest>> createAccountRequest(
+      final int nbAccounts,
+      final int accountLimit,
+      final int storageLimit,
+      final boolean withData,
+      final SnapWorldDownloadState downloadState,
+      final boolean includeProof) {
     final WorldStateStorage worldStateStorage =
-        new InMemoryKeyValueStorageProvider().createWorldStateStorage(DataStorageFormat.FOREST);
+        new InMemoryKeyValueStorageProvider().createWorldStateStorage(DataStorageFormat.BONSAI);
 
     final WorldStateProofProvider worldStateProofProvider =
         new WorldStateProofProvider(worldStateStorage);
 
-    final MerkleTrie<Bytes, Bytes> trie = TrieGenerator.generateTrie(worldStateStorage, 1);
+    final MerkleTrie<Bytes, Bytes> trie = TrieGenerator.generateTrie(worldStateStorage, nbAccounts);
     final RangeStorageEntriesCollector collector =
         RangeStorageEntriesCollector.createCollector(
-            Bytes32.ZERO, RangeManager.MAX_RANGE, 1, Integer.MAX_VALUE);
+            Bytes32.ZERO, RangeManager.MAX_RANGE, accountLimit, Integer.MAX_VALUE);
     final TrieIterator<Bytes> visitor = RangeStorageEntriesCollector.createVisitor(collector);
     final TreeMap<Bytes32, Bytes> accounts =
         (TreeMap<Bytes32, Bytes>)
@@ -81,7 +90,10 @@ public static List<Task<SnapDataRequest>> createAccountRequest(final boolean wit
             rootHash,
             accountHash,
             stateTrieAccountValue.getStorageRoot(),
-            withData);
+            withData,
+            storageLimit,
+            downloadState,
+            includeProof);
     final BytecodeRequest bytecodeRequest =
         createBytecodeDataRequest(
             worldStateStorage,
@@ -102,11 +114,14 @@ private static StorageRangeDataRequest createStorageRangeDataRequest(
       final Hash rootHash,
       final Hash accountHash,
       final Bytes32 storageRoot,
-      final boolean withData) {
+      final boolean withData,
+      final int storageLimit,
+      final SnapWorldDownloadState downloadState,
+      final boolean includeProof) {
 
     final RangeStorageEntriesCollector collector =
         RangeStorageEntriesCollector.createCollector(
-            Bytes32.ZERO, RangeManager.MAX_RANGE, 100, Integer.MAX_VALUE);
+            Bytes32.ZERO, RangeManager.MAX_RANGE, storageLimit, Integer.MAX_VALUE);
     final StoredMerklePatriciaTrie<Bytes, Bytes> storageTrie =
         new StoredMerklePatriciaTrie<>(
             (location, hash) ->
@@ -128,7 +143,16 @@ private static StorageRangeDataRequest createStorageRangeDataRequest(
             rootHash, accountHash, storageRoot, RangeManager.MIN_RANGE, RangeManager.MAX_RANGE);
     if (withData) {
       request.setProofValid(true);
-      request.addResponse(null, worldStateProofProvider, slots, new ArrayDeque<>());
+      ArrayDeque<Bytes> proofNodes = new ArrayDeque<>();
+      if (includeProof) {
+        proofNodes.addAll(
+            worldStateProofProvider.getStorageProofRelatedNodes(
+                storageRoot, accountHash, RangeManager.MIN_RANGE));
+        proofNodes.addAll(
+            worldStateProofProvider.getStorageProofRelatedNodes(
+                storageRoot, accountHash, slots.lastKey()));
+      }
+      request.addResponse(downloadState, worldStateProofProvider, slots, proofNodes);
     }
     return request;
   }