Add WKWebView SPI to fetch and restore Session Storage data

https://bugs.webkit.org/show_bug.cgi?id=283503
rdar://140356755

Reviewed by Sihui Liu and Wenson Hsieh.

Some users may want Session Storage to be restored after a software
update. So we want to allow clients (Safari) to fetch this storage
and then later restore it.

The restoration process will have two parts:
1. Client fetches the session storage data from WebKit and holds onto it.
2. Client gives back the data to WebKit to restore it.

This patch adds SPI for fetch and restore--and will eventually go through
API review.

The SPI works for both ephemeral and persistent data stores, so there is
an API test for each. The SPI also works for both first party storage and
for third party storage, so there are API tests for each.

* Source/WebKit/NetworkProcess/NetworkProcess.cpp:
(WebKit::NetworkProcess::fetchSessionStorage):
(WebKit::NetworkProcess::restoreSessionStorage):
* Source/WebKit/NetworkProcess/NetworkProcess.h:
* Source/WebKit/NetworkProcess/NetworkProcess.messages.in:
* Source/WebKit/NetworkProcess/storage/LocalStorageManager.cpp:
(WebKit::LocalStorageManager::setStorageMap):
(WebKit::LocalStorageManager::populateStorageArea): Deleted.
* Source/WebKit/NetworkProcess/storage/LocalStorageManager.h:
* Source/WebKit/NetworkProcess/storage/NetworkStorageManager.cpp:
(WebKit::NetworkStorageManager::fetchSessionStorageForWebPage):
(WebKit::NetworkStorageManager::restoreSessionStorageForWebPage):
(WebKit::NetworkStorageManager::restoreLocalStorage):
* Source/WebKit/NetworkProcess/storage/NetworkStorageManager.h:
* Source/WebKit/NetworkProcess/storage/SessionStorageManager.cpp:
(WebKit::SessionStorageManager::fetchStorageMap):
(WebKit::SessionStorageManager::setStorageMap):
* Source/WebKit/NetworkProcess/storage/SessionStorageManager.h:
* Source/WebKit/UIProcess/API/Cocoa/WKWebView.mm:
(-[WKWebView _fetchDataOfTypes:completionHandler:]):
(-[WKWebView _restoreData:completionHandler:]):
* Source/WebKit/UIProcess/API/Cocoa/WKWebViewPrivate.h:
* Source/WebKit/UIProcess/API/Cocoa/WKWebsiteDataStore.mm:
(-[WKWebsiteDataStore _fetchDataOfTypes:completionHandler:]):
(-[WKWebsiteDataStore _restoreData:completionHandler:]):
* Source/WebKit/UIProcess/WebPageProxy.cpp:
(WebKit::WebPageProxy::fetchSessionStorage):
(WebKit::WebPageProxy::restoreSessionStorage):
* Source/WebKit/UIProcess/WebPageProxy.h:
* Source/WebKit/UIProcess/WebsiteData/WebsiteDataStore.cpp:
(WebKit::WebsiteDataStore::fetchLocalStorage):
* Tools/TestWebKitAPI/SourcesCocoa.txt:
* Tools/TestWebKitAPI/TestWebKitAPI.xcodeproj/project.pbxproj:
* Tools/TestWebKitAPI/Tests/WebKitCocoa/RestoreLocalStorage.mm:
(TEST(WebKit, RestoreLocalStorageFromPersistentDataStore)):
(TEST(WebKit, RestoreLocalStorageFromEphemeralDataStore)):
(RestoreLocalStorageFromPersistentDataStoreThirdPartyIFrame)):
(RestoreLocalStorageFromEphemeralDataStoreThirdPartyIFrame)):
* Tools/TestWebKitAPI/Tests/WebKitCocoa/RestoreSessionStorage.mm: Copied froTools/TestWebKitAPI/Tests/WebKitCocoa/RestoreLocalStorage.mm.
(testRestoreSessionStorage):
(TEST(WebKit, RestoreSessionStorageFromPersistentDataStore)):
(TEST(WebKit, RestoreSessionStorageFromEphemeralDataStore)):
(-[RestoreSessionStorageMessageHandler userContentController:didReceiveScriptMessage:]):
(postMessage):
(item):
(RestoreSessionStorageFromPersistentDataStoreThirdPartyIFrame)):
(RestoreSessionStorageFromEphemeralDataStoreThirdPartyIFrame)):
* Tools/TestWebKitAPI/Tests/WebKitCocoa/SiteIsolation.mm:
* Tools/TestWebKitAPI/Tests/WebKitCocoa/TextWidth.mm:

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

Source/WebKit/NetworkProcess/NetworkProcess.cpp

@@ -3127,4 +3127,26 @@ void NetworkProcess::restoreLocalStorage(PAL::SessionID sessionID, HashMap<WebCo
     session->protectedStorageManager()->restoreLocalStorage(WTFMove(localStorageMap), WTFMove(completionHandler));
 }
 
+void NetworkProcess::fetchSessionStorage(PAL::SessionID sessionID, WebPageProxyIdentifier pageID, CompletionHandler<void(HashMap<WebCore::ClientOrigin, HashMap<String, String>>&&)>&& completionHandler)
+{
+    CheckedPtr session = networkSession(sessionID);
+    if (!session) {
+        completionHandler({ });
+        return;
+    }
+
+    session->protectedStorageManager()->fetchSessionStorageForWebPage(pageID, WTFMove(completionHandler));
+}
+
+void NetworkProcess::restoreSessionStorage(PAL::SessionID sessionID, WebPageProxyIdentifier pageID, HashMap<WebCore::ClientOrigin, HashMap<String, String>>&& sessionStorageMap, CompletionHandler<void(bool)>&& completionHandler)
+{
+    CheckedPtr session = networkSession(sessionID);
+    if (!session) {
+        completionHandler(false);
+        return;
+    }
+
+    session->protectedStorageManager()->restoreSessionStorageForWebPage(pageID, WTFMove(sessionStorageMap), WTFMove(completionHandler));
+}
+
 } // namespace WebKit

Source/WebKit/NetworkProcess/NetworkProcess.h

@@ -455,6 +455,9 @@ class NetworkProcess final : public AuxiliaryProcess, private DownloadManager::C
     void fetchLocalStorage(PAL::SessionID, CompletionHandler<void(HashMap<WebCore::ClientOrigin, HashMap<String, String>>&&)>&&);
     void restoreLocalStorage(PAL::SessionID, HashMap<WebCore::ClientOrigin, HashMap<String, String>>&&, CompletionHandler<void(bool)>&&);
 
+    void fetchSessionStorage(PAL::SessionID, WebPageProxyIdentifier, CompletionHandler<void(HashMap<WebCore::ClientOrigin, HashMap<String, String>>&&)>&&);
+    void restoreSessionStorage(PAL::SessionID, WebPageProxyIdentifier, HashMap<WebCore::ClientOrigin, HashMap<String, String>>&&, CompletionHandler<void(bool)>&&);
+
 private:
     void platformInitializeNetworkProcess(const NetworkProcessCreationParameters&);
 

Source/WebKit/NetworkProcess/NetworkProcess.messages.in

@@ -262,4 +262,7 @@ messages -> NetworkProcess : AuxiliaryProcess WantsAsyncDispatchMessage {
 
     FetchLocalStorage(PAL::SessionID sessionID) -> (HashMap<WebCore::ClientOrigin, HashMap<String, String>> localStorageMap)
     RestoreLocalStorage(PAL::SessionID sessionID, HashMap<WebCore::ClientOrigin, HashMap<String, String>> localStorageMap) -> (bool succeeded)
+
+    FetchSessionStorage(PAL::SessionID sessionID, WebKit::WebPageProxyIdentifier pageID) -> (HashMap<WebCore::ClientOrigin, HashMap<String, String>> sessionStorageMap)
+    RestoreSessionStorage(PAL::SessionID sessionID, WebKit::WebPageProxyIdentifier pageID, HashMap<WebCore::ClientOrigin, HashMap<String, String>> localStorageMap) -> (bool succeeded)
 }

Source/WebKit/NetworkProcess/storage/LocalStorageManager.cpp

@@ -261,7 +261,7 @@ HashMap<String, String> LocalStorageManager::fetchStorageMap() const
     return { };
 }
 
-bool LocalStorageManager::populateStorageArea(WebCore::ClientOrigin clientOrigin, HashMap<String, String>&& storageMap, Ref<WorkQueue>&& workQueue)
+bool LocalStorageManager::setStorageMap(WebCore::ClientOrigin clientOrigin, HashMap<String, String>&& storageMap, Ref<WorkQueue>&& workQueue)
 {
     bool succeeded = true;
 

Source/WebKit/NetworkProcess/storage/LocalStorageManager.h

@@ -66,7 +66,7 @@ class LocalStorageManager {
     void disconnectFromStorageArea(IPC::Connection::UniqueID, StorageAreaIdentifier);
 
     HashMap<String, String> fetchStorageMap() const;
-    bool populateStorageArea(WebCore::ClientOrigin, HashMap<String, String>&&, Ref<WorkQueue>&&);
+    bool setStorageMap(WebCore::ClientOrigin, HashMap<String, String>&&, Ref<WorkQueue>&&);
 
 private:
     void connectionClosedForLocalStorageArea(IPC::Connection::UniqueID);

Source/WebKit/NetworkProcess/storage/NetworkStorageManager.cpp

@@ -830,6 +830,58 @@ void NetworkStorageManager::cloneSessionStorageNamespace(StorageNamespaceIdentif
     }
 }
 
+void NetworkStorageManager::fetchSessionStorageForWebPage(WebPageProxyIdentifier pageIdentifier, CompletionHandler<void(HashMap<WebCore::ClientOrigin, HashMap<String, String>>&&)>&& completionHandler)
+{
+    ASSERT(RunLoop::isMain());
+    ASSERT(!m_closed);
+
+    protectedWorkQueue()->dispatch([this, protectedThis = Ref { *this }, pageIdentifier, completionHandler = WTFMove(completionHandler)]() mutable {
+        assertIsCurrent(workQueue());
+
+        HashMap<WebCore::ClientOrigin, HashMap<String, String>> sessionStorageMap;
+        StorageNamespaceIdentifier storageNameSpaceIdentifier { pageIdentifier.toUInt64() };
+
+        for (auto& [origin, originStorageManager] : m_originStorageManagers) {
+            auto* sessionStorageManager = originStorageManager->existingSessionStorageManager();
+            if (!sessionStorageManager)
+                continue;
+
+            auto storageMap = sessionStorageManager->fetchStorageMap(storageNameSpaceIdentifier);
+            if (!storageMap.isEmpty())
+                sessionStorageMap.add(origin, WTFMove(storageMap));
+        }
+
+        RunLoop::protectedMain()->dispatch([completionHandler = WTFMove(completionHandler), sessionStorageMap = crossThreadCopy(WTFMove(sessionStorageMap))] mutable {
+            completionHandler(WTFMove(sessionStorageMap));
+        });
+    });
+}
+
+void NetworkStorageManager::restoreSessionStorageForWebPage(WebPageProxyIdentifier pageIdentifier, HashMap<WebCore::ClientOrigin, HashMap<String, String>>&& sessionStorageMap, CompletionHandler<void(bool)>&& completionHandler)
+{
+    ASSERT(RunLoop::isMain());
+    ASSERT(!m_closed);
+
+    protectedWorkQueue()->dispatch([this, protectedThis = Ref { *this }, pageIdentifier, sessionStorageMap = crossThreadCopy(WTFMove(sessionStorageMap)), completionHandler = WTFMove(completionHandler)]() mutable {
+        assertIsCurrent(workQueue());
+
+        bool succeeded = true;
+        StorageNamespaceIdentifier storageNameSpaceIdentifier { pageIdentifier.toUInt64() };
+
+        for (auto& [clientOrigin, storageMap] : sessionStorageMap) {
+            auto& sessionStorageManager = originStorageManager(clientOrigin, ShouldWriteOriginFile::Yes).sessionStorageManager(*m_storageAreaRegistry);
+            auto result = sessionStorageManager.setStorageMap(storageNameSpaceIdentifier, clientOrigin, WTFMove(storageMap));
+
+            if (!result)
+                succeeded = false;
+        }
+
+        RunLoop::protectedMain()->dispatch([protectedThis = WTFMove(protectedThis), completionHandler = WTFMove(completionHandler), succeeded] mutable {
+            completionHandler(succeeded);
+        });
+    });
+}
+
 void NetworkStorageManager::didIncreaseQuota(WebCore::ClientOrigin&& origin, QuotaIncreaseRequestIdentifier identifier, std::optional<uint64_t> newQuota)
 {
     ASSERT(RunLoop::isMain());
@@ -1353,7 +1405,7 @@ void NetworkStorageManager::restoreLocalStorage(HashMap<WebCore::ClientOrigin, H
 
         for (auto& [clientOrigin, storageMap] : localStorageMap) {
             auto& localStorageManager = originStorageManager(clientOrigin, ShouldWriteOriginFile::Yes).localStorageManager(*m_storageAreaRegistry);
-            auto result = localStorageManager.populateStorageArea(clientOrigin, WTFMove(storageMap), protectedWorkQueue());
+            auto result = localStorageManager.setStorageMap(clientOrigin, WTFMove(storageMap), protectedWorkQueue());
 
             if (!result)
                 succeeded = false;

Source/WebKit/NetworkProcess/storage/NetworkStorageManager.h

@@ -111,6 +111,8 @@ class NetworkStorageManager final : public IPC::WorkQueueMessageReceiver, public
     void resetStoragePersistedState(CompletionHandler<void()>&&);
     void clearStorageForWebPage(WebPageProxyIdentifier);
     void cloneSessionStorageForWebPage(WebPageProxyIdentifier, WebPageProxyIdentifier);
+    void fetchSessionStorageForWebPage(WebPageProxyIdentifier, CompletionHandler<void(HashMap<WebCore::ClientOrigin, HashMap<String, String>>&&)>&&);
+    void restoreSessionStorageForWebPage(WebPageProxyIdentifier, HashMap<WebCore::ClientOrigin, HashMap<String, String>>&&, CompletionHandler<void(bool)>&&);
     void didIncreaseQuota(WebCore::ClientOrigin&&, QuotaIncreaseRequestIdentifier, std::optional<uint64_t> newQuota);
     enum class ShouldComputeSize : bool { No, Yes };
     void fetchData(OptionSet<WebsiteDataType>, ShouldComputeSize, CompletionHandler<void(Vector<WebsiteData::Entry>&&)>&&);

Source/WebKit/NetworkProcess/storage/SessionStorageManager.cpp

@@ -129,4 +129,36 @@ void SessionStorageManager::cloneStorageArea(StorageNamespaceIdentifier sourceNa
         addStorageArea(storageArea->clone(), targetNamespaceIdentifier);
 }
 
+HashMap<String, String> SessionStorageManager::fetchStorageMap(StorageNamespaceIdentifier namespaceIdentifier)
+{
+    auto identifier = m_storageAreasByNamespace.getOptional(namespaceIdentifier);
+    if (!identifier)
+        return { };
+
+    RefPtr storageArea = m_storageAreas.get(*identifier);
+    if (!storageArea)
+        return { };
+
+    return storageArea->allItems();
+}
+
+bool SessionStorageManager::setStorageMap(StorageNamespaceIdentifier storageNamespaceIdentifier, WebCore::ClientOrigin clientOrigin, HashMap<String, String>&& storageMap)
+{
+    auto identifier = m_storageAreasByNamespace.getOptional(storageNamespaceIdentifier);
+    if (!identifier)
+        identifier = addStorageArea(MemoryStorageArea::create(clientOrigin), storageNamespaceIdentifier);
+
+    RefPtr storageArea = m_storageAreas.get(*identifier);
+    if (!storageArea)
+        return false;
+
+    bool succeeded = true;
+    for (auto& [key, value] : storageMap) {
+        if (!storageArea->setItem({ }, { }, WTFMove(key), WTFMove(value), { }))
+            succeeded = false;
+    }
+
+    return succeeded;
+}
+
 } // namespace WebKit

Source/WebKit/NetworkProcess/storage/SessionStorageManager.h

@@ -55,6 +55,9 @@ class SessionStorageManager {
     void disconnectFromStorageArea(IPC::Connection::UniqueID, StorageAreaIdentifier);
     void cloneStorageArea(StorageNamespaceIdentifier, StorageNamespaceIdentifier);
 
+    HashMap<String, String> fetchStorageMap(StorageNamespaceIdentifier);
+    bool setStorageMap(StorageNamespaceIdentifier, WebCore::ClientOrigin, HashMap<String, String>&&);
+
 private:
     StorageAreaIdentifier addStorageArea(Ref<MemoryStorageArea>&&, StorageNamespaceIdentifier);
 

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

@@ -181,6 +181,8 @@
 #import <wtf/cocoa/SpanCocoa.h>
 #import <wtf/cocoa/TypeCastsCocoa.h>
 #import <wtf/cocoa/VectorCocoa.h>
+#import <wtf/persistence/PersistentDecoder.h>
+#import <wtf/persistence/PersistentEncoder.h>
 #import <wtf/spi/darwin/dyldSPI.h>
 #import <wtf/text/MakeString.h>
 #import <wtf/text/StringToIntegerConversion.h>
@@ -5245,6 +5247,111 @@ - (_WKWebProcessState)_webProcessState
     }
 }
 
+namespace WebKit {
+enum class WebViewDataType : uint32_t {
+    SessionStorage
+};
+}
+
+namespace WTF {
+template<> struct EnumTraitsForPersistence<WebKit::WebViewDataType> {
+    using values = EnumValues<
+        WebKit::WebViewDataType,
+        WebKit::WebViewDataType::SessionStorage
+    >;
+};
+}
+
+struct WKWebViewData {
+    std::optional<HashMap<WebCore::ClientOrigin, HashMap<String, String>>> sessionStorage;
+};
+
+- (void)_fetchDataOfTypes:(_WKWebViewDataType)dataTypes completionHandler:(void (^)(NSData *))completionHandler
+{
+    Vector<WebKit::WebViewDataType> dataTypesToEncode;
+    if (dataTypes & _WKWebViewDataTypeSessionStorage)
+        dataTypesToEncode.append(WebKit::WebViewDataType::SessionStorage);
+
+    auto data = Box<WKWebViewData>::create();
+
+    auto callbackAggregator = CallbackAggregator::create([completionHandler = makeBlockPtr(completionHandler), dataTypesToEncode = WTFMove(dataTypesToEncode), data] {
+        WTF::Persistence::Encoder encoder;
+        constexpr unsigned currentWKWebViewDataSerializationVersion = 1;
+        encoder << currentWKWebViewDataSerializationVersion;
+        encoder << dataTypesToEncode;
+
+        for (auto& dataTypeToEncode : dataTypesToEncode) {
+            switch (dataTypeToEncode) {
+            case WebKit::WebViewDataType::SessionStorage:
+                encoder << data->sessionStorage.value();
+                break;
+            default:
+                ASSERT_NOT_REACHED();
+                break;
+            }
+        }
+
+        completionHandler(toNSData(encoder.span()).get());
+    });
+
+    if (dataTypes & _WKWebViewDataTypeSessionStorage) {
+        RefPtr page = [self _protectedPage];
+        page->fetchSessionStorage([callbackAggregator, protectedPage = page, data](auto&& sessionStorage) {
+            data->sessionStorage = WTFMove(sessionStorage);
+        });
+    }
+}
+
+- (void)_restoreData:(NSData *)data completionHandler:(void(^)(BOOL))completionHandler
+{
+    WTF::Persistence::Decoder decoder(span(data));
+
+    std::optional<unsigned> currentWKWebViewDataSerializationVersion;
+    decoder >> currentWKWebViewDataSerializationVersion;
+    if (!currentWKWebViewDataSerializationVersion) {
+        completionHandler(NO);
+        return;
+    }
+
+    std::optional<Vector<WebKit::WebViewDataType>> encodedDataTypes;
+    decoder >> encodedDataTypes;
+    if (!encodedDataTypes) {
+        completionHandler(NO);
+        return;
+    }
+
+    auto succeeded = Box<bool>::create(true);
+    auto callbackAggregator = CallbackAggregator::create([completionHandler = makeBlockPtr(completionHandler), succeeded] {
+        completionHandler(*succeeded);
+    });
+
+    for (auto& encodedDataType : *encodedDataTypes) {
+        switch (encodedDataType) {
+        case WebKit::WebViewDataType::SessionStorage: {
+            std::optional<HashMap<WebCore::ClientOrigin, HashMap<String, String>>> sessionStorage;
+            decoder >> sessionStorage;
+
+            if (!sessionStorage) {
+                *succeeded = false;
+                return;
+            }
+
+            if (!sessionStorage->isEmpty()) {
+                RefPtr page = [self _protectedPage];
+                page->restoreSessionStorage(WTFMove(*sessionStorage), [callbackAggregator, succeeded](bool restoreSucceeded) {
+                    if (!restoreSucceeded)
+                        *succeeded = false;
+                });
+            }
+            break;
+        }
+        default:
+            ASSERT_NOT_REACHED();
+            break;
+        }
+    }
+}
+
 @end
 
 @implementation WKWebView (WKDeprecated)

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

@@ -601,6 +601,13 @@ typedef NS_OPTIONS(NSUInteger, WKDisplayCaptureSurfaces) {
 // This property is KVO compliant.
 @property (nonatomic, readonly) _WKWebProcessState _webProcessState WK_API_AVAILABLE(macos(WK_MAC_TBA), ios(WK_IOS_TBA), visionos(WK_XROS_TBA));
 
+typedef NS_OPTIONS(NSUInteger, _WKWebViewDataType) {
+    _WKWebViewDataTypeSessionStorage = 1 << 0
+} WK_API_AVAILABLE(macos(WK_MAC_TBA), ios(WK_IOS_TBA), visionos(WK_XROS_TBA));
+
+- (void)_fetchDataOfTypes:(_WKWebViewDataType)dataTypes completionHandler:(WK_SWIFT_UI_ACTOR void (^)(NSData *))completionHandler WK_API_AVAILABLE(macos(WK_MAC_TBA), ios(WK_IOS_TBA), visionos(WK_XROS_TBA));
+- (void)_restoreData:(NSData *)data completionHandler:(WK_SWIFT_UI_ACTOR void(^)(BOOL))completionHandler WK_API_AVAILABLE(macos(WK_MAC_TBA), ios(WK_IOS_TBA), visionos(WK_XROS_TBA));
+
 @end
 
 #if TARGET_OS_IPHONE

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

@@ -1470,7 +1470,7 @@ - (void)_fetchDataOfTypes:(NSSet<NSString *> *)dataTypes completionHandler:(void
     });
 
     if ([dataTypes containsObject:WKWebsiteDataTypeLocalStorage]) {
-        _websiteDataStore->fetchLocalStorage([callbackAggregator, data](HashMap<WebCore::ClientOrigin, HashMap<String, String>>&& localStorage) {
+        _websiteDataStore->fetchLocalStorage([callbackAggregator, data](auto&& localStorage) {
             data->localStorage = WTFMove(localStorage);
         });
     }
@@ -1480,9 +1480,9 @@ - (void)_restoreData:(NSData *)data completionHandler:(void(^)(BOOL))completionH
 {
     WTF::Persistence::Decoder decoder(span(data));
 
-    std::optional<unsigned> currentLocalStorageSerializationVersion;
-    decoder >> currentLocalStorageSerializationVersion;
-    if (!currentLocalStorageSerializationVersion) {
+    std::optional<unsigned> currentWKWebsiteDataSerializationVersion;
+    decoder >> currentWKWebsiteDataSerializationVersion;
+    if (!currentWKWebsiteDataSerializationVersion) {
         completionHandler(NO);
         return;
     }

Source/WebKit/UIProcess/WebPageProxy.cpp

@@ -15651,6 +15651,19 @@ FloatPoint WebPageProxy::mainFrameScrollPosition() const
 
 #endif // PLATFORM(COCOA) && ENABLE(ASYNC_SCROLLING)
 
+void WebPageProxy::fetchSessionStorage(CompletionHandler<void(HashMap<WebCore::ClientOrigin, HashMap<String, String>>&&)>&& completionHandler)
+{
+    if (RefPtr networkProcess = websiteDataStore().networkProcessIfExists())
+        networkProcess->sendWithAsyncReply(Messages::NetworkProcess::FetchSessionStorage(sessionID(), identifier()), WTFMove(completionHandler));
+    else
+        completionHandler({ });
+}
+
+void WebPageProxy::restoreSessionStorage(HashMap<WebCore::ClientOrigin, HashMap<String, String>>&& sessionStorage, CompletionHandler<void(bool)>&& completionHandler)
+{
+    protectedWebsiteDataStore()->protectedNetworkProcess()->sendWithAsyncReply(Messages::NetworkProcess::RestoreSessionStorage(sessionID(), identifier(), WTFMove(sessionStorage)), WTFMove(completionHandler));
+}
+
 } // namespace WebKit
 
 #undef WEBPAGEPROXY_RELEASE_LOG

Source/WebKit/UIProcess/WebPageProxy.h

@@ -2614,6 +2614,9 @@ class WebPageProxy final : public API::ObjectImpl<API::Object::Type::Page>, publ
     WebCore::FloatPoint mainFrameScrollPosition() const;
 #endif
 
+    void fetchSessionStorage(CompletionHandler<void(HashMap<WebCore::ClientOrigin, HashMap<String, String>>&&)>&&);
+    void restoreSessionStorage(HashMap<WebCore::ClientOrigin, HashMap<String, String>>&&, CompletionHandler<void(bool)>&&);
+
 private:
     void getWebCryptoMasterKey(CompletionHandler<void(std::optional<Vector<uint8_t>>&&)>&&);
     WebPageProxy(PageClient&, WebProcessProxy&, Ref<API::PageConfiguration>&&);

Source/WebKit/UIProcess/WebsiteData/WebsiteDataStore.cpp

@@ -2718,7 +2718,10 @@ void WebsiteDataStore::setRestrictedOpenerTypeForDomainForTesting(const WebCore:
 
 void WebsiteDataStore::fetchLocalStorage(CompletionHandler<void(HashMap<WebCore::ClientOrigin, HashMap<String, String>>&&)>&& completionHandler)
 {
-    protectedNetworkProcess()->fetchLocalStorage(m_sessionID, WTFMove(completionHandler));
+    if (RefPtr networkProcess = networkProcessIfExists())
+        networkProcess->fetchLocalStorage(m_sessionID, WTFMove(completionHandler));
+    else
+        completionHandler({ });
 }
 
 void WebsiteDataStore::restoreLocalStorage(HashMap<WebCore::ClientOrigin, HashMap<String, String>>&& localStorage, CompletionHandler<void(bool)>&& completionHandler)

Tools/TestWebKitAPI/SourcesCocoa.txt

@@ -236,6 +236,7 @@ Tests/WebKitCocoa/ResponsivenessTimerDoesntFireEarly.mm
 Tests/WebKitCocoa/RestoreLocalStorage.mm
 Tests/WebKitCocoa/RestoreScrollPosition.mm
 Tests/WebKitCocoa/RestoreSessionStateWithoutNavigation.mm
+Tests/WebKitCocoa/RestoreSessionStorage.mm
 Tests/WebKitCocoa/RunOpenPanel.mm
 Tests/WebKitCocoa/RunScriptAfterDocumentLoad.mm
 Tests/WebKitCocoa/SafeBrowsing.mm