From 67637fa3f0a23671394d5a552376944088007ccc Mon Sep 17 00:00:00 2001
From: Jason Frame <jason.frame@consensys.net>
Date: Thu, 27 Jun 2024 09:35:31 +1000
Subject: [PATCH] Snapserver GetTrieNodes request to handle short hash for
 storage (#7264)

Signed-off-by: Jason Frame <jason.frame@consensys.net>
---
 .../ethereum/eth/manager/snap/SnapServer.java |  2 +-
 .../eth/manager/snap/SnapServerTest.java      | 25 ++++++++++++++++++-
 2 files changed, 25 insertions(+), 2 deletions(-)

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
index 8d0add16f98..49f725d825f 100644
--- 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
@@ -530,7 +530,7 @@ MessageData constructGetTrieNodesResponse(final MessageData message) {
                     // otherwise the first element should be account hash, and subsequent paths
                     // are compact encoded account storage paths
 
-                    final Bytes accountPrefix = triePath.get(0);
+                    final Bytes accountPrefix = Bytes32.leftPad(triePath.getFirst());
 
                     List<Bytes> storagePaths = triePath.subList(1, triePath.size());
                     for (var path : storagePaths) {
diff --git a/ethereum/eth/src/test/java/org/hyperledger/besu/ethereum/eth/manager/snap/SnapServerTest.java b/ethereum/eth/src/test/java/org/hyperledger/besu/ethereum/eth/manager/snap/SnapServerTest.java
index 1b4cc0618cf..2103af41242 100644
--- a/ethereum/eth/src/test/java/org/hyperledger/besu/ethereum/eth/manager/snap/SnapServerTest.java
+++ b/ethereum/eth/src/test/java/org/hyperledger/besu/ethereum/eth/manager/snap/SnapServerTest.java
@@ -534,6 +534,24 @@ public void assertStorageTriePathRequest() {
     assertThat(trieNodes.size()).isEqualTo(4);
   }
 
+  @Test
+  public void assertStorageTrieShortAccountHashPathRequest() {
+    Bytes accountShortHash = Bytes.fromHexStringLenient("0x40");
+    Hash accountFullHash = Hash.wrap(Bytes32.leftPad(accountShortHash));
+    SnapTestAccount testAccount = createTestContractAccount(accountFullHash, inMemoryStorage);
+    insertTestAccounts(testAccount);
+    var pathToSlot11 = CompactEncoding.encode(Bytes.fromHexStringLenient("0x0101"));
+    var pathToSlot12 = CompactEncoding.encode(Bytes.fromHexStringLenient("0x0102"));
+    var trieNodeRequest =
+        requestTrieNodes(
+            storageTrie.getRootHash(),
+            List.of(List.of(accountShortHash, pathToSlot11, pathToSlot12)));
+    assertThat(trieNodeRequest).isNotNull();
+    List<Bytes> trieNodes = trieNodeRequest.nodes(false);
+    assertThat(trieNodes).isNotNull();
+    assertThat(trieNodes.size()).isEqualTo(2);
+  }
+
   @Test
   public void assertStorageTrieLimitRequest() {
     insertTestAccounts(acct1, acct2, acct3, acct4);
@@ -671,7 +689,12 @@ static SnapTestAccount createTestAccount(final String hexAddr) {
 
   static SnapTestAccount createTestContractAccount(
       final String hexAddr, final BonsaiWorldStateKeyValueStorage storage) {
-    Hash acctHash = Hash.wrap(Bytes32.rightPad(Bytes.fromHexString(hexAddr)));
+    final Hash acctHash = Hash.wrap(Bytes32.rightPad(Bytes.fromHexString(hexAddr)));
+    return createTestContractAccount(acctHash, storage);
+  }
+
+  static SnapTestAccount createTestContractAccount(
+      final Hash acctHash, final BonsaiWorldStateKeyValueStorage storage) {
     MerkleTrie<Bytes32, Bytes> trie =
         new StoredMerklePatriciaTrie<>(
             (loc, hash) -> storage.getAccountStorageTrieNode(acctHash, loc, hash),