Use WKDataTask for SystemPreview downloads

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