Use runtime setting to enable ExtensionKit assertions

https://bugs.webkit.org/show_bug.cgi?id=264254 rdar://117997757 Reviewed by Brent Fulgham. Use runtime setting to enable ExtensionKit assertions. This setting is not enabled by default. * Source/WebKit/UIProcess/AuxiliaryProcessProxy.h: (WebKit::AuxiliaryProcessProxy::manageProcessesAsExtensions): * Source/WebKit/UIProcess/Cocoa/ProcessAssertionCocoa.mm: (WebKit::ProcessAssertion::ProcessAssertion): (WebKit::ProcessAssertion::~ProcessAssertion): * Source/WebKit/UIProcess/Launcher/cocoa/ProcessLauncherCocoa.mm: (WebKit::ProcessLauncher::launchProcess): * Source/WebKit/UIProcess/ProcessAssertion.h: * Source/WebKit/UIProcess/ProcessThrottler.cpp: (WebKit::ProcessThrottler::setThrottleState): Canonical link: https://commits.webkit.org/270318@main
WebKit · Nov 7, 2023 · 0df9746 · 0df9746
1 parent a48a890
commit 0df9746
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 30 deletions.
diff --git a/Source/WebKit/UIProcess/AuxiliaryProcessProxy.h b/Source/WebKit/UIProcess/AuxiliaryProcessProxy.h
@@ -184,6 +184,7 @@ class AuxiliaryProcessProxy
 #if USE(EXTENSIONKIT)
     RetainPtr<_SEExtensionProcess> extensionProcess() const;
     static void setManageProcessesAsExtensions(bool manageProcessesAsExtensions) { s_manageProcessesAsExtensions = manageProcessesAsExtensions; }
+    static bool manageProcessesAsExtensions() { return s_manageProcessesAsExtensions; }
 #endif
 
 protected:

diff --git a/Source/WebKit/UIProcess/Cocoa/ProcessAssertionCocoa.mm b/Source/WebKit/UIProcess/Cocoa/ProcessAssertionCocoa.mm
@@ -34,6 +34,7 @@
 #import "WebProcessPool.h"
 #import <wtf/HashMap.h>
 #import <wtf/RunLoop.h>
+#import <wtf/SoftLinking.h>
 #import <wtf/ThreadSafeWeakHashSet.h>
 #import <wtf/Vector.h>
 #import <wtf/WeakHashSet.h>
@@ -48,6 +49,9 @@
 
 #if USE(EXTENSIONKIT)
 #import "ExtensionKitSPI.h"
+
+SOFT_LINK_FRAMEWORK_OPTIONAL(ServiceExtensions);
+SOFT_LINK_CLASS_OPTIONAL(ServiceExtensions, _SECapabilities);
 #endif
 
 static WorkQueue& assertionsWorkQueue()
@@ -367,6 +371,25 @@ - (void)assertion:(RBSAssertion *)assertion didInvalidateWithError:(NSError *)er
     init(environmentIdentifier);
 }
 
+ProcessAssertion::ProcessAssertion(AuxiliaryProcessProxy& process, const String& reason, ProcessAssertionType assertionType)
+    : m_assertionType(assertionType)
+    , m_pid(process.processID())
+    , m_reason(reason)
+{
+#if USE(EXTENSIONKIT)
+    if (AuxiliaryProcessProxy::manageProcessesAsExtensions()) {
+        NSString *runningBoardAssertionName = runningBoardNameForAssertionType(m_assertionType);
+        NSString *runningBoardDomain = runningBoardDomainForAssertionType(m_assertionType);
+        m_capabilities = [get_SECapabilitiesClass() assertionWithDomain:runningBoardDomain name:runningBoardAssertionName];
+        m_process = process.extensionProcess();
+        if (m_capabilities)
+            return;
+        RELEASE_LOG(ProcessSuspension, "%p - ProcessAssertion() Failed to create capability %@", this, runningBoardAssertionName);
+    }
+#endif
+    init(process.environmentIdentifier());
+}
+
 void ProcessAssertion::init(const String& environmentIdentifier)
 {
     NSString *runningBoardAssertionName = runningBoardNameForAssertionType(m_assertionType);
@@ -398,25 +421,6 @@ - (void)assertion:(RBSAssertion *)assertion didInvalidateWithError:(NSError *)er
     };
 }
 
-ProcessAssertion::ProcessAssertion(AuxiliaryProcessProxy& process, const String& reason, ProcessAssertionType assertionType)
-    : m_assertionType(assertionType)
-    , m_pid(process.processID())
-    , m_reason(reason)
-{
-#if USE(EXTENSIONKIT_ASSERTIONS)
-    NSString *runningBoardAssertionName = runningBoardNameForAssertionType(assertionType);
-    NSString *runningBoardDomain = runningBoardDomainForAssertionType(assertionType);
-    auto capability = adoptNS([_SECapabilities assertionWithDomain:runningBoardDomain name:runningBoardAssertionName]);
-    auto extensionProcess = process.extensionProcess();
-    NSError* error = nil;
-    m_grant = [extensionProcess.get() grantCapabilities:capability.get() error:&error];
-    if (!m_grant)
-        RELEASE_LOG(ProcessSuspension, "%p - ProcessAssertion() Failed to grant capability with error %@", this, error);
-#else
-    init(process.environmentIdentifier());
-#endif
-}
-
 double ProcessAssertion::remainingRunTimeInSeconds(ProcessID pid)
 {
     RBSProcessIdentifier *processIdentifier = [RBSProcessIdentifier identifierWithPid:pid];
@@ -449,7 +453,17 @@ - (void)assertion:(RBSAssertion *)assertion didInvalidateWithError:(NSError *)er
 void ProcessAssertion::acquireSync()
 {
     RELEASE_LOG(ProcessSuspension, "%p - ProcessAssertion::acquireSync Trying to take RBS assertion '%{public}s' for process with PID=%d", this, m_reason.utf8().data(), m_pid);
-
+#if USE(EXTENSIONKIT)
+    if (m_process) {
+        NSError *error = nil;
+        m_grant = [m_process.get() grantCapabilities:m_capabilities.get() error:&error];
+        if (m_grant) {
+            RELEASE_LOG(ProcessSuspension, "%p - ProcessAssertion() Successfully granted capability", this);
+            return;
+        }
+        RELEASE_LOG(ProcessSuspension, "%p - ProcessAssertion() Failed to grant capability with error %@", this, error);
+    }
+#endif
     NSError *acquisitionError = nil;
     if (![m_rbsAssertion acquireWithError:&acquisitionError]) {
         RELEASE_LOG_ERROR(ProcessSuspension, "%p - ProcessAssertion::acquireSync Failed to acquire RBS assertion '%{public}s' for process with PID=%d, error: %{public}@", this, m_reason.utf8().data(), m_pid, acquisitionError);
@@ -473,7 +487,7 @@ - (void)assertion:(RBSAssertion *)assertion didInvalidateWithError:(NSError *)er
         [m_rbsAssertion invalidate];
     }
 
-#if USE(EXTENSIONKIT_ASSERTIONS)
+#if USE(EXTENSIONKIT)
     [m_grant invalidateWithError:nil];
 #endif
 }

diff --git a/Source/WebKit/UIProcess/Launcher/cocoa/ProcessLauncherCocoa.mm b/Source/WebKit/UIProcess/Launcher/cocoa/ProcessLauncherCocoa.mm
@@ -158,6 +158,7 @@ static void launchWithExtensionKit(ProcessLauncher& processLauncher, ProcessLaun
                 launcher->m_xpcConnection = adoptOSObject(xpc_connection_create(name, nullptr));
                 launcher->finishLaunchingProcess(name);
             });
+            return;
         }
         callOnMainRunLoop([weakProcessLauncher = weakProcessLauncher, name = name, process = RetainPtr<_SEExtensionProcess>(process)] {
             auto connection = [process makeLibXPCConnectionError:nil];

diff --git a/Source/WebKit/UIProcess/ProcessAssertion.h b/Source/WebKit/UIProcess/ProcessAssertion.h
@@ -42,8 +42,9 @@ OBJC_CLASS RBSAssertion;
 OBJC_CLASS WKRBSAssertionDelegate;
 #endif // USE(RUNNINGBOARD)
 
-#if USE(EXTENSIONKIT_ASSERTIONS)
+#if USE(EXTENSIONKIT)
 OBJC_CLASS _SEExtensionProcess;
+OBJC_CLASS _SECapabilities;
 OBJC_PROTOCOL(_SEGrant);
 #endif
 
@@ -67,7 +68,7 @@ class ProcessAssertion : public ThreadSafeRefCountedAndCanMakeThreadSafeWeakPtr<
     WTF_MAKE_FAST_ALLOCATED;
 public:
     enum class Mode : bool { Sync, Async };
-#if USE(EXTENSIONKIT_ASSERTIONS)
+#if USE(EXTENSIONKIT)
     static Ref<ProcessAssertion> create(RetainPtr<_SEExtensionProcess>, const String& reason, ProcessAssertionType, Mode = Mode::Async, const String& environmentIdentifier = emptyString(), CompletionHandler<void()>&& acquisisionHandler = nullptr);
 #endif
     static Ref<ProcessAssertion> create(ProcessID, const String& reason, ProcessAssertionType, Mode = Mode::Async, const String& environmentIdentifier = emptyString(), CompletionHandler<void()>&& acquisisionHandler = nullptr);
@@ -111,8 +112,10 @@ class ProcessAssertion : public ThreadSafeRefCountedAndCanMakeThreadSafeWeakPtr<
 #endif
     Function<void()> m_prepareForInvalidationHandler;
     Function<void()> m_invalidationHandler;
-#if USE(EXTENSIONKIT_ASSERTIONS)
+#if USE(EXTENSIONKIT)
+    RetainPtr<_SECapabilities> m_capabilities;
     RetainPtr<_SEGrant> m_grant;
+    RetainPtr<_SEExtensionProcess> m_process;
 #endif
 };
 

diff --git a/Source/WebKit/UIProcess/ProcessThrottler.cpp b/Source/WebKit/UIProcess/ProcessThrottler.cpp
@@ -256,13 +256,8 @@ void ProcessThrottler::setThrottleState(ProcessThrottleState newState)
                     weakThis->uiAssertionWillExpireImminently();
             });
             m_assertion = WTFMove(assertion);
-        } else {
-#if USE(EXTENSIONKIT_ASSERTIONS)
+        } else
             m_assertion = ProcessAssertion::create(*m_processProxy, assertionName(newType), newType, ProcessAssertion::Mode::Async, [previousAssertion = WTFMove(previousAssertion)] { });
-#else
-            m_assertion = ProcessAssertion::create(m_processProxy->processID(), assertionName(newType), newType, ProcessAssertion::Mode::Async, m_process.environmentIdentifier(), [previousAssertion = WTFMove(previousAssertion)] { });
-#endif
-        }
     }
     m_assertion->setInvalidationHandler([weakThis = WeakPtr { *this }] {
         if (weakThis)