Skip to content

Commit

Permalink
Wire up some more
Browse files Browse the repository at this point in the history
  • Loading branch information
@oleavr
oleavr committed Sep 25, 2024
1 parent 46d401b commit eebbe13
Show file tree
Hide file tree
Showing 5 changed files with 142 additions and 45 deletions.
142 changes: 104 additions & 38 deletions agents/tracer/agent.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,9 +57,14 @@ class Agent {
throw new Error("invalid target ID");
}

const newHandler = this.parseHandler(name, script);
handler[0] = newHandler[0];
handler[1] = newHandler[1];
if (handler.length === 3) {
const newHandler = this.parseFunctionHandler(name, script);
handler[0] = newHandler[0];
handler[1] = newHandler[1];
} else {
const newHandler = this.parseInstructionHandler(name, script);
handler[0] = newHandler[0];
}
}

updateHandlerConfig(id: TraceTargetId, config: HandlerConfig) {
Expand Down Expand Up @@ -138,10 +143,9 @@ class Agent {
const croppedMethods = new Map([[methodName, methodNameOrSignature]]);
const croppedClass: JavaTargetClass = { methods: croppedMethods };
const croppedGroup: JavaTargetGroup = { loader: group.loader, classes: new Map([[className, croppedClass]]) };
return {
native: new Map(),
java: [croppedGroup],
};
const croppedPlan = new TracePlan();
croppedPlan.java.push(croppedGroup);
return croppedPlan;
}
candidateId--;
}
Expand All @@ -151,10 +155,9 @@ class Agent {
candidateId = 1;
for (const [k, v] of plan.native.entries()) {
if (candidateId === id) {
return {
native: new Map([[k, v]]),
java: [],
};
const croppedPlan = new TracePlan();
croppedPlan.native.set(k, v);
return croppedPlan;
}
candidateId++;
}
Expand Down Expand Up @@ -186,10 +189,7 @@ class Agent {

private async createPlan(spec: TraceSpec,
onJavaReady: (plan: TracePlan) => Promise<void> = async () => {}): Promise<TracePlanRequest> {
const plan: TracePlan = {
native: new Map<NativeId, NativeTarget>(),
java: []
};
const plan = new TracePlan();

const javaEntries: [TraceSpecOperation, TraceSpecPattern][] = [];
for (const [operation, scope, pattern] of spec) {
Expand All @@ -213,6 +213,11 @@ class Agent {
this.includeRelativeFunction(pattern, plan);
}
break;
case "absolute-instruction":
if (operation === "include") {
this.includeAbsoluteInstruction(ptr(pattern), plan);
}
break;
case "imports":
if (operation === "include") {
this.includeImports(pattern, plan);
Expand Down Expand Up @@ -283,13 +288,17 @@ class Agent {
}

private async traceNativeTargets(targets: NativeTargets): Promise<TraceTargetId[]> {
const insnGroups = new Map<string, NativeItem[]>();
const cGroups = new Map<string, NativeItem[]>();
const objcGroups = new Map<string, NativeItem[]>();
const swiftGroups = new Map<string, NativeItem[]>();

for (const [id, [type, scope, name]] of targets.entries()) {
let entries: Map<string, NativeItem[]>;
switch (type) {
case "insn":
entries = insnGroups;
break;
case "c":
entries = cGroups;
break;
Expand All @@ -311,6 +320,7 @@ class Agent {
}

const [cIds, objcIds, swiftIds] = await Promise.all([
this.traceNativeEntries("insn", insnGroups),
this.traceNativeEntries("c", cGroups),
this.traceNativeEntries("objc", objcGroups),
this.traceNativeEntries("swift", swiftGroups),
Expand All @@ -319,7 +329,7 @@ class Agent {
return [...cIds, ...objcIds, ...swiftIds];
}

private async traceNativeEntries(flavor: "c" | "objc" | "swift", groups: NativeTargetScopes): Promise<TraceTargetId[]> {
private async traceNativeEntries(flavor: NativeTargetType, groups: NativeTargetScopes): Promise<TraceTargetId[]> {
if (groups.size === 0) {
return [];
}
Expand All @@ -336,7 +346,7 @@ class Agent {
scopes.push({
name,
members: items.map(item => item[0]),
addresses: items.map(item => item[1].toString()),
addresses: items.map(item => item[1].toString())
});
this.nextId += items.length;
}
Expand All @@ -345,16 +355,21 @@ class Agent {

const ids: TraceTargetId[] = [];
let offset = 0;
const isInstruction = flavor === "insn";
for (const items of groups.values()) {
for (const [name, address] of items) {
const id = baseId + offset;
const displayName = (typeof name === "string") ? name : name[1];

const handler = this.parseHandler(displayName, scripts[offset]);
const handler = isInstruction
? this.parseInstructionHandler(displayName, scripts[offset])
: this.parseFunctionHandler(displayName, scripts[offset]);
this.handlers.set(id, handler);

try {
Interceptor.attach(address, this.makeNativeListenerCallbacks(id, handler));
Interceptor.attach(address, isInstruction
? this.makeNativeInstructionListener(id, handler as TraceInstructionHandler)
: this.makeNativeFunctionListener(id, handler as TraceFunctionHandler));
} catch (e: any) {
send({
type: "agent:warning",
Expand Down Expand Up @@ -406,7 +421,7 @@ class Agent {
for (const [bareName, fullName] of methods.entries()) {
const id = baseId + offset;

const handler = this.parseHandler(fullName, scripts[offset]);
const handler = this.parseFunctionHandler(fullName, scripts[offset]);
this.handlers.set(id, handler);

const dispatcher: Java.MethodDispatcher = C[bareName];
Expand All @@ -425,7 +440,7 @@ class Agent {
});
}

private makeNativeListenerCallbacks(id: TraceTargetId, handler: TraceHandler): InvocationListenerCallbacks {
private makeNativeFunctionListener(id: TraceTargetId, handler: TraceFunctionHandler): InvocationListenerCallbacks {
const agent = this;

return {
Expand All @@ -440,15 +455,24 @@ class Agent {
};
}

private makeJavaMethodWrapper(id: TraceTargetId, method: Java.Method, handler: TraceHandler): Java.MethodImplementation {
private makeNativeInstructionListener(id: TraceTargetId, handler: TraceInstructionHandler): InstructionProbeCallback {
const agent = this;

return function (args) {
const [onHit, config] = handler;
agent.invokeNativeHandler(id, onHit, config, this, args, "|");
};
}

private makeJavaMethodWrapper(id: TraceTargetId, method: Java.Method, handler: TraceFunctionHandler): Java.MethodImplementation {
const agent = this;

return function (...args: any[]) {
return agent.handleJavaInvocation(id, method, handler, this, args);
};
}

private handleJavaInvocation(id: TraceTargetId, method: Java.Method, handler: TraceHandler, instance: Java.Wrapper, args: any[]): any {
private handleJavaInvocation(id: TraceTargetId, method: Java.Method, handler: TraceFunctionHandler, instance: Java.Wrapper, args: any[]): any {
const [onEnter, onLeave, config] = handler;

this.invokeJavaHandler(id, onEnter, config, instance, args, ">");
Expand All @@ -460,8 +484,8 @@ class Agent {
return (replacementRetval !== undefined) ? replacementRetval : retval;
}

private invokeNativeHandler(id: TraceTargetId, callback: TraceEnterHandler | TraceLeaveHandler, config: HandlerConfig,
context: InvocationContext, param: any, cutPoint: CutPoint) {
private invokeNativeHandler(id: TraceTargetId, callback: TraceEnterHandler | TraceLeaveHandler | TraceProbeHandler,
config: HandlerConfig, context: InvocationContext, param: any, cutPoint: CutPoint) {
const timestamp = Date.now() - this.started;
const threadId = context.threadId;
const depth = this.updateDepth(threadId, cutPoint);
Expand Down Expand Up @@ -503,7 +527,7 @@ class Agent {
let depth = depthEntries.get(threadId) ?? 0;
if (cutPoint === ">") {
depthEntries.set(threadId, depth + 1);
} else {
} else if (cutPoint === "<") {
depth--;
if (depth !== 0) {
depthEntries.set(threadId, depth);
Expand All @@ -515,20 +539,37 @@ class Agent {
return depth;
}

private parseHandler(name: string, script: string): TraceHandler {
const id = `/handlers/${name}.js`;
private parseFunctionHandler(name: string, script: string): TraceFunctionHandler {
try {
const h = Script.evaluate(id, script);
const h = this.parseHandlerScript(name, script);
return [h.onEnter ?? noop, h.onLeave ?? noop, makeDefaultHandlerConfig()];
} catch (e: any) {
send({
type: "agent:warning",
message: `${id}: ${e.message}`
message: `${name}: ${e.message}`
});
return [noop, noop, makeDefaultHandlerConfig()];
}
}

private parseInstructionHandler(name: string, script: string): TraceInstructionHandler {
try {
const onHit = this.parseHandlerScript(name, script);
return [onHit, makeDefaultHandlerConfig()];
} catch (e: any) {
send({
type: "agent:warning",
message: `${name}: ${e.message}`
});
return [noop, makeDefaultHandlerConfig()];
}
}

private parseHandlerScript(name: string, script: string): any {
const id = `/handlers/${name}.js`;
return Script.evaluate(id, script);
}

private includeModule(pattern: string, plan: TracePlan) {
const { native } = plan;
for (const m of this.getModuleResolver().enumerateMatches(`exports:${pattern}!*`)) {
Expand Down Expand Up @@ -565,6 +606,15 @@ class Agent {
plan.native.set(address.toString(), ["c", e.module, `sub_${e.offset.toString(16)}`]);
}

private includeAbsoluteInstruction(address: NativePointer, plan: TracePlan) {
const module = plan.modules.find(address);
if (module !== null) {
plan.native.set(address.toString(), ["insn", module.name, `insn_${address.sub(module.base).toString(16)}`]);
} else {
plan.native.set(address.toString(), ["insn", "", `insn_${address.toString(16)}`]);
}
}

private includeImports(pattern: string, plan: TracePlan) {
let matches: ApiResolverMatch[];
if (pattern === null) {
Expand Down Expand Up @@ -946,6 +996,7 @@ type TraceSpecScope =
| "module"
| "function"
| "relative-function"
| "absolute-instruction"
| "imports"
| "objc-method"
| "swift-func"
Expand All @@ -959,19 +1010,31 @@ interface TracePlanRequest {
ready: Promise<void>;
}

interface TracePlan {
native: NativeTargets;
java: JavaTargetGroup[];
class TracePlan {
native: NativeTargets = new Map<NativeId, NativeTarget>();
java: JavaTargetGroup[] = [];

#cachedModules: ModuleMap | null = null;

get modules(): ModuleMap {
let modules = this.#cachedModules;
if (modules === null) {
modules = new ModuleMap();
this.#cachedModules = modules;
}
return modules;
}
}

type TargetType = "c" | "objc" | "swift" | "java";
type TargetType = NativeTargetType | "java";
type ScopeName = string;
type MemberName = string | [string, string];

type NativeTargetType = "insn" | "c" | "objc" | "swift";
type NativeTargets = Map<NativeId, NativeTarget>;
type NativeTarget = ["c" | "objc" | "swift", ScopeName, MemberName];
type NativeTarget = [type: NativeTargetType, scope: ScopeName, name: MemberName];
type NativeTargetScopes = Map<ScopeName, NativeItem[]>;
type NativeItem = [MemberName, NativePointer];
type NativeItem = [name: MemberName, address: NativePointer];
type NativeId = string;

interface JavaTargetGroup {
Expand Down Expand Up @@ -1016,11 +1079,14 @@ type Caller = string | null;
type Backtrace = string[] | null;
type Message = string;

type TraceHandler = [onEnter: TraceEnterHandler, onLeave: TraceLeaveHandler, config: HandlerConfig];
type TraceHandler = TraceFunctionHandler | TraceInstructionHandler;
type TraceFunctionHandler = [onEnter: TraceEnterHandler, onLeave: TraceLeaveHandler, config: HandlerConfig];
type TraceInstructionHandler = [onHit: TraceProbeHandler, config: HandlerConfig];
type TraceEnterHandler = (log: LogHandler, args: any[], state: TraceState) => void;
type TraceLeaveHandler = (log: LogHandler, retval: any, state: TraceState) => any;
type TraceProbeHandler = (log: LogHandler, args: any[], state: TraceState) => void;

type CutPoint = ">" | "<";
type CutPoint = ">" | "|" | "<";

type LogHandler = (...message: string[]) => void;

Expand Down
8 changes: 5 additions & 3 deletions apps/tracer/src/App.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,8 @@ export default function App() {
stageItems,
stagedItems,
commitItems,

addInstructionHook,
} = useModel();
const captureBacktracesSwitchRef = useRef<HTMLInputElement>(null);
const [selectedTabId, setSelectedTabId] = useState("events");
Expand Down Expand Up @@ -78,13 +80,13 @@ export default function App() {
);

const disassemblyView = (
<DisassemblyView target={disassemblyTarget} />
<DisassemblyView target={disassemblyTarget} onAddInstructionHook={addInstructionHook} />
);

return (
<>
<Resplit.Root className="app-content" direction="vertical">
<Resplit.Pane className="top-area" order={0} initialSize="0.7fr">
<Resplit.Pane className="top-area" order={0} initialSize="0.5fr">
<section className="navigation-area">
<HandlerList
handlers={handlers}
Expand Down Expand Up @@ -151,7 +153,7 @@ export default function App() {
</section>
</Resplit.Pane>
<Resplit.Splitter className="app-splitter" order={1} size="5px" />
<Resplit.Pane className="bottom-area" order={2} initialSize="0.3fr">
<Resplit.Pane className="bottom-area" order={2} initialSize="0.5fr">
<Tabs className="bottom-tabs" selectedTabId={selectedTabId} onChange={tabId => setSelectedTabId(tabId as string)} animate={false}>
<Tab id="events" title="Events" panel={eventView} panelClassName="bottom-tab-panel" />
<Tab id="disassembly" title="Disassembly" panel={disassemblyView} panelClassName="bottom-tab-panel" />
Expand Down
14 changes: 12 additions & 2 deletions apps/tracer/src/DisassemblyView.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react";

export interface DisassemblyViewProps {
target?: DisassemblyTarget;
onAddInstructionHook: AddInstructionHookHandler;
}

export type DisassemblyTarget = FunctionTarget | InstructionTarget;
Expand All @@ -20,7 +21,9 @@ export interface InstructionTarget {
address: string;
}

export default function DisassemblyView({ target }: DisassemblyViewProps = {}) {
export type AddInstructionHookHandler = (address: bigint) => void;

export default function DisassemblyView({ target, onAddInstructionHook }: DisassemblyViewProps) {
const [r2Output, setR2Output] = useState<string[]>([]);
const [isLoading, setIsLoading] = useState(false);
const highlightedAddressAnchorRef = useRef<HTMLAnchorElement | null>(null);
Expand Down Expand Up @@ -79,7 +82,14 @@ export default function DisassemblyView({ target }: DisassemblyViewProps = {}) {

const addressMenu = useMemo(() => (
<Menu>
<MenuItem text="Add instruction-level hook" icon="add" />
<MenuItem
text="Add instruction-level hook here"
icon="add"
onClick={() => {
const address = BigInt(highlightedAddressAnchorRef.current!.innerText);
onAddInstructionHook(address);
}}
/>
</Menu>
), [handleAddressMenuClose]);

Expand Down
Loading

0 comments on commit eebbe13

Please sign in to comment.