Small: Use a weak pointer for callbacks to V4LocalDatabaseManager

components/safe_browsing_db/v4_local_database_manager_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 "base/files/scoped_temp_dir.h"
 #include "base/memory/ptr_util.h"
 #include "base/memory/ref_counted.h"
 #include "base/run_loop.h"
 #include "base/test/test_simple_task_runner.h"
 #include "components/safe_browsing_db/v4_database.h"
@@ skipping 154 matching lines @@
 
   void SetTaskRunnerForTest() {
     v4_local_database_manager_->SetTaskRunnerForTest(task_runner_);
   }
 
   void StartLocalDatabaseManager() {
     v4_local_database_manager_->StartOnIOThread(NULL, V4ProtocolConfig());
   }
 
   void StopLocalDatabaseManager() {
-    v4_local_database_manager_->StopOnIOThread(true);
+    if (v4_local_database_manager_) {
+      v4_local_database_manager_->StopOnIOThread(true);
+    }
 
     // Force destruction of the database.
-    task_runner_->RunPendingTasks();
-    base::RunLoop().RunUntilIdle();
+    WaitForTasksOnTaskRunner();
   }
 
   void WaitForTasksOnTaskRunner() {
     // Wait for tasks on the task runner so we're sure that the
     // V4LocalDatabaseManager has read the data from disk.
     task_runner_->RunPendingTasks();
     base::RunLoop().RunUntilIdle();
   }
 
   base::ScopedTempDir base_dir_;
@@ skipping 135 matching lines @@
   EXPECT_FALSE(FakeV4LocalDatabaseManager::PerformFullHashCheckCalled(
       v4_local_database_manager_));
 
   // Wait for PerformFullHashCheck to complete.
   WaitForTasksOnTaskRunner();
 
   EXPECT_TRUE(FakeV4LocalDatabaseManager::PerformFullHashCheckCalled(
       v4_local_database_manager_));
 }
 
+TEST_F(V4LocalDatabaseManagerTest, UsingWeakPtrDropsCallback) {
+  // StopLocalDatabaseManager before resetting it because that's what
+  // ~V4LocalDatabaseManager expects.
+  StopLocalDatabaseManager();
+  v4_local_database_manager_ =
+      make_scoped_refptr(new FakeV4LocalDatabaseManager(base_dir_.GetPath()));
+  SetTaskRunnerForTest();
+  StartLocalDatabaseManager();
+  WaitForTasksOnTaskRunner();
+  net::TestURLFetcherFactory factory;
+
+  StoreAndHashPrefixes store_and_hash_prefixes;
+  store_and_hash_prefixes.emplace_back(GetUrlMalwareId(), HashPrefix("aaaa"));
+  ReplaceV4Database(store_and_hash_prefixes);
+
+  // The fake database returns a matched hash prefix.
+  EXPECT_FALSE(v4_local_database_manager_->CheckBrowseUrl(
+      GURL("http://example.com/a/"), nullptr));
+  v4_local_database_manager_->StopOnIOThread(true);
+
+  // Release the V4LocalDatabaseManager object right away before the callback
+  // gets called. When the callback gets called, without using a weak-ptr
+  // factory, this leads to a use after free. However, using the weak-ptr means
+  // that the callback is simply dropped.
+  v4_local_database_manager_ = nullptr;
+
+  // Wait for the tasks scheduled by StopOnIOThread to complete.
+  WaitForTasksOnTaskRunner();
+}
+
 }  // namespace safe_browsing