Merge 3164f05 into 34cdd8a

trufflesuite · Oct 8, 2019 · 0a77f10 · 0a77f10
2 parents 34cdd8a + 3164f05
commit 0a77f10
Showing 1 changed file with 34 additions and 10 deletions.
diff --git a/lib/utils/forkedblockchain.js b/lib/utils/forkedblockchain.js
@@ -15,26 +15,46 @@ var inherits = require("util").inherits;
 
 inherits(ForkedBlockchain, BlockchainDouble);
 
+const httpReg = /^https?:/i;
+const protocolReg = /^[A-Za-z][A-Za-z0-9+\-.]*:/;
+const validProtocolReg = /^(?:http|ws)s?:/i;
+
 function ForkedBlockchain(options) {
   this.options = options || {};
 
-  if (options.fork == null) {
+  if (options.fork == null || (typeof options.fork === "string" && options.fork.trim().length === 0)) {
     throw new Error("ForkedBlockchain must be passed a fork parameter.");
   }
 
-  this.fork = options.fork;
-  this.forkBlockNumber = options.fork_block_number;
   this.forkVersion = null;
 
-  if (typeof this.fork === "string") {
-    if (this.fork.indexOf("@") >= 0) {
-      var split = this.fork.split("@");
-      this.fork = split[0];
-      this.forkBlockNumber = parseInt(split[1]);
+  if (typeof options.fork === "string") {
+    if (options.fork.indexOf("@") >= 0) {
+      const split = options.fork.split("@", 2);
+      options.fork = split[0];
+      options.fork_block_number = parseInt(split[1]);
+    }
+
+    let fork;
+    if (!protocolReg.test(options.fork)) {
+      // we don't have a protocol at all, assume ws
+      options.fork = "ws://" + options.fork;
+      fork = new Web3.providers.HttpProvider(options.fork);
+    } else if (validProtocolReg.test(options.fork)) {
+      if (httpReg.test(options.fork)) {
+        fork = new Web3.providers.HttpProvider(options.fork);
+      } else {
+        fork = new Web3.providers.WebsocketProvider(options.fork);
+      }
+    } else {
+      throw new Error(`Invalid scheme for fork url: ${options.fork}. Supported schemes are: http, https, ws, and wss.`);
     }
 
-    this.fork = new Web3.providers.HttpProvider(this.fork);
+    this.fork = fork;
+  } else {
+    this.fork = options.fork;
   }
+  this.forkBlockNumber = options.fork_block_number;
 
   this.time = options.time;
   this.storageTrieCache = {};
@@ -194,14 +214,17 @@ ForkedBlockchain.prototype.getFallbackBlock = function(numberOrHash, cb) {
     // buffers instead of numbers.
     numberOrHash = to.rpcDataHexString(numberOrHash);
   }
+  if (typeof numberOrHash === "string" && numberOrHash.length !== 66) {
+    numberOrHash = to.number(numberOrHash);
+  }
 
   self.web3.eth.getBlock(numberOrHash, true, function(err, json) {
     if (err) {
       return cb(err);
     }
 
     if (json == null) {
-      return cb();
+      return cb(new Error("Block not found"));
     }
 
     var block = new Block();
@@ -439,6 +462,7 @@ ForkedBlockchain.prototype.getTransactionReceipt = function(hash, callback) {
           }
 
           var logs = receiptJson.logs.map(function(log) {
+            log.block = result.block;
             return new Log(log);
           });