[WIP] tracer: Add support for Swift tracing

Co-authored-by: Håvard Sørbø <[email protected]>

agents/tracer/agent.ts

@@ -10,6 +10,7 @@ class Agent {
 
     private cachedModuleResolver: ApiResolver | null = null;
     private cachedObjcResolver: ApiResolver | null = null;
+    private cachedSwiftResolver: ApiResolver | null = null;
 
     init(stage: Stage, parameters: TraceParameters, initScripts: InitScript[], spec: TraceSpec) {
         const g = global as any as TraceScriptGlobals;
@@ -88,6 +89,13 @@ class Agent {
                         this.excludeObjCMethod(pattern, plan);
                     }
                     break;
+                case "swift-method":
+                    if (operation === "include") {
+                        this.includeSwiftMethod(pattern, plan);
+                    } else {
+                        this.excludeSwiftMethod(pattern, plan);
+                    }
+                    break;
                 case "java-method":
                     javaEntries.push([operation, pattern]);
                     break;
@@ -146,9 +154,21 @@ class Agent {
     private async traceNativeTargets(targets: NativeTargets) {
         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()) {
-            const entries = (type === "objc") ? objcGroups : cGroups;
+            let entries: Map<string, NativeItem[]>;
+            switch (type) {
+                case "c":
+                    entries = cGroups;
+                    break;
+                case "objc":
+                    entries = objcGroups;
+                    break;
+                case "swift":
+                    entries = swiftGroups;
+                    break;
+            }
 
             let group = entries.get(scope);
             if (group === undefined) {
@@ -161,11 +181,12 @@ class Agent {
 
         return await Promise.all([
             this.traceNativeEntries("c", cGroups),
-            this.traceNativeEntries("objc", objcGroups)
+            this.traceNativeEntries("objc", objcGroups),
+            this.traceNativeEntries("swift", swiftGroups),
         ]);
     }
 
-    private async traceNativeEntries(flavor: "c" | "objc", groups: NativeTargetScopes) {
+    private async traceNativeEntries(flavor: "c" | "objc" | "swift", groups: NativeTargetScopes) {
         if (groups.size === 0) {
             return;
         }
@@ -425,6 +446,20 @@ class Agent {
         }
     }
 
+    private includeSwiftMethod(pattern: string, plan: TracePlan) {
+        const { native } = plan;
+        for (const m of this.getSwiftResolver().enumerateMatches(pattern)) {
+            native.set(m.address.toString(), swiftMethodTargetFromMatch(m));
+        }
+    }
+
+    private excludeSwiftMethod(pattern: string, plan: TracePlan) {
+        const { native } = plan;
+        for (const m of this.getSwiftResolver().enumerateMatches(pattern)) {
+            native.delete(m.address.toString());
+        }
+    }
+
     private includeJavaMethod(pattern: string, plan: TracePlan) {
         const existingGroups = plan.java;
 
@@ -561,6 +596,19 @@ class Agent {
         }
         return resolver;
     }
+
+    private getSwiftResolver(): ApiResolver {
+        let resolver = this.cachedSwiftResolver;
+        if (resolver === null) {
+            try {
+                resolver = new ApiResolver("swift" as ApiResolverType); // FIXME: Update typings.
+            } catch (e: any) {
+                throw new Error("Swift runtime is not available");
+            }
+            this.cachedSwiftResolver = resolver;
+        }
+        return resolver;
+    }
 }
 
 async function getHandlers(request: HandlerRequest): Promise<HandlerResponse> {
@@ -646,6 +694,12 @@ function objcMethodTargetFromMatch(m: ApiResolverMatch): NativeTarget {
     return ["objc", className, [methodName, name]];
 }
 
+function swiftMethodTargetFromMatch(m: ApiResolverMatch): NativeTarget {
+    const { name } = m;
+    const [modulePath, methodName] = name.split("!", 2);
+    return ["swift", modulePath, methodName];
+}
+
 function debugSymbolTargetFromAddress(address: NativePointer): NativeTarget {
     const symbol = DebugSymbol.fromAddress(address);
     return ["c", symbol.moduleName ?? "", symbol.name!];
@@ -739,6 +793,7 @@ type TraceSpecScope =
     | "relative-function"
     | "imports"
     | "objc-method"
+    | "swift-method"
     | "java-method"
     | "debug-symbol"
     ;
@@ -749,12 +804,12 @@ interface TracePlan {
     java: JavaTargetGroup[];
 }
 
-type TargetType = "c" | "objc" | "java";
+type TargetType = "c" | "objc" | "swift" | "java";
 type ScopeName = string;
 type MemberName = string | [string, string]
 
 type NativeTargets = Map<NativeId, NativeTarget>;
-type NativeTarget = ["c" | "objc", ScopeName, MemberName];
+type NativeTarget = ["c" | "objc" | "swift", ScopeName, MemberName];
 type NativeTargetScopes = Map<ScopeName, NativeItem[]>;
 type NativeItem = [MemberName, NativePointer];
 type NativeId = string;

frida_tools/tracer.py

@@ -66,6 +66,20 @@ def _add_options(self, parser: argparse.ArgumentParser) -> None:
                 metavar="OBJC_METHOD",
                 type=pb.exclude_objc_method,
             )
+            parser.add_argument(
+                "-y",
+                "--include-swift-method",
+                help="include SWIFT_METHOD",
+                metavar="SWIFT_METHOD",
+                type=pb.include_swift_method,
+            )
+            parser.add_argument(
+                "-Y",
+                "--exclude-swift-method",
+                help="exclude SWIFT_METHOD",
+                metavar="SWIFT_METHOD",
+                type=pb.exclude_swift_method,
+            )
             parser.add_argument(
                 "-j",
                 "--include-java-method",
@@ -278,6 +292,16 @@ def exclude_objc_method(self, *function_name_globs: str) -> "TracerProfileBuilde
             self._spec.append(("exclude", "objc-method", f))
         return self
 
+    def include_swift_method(self, *function_name_globs: str) -> "TracerProfileBuilder":
+        for f in function_name_globs:
+            self._spec.append(("include", "swift-method", f))
+        return self
+
+    def exclude_swift_method(self, *function_name_globs: str) -> "TracerProfileBuilder":
+        for f in function_name_globs:
+            self._spec.append(("exclude", "swift-method", f))
+        return self
+
     def include_java_method(self, *function_name_globs: str) -> "TracerProfileBuilder":
         for f in function_name_globs:
             self._spec.append(("include", "java-method", f))
@@ -499,6 +523,12 @@ def objc_arg(m):
             log_str = "`" + re.sub(r":", objc_arg, target.display_name) + "`"
             if log_str.endswith("} ]`"):
                 log_str = log_str[:-3] + "]`"
+        elif target.flavor == "swift":
+            if decorate:
+                module_string = " [%s]" % os.path.basename(target.scope)
+            else:
+                module_string = ""
+            log_str = "'%(name)s()%(module_string)s'" % {"name": target.name, "module_string": module_string}
         else:
             for man_section in (2, 3):
                 args = []