Clear up session only storage on localstorage shutdown

content/browser/dom_storage/local_storage_context_mojo_unittest.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/files/file_enumerator.h"
 #include "base/files/scoped_temp_dir.h"
 #include "base/run_loop.h"
 #include "base/strings/utf_string_conversions.h"
@@ skipping 11 matching lines @@
 #include "mojo/public/cpp/bindings/associated_binding.h"
 #include "mojo/public/cpp/bindings/binding.h"
 #include "mojo/public/cpp/bindings/binding_set.h"
 #include "services/file/file_service.h"
 #include "services/file/public/interfaces/constants.mojom.h"
 #include "services/file/user_id_map.h"
 #include "services/service_manager/public/cpp/interface_factory.h"
 #include "services/service_manager/public/cpp/service_context.h"
 #include "services/service_manager/public/cpp/service_test.h"
 #include "services/service_manager/public/interfaces/service_factory.mojom.h"
+#include "storage/browser/test/mock_special_storage_policy.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "third_party/leveldatabase/env_chromium.h"
 #include "third_party/leveldatabase/src/include/leveldb/db.h"
 
 using leveldb::StdStringToUint8Vector;
 using leveldb::Uint8VectorToStdString;
 
 namespace content {
 
 namespace {
@@ skipping 69 matching lines @@
 };
 
 }  // namespace
 
 class LocalStorageContextMojoTest : public testing::Test {
  public:
   LocalStorageContextMojoTest()
       : db_(&mock_data_),
         db_binding_(&db_),
         task_runner_(new MockDOMStorageTaskRunner(
-            base::ThreadTaskRunnerHandle::Get().get())) {
+            base::ThreadTaskRunnerHandle::Get().get())),
+        mock_special_storage_policy_(new MockSpecialStoragePolicy()) {
     EXPECT_TRUE(temp_path_.CreateUniqueTempDir());
     dom_storage_context_ = new DOMStorageContextImpl(
         temp_path_.GetPath(), base::FilePath(), nullptr, task_runner_);
   }
 
   ~LocalStorageContextMojoTest() override {
     if (dom_storage_context_)
       dom_storage_context_->Shutdown();
+    if (context_)
+      context_->Shutdown();
   }
 
   LocalStorageContextMojo* context() {
     if (!context_) {
-      context_ = base::MakeUnique<LocalStorageContextMojo>(
+      context_ = new LocalStorageContextMojo(
           nullptr, task_runner_, temp_path_.GetPath(),
-          base::FilePath(FILE_PATH_LITERAL("leveldb")));
+          base::FilePath(FILE_PATH_LITERAL("leveldb")),
+          special_storage_policy());
       db_binding_.Bind(context_->DatabaseRequestForTesting());
     }
     return context_.get();
   }
 
+  void ShutdownContext() {
+    context_->Shutdown();
+    context_ = nullptr;
+    base::RunLoop().RunUntilIdle();
+  }
+
   DOMStorageNamespace* local_storage_namespace() {
     return dom_storage_context_->GetStorageNamespace(kLocalStorageNamespaceId);
   }
 
+  MockSpecialStoragePolicy* special_storage_policy() {
+    return mock_special_storage_policy_.get();
+  }
+
   void FlushAndPurgeDOMStorageMemory() {
     dom_storage_context_->Flush();
     base::RunLoop().RunUntilIdle();
     dom_storage_context_->PurgeMemory(DOMStorageContextImpl::PURGE_AGGRESSIVE);
   }
 
   const std::map<std::vector<uint8_t>, std::vector<uint8_t>>& mock_data() {
     return mock_data_;
   }
 
@@ skipping 13 matching lines @@
  private:
   TestBrowserThreadBundle thread_bundle_;
   base::ScopedTempDir temp_path_;
   std::map<std::vector<uint8_t>, std::vector<uint8_t>> mock_data_;
   MockLevelDBDatabase db_;
   mojo::AssociatedBinding<leveldb::mojom::LevelDBDatabase> db_binding_;
 
   scoped_refptr<MockDOMStorageTaskRunner> task_runner_;
   scoped_refptr<DOMStorageContextImpl> dom_storage_context_;
 
-  std::unique_ptr<LocalStorageContextMojo> context_;
+  scoped_refptr<LocalStorageContextMojo> context_;
+
+  scoped_refptr<MockSpecialStoragePolicy> mock_special_storage_policy_;
 
   DISALLOW_COPY_AND_ASSIGN(LocalStorageContextMojoTest);
 };
 
 TEST_F(LocalStorageContextMojoTest, Basic) {
   auto key = StdStringToUint8Vector("key");
   auto value = StdStringToUint8Vector("value");
 
   mojom::LevelDBWrapperPtr wrapper;
   context()->OpenLocalStorage(url::Origin(GURL("http://foobar.com")),
@@ skipping 411 matching lines @@
   run_loop.Run();
   EXPECT_TRUE(success);
   EXPECT_EQ(LocalStorageContextMojo::MigrateString(value), result);
 
   // Origin1 should no longer exist in old storage.
   area = local->OpenStorageArea(origin1.GetURL());
   ASSERT_EQ(0u, area->Length());
   local->CloseStorageArea(area);
 }
 
+TEST_F(LocalStorageContextMojoTest, ShutdownClearsData) {
+  url::Origin origin1(GURL("http://foobar.com"));
+  url::Origin origin2(GURL("http://example.com"));
+  auto key1 = StdStringToUint8Vector("key1");
+  auto key2 = StdStringToUint8Vector("key");
+  auto value = StdStringToUint8Vector("value");
+
+  mojom::LevelDBWrapperPtr wrapper;
+  context()->OpenLocalStorage(origin1, MakeRequest(&wrapper));
+  wrapper->Put(key1, value, "source", base::Bind(&NoOpSuccess));
+  wrapper->Put(key2, value, "source", base::Bind(&NoOpSuccess));
+  wrapper.reset();
+
+  context()->OpenLocalStorage(origin2, MakeRequest(&wrapper));
+  wrapper->Put(key2, value, "source", base::Bind(&NoOpSuccess));
+  wrapper.reset();
+
+  // Make sure all data gets committed to the DB.
+  base::RunLoop().RunUntilIdle();
+
+  special_storage_policy()->AddSessionOnly(origin1.GetURL());
+  ShutdownContext();
+
+  // Data from origin2 should exist, including meta-data, but nothing should
+  // exist for origin1.
+  EXPECT_EQ(3u, mock_data().size());
+  for (const auto& it : mock_data()) {
+    if (Uint8VectorToStdString(it.first) == "VERSION")
+      continue;
+    EXPECT_EQ(std::string::npos,
+              Uint8VectorToStdString(it.first).find(origin1.Serialize()));
+    EXPECT_NE(std::string::npos,
+              Uint8VectorToStdString(it.first).find(origin2.Serialize()));
+  }
+}
+
 namespace {
 
 class ServiceTestClient : public service_manager::test::ServiceTestClient,
                           public service_manager::mojom::ServiceFactory,
                           public service_manager::InterfaceFactory<
                               service_manager::mojom::ServiceFactory> {
  public:
   explicit ServiceTestClient(service_manager::test::ServiceTest* test)
       : service_manager::test::ServiceTestClient(test) {
     registry_.AddInterface<service_manager::mojom::ServiceFactory>(this);
@@ skipping 96 matching lines @@
   }
 
  private:
   TestBrowserThreadBundle thread_bundle_;
   base::ScopedTempDir temp_path_;
 
   DISALLOW_COPY_AND_ASSIGN(LocalStorageContextMojoTestWithService);
 };
 
 TEST_F(LocalStorageContextMojoTestWithService, InMemory) {
-  auto context = base::MakeUnique<LocalStorageContextMojo>(
-      connector(), nullptr, base::FilePath(), base::FilePath());
+  auto context = base::MakeShared<LocalStorageContextMojo>(
+      connector(), nullptr, base::FilePath(), base::FilePath(), nullptr);
   auto key = StdStringToUint8Vector("key");
   auto value = StdStringToUint8Vector("value");
 
   mojom::LevelDBWrapperPtr wrapper;
   context->OpenLocalStorage(url::Origin(GURL("http://foobar.com")),
                             MakeRequest(&wrapper));
 
   DoTestPut(context.get(), key, value);
   std::vector<uint8_t> result;
   EXPECT_TRUE(DoTestGet(context.get(), key, &result));
   EXPECT_EQ(value, result);
 
-  context.reset();
+  context->Shutdown();
+  context = nullptr;
   base::RunLoop().RunUntilIdle();
 
   // Should not have created any files.
   EXPECT_TRUE(FirstEntryInDir().empty());
 
   // Re-opening should get fresh data.
-  context = base::MakeUnique<LocalStorageContextMojo>(
-      connector(), nullptr, base::FilePath(), base::FilePath());
+  context = base::MakeShared<LocalStorageContextMojo>(
+      connector(), nullptr, base::FilePath(), base::FilePath(), nullptr);
   EXPECT_FALSE(DoTestGet(context.get(), key, &result));
+  context->Shutdown();
 }
 
 TEST_F(LocalStorageContextMojoTestWithService, InMemoryInvalidPath) {
-  auto context = base::MakeUnique<LocalStorageContextMojo>(
+  auto context = base::MakeShared<LocalStorageContextMojo>(
       connector(), nullptr, base::FilePath(),
-      base::FilePath(FILE_PATH_LITERAL("../../")));
+      base::FilePath(FILE_PATH_LITERAL("../../")), nullptr);
   auto key = StdStringToUint8Vector("key");
   auto value = StdStringToUint8Vector("value");
 
   mojom::LevelDBWrapperPtr wrapper;
   context->OpenLocalStorage(url::Origin(GURL("http://foobar.com")),
                             MakeRequest(&wrapper));
 
   DoTestPut(context.get(), key, value);
   std::vector<uint8_t> result;
   EXPECT_TRUE(DoTestGet(context.get(), key, &result));
   EXPECT_EQ(value, result);
 
-  context.reset();
+  context->Shutdown();
+  context = nullptr;
   base::RunLoop().RunUntilIdle();
 
   // Should not have created any files.
   EXPECT_TRUE(FirstEntryInDir().empty());
 }
 
 TEST_F(LocalStorageContextMojoTestWithService, OnDisk) {
   base::FilePath test_path(FILE_PATH_LITERAL("test_path"));
-  auto context = base::MakeUnique<LocalStorageContextMojo>(
-      connector(), nullptr, base::FilePath(), test_path);
+  auto context = base::MakeShared<LocalStorageContextMojo>(
+      connector(), nullptr, base::FilePath(), test_path, nullptr);
   auto key = StdStringToUint8Vector("key");
   auto value = StdStringToUint8Vector("value");
 
   DoTestPut(context.get(), key, value);
   std::vector<uint8_t> result;
   EXPECT_TRUE(DoTestGet(context.get(), key, &result));
   EXPECT_EQ(value, result);
 
-  context.reset();
+  context->Shutdown();
+  context = nullptr;
   base::RunLoop().RunUntilIdle();
 
   // Should have created files.
   EXPECT_EQ(test_path, FirstEntryInDir().BaseName());
 
   // Should be able to re-open.
-  context = base::MakeUnique<LocalStorageContextMojo>(
-      connector(), nullptr, base::FilePath(), test_path);
+  context = base::MakeShared<LocalStorageContextMojo>(
+      connector(), nullptr, base::FilePath(), test_path, nullptr);
   EXPECT_TRUE(DoTestGet(context.get(), key, &result));
   EXPECT_EQ(value, result);
+  context->Shutdown();
 }
 
 TEST_F(LocalStorageContextMojoTestWithService, InvalidVersionOnDisk) {
   base::FilePath test_path(FILE_PATH_LITERAL("test_path"));
 
   // Create context and add some data to it.
-  auto context = base::MakeUnique<LocalStorageContextMojo>(
-      connector(), nullptr, base::FilePath(), test_path);
+  auto context = base::MakeShared<LocalStorageContextMojo>(
+      connector(), nullptr, base::FilePath(), test_path, nullptr);
   auto key = StdStringToUint8Vector("key");
   auto value = StdStringToUint8Vector("value");
 
   DoTestPut(context.get(), key, value);
   std::vector<uint8_t> result;
   EXPECT_TRUE(DoTestGet(context.get(), key, &result));
   EXPECT_EQ(value, result);
 
-  context.reset();
+  context->Shutdown();
+  context = nullptr;
   base::RunLoop().RunUntilIdle();
 
   {
     // Mess up version number in database.
     leveldb_env::ChromiumEnv env;
     leveldb::DB* db = nullptr;
     leveldb::Options options;
     options.env = &env;
     base::FilePath db_path =
         temp_path().Append(test_path).Append(FILE_PATH_LITERAL("leveldb"));
     ASSERT_TRUE(leveldb::DB::Open(options, db_path.AsUTF8Unsafe(), &db).ok());
     std::unique_ptr<leveldb::DB> db_owner(db);
     ASSERT_TRUE(db->Put(leveldb::WriteOptions(), "VERSION", "argh").ok());
   }
 
   // Make sure data is gone.
-  context = base::MakeUnique<LocalStorageContextMojo>(
-      connector(), nullptr, base::FilePath(), test_path);
+  context = base::MakeShared<LocalStorageContextMojo>(
+      connector(), nullptr, base::FilePath(), test_path, nullptr);
   EXPECT_FALSE(DoTestGet(context.get(), key, &result));
 
   // Write data again.
   DoTestPut(context.get(), key, value);
 
-  context.reset();
+  context->Shutdown();
+  context = nullptr;
   base::RunLoop().RunUntilIdle();
 
   // Data should have been preserved now.
-  context = base::MakeUnique<LocalStorageContextMojo>(
-      connector(), nullptr, base::FilePath(), test_path);
+  context = base::MakeShared<LocalStorageContextMojo>(
+      connector(), nullptr, base::FilePath(), test_path, nullptr);
   EXPECT_TRUE(DoTestGet(context.get(), key, &result));
   EXPECT_EQ(value, result);
+  context->Shutdown();
 }
 
 }  // namespace content