Stop making an assumption that contract code is immutable

1 files changed, 21 insertions, 3 deletions

diff --git a/packages/sol-cov/src/coverage_subprovider.ts b/packages/sol-cov/src/coverage_subprovider.ts
index ef425ee81..d3783abb2 100644
--- a/packages/sol-cov/src/coverage_subprovider.ts
+++ b/packages/sol-cov/src/coverage_subprovider.ts
@@ -5,6 +5,7 @@ import * as Web3 from 'web3';
 
 import { constants } from './constants';
 import { CoverageManager } from './coverage_manager';
+import { TraceInfoExistingContract, TraceInfoNewContract } from './types';
 
 /*
  * This class implements the web3-provider-engine subprovider interface and collects traces of all transactions that were sent and all calls that were executed.
@@ -84,15 +85,32 @@ export class CoverageSubprovider extends Subprovider {
         cb();
     }
     private async _recordTxTraceAsync(address: string, data: string | undefined, txHash: string): Promise<void> {
-        this._coverageManager.setTxDataByHash(txHash, data || '');
-        const payload = {
+        let payload = {
             method: 'debug_traceTransaction',
             params: [txHash, { disableMemory: true, disableStack: true, disableStorage: true }], // TODO For now testrpc just ignores those parameters https://github.com/trufflesuite/ganache-cli/issues/489
         };
         const jsonRPCResponsePayload = await this.emitPayloadAsync(payload);
         const trace: Web3.TransactionTrace = jsonRPCResponsePayload.result;
         const coveredPcs = _.map(trace.structLogs, log => log.pc);
-        this._coverageManager.appendTraceInfo(address, { coveredPcs, txHash });
+        if (address === constants.NEW_CONTRACT) {
+            const traceInfo: TraceInfoNewContract = {
+                coveredPcs,
+                txHash,
+                address,
+                bytecode: data as string,
+            };
+            this._coverageManager.appendTraceInfo(traceInfo);
+        } else {
+            payload = { method: 'eth_getCode', params: [address, 'latest'] };
+            const runtimeBytecode = (await this.emitPayloadAsync(payload)).result;
+            const traceInfo: TraceInfoExistingContract = {
+                coveredPcs,
+                txHash,
+                address,
+                runtimeBytecode,
+            };
+            this._coverageManager.appendTraceInfo(traceInfo);
+        }
     }
     private async _recordCallTraceAsync(callData: Partial<Web3.CallData>, blockNumber: Web3.BlockParam): Promise<void> {
         const snapshotId = Number((await this.emitPayloadAsync({ method: 'evm_snapshot' })).result);