From 65be9425dc24529d589c24a6ef06a201ec1a6594 Mon Sep 17 00:00:00 2001
From: Pawel Lampe <pawel.lampe@gmail.com>
Date: Thu, 9 Nov 2023 10:03:52 +0100
Subject: [PATCH] Restore local storage after corruption

---
 .../storage/SQLiteStorageArea.cpp             | 34 +++++++++++++------
 .../storage/SQLiteStorageArea.h               |  1 +
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/Source/WebKit/NetworkProcess/storage/SQLiteStorageArea.cpp b/Source/WebKit/NetworkProcess/storage/SQLiteStorageArea.cpp
index 1ee094bb5c5cc..7f5a37cea6567 100644
--- a/Source/WebKit/NetworkProcess/storage/SQLiteStorageArea.cpp
+++ b/Source/WebKit/NetworkProcess/storage/SQLiteStorageArea.cpp
@@ -324,6 +324,17 @@ HashMap<String, String> SQLiteStorageArea::allItems()
 }
 
 Expected<void, StorageError> SQLiteStorageArea::setItem(IPC::Connection::UniqueID connection, StorageAreaImplIdentifier storageAreaImplID, String&& key, String&& value, const String& urlString)
+{
+    String oldValue;
+    if (auto valueOrError = getItem(key))
+        oldValue = valueOrError.value();
+    auto settingOutcome = setItem(key, value, true);
+    if (settingOutcome)
+        dispatchEvents(connection, storageAreaImplID, key, oldValue, value, urlString);
+    return settingOutcome;
+}
+
+Expected<void, StorageError> SQLiteStorageArea::setItem(const String& key, const String& value, bool handleDatabaseCorruption)
 {
     ASSERT(!isMainRunLoop());
 
@@ -334,9 +345,6 @@ Expected<void, StorageError> SQLiteStorageArea::setItem(IPC::Connection::UniqueI
         return makeUnexpected(StorageError::QuotaExceeded);
 
     startTransactionIfNecessary();
-    String oldValue;
-    if (auto valueOrError = getItem(key))
-        oldValue = valueOrError.value();
 
     auto statement = cachedStatement(StatementType::SetItem);
     if (!statement || statement->bindText(1, key) || statement->bindBlob(2, value)) {
@@ -346,12 +354,13 @@ Expected<void, StorageError> SQLiteStorageArea::setItem(IPC::Connection::UniqueI
 
     const auto result = statement->step();
     if (result != SQLITE_DONE) {
-        RELEASE_LOG_ERROR(Storage, "SQLiteStorageArea::setItem failed on stepping statement (%d) - %s", m_database->lastError(), m_database->lastErrorMsg());
-        handleDatabaseCorruptionIfNeeded(result);
-        return makeUnexpected(StorageError::Database);
+        if (!handleDatabaseCorruption || !handleDatabaseCorruptionIfNeeded(result))
+            return makeUnexpected(StorageError::Database);
+        statement = cachedStatement(StatementType::SetItem);
+        if (!statement || statement->bindText(1, key) || statement->bindBlob(2, value) || statement->step() != SQLITE_DONE)
+            return makeUnexpected(StorageError::Database);
     }
 
-    dispatchEvents(connection, storageAreaImplID, key, oldValue, value, urlString);
     updateCacheIfNeeded(key, value);
 
     return { };
@@ -454,11 +463,16 @@ bool SQLiteStorageArea::handleDatabaseCorruptionIfNeeded(int databaseError)
     if (databaseError != SQLITE_CORRUPT && databaseError != SQLITE_NOTADB)
         return false;
 
-    m_database = nullptr;
-    m_cache = std::nullopt;
-    m_cacheSize = std::nullopt;
+    HashMap<String, Value> cache(WTFMove(*m_cache));
+    close();
     RELEASE_LOG(Storage, "SQLiteStorageArea::handleDatabaseCorruption deletes corrupted database file '%s'", m_path.utf8().data());
     WebCore::SQLiteFileSystem::deleteDatabaseFile(m_path);
+
+    // Reconstruct database based on cache.
+    for (auto& [key, value] : cache)
+        if (auto* valueString = std::get_if<String>(&value))
+            setItem(key, *valueString, false);
+
     return true;
 }
 
diff --git a/Source/WebKit/NetworkProcess/storage/SQLiteStorageArea.h b/Source/WebKit/NetworkProcess/storage/SQLiteStorageArea.h
index b1282c3b13191..6bcf427b0ef28 100644
--- a/Source/WebKit/NetworkProcess/storage/SQLiteStorageArea.h
+++ b/Source/WebKit/NetworkProcess/storage/SQLiteStorageArea.h
@@ -48,6 +48,7 @@ class SQLiteStorageArea final : public StorageAreaBase {
     void clear() final;
     HashMap<String, String> allItems() final;
     Expected<void, StorageError> setItem(IPC::Connection::UniqueID, StorageAreaImplIdentifier, String&& key, String&& value, const String& urlString) final;
+    Expected<void, StorageError> setItem(const String& key, const String& value, bool handleDatabaseCorruption);
     Expected<void, StorageError> removeItem(IPC::Connection::UniqueID, StorageAreaImplIdentifier, const String& key, const String& urlString) final;
     Expected<void, StorageError> clear(IPC::Connection::UniqueID, StorageAreaImplIdentifier, const String& urlString) final;