diff --git a/lib/forking/forked_blockchain.js b/lib/forking/forked_blockchain.js index 713075a5b6..a63c6fc64c 100644 --- a/lib/forking/forked_blockchain.js +++ b/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); + }); + } }); }); }; diff --git a/lib/forking/forked_storage_trie.js b/lib/forking/forked_storage_trie.js index f63aecb68d..6e9c34f682 100644 --- a/lib/forking/forked_storage_trie.js +++ b/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) { diff --git a/test/forking/forking.js b/test/forking/forking.js index ec3bcaa6c9..7d9a97a201 100644 --- a/test/forking/forking.js +++ b/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);