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

Overall design document: https://goto.google.com/design-doc-v4store

BUG=543161

Review-Url: https://codereview.chromium.org/2160193003
Cr-Commit-Position: refs/heads/master@{#406686}

components/safe_browsing_db/v4_database.cc

@@ -168,4 +168,19 @@ std::unique_ptr<StoreStateMap> V4Database::GetStoreStateMap() {
   return store_state_map;
 }
 
+void V4Database::GetStoresMatchingFullHash(
+    const FullHash& full_hash,
+    const base::hash_set<UpdateListIdentifier>& stores_to_look,
+    MatchedHashPrefixMap* matched_hash_prefix_map) {
+  for (const UpdateListIdentifier& identifier : stores_to_look) {
+    const auto& store_pair = store_map_->find(identifier);
+    DCHECK(store_pair != store_map_->end());
+    const std::unique_ptr<V4Store>& store = store_pair->second;
+    HashPrefix hash_prefix = store->GetMatchingHashPrefix(full_hash);
+    if (!hash_prefix.empty()) {
+      (*matched_hash_prefix_map)[identifier] = hash_prefix;
+    }
+  }
+}
+
 }  // namespace safe_browsing

components/safe_browsing_db/v4_database.h

@@ -30,6 +30,10 @@ typedef base::Callback<void()> DatabaseUpdatedCallback;
 // as manage their storage on disk.
 typedef base::hash_map<UpdateListIdentifier, std::unique_ptr<V4Store>> StoreMap;
 
+// Map of identifier for any store that had a hash prefix matching the given
+// full hash to the matching hash prefix.
+typedef base::hash_map<UpdateListIdentifier, HashPrefix> MatchedHashPrefixMap;
+
 // Factory for creating V4Database. Tests implement this factory to create fake
 // databases for testing.
 class V4DatabaseFactory {
@@ -74,6 +78,14 @@ class V4Database {
   // Returns the current state of each of the stores being managed.
   std::unique_ptr<StoreStateMap> GetStoreStateMap();
 
+  // Searches for a hash prefix matching the |full_hash| in stores in the
+  // database, filtered by |stores_to_look|, and returns the identifier of the
+  // store along with the matching hash prefix in |matched_hash_prefix_map|.
+  void GetStoresMatchingFullHash(
+      const FullHash& full_hash,
+      const base::hash_set<UpdateListIdentifier>& stores_to_look,
+      MatchedHashPrefixMap* matched_hash_prefix_map);
+
   // Deletes the current database and creates a new one.
   virtual bool ResetDatabase();
 
@@ -90,6 +102,7 @@ class V4Database {
   FRIEND_TEST_ALL_PREFIXES(V4DatabaseTest, TestApplyUpdateWithNoNewState);
   FRIEND_TEST_ALL_PREFIXES(V4DatabaseTest, TestApplyUpdateWithEmptyUpdate);
   FRIEND_TEST_ALL_PREFIXES(V4DatabaseTest, TestApplyUpdateWithInvalidUpdate);
+  FRIEND_TEST_ALL_PREFIXES(V4DatabaseTest, TestSomeStoresMatchFullHash);
 
   // Makes the passed |factory| the factory used to instantiate a V4Store. Only
   // for tests.

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_should_match_(hash_prefix_matches),
         reset_succeeds_(reset_succeeds) {}
 
   bool Reset() override { return reset_succeeds_; }
 
+  HashPrefix GetMatchingHashPrefix(const FullHash& full_hash) override {
+    return hash_prefix_should_match_ ? full_hash : HashPrefix();
+  }
+
+  void set_hash_prefix_matches(bool hash_prefix_matches) {
+    hash_prefix_should_match_ = hash_prefix_matches;
+  }
+
  private:
+  bool hash_prefix_should_match_;
   bool reset_succeeds_;
 };
 
@@ -36,24 +47,30 @@ 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_should_match_(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_should_match_);
   }
 
  private:
+  bool hash_prefix_should_match_;
   bool next_store_reset_fails_;
 };
 
 class V4DatabaseTest : public PlatformTest {
  public:
-  V4DatabaseTest() : task_runner_(new base::TestSimpleTaskRunner) {}
+  V4DatabaseTest()
+      : task_runner_(new base::TestSimpleTaskRunner),
+        linux_malware_id_(LINUX_PLATFORM, URL, MALWARE_THREAT),
+        win_malware_id_(WINDOWS_PLATFORM, URL, MALWARE_THREAT) {}
 
   void SetUp() override {
     PlatformTest::SetUp();
@@ -80,21 +97,21 @@ 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());
   }
 
   void SetupInfoMapAndExpectedState() {
-    UpdateListIdentifier win_malware_id(WINDOWS_PLATFORM, URL, MALWARE_THREAT);
-    store_file_name_map_[win_malware_id] = "win_url_malware";
-    expected_identifiers_.push_back(win_malware_id);
+    store_file_name_map_[win_malware_id_] = "win_url_malware";
+    expected_identifiers_.push_back(win_malware_id_);
     expected_store_paths_.push_back(
         database_dirname_.AppendASCII("win_url_malware.fake"));
 
-    UpdateListIdentifier linux_malware_id(LINUX_PLATFORM, URL, MALWARE_THREAT);
-    store_file_name_map_[linux_malware_id] = "linux_url_malware";
-    expected_identifiers_.push_back(linux_malware_id);
+    store_file_name_map_[linux_malware_id_] = "linux_url_malware";
+    expected_identifiers_.push_back(linux_malware_id_);
     expected_store_paths_.push_back(
         database_dirname_.AppendASCII("linux_url_malware.fake"));
   }
@@ -192,6 +209,7 @@ class V4DatabaseTest : public PlatformTest {
   NewDatabaseReadyCallback callback_db_ready_;
   StoreStateMap expected_store_state_map_;
   base::hash_map<UpdateListIdentifier, V4Store*> old_stores_map_;
+  const UpdateListIdentifier linux_malware_id_, win_malware_id_;
 };
 
 // Test to set up the database with fake stores.
@@ -346,4 +364,104 @@ 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_);
+
+  base::hash_set<UpdateListIdentifier> stores_to_look(
+      {linux_malware_id_, win_malware_id_});
+  MatchedHashPrefixMap matched_hash_prefix_map;
+  v4_database_->GetStoresMatchingFullHash("anything", stores_to_look,
+                                          &matched_hash_prefix_map);
+  EXPECT_EQ(2u, matched_hash_prefix_map.size());
+  EXPECT_FALSE(matched_hash_prefix_map[linux_malware_id_].empty());
+  EXPECT_FALSE(matched_hash_prefix_map[win_malware_id_].empty());
+}
+
+// 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_);
+
+  base::hash_set<UpdateListIdentifier> stores_to_look(
+      {linux_malware_id_, win_malware_id_});
+  MatchedHashPrefixMap matched_hash_prefix_map;
+  v4_database_->GetStoresMatchingFullHash("anything", stores_to_look,
+                                          &matched_hash_prefix_map);
+  EXPECT_TRUE(matched_hash_prefix_map.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_);
+
+  // Set the store corresponding to linux_malware_id_ to match the full hash.
+  FakeV4Store* store = static_cast<FakeV4Store*>(
+      v4_database_->store_map_->at(linux_malware_id_).get());
+  store->set_hash_prefix_matches(true);
+
+  base::hash_set<UpdateListIdentifier> stores_to_look(
+      {linux_malware_id_, win_malware_id_});
+  MatchedHashPrefixMap matched_hash_prefix_map;
+  v4_database_->GetStoresMatchingFullHash("anything", stores_to_look,
+                                          &matched_hash_prefix_map);
+  EXPECT_EQ(1u, matched_hash_prefix_map.size());
+  EXPECT_FALSE(matched_hash_prefix_map[linux_malware_id_].empty());
+}
+
+// Test to ensure the case that only some stores are reported to match a given
+// full hash because of stores_to_look.
+TEST_F(V4DatabaseTest, TestSomeStoresMatchFullHashBecauseOfStoresToMatch) {
+  // Setup all stores to match the full hash.
+  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_);
+
+  base::hash_set<UpdateListIdentifier> stores_to_look({linux_malware_id_});
+  // Don't add win_malware_id_ to the stores_to_look.
+  MatchedHashPrefixMap matched_hash_prefix_map;
+  v4_database_->GetStoresMatchingFullHash("anything", stores_to_look,
+                                          &matched_hash_prefix_map);
+  EXPECT_EQ(1u, matched_hash_prefix_map.size());
+  EXPECT_FALSE(matched_hash_prefix_map[linux_malware_id_].empty());
+}
+
 }  // namespace safe_browsing

components/safe_browsing_db/v4_store.h

@@ -169,7 +169,7 @@ class V4Store {
 
   // If a hash prefix in this store matches |full_hash|, returns that hash
   // prefix; otherwise returns an empty hash prefix.
-  HashPrefix GetMatchingHashPrefix(const FullHash& full_hash);
+  virtual HashPrefix GetMatchingHashPrefix(const FullHash& full_hash);
 
   std::string DebugString() const;