Defer trace output display

WorksButNotTested · Sep 19, 2024 · 328b613 · 328b613
1 parent 22e100c
commit 328b613
Show file tree

Hide file tree

Showing 7 changed files with 167 additions and 89 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "frida-cshell",
-  "version": "1.4.5",
+  "version": "1.4.6",
   "description": "Frida's CShell",
   "scripts": {
     "prepare": "npm run build && npm run version && npm run package && npm run copy",

diff --git a/src/breakpoints/bp.ts b/src/breakpoints/bp.ts
@@ -7,7 +7,7 @@ import { Parser } from '../io/parser.js';
 import { Regs } from './regs.js';
 import { Format } from '../misc/format.js';
 import { BlockTrace } from '../traces/block.js';
-import { Trace, Traces } from '../traces/trace.js';
+import { Trace, TraceData, Traces } from '../traces/trace.js';
 import { Var } from '../vars/var.js';
 import { Vars } from '../vars/vars.js';
 import { CallTrace } from '../traces/call.js';
@@ -210,11 +210,8 @@ export class Bp {
       if (this._trace === null) return;
       this._trace.stop();
 
-      Output.writeln(Output.blue('-'.repeat(80)));
-      this._trace.display();
-      Output.writeln(Output.blue('-'.repeat(80)));
       Output.clearLine();
-
+      Output.writeln(Output.blue('-'.repeat(80)));
       Output.writeln(
         [
           `${Output.yellow('|')} Stop Trace`,
@@ -227,13 +224,35 @@ export class Bp {
       );
       Output.writeln(Output.yellow('-'.repeat(80)));
 
+      const data = this._trace.data();
+      setTimeout(() => this.displayTraceData(data));
+
       Traces.delete(threadId);
     } finally {
       Input.prompt();
       Regs.clear();
     }
   }
 
+  private displayTraceData(trace: TraceData) {
+    Output.clearLine();
+    Output.writeln(Output.yellow('-'.repeat(80)));
+    Output.writeln(`${Output.yellow('|')} Displaying trace:`);
+    Output.writeln(Output.yellow('-'.repeat(80)));
+    Input.suppressIntercept(true);
+    Output.setIndent(true);
+    Output.writeln();
+    try {
+      trace.display();
+      Output.writeln();
+    } finally {
+      Output.setIndent(false);
+      Input.suppressIntercept(false);
+      Output.writeln(Output.yellow('-'.repeat(80)));
+      Input.prompt();
+    }
+  }
+
   private breakCode(
     threadId: ThreadId,
     ctx: CpuContext,

diff --git a/src/traces/block.ts b/src/traces/block.ts
@@ -1,45 +1,20 @@
 import { Output } from '../io/output.js';
-import { Trace, Traces } from './trace.js';
+import { Trace, TraceData, Traces } from './trace.js';
 
-export class BlockTrace implements Trace {
+class BlockTraceData implements TraceData {
   private static readonly MAX_BLOCKS = 1024;
-  private threadId: ThreadId;
   private trace: ArrayBuffer = new ArrayBuffer(0);
   private depth: number;
 
-  private constructor(threadId: ThreadId, depth: number, unique: boolean) {
-    this.threadId = threadId;
+  public constructor(depth: number) {
     this.depth = depth;
-    Stalker.follow(threadId, {
-      events: {
-        call: true,
-        ret: true,
-        block: !unique,
-        compile: unique,
-      },
-      onReceive: (events: ArrayBuffer) => {
-        const newBuffer = new Uint8Array(
-          this.trace.byteLength + events.byteLength,
-        );
-        newBuffer.set(new Uint8Array(this.trace), 0);
-        newBuffer.set(new Uint8Array(events), this.trace.byteLength);
-        this.trace = newBuffer.buffer as ArrayBuffer;
-      },
-    });
   }
 
-  public static create(
-    threadId: ThreadId,
-    depth: number,
-    unique: boolean = false,
-  ): BlockTrace {
-    if (Traces.has(threadId)) {
-      throw new Error(`trace already exists for threadId: ${threadId}`);
-    }
-
-    const trace = new BlockTrace(threadId, depth, unique);
-    Traces.set(threadId, trace);
-    return trace;
+  public append(events: ArrayBuffer) {
+    const newBuffer = new Uint8Array(this.trace.byteLength + events.byteLength);
+    newBuffer.set(new Uint8Array(this.trace), 0);
+    newBuffer.set(new Uint8Array(events), this.trace.byteLength);
+    this.trace = newBuffer.buffer as ArrayBuffer;
   }
 
   public display() {
@@ -55,15 +30,15 @@ export class BlockTrace implements Trace {
       switch (e.length) {
         case 3: {
           const [kind, start, _end] = e;
-          if (currentDepth >= this.depth) break;
+          if (!first && currentDepth >= this.depth) break;
           if (kind !== 'block' && kind !== 'compile') break;
           const name = Traces.getAddressString(start as NativePointer);
           if (name === null) break;
           if (first) {
             currentDepth = 0;
             first = false;
           }
-          if (numOutput >= BlockTrace.MAX_BLOCKS) {
+          if (numOutput >= BlockTraceData.MAX_BLOCKS) {
             Output.writeln(Output.red(`TRACE TRUNCATED`));
             return;
           }
@@ -89,9 +64,50 @@ export class BlockTrace implements Trace {
       }
     }
   }
+}
+
+export class BlockTrace implements Trace {
+  private threadId: ThreadId;
+  private trace: BlockTraceData;
+
+  private constructor(threadId: ThreadId, depth: number, unique: boolean) {
+    this.threadId = threadId;
+    this.trace = new BlockTraceData(depth);
+    Stalker.follow(threadId, {
+      events: {
+        call: true,
+        ret: true,
+        block: !unique,
+        compile: unique,
+      },
+      onReceive: (events: ArrayBuffer) => {
+        if (this.trace !== null) {
+          this.trace.append(events);
+        }
+      },
+    });
+  }
+
+  public static create(
+    threadId: ThreadId,
+    depth: number,
+    unique: boolean = false,
+  ): BlockTrace {
+    if (Traces.has(threadId)) {
+      throw new Error(`trace already exists for threadId: ${threadId}`);
+    }
+
+    const trace = new BlockTrace(threadId, depth, unique);
+    Traces.set(threadId, trace);
+    return trace;
+  }
 
   public stop() {
     Stalker.unfollow(this.threadId);
     Stalker.flush();
   }
+
+  public data(): BlockTraceData {
+    return this.trace;
+  }
 }
diff --git a/src/traces/call.ts b/src/traces/call.ts
@@ -1,39 +1,20 @@
 import { Output } from '../io/output.js';
-import { Trace, Traces } from './trace.js';
+import { Trace, TraceData, Traces } from './trace.js';
 
-export class CallTrace implements Trace {
+class CallTraceData implements TraceData {
   private static readonly MAX_CALLS = 1024;
-  private threadId: ThreadId;
   private trace: ArrayBuffer = new ArrayBuffer(0);
   private depth: number;
 
-  private constructor(threadId: ThreadId, depth: number) {
-    this.threadId = threadId;
+  public constructor(depth: number) {
     this.depth = depth;
-    Stalker.follow(threadId, {
-      events: {
-        call: true,
-        ret: true,
-      },
-      onReceive: (events: ArrayBuffer) => {
-        const newBuffer = new Uint8Array(
-          this.trace.byteLength + events.byteLength,
-        );
-        newBuffer.set(new Uint8Array(this.trace), 0);
-        newBuffer.set(new Uint8Array(events), this.trace.byteLength);
-        this.trace = newBuffer.buffer as ArrayBuffer;
-      },
-    });
   }
 
-  public static create(threadId: ThreadId, depth: number): CallTrace {
-    if (Traces.has(threadId)) {
-      throw new Error(`trace already exists for threadId: ${threadId}`);
-    }
-
-    const trace = new CallTrace(threadId, depth);
-    Traces.set(threadId, trace);
-    return trace;
+  public append(events: ArrayBuffer) {
+    const newBuffer = new Uint8Array(this.trace.byteLength + events.byteLength);
+    newBuffer.set(new Uint8Array(this.trace), 0);
+    newBuffer.set(new Uint8Array(events), this.trace.byteLength);
+    this.trace = newBuffer.buffer as ArrayBuffer;
   }
 
   public display() {
@@ -49,7 +30,7 @@ export class CallTrace implements Trace {
       const [kind, from, to, _depth] = e;
       if (kind === 'call') {
         currentDepth += 1;
-        if (currentDepth >= this.depth) continue;
+        if (!first && currentDepth >= this.depth) continue;
         const toName = Traces.getAddressString(to as NativePointer);
         if (toName === null) continue;
 
@@ -61,7 +42,7 @@ export class CallTrace implements Trace {
           currentDepth = 1;
           first = false;
         }
-        if (numOutput >= CallTrace.MAX_CALLS) {
+        if (numOutput >= CallTraceData.MAX_CALLS) {
           Output.writeln(Output.red(`TRACE TRUNCATED`));
           return;
         }
@@ -76,9 +57,44 @@ export class CallTrace implements Trace {
       }
     }
   }
+}
+
+export class CallTrace implements Trace {
+  private threadId: ThreadId;
+  private trace: CallTraceData;
+
+  private constructor(threadId: ThreadId, depth: number) {
+    this.threadId = threadId;
+    this.trace = new CallTraceData(depth);
+    Stalker.follow(threadId, {
+      events: {
+        call: true,
+        ret: true,
+      },
+      onReceive: (events: ArrayBuffer) => {
+        if (this.trace !== null) {
+          this.trace.append(events);
+        }
+      },
+    });
+  }
+
+  public static create(threadId: ThreadId, depth: number): CallTrace {
+    if (Traces.has(threadId)) {
+      throw new Error(`trace already exists for threadId: ${threadId}`);
+    }
+
+    const trace = new CallTrace(threadId, depth);
+    Traces.set(threadId, trace);
+    return trace;
+  }
 
   public stop() {
     Stalker.unfollow(this.threadId);
     Stalker.flush();
   }
+
+  public data(): CallTraceData {
+    return this.trace;
+  }
 }
diff --git a/src/traces/coverage/trace.ts b/src/traces/coverage/trace.ts
@@ -1,18 +1,33 @@
 import { Output } from '../../io/output.js';
-import { Trace, Traces } from '../trace.js';
+import { Trace, TraceData, Traces } from '../trace.js';
 import { Coverage, CoverageSession } from './coverage.js';
 
-export class CoverageTrace implements Trace {
-  private threadId: ThreadId;
+class CoverageTraceData implements TraceData {
   private filename: string;
+  private threadId: ThreadId;
+
+  public constructor(threadId: ThreadId, filename: string) {
+    this.threadId = threadId;
+    this.filename = filename;
+  }
+
+  public display() {
+    Output.writeln(
+      `Wrote coverage for thread: ${Output.yellow(this.threadId.toString())} to: ${Output.green(this.filename)}`,
+    );
+  }
+}
+export class CoverageTrace implements Trace {
   private file: File;
   private coverage: CoverageSession;
   private stopped: boolean = false;
+  private trace: CoverageTraceData;
 
   private constructor(threadId: ThreadId) {
-    this.threadId = threadId;
-    this.filename = CoverageTrace.getRandomFileName();
-    this.file = new File(this.filename, 'wb+');
+    const filename = CoverageTrace.getRandomFileName();
+    this.trace = new CoverageTraceData(threadId, filename);
+
+    this.file = new File(filename, 'wb+');
     this.coverage = Coverage.start({
       moduleFilter: m => Coverage.allModules(m),
       onCoverage: coverageData => {
@@ -32,12 +47,6 @@ export class CoverageTrace implements Trace {
     return trace;
   }
 
-  public display() {
-    Output.writeln(
-      `Wrote coverage for thread: ${Output.yellow(this.threadId.toString())} to: ${Output.green(this.filename)}`,
-    );
-  }
-
   public stop() {
     if (this.stopped) return;
     this.coverage.stop();
@@ -61,4 +70,8 @@ export class CoverageTrace implements Trace {
     const filename = `/tmp/${rand}.cov`;
     return filename;
   }
+
+  public data(): CoverageTraceData {
+    return this.trace;
+  }
 }