From d5a9811f0c926da5ba67ada131029cf344a25d2d Mon Sep 17 00:00:00 2001
From: osy <50960678+osy@users.noreply.github.com>
Date: Sun, 30 Jul 2023 22:59:50 -0700
Subject: [PATCH] process: hold token to avoid XPC service termination

When under memory pressure, XNU can decide to terminate the XPC helper
unless there is a pending response. We force it to have a pending response
by providing a NOP callback which will only be triggered after QEMU
terminates.

Fixes #5489
---
 QEMUHelper/QEMUHelper.m         | 26 ++++++++++++++++++++++++--
 QEMUHelper/QEMUHelperProtocol.h |  8 ++++++++
 Services/UTMProcess.m           |  4 ++++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/QEMUHelper/QEMUHelper.m b/QEMUHelper/QEMUHelper.m
index 14c2b862..9002ffe6 100644
--- a/QEMUHelper/QEMUHelper.m
+++ b/QEMUHelper/QEMUHelper.m
@@ -20,8 +20,9 @@
 
 @interface QEMUHelper ()
 
-@property NSMutableArray<NSURL *> *urls;
-@property NSTask *childTask;
+@property (nonatomic) NSMutableArray<NSURL *> *urls;
+@property (nonatomic, nullable) NSTask *childTask;
+@property (nonatomic, nullable) tokenCallback_t activeToken;
 
 @end
 
@@ -136,9 +137,11 @@
     task.terminationHandler = ^(NSTask *task) {
         _self.childTask = nil;
         [_self.connection.remoteObjectProxy processHasExited:task.terminationStatus message:nil];
+        [_self invalidateToken];
     };
     if (![task launchAndReturnError:&err]) {
         NSLog(@"Error starting QEMU: %@", err);
+        [self invalidateToken];
         completion(NO, err.localizedDescription);
     } else {
         self.childTask = task;
@@ -149,6 +152,25 @@
 - (void)terminate {
     [self.childTask terminate];
     self.childTask = nil;
+    [self invalidateToken];
+}
+
+- (void)assertActiveWithToken:(tokenCallback_t)token {
+    @synchronized (self) {
+        if (self.activeToken) {
+            self.activeToken(NO);
+        }
+        self.activeToken = token;
+    }
+}
+
+- (void)invalidateToken {
+    @synchronized (self) {
+        if (self.activeToken) {
+            self.activeToken(YES);
+        }
+        self.activeToken = nil;
+    }
 }
 
 @end
diff --git a/QEMUHelper/QEMUHelperProtocol.h b/QEMUHelper/QEMUHelperProtocol.h
index 55dfd389..87e51b31 100644
--- a/QEMUHelper/QEMUHelperProtocol.h
+++ b/QEMUHelper/QEMUHelperProtocol.h
@@ -16,6 +16,8 @@
 
 #import <Foundation/Foundation.h>
 
+typedef void (^tokenCallback_t)(BOOL);
+
 NS_ASSUME_NONNULL_BEGIN
 
 // The protocol that this service will vend as its API. This header file will also need to be visible to the process hosting the service.
@@ -29,6 +31,12 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)startQemu:(NSString *)binName standardOutput:(NSFileHandle *)standardOutput standardError:(NSFileHandle *)standardError libraryBookmark:(NSData *)libBookmark argv:(NSArray<NSString *> *)argv completion:(void(^)(BOOL,NSString *))completion;
 - (void)terminate;
 
+/// Helper holds on to the token to keep this XPC service active
+///
+/// If this is not called after `startQemu`, XNU may terminate this helper as idle.
+/// - Parameter token: Token to hold, result may be discarded.
+- (void)assertActiveWithToken:(tokenCallback_t)token;
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/Services/UTMProcess.m b/Services/UTMProcess.m
index ac594efa..ecf368c2 100644
--- a/Services/UTMProcess.m
+++ b/Services/UTMProcess.m
@@ -234,6 +234,10 @@ static int defaultEntry(UTMProcess *self, int argc, const char *argv[], const ch
     NSFileHandle *standardError = self.standardError.fileHandleForWriting;
     [_connection.remoteObjectProxy setEnvironment:self.environment];
     [_connection.remoteObjectProxy setCurrentDirectoryPath:self.currentDirectoryUrl.path];
+    // this is needed to prevent XNU from terminating an idle XPC helper
+    [_connection.remoteObjectProxy assertActiveWithToken:^(BOOL ignored) {
+        // do nothing
+    }];
     [[_connection remoteObjectProxyWithErrorHandler:^(NSError * _Nonnull error) {
         if (error.domain == NSCocoaErrorDomain && error.code == NSXPCConnectionInvalid) {
             // inhibit this error since we always see it on quit