fix(ios): find the agent again after the daemon exits (#562)

Author: mobile-kevin

agents/ios/server.h

@@ -5,6 +5,7 @@ typedef NSData * _Nullable (^RpcHandler)(NSData * _Nonnull body);
 @interface MobileServer : NSObject
 - (instancetype _Nonnull)initWithHandler:(RpcHandler _Nonnull)handler;
 @property (readonly) int port;
++ (int)activePort;  // port the app's server listens on, 0 if none; read by the debugger after the CLI lost it
 - (BOOL)bind;       // binds and listens; call once, returns immediately
 - (void)run;        // accept loop — blocks the calling thread
 @end

agents/ios/server.m

@@ -18,12 +18,16 @@ static BOOL send_all(int fd, NSData *data) {
     return YES;
 }
 
+static int activePort = 0;
+
 @implementation MobileServer {
     RpcHandler _handler;
     int _serverFd;
     int _port;
 }
 
++ (int)activePort { return activePort; }
+
 - (instancetype)initWithHandler:(RpcHandler)handler {
     self = [super init];
     _handler = [handler copy];
@@ -74,14 +78,18 @@ - (BOOL)bind {
         _serverFd = -1;
         return NO;
     }
+    activePort = _port;
     NSLog(@"[mobilecli] bound to 127.0.0.1:%d", _port);
     return YES;
 }
 
 - (void)run {
     while (1) {
         int clientFd = accept(_serverFd, NULL, NULL);
-        if (clientFd < 0) continue;
+        if (clientFd < 0) {
+            usleep(100000); // a failed listening socket must not spin a core in the app
+            continue;
+        }
         dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
             [self handleConnection:clientFd];
         });

devices/ios/machomap/exporttrie.go

@@ -37,6 +37,8 @@ func NewExportTrieResolver(c *RSP) *ExportTrieResolver {
 			"class_getInstanceVariable":     "/usr/lib/libobjc.A.dylib",
 			"ivar_getOffset":                "/usr/lib/libobjc.A.dylib",
 			"class_getMethodImplementation": "/usr/lib/libobjc.A.dylib",
+			"class_getClassMethod":          "/usr/lib/libobjc.A.dylib",
+			"method_getImplementation":      "/usr/lib/libobjc.A.dylib",
 		},
 		loadAddr: map[string]uint64{},
 	}

devices/ios/machomap/runtime.go

@@ -1,9 +1,18 @@
 package machomap
 
+import "errors"
+
 // Helpers for driving the target after a map: calling an exported function, and
-// checking whether a class is already registered (used to keep injection
+// asking a class that may already be registered (used to keep injection
 // idempotent — mapping the same image twice fails on duplicate class pairs).
 
+// ErrClassMethodNotFound means the class is registered but lacks the class method.
+var ErrClassMethodNotFound = errors.New("class method not found")
+
+var classMethodHelpers = []string{
+	"mmap", "objc_getClass", "sel_registerName", "class_getClassMethod", "method_getImplementation",
+}
+
 // CallFunction borrows a thread and calls fn(args...) once, returning x0. Suitable
 // for simple exported getters/setters; it uses the default return breakpoint, so
 // callers should not use it for re-entrant functions (Map sets up a trampoline
@@ -17,36 +26,63 @@ func CallFunction(c *RSP, fn uint64, args ...uint64) (uint64, error) {
 	return t.Call(fn, args...)
 }
 
-// ClassIsRegistered reports whether objc_getClass(name) is non-nil in the target.
-func ClassIsRegistered(c *RSP, res Resolver, name string) (bool, error) {
-	h, err := res.Resolve([]string{"mmap", "objc_getClass"})
+// CallClassMethod calls +[class selector] in the target and returns x0. The second
+// result is false when the class is not registered.
+func CallClassMethod(c *RSP, res Resolver, class, selector string) (uint64, bool, error) {
+	h, err := res.Resolve(classMethodHelpers)
 	if err != nil {
-		return false, err
+		return 0, false, err
 	}
 	t, err := NewTarget(c)
 	if err != nil {
-		return false, err
+		return 0, false, err
 	}
 	defer func() { _ = t.Restore() }()
 	scratch, err := t.Call(h["mmap"], 0, 0x1000, protRead|protWrite, mapAnon|mapPrivate, neg1, 0)
 	if err != nil {
-		return false, err
+		return 0, false, err
 	}
-	if err := c.WriteMem(scratch, append([]byte(name), 0), 0); err != nil {
-		return false, err
+	names := append(append([]byte(class), 0), append([]byte(selector), 0)...)
+	if err := c.WriteMem(scratch, names, 0); err != nil {
+		return 0, false, err
 	}
-	// objc_getClass can wait on a lock until every thread runs, and the default
+	// the objc runtime can wait on a lock until every thread runs, and the default
 	// return breakpoint sits in shared code that another thread could then hit
 	landingPage, err := t.Call(h["mmap"], 0, 0x1000, protRead|protExec, mapAnon|mapPrivate, neg1, 0)
 	if err != nil {
-		return false, err
+		return 0, false, err
 	}
 	if err := t.SetLanding(landingPage); err != nil {
-		return false, err
+		return 0, false, err
+	}
+	return callClassMethod(t, h, scratch, scratch+uint64(len(class))+1)
+}
+
+// callClassMethod looks the method up before calling it, because sending an
+// unknown selector would crash the app.
+func callClassMethod(t *Target, h map[string]uint64, className, selectorName uint64) (uint64, bool, error) {
+	cls, err := t.Call(h["objc_getClass"], className)
+	if err != nil || cls == 0 {
+		return 0, false, err
+	}
+	sel, err := t.Call(h["sel_registerName"], selectorName)
+	if err != nil {
+		return 0, false, err
+	}
+	method, err := t.Call(h["class_getClassMethod"], cls, sel)
+	if err != nil {
+		return 0, false, err
+	}
+	if method == 0 {
+		return 0, true, ErrClassMethodNotFound
+	}
+	implementation, err := t.Call(h["method_getImplementation"], method)
+	if err != nil {
+		return 0, false, err
 	}
-	cls, err := t.Call(h["objc_getClass"], scratch)
+	value, err := t.Call(implementation, cls, sel)
 	if err != nil {
-		return false, err
+		return 0, false, err
 	}
-	return cls != 0, nil
+	return value, true, nil
 }

devices/ios/machomap/runtime_test.go

@@ -0,0 +1,73 @@
+package machomap
+
+import (
+	"errors"
+	"fmt"
+	"slices"
+	"testing"
+)
+
+const (
+	fakeClass          = 0xc1a55
+	fakeSelector       = 0x5e1
+	fakeMethod         = 0x3e7
+	fakeImplementation = 0x1a9
+	fakePort           = 49152
+)
+
+var fakeObjcHelpers = map[string]uint64{
+	"objc_getClass":            0x4000,
+	"sel_registerName":         0x4100,
+	"class_getClassMethod":     0x4200,
+	"method_getImplementation": 0x4300,
+}
+
+func wasWritten(server *fakeDebugserver, register int, value uint64) bool {
+	return slices.Contains(server.packets, fmt.Sprintf("P%x=%s;thread:%x;", register, registerValue(value), fakeThread))
+}
+
+func wasCalled(server *fakeDebugserver, function uint64) bool {
+	return wasWritten(server, regPC, function)
+}
+
+func TestClassMethodReturnsTheValueOfItsImplementation(t *testing.T) {
+	server := &fakeDebugserver{returnValues: []uint64{fakeClass, fakeSelector, fakeMethod, fakeImplementation, fakePort}}
+	target := borrowedThreadOn(t, server)
+
+	value, registered, err := callClassMethod(target, fakeObjcHelpers, 1, 2)
+
+	if err != nil || !registered || value != fakePort {
+		t.Errorf("value = %d, registered = %v, err = %v; want %d, true and no error", value, registered, err, fakePort)
+	}
+	if !wasCalled(server, fakeImplementation) {
+		t.Error("the implementation of the class method was not called")
+	}
+	if !wasWritten(server, 0, fakeClass) || !wasWritten(server, 1, fakeSelector) {
+		t.Error("the class and selector were not passed as the first two arguments")
+	}
+}
+
+func TestClassMethodOfAnUnregisteredClassReportsNoClass(t *testing.T) {
+	server := &fakeDebugserver{returnValues: []uint64{0}}
+	target := borrowedThreadOn(t, server)
+
+	_, registered, err := callClassMethod(target, fakeObjcHelpers, 1, 2)
+
+	if err != nil || registered {
+		t.Errorf("registered = %v, err = %v; want false and no error", registered, err)
+	}
+}
+
+func TestMissingClassMethodIsNeverCalled(t *testing.T) {
+	server := &fakeDebugserver{returnValues: []uint64{fakeClass, fakeSelector, 0}}
+	target := borrowedThreadOn(t, server)
+
+	_, registered, err := callClassMethod(target, fakeObjcHelpers, 1, 2)
+
+	if !registered || !errors.Is(err, ErrClassMethodNotFound) {
+		t.Errorf("registered = %v, err = %v; want true and the missing method error", registered, err)
+	}
+	if wasCalled(server, fakeObjcHelpers["method_getImplementation"]) {
+		t.Error("the missing method was still looked up")
+	}
+}

devices/ios/machomap/target_test.go

@@ -29,6 +29,16 @@ type fakeDebugserver struct {
 	packets []string
 
 	refusesMemoryWrites bool
+	returnValues        []uint64 // one per call, in order; fakeReturned once they run out
+}
+
+func (f *fakeDebugserver) nextReturnValue() uint64 {
+	if len(f.returnValues) == 0 {
+		return fakeReturned
+	}
+	value := f.returnValues[0]
+	f.returnValues = f.returnValues[1:]
+	return value
 }
 
 func (f *fakeDebugserver) reply(payload string) {
@@ -94,7 +104,7 @@ func (f *fakeDebugserver) serve() {
 		case readPC:
 			f.reply(registerValue(f.pc))
 		case readX0:
-			f.reply(registerValue(fakeReturned))
+			f.reply(registerValue(f.nextReturnValue()))
 		default:
 			f.reply("OK")
 		}

devices/ios_device_agent.go

@@ -143,7 +143,7 @@ func (d *IOSDevice) ensureIOSDeviceAgentReady() (int, error) {
 		return agent.localPort, nil
 	}
 	if !ok {
-		devicePort, err := d.injectFreshAgent(activeApp.BundleID, activeApp.ProcessID)
+		devicePort, err := d.findOrInjectAgent(activeApp.BundleID, activeApp.ProcessID)
 		if err != nil {
 			return 0, err
 		}
@@ -153,9 +153,10 @@ func (d *IOSDevice) ensureIOSDeviceAgentReady() (int, error) {
 	return d.connectToAgent(agent)
 }
 
-// injectFreshAgent performs the slow path: map the agent into the app over the
-// debug proxy. Returns the device port the agent's server bound.
-func (d *IOSDevice) injectFreshAgent(bundleID string, pid int) (int, error) {
+// findOrInjectAgent performs the slow path: map the agent into the app over the
+// debug proxy, unless the app already has one from an earlier daemon. Returns the
+// device port the agent's server bound.
+func (d *IOSDevice) findOrInjectAgent(bundleID string, pid int) (int, error) {
 	utils.Verbose("getting enhanced device info")
 	device, err := d.getEnhancedDevice()
 	if err != nil {
@@ -167,12 +168,12 @@ func (d *IOSDevice) injectFreshAgent(bundleID string, pid int) (int, error) {
 	}
 	utils.Verbose("debug proxy port from RSD: %d", proxyPort)
 
-	utils.Verbose("mapping agent dylib into %s (pid %d)", bundleID, pid)
+	utils.Verbose("finding or mapping the agent in %s (pid %d)", bundleID, pid)
 	devicePort, err := injectAgentNative(device, proxyPort, pid)
 	if err != nil {
-		return 0, fmt.Errorf("inject agent: %w", err)
+		return 0, fmt.Errorf("set up agent: %w", err)
 	}
-	utils.Verbose("agent started on device port %d", devicePort)
+	utils.Verbose("agent is on device port %d", devicePort)
 	return devicePort, nil
 }
 
@@ -199,7 +200,7 @@ func (d *IOSDevice) connectToAgent(agent *injectedAgent) (int, error) {
 	for !isAgentReady(localPort) {
 		if time.Now().After(deadline) {
 			_ = pf.Stop()
-			return 0, fmt.Errorf("iOS device agent did not respond within 3s")
+			return 0, fmt.Errorf("iOS device agent did not respond within 3s; relaunch the app")
 		}
 		time.Sleep(100 * time.Millisecond)
 	}

devices/ios_device_agent_native.go

@@ -71,9 +71,24 @@ func probeError(err error) error {
 	if errors.Is(err, machomap.ErrUnexpectedStop) {
 		return fmt.Errorf("%w: %w", errAppBusyDuringProbe, err)
 	}
+	if errors.Is(err, machomap.ErrClassMethodNotFound) {
+		return fmt.Errorf("the agent in this app is from an older mobilecli; relaunch the app: %w", err)
+	}
 	return fmt.Errorf("probe existing agent: %w", err)
 }
 
+// existingAgentPort reads the port an already injected agent reported, where 0
+// means it has no server. The value comes from a C int in a 64-bit register, so
+// only the low bits are valid.
+func existingAgentPort(pid int, reported uint64) (int, error) {
+	port := int(uint32(reported)) // #nosec G115 -- the method returns a C int in a 64-bit register
+	if port == 0 {
+		return 0, fmt.Errorf("agent in pid %d has no server; relaunch the app", pid)
+	}
+	utils.Verbose("found existing agent on device port %d", port)
+	return port, nil
+}
+
 func injectAgentOnce(device goios.DeviceEntry, proxyPort, pid int) (int, error) {
 	conn, err := goios.ConnectTUNDevice(device.Address, proxyPort, device)
 	if err != nil {
@@ -127,14 +142,15 @@ func mapAgentOverConn(conn io.ReadWriter, pid int) (int, error) {
 	resolver := machomap.NewExportTrieResolver(rsp)
 
 	// Idempotency: never map the same image twice into one process — duplicate
-	// ObjC class registration fails and leaves the image half-wired. If the agent
-	// is already present (e.g. after a daemon restart while the app kept running),
-	// the in-process port cache is gone and the agent's ephemeral port cannot be
-	// recovered, so ask the caller to relaunch the app.
-	if already, err := machomap.ClassIsRegistered(rsp, resolver, "MobileServer"); err != nil {
+	// ObjC class registration fails and leaves the image half-wired. An agent that
+	// is already present (e.g. the daemon exited while the app kept running) still
+	// knows its port, so ask it.
+	activePort, registered, err := machomap.CallClassMethod(rsp, resolver, "MobileServer", "activePort")
+	if err != nil {
 		return 0, probeError(err)
-	} else if already {
-		return 0, fmt.Errorf("agent already injected in pid %d; relaunch the app to re-inject", pid)
+	}
+	if registered {
+		return existingAgentPort(pid, activePort)
 	}
 
 	res, err := machomap.Map(rsp, agents.IOSAgentDevDylib, resolver, machomap.Options{RunInit: true})

devices/ios_device_agent_native_test.go

@@ -204,6 +204,35 @@ func TestInjectionFailsAtOnceWhenTheAppIsNotDebuggable(t *testing.T) {
 	}
 }
 
+func TestExistingAgentPortIgnoresTheUpperHalfOfTheRegister(t *testing.T) {
+	port, err := existingAgentPort(1, 0xdeadbeef0000c000)
+
+	if err != nil || port != 0xc000 {
+		t.Errorf("port = %d, err = %v; want %d and no error", port, err, 0xc000)
+	}
+}
+
+func TestExistingAgentWithoutAServerAsksForARelaunch(t *testing.T) {
+	_, err := existingAgentPort(1, 0)
+
+	if err == nil || !strings.Contains(err.Error(), "relaunch the app") {
+		t.Errorf("err = %v, want one that asks for a relaunch", err)
+	}
+}
+
+func TestAgentFromAnOlderVersionAsksForARelaunchAndIsNotRetried(t *testing.T) {
+	calls := 0
+
+	_, err := retryWhileAppIsBusy(injectionThatFails(probeError(machomap.ErrClassMethodNotFound), -1, &calls))
+
+	if err == nil || !strings.Contains(err.Error(), "relaunch the app") {
+		t.Errorf("err = %v, want one that asks for a relaunch", err)
+	}
+	if calls != 1 {
+		t.Errorf("attached %d times, want 1", calls)
+	}
+}
+
 func TestEachDeviceHasItsOwnSetupLock(t *testing.T) {
 	firstLookup := deviceAgentLock("device-a")
 	secondLookup := deviceAgentLock("device-a")