PVer4: Drop hash prefixes based on the removals field in the response

components/safe_browsing_db/v4_store_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/file_util.h"
 #include "base/files/scoped_temp_dir.h"
 #include "base/memory/ptr_util.h"
 #include "base/run_loop.h"
 #include "base/test/test_simple_task_runner.h"
 #include "base/time/time.h"
 #include "components/safe_browsing_db/v4_store.h"
 #include "components/safe_browsing_db/v4_store.pb.h"
 #include "content/public/test/test_browser_thread_bundle.h"
 #include "testing/platform_test.h"
 
 namespace safe_browsing {
 
+using ::google::protobuf::RepeatedField;
+using ::google::protobuf::int32;
+
 class V4StoreTest : public PlatformTest {
  public:
   V4StoreTest() : task_runner_(new base::TestSimpleTaskRunner) {}
 
   void SetUp() override {
     PlatformTest::SetUp();
 
     ASSERT_TRUE(temp_dir_.CreateUniqueTempDir());
     store_path_ = temp_dir_.path().AppendASCII("V4StoreTest.store");
     DVLOG(1) << "store_path_: " << store_path_.value();
@@ skipping 101 matching lines @@
 }
 
 TEST_F(V4StoreTest, TestWriteFullResponseType) {
   std::unique_ptr<ListUpdateResponse> list_update_response(
       new ListUpdateResponse);
   list_update_response->set_response_type(ListUpdateResponse::FULL_UPDATE);
   list_update_response->set_new_client_state("test_client_state");
   EXPECT_EQ(WRITE_SUCCESS, V4Store(task_runner_, store_path_)
                                .WriteToDisk(std::move(list_update_response)));
 
-  std::unique_ptr<V4Store> read_store(new V4Store(task_runner_, store_path_));
-  EXPECT_EQ(READ_SUCCESS, read_store->ReadFromDisk());
-  EXPECT_EQ("test_client_state", read_store->state_);
-  EXPECT_TRUE(read_store->hash_prefix_map_.empty());
+  V4Store read_store(task_runner_, store_path_);
+  EXPECT_EQ(READ_SUCCESS, read_store.ReadFromDisk());
+  EXPECT_EQ("test_client_state", read_store.state_);
+  EXPECT_TRUE(read_store.hash_prefix_map_.empty());
 }
 
 TEST_F(V4StoreTest, TestAddUnlumpedHashesWithInvalidAddition) {
   HashPrefixMap prefix_map;
   EXPECT_EQ(ADDITIONS_SIZE_UNEXPECTED_FAILURE,
             V4Store::AddUnlumpedHashes(5, "a", &prefix_map));
   EXPECT_TRUE(prefix_map.empty());
 }
 
 TEST_F(V4StoreTest, TestAddUnlumpedHashesWithEmptyString) {
@@ skipping 55 matching lines @@
             V4Store::AddUnlumpedHashes(4, "abcdefgh", &prefix_map_old));
   EXPECT_EQ(APPLY_UPDATE_SUCCESS,
             V4Store::AddUnlumpedHashes(5, "54321abcde", &prefix_map_old));
   HashPrefixMap prefix_map_additions;
   EXPECT_EQ(
       APPLY_UPDATE_SUCCESS,
       V4Store::AddUnlumpedHashes(4, "----1111bbbb", &prefix_map_additions));
   EXPECT_EQ(APPLY_UPDATE_SUCCESS,
             V4Store::AddUnlumpedHashes(5, "22222bcdef", &prefix_map_additions));
 
-  std::unique_ptr<V4Store> store(new V4Store(task_runner_, store_path_));
-  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
-            store->MergeUpdate(prefix_map_old, prefix_map_additions));
-  const HashPrefixMap& prefix_map = store->hash_prefix_map_;
+  V4Store store(task_runner_, store_path_);
+  RepeatedField<int32> raw_removals;
+  EXPECT_EQ(
+      APPLY_UPDATE_SUCCESS,
+      store.MergeUpdate(prefix_map_old, prefix_map_additions, raw_removals));
+  const HashPrefixMap& prefix_map = store.hash_prefix_map_;
   EXPECT_EQ(2u, prefix_map.size());
 
   PrefixSize prefix_size = 4;
   HashPrefixes hash_prefixes = prefix_map.at(prefix_size);
   EXPECT_EQ(5 * prefix_size, hash_prefixes.size());
   EXPECT_EQ("----", hash_prefixes.substr(0 * prefix_size, prefix_size));
   EXPECT_EQ("1111", hash_prefixes.substr(1 * prefix_size, prefix_size));
   EXPECT_EQ("abcd", hash_prefixes.substr(2 * prefix_size, prefix_size));
   EXPECT_EQ("bbbb", hash_prefixes.substr(3 * prefix_size, prefix_size));
   EXPECT_EQ("efgh", hash_prefixes.substr(4 * prefix_size, prefix_size));
 
   prefix_size = 5;
   hash_prefixes = prefix_map.at(prefix_size);
   EXPECT_EQ(4 * prefix_size, hash_prefixes.size());
   EXPECT_EQ("22222", hash_prefixes.substr(0 * prefix_size, prefix_size));
   EXPECT_EQ("54321", hash_prefixes.substr(1 * prefix_size, prefix_size));
   EXPECT_EQ("abcde", hash_prefixes.substr(2 * prefix_size, prefix_size));
   EXPECT_EQ("bcdef", hash_prefixes.substr(3 * prefix_size, prefix_size));
 }
 
 TEST_F(V4StoreTest, TestMergeUpdatesWithDifferentSizesInEachMap) {
   HashPrefixMap prefix_map_old;
   EXPECT_EQ(APPLY_UPDATE_SUCCESS,
             V4Store::AddUnlumpedHashes(4, "1111abcdefgh", &prefix_map_old));
   HashPrefixMap prefix_map_additions;
   EXPECT_EQ(APPLY_UPDATE_SUCCESS,
             V4Store::AddUnlumpedHashes(5, "22222bcdef", &prefix_map_additions));
 
-  std::unique_ptr<V4Store> store(new V4Store(task_runner_, store_path_));
-  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
-            store->MergeUpdate(prefix_map_old, prefix_map_additions));
-  const HashPrefixMap& prefix_map = store->hash_prefix_map_;
+  V4Store store(task_runner_, store_path_);
+  RepeatedField<int32> raw_removals;
+  EXPECT_EQ(
+      APPLY_UPDATE_SUCCESS,
+      store.MergeUpdate(prefix_map_old, prefix_map_additions, raw_removals));
+  const HashPrefixMap& prefix_map = store.hash_prefix_map_;
   EXPECT_EQ(2u, prefix_map.size());
 
   PrefixSize prefix_size = 4;
   HashPrefixes hash_prefixes = prefix_map.at(prefix_size);
   EXPECT_EQ(3 * prefix_size, hash_prefixes.size());
   EXPECT_EQ("1111abcdefgh", hash_prefixes);
 
   prefix_size = 5;
   hash_prefixes = prefix_map.at(prefix_size);
   EXPECT_EQ(2 * prefix_size, hash_prefixes.size());
   EXPECT_EQ("22222bcdef", hash_prefixes);
 }
 
 TEST_F(V4StoreTest, TestMergeUpdatesOldMapRunsOutFirst) {
   HashPrefixMap prefix_map_old;
   EXPECT_EQ(APPLY_UPDATE_SUCCESS,
             V4Store::AddUnlumpedHashes(4, "00001111", &prefix_map_old));
   HashPrefixMap prefix_map_additions;
   EXPECT_EQ(APPLY_UPDATE_SUCCESS,
             V4Store::AddUnlumpedHashes(4, "2222", &prefix_map_additions));
 
-  std::unique_ptr<V4Store> store(new V4Store(task_runner_, store_path_));
-  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
-            store->MergeUpdate(prefix_map_old, prefix_map_additions));
-  const HashPrefixMap& prefix_map = store->hash_prefix_map_;
+  V4Store store(task_runner_, store_path_);
+  RepeatedField<int32> raw_removals;
+  EXPECT_EQ(
+      APPLY_UPDATE_SUCCESS,
+      store.MergeUpdate(prefix_map_old, prefix_map_additions, raw_removals));
+  const HashPrefixMap& prefix_map = store.hash_prefix_map_;
   EXPECT_EQ(1u, prefix_map.size());
 
   PrefixSize prefix_size = 4;
   HashPrefixes hash_prefixes = prefix_map.at(prefix_size);
   EXPECT_EQ(3 * prefix_size, hash_prefixes.size());
   EXPECT_EQ("0000", hash_prefixes.substr(0 * prefix_size, prefix_size));
   EXPECT_EQ("1111", hash_prefixes.substr(1 * prefix_size, prefix_size));
   EXPECT_EQ("2222", hash_prefixes.substr(2 * prefix_size, prefix_size));
 }
 
 TEST_F(V4StoreTest, TestMergeUpdatesAdditionsMapRunsOutFirst) {
   HashPrefixMap prefix_map_old;
   EXPECT_EQ(APPLY_UPDATE_SUCCESS,
             V4Store::AddUnlumpedHashes(4, "2222", &prefix_map_old));
   HashPrefixMap prefix_map_additions;
   EXPECT_EQ(APPLY_UPDATE_SUCCESS,
             V4Store::AddUnlumpedHashes(4, "00001111", &prefix_map_additions));
 
-  std::unique_ptr<V4Store> store(new V4Store(task_runner_, store_path_));
-  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
-            store->MergeUpdate(prefix_map_old, prefix_map_additions));
-  const HashPrefixMap& prefix_map = store->hash_prefix_map_;
+  V4Store store(task_runner_, store_path_);
+  RepeatedField<int32> raw_removals;
+  EXPECT_EQ(
+      APPLY_UPDATE_SUCCESS,
+      store.MergeUpdate(prefix_map_old, prefix_map_additions, raw_removals));
+  const HashPrefixMap& prefix_map = store.hash_prefix_map_;
   EXPECT_EQ(1u, prefix_map.size());
 
   PrefixSize prefix_size = 4;
   HashPrefixes hash_prefixes = prefix_map.at(prefix_size);
   EXPECT_EQ(3 * prefix_size, hash_prefixes.size());
   EXPECT_EQ("0000", hash_prefixes.substr(0 * prefix_size, prefix_size));
   EXPECT_EQ("1111", hash_prefixes.substr(1 * prefix_size, prefix_size));
   EXPECT_EQ("2222", hash_prefixes.substr(2 * prefix_size, prefix_size));
 }
 
 TEST_F(V4StoreTest, TestMergeUpdatesFailsForRepeatedHashPrefix) {
   HashPrefixMap prefix_map_old;
   EXPECT_EQ(APPLY_UPDATE_SUCCESS,
             V4Store::AddUnlumpedHashes(4, "2222", &prefix_map_old));
   HashPrefixMap prefix_map_additions;
   EXPECT_EQ(APPLY_UPDATE_SUCCESS,
             V4Store::AddUnlumpedHashes(4, "2222", &prefix_map_additions));
 
-  std::unique_ptr<V4Store> store(new V4Store(task_runner_, store_path_));
-  EXPECT_EQ(ADDITIONS_HAS_EXISTING_PREFIX_FAILURE,
-            store->MergeUpdate(prefix_map_old, prefix_map_additions));
+  V4Store store(task_runner_, store_path_);
+  RepeatedField<int32> raw_removals;
+  EXPECT_EQ(
+      ADDITIONS_HAS_EXISTING_PREFIX_FAILURE,
+      store.MergeUpdate(prefix_map_old, prefix_map_additions, raw_removals));
+}
+
+TEST_F(V4StoreTest, TestMergeUpdatesFailsWhenRemovalsIndexTooLarge) {
+  HashPrefixMap prefix_map_old;
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(4, "2222", &prefix_map_old));
+  HashPrefixMap prefix_map_additions;
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(4, "11113333", &prefix_map_additions));
+
+  // Even though the merged map could have size 3 without removals, the
+  // removals index should only count the entries in the old map.
+  V4Store store(task_runner_, store_path_);
+  RepeatedField<int32> raw_removals;
+  // old_store: ["2222"]
+  raw_removals.Add(1);
+  EXPECT_EQ(
+      REMOVALS_INDEX_TOO_LARGE,
+      store.MergeUpdate(prefix_map_old, prefix_map_additions, raw_removals));
+}
+
+TEST_F(V4StoreTest, TestMergeUpdatesRemovesOnlyElement) {
+  HashPrefixMap prefix_map_old;
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(4, "2222", &prefix_map_old));
+  HashPrefixMap prefix_map_additions;
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(5, "1111133333", &prefix_map_additions));
+
+  V4Store store(task_runner_, store_path_);
+  RepeatedField<int32> raw_removals;
+  // old_store: ["2222"]
+  raw_removals.Add(0);  // Removes "2222"
+  EXPECT_EQ(
+      APPLY_UPDATE_SUCCESS,
+      store.MergeUpdate(prefix_map_old, prefix_map_additions, raw_removals));
+  const HashPrefixMap& prefix_map = store.hash_prefix_map_;
+  // The size is 2 since we reserve space anyway.
+  EXPECT_EQ(2u, prefix_map.size());
+  EXPECT_TRUE(prefix_map.at(4).empty());
+  EXPECT_EQ("1111133333", prefix_map.at(5));
+}
+
+TEST_F(V4StoreTest, TestMergeUpdatesRemovesFirstElement) {
+  HashPrefixMap prefix_map_old;
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(4, "22224444", &prefix_map_old));
+  HashPrefixMap prefix_map_additions;
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(5, "1111133333", &prefix_map_additions));
+
+  V4Store store(task_runner_, store_path_);
+  RepeatedField<int32> raw_removals;
+  // old_store: ["2222", "4444"]
+  raw_removals.Add(0);  // Removes "2222"
+  EXPECT_EQ(
+      APPLY_UPDATE_SUCCESS,
+      store.MergeUpdate(prefix_map_old, prefix_map_additions, raw_removals));
+  const HashPrefixMap& prefix_map = store.hash_prefix_map_;
+  // The size is 2 since we reserve space anyway.
+  EXPECT_EQ(2u, prefix_map.size());
+  EXPECT_EQ("4444", prefix_map.at(4));
+  EXPECT_EQ("1111133333", prefix_map.at(5));
+}
+
+TEST_F(V4StoreTest, TestMergeUpdatesRemovesMiddleElement) {
+  HashPrefixMap prefix_map_old;
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(4, "222233334444", &prefix_map_old));
+  HashPrefixMap prefix_map_additions;
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(5, "1111133333", &prefix_map_additions));
+
+  V4Store store(task_runner_, store_path_);
+  RepeatedField<int32> raw_removals;
+  // old_store: ["2222", "3333", 4444"]
+  raw_removals.Add(1);  // Removes "3333"
+  EXPECT_EQ(
+      APPLY_UPDATE_SUCCESS,
+      store.MergeUpdate(prefix_map_old, prefix_map_additions, raw_removals));
+  const HashPrefixMap& prefix_map = store.hash_prefix_map_;
+  // The size is 2 since we reserve space anyway.
+  EXPECT_EQ(2u, prefix_map.size());
+  EXPECT_EQ("22224444", prefix_map.at(4));
+  EXPECT_EQ("1111133333", prefix_map.at(5));
+}
+
+TEST_F(V4StoreTest, TestMergeUpdatesRemovesLastElement) {
+  HashPrefixMap prefix_map_old;
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(4, "222233334444", &prefix_map_old));
+  HashPrefixMap prefix_map_additions;
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(5, "1111133333", &prefix_map_additions));
+
+  V4Store store(task_runner_, store_path_);
+  RepeatedField<int32> raw_removals;
+  // old_store: ["2222", "3333", 4444"]
+  raw_removals.Add(2);  // Removes "4444"
+  EXPECT_EQ(
+      APPLY_UPDATE_SUCCESS,
+      store.MergeUpdate(prefix_map_old, prefix_map_additions, raw_removals));
+  const HashPrefixMap& prefix_map = store.hash_prefix_map_;
+  // The size is 2 since we reserve space anyway.
+  EXPECT_EQ(2u, prefix_map.size());
+  EXPECT_EQ("22223333", prefix_map.at(4));
+  EXPECT_EQ("1111133333", prefix_map.at(5));
+}
+
+TEST_F(V4StoreTest, TestMergeUpdatesRemovesWhenOldHasDifferentSizes) {
+  HashPrefixMap prefix_map_old;
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(4, "222233334444", &prefix_map_old));
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(5, "aaaaabbbbb", &prefix_map_old));
+  HashPrefixMap prefix_map_additions;
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(5, "1111133333", &prefix_map_additions));
+
+  V4Store store(task_runner_, store_path_);
+  RepeatedField<int32> raw_removals;
+  // old_store: ["2222", "3333", 4444", "aaaaa", "bbbbb"]
+  raw_removals.Add(3);  // Removes "aaaaa"
+  EXPECT_EQ(
+      APPLY_UPDATE_SUCCESS,
+      store.MergeUpdate(prefix_map_old, prefix_map_additions, raw_removals));
+  const HashPrefixMap& prefix_map = store.hash_prefix_map_;
+  // The size is 2 since we reserve space anyway.
+  EXPECT_EQ(2u, prefix_map.size());
+  EXPECT_EQ("222233334444", prefix_map.at(4));
+  EXPECT_EQ("1111133333bbbbb", prefix_map.at(5));
+}
+
+TEST_F(V4StoreTest, TestMergeUpdatesRemovesMultipleAcrossDifferentSizes) {
+  HashPrefixMap prefix_map_old;
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(4, "22223333aaaa", &prefix_map_old));
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(5, "3333344444bbbbb", &prefix_map_old));
+  HashPrefixMap prefix_map_additions;
+  EXPECT_EQ(APPLY_UPDATE_SUCCESS,
+            V4Store::AddUnlumpedHashes(5, "11111", &prefix_map_additions));
+
+  V4Store store(task_runner_, store_path_);
+  RepeatedField<int32> raw_removals;
+  // old_store: ["2222", "3333", "33333", "44444", "aaaa", "bbbbb"]
+  raw_removals.Add(1);  // Removes "3333"
+  raw_removals.Add(3);  // Removes "44444"
+  EXPECT_EQ(
+      APPLY_UPDATE_SUCCESS,
+      store.MergeUpdate(prefix_map_old, prefix_map_additions, raw_removals));
+  const HashPrefixMap& prefix_map = store.hash_prefix_map_;
+  // The size is 2 since we reserve space anyway.
+  EXPECT_EQ(2u, prefix_map.size());
+  EXPECT_EQ("2222aaaa", prefix_map.at(4));
+  EXPECT_EQ("1111133333bbbbb", prefix_map.at(5));
 }
 
 TEST_F(V4StoreTest, TestReadFullResponseWithValidHashPrefixMap) {
   std::unique_ptr<ListUpdateResponse> lur(new ListUpdateResponse);
   lur->set_response_type(ListUpdateResponse::FULL_UPDATE);
   lur->set_new_client_state("test_client_state");
   lur->set_platform_type(WINDOWS_PLATFORM);
   lur->set_threat_entry_type(URL);
   lur->set_threat_type(MALWARE_THREAT);
   ThreatEntrySet* additions = lur->add_additions();
@@ skipping 33 matching lines @@
   EXPECT_EQ(WRITE_SUCCESS,
             V4Store(task_runner_, store_path_).WriteToDisk(std::move(lur)));
 
   V4Store read_store(task_runner_, store_path_);
   EXPECT_EQ(HASH_PREFIX_MAP_GENERATION_FAILURE, read_store.ReadFromDisk());
   EXPECT_TRUE(read_store.state_.empty());
   EXPECT_TRUE(read_store.hash_prefix_map_.empty());
 }
 
 }  // namespace safe_browsing