PVer4: V4Database: Return stores that have a hash prefix for a given full hash

components/safe_browsing_db/v4_database_unittest.cc

Index: components/safe_browsing_db/v4_database_unittest.cc
diff --git a/components/safe_browsing_db/v4_database_unittest.cc b/components/safe_browsing_db/v4_database_unittest.cc
index c841cbfb295224293f2db242cd8ff9eb074abd98..1cac1aff663c1eee98469ea9e980560ebb122549 100644
--- a/components/safe_browsing_db/v4_database_unittest.cc
+++ b/components/safe_browsing_db/v4_database_unittest.cc
@@ -19,15 +19,26 @@ class FakeV4Store : public V4Store {
  public:
   FakeV4Store(const scoped_refptr<base::SequencedTaskRunner>& task_runner,
               const base::FilePath& store_path,
-              const bool reset_succeeds)
+              const bool reset_succeeds,
+              const bool hash_prefix_matches)
       : V4Store(
             task_runner,
             base::FilePath(store_path.value() + FILE_PATH_LITERAL(".fake"))),
+        hash_prefix_matches_(hash_prefix_matches),
         reset_succeeds_(reset_succeeds) {}
 
   bool Reset() override { return reset_succeeds_; }
 
+  HashPrefix GetMatchingHashPrefix(const FullHash& full_hash) override {
+    return hash_prefix_matches_ ? full_hash : HashPrefix();
+  }
+
+  void set_hash_prefix_matches(bool hash_prefix_matches) {
+    hash_prefix_matches_ = hash_prefix_matches;
+  }
+
  private:
+  bool hash_prefix_matches_;
   bool reset_succeeds_;
 };
 
@@ -36,18 +47,21 @@ class FakeV4Store : public V4Store {
 // always return true. This is used to test the Reset() method in V4Database.
 class FakeV4StoreFactory : public V4StoreFactory {
  public:
-  FakeV4StoreFactory(bool next_store_reset_fails)
-      : next_store_reset_fails_(next_store_reset_fails) {}
+  FakeV4StoreFactory(bool next_store_reset_fails, bool hash_prefix_matches)
+      : hash_prefix_matches_(hash_prefix_matches),
+        next_store_reset_fails_(next_store_reset_fails) {}
 
   V4Store* CreateV4Store(
       const scoped_refptr<base::SequencedTaskRunner>& task_runner,
       const base::FilePath& store_path) override {
     bool reset_succeeds = !next_store_reset_fails_;
     next_store_reset_fails_ = false;
-    return new FakeV4Store(task_runner, store_path, reset_succeeds);
+    return new FakeV4Store(task_runner, store_path, reset_succeeds,
+                           hash_prefix_matches_);
   }
 
  private:
+  bool hash_prefix_matches_;
   bool next_store_reset_fails_;
 };
 
@@ -80,8 +94,10 @@ class V4DatabaseTest : public PlatformTest {
     PlatformTest::TearDown();
   }
 
-  void RegisterFactory(bool fails_first_reset) {
-    factory_.reset(new FakeV4StoreFactory(fails_first_reset));
+  void RegisterFactory(bool fails_first_reset,
+                       bool hash_prefix_matches = true) {
+    factory_.reset(
+        new FakeV4StoreFactory(fails_first_reset, hash_prefix_matches));
     V4Database::RegisterStoreFactoryForTest(factory_.get());
   }
 
@@ -346,4 +362,77 @@ TEST_F(V4DatabaseTest, TestApplyUpdateWithInvalidUpdate) {
   VerifyExpectedStoresState(false);
 }
 
+// Test to ensure the case that all stores match a given full hash.
+TEST_F(V4DatabaseTest, TestAllStoresMatchFullHash) {
+  bool hash_prefix_matches = true;
+  expected_resets_successfully_ = true;
+  RegisterFactory(!expected_resets_successfully_, hash_prefix_matches);
+
+  V4Database::Create(task_runner_, database_dirname_, store_file_name_map_,
+                     callback_db_ready_);
+  created_but_not_called_back_ = true;
+  task_runner_->RunPendingTasks();
+
+  base::RunLoop().RunUntilIdle();
+  EXPECT_EQ(true, created_and_called_back_);
+
+  StoresMatched stores_matched;
+  v4_database_->GetStoresMatchingFullHash("anything", &stores_matched);
+  EXPECT_EQ(2u, stores_matched.size());
+  EXPECT_TRUE(
+      stores_matched.end() !=
+      std::find(stores_matched.begin(), stores_matched.end(),
+                UpdateListIdentifier(WINDOWS_PLATFORM, URL, MALWARE_THREAT)));
+  EXPECT_TRUE(
+      stores_matched.end() !=
+      std::find(stores_matched.begin(), stores_matched.end(),
+                UpdateListIdentifier(LINUX_PLATFORM, URL, MALWARE_THREAT)));
+}
+
+// Test to ensure the case that no stores match a given full hash.
+TEST_F(V4DatabaseTest, TestNoStoreMatchesFullHash) {
+  bool hash_prefix_matches = false;
+  expected_resets_successfully_ = true;
+  RegisterFactory(!expected_resets_successfully_, hash_prefix_matches);
+
+  V4Database::Create(task_runner_, database_dirname_, store_file_name_map_,
+                     callback_db_ready_);
+  created_but_not_called_back_ = true;
+  task_runner_->RunPendingTasks();
+
+  base::RunLoop().RunUntilIdle();
+  EXPECT_EQ(true, created_and_called_back_);
+
+  StoresMatched stores_matched;
+  v4_database_->GetStoresMatchingFullHash("anything", &stores_matched);
+  EXPECT_TRUE(stores_matched.empty());
+}
+
+// Test to ensure the case that some stores match a given full hash.
+TEST_F(V4DatabaseTest, TestSomeStoresMatchFullHash) {
+  // Setup stores to not match the full hash.
+  bool hash_prefix_matches = false;
+  expected_resets_successfully_ = true;
+  RegisterFactory(!expected_resets_successfully_, hash_prefix_matches);
+
+  V4Database::Create(task_runner_, database_dirname_, store_file_name_map_,
+                     callback_db_ready_);
+  created_but_not_called_back_ = true;
+  task_runner_->RunPendingTasks();
+
+  base::RunLoop().RunUntilIdle();
+  EXPECT_EQ(true, created_and_called_back_);
+
+  UpdateListIdentifier identifier(LINUX_PLATFORM, URL, MALWARE_THREAT);
+  FakeV4Store* store =
+      static_cast<FakeV4Store*>(v4_database_->store_map_->at(identifier).get());
+  // Set the store corresponding to identifier to match the full hash.
+  store->set_hash_prefix_matches(true);
+
+  StoresMatched stores_matched;
+  v4_database_->GetStoresMatchingFullHash("anything", &stores_matched);
+  EXPECT_EQ(1u, stores_matched.size());
+  EXPECT_EQ(identifier, stores_matched[0]);
+}
+
 }  // namespace safe_browsing