fix trie by asking peers

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/BonsaiInMemoryWorldStateKeyValueStorage.java

@@ -14,10 +14,13 @@
  */
 package org.hyperledger.besu.ethereum.bonsai;
 
+import org.hyperledger.besu.ethereum.worldstate.PeerTrieNodeFinder;
 import org.hyperledger.besu.ethereum.worldstate.WorldStateStorage;
 import org.hyperledger.besu.plugin.services.storage.KeyValueStorage;
 import org.hyperledger.besu.plugin.services.storage.KeyValueStorageTransaction;
 
+import java.util.Optional;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -32,8 +35,15 @@ public BonsaiInMemoryWorldStateKeyValueStorage(
       final KeyValueStorage codeStorage,
       final KeyValueStorage storageStorage,
       final KeyValueStorage trieBranchStorage,
-      final KeyValueStorage trieLogStorage) {
-    super(accountStorage, codeStorage, storageStorage, trieBranchStorage, trieLogStorage);
+      final KeyValueStorage trieLogStorage,
+      final Optional<PeerTrieNodeFinder> fallbackNodeFinder) {
+    super(
+        accountStorage,
+        codeStorage,
+        storageStorage,
+        trieBranchStorage,
+        trieLogStorage,
+        fallbackNodeFinder);
   }
 
   @Override

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/BonsaiLayeredWorldState.java

@@ -259,14 +259,15 @@ public Stream<StreamableAccount> streamAccounts(final Bytes32 startKeyHash, fina
   public MutableWorldState copy() {
     final BonsaiPersistedWorldState bonsaiPersistedWorldState =
         ((BonsaiPersistedWorldState) archive.getMutable());
-    return new BonsaiInMemoryWorldState(
-        archive,
+    BonsaiInMemoryWorldStateKeyValueStorage bonsaiInMemoryWorldStateKeyValueStorage =
         new BonsaiInMemoryWorldStateKeyValueStorage(
             bonsaiPersistedWorldState.getWorldStateStorage().accountStorage,
             bonsaiPersistedWorldState.getWorldStateStorage().codeStorage,
             bonsaiPersistedWorldState.getWorldStateStorage().storageStorage,
             bonsaiPersistedWorldState.getWorldStateStorage().trieBranchStorage,
-            bonsaiPersistedWorldState.getWorldStateStorage().trieLogStorage));
+            bonsaiPersistedWorldState.getWorldStateStorage().trieLogStorage,
+            bonsaiPersistedWorldState.getWorldStateStorage().getMaybeFallbackNodeFinder());
+    return new BonsaiInMemoryWorldState(archive, bonsaiInMemoryWorldStateKeyValueStorage);
   }
 
   @Override

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/BonsaiPersistedWorldState.java

@@ -75,14 +75,16 @@ public BonsaiWorldStateArchive getArchive() {
 
   @Override
   public MutableWorldState copy() {
-    return new BonsaiInMemoryWorldState(
-        archive,
+    BonsaiInMemoryWorldStateKeyValueStorage bonsaiInMemoryWorldStateKeyValueStorage =
         new BonsaiInMemoryWorldStateKeyValueStorage(
             worldStateStorage.accountStorage,
             worldStateStorage.codeStorage,
             worldStateStorage.storageStorage,
             worldStateStorage.trieBranchStorage,
-            worldStateStorage.trieLogStorage));
+            worldStateStorage.trieLogStorage,
+            getWorldStateStorage().getMaybeFallbackNodeFinder());
+
+    return new BonsaiInMemoryWorldState(archive, bonsaiInMemoryWorldStateKeyValueStorage);
   }
 
   @Override

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/BonsaiWorldStateArchive.java

@@ -16,6 +16,7 @@
 
 package org.hyperledger.besu.ethereum.bonsai;
 
+import static com.google.common.base.Preconditions.checkNotNull;
 import static org.hyperledger.besu.datatypes.Hash.fromPlugin;
 
 import org.hyperledger.besu.datatypes.Address;
@@ -26,6 +27,7 @@
 import org.hyperledger.besu.ethereum.core.MutableWorldState;
 import org.hyperledger.besu.ethereum.proof.WorldStateProof;
 import org.hyperledger.besu.ethereum.storage.StorageProvider;
+import org.hyperledger.besu.ethereum.worldstate.PeerTrieNodeFinder;
 import org.hyperledger.besu.ethereum.worldstate.WorldStateArchive;
 import org.hyperledger.besu.evm.worldstate.WorldState;
 
@@ -243,4 +245,9 @@ public Optional<WorldStateProof> getAccountProof(
     // FIXME we can do proofs for layered tries and the persisted trie
     return Optional.empty();
   }
+
+  public void useFallbackNodeFinder(final Optional<PeerTrieNodeFinder> fallbackNodeFinder) {
+    checkNotNull(fallbackNodeFinder);
+    worldStateStorage.useFallbackNodeFinder(fallbackNodeFinder);
+  }
 }

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/BonsaiWorldStateKeyValueStorage.java

@@ -14,12 +14,15 @@
  */
 package org.hyperledger.besu.ethereum.bonsai;
 
+import static com.google.common.base.Preconditions.checkNotNull;
+
 import org.hyperledger.besu.datatypes.Hash;
 import org.hyperledger.besu.ethereum.storage.StorageProvider;
 import org.hyperledger.besu.ethereum.storage.keyvalue.KeyValueSegmentIdentifier;
 import org.hyperledger.besu.ethereum.trie.MerklePatriciaTrie;
 import org.hyperledger.besu.ethereum.trie.StoredMerklePatriciaTrie;
 import org.hyperledger.besu.ethereum.trie.StoredNodeFactory;
+import org.hyperledger.besu.ethereum.worldstate.PeerTrieNodeFinder;
 import org.hyperledger.besu.ethereum.worldstate.StateTrieAccountValue;
 import org.hyperledger.besu.ethereum.worldstate.WorldStateStorage;
 import org.hyperledger.besu.plugin.services.storage.KeyValueStorage;
@@ -46,16 +49,16 @@ public class BonsaiWorldStateKeyValueStorage implements WorldStateStorage {
   protected final KeyValueStorage trieBranchStorage;
   protected final KeyValueStorage trieLogStorage;
 
+  private Optional<PeerTrieNodeFinder> maybeFallbackNodeFinder;
+
   public BonsaiWorldStateKeyValueStorage(final StorageProvider provider) {
-    accountStorage =
-        provider.getStorageBySegmentIdentifier(KeyValueSegmentIdentifier.ACCOUNT_INFO_STATE);
-    codeStorage = provider.getStorageBySegmentIdentifier(KeyValueSegmentIdentifier.CODE_STORAGE);
-    storageStorage =
-        provider.getStorageBySegmentIdentifier(KeyValueSegmentIdentifier.ACCOUNT_STORAGE_STORAGE);
-    trieBranchStorage =
-        provider.getStorageBySegmentIdentifier(KeyValueSegmentIdentifier.TRIE_BRANCH_STORAGE);
-    trieLogStorage =
-        provider.getStorageBySegmentIdentifier(KeyValueSegmentIdentifier.TRIE_LOG_STORAGE);
+    this(
+        provider.getStorageBySegmentIdentifier(KeyValueSegmentIdentifier.ACCOUNT_INFO_STATE),
+        provider.getStorageBySegmentIdentifier(KeyValueSegmentIdentifier.CODE_STORAGE),
+        provider.getStorageBySegmentIdentifier(KeyValueSegmentIdentifier.ACCOUNT_STORAGE_STORAGE),
+        provider.getStorageBySegmentIdentifier(KeyValueSegmentIdentifier.TRIE_BRANCH_STORAGE),
+        provider.getStorageBySegmentIdentifier(KeyValueSegmentIdentifier.TRIE_LOG_STORAGE),
+        Optional.empty());
   }
 
   public BonsaiWorldStateKeyValueStorage(
@@ -64,11 +67,28 @@ public BonsaiWorldStateKeyValueStorage(
       final KeyValueStorage storageStorage,
       final KeyValueStorage trieBranchStorage,
       final KeyValueStorage trieLogStorage) {
+    this(
+        accountStorage,
+        codeStorage,
+        storageStorage,
+        trieBranchStorage,
+        trieLogStorage,
+        Optional.empty());
+  }
+
+  public BonsaiWorldStateKeyValueStorage(
+      final KeyValueStorage accountStorage,
+      final KeyValueStorage codeStorage,
+      final KeyValueStorage storageStorage,
+      final KeyValueStorage trieBranchStorage,
+      final KeyValueStorage trieLogStorage,
+      final Optional<PeerTrieNodeFinder> fallbackNodeFinder) {
     this.accountStorage = accountStorage;
     this.codeStorage = codeStorage;
     this.storageStorage = storageStorage;
     this.trieBranchStorage = trieBranchStorage;
     this.trieLogStorage = trieLogStorage;
+    this.maybeFallbackNodeFinder = fallbackNodeFinder;
   }
 
   @Override
@@ -104,7 +124,17 @@ public Optional<Bytes> getAccountStateTrieNode(final Bytes location, final Bytes
     if (nodeHash.equals(MerklePatriciaTrie.EMPTY_TRIE_NODE_HASH)) {
       return Optional.of(MerklePatriciaTrie.EMPTY_TRIE_NODE);
     } else {
-      return trieBranchStorage.get(location.toArrayUnsafe()).map(Bytes::wrap);
+      final Optional<Bytes> value =
+          trieBranchStorage.get(location.toArrayUnsafe()).map(Bytes::wrap);
+      if (value.isPresent()) {
+        return value
+            .filter(b -> Hash.hash(b).equals(nodeHash))
+            .or(
+                () ->
+                    maybeFallbackNodeFinder.flatMap(
+                        finder -> finder.getAccountStateTrieNode(location, nodeHash)));
+      }
+      return Optional.empty();
     }
   }
 
@@ -114,9 +144,20 @@ public Optional<Bytes> getAccountStorageTrieNode(
     if (nodeHash.equals(MerklePatriciaTrie.EMPTY_TRIE_NODE_HASH)) {
       return Optional.of(MerklePatriciaTrie.EMPTY_TRIE_NODE);
     } else {
-      return trieBranchStorage
-          .get(Bytes.concatenate(accountHash, location).toArrayUnsafe())
-          .map(Bytes::wrap);
+      final Optional<Bytes> value =
+          trieBranchStorage
+              .get(Bytes.concatenate(accountHash, location).toArrayUnsafe())
+              .map(Bytes::wrap);
+      if (value.isPresent()) {
+        return value
+            .filter(b -> Hash.hash(b).equals(nodeHash))
+            .or(
+                () ->
+                    maybeFallbackNodeFinder.flatMap(
+                        finder ->
+                            finder.getAccountStorageTrieNode(accountHash, location, nodeHash)));
+      }
+      return Optional.empty();
     }
   }
 
@@ -218,6 +259,15 @@ public void removeNodeAddedListener(final long id) {
     throw new RuntimeException("removeNodeAddedListener not available");
   }
 
+  public Optional<PeerTrieNodeFinder> getMaybeFallbackNodeFinder() {
+    return maybeFallbackNodeFinder;
+  }
+
+  public void useFallbackNodeFinder(final Optional<PeerTrieNodeFinder> maybeFallbackNodeFinder) {
+    checkNotNull(maybeFallbackNodeFinder);
+    this.maybeFallbackNodeFinder = maybeFallbackNodeFinder;
+  }
+
   public static class Updater implements WorldStateStorage.Updater {
 
     private final KeyValueStorageTransaction accountStorageTransaction;

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/worldstate/PeerTrieNodeFinder.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright contributors to Hyperledger Besu
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.hyperledger.besu.ethereum.worldstate;
+
+import org.hyperledger.besu.datatypes.Hash;
+
+import java.util.Optional;
+
+import org.apache.tuweni.bytes.Bytes;
+import org.apache.tuweni.bytes.Bytes32;
+
+public interface PeerTrieNodeFinder {
+
+  Optional<Bytes> getAccountStateTrieNode(final Bytes location, final Bytes32 nodeHash);
+
+  Optional<Bytes> getAccountStorageTrieNode(Hash accountHash, Bytes location, Bytes32 nodeHash);
+}

ethereum/core/src/test/java/org/hyperledger/besu/ethereum/bonsai/BonsaiWorldStateKeyValueStorageTest.java

@@ -18,9 +18,11 @@
 import static org.hyperledger.besu.ethereum.bonsai.BonsaiWorldStateKeyValueStorage.WORLD_ROOT_HASH_KEY;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import org.hyperledger.besu.datatypes.Address;
 import org.hyperledger.besu.datatypes.Hash;
@@ -30,8 +32,10 @@
 import org.hyperledger.besu.ethereum.trie.MerklePatriciaTrie;
 import org.hyperledger.besu.ethereum.trie.StorageEntriesCollector;
 import org.hyperledger.besu.ethereum.trie.StoredMerklePatriciaTrie;
+import org.hyperledger.besu.ethereum.worldstate.PeerTrieNodeFinder;
 import org.hyperledger.besu.ethereum.worldstate.StateTrieAccountValue;
 
+import java.util.Optional;
 import java.util.TreeMap;
 
 import org.apache.tuweni.bytes.Bytes;
@@ -291,6 +295,62 @@ public void isWorldStateAvailable_afterCallingSaveWorldstate() {
     assertThat(storage.isWorldStateAvailable(Bytes32.wrap(nodeHashKey), Hash.EMPTY)).isTrue();
   }
 
+  @Test
+  public void getAccountStateTrieNode_callFallbackMechanismForInvalidNode() {
+
+    PeerTrieNodeFinder peerTrieNodeFinder = mock(PeerTrieNodeFinder.class);
+
+    final Bytes location = Bytes.fromHexString("0x01");
+    final Bytes bytesInDB = Bytes.fromHexString("0x123456");
+
+    final Hash hashToFind = Hash.hash(Bytes.of(1));
+    final Bytes bytesToFind = Bytes.fromHexString("0x123457");
+
+    final BonsaiWorldStateKeyValueStorage storage = emptyStorage();
+
+    when(peerTrieNodeFinder.getAccountStateTrieNode(location, hashToFind))
+        .thenReturn(Optional.of(bytesToFind));
+    storage.useFallbackNodeFinder(Optional.of(peerTrieNodeFinder));
+
+    storage.updater().putAccountStateTrieNode(location, Hash.hash(bytesInDB), bytesInDB).commit();
+
+    Optional<Bytes> accountStateTrieNodeResult =
+        storage.getAccountStateTrieNode(location, hashToFind);
+
+    verify(peerTrieNodeFinder).getAccountStateTrieNode(location, hashToFind);
+    assertThat(accountStateTrieNodeResult).contains(bytesToFind);
+  }
+
+  @Test
+  public void getAccountStorageTrieNode_callFallbackMechanismForInvalidNode() {
+
+    PeerTrieNodeFinder peerTrieNodeFinder = mock(PeerTrieNodeFinder.class);
+
+    final Hash account = Hash.hash(Bytes32.ZERO);
+    final Bytes location = Bytes.fromHexString("0x01");
+    final Bytes bytesInDB = Bytes.fromHexString("0x123456");
+
+    final Hash hashToFind = Hash.hash(Bytes.of(1));
+    final Bytes bytesToFind = Bytes.fromHexString("0x123457");
+
+    final BonsaiWorldStateKeyValueStorage storage = emptyStorage();
+
+    when(peerTrieNodeFinder.getAccountStorageTrieNode(account, location, hashToFind))
+        .thenReturn(Optional.of(bytesToFind));
+    storage.useFallbackNodeFinder(Optional.of(peerTrieNodeFinder));
+
+    storage
+        .updater()
+        .putAccountStorageTrieNode(account, location, Hash.hash(bytesInDB), bytesInDB)
+        .commit();
+
+    Optional<Bytes> accountStateTrieNodeResult =
+        storage.getAccountStorageTrieNode(account, location, hashToFind);
+
+    verify(peerTrieNodeFinder).getAccountStorageTrieNode(account, location, hashToFind);
+    assertThat(accountStateTrieNodeResult).contains(bytesToFind);
+  }
+
   private BonsaiWorldStateKeyValueStorage emptyStorage() {
     return new BonsaiWorldStateKeyValueStorage(new InMemoryKeyValueStorageProvider());
   }