Merge e53379c into a8a4d8f

lib/forking/forked_blockchain.js

@@ -403,11 +403,31 @@ ForkedBlockchain.prototype.getStorage = function(address, key, number, callback)
     if (err) {
       return callback(err);
     }
-    this.getEffectiveBlockNumber(number, (err, number) => {
+    this.getEffectiveBlockNumber(number, (err, blockNumber) => {
       if (err) {
         return callback(err);
       }
-      trie.get(utils.setLengthLeft(utils.toBuffer(key), 32), number, callback);
+
+      if (blockNumber > this.forkBlockNumber) {
+        // only hit the ForkedStorageTrieBase if we're not looking
+        // for something that's on the forked chain
+        trie.get(utils.setLengthLeft(utils.toBuffer(key), 32), blockNumber, callback);
+      } else {
+        // we're looking for something prior to forking, so let's
+        // hit eth_getStorageAt
+        this.web3.eth.getStorageAt(to.rpcDataHexString(address), to.rpcDataHexString(key), blockNumber, function(
+          err,
+          value
+        ) {
+          if (err) {
+            return callback(err);
+          }
+
+          value = utils.rlp.encode(value);
+
+          callback(null, value);
+        });
+      }
     });
   });
 };

lib/forking/forked_storage_trie.js

@@ -2,6 +2,7 @@ const Sublevel = require("level-sublevel");
 const MerklePatriciaTree = require("merkle-patricia-tree");
 const BaseTrie = require("merkle-patricia-tree/baseTrie");
 const checkpointInterface = require("merkle-patricia-tree/checkpoint-interface");
+const Account = require("ethereumjs-account").default;
 var utils = require("ethereumjs-util");
 var inherits = require("util").inherits;
 var Web3 = require("web3");
@@ -24,11 +25,13 @@ function ForkedStorageBaseTrie(db, root, options) {
 // Note: This overrides a standard method whereas the other methods do not.
 ForkedStorageBaseTrie.prototype.get = function(key, blockNumber, callback) {
   var self = this;
+  let blockNumberProvided = true;
 
   // Allow an optional blockNumber
   if (typeof blockNumber === "function") {
     callback = blockNumber;
     blockNumber = this.forkBlockNumber;
+    blockNumberProvided = false;
   }
 
   key = utils.toBuffer(key);
@@ -40,12 +43,49 @@ ForkedStorageBaseTrie.prototype.get = function(key, blockNumber, callback) {
     }
 
     if (exists) {
-      // TODO: just because we have the key doesn't mean we're at the right
-      // block number/root to send it. We need to check the block number
-      // before using the data in our own trie.
-      MerklePatriciaTree.prototype.get.call(self, key, function(err, r) {
-        callback(err, r);
-      });
+      // I'm checking to see if a blockNumber is provided because the below
+      // logic breaks for things like nonce lookup, in which we should just
+      // use the root trie as is. I'm guessing there's a cleaner architecture
+      // that doesn't require such checks
+      if (blockNumberProvided) {
+        // this logic is heavily influenced by BlockchainDouble.prototype.getStorage
+        // but some adjustments were necessary due to the ForkedStorageTrieBase context
+        self.blockchain.getBlock(blockNumber, function(err, block) {
+          if (err) {
+            return callback(err);
+          }
+
+          // Manipulate the state root in place to maintain checkpoints
+          const currentStateRoot = self.root;
+          self.root = block.header.stateRoot;
+
+          MerklePatriciaTree.prototype.get.call(self, utils.toBuffer(self.address), function(err, data) {
+            if (err != null) {
+              // Put the stateRoot back if there's an error
+              self.root = currentStateRoot;
+              return callback(err);
+            }
+
+            const account = new Account(data);
+
+            self.root = account.stateRoot;
+            MerklePatriciaTree.prototype.get.call(self, key, function(err, value) {
+              // Finally, put the stateRoot back for good
+              self.root = currentStateRoot;
+
+              if (err != null) {
+                return callback(err, value);
+              }
+
+              callback(null, value);
+            });
+          });
+        });
+      } else {
+        MerklePatriciaTree.prototype.get.call(self, key, function(err, r) {
+          callback(err, r);
+        });
+      }
     } else {
       self.keyIsDeleted(key, (err, deleted) => {
         if (err) {

test/forking/forking.js

@@ -44,6 +44,7 @@ describe("Forking", function() {
   var forkBlockNumber;
 
   var initialDeployTransactionHash;
+  var variableChangedBlockNumber;
 
   before("set up test data", function() {
     this.timeout(10000);
@@ -184,7 +185,8 @@ describe("Forking", function() {
       });
     });
 
-    await mainExample.methods.setValue(7).send({ from: mainAccounts[0] });
+    const receipt = await mainExample.methods.setValue(7).send({ from: mainAccounts[0] });
+    variableChangedBlockNumber = receipt.blockNumber;
     await eventData;
   });
 
@@ -324,6 +326,15 @@ describe("Forking", function() {
     assert.strictEqual(mainWeb3.utils.hexToNumber(result), 7);
   });
 
+  it("should get the correct storage values based on block", async() => {
+    const result = await mainWeb3.eth.getStorageAt(
+      thirdContractAddress,
+      contract.position_of_value,
+      variableChangedBlockNumber - 1
+    );
+    assert.strictEqual(mainWeb3.utils.hexToNumber(result), 5);
+  });
+
   it("should get storage values on the forked provider via the main provider", async() => {
     const result = await mainWeb3.eth.getStorageAt(contractAddress, contract.position_of_value);
     assert.strictEqual(mainWeb3.utils.hexToNumber(result), 7);