Clear up session only storage on localstorage shutdown

content/browser/dom_storage/local_storage_context_mojo.cc

Index: content/browser/dom_storage/local_storage_context_mojo.cc
diff --git a/content/browser/dom_storage/local_storage_context_mojo.cc b/content/browser/dom_storage/local_storage_context_mojo.cc
index f4aca0ad42955b3665c6ce3b8b79e8a84189104a..2be6213f8ce2f40c4be1d0821cb30765217a20f4 100644
--- a/content/browser/dom_storage/local_storage_context_mojo.cc
+++ b/content/browser/dom_storage/local_storage_context_mojo.cc
@@ -20,6 +20,7 @@
 #include "services/file/public/interfaces/constants.mojom.h"
 #include "services/service_manager/public/cpp/connector.h"
 #include "sql/connection.h"
+#include "storage/browser/quota/special_storage_policy.h"
 #include "third_party/leveldatabase/env_chromium.h"
 
 namespace content {
@@ -237,15 +238,15 @@ LocalStorageContextMojo::LocalStorageContextMojo(
     service_manager::Connector* connector,
     scoped_refptr<DOMStorageTaskRunner> task_runner,
     const base::FilePath& old_localstorage_path,
-    const base::FilePath& subdirectory)
-    : connector_(connector),
+    const base::FilePath& subdirectory,
+    scoped_refptr<storage::SpecialStoragePolicy> special_storage_policy)
+    : connector_(connector ? connector->Clone() : nullptr),
       subdirectory_(subdirectory),
+      special_storage_policy_(std::move(special_storage_policy)),
       task_runner_(std::move(task_runner)),
       old_localstorage_path_(old_localstorage_path),
       weak_ptr_factory_(this) {}
 
-LocalStorageContextMojo::~LocalStorageContextMojo() {}
-
 void LocalStorageContextMojo::OpenLocalStorage(
     const url::Origin& origin,
     mojom::LevelDBWrapperRequest request) {
@@ -299,6 +300,42 @@ void LocalStorageContextMojo::Flush() {
     it.second->level_db_wrapper()->ScheduleImmediateCommit();
 }
 
+void LocalStorageContextMojo::ShutdownAndDelete() {
+  DCHECK_NE(connection_state_, CONNECTION_SHUTDOWN);
+
+  // Nothing to do if no connection to the database was ever finished.
+  if (connection_state_ != CONNECTION_FINISHED) {
+    connection_state_ = CONNECTION_SHUTDOWN;
+    OnShutdownComplete(leveldb::mojom::DatabaseError::OK);
+    return;
+  }
+
+  connection_state_ = CONNECTION_SHUTDOWN;
+
+  // Flush any uncommitted data.
+  for (const auto& it : level_db_wrappers_)
+    it.second->level_db_wrapper()->ScheduleImmediateCommit();
+
+  // Respect the content policy settings about what to
+  // keep and what to discard.
+  if (force_keep_session_state_) {
+    OnShutdownComplete(leveldb::mojom::DatabaseError::OK);
+    return;  // Keep everything.
+  }
+
+  bool has_session_only_origins =
+      special_storage_policy_.get() &&
+      special_storage_policy_->HasSessionOnlyOrigins();
+
+  if (has_session_only_origins) {
+    RetrieveStorageUsage(
+        base::BindOnce(&LocalStorageContextMojo::OnGotStorageUsageForShutdown,
+                       base::Unretained(this)));
+  } else {
+    OnShutdownComplete(leveldb::mojom::DatabaseError::OK);
+  }
+}
+
 void LocalStorageContextMojo::PurgeMemory() {
   for (const auto& it : level_db_wrappers_)
     it.second->level_db_wrapper()->PurgeMemory();
@@ -327,7 +364,13 @@ std::vector<uint8_t> LocalStorageContextMojo::MigrateString(
   return result;
 }
 
+LocalStorageContextMojo::~LocalStorageContextMojo() {
+  DCHECK_EQ(connection_state_, CONNECTION_SHUTDOWN);
+}
+
 void LocalStorageContextMojo::RunWhenConnected(base::OnceClosure callback) {
+  DCHECK_NE(connection_state_, CONNECTION_SHUTDOWN);
+
   // If we don't have a filesystem_connection_, we'll need to establish one.
   if (connection_state_ == NO_CONNECTION) {
     connection_state_ = CONNECTION_IN_PROGRESS;
@@ -609,4 +652,30 @@ void LocalStorageContextMojo::OnGotStorageUsageForDeletePhysicalOrigin(
   DeleteStorage(origin);
 }
 
+void LocalStorageContextMojo::OnGotStorageUsageForShutdown(
+    std::vector<LocalStorageUsageInfo> usage) {
+  std::vector<leveldb::mojom::BatchedOperationPtr> operations;
+  for (const auto& info : usage) {
+    if (special_storage_policy_->IsStorageProtected(info.origin))
+      continue;
+    if (!special_storage_policy_->IsStorageSessionOnly(info.origin))
+      continue;
+
+    AddDeleteOriginOperations(&operations, url::Origin(info.origin));
+  }
+
+  if (!operations.empty()) {
+    database_->Write(std::move(operations),
+                     base::Bind(&LocalStorageContextMojo::OnShutdownComplete,
+                                base::Unretained(this)));
+  } else {
+    OnShutdownComplete(leveldb::mojom::DatabaseError::OK);
+  }
+}
+
+void LocalStorageContextMojo::OnShutdownComplete(
+    leveldb::mojom::DatabaseError error) {
+  delete this;
+}
+
 }  // namespace content