De-duplicate code by refactoring subprovider classes

5 files changed, 89 insertions, 204 deletions

diff --git a/packages/sol-cov/src/coverage_subprovider.ts b/packages/sol-cov/src/coverage_subprovider.ts
index 0fa7f873e..065a48434 100644
--- a/packages/sol-cov/src/coverage_subprovider.ts
+++ b/packages/sol-cov/src/coverage_subprovider.ts
@@ -2,8 +2,8 @@ import * as _ from 'lodash';
 
 import { AbstractArtifactAdapter } from './artifact_adapters/abstract_artifact_adapter';
 import { collectCoverageEntries } from './collect_coverage_entries';
-import { TraceCollectionSubprovider } from './trace_collection_subprovider';
 import { SingleFileSubtraceHandler, TraceCollector } from './trace_collector';
+import { TraceInfoSubprovider } from './trace_info_subprovider';
 import {
     BranchCoverage,
     ContractData,
@@ -22,7 +22,7 @@ import { utils } from './utils';
  * This class implements the [web3-provider-engine](https://github.com/MetaMask/provider-engine) subprovider interface.
  * It's used to compute your code coverage while running solidity tests.
  */
-export class CoverageSubprovider extends TraceCollectionSubprovider {
+export class CoverageSubprovider extends TraceInfoSubprovider {
     private _coverageCollector: TraceCollector;
     /**
      * Instantiates a CoverageSubprovider instance
@@ -39,7 +39,7 @@ export class CoverageSubprovider extends TraceCollectionSubprovider {
         super(defaultFromAddress, traceCollectionSubproviderConfig);
         this._coverageCollector = new TraceCollector(artifactAdapter, isVerbose, coverageHandler);
     }
-    public async handleTraceInfoAsync(traceInfo: TraceInfo): Promise<void> {
+    protected async _handleTraceInfoAsync(traceInfo: TraceInfo): Promise<void> {
         await this._coverageCollector.computeSingleTraceCoverageAsync(traceInfo);
     }
     /**
diff --git a/packages/sol-cov/src/profiler_subprovider.ts b/packages/sol-cov/src/profiler_subprovider.ts
index 62ed1b472..9f98da524 100644
--- a/packages/sol-cov/src/profiler_subprovider.ts
+++ b/packages/sol-cov/src/profiler_subprovider.ts
@@ -2,8 +2,8 @@ import * as _ from 'lodash';
 
 import { AbstractArtifactAdapter } from './artifact_adapters/abstract_artifact_adapter';
 import { collectCoverageEntries } from './collect_coverage_entries';
-import { TraceCollectionSubprovider } from './trace_collection_subprovider';
 import { SingleFileSubtraceHandler, TraceCollector } from './trace_collector';
+import { TraceInfoSubprovider } from './trace_info_subprovider';
 import { ContractData, Coverage, SourceRange, Subtrace, TraceInfo } from './types';
 import { utils } from './utils';
 
@@ -11,7 +11,7 @@ import { utils } from './utils';
  * This class implements the [web3-provider-engine](https://github.com/MetaMask/provider-engine) subprovider interface.
  * ProfilerSubprovider is used to profile Solidity code while running tests.
  */
-export class ProfilerSubprovider extends TraceCollectionSubprovider {
+export class ProfilerSubprovider extends TraceInfoSubprovider {
     private _profilerCollector: TraceCollector;
     /**
      * Instantiates a ProfilerSubprovider instance
@@ -28,7 +28,7 @@ export class ProfilerSubprovider extends TraceCollectionSubprovider {
         super(defaultFromAddress, traceCollectionSubproviderConfig);
         this._profilerCollector = new TraceCollector(artifactAdapter, isVerbose, profilerHandler);
     }
-    public async handleTraceInfoAsync(traceInfo: TraceInfo): Promise<void> {
+    protected async _handleTraceInfoAsync(traceInfo: TraceInfo): Promise<void> {
         await this._profilerCollector.computeSingleTraceCoverageAsync(traceInfo);
     }
     /**
diff --git a/packages/sol-cov/src/revert_trace_subprovider.ts b/packages/sol-cov/src/revert_trace_subprovider.ts
index 68e752aee..f6501757a 100644
--- a/packages/sol-cov/src/revert_trace_subprovider.ts
+++ b/packages/sol-cov/src/revert_trace_subprovider.ts
@@ -1,149 +1,43 @@
-import { BlockchainLifecycle } from '@0xproject/dev-utils';
-import { Callback, ErrorCallback, NextCallback, Subprovider } from '@0xproject/subproviders';
-import { Web3Wrapper } from '@0xproject/web3-wrapper';
-import { CallData, JSONRPCRequestPayload, Provider, TxData } from 'ethereum-types';
 import { stripHexPrefix } from 'ethereumjs-util';
 import * as _ from 'lodash';
 import { getLogger, levels, Logger } from 'loglevel';
-import { Lock } from 'semaphore-async-await';
 
 import { AbstractArtifactAdapter } from './artifact_adapters/abstract_artifact_adapter';
 import { constants } from './constants';
 import { getRevertTrace } from './revert_trace';
 import { parseSourceMap } from './source_maps';
-import { BlockParamLiteral, ContractData, EvmCallStack, SourceRange } from './types';
+import { TraceCollectionSubprovider } from './trace_collection_subprovider';
+import { ContractData, EvmCallStack, SourceRange } from './types';
 import { utils } from './utils';
 
-interface MaybeFakeTxData extends TxData {
-    isFakeTransaction?: boolean;
-}
-
-const BLOCK_GAS_LIMIT = 6000000;
-
 /**
  * This class implements the [web3-provider-engine](https://github.com/MetaMask/provider-engine) subprovider interface.
  * It is used to report call stack traces whenever a revert occurs.
  */
-export class RevertTraceSubprovider extends Subprovider {
+export class RevertTraceSubprovider extends TraceCollectionSubprovider {
     // Lock is used to not accept normal transactions while doing call/snapshot magic because they'll be reverted later otherwise
-    private _lock = new Lock();
-    private _defaultFromAddress: string;
-    private _web3Wrapper!: Web3Wrapper;
-    private _isEnabled = true;
-    private _artifactAdapter: AbstractArtifactAdapter;
     private _contractsData!: ContractData[];
+    private _artifactAdapter: AbstractArtifactAdapter;
     private _logger: Logger;
 
     /**
-     * Instantiates a TraceCollectionSubprovider instance
+     * Instantiates a RevertTraceSubprovider instance
+     * @param artifactAdapter Adapter for used artifacts format (0x, truffle, giveth, etc.)
      * @param defaultFromAddress default from address to use when sending transactions
+     * @param isVerbose If true, we will log any unknown transactions. Otherwise we will ignore them
      */
-    constructor(artifactAdapter: AbstractArtifactAdapter, defaultFromAddress: string, isVerbose: boolean) {
-        super();
+    constructor(artifactAdapter: AbstractArtifactAdapter, defaultFromAddress: string, isVerbose: boolean = true) {
+        const traceCollectionSubproviderConfig = {
+            shouldCollectTransactionTraces: true,
+            shouldCollectGasEstimateTraces: true,
+            shouldCollectCallTraces: true,
+        };
+        super(defaultFromAddress, traceCollectionSubproviderConfig);
         this._artifactAdapter = artifactAdapter;
-        this._defaultFromAddress = defaultFromAddress;
         this._logger = getLogger('sol-cov');
         this._logger.setLevel(isVerbose ? levels.TRACE : levels.ERROR);
     }
-    /**
-     * Starts trace collection
-     */
-    public start(): void {
-        this._isEnabled = true;
-    }
-    /**
-     * Stops trace collection
-     */
-    public stop(): void {
-        this._isEnabled = false;
-    }
-    /**
-     * This method conforms to the web3-provider-engine interface.
-     * It is called internally by the ProviderEngine when it is this subproviders
-     * turn to handle a JSON RPC request.
-     * @param payload JSON RPC payload
-     * @param next Callback to call if this subprovider decides not to handle the request
-     * @param end Callback to call if subprovider handled the request and wants to pass back the request.
-     */
-    // tslint:disable-next-line:prefer-function-over-method async-suffix
-    public async handleRequest(payload: JSONRPCRequestPayload, next: NextCallback, _end: ErrorCallback): Promise<void> {
-        if (this._isEnabled) {
-            switch (payload.method) {
-                case 'eth_sendTransaction':
-                    const txData = payload.params[0];
-                    next(this._onTransactionSentAsync.bind(this, txData));
-                    return;
-
-                case 'eth_call':
-                    const callData = payload.params[0];
-                    next(this._onCallOrGasEstimateExecutedAsync.bind(this, callData));
-                    return;
-
-                case 'eth_estimateGas':
-                    const estimateGasData = payload.params[0];
-                    next(this._onCallOrGasEstimateExecutedAsync.bind(this, estimateGasData));
-                    return;
-
-                default:
-                    next();
-                    return;
-            }
-        } else {
-            next();
-            return;
-        }
-    }
-    /**
-     * Set's the subprovider's engine to the ProviderEngine it is added to.
-     * This is only called within the ProviderEngine source code, do not call
-     * directly.
-     */
-    public setEngine(engine: Provider): void {
-        super.setEngine(engine);
-        this._web3Wrapper = new Web3Wrapper(engine);
-    }
-    private async _onTransactionSentAsync(
-        txData: MaybeFakeTxData,
-        err: Error | null,
-        txHash: string | undefined,
-        cb: Callback,
-    ): Promise<void> {
-        if (!txData.isFakeTransaction) {
-            // This transaction is a usual transaction. Not a call executed as one.
-            // And we don't want it to be executed within a snapshotting period
-            await this._lock.acquire();
-        }
-        const NULL_ADDRESS = '0x0';
-        if (_.isNull(err)) {
-            const toAddress =
-                _.isUndefined(txData.to) || txData.to === NULL_ADDRESS ? constants.NEW_CONTRACT : txData.to;
-            await this._recordTxTraceAsync(toAddress, txHash as string);
-        } else {
-            const latestBlock = await this._web3Wrapper.getBlockWithTransactionDataAsync(BlockParamLiteral.Latest);
-            const transactions = latestBlock.transactions;
-            for (const transaction of transactions) {
-                const toAddress =
-                    _.isUndefined(txData.to) || txData.to === NULL_ADDRESS ? constants.NEW_CONTRACT : txData.to;
-                await this._recordTxTraceAsync(toAddress, transaction.hash);
-            }
-        }
-        if (!txData.isFakeTransaction) {
-            // This transaction is a usual transaction. Not a call executed as one.
-            // And we don't want it to be executed within a snapshotting period
-            this._lock.release();
-        }
-        cb();
-    }
-    private async _onCallOrGasEstimateExecutedAsync(
-        callData: Partial<CallData>,
-        _err: Error | null,
-        _callResult: string,
-        cb: Callback,
-    ): Promise<void> {
-        await this._recordCallOrGasEstimateTraceAsync(callData);
-        cb();
-    }
-    private async _recordTxTraceAsync(address: string, txHash: string): Promise<void> {
+    protected async _recordTxTraceAsync(address: string, data: string | undefined, txHash: string): Promise<void> {
         await this._web3Wrapper.awaitTransactionMinedAsync(txHash, 0);
         const trace = await this._web3Wrapper.getTransactionTraceAsync(txHash, {
             disableMemory: true,
@@ -157,28 +51,6 @@ export class RevertTraceSubprovider extends Subprovider {
             await this._printStackTraceAsync(evmCallStack);
         }
     }
-    private async _recordCallOrGasEstimateTraceAsync(callData: Partial<CallData>): Promise<void> {
-        // We don't want other transactions to be exeucted during snashotting period, that's why we lock the
-        // transaction execution for all transactions except our fake ones.
-        await this._lock.acquire();
-        const blockchainLifecycle = new BlockchainLifecycle(this._web3Wrapper);
-        await blockchainLifecycle.startAsync();
-        const fakeTxData: MaybeFakeTxData = {
-            gas: BLOCK_GAS_LIMIT,
-            isFakeTransaction: true, // This transaction (and only it) is allowed to come through when the lock is locked
-            ...callData,
-            from: callData.from || this._defaultFromAddress,
-        };
-        try {
-            const txHash = await this._web3Wrapper.sendTransactionAsync(fakeTxData);
-            await this._web3Wrapper.awaitTransactionMinedAsync(txHash, 0);
-        } catch (err) {
-            // Even if this transaction failed - we've already recorded it's trace.
-            _.noop();
-        }
-        await blockchainLifecycle.revertAsync();
-        this._lock.release();
-    }
     private async _printStackTraceAsync(evmCallStack: EvmCallStack): Promise<void> {
         const sourceRanges: SourceRange[] = [];
         if (_.isUndefined(this._contractsData)) {
diff --git a/packages/sol-cov/src/trace_collection_subprovider.ts b/packages/sol-cov/src/trace_collection_subprovider.ts
index 742735935..9866472b9 100644
--- a/packages/sol-cov/src/trace_collection_subprovider.ts
+++ b/packages/sol-cov/src/trace_collection_subprovider.ts
@@ -6,8 +6,7 @@ import * as _ from 'lodash';
 import { Lock } from 'semaphore-async-await';
 
 import { constants } from './constants';
-import { getTracesByContractAddress } from './trace';
-import { BlockParamLiteral, TraceInfo, TraceInfoExistingContract, TraceInfoNewContract } from './types';
+import { BlockParamLiteral } from './types';
 
 interface MaybeFakeTxData extends TxData {
     isFakeTransaction?: boolean;
@@ -27,13 +26,14 @@ export interface TraceCollectionSubproviderConfig {
 
 /**
  * This class implements the [web3-provider-engine](https://github.com/MetaMask/provider-engine) subprovider interface.
- * It collects traces of all transactions that were sent and all calls that were executed through JSON RPC.
+ * It collects traces of all transactions that were sent and all calls that were executed through JSON RPC. It must
+ * be extended by implementing the _recordTxTraceAsync method which is called for every transaction.
  */
 export abstract class TraceCollectionSubprovider extends Subprovider {
+    protected _web3Wrapper!: Web3Wrapper;
     // Lock is used to not accept normal transactions while doing call/snapshot magic because they'll be reverted later otherwise
     private _lock = new Lock();
     private _defaultFromAddress: string;
-    private _web3Wrapper!: Web3Wrapper;
     private _isEnabled = true;
     private _config: TraceCollectionSubproviderConfig;
     /**
@@ -58,11 +58,6 @@ export abstract class TraceCollectionSubprovider extends Subprovider {
         this._isEnabled = false;
     }
     /**
-     * Called for each subtrace.
-     * @param traceInfo Trace info for this subtrace.
-     */
-    public abstract handleTraceInfoAsync(traceInfo: TraceInfo): Promise<void>;
-    /**
      * This method conforms to the web3-provider-engine interface.
      * It is called internally by the ProviderEngine when it is this subproviders
      * turn to handle a JSON RPC request.
@@ -119,6 +114,11 @@ export abstract class TraceCollectionSubprovider extends Subprovider {
         super.setEngine(engine);
         this._web3Wrapper = new Web3Wrapper(engine);
     }
+    protected abstract async _recordTxTraceAsync(
+        address: string,
+        data: string | undefined,
+        txHash: string,
+    ): Promise<void>;
     private async _onTransactionSentAsync(
         txData: MaybeFakeTxData,
         err: Error | null,
@@ -160,52 +160,6 @@ export abstract class TraceCollectionSubprovider extends Subprovider {
         await this._recordCallOrGasEstimateTraceAsync(callData);
         cb();
     }
-    private async _recordTxTraceAsync(address: string, data: string | undefined, txHash: string): Promise<void> {
-        await this._web3Wrapper.awaitTransactionMinedAsync(txHash, 0);
-        const trace = await this._web3Wrapper.getTransactionTraceAsync(txHash, {
-            disableMemory: true,
-            disableStack: false,
-            disableStorage: true,
-        });
-        const tracesByContractAddress = getTracesByContractAddress(trace.structLogs, address);
-        const subcallAddresses = _.keys(tracesByContractAddress);
-        if (address === constants.NEW_CONTRACT) {
-            for (const subcallAddress of subcallAddresses) {
-                let traceInfo: TraceInfoNewContract | TraceInfoExistingContract;
-                if (subcallAddress === 'NEW_CONTRACT') {
-                    const traceForThatSubcall = tracesByContractAddress[subcallAddress];
-                    traceInfo = {
-                        subtrace: traceForThatSubcall,
-                        txHash,
-                        address: subcallAddress,
-                        bytecode: data as string,
-                    };
-                } else {
-                    const runtimeBytecode = await this._web3Wrapper.getContractCodeAsync(subcallAddress);
-                    const traceForThatSubcall = tracesByContractAddress[subcallAddress];
-                    traceInfo = {
-                        subtrace: traceForThatSubcall,
-                        txHash,
-                        address: subcallAddress,
-                        runtimeBytecode,
-                    };
-                }
-                await this.handleTraceInfoAsync(traceInfo);
-            }
-        } else {
-            for (const subcallAddress of subcallAddresses) {
-                const runtimeBytecode = await this._web3Wrapper.getContractCodeAsync(subcallAddress);
-                const traceForThatSubcall = tracesByContractAddress[subcallAddress];
-                const traceInfo: TraceInfoExistingContract = {
-                    subtrace: traceForThatSubcall,
-                    txHash,
-                    address: subcallAddress,
-                    runtimeBytecode,
-                };
-                await this.handleTraceInfoAsync(traceInfo);
-            }
-        }
-    }
     private async _recordCallOrGasEstimateTraceAsync(callData: Partial<CallData>): Promise<void> {
         // We don't want other transactions to be exeucted during snashotting period, that's why we lock the
         // transaction execution for all transactions except our fake ones.
diff --git a/packages/sol-cov/src/trace_info_subprovider.ts b/packages/sol-cov/src/trace_info_subprovider.ts
new file mode 100644
index 000000000..635a68f58
--- /dev/null
+++ b/packages/sol-cov/src/trace_info_subprovider.ts
@@ -0,0 +1,59 @@
+import * as _ from 'lodash';
+
+import { constants } from './constants';
+import { getTracesByContractAddress } from './trace';
+import { TraceCollectionSubprovider } from './trace_collection_subprovider';
+import { TraceInfo, TraceInfoExistingContract, TraceInfoNewContract } from './types';
+
+// TraceInfoSubprovider is extended by subproviders which need to work with one
+// TraceInfo at a time. It has one abstract method: _handleTraceInfoAsync, which
+// is called for each TraceInfo.
+export abstract class TraceInfoSubprovider extends TraceCollectionSubprovider {
+    protected abstract _handleTraceInfoAsync(traceInfo: TraceInfo): Promise<void>;
+    protected async _recordTxTraceAsync(address: string, data: string | undefined, txHash: string): Promise<void> {
+        await this._web3Wrapper.awaitTransactionMinedAsync(txHash, 0);
+        const trace = await this._web3Wrapper.getTransactionTraceAsync(txHash, {
+            disableMemory: true,
+            disableStack: false,
+            disableStorage: true,
+        });
+        const tracesByContractAddress = getTracesByContractAddress(trace.structLogs, address);
+        const subcallAddresses = _.keys(tracesByContractAddress);
+        if (address === constants.NEW_CONTRACT) {
+            for (const subcallAddress of subcallAddresses) {
+                let traceInfo: TraceInfoNewContract | TraceInfoExistingContract;
+                if (subcallAddress === 'NEW_CONTRACT') {
+                    const traceForThatSubcall = tracesByContractAddress[subcallAddress];
+                    traceInfo = {
+                        subtrace: traceForThatSubcall,
+                        txHash,
+                        address: subcallAddress,
+                        bytecode: data as string,
+                    };
+                } else {
+                    const runtimeBytecode = await this._web3Wrapper.getContractCodeAsync(subcallAddress);
+                    const traceForThatSubcall = tracesByContractAddress[subcallAddress];
+                    traceInfo = {
+                        subtrace: traceForThatSubcall,
+                        txHash,
+                        address: subcallAddress,
+                        runtimeBytecode,
+                    };
+                }
+                await this._handleTraceInfoAsync(traceInfo);
+            }
+        } else {
+            for (const subcallAddress of subcallAddresses) {
+                const runtimeBytecode = await this._web3Wrapper.getContractCodeAsync(subcallAddress);
+                const traceForThatSubcall = tracesByContractAddress[subcallAddress];
+                const traceInfo: TraceInfoExistingContract = {
+                    subtrace: traceForThatSubcall,
+                    txHash,
+                    address: subcallAddress,
+                    runtimeBytecode,
+                };
+                await this._handleTraceInfoAsync(traceInfo);
+            }
+        }
+    }
+}