Clear up session only storage on localstorage shutdown

content/browser/dom_storage/local_storage_context_mojo.cc

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "content/browser/dom_storage/local_storage_context_mojo.h"
 
 #include "base/memory/ptr_util.h"
 #include "base/metrics/histogram_macros.h"
 #include "base/single_thread_task_runner.h"
 #include "base/strings/string_number_conversions.h"
 #include "components/leveldb/public/cpp/util.h"
 #include "components/leveldb/public/interfaces/leveldb.mojom.h"
 #include "content/browser/dom_storage/dom_storage_area.h"
 #include "content/browser/dom_storage/dom_storage_database.h"
 #include "content/browser/dom_storage/dom_storage_task_runner.h"
 #include "content/browser/dom_storage/local_storage_database.pb.h"
 #include "content/browser/leveldb_wrapper_impl.h"
 #include "content/common/dom_storage/dom_storage_types.h"
 #include "content/public/browser/local_storage_usage_info.h"
 #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 {
 
 // LevelDB database schema
 // =======================
 //
 // Version 1 (in sorted order):
 //   key: "VERSION"
 //   value: "1"
@@ skipping 197 matching lines @@
   // called and need access  to the LevelDBWrapperImpl instance. The unique_ptr
   // could already be null, but this field should still be valid.
   LevelDBWrapperImpl* level_db_wrapper_ptr_;
   bool deleted_old_data_ = false;
 };
 
 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) {
   RunWhenConnected(base::BindOnce(&LocalStorageContextMojo::BindLocalStorage,
                                   weak_ptr_factory_.GetWeakPtr(), origin,
                                   std::move(request)));
 }
 
 void LocalStorageContextMojo::GetStorageUsage(
     GetStorageUsageCallback callback) {
@@ skipping 33 matching lines @@
 void LocalStorageContextMojo::Flush() {
   if (connection_state_ != CONNECTION_FINISHED) {
     RunWhenConnected(base::BindOnce(&LocalStorageContextMojo::Flush,
                                     weak_ptr_factory_.GetWeakPtr()));
     return;
   }
   for (const auto& it : level_db_wrappers_)
     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();
 }
 
 leveldb::mojom::LevelDBDatabaseAssociatedRequest
 LocalStorageContextMojo::DatabaseRequestForTesting() {
   DCHECK_EQ(connection_state_, NO_CONNECTION);
   connection_state_ = CONNECTION_IN_PROGRESS;
   leveldb::mojom::LevelDBDatabaseAssociatedRequest request =
       MakeIsolatedRequest(&database_);
   OnDatabaseOpened(true, leveldb::mojom::DatabaseError::OK);
   return request;
 }
 
 // static
 std::vector<uint8_t> LocalStorageContextMojo::MigrateString(
     const base::string16& input) {
   static const uint8_t kUTF16Format = 0;
 
   const uint8_t* data = reinterpret_cast<const uint8_t*>(input.data());
   std::vector<uint8_t> result;
   result.reserve(input.size() * sizeof(base::char16) + 1);
   result.push_back(kUTF16Format);
   result.insert(result.end(), data, data + input.size() * sizeof(base::char16));
   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;
     InitiateConnection();
   }
 
   if (connection_state_ == CONNECTION_IN_PROGRESS) {
     // Queue this OpenLocalStorage call for when we have a level db pointer.
     on_database_opened_callbacks_.push_back(std::move(callback));
     return;
@@ skipping 261 matching lines @@
     std::vector<LocalStorageUsageInfo> usage) {
   for (const auto& info : usage) {
     url::Origin origin_candidate(info.origin);
     if (!origin_candidate.IsSameOriginWith(origin) &&
         origin_candidate.IsSamePhysicalOriginWith(origin))
       DeleteStorage(origin_candidate);
   }
   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