REGRESSION (271259@main?): [ iOS Sonoma Debug arm64 ] TestWebKitAPI.W…

…KWebExtensionAPITabs.SendMessageWithPromiseReply timeout https://webkit.org/b/266518 rdar://problem/119740501 Reviewed by Brian Weinstein. Adopt EagerCallbackAggregator for the reply handler. This was added on the UIProcess side for the message sending in 271636@main. It can also be used here, with the help of an ObjC wrapper object that can keep the aggregator in scope for the required blocks. We can't use BlockPtr for these since the JSValue ObjC translation requires blocks to have a signature to translate them into JavaScript functions, which is only for compiled blocks. Without this, the aggregator falls out of scope and is released. Added additional tests to make sure async replies work in addition to sync replies. Also adopted optional chaining for some tests since an exception is thrown when response is null. * Source/WebKit/WebProcess/Extensions/API/Cocoa/WebExtensionAPIRuntimeCocoa.mm: (-[_WKReplyCallbackAggregator initWithAggregator:]): Added. (-[_WKReplyCallbackAggregator aggregator]): Added. (WebKit::WebExtensionContextProxy::internalDispatchRuntimeMessageEvent): * Tools/TestWebKitAPI/Tests/WebKitCocoa/WKWebExtensionAPIRuntime.mm: (TestWebKitAPI::TEST): * Tools/TestWebKitAPI/Tests/WebKitCocoa/WKWebExtensionAPITabs.mm: (TestWebKitAPI::TEST): Canonical link: https://commits.webkit.org/272236@main
WebKit · Dec 18, 2023 · fdd1ff7 · fdd1ff7
1 parent df351d0
commit fdd1ff7
Show file tree

Hide file tree

Showing 3 changed files with 318 additions and 40 deletions.
diff --git a/Source/WebKit/WebProcess/Extensions/API/Cocoa/WebExtensionAPIRuntimeCocoa.mm b/Source/WebKit/WebProcess/Extensions/API/Cocoa/WebExtensionAPIRuntimeCocoa.mm
@@ -45,6 +45,7 @@
 #import "WebProcess.h"
 #import <WebCore/SecurityOrigin.h>
 #import <wtf/BlockPtr.h>
+#import <wtf/CallbackAggregator.h>
 
 static NSString * const idKey = @"id";
 static NSString * const frameIdKey = @"frameId";
@@ -57,6 +58,41 @@
 
 namespace WebKit {
 
+using ReplyCallbackAggregator = EagerCallbackAggregator<void(id)>;
+
+}
+
+@interface _WKReplyCallbackAggregator : NSObject
+
+- (instancetype)initWithAggregator:(WebKit::ReplyCallbackAggregator&)aggregator;
+
+@property (nonatomic, readonly) WebKit::ReplyCallbackAggregator& aggregator;
+
+@end
+
+@implementation _WKReplyCallbackAggregator {
+    RefPtr<WebKit::ReplyCallbackAggregator> _aggregator;
+}
+
+- (instancetype)initWithAggregator:(WebKit::ReplyCallbackAggregator&)aggregator
+{
+    if (!(self = [super init]))
+        return nil;
+
+    _aggregator = &aggregator;
+
+    return self;
+}
+
+- (WebKit::ReplyCallbackAggregator&)aggregator
+{
+    return *_aggregator;
+}
+
+@end
+
+namespace WebKit {
+
 JSValue *WebExtensionAPIRuntimeBase::reportError(NSString *errorMessage, JSGlobalContextRef contextRef, Function<void()>&& handler)
 {
     ASSERT(errorMessage.length);
@@ -427,13 +463,7 @@
     id message = parseJSON(messageJSON, { JSONOptions::FragmentsAllowed });
     auto *senderInfo = toWebAPI(senderParameters);
 
-    __block bool anyListenerHandledMessage = false;
-    bool sentReply = false;
-
-    __block auto handleReply = [&sentReply, completionHandler = WTFMove(completionHandler)](id replyMessage) mutable {
-        if (sentReply)
-            return;
-
+    auto callbackAggregator = ReplyCallbackAggregator::create([completionHandler = WTFMove(completionHandler)](id replyMessage) mutable {
         NSString *replyMessageJSON;
         if (JSValue *value = dynamic_objc_cast<JSValue>(replyMessage))
             replyMessageJSON = value._toJSONString;
@@ -443,11 +473,14 @@
         if (replyMessageJSON.length > webExtensionMaxMessageLength)
             replyMessageJSON = nil;
 
-        sentReply = true;
         completionHandler(replyMessageJSON ? std::optional(String(replyMessageJSON)) : std::nullopt);
-    };
+    }, nil);
+
+    // This ObjC wrapper is need for the inner reply block, which is required to be a compiled block.
+    auto *callbackAggregatorWrapper = [[_WKReplyCallbackAggregator alloc] initWithAggregator:callbackAggregator];
 
-    enumerateFramesAndNamespaceObjects(makeBlockPtr(^(WebFrame& frame, WebExtensionAPINamespace& namespaceObject) {
+    bool anyListenerHandledMessage = false;
+    enumerateFramesAndNamespaceObjects([&, callbackAggregatorWrapper = RetainPtr { callbackAggregatorWrapper }](WebFrame& frame, WebExtensionAPINamespace& namespaceObject) {
         // Skip all frames that don't match if a target frame identifier is specified.
         if (frameIdentifier && !matchesFrame(frameIdentifier.value(), frame))
             return;
@@ -457,8 +490,11 @@
             return;
 
         for (auto& listener : namespaceObject.runtime().onMessage().listeners()) {
+            // Using BlockPtr for this call does not work, since JSValue needs a compiled block
+            // with a signature to translate the JS function arguments. Having the block capture
+            // callbackAggregatorWrapper ensures that callbackAggregator remains in scope.
             id returnValue = listener->call(message, senderInfo, ^(id replyMessage) {
-                handleReply(replyMessage);
+                callbackAggregatorWrapper.get().aggregator(replyMessage);
             });
 
             if (dynamic_objc_cast<NSNumber>(returnValue).boolValue) {
@@ -476,13 +512,13 @@
                 if (error)
                     return;
 
-                handleReply(replyMessage);
+                callbackAggregatorWrapper.get().aggregator(replyMessage);
             }];
         }
-    }).get(), world);
+    }, world);
 
-    if (!sentReply && !anyListenerHandledMessage)
-        handleReply(nil);
+    if (!anyListenerHandledMessage)
+        callbackAggregator.get()(nil);
 }
 
 void WebExtensionContextProxy::dispatchRuntimeMessageEvent(WebExtensionContentWorldType contentWorldType, const String& messageJSON, std::optional<WebExtensionFrameIdentifier> frameIdentifier, const WebExtensionMessageSenderParameters& senderParameters, CompletionHandler<void(std::optional<String> replyJSON)>&& completionHandler)

diff --git a/Tools/TestWebKitAPI/Tests/WebKitCocoa/WKWebExtensionAPIRuntime.mm b/Tools/TestWebKitAPI/Tests/WebKitCocoa/WKWebExtensionAPIRuntime.mm
@@ -394,7 +394,7 @@
         [NSString stringWithFormat:@"const expectedOrigin = '%@'", [urlString substringToIndex:urlString.length - 1]],
 
         @"browser.runtime.onMessage.addListener((message, sender, sendResponse) => {",
-        @"  browser.test.assertEq(message.content, 'Hello', 'Should receive the correct message from the content script')",
+        @"  browser.test.assertEq(message?.content, 'Hello', 'Should receive the correct message from the content script')",
 
         @"  browser.test.assertEq(typeof sender, 'object', 'sender should be an object')",
 
@@ -416,7 +416,65 @@
         @"(async () => {",
         @"  const response = await browser.runtime.sendMessage({ content: 'Hello' })",
 
-        @"  browser.test.assertEq(response.content, 'Received', 'Should get the response from the background script')",
+        @"  browser.test.assertEq(response?.content, 'Received', 'Should get the response from the background script')",
+
+        @"  browser.test.notifyPass()",
+        @"})()"
+    ]);
+
+    auto extension = adoptNS([[_WKWebExtension alloc] _initWithManifestDictionary:runtimeContentScriptManifest resources:@{ @"background.js": backgroundScript, @"content.js": contentScript }]);
+    auto manager = adoptNS([[TestWebExtensionManager alloc] initForExtension:extension.get()]);
+
+    [manager.get().context setPermissionStatus:_WKWebExtensionContextPermissionStatusGrantedExplicitly forURL:urlRequest.URL];
+
+    [manager loadAndRun];
+
+    EXPECT_NS_EQUAL(manager.get().yieldMessage, @"Load Tab");
+
+    [manager.get().defaultTab.mainWebView loadRequest:urlRequest];
+
+    [manager run];
+}
+
+TEST(WKWebExtensionAPIRuntime, SendMessageFromContentScriptWithAsyncReply)
+{
+    TestWebKitAPI::HTTPServer server({
+        { "/"_s, { { { "Content-Type"_s, "text/html"_s } }, ""_s } },
+    }, TestWebKitAPI::HTTPServer::Protocol::Http);
+
+    auto *urlRequest = server.requestWithLocalhost();
+    auto *urlString = urlRequest.URL.absoluteString;
+
+    auto *backgroundScript = Util::constructScript(@[
+        [NSString stringWithFormat:@"const expectedURL = '%@'", urlString],
+        [NSString stringWithFormat:@"const expectedOrigin = '%@'", [urlString substringToIndex:urlString.length - 1]],
+
+        @"browser.runtime.onMessage.addListener((message, sender, sendResponse) => {",
+        @"  browser.test.assertEq(message?.content, 'Hello', 'Should receive the correct message from the content script')",
+
+        @"  browser.test.assertEq(typeof sender, 'object', 'sender should be an object')",
+
+        @"  browser.test.assertEq(typeof sender.tab, 'object', 'sender.tab should be an object')",
+        @"  browser.test.assertEq(sender.tab.url, expectedURL, 'sender.tab.url should be the expected URL')",
+
+        @"  browser.test.assertEq(sender.url, expectedURL, 'sender.url should be the expected URL')",
+        @"  browser.test.assertEq(sender.origin, expectedOrigin, 'sender.origin should be the expected origin')",
+
+        @"  browser.test.assertEq(sender.frameId, 0, 'sender.frameId should be 0')",
+
+        @"  setTimeout(() => sendResponse({ content: 'Received' }), 1000)",
+
+        @"  return true",
+        @"})",
+
+        @"browser.test.yield('Load Tab')"
+    ]);
+
+    auto *contentScript = Util::constructScript(@[
+        @"(async () => {",
+        @"  const response = await browser.runtime.sendMessage({ content: 'Hello' })",
+
+        @"  browser.test.assertEq(response?.content, 'Received', 'Should get the response from the background script')",
 
         @"  browser.test.notifyPass()",
         @"})()"
@@ -450,7 +508,7 @@
         [NSString stringWithFormat:@"const expectedOrigin = '%@'", [urlString substringToIndex:urlString.length - 1]],
 
         @"browser.runtime.onMessage.addListener((message, sender, sendResponse) => {",
-        @"  browser.test.assertEq(message.content, 'Hello', 'Should receive the correct message from the content script')",
+        @"  browser.test.assertEq(message?.content, 'Hello', 'Should receive the correct message from the content script')",
 
         @"  browser.test.assertEq(typeof sender, 'object', 'sender should be an object')",
 
@@ -472,7 +530,65 @@
         @"(async () => {",
         @"  const response = await browser.runtime.sendMessage({ content: 'Hello' })",
 
-        @"  browser.test.assertEq(response.content, 'Received', 'Should get the response from the background script')",
+        @"  browser.test.assertEq(response?.content, 'Received', 'Should get the response from the background script')",
+
+        @"  browser.test.notifyPass()",
+        @"})()"
+    ]);
+
+    auto extension = adoptNS([[_WKWebExtension alloc] _initWithManifestDictionary:runtimeContentScriptManifest resources:@{ @"background.js": backgroundScript, @"content.js": contentScript }]);
+    auto manager = adoptNS([[TestWebExtensionManager alloc] initForExtension:extension.get()]);
+
+    [manager.get().context setPermissionStatus:_WKWebExtensionContextPermissionStatusGrantedExplicitly forURL:urlRequest.URL];
+
+    [manager loadAndRun];
+
+    EXPECT_NS_EQUAL(manager.get().yieldMessage, @"Load Tab");
+
+    [manager.get().defaultTab.mainWebView loadRequest:urlRequest];
+
+    [manager run];
+}
+
+TEST(WKWebExtensionAPIRuntime, SendMessageFromContentScriptWithAsyncPromiseReply)
+{
+    TestWebKitAPI::HTTPServer server({
+        { "/"_s, { { { "Content-Type"_s, "text/html"_s } }, ""_s } },
+    }, TestWebKitAPI::HTTPServer::Protocol::Http);
+
+    auto *urlRequest = server.requestWithLocalhost();
+    auto *urlString = urlRequest.URL.absoluteString;
+
+    auto *backgroundScript = Util::constructScript(@[
+        [NSString stringWithFormat:@"const expectedURL = '%@'", urlString],
+        [NSString stringWithFormat:@"const expectedOrigin = '%@'", [urlString substringToIndex:urlString.length - 1]],
+
+        @"browser.runtime.onMessage.addListener((message, sender, sendResponse) => {",
+        @"  browser.test.assertEq(message?.content, 'Hello', 'Should receive the correct message from the content script')",
+
+        @"  browser.test.assertEq(typeof sender, 'object', 'sender should be an object')",
+
+        @"  browser.test.assertEq(typeof sender.tab, 'object', 'sender.tab should be an object')",
+        @"  browser.test.assertEq(sender.tab.url, expectedURL, 'sender.tab.url should be the expected URL')",
+
+        @"  browser.test.assertEq(sender.url, expectedURL, 'sender.url should be the expected URL')",
+        @"  browser.test.assertEq(sender.origin, expectedOrigin, 'sender.origin should be the expected origin')",
+
+        @"  browser.test.assertEq(sender.frameId, 0, 'sender.frameId should be 0')",
+
+        @"  return new Promise((resolve) => {",
+        @"    setTimeout(() => resolve({ content: 'Received' }), 1000)",
+        @"  })",
+        @"})",
+
+        @"browser.test.yield('Load Tab')"
+    ]);
+
+    auto *contentScript = Util::constructScript(@[
+        @"(async () => {",
+        @"  const response = await browser.runtime.sendMessage({ content: 'Hello' })",
+
+        @"  browser.test.assertEq(response?.content, 'Received', 'Should get the response from the background script')",
 
         @"  browser.test.notifyPass()",
         @"})()"
@@ -506,7 +622,7 @@
         [NSString stringWithFormat:@"const expectedOrigin = '%@'", [urlString substringToIndex:urlString.length - 1]],
 
         @"browser.runtime.onMessage.addListener((message, sender) => {",
-        @"  browser.test.assertEq(message.content, 'Hello', 'Should receive the correct message from the content script')",
+        @"  browser.test.assertEq(message?.content, 'Hello', 'Should receive the correct message from the content script')",
 
         @"  browser.test.assertEq(typeof sender, 'object', 'sender should be an object')",
 
@@ -517,6 +633,8 @@
         @"  browser.test.assertEq(sender.origin, expectedOrigin, 'sender.origin should be the expected origin')",
 
         @"  browser.test.assertEq(sender.frameId, 0, 'sender.frameId should be 0')",
+
+        @"  return false",
         @"})",
 
         @"browser.test.yield('Load Tab')"