Fix basic WKURLSchemeHandler bugs.
authorbeidson@apple.com <beidson@apple.com@268f45cc-cd09-0410-ab3c-d52691b4dbfc>
Sat, 15 Apr 2017 00:52:10 +0000 (00:52 +0000)
committerbeidson@apple.com <beidson@apple.com@268f45cc-cd09-0410-ab3c-d52691b4dbfc>
Sat, 15 Apr 2017 00:52:10 +0000 (00:52 +0000)
<rdar://problem/30647559> and https://bugs.webkit.org/show_bug.cgi?id=170862

Reviewed by Andy Estes.

Source/WebCore:

Covered by new API tests.

* loader/SubresourceLoader.cpp:
(WebCore::SubresourceLoader::didReceiveDataOrBuffer):

Source/WebKit2:

* UIProcess/Cocoa/NavigationState.mm:
(WebKit::NavigationState::NavigationClient::decidePolicyForNavigationAction):

* UIProcess/WebURLSchemeHandlerTask.cpp:
(WebKit::WebURLSchemeHandlerTask::didReceiveResponse):

* WebProcess/Network/WebLoaderStrategy.cpp:
(WebKit::WebLoaderStrategy::addURLSchemeHandlerTaskProxy):
(WebKit::WebLoaderStrategy::removeURLSchemeHandlerTaskProxy):
(WebKit::WebLoaderStrategy::remove):
* WebProcess/Network/WebLoaderStrategy.h:

* WebProcess/WebPage/WebURLSchemeHandlerProxy.cpp:
(WebKit::WebURLSchemeHandlerProxy::startNewTask):
(WebKit::WebURLSchemeHandlerProxy::taskDidComplete):
(WebKit::WebURLSchemeHandlerProxy::taskDidStopLoading):
* WebProcess/WebPage/WebURLSchemeHandlerProxy.h:

* WebProcess/WebPage/WebURLSchemeHandlerTaskProxy.cpp:
(WebKit::WebURLSchemeHandlerTaskProxy::WebURLSchemeHandlerTaskProxy):
(WebKit::WebURLSchemeHandlerTaskProxy::stopLoading):
(WebKit::WebURLSchemeHandlerTaskProxy::didReceiveResponse):
(WebKit::WebURLSchemeHandlerTaskProxy::didReceiveData):
(WebKit::WebURLSchemeHandlerTaskProxy::didComplete):
(WebKit::WebURLSchemeHandlerTaskProxy::hasLoader):
* WebProcess/WebPage/WebURLSchemeHandlerTaskProxy.h:
(WebKit::WebURLSchemeHandlerTaskProxy::identifier):

Tools:

* TestWebKitAPI/TestWebKitAPI.xcodeproj/project.pbxproj:
* TestWebKitAPI/Tests/WebKit2Cocoa/WKURLSchemeHandler-1.mm: Added.
(-[SchemeHandler initWithData:mimeType:]):
(-[SchemeHandler dealloc]):
(-[SchemeHandler webView:startTask:]):
(-[SchemeHandler webView:stopTask:]):
(TEST):

git-svn-id: http://svn.webkit.org/repository/webkit/trunk@215384 268f45cc-cd09-0410-ab3c-d52691b4dbfc

14 files changed:
Source/WebCore/ChangeLog
Source/WebCore/loader/SubresourceLoader.cpp
Source/WebKit2/ChangeLog
Source/WebKit2/UIProcess/Cocoa/NavigationState.mm
Source/WebKit2/UIProcess/WebURLSchemeHandlerTask.cpp
Source/WebKit2/WebProcess/Network/WebLoaderStrategy.cpp
Source/WebKit2/WebProcess/Network/WebLoaderStrategy.h
Source/WebKit2/WebProcess/WebPage/WebURLSchemeHandlerProxy.cpp
Source/WebKit2/WebProcess/WebPage/WebURLSchemeHandlerProxy.h
Source/WebKit2/WebProcess/WebPage/WebURLSchemeHandlerTaskProxy.cpp
Source/WebKit2/WebProcess/WebPage/WebURLSchemeHandlerTaskProxy.h
Tools/ChangeLog
Tools/TestWebKitAPI/TestWebKitAPI.xcodeproj/project.pbxproj
Tools/TestWebKitAPI/Tests/WebKit2Cocoa/WKURLSchemeHandler-1.mm [new file with mode: 0644]

index 3341b91..8b2f5ac 100644 (file)
@@ -1,3 +1,15 @@
+2017-04-14  Brady Eidson  <beidson@apple.com>
+
+        Fix basic WKURLSchemeHandler bugs.
+        <rdar://problem/30647559> and https://bugs.webkit.org/show_bug.cgi?id=170862
+
+        Reviewed by Andy Estes.
+
+        Covered by new API tests.
+
+        * loader/SubresourceLoader.cpp:
+        (WebCore::SubresourceLoader::didReceiveDataOrBuffer):
+
 2017-04-14  Jiewen Tan  <jiewen_tan@apple.com>
 
         [WebCrypto] Support HKDF
index 8ba561c..275b06a 100644 (file)
@@ -390,6 +390,8 @@ void SubresourceLoader::didReceiveBuffer(Ref<SharedBuffer>&& buffer, long long e
 
 void SubresourceLoader::didReceiveDataOrBuffer(const char* data, int length, RefPtr<SharedBuffer>&& buffer, long long encodedDataLength, DataPayloadType dataPayloadType)
 {
+    ASSERT(m_resource);
+
     if (m_resource->response().httpStatusCode() >= 400 && !m_resource->shouldIgnoreHTTPStatusCodeErrors())
         return;
     ASSERT(!m_resource->resourceToRevalidate());
index 69d6f04..698020a 100644 (file)
@@ -1,3 +1,38 @@
+2017-04-14  Brady Eidson  <beidson@apple.com>
+
+        Fix basic WKURLSchemeHandler bugs.
+        <rdar://problem/30647559> and https://bugs.webkit.org/show_bug.cgi?id=170862
+
+        Reviewed by Andy Estes.
+
+        * UIProcess/Cocoa/NavigationState.mm:
+        (WebKit::NavigationState::NavigationClient::decidePolicyForNavigationAction):
+
+        * UIProcess/WebURLSchemeHandlerTask.cpp:
+        (WebKit::WebURLSchemeHandlerTask::didReceiveResponse):
+
+        * WebProcess/Network/WebLoaderStrategy.cpp:
+        (WebKit::WebLoaderStrategy::addURLSchemeHandlerTaskProxy):
+        (WebKit::WebLoaderStrategy::removeURLSchemeHandlerTaskProxy):
+        (WebKit::WebLoaderStrategy::remove):
+        * WebProcess/Network/WebLoaderStrategy.h:
+
+        * WebProcess/WebPage/WebURLSchemeHandlerProxy.cpp:
+        (WebKit::WebURLSchemeHandlerProxy::startNewTask):
+        (WebKit::WebURLSchemeHandlerProxy::taskDidComplete):
+        (WebKit::WebURLSchemeHandlerProxy::taskDidStopLoading):
+        * WebProcess/WebPage/WebURLSchemeHandlerProxy.h:
+
+        * WebProcess/WebPage/WebURLSchemeHandlerTaskProxy.cpp:
+        (WebKit::WebURLSchemeHandlerTaskProxy::WebURLSchemeHandlerTaskProxy):
+        (WebKit::WebURLSchemeHandlerTaskProxy::stopLoading):
+        (WebKit::WebURLSchemeHandlerTaskProxy::didReceiveResponse):
+        (WebKit::WebURLSchemeHandlerTaskProxy::didReceiveData):
+        (WebKit::WebURLSchemeHandlerTaskProxy::didComplete):
+        (WebKit::WebURLSchemeHandlerTaskProxy::hasLoader):
+        * WebProcess/WebPage/WebURLSchemeHandlerTaskProxy.h:
+        (WebKit::WebURLSchemeHandlerTaskProxy::identifier):
+
 2017-04-14  Jer Noble  <jer.noble@apple.com>
 
         [MediaSource] Push capabilities across process boundary during UIProcess capture.
index 96b9602..b66e1a1 100644 (file)
@@ -311,7 +311,7 @@ void NavigationState::NavigationClient::decidePolicyForNavigationAction(WebPageP
         RefPtr<API::NavigationAction> localNavigationAction = &navigationAction;
         RefPtr<WebFramePolicyListenerProxy> localListener = WTFMove(listener);
 
-        tryAppLink(WTFMove(localNavigationAction), mainFrameURLString, [localListener, localNavigationAction = RefPtr<API::NavigationAction>(&navigationAction)] (bool followedLinkToApp) {
+        tryAppLink(WTFMove(localNavigationAction), mainFrameURLString, [webPage = RefPtr<WebPageProxy>(&webPageProxy), localListener, localNavigationAction = RefPtr<API::NavigationAction>(&navigationAction)] (bool followedLinkToApp) {
             if (followedLinkToApp) {
                 localListener->ignore();
                 return;
@@ -323,7 +323,7 @@ void NavigationState::NavigationClient::decidePolicyForNavigationAction(WebPageP
             }
 
             RetainPtr<NSURLRequest> nsURLRequest = adoptNS(wrapper(API::URLRequest::create(localNavigationAction->request()).leakRef()));
-            if ([NSURLConnection canHandleRequest:nsURLRequest.get()]) {
+            if ([NSURLConnection canHandleRequest:nsURLRequest.get()] || webPage->urlSchemeHandlerForScheme([nsURLRequest URL].scheme)) {
                 if (localNavigationAction->shouldPerformDownload())
                     localListener->download();
                 else
index 120ab3f..0deab3e 100644 (file)
@@ -60,6 +60,8 @@ WebURLSchemeHandlerTask::ExceptionType WebURLSchemeHandlerTask::didReceiveRespon
         return WebURLSchemeHandlerTask::ExceptionType::DataAlreadySent;
 
     m_responseSent = true;
+
+    response.includeCertificateInfo();
     m_page->send(Messages::WebPage::URLSchemeHandlerTaskDidReceiveResponse(m_urlSchemeHandler->identifier(), m_identifier, response));
     return WebURLSchemeHandlerTask::ExceptionType::None;
 }
index 7a9b6f0..fbd3dc8 100644 (file)
@@ -264,18 +264,30 @@ void WebLoaderStrategy::startLocalLoad(WebCore::ResourceLoader& resourceLoader)
     m_webResourceLoaders.set(resourceLoader.identifier(), WebResourceLoader::create(resourceLoader, { }));
 }
 
+void WebLoaderStrategy::addURLSchemeHandlerTaskProxy(WebURLSchemeHandlerTaskProxy& task)
+{
+    auto result = m_urlSchemeHandlerTasks.add(task.identifier(), &task);
+    ASSERT_UNUSED(result, result.isNewEntry);
+}
+
+void WebLoaderStrategy::removeURLSchemeHandlerTaskProxy(WebURLSchemeHandlerTaskProxy& task)
+{
+    m_urlSchemeHandlerTasks.remove(task.identifier());
+}
+
 void WebLoaderStrategy::remove(ResourceLoader* resourceLoader)
 {
     ASSERT(resourceLoader);
     LOG(NetworkScheduling, "(WebProcess) WebLoaderStrategy::remove, url '%s'", resourceLoader->url().string().utf8().data());
 
-    if (m_internallyFailedResourceLoaders.contains(resourceLoader)) {
-        m_internallyFailedResourceLoaders.remove(resourceLoader);
+    if (auto task = m_urlSchemeHandlerTasks.take(resourceLoader->identifier())) {
+        ASSERT(!m_internallyFailedResourceLoaders.contains(resourceLoader));
+        task->stopLoading();
         return;
     }
 
-    if (auto task = m_urlSchemeHandlerTasks.take(resourceLoader->identifier())) {
-        task->stopLoading();
+    if (m_internallyFailedResourceLoaders.contains(resourceLoader)) {
+        m_internallyFailedResourceLoaders.remove(resourceLoader);
         return;
     }
     
index 533e4df..51b2a82 100644 (file)
@@ -66,6 +66,9 @@ public:
 
     void networkProcessCrashed();
 
+    void addURLSchemeHandlerTaskProxy(WebURLSchemeHandlerTaskProxy&);
+    void removeURLSchemeHandlerTaskProxy(WebURLSchemeHandlerTaskProxy&);
+
 private:
     void scheduleLoad(WebCore::ResourceLoader&, WebCore::CachedResource*, bool shouldClearReferrerOnHTTPSToHTTPRedirect);
     void scheduleInternallyFailedLoad(WebCore::ResourceLoader&);
@@ -76,7 +79,7 @@ private:
     RunLoop::Timer<WebLoaderStrategy> m_internallyFailedLoadTimer;
     
     HashMap<unsigned long, RefPtr<WebResourceLoader>> m_webResourceLoaders;
-    HashMap<unsigned long, std::unique_ptr<WebURLSchemeHandlerTaskProxy>> m_urlSchemeHandlerTasks;
+    HashMap<unsigned long, WebURLSchemeHandlerTaskProxy*> m_urlSchemeHandlerTasks;
 };
 
 } // namespace WebKit
index 7e0a8d1..941b797 100644 (file)
@@ -27,6 +27,8 @@
 #include "WebURLSchemeHandlerProxy.h"
 
 #include "WebErrors.h"
+#include "WebLoaderStrategy.h"
+#include "WebProcess.h"
 #include <WebCore/ResourceLoader.h>
 
 using namespace WebCore;
@@ -49,6 +51,7 @@ void WebURLSchemeHandlerProxy::startNewTask(ResourceLoader& loader)
     auto result = m_tasks.add(loader.identifier(), std::make_unique<WebURLSchemeHandlerTaskProxy>(*this, loader));
     ASSERT(result.isNewEntry);
 
+    WebProcess::singleton().webLoaderStrategy().addURLSchemeHandlerTaskProxy(*result.iterator->value);
     result.iterator->value->startLoading();
 }
 
@@ -77,7 +80,14 @@ void WebURLSchemeHandlerProxy::taskDidComplete(uint64_t taskIdentifier, const Re
     if (!task)
         return;
 
+    WebProcess::singleton().webLoaderStrategy().removeURLSchemeHandlerTaskProxy(*task);
     task->didComplete(error);
 }
 
+void WebURLSchemeHandlerProxy::taskDidStopLoading(WebURLSchemeHandlerTaskProxy& task)
+{
+    ASSERT(m_tasks.get(task.identifier()) == &task);
+    m_tasks.remove(task.identifier());
+}
+
 } // namespace WebKit
index cefc3af..f297300 100644 (file)
@@ -52,6 +52,7 @@ public:
     void taskDidReceiveResponse(uint64_t taskIdentifier, const WebCore::ResourceResponse&);
     void taskDidReceiveData(uint64_t taskIdentifier, size_t, const uint8_t* data);
     void taskDidComplete(uint64_t taskIdentifier, const WebCore::ResourceError&);
+    void taskDidStopLoading(WebURLSchemeHandlerTaskProxy&);
 
 private:
     WebPage& m_webPage;
index c8e636d..181ba16 100644 (file)
@@ -42,6 +42,7 @@ WebURLSchemeHandlerTaskProxy::WebURLSchemeHandlerTaskProxy(WebURLSchemeHandlerPr
     : m_urlSchemeHandler(handler)
     , m_coreLoader(&loader)
     , m_request(loader.request())
+    , m_identifier(loader.identifier())
 {
 }
 
@@ -53,16 +54,17 @@ void WebURLSchemeHandlerTaskProxy::startLoading()
 
 void WebURLSchemeHandlerTaskProxy::stopLoading()
 {
-    if (!m_coreLoader)
-        return;
-
+    ASSERT(m_coreLoader);
     m_urlSchemeHandler.page().send(Messages::WebPageProxy::StopURLSchemeHandlerTask(m_urlSchemeHandler.identifier(), m_coreLoader->identifier()));
     m_coreLoader = nullptr;
+
+    // This line will result in this being deleted.
+    m_urlSchemeHandler.taskDidStopLoading(*this);
 }
 
 void WebURLSchemeHandlerTaskProxy::didReceiveResponse(const ResourceResponse& response)
 {
-    if (!m_coreLoader)
+    if (!hasLoader())
         return;
 
     m_coreLoader->didReceiveResponse(response);
@@ -70,7 +72,7 @@ void WebURLSchemeHandlerTaskProxy::didReceiveResponse(const ResourceResponse& re
 
 void WebURLSchemeHandlerTaskProxy::didReceiveData(size_t size, const uint8_t* data)
 {
-    if (!m_coreLoader)
+    if (!hasLoader())
         return;
 
     m_coreLoader->didReceiveData(reinterpret_cast<const char*>(data), size, 0, DataPayloadType::DataPayloadBytes);
@@ -78,7 +80,7 @@ void WebURLSchemeHandlerTaskProxy::didReceiveData(size_t size, const uint8_t* da
 
 void WebURLSchemeHandlerTaskProxy::didComplete(const ResourceError& error)
 {
-    if (!m_coreLoader)
+    if (!hasLoader())
         return;
 
     if (error.isNull())
@@ -89,4 +91,12 @@ void WebURLSchemeHandlerTaskProxy::didComplete(const ResourceError& error)
     m_coreLoader = nullptr;
 }
 
+bool WebURLSchemeHandlerTaskProxy::hasLoader()
+{
+    if (m_coreLoader && m_coreLoader->reachedTerminalState())
+        m_coreLoader = nullptr;
+
+    return m_coreLoader;
+}
+
 } // namespace WebKit
index f97f263..19e9f15 100644 (file)
@@ -51,11 +51,15 @@ public:
     void didReceiveData(size_t, const uint8_t* data);
     void didComplete(const WebCore::ResourceError&);
 
+    unsigned long identifier() const { return m_identifier; }
+
 private:
+    bool hasLoader();
+
     WebURLSchemeHandlerProxy& m_urlSchemeHandler;
     RefPtr<WebCore::ResourceLoader> m_coreLoader;
     WebCore::ResourceRequest m_request;
-
+    unsigned long m_identifier;
 };
 
 } // namespace WebKit
index fcd1605..899c80d 100644 (file)
@@ -1,3 +1,18 @@
+2017-04-14  Brady Eidson  <beidson@apple.com>
+
+        Fix basic WKURLSchemeHandler bugs.
+        <rdar://problem/30647559> and https://bugs.webkit.org/show_bug.cgi?id=170862
+
+        Reviewed by Andy Estes.
+
+        * TestWebKitAPI/TestWebKitAPI.xcodeproj/project.pbxproj:
+        * TestWebKitAPI/Tests/WebKit2Cocoa/WKURLSchemeHandler-1.mm: Added.
+        (-[SchemeHandler initWithData:mimeType:]):
+        (-[SchemeHandler dealloc]):
+        (-[SchemeHandler webView:startTask:]):
+        (-[SchemeHandler webView:stopTask:]):
+        (TEST):
+
 2017-04-14  Bill Ming  <mbbill@gmail.com>
 
         webkit-patch failed to detect git repository
index 8e82604..ced8ea8 100644 (file)
                51B1EE971C80FAEF0064FB98 /* IndexedDBPersistence-2.html in Copy Resources */ = {isa = PBXBuildFile; fileRef = 51B1EE951C80FADD0064FB98 /* IndexedDBPersistence-2.html */; };
                51BCEE4E1C84F53B0042C82E /* IndexedDBMultiProcess-1.html in Copy Resources */ = {isa = PBXBuildFile; fileRef = 51BCEE4C1C84F52C0042C82E /* IndexedDBMultiProcess-1.html */; };
                51BCEE4F1C84F53B0042C82E /* IndexedDBMultiProcess-2.html in Copy Resources */ = {isa = PBXBuildFile; fileRef = 51BCEE4D1C84F52C0042C82E /* IndexedDBMultiProcess-2.html */; };
+               51C683DE1EA134E800650183 /* WKURLSchemeHandler-1.mm in Sources */ = {isa = PBXBuildFile; fileRef = 51C683DD1EA134DB00650183 /* WKURLSchemeHandler-1.mm */; };
                51CD1C6C1B38CE4300142CA5 /* ModalAlerts.mm in Sources */ = {isa = PBXBuildFile; fileRef = 51CD1C6A1B38CE3600142CA5 /* ModalAlerts.mm */; };
                51CD1C721B38D48400142CA5 /* modal-alerts-in-new-about-blank-window.html in Copy Resources */ = {isa = PBXBuildFile; fileRef = 51CD1C711B38D48400142CA5 /* modal-alerts-in-new-about-blank-window.html */; };
                51D124981E763B02002B2820 /* WKHTTPCookieStore.mm in Sources */ = {isa = PBXBuildFile; fileRef = 51D124971E763AF8002B2820 /* WKHTTPCookieStore.mm */; };
                51BCEE491C84F4AF0042C82E /* IndexedDBMultiProcess.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = IndexedDBMultiProcess.mm; sourceTree = "<group>"; };
                51BCEE4C1C84F52C0042C82E /* IndexedDBMultiProcess-1.html */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text.html; path = "IndexedDBMultiProcess-1.html"; sourceTree = "<group>"; };
                51BCEE4D1C84F52C0042C82E /* IndexedDBMultiProcess-2.html */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text.html; path = "IndexedDBMultiProcess-2.html"; sourceTree = "<group>"; };
+               51C683DD1EA134DB00650183 /* WKURLSchemeHandler-1.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = "WKURLSchemeHandler-1.mm"; sourceTree = "<group>"; };
                51CB4AD71B3A079C00C1B1C6 /* ModalAlertsSPI.cpp */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.cpp; path = ModalAlertsSPI.cpp; sourceTree = "<group>"; };
                51CD1C6A1B38CE3600142CA5 /* ModalAlerts.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ModalAlerts.mm; sourceTree = "<group>"; };
                51CD1C711B38D48400142CA5 /* modal-alerts-in-new-about-blank-window.html */ = {isa = PBXFileReference; lastKnownFileType = text.html; path = "modal-alerts-in-new-about-blank-window.html"; sourceTree = "<group>"; };
                                A14AAB611E78D7DE00C1ADC2 /* WKPDFView.mm */,
                                2D00065D1C1F58940088E6A7 /* WKPDFViewResizeCrash.mm */,
                                5E4B1D2C1D404C6100053621 /* WKScrollViewDelegateCrash.mm */,
+                               51C683DD1EA134DB00650183 /* WKURLSchemeHandler-1.mm */,
                                5CE354D81E70D9C300BEFE3B /* WKUserContentExtensionStore.mm */,
                                2EFF06D61D8AF34A0004BB30 /* WKWebViewCandidateTests.mm */,
                                7C417F311D19E14800B8EF53 /* WKWebViewDefaultNavigationDelegate.mm */,
                                CE06DF9B1E1851F200E570C9 /* SecurityOrigin.cpp in Sources */,
                                5769C50B1D9B0002000847FB /* SerializedCryptoKeyWrap.mm in Sources */,
                                7CCE7ECB1A411A7E00447C4C /* SetAndUpdateCacheModel.mm in Sources */,
+                               51C683DE1EA134E800650183 /* WKURLSchemeHandler-1.mm in Sources */,
                                7CCE7ECC1A411A7E00447C4C /* SetDocumentURI.mm in Sources */,
                                7C83E0521D0A641800FEBCF3 /* SharedBuffer.cpp in Sources */,
                                A17991881E1C994E00A505ED /* SharedBuffer.mm in Sources */,
diff --git a/Tools/TestWebKitAPI/Tests/WebKit2Cocoa/WKURLSchemeHandler-1.mm b/Tools/TestWebKitAPI/Tests/WebKit2Cocoa/WKURLSchemeHandler-1.mm
new file mode 100644 (file)
index 0000000..56ccf9f
--- /dev/null
@@ -0,0 +1,151 @@
+/*
+ * Copyright (C) 2017 Apple Inc. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ * 1. Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ * 2. Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY APPLE INC. AND ITS CONTRIBUTORS ``AS IS''
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
+ * THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL APPLE INC. OR ITS CONTRIBUTORS
+ * BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
+ * THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+#import "config.h"
+
+#import "PlatformUtilities.h"
+#import "Test.h"
+#import <WebKit/WKURLSchemeHandler.h>
+#import <WebKit/WKURLSchemeHandlerTask.h>
+#import <WebKit/WKWebViewConfigurationPrivate.h>
+#import <WebKit/WebKit.h>
+#import <wtf/RetainPtr.h>
+
+#if WK_API_ENABLED
+
+static bool receivedScriptMessage;
+
+@interface SchemeHandler : NSObject <WKURLSchemeHandler>
+@property (readonly) NSMutableArray<NSURL *> *startedURLs;
+@property (readonly) NSMutableArray<NSURL *> *stoppedURLs;
+- (instancetype)initWithData:(NSData *)data mimeType:(NSString *)inMIMEType;
+@end
+
+@implementation SchemeHandler {
+    RetainPtr<NSData> resourceData;
+    RetainPtr<NSString> mimeType;
+}
+
+- (instancetype)initWithData:(NSData *)data mimeType:(NSString *)inMIMEType
+{
+    self = [super init];
+    if (!self)
+        return nil;
+
+    resourceData = data;
+    mimeType = inMIMEType;
+    _startedURLs = [[NSMutableArray alloc] init];
+    _stoppedURLs = [[NSMutableArray alloc] init];
+
+    return self;
+}
+
+- (void)dealloc
+{
+    [_startedURLs release];
+    [_stoppedURLs release];
+    [super dealloc];
+}
+
+- (void)webView:(WKWebView *)webView startTask:(id <WKURLSchemeHandlerTask>)task
+{
+    [_startedURLs addObject:task.request.URL];
+
+    // Always fail the image load.
+    if ([task.request.URL.absoluteString isEqualToString:@"testing:image"]) {
+        [task didFailWithError:[NSError errorWithDomain:@"TestWebKitAPI" code:1 userInfo:nil]];
+        receivedScriptMessage = true;
+        return;
+    }
+
+    RetainPtr<NSURLResponse> response = adoptNS([[NSURLResponse alloc] initWithURL:task.request.URL MIMEType:mimeType.get() expectedContentLength:1 textEncodingName:nil]);
+    [task didReceiveResponse:response.get()];
+    [task didReceiveData:resourceData.get()];
+    [task didFinish];
+}
+
+- (void)webView:(WKWebView *)webView stopTask:(id <WKURLSchemeHandlerTask>)task
+{
+    [_stoppedURLs addObject:task.request.URL];
+
+    receivedScriptMessage = true;
+}
+
+@end
+
+static const char mainBytes[] =
+"<html>" \
+"<img src='testing:image'>" \
+"</html>";
+
+TEST(URLSchemeHandler, Basic)
+{
+    receivedScriptMessage = false;
+
+    RetainPtr<WKWebViewConfiguration> configuration = adoptNS([[WKWebViewConfiguration alloc] init]);
+
+    RetainPtr<SchemeHandler> handler = adoptNS([[SchemeHandler alloc] initWithData:[NSData dataWithBytesNoCopy:(void*)mainBytes length:sizeof(mainBytes)] mimeType:@"text/html"]);
+    [configuration setURLSchemeHandler:handler.get() forURLScheme:@"testing"];
+
+    RetainPtr<WKWebView> webView = adoptNS([[WKWebView alloc] initWithFrame:NSMakeRect(0, 0, 800, 600) configuration:configuration.get()]);
+
+    NSURLRequest *request = [NSURLRequest requestWithURL:[NSURL URLWithString:@"testing:main"]];
+    [webView loadRequest:request];
+
+    TestWebKitAPI::Util::run(&receivedScriptMessage);
+
+    EXPECT_EQ([handler.get().startedURLs count], 2u);
+    EXPECT_TRUE([[handler.get().startedURLs objectAtIndex:0] isEqual:[NSURL URLWithString:@"testing:main"]]);
+    EXPECT_TRUE([[handler.get().startedURLs objectAtIndex:1] isEqual:[NSURL URLWithString:@"testing:image"]]);
+    EXPECT_EQ([handler.get().stoppedURLs count], 0u);
+}
+
+TEST(URLSchemeHandler, NoMIMEType)
+{
+    // Since there's no MIMEType, and no NavigationDelegate to tell WebKit to do the load anyways, WebKit will ignore (silently fail) the load.
+    // This test makes sure that is communicated back to the URLSchemeHandler.
+
+    receivedScriptMessage = false;
+
+    RetainPtr<WKWebViewConfiguration> configuration = adoptNS([[WKWebViewConfiguration alloc] init]);
+
+    RetainPtr<SchemeHandler> handler = adoptNS([[SchemeHandler alloc] initWithData:[NSData dataWithBytesNoCopy:(void*)mainBytes length:sizeof(mainBytes)] mimeType:nil]);
+    [configuration setURLSchemeHandler:handler.get() forURLScheme:@"testing"];
+
+    RetainPtr<WKWebView> webView = adoptNS([[WKWebView alloc] initWithFrame:NSMakeRect(0, 0, 800, 600) configuration:configuration.get()]);
+
+    NSURLRequest *request = [NSURLRequest requestWithURL:[NSURL URLWithString:@"testing:main"]];
+    [webView loadRequest:request];
+
+    TestWebKitAPI::Util::run(&receivedScriptMessage);
+
+    EXPECT_EQ([handler.get().startedURLs count], 1u);
+    EXPECT_TRUE([[handler.get().startedURLs objectAtIndex:0] isEqual:[NSURL URLWithString:@"testing:main"]]);
+    EXPECT_EQ([handler.get().stoppedURLs count], 1u);
+    EXPECT_TRUE([[handler.get().stoppedURLs objectAtIndex:0] isEqual:[NSURL URLWithString:@"testing:main"]]);
+}
+
+
+#endif