Use WKDataTask for SystemPreview downloads

https://bugs.webkit.org/show_bug.cgi?id=255954
rdar://108524222

Reviewed by Tim Horton.

Migrate away from LegacyDownloadClient and instead handle System Previews
directly. This will allow us to remove the horrible code in the download
client, as well eventually allow 3rd-party engines to use ARQL.

The fix is to intercept the HTMLAnchorElement click for a system
preview, and send that directly to the SystemPreviewController in
the UI Process. That object then creates a WKDataTask for the download
and presents ARQL as normal.

A followup patch will remove the old code.

* Source/WebCore/html/HTMLAnchorElement.cpp:
* Source/WebCore/page/ChromeClient.h:
* Source/WebCore/page/Page.cpp:
* Source/WebCore/page/Page.h:
* Source/WebKit/UIProcess/API/Cocoa/WKWebViewPrivateForTesting.h:
* Source/WebKit/UIProcess/API/Cocoa/WKWebViewTesting.mm:
* Source/WebKit/UIProcess/Cocoa/SystemPreviewControllerCocoa.mm:
* Source/WebKit/UIProcess/SystemPreviewController.h:
* Source/WebKit/UIProcess/WebPageProxy.cpp:
* Source/WebKit/UIProcess/WebPageProxy.h:
* Source/WebKit/UIProcess/WebPageProxy.messages.in:
* Source/WebKit/WebProcess/WebCoreSupport/WebChromeClient.cpp:
* Source/WebKit/WebProcess/WebCoreSupport/WebChromeClient.h:
* Tools/TestWebKitAPI/Tests/WebKitCocoa/SystemPreview.mm:
(-[TestSystemPreviewUIDelegate _presentingViewControllerForWebView:]):
(TestWebKitAPI::TEST):
* Tools/TestWebKitAPI/Tests/WebKitCocoa/UnitBox.usdz: Added.
* Tools/TestWebKitAPI/Tests/WebKitCocoa/system-preview-trigger.html:
* Tools/TestWebKitAPI/Tests/WebKitCocoa/system-preview.html: Added.

Canonical link: https://commits.webkit.org/263393@main

Source/WebCore/PAL/pal/spi/cocoa/FoundationSPI.h

@@ -30,3 +30,6 @@
 - (NSDictionary *)detail;
 @end
 
+@interface NSHTTPURLResponse ()
++ (BOOL)isErrorStatusCode:(NSInteger)statusCode;
+@end

Source/WebCore/html/HTMLAnchorElement.cpp

@@ -611,6 +611,10 @@ void HTMLAnchorElement::handleClick(Event& event)
         systemPreviewInfo.element.webPageIdentifier = valueOrDefault(document().frame()->loader().pageID());
         if (auto* child = firstElementChild())
             systemPreviewInfo.previewRect = child->boundsInRootViewSpace();
+
+        if (auto* page = document().page())
+            page->handleSystemPreview(WTFMove(completedURL), WTFMove(systemPreviewInfo));
+        return;
     }
 #endif
 

Source/WebCore/page/ChromeClient.h

@@ -637,6 +637,10 @@ class ChromeClient {
     virtual void abortApplePayAMSUISession() { }
 #endif
 
+#if USE(SYSTEM_PREVIEW)
+    virtual void handleSystemPreview(const URL&, const SystemPreviewInfo&) { }
+#endif
+
     virtual void requestCookieConsent(CompletionHandler<void(CookieConsentDecisionResult)>&&) = 0;
 
     virtual const AtomString& searchStringForModalContainerObserver() const { return nullAtom(); }

Source/WebCore/page/Page.cpp

@@ -3903,6 +3903,13 @@ void Page::abortApplePayAMSUISession(ApplePayAMSUIPaymentHandler& paymentHandler
 
 #endif // ENABLE(APPLE_PAY_AMS_UI)
 
+#if USE(SYSTEM_PREVIEW)
+void Page::handleSystemPreview(const URL& url, const SystemPreviewInfo& systemPreviewInfo)
+{
+    chrome().client().handleSystemPreview(url, systemPreviewInfo);
+}
+#endif
+
 #if ENABLE(MEDIA_SESSION_COORDINATOR)
 void Page::setMediaSessionCoordinator(Ref<MediaSessionCoordinatorPrivate>&& mediaSessionCoordinator)
 {

Source/WebCore/page/Page.h

@@ -601,6 +601,10 @@ class Page : public Supplementable<Page>, public CanMakeWeakPtr<Page> {
     void abortApplePayAMSUISession(ApplePayAMSUIPaymentHandler&);
 #endif
 
+#if USE(SYSTEM_PREVIEW)
+    void handleSystemPreview(const URL&, const SystemPreviewInfo&);
+#endif
+
 #if ENABLE(WEB_AUTHN)
     AuthenticatorCoordinator& authenticatorCoordinator() { return m_authenticatorCoordinator.get(); }
 #endif

Source/WebKit/UIProcess/API/Cocoa/WKWebViewPrivateForTesting.h

@@ -136,6 +136,8 @@ struct WKAppPrivacyReportTestingData {
 
 - (void)_setConnectedToHardwareConsoleForTesting:(BOOL)connected;
 
+- (void)_setSystemPreviewCompletionHandlerForLoadTesting:(void(^)(bool))completionHandler;
+
 @end
 
 typedef NS_ENUM(NSInteger, _WKMediaSessionReadyState) {

Source/WebKit/UIProcess/API/Cocoa/WKWebViewTesting.mm

@@ -577,6 +577,13 @@ - (void)_setConnectedToHardwareConsoleForTesting:(BOOL)connected
 #endif
 }
 
+- (void)_setSystemPreviewCompletionHandlerForLoadTesting:(void(^)(bool))completionHandler
+{
+#if USE(SYSTEM_PREVIEW)
+    _page->setSystemPreviewCompletionHandlerForLoadTesting(makeBlockPtr(completionHandler));
+#endif
+}
+
 - (void)_createMediaSessionCoordinatorForTesting:(id <_WKMediaSessionCoordinator>)privateCoordinator completionHandler:(void(^)(BOOL))completionHandler
 {
 #if ENABLE(MEDIA_SESSION_COORDINATOR)

Source/WebKit/UIProcess/Cocoa/SystemPreviewControllerCocoa.mm

@@ -31,12 +31,15 @@
 #import "APIUIClient.h"
 #import "WebPageProxy.h"
 #import "WebProcessProxy.h"
+#import "_WKDataTaskDelegate.h"
+#import "_WKDataTaskInternal.h"
 #import <MobileCoreServices/MobileCoreServices.h>
 #import <QuickLook/QuickLook.h>
 #import <UIKit/UIViewController.h>
 #import <WebCore/MIMETypeRegistry.h>
 #import <WebCore/UTIUtilities.h>
 #import <pal/ios/QuickLookSoftLink.h>
+#import <pal/spi/cocoa/FoundationSPI.h>
 #import <pal/spi/ios/QuickLookSPI.h>
 #import <wtf/WeakObjCPtr.h>
 
@@ -253,8 +256,202 @@ - (UIImage *)previewController:(QLPreviewController *)controller transitionImage
 
 @end
 
+@interface _WKSystemPreviewDataTaskDelegate : NSObject <_WKDataTaskDelegate> {
+    WebKit::SystemPreviewController* _previewController;
+    long long _expectedContentLength;
+    RetainPtr<NSMutableData> _data;
+    RetainPtr<NSString> _suggestedFilename;
+};
+@end
+
+@implementation _WKSystemPreviewDataTaskDelegate
+
+- (id)initWithSystemPreviewController:(WebKit::SystemPreviewController*)previewController
+{
+    if (!(self = [super init]))
+        return nil;
+
+    _previewController = previewController;
+    return self;
+}
+
+- (BOOL)isValidMIMEType:(NSString *)MIMEType
+{
+    return WebCore::MIMETypeRegistry::isUSDMIMEType(MIMEType);
+}
+
+- (void)dataTask:(_WKDataTask *)dataTask didReceiveResponse:(NSURLResponse *)response decisionHandler:(void (^)(_WKDataTaskResponsePolicy))decisionHandler
+{
+    if ([response isKindOfClass:[NSHTTPURLResponse class]]) {
+        NSHTTPURLResponse *HTTPResponse = (NSHTTPURLResponse *)response;
+        if ([NSHTTPURLResponse isErrorStatusCode:HTTPResponse.statusCode]) {
+            RELEASE_LOG(SystemPreview, "cancelling subresource load due to error status code: %ld", (long)HTTPResponse.statusCode);
+            decisionHandler(_WKDataTaskResponsePolicyCancel);
+            _previewController->loadFailed();
+            return;
+        }
+    }
+
+    if (![self isValidMIMEType:response.MIMEType]) {
+        RELEASE_LOG(SystemPreview, "cancelling subresource load due to unhandled MIME type: \"%@\"", response.MIMEType);
+        decisionHandler(_WKDataTaskResponsePolicyCancel);
+        _previewController->loadFailed();
+        return;
+    }
+
+    _expectedContentLength = response.expectedContentLength;
+    if (_expectedContentLength == NSURLResponseUnknownLength)
+        _expectedContentLength = 0;
+
+    _data = adoptNS([[NSMutableData alloc] initWithCapacity:_expectedContentLength]);
+    _suggestedFilename = adoptNS([response.suggestedFilename copy]);
+    decisionHandler(_WKDataTaskResponsePolicyAllow);
+}
+
+- (void)dataTask:(_WKDataTask *)dataTask didReceiveData:(NSData *)data
+{
+    ASSERT(_data);
+    [_data appendData:data];
+    if (_expectedContentLength)
+        _previewController->updateProgress((float)_data.get().length / _expectedContentLength);
+}
+
+- (void)dataTask:(_WKDataTask *)dataTask didCompleteWithError:(NSError *)error
+{
+    if (error) {
+        _previewController->loadFailed();
+        return;
+    }
+
+    [self completeLoad];
+}
+
+- (void)completeLoad
+{
+    FileSystem::PlatformFileHandle fileHandle;
+    auto filePath = FileSystem::openTemporaryFile("SystemPreview"_s, fileHandle, ".usdz"_s);
+    ASSERT(FileSystem::isHandleValid(fileHandle));
+
+    size_t byteCount = FileSystem::writeToFile(fileHandle, [_data bytes], [_data length]);
+    FileSystem::closeFile(fileHandle);
+
+    if (byteCount != _data.get().length) {
+        _previewController->loadFailed();
+        return;
+    }
+
+    _previewController->loadCompleted(URL::fileURLWithFileSystemPath(filePath));
+}
+
+@end
+
 namespace WebKit {
 
+void SystemPreviewController::begin(const URL& url, const WebCore::SystemPreviewInfo& systemPreviewInfo)
+{
+    ASSERT(!m_qlPreviewController);
+    if (m_qlPreviewController)
+        return;
+
+    UIViewController *presentingViewController = m_webPageProxy.uiClient().presentingViewController();
+
+    if (!presentingViewController)
+        return;
+
+    m_systemPreviewInfo = systemPreviewInfo;
+
+    RELEASE_LOG(SystemPreview, "SystemPreview began on %lld", m_systemPreviewInfo.element.elementIdentifier.toUInt64());
+
+    auto request = WebCore::ResourceRequest(url);
+    WeakPtr weakThis { *this };
+    m_webPageProxy.dataTaskWithRequest(WTFMove(request), [weakThis] (Ref<API::DataTask>&& task) {
+        if (!weakThis)
+            return;
+
+        auto strongThis = weakThis.get();
+
+        _WKDataTask *dataTask = wrapper(task);
+        strongThis->m_wkSystemPreviewDataTaskDelegate = adoptNS([[_WKSystemPreviewDataTaskDelegate alloc] initWithSystemPreviewController:strongThis]);
+        [dataTask setDelegate:strongThis->m_wkSystemPreviewDataTaskDelegate.get()];
+        strongThis->takeActivityToken();
+    });
+
+    m_qlPreviewController = adoptNS([PAL::allocQLPreviewControllerInstance() init]);
+
+    m_qlPreviewControllerDelegate = adoptNS([[_WKPreviewControllerDelegate alloc] initWithSystemPreviewController:this]);
+    [m_qlPreviewController setDelegate:m_qlPreviewControllerDelegate.get()];
+
+    m_qlPreviewControllerDataSource = adoptNS([[_WKPreviewControllerDataSource alloc] initWithSystemPreviewController:this MIMEType:@"model/vnd.usdz+zip" originatingPageURL:url]);
+    [m_qlPreviewController setDataSource:m_qlPreviewControllerDataSource.get()];
+
+    [presentingViewController presentViewController:m_qlPreviewController.get() animated:YES completion:nullptr];
+}
+
+void SystemPreviewController::loadCompleted(const URL& downloadedFile)
+{
+    RELEASE_LOG(SystemPreview, "SystemPreview load has finished on %lld", m_systemPreviewInfo.element.elementIdentifier.toUInt64());
+
+#if HAVE(UIKIT_WEBKIT_INTERNALS)
+    ASSERT(equalIgnoringFragmentIdentifier(m_destinationURL, url));
+    NSURL *nsurl = (NSURL *)url;
+    if ([getASVLaunchPreviewClass() respondsToSelector:@selector(launchPreviewApplicationWithURLs:completion:)])
+        [getASVLaunchPreviewClass() launchPreviewApplicationWithURLs:@[nsurl] completion:^(NSError *error) { }];
+#else
+    if (m_qlPreviewControllerDataSource)
+        [m_qlPreviewControllerDataSource finish:downloadedFile];
+#endif
+    releaseActivityTokenIfNecessary();
+
+    if (m_testingCallback)
+        m_testingCallback(true);
+}
+
+void SystemPreviewController::loadFailed()
+{
+    RELEASE_LOG(SystemPreview, "SystemPreview failed on %lld", m_systemPreviewInfo.element.elementIdentifier.toUInt64());
+
+#if !HAVE(UIKIT_WEBKIT_INTERNALS)
+    if (m_qlPreviewControllerDataSource)
+        [m_qlPreviewControllerDataSource.get() failWithError:nil];
+
+    if (m_qlPreviewController)
+        [m_qlPreviewController.get() dismissViewControllerAnimated:YES completion:nullptr];
+
+    m_qlPreviewControllerDelegate = nullptr;
+    m_qlPreviewControllerDataSource = nullptr;
+    m_qlPreviewController = nullptr;
+    m_wkSystemPreviewDataTaskDelegate = nullptr;
+#endif
+    releaseActivityTokenIfNecessary();
+
+    if (m_testingCallback)
+        m_testingCallback(false);
+}
+
+void SystemPreviewController::takeActivityToken()
+{
+#if USE(RUNNINGBOARD)
+    RELEASE_LOG(ProcessSuspension, "%p - UIProcess is taking a background assertion because it is downloading a system preview", this);
+    ASSERT(!m_activity);
+    m_activity = page().process().throttler().backgroundActivity("System preview download"_s).moveToUniquePtr();
+#endif
+}
+
+void SystemPreviewController::releaseActivityTokenIfNecessary()
+{
+#if USE(RUNNINGBOARD)
+    if (m_activity) {
+        RELEASE_LOG(ProcessSuspension, "%p UIProcess is releasing a background assertion because a system preview download completed", this);
+        m_activity = nullptr;
+    }
+#endif
+}
+
+void SystemPreviewController::setCompletionHandlerForLoadTesting(CompletionHandler<void(bool)>&& handler)
+{
+    m_testingCallback = WTFMove(handler);
+}
+
 void SystemPreviewController::start(URL originatingPageURL, const String& mimeType, const WebCore::SystemPreviewInfo& systemPreviewInfo)
 {
 #if HAVE(UIKIT_WEBKIT_INTERNALS)
@@ -285,7 +482,7 @@ - (UIImage *)previewController:(QLPreviewController *)controller transitionImage
 
     m_originatingPageURL = originatingPageURL;
 
-    RELEASE_LOG(SystemPreview, "SystemPreview began on %lld", m_systemPreviewInfo.element.elementIdentifier.toUInt64());
+    RELEASE_LOG(SystemPreview, "SystemPreview started on %lld", m_systemPreviewInfo.element.elementIdentifier.toUInt64());
 }
 
 void SystemPreviewController::setDestinationURL(URL url)

Source/WebKit/UIProcess/SystemPreviewController.h

@@ -27,30 +27,40 @@
 
 #if USE(SYSTEM_PREVIEW)
 
+#include "ProcessThrottler.h"
 #include <WebCore/FrameLoaderTypes.h>
 #include <WebCore/IntRect.h>
 #include <WebCore/ResourceError.h>
 #include <wtf/RetainPtr.h>
 #include <wtf/URL.h>
+#include <wtf/WeakPtr.h>
 
 OBJC_CLASS NSString;
 #if USE(QUICK_LOOK)
 OBJC_CLASS QLPreviewController;
 OBJC_CLASS _WKPreviewControllerDataSource;
 OBJC_CLASS _WKPreviewControllerDelegate;
+OBJC_CLASS _WKSystemPreviewDataTaskDelegate;
 #endif
 
 namespace WebKit {
 
 class WebPageProxy;
 
-class SystemPreviewController {
+class SystemPreviewController : public CanMakeWeakPtr<SystemPreviewController> {
     WTF_MAKE_FAST_ALLOCATED;
 public:
     explicit SystemPreviewController(WebPageProxy&);
 
     bool canPreview(const String& mimeType) const;
 
+    // New methods that use WKDataTask.
+    void begin(const URL&, const WebCore::SystemPreviewInfo&);
+    void loadCompleted(const URL& downloadedFile);
+    void loadFailed();
+    void end();
+
+    // Old methods that use LegacyDownloadClient.
     void start(URL originatingPageURL, const String& mimeType, const WebCore::SystemPreviewInfo&);
     void setDestinationURL(URL);
     void updateProgress(float);
@@ -64,8 +74,13 @@ class SystemPreviewController {
     void triggerSystemPreviewAction();
 
     void triggerSystemPreviewActionWithTargetForTesting(uint64_t elementID, NSString* documentID, uint64_t pageID);
+    void setCompletionHandlerForLoadTesting(CompletionHandler<void(bool)>&&);
 
 private:
+
+    void takeActivityToken();
+    void releaseActivityTokenIfNecessary();
+
     WebPageProxy& m_webPageProxy;
     WebCore::SystemPreviewInfo m_systemPreviewInfo;
     URL m_destinationURL;
@@ -74,7 +89,12 @@ class SystemPreviewController {
     RetainPtr<QLPreviewController> m_qlPreviewController;
     RetainPtr<_WKPreviewControllerDelegate> m_qlPreviewControllerDelegate;
     RetainPtr<_WKPreviewControllerDataSource> m_qlPreviewControllerDataSource;
+    RetainPtr<_WKSystemPreviewDataTaskDelegate> m_wkSystemPreviewDataTaskDelegate;
 #endif
+
+    std::unique_ptr<ProcessThrottler::BackgroundActivity> m_activity;
+    CompletionHandler<void(bool)> m_testingCallback;
+
 };
 
 }

Source/WebKit/UIProcess/WebPageProxy.cpp

@@ -12570,6 +12570,20 @@ FloatSize WebPageProxy::viewportSizeForCSSViewportUnits() const
     return valueOrDefault(internals().viewportSizeForCSSViewportUnits);
 }
 
+#if USE(SYSTEM_PREVIEW)
+void WebPageProxy::handleSystemPreview(const URL& url, const SystemPreviewInfo& systemPreviewInfo)
+{
+    if (m_systemPreviewController)
+        m_systemPreviewController->begin(url, systemPreviewInfo);
+}
+
+void WebPageProxy::setSystemPreviewCompletionHandlerForLoadTesting(CompletionHandler<void(bool)>&& handler)
+{
+    if (m_systemPreviewController)
+        m_systemPreviewController->setCompletionHandlerForLoadTesting(WTFMove(handler));
+}
+#endif
+
 } // namespace WebKit
 
 #undef WEBPAGEPROXY_RELEASE_LOG