Fix graphql storageAt for bonsai (#5548)

* Fix code and storage retrieve for bonsai getting a new world state snapshot to fetch code/storage from in the AccountAdapter

Signed-off-by: Gabriel Fukushima <gabrielfukushima@gmail.com>

---------

Signed-off-by: Gabriel Fukushima <gabrielfukushima@gmail.com>

ethereum/api/src/main/java/org/hyperledger/besu/ethereum/api/graphql/internal/pojoadapter/AccountAdapter.java

@@ -16,6 +16,8 @@ package org.hyperledger.besu.ethereum.api.graphql.internal.pojoadapter;
 
 import org.hyperledger.besu.datatypes.Address;
 import org.hyperledger.besu.datatypes.Wei;
+import org.hyperledger.besu.ethereum.api.query.BlockchainQueries;
+import org.hyperledger.besu.ethereum.bonsai.BonsaiAccount;
 import org.hyperledger.besu.evm.account.Account;
 import org.hyperledger.besu.evm.account.AccountState;
 
@@ -31,14 +33,25 @@ public class AccountAdapter extends AdapterBase {
 
   private final Optional<Account> account;
   private final Address address;
+  private final Optional<Long> blockNumber;
 
   public AccountAdapter(final Account account) {
-    this(account == null ? null : account.getAddress(), account);
+    this(account == null ? null : account.getAddress(), account, Optional.empty());
+  }
+
+  public AccountAdapter(final Account account, final Optional<Long> blockNumber) {
+    this(account == null ? null : account.getAddress(), account, blockNumber);
   }
 
   public AccountAdapter(final Address address, final Account account) {
+    this(address, account, Optional.empty());
+  }
+
+  public AccountAdapter(
+      final Address address, final Account account, final Optional<Long> blockNumber) {
     this.account = Optional.ofNullable(account);
     this.address = address;
+    this.blockNumber = blockNumber;
   }
 
   public Address getAddress() {
@@ -53,14 +66,34 @@ public class AccountAdapter extends AdapterBase {
     return account.map(AccountState::getNonce).orElse(0L);
   }
 
-  public Bytes getCode() {
-    return account.map(AccountState::getCode).orElse(Bytes.EMPTY);
+  public Bytes getCode(final DataFetchingEnvironment environment) {
+
+    if (account.get() instanceof BonsaiAccount) {
+      final BlockchainQueries query = getBlockchainQueries(environment);
+      return query
+          .getAndMapWorldState(
+              blockNumber.orElse(query.headBlockNumber()),
+              ws -> Optional.of(ws.get(account.get().getAddress()).getCode()))
+          .get();
+    } else {
+      return account.map(AccountState::getCode).orElse(Bytes.EMPTY);
+    }
   }
 
   public Bytes32 getStorage(final DataFetchingEnvironment environment) {
+    final BlockchainQueries query = getBlockchainQueries(environment);
     final Bytes32 slot = environment.getArgument("slot");
-    return account
-        .map(a -> (Bytes32) a.getStorageValue(UInt256.fromBytes(slot)))
-        .orElse(Bytes32.ZERO);
+
+    if (account.get() instanceof BonsaiAccount) {
+      return query
+          .getAndMapWorldState(
+              blockNumber.orElse(query.headBlockNumber()),
+              ws -> Optional.of((Bytes32) ws.get(address).getStorageValue(UInt256.fromBytes(slot))))
+          .get();
+    } else {
+      return account
+          .map(a -> (Bytes32) a.getStorageValue(UInt256.fromBytes(slot)))
+          .orElse(Bytes32.ZERO);
+    }
   }
 }

ethereum/api/src/main/java/org/hyperledger/besu/ethereum/api/graphql/internal/pojoadapter/BlockAdapterBase.java

@@ -31,7 +31,6 @@ import org.hyperledger.besu.ethereum.mainnet.TransactionValidationParams;
 import org.hyperledger.besu.ethereum.rlp.BytesValueRLPOutput;
 import org.hyperledger.besu.ethereum.transaction.CallParameter;
 import org.hyperledger.besu.ethereum.transaction.TransactionSimulator;
-import org.hyperledger.besu.evm.account.Account;
 import org.hyperledger.besu.evm.log.LogTopic;
 import org.hyperledger.besu.evm.tracing.OperationTracer;
 
@@ -143,18 +142,12 @@ public class BlockAdapterBase extends AdapterBase {
   }
 
   public AccountAdapter getAccount(final DataFetchingEnvironment environment) {
-
     final BlockchainQueries query = getBlockchainQueries(environment);
-    final Address address = environment.getArgument("address");
     final long bn = header.getNumber();
+    final Address address = environment.getArgument("address");
     return query
         .getAndMapWorldState(
-            bn,
-            ws -> {
-              Account account = ws.get(address);
-              if (account != null) account.getCode();
-              return Optional.of(new AccountAdapter(account));
-            })
+            bn, ws -> Optional.of(new AccountAdapter(ws.get(address), Optional.of(bn))))
         .get();
   }
 

ethereum/api/src/main/java/org/hyperledger/besu/ethereum/api/graphql/internal/pojoadapter/EmptyAccountAdapter.java

@@ -45,7 +45,7 @@ public class EmptyAccountAdapter extends AccountAdapter {
   }
 
   @Override
-  public Bytes getCode() {
+  public Bytes getCode(final DataFetchingEnvironment environment) {
     return Bytes.EMPTY;
   }
 

ethereum/api/src/test/java/org/hyperledger/besu/ethereum/api/graphql/AbstractEthGraphQLHttpServiceTest.java

@@ -75,7 +75,7 @@ public abstract class AbstractEthGraphQLHttpServiceTest {
   public static void setupConstants() {
     blockchainSetupUtil =
         BlockchainSetupUtil.createForEthashChain(
-            BlockTestUtil.getHiveTestChainResources(), DataStorageFormat.FOREST);
+            BlockTestUtil.getHiveTestChainResources(), DataStorageFormat.BONSAI);
     blockchainSetupUtil.importAllBlocks();
   }