add get proof for bonsai

ethereum/api/src/main/java/org/hyperledger/besu/ethereum/api/jsonrpc/internal/methods/EthGetProof.java

@@ -25,7 +25,8 @@
 import org.hyperledger.besu.ethereum.api.jsonrpc.internal.response.RpcErrorType;
 import org.hyperledger.besu.ethereum.api.jsonrpc.internal.results.proof.GetProofResult;
 import org.hyperledger.besu.ethereum.api.query.BlockchainQueries;
-import org.hyperledger.besu.ethereum.proof.WorldStateProof;
+import org.hyperledger.besu.ethereum.chain.Blockchain;
+import org.hyperledger.besu.ethereum.worldstate.WorldStateArchive;
 
 import java.util.Arrays;
 import java.util.List;
@@ -57,27 +58,30 @@ protected Object resultByBlockHash(
     final Address address = requestContext.getRequiredParameter(0, Address.class);
     final List<UInt256> storageKeys = getStorageKeys(requestContext);
 
-    return getBlockchainQueries()
-        .getAndMapWorldState(
-            blockHash,
-            worldState -> {
-              Optional<WorldStateProof> proofOptional =
-                  getBlockchainQueries()
-                      .getWorldStateArchive()
-                      .getAccountProof(worldState.rootHash(), address, storageKeys);
-              return proofOptional
-                  .map(
-                      proof ->
-                          (JsonRpcResponse)
-                              new JsonRpcSuccessResponse(
-                                  requestContext.getRequest().getId(),
-                                  GetProofResult.buildGetProofResult(address, proof)))
-                  .or(
-                      () ->
-                          Optional.of(
-                              new JsonRpcErrorResponse(
-                                  requestContext.getRequest().getId(),
-                                  RpcErrorType.NO_ACCOUNT_FOUND)));
+    final Blockchain blockchain = getBlockchainQueries().getBlockchain();
+    final WorldStateArchive worldStateArchive = getBlockchainQueries().getWorldStateArchive();
+    return blockchain
+        .getBlockHeader(blockHash)
+        .flatMap(
+            blockHeader -> {
+              return worldStateArchive.getAccountProof(
+                  blockHeader,
+                  address,
+                  storageKeys,
+                  maybeWorldStateProof ->
+                      maybeWorldStateProof
+                          .map(
+                              proof ->
+                                  (JsonRpcResponse)
+                                      new JsonRpcSuccessResponse(
+                                          requestContext.getRequest().getId(),
+                                          GetProofResult.buildGetProofResult(address, proof)))
+                          .or(
+                              () ->
+                                  Optional.of(
+                                      new JsonRpcErrorResponse(
+                                          requestContext.getRequest().getId(),
+                                          RpcErrorType.NO_ACCOUNT_FOUND))));
             })
         .orElse(
             new JsonRpcErrorResponse(

ethereum/api/src/test/java/org/hyperledger/besu/ethereum/api/jsonrpc/bonsai/EthJsonRpcHttpBySpecTest.java

@@ -28,7 +28,6 @@ public void setup() throws Exception {
   }
 
   public static Object[][] specs() {
-    return findSpecFiles(
-        new String[] {"eth"}, "getProof"); // getProof is not working with bonsai trie
+    return findSpecFiles(new String[] {"eth"});
   }
 }

ethereum/api/src/test/java/org/hyperledger/besu/ethereum/api/jsonrpc/internal/methods/EthGetProofTest.java

@@ -17,10 +17,9 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.hyperledger.besu.evm.account.Account.MAX_NONCE;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyList;
+import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
@@ -40,7 +39,6 @@
 import org.hyperledger.besu.ethereum.chain.Blockchain;
 import org.hyperledger.besu.ethereum.chain.ChainHead;
 import org.hyperledger.besu.ethereum.core.BlockHeader;
-import org.hyperledger.besu.ethereum.core.MutableWorldState;
 import org.hyperledger.besu.ethereum.proof.WorldStateProof;
 import org.hyperledger.besu.ethereum.worldstate.StateTrieAccountValue;
 import org.hyperledger.besu.ethereum.worldstate.WorldStateArchive;
@@ -91,7 +89,8 @@ public void setUp() {
     when(blockchain.getChainHead()).thenReturn(chainHead);
     when(chainHead.getBlockHeader()).thenReturn(blockHeader);
     when(blockHeader.getBlockHash()).thenReturn(Hash.ZERO);
-    when(blockchain.getBlockHeader(Hash.ZERO)).thenReturn(Optional.of(mock(BlockHeader.class)));
+    when(blockchainQueries.getBlockHashByNumber(blockNumber)).thenReturn(Optional.of(Hash.ZERO));
+    when(blockchain.getBlockHeader(Hash.ZERO)).thenReturn(Optional.of(blockHeader));
     method = spy(new EthGetProof(blockchainQueries));
   }
 
@@ -130,8 +129,7 @@ void errorWhenNoBlockNumberSupplied() {
   @Test
   void errorWhenAccountNotFound() {
     generateWorldState();
-    when(archive.getAccountProof(any(Hash.class), any(Address.class), any()))
-        .thenReturn(Optional.empty());
+
     final JsonRpcErrorResponse expectedResponse =
         new JsonRpcErrorResponse(null, RpcErrorType.NO_ACCOUNT_FOUND);
 
@@ -151,13 +149,12 @@ void errorWhenWorldStateUnavailable() {
 
     final JsonRpcErrorResponse expectedResponse =
         new JsonRpcErrorResponse(null, RpcErrorType.WORLD_STATE_UNAVAILABLE);
-    when(archive.getMutable(any(BlockHeader.class), anyBoolean())).thenReturn(Optional.empty());
 
     final JsonRpcRequestContext request =
         requestWithParams(
             Address.fromHexString("0x0000000000000000000000000000000000000000"),
             new String[] {storageKey.toString()},
-            String.valueOf(blockNumber));
+            String.valueOf(2));
 
     final JsonRpcErrorResponse response = (JsonRpcErrorResponse) method.response(request);
 
@@ -194,8 +191,6 @@ private GetProofResult generateWorldState() {
     final Hash codeHash =
         Hash.fromHexString("0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470");
     final long nonce = MAX_NONCE - 1;
-    final Hash rootHash =
-        Hash.fromHexString("0x56e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b431");
     final Hash storageRoot =
         Hash.fromHexString("0x56e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421");
 
@@ -222,19 +217,22 @@ private GetProofResult generateWorldState() {
                     "0x2222222222222222222222222222222222222222222222222222222222222222")));
     when(worldStateProof.getStorageValue(storageKey)).thenReturn(UInt256.ZERO);
 
-    when(archive.getAccountProof(eq(rootHash), eq(address), anyList()))
-        .thenReturn(Optional.of(worldStateProof));
-
-    final MutableWorldState mutableWorldState = mock(MutableWorldState.class);
-    when(mutableWorldState.rootHash()).thenReturn(rootHash);
-    doAnswer(
-            invocation ->
-                invocation
-                    .<Function<MutableWorldState, Optional<? extends JsonRpcResponse>>>getArgument(
-                        1)
-                    .apply(mutableWorldState))
-        .when(blockchainQueries)
-        .getAndMapWorldState(any(), any());
+    when(archive.getAccountProof(eq(blockHeader), eq(address), anyList(), any()))
+        .thenAnswer(
+            invocation -> {
+              Function<Optional<WorldStateProof>, Optional<JsonRpcResponse>> realMapper =
+                  invocation.getArgument(3);
+              return realMapper.apply(Optional.of(worldStateProof));
+            });
+
+    when(archive.getAccountProof(
+            eq(blockHeader), argThat(arg -> !arg.equals(address)), anyList(), any()))
+        .thenAnswer(
+            invocation -> {
+              Function<Optional<WorldStateProof>, Optional<JsonRpcResponse>> realMapper =
+                  invocation.getArgument(3);
+              return realMapper.apply(Optional.empty());
+            });
 
     return GetProofResult.buildGetProofResult(address, worldStateProof);
   }

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

@@ -30,6 +30,7 @@
 import org.hyperledger.besu.ethereum.core.BlockHeader;
 import org.hyperledger.besu.ethereum.core.MutableWorldState;
 import org.hyperledger.besu.ethereum.proof.WorldStateProof;
+import org.hyperledger.besu.ethereum.proof.WorldStateProofProvider;
 import org.hyperledger.besu.ethereum.rlp.RLP;
 import org.hyperledger.besu.ethereum.storage.StorageProvider;
 import org.hyperledger.besu.ethereum.trie.MerkleTrieException;
@@ -64,7 +65,6 @@ public class BonsaiWorldStateProvider implements WorldStateArchive {
   private final TrieLogManager trieLogManager;
   private final BonsaiWorldState persistedState;
   private final BonsaiWorldStateKeyValueStorage worldStateStorage;
-
   private final CachedMerkleTrieLoader cachedMerkleTrieLoader;
 
   public BonsaiWorldStateProvider(
@@ -363,16 +363,27 @@ public void resetArchiveStateTo(final BlockHeader blockHeader) {
   }
 
   @Override
-  public Optional<Bytes> getNodeData(final Hash hash) {
+  public <U> Optional<U> getAccountProof(
+      final BlockHeader blockHeader,
+      final Address accountAddress,
+      final List<UInt256> accountStorageKeys,
+      final Function<Optional<WorldStateProof>, ? extends Optional<U>> mapper) {
+    try (BonsaiWorldState ws = (BonsaiWorldState) getMutable(blockHeader, false).orElse(null)) {
+      if (ws != null) {
+        final WorldStateProofProvider worldStateProofProvider =
+            new WorldStateProofProvider(ws.getWorldStateStorage());
+        return mapper.apply(
+            worldStateProofProvider.getAccountProof(
+                ws.getWorldStateRootHash(), accountAddress, accountStorageKeys));
+      }
+    } catch (Exception ex) {
+      LOG.error("failed proof query for " + blockHeader.getBlockHash().toShortHexString(), ex);
+    }
     return Optional.empty();
   }
 
   @Override
-  public Optional<WorldStateProof> getAccountProof(
-      final Hash worldStateRoot,
-      final Address accountAddress,
-      final List<UInt256> accountStorageKeys) {
-    // FIXME we can do proofs for layered tries and the persisted trie
+  public Optional<Bytes> getNodeData(final Hash hash) {
     return Optional.empty();
   }
 

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

@@ -26,6 +26,7 @@
 
 import java.util.List;
 import java.util.Optional;
+import java.util.function.Function;
 
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.units.bigints.UInt256;
@@ -89,11 +90,14 @@ public WorldStateStorage getWorldStateStorage() {
   }
 
   @Override
-  public Optional<WorldStateProof> getAccountProof(
-      final Hash worldStateRoot,
+  public <U> Optional<U> getAccountProof(
+      final BlockHeader blockHeader,
       final Address accountAddress,
-      final List<UInt256> accountStorageKeys) {
-    return worldStateProof.getAccountProof(worldStateRoot, accountAddress, accountStorageKeys);
+      final List<UInt256> accountStorageKeys,
+      final Function<Optional<WorldStateProof>, ? extends Optional<U>> mapper) {
+    return mapper.apply(
+        worldStateProof.getAccountProof(
+            blockHeader.getStateRoot(), accountAddress, accountStorageKeys));
   }
 
   @Override

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

@@ -25,6 +25,7 @@
 import java.io.Closeable;
 import java.util.List;
 import java.util.Optional;
+import java.util.function.Function;
 
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.units.bigints.UInt256;
@@ -51,6 +52,19 @@ public interface WorldStateArchive extends Closeable {
 
   Optional<Bytes> getNodeData(Hash hash);
 
-  Optional<WorldStateProof> getAccountProof(
-      Hash worldStateRoot, Address accountAddress, List<UInt256> accountStorageKeys);
+  /**
+   * Retrieves an account proof based on the provided parameters.
+   *
+   * @param blockHeader The header of the block for which to retrieve the account proof.
+   * @param accountAddress The address of the account for which to retrieve the proof.
+   * @param accountStorageKeys The storage keys of the account for which to retrieve the proof.
+   * @param mapper A function to map the retrieved WorldStateProof to a desired type.
+   * @return An Optional containing the mapped result if the account proof is successfully retrieved
+   *     and mapped, or an empty Optional otherwise.
+   */
+  <U> Optional<U> getAccountProof(
+      final BlockHeader blockHeader,
+      final Address accountAddress,
+      final List<UInt256> accountStorageKeys,
+      final Function<Optional<WorldStateProof>, ? extends Optional<U>> mapper);
 }