From a8a4d8f04a357380963c62110fc7659b4f03c4ec Mon Sep 17 00:00:00 2001
From: Mike Seese <seesemichaelj@gmail.com>
Date: Tue, 18 Aug 2020 12:44:59 -0700
Subject: [PATCH] fix(forking): invalidate `_deleted` in `ForkedStorageTrie` if
 a subsequent `put` happens (#612)

---
 lib/forking/forked_storage_trie.js       | 17 ++++++++
 test/contracts/forking/StorageDelete.sol | 18 ++++++++
 test/forking/delete.js                   | 54 ++++++++++++++++++++++++
 3 files changed, 89 insertions(+)
 create mode 100644 test/contracts/forking/StorageDelete.sol
 create mode 100644 test/forking/delete.js

diff --git a/lib/forking/forked_storage_trie.js b/lib/forking/forked_storage_trie.js
index c64bb48ff2..f63aecb68d 100644
--- a/lib/forking/forked_storage_trie.js
+++ b/lib/forking/forked_storage_trie.js
@@ -97,6 +97,23 @@ ForkedStorageBaseTrie.prototype.keyExists = function(key, callback) {
   });
 };
 
+const originalPut = ForkedStorageBaseTrie.prototype.put;
+ForkedStorageBaseTrie.prototype.put = function(key, value, callback) {
+  let rpcKey = to.rpcDataHexString(key);
+  if (this.address) {
+    rpcKey = `${to.rpcDataHexString(this.address)};${rpcKey}`;
+  }
+  this._deleted.get(rpcKey, (_, result) => {
+    if (result === 1) {
+      this._deleted.put(rpcKey, 0, () => {
+        originalPut.call(this, key, value, callback);
+      });
+    } else {
+      originalPut.call(this, key, value, callback);
+    }
+  });
+};
+
 ForkedStorageBaseTrie.prototype.keyIsDeleted = function(key, callback) {
   let rpcKey = to.rpcDataHexString(key);
   if (this.address) {
diff --git a/test/contracts/forking/StorageDelete.sol b/test/contracts/forking/StorageDelete.sol
new file mode 100644
index 0000000000..c40a5b70b4
--- /dev/null
+++ b/test/contracts/forking/StorageDelete.sol
@@ -0,0 +1,18 @@
+pragma solidity ^0.6.0;
+
+contract StorageDelete {
+  bool entered;
+
+  constructor() public {
+    entered = true;
+  }
+
+  function test() public nonReentrant {}
+
+  modifier nonReentrant() {
+    require(entered, "re-entered");
+    entered = false;
+    _;
+    entered = true;
+  }
+}
diff --git a/test/forking/delete.js b/test/forking/delete.js
new file mode 100644
index 0000000000..6be8048d51
--- /dev/null
+++ b/test/forking/delete.js
@@ -0,0 +1,54 @@
+const assert = require("assert");
+const bootstrap = require("../helpers/contract/bootstrap");
+const intializeTestProvider = require("../helpers/web3/initializeTestProvider");
+
+/**
+ * NOTE: Naming in these tests is a bit confusing. Here, the "main chain"
+ * is the main chain the tests interact with; and the "forked chain" is the
+ * chain that _was forked_. This is in contrast to general naming, where the
+ * main chain represents the main chain to be forked (like the Ethereum live
+ * network) and the fork chaing being "the fork".
+ */
+
+describe("Forking Deletion", () => {
+  let forkedContext;
+  let mainContext;
+  const logger = {
+    log: function(msg) {}
+  };
+
+  before("Set up forked provider with web3 instance and deploy a contract", async function() {
+    this.timeout(5000);
+
+    const contractRef = {
+      contractFiles: ["StorageDelete"],
+      contractSubdirectory: "forking"
+    };
+
+    const ganacheProviderOptions = {
+      logger,
+      seed: "main provider"
+    };
+
+    forkedContext = await bootstrap(contractRef, ganacheProviderOptions);
+  });
+
+  before("Set up main provider and web3 instance", async function() {
+    const { provider: forkedProvider } = forkedContext;
+    mainContext = await intializeTestProvider({
+      fork: forkedProvider,
+      logger,
+      seed: "forked provider"
+    });
+  });
+
+  it("successfully manages storage slot deletion", async() => {
+    const { instance: forkedInstance, abi } = forkedContext;
+    const { web3: mainWeb3 } = mainContext;
+
+    const instance = new mainWeb3.eth.Contract(abi, forkedInstance._address);
+
+    assert.ok(await instance.methods.test().call());
+    assert.ok(await instance.methods.test().call());
+  });
+});