Fix Snap Server Account Range tests (#7552)

Signed-off-by: Gabriel-Trintinalia <[email protected]>
hyperledger · Sep 18, 2024 · 578104e · 578104e
1 parent 33085dd
commit 578104e
Show file tree

Hide file tree

Showing 10 changed files with 772 additions and 15 deletions.
diff --git a/...um/core/src/test-support/java/org/hyperledger/besu/ethereum/core/BlockchainSetupUtil.java b/...um/core/src/test-support/java/org/hyperledger/besu/ethereum/core/BlockchainSetupUtil.java
@@ -85,6 +85,13 @@ public Blockchain importAllBlocks() {
     return blockchain;
   }
 
+  public Blockchain importAllBlocks(
+      final HeaderValidationMode headerValidationMode,
+      final HeaderValidationMode ommerValidationMode) {
+    importBlocks(blocks, headerValidationMode, ommerValidationMode);
+    return blockchain;
+  }
+
   public void importFirstBlocks(final int count) {
     importBlocks(blocks.subList(0, count));
   }
@@ -126,6 +133,10 @@ public static BlockchainSetupUtil forUpgradedFork() {
     return createForEthashChain(BlockTestUtil.getUpgradedForkResources(), DataStorageFormat.FOREST);
   }
 
+  public static BlockchainSetupUtil forSnapTesting(final DataStorageFormat storageFormat) {
+    return createForEthashChain(BlockTestUtil.getSnapTestChainResources(), storageFormat);
+  }
+
   public static BlockchainSetupUtil createForEthashChain(
       final ChainResources chainResources, final DataStorageFormat storageFormat) {
     return create(
@@ -241,14 +252,22 @@ public TransactionPool getTransactionPool() {
   }
 
   private void importBlocks(final List<Block> blocks) {
+    importBlocks(blocks, HeaderValidationMode.FULL, HeaderValidationMode.FULL);
+  }
+
+  private void importBlocks(
+      final List<Block> blocks,
+      final HeaderValidationMode headerValidationMode,
+      final HeaderValidationMode ommerValidationMode) {
     for (final Block block : blocks) {
       if (block.getHeader().getNumber() == BlockHeader.GENESIS_BLOCK_NUMBER) {
         continue;
       }
       final ProtocolSpec protocolSpec = protocolSchedule.getByBlockHeader(block.getHeader());
       final BlockImporter blockImporter = protocolSpec.getBlockImporter();
       final BlockImportResult result =
-          blockImporter.importBlock(protocolContext, block, HeaderValidationMode.FULL);
+          blockImporter.importBlock(
+              protocolContext, block, headerValidationMode, ommerValidationMode);
       if (!result.isImported()) {
         throw new IllegalStateException("Unable to import block " + block.getHeader().getNumber());
       }

diff --git a/ethereum/eth/src/main/java/org/hyperledger/besu/ethereum/eth/manager/snap/SnapServer.java b/ethereum/eth/src/main/java/org/hyperledger/besu/ethereum/eth/manager/snap/SnapServer.java
@@ -31,6 +31,7 @@
 import org.hyperledger.besu.ethereum.p2p.rlpx.wire.MessageData;
 import org.hyperledger.besu.ethereum.proof.WorldStateProofProvider;
 import org.hyperledger.besu.ethereum.rlp.BytesValueRLPOutput;
+import org.hyperledger.besu.ethereum.rlp.RLP;
 import org.hyperledger.besu.ethereum.trie.CompactEncoding;
 import org.hyperledger.besu.ethereum.trie.MerkleTrie;
 import org.hyperledger.besu.ethereum.trie.diffbased.bonsai.BonsaiWorldStateProvider;
@@ -240,12 +241,9 @@ MessageData constructGetAccountRangeResponse(final MessageData message) {
                         stopWatch,
                         maxResponseBytes,
                         (pair) -> {
-                          var rlpOutput = new BytesValueRLPOutput();
-                          rlpOutput.startList();
-                          rlpOutput.writeBytes(pair.getFirst());
-                          rlpOutput.writeRLPBytes(pair.getSecond());
-                          rlpOutput.endList();
-                          return rlpOutput.encodedSize();
+                          Bytes bytes =
+                              AccountRangeMessage.toSlimAccount(RLP.input(pair.getSecond()));
+                          return Hash.SIZE + bytes.size();
                         });
 
                 final Bytes32 endKeyBytes = range.endKeyHash();
@@ -257,11 +255,15 @@ MessageData constructGetAccountRangeResponse(final MessageData message) {
                     storage.streamFlatAccounts(range.startKeyHash(), shouldContinuePredicate);
 
                 if (accounts.isEmpty() && shouldContinuePredicate.shouldContinue.get()) {
+                  var fromNextHash =
+                      range.endKeyHash().compareTo(range.startKeyHash()) >= 0
+                          ? range.endKeyHash()
+                          : range.startKeyHash();
                   // fetch next account after range, if it exists
                   LOGGER.debug(
                       "found no accounts in range, taking first value starting from {}",
-                      asLogHash(range.endKeyHash()));
-                  accounts = storage.streamFlatAccounts(range.endKeyHash(), UInt256.MAX_VALUE, 1L);
+                      asLogHash(fromNextHash));
+                  accounts = storage.streamFlatAccounts(fromNextHash, UInt256.MAX_VALUE, 1L);
                 }
 
                 final var worldStateProof =

diff --git a/...th/src/main/java/org/hyperledger/besu/ethereum/eth/messages/snap/AccountRangeMessage.java b/...th/src/main/java/org/hyperledger/besu/ethereum/eth/messages/snap/AccountRangeMessage.java
@@ -14,6 +14,7 @@
  */
 package org.hyperledger.besu.ethereum.eth.messages.snap;
 
+import org.hyperledger.besu.datatypes.Hash;
 import org.hyperledger.besu.ethereum.p2p.rlpx.wire.AbstractSnapMessageData;
 import org.hyperledger.besu.ethereum.p2p.rlpx.wire.MessageData;
 import org.hyperledger.besu.ethereum.rlp.BytesValueRLPInput;
@@ -28,6 +29,7 @@
 import java.util.Optional;
 import java.util.TreeMap;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Maps;
 import kotlin.collections.ArrayDeque;
 import org.apache.tuweni.bytes.Bytes;
@@ -117,7 +119,8 @@ public AccountRangeData accountData(final boolean withRequestId) {
     return ImmutableAccountRangeData.builder().accounts(accounts).proofs(proofs).build();
   }
 
-  private Bytes toFullAccount(final RLPInput rlpInput) {
+  @VisibleForTesting
+  public static Bytes toFullAccount(final RLPInput rlpInput) {
     final StateTrieAccountValue accountValue = StateTrieAccountValue.readFrom(rlpInput);
 
     final BytesValueRLPOutput rlpOutput = new BytesValueRLPOutput();
@@ -131,6 +134,26 @@ private Bytes toFullAccount(final RLPInput rlpInput) {
     return rlpOutput.encoded();
   }
 
+  public static Bytes toSlimAccount(final RLPInput rlpInput) {
+    StateTrieAccountValue accountValue = StateTrieAccountValue.readFrom(rlpInput);
+    var rlpOutput = new BytesValueRLPOutput();
+    rlpOutput.startList();
+    rlpOutput.writeLongScalar(accountValue.getNonce());
+    rlpOutput.writeUInt256Scalar(accountValue.getBalance());
+    if (accountValue.getStorageRoot().equals(Hash.EMPTY_TRIE_HASH)) {
+      rlpOutput.writeNull();
+    } else {
+      rlpOutput.writeBytes(accountValue.getStorageRoot());
+    }
+    if (accountValue.getCodeHash().equals(Hash.EMPTY)) {
+      rlpOutput.writeNull();
+    } else {
+      rlpOutput.writeBytes(accountValue.getCodeHash());
+    }
+    rlpOutput.endList();
+    return rlpOutput.encoded();
+  }
+
   @Value.Immutable
   public interface AccountRangeData {
 

diff --git a/...src/main/java/org/hyperledger/besu/ethereum/eth/messages/snap/GetAccountRangeMessage.java b/...src/main/java/org/hyperledger/besu/ethereum/eth/messages/snap/GetAccountRangeMessage.java
@@ -24,6 +24,7 @@
 import java.math.BigInteger;
 import java.util.Optional;
 
+import com.google.common.annotations.VisibleForTesting;
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes32;
 import org.immutables.value.Value;
@@ -48,12 +49,21 @@ public static GetAccountRangeMessage readFrom(final MessageData message) {
 
   public static GetAccountRangeMessage create(
       final Hash worldStateRootHash, final Bytes32 startKeyHash, final Bytes32 endKeyHash) {
+    return create(worldStateRootHash, startKeyHash, endKeyHash, SIZE_REQUEST);
+  }
+
+  @VisibleForTesting
+  public static GetAccountRangeMessage create(
+      final Hash worldStateRootHash,
+      final Bytes32 startKeyHash,
+      final Bytes32 endKeyHash,
+      final BigInteger sizeRequest) {
     final BytesValueRLPOutput tmp = new BytesValueRLPOutput();
     tmp.startList();
     tmp.writeBytes(worldStateRootHash);
     tmp.writeBytes(startKeyHash);
     tmp.writeBytes(endKeyHash);
-    tmp.writeBigIntegerScalar(SIZE_REQUEST);
+    tmp.writeBigIntegerScalar(sizeRequest);
     tmp.endList();
     return new GetAccountRangeMessage(tmp.encoded());
   }