Merge the hash prefixes from the old store and the additions in the partial

components/safe_browsing_db/v4_store.h

 // 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.
 
 #ifndef COMPONENTS_SAFE_BROWSING_DB_V4_STORE_H_
 #define COMPONENTS_SAFE_BROWSING_DB_V4_STORE_H_
 
 #include "base/files/file_path.h"
 #include "base/memory/ref_counted.h"
 #include "base/sequenced_task_runner.h"
 #include "base/single_thread_task_runner.h"
 #include "components/safe_browsing_db/v4_protocol_manager_util.h"
 
 namespace safe_browsing {
 
 class V4Store;
 
 typedef base::Callback<void(std::unique_ptr<V4Store>)>
     UpdatedStoreReadyCallback;
 
+// The size of the hash prefix, in bytes. It should be between 4 to 32 (full
+// hash).
+typedef size_t PrefixSize;
+
+// A hash prefix sent by the SafeBrowsing PVer4 service.
+typedef std::string HashPrefix;
+
+// The sorted list of hash prefixes.
+typedef std::string HashPrefixes;
+
+// Stores the list of sorted hash prefixes, by size.
+// For instance: {4: ["abcd", "bcde", "cdef", "gggg"], 5: ["fffff"]}
+typedef base::hash_map<PrefixSize, HashPrefixes> HashPrefixMap;
+
+// Stores the index of the last element merged from the HashPrefixMap for a
+// given prefix size. For instance: {4:3, 5:1} means that we have already merged
+// 3 hash prefixes of length 4, and 1 hash prefix of length 5.
+typedef base::hash_map<PrefixSize, size_t> CounterMap;
+
 // Enumerate different failure events while parsing the file read from disk for
 // histogramming purposes.  DO NOT CHANGE THE ORDERING OF THESE VALUES.
 enum StoreReadResult {
   // No errors.
   READ_SUCCESS = 0,
 
   // Reserved for errors in parsing this enum.
   UNEXPECTED_READ_FAILURE = 1,
 
   // The contents of the file could not be read.
@@ skipping 41 matching lines @@
   UNEXPECTED_BYTES_WRITTEN_FAILURE = 3,
 
   // Renaming the temporary file to store file failed.
   UNABLE_TO_RENAME_FAILURE = 4,
 
   // Memory space for histograms is determined by the max.  ALWAYS
   // ADD NEW VALUES BEFORE THIS ONE.
   STORE_WRITE_RESULT_MAX
 };
 
+// Enumerate different events while merging the update fetched fom the server
+// for histogramming purposes.
+// DO NOT CHANGE THE ORDERING OF THESE VALUES.
+enum MergeResult {
+  // No errors.
+  MERGE_SUCCESS = 0,
+
+  // Reserved for errors in parsing this enum.
+  UNEXPECTED_MERGE_FAILURE = 1,
+
+  // Prefix size smaller than 4 (which is the lowest expected).
+  PREFIX_SIZE_TOO_SMALL_FAILURE = 2,
+
+  // Prefix size larger than 32 (length of a full SHA256 hash).
+  PREFIX_SIZE_TOO_LARGE_FAILURE = 3,
+
+  // The number of bytes in additions isn't a multiple of prefix size.
+  ADDITIONS_SIZE_UNEXPECTED_FAILURE = 4,
+
+  // Memory space for histograms is determined by the max.  ALWAYS
+  // ADD NEW VALUES BEFORE THIS ONE.
+  MERGE_RESULT_MAX
+};
+
 // Factory for creating V4Store. Tests implement this factory to create fake
 // stores for testing.
 class V4StoreFactory {
  public:
   virtual ~V4StoreFactory() {}
   virtual V4Store* CreateV4Store(
       const scoped_refptr<base::SequencedTaskRunner>& task_runner,
       const base::FilePath& store_path);
 };
 
@@ skipping 28 matching lines @@
   FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestReadFromAbsentFile);
   FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestReadFromInvalidContentsFile);
   FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestReadFromUnexpectedMagicFile);
   FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestReadFromLowVersionFile);
   FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestReadFromNoHashPrefixInfoFile);
   FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestReadFromNoHashPrefixesFile);
   FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestWriteNoResponseType);
   FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestWritePartialResponseType);
   FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestWriteFullResponseType);
   FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestReadFromFileWithUnknownProto);
+  FRIEND_TEST_ALL_PREFIXES(V4StoreTest,
+                           TestAddUnlumpedHashesWithInvalidAddition);
+  FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestAddUnlumpedHashes);
+  FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestAddUnlumpedHashesWithEmptyString);
+  FRIEND_TEST_ALL_PREFIXES(V4StoreTest,
+                           TestGetNextSmallestPrefixSizeWithEmptyPrefixMap);
+  FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestGetNextSmallestPrefixSize);
+  FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestGetNextUnmergedPrefix);
+  FRIEND_TEST_ALL_PREFIXES(V4StoreTest, TestMergeUpdates);
+
+  // If |prefix_size| is within expected range, and |raw_hashes| is not invalid,
+  // then it sets |raw_hashes| as the value at key |prefix_size| in
+  // |additions_map|
+  static MergeResult AddUnlumpedHashes(PrefixSize prefix_size,
+                                       const std::string& raw_hashes,
+                                       HashPrefixMap* additions_map);
+
+  // Get the size of the next unmerged hash prefix in dictionary order from
+  // |hash_prefix_map|. |counter_map| is used to determine which hash prefixes
+  // have been merged already. Returns true if there are any unmerged hash
+  // prefixes in the list.
+  static bool GetNextSmallestPrefixSize(const HashPrefixMap& hash_prefix_map,
+                                        const CounterMap& counter_map,
+                                        PrefixSize* smallest_prefix_size);
+
+  // Returns the next hash prefix of length |prefix_size| from |hash_prefix_map|
+  // that hasn't been merged already. |counter_map| is used to determine the
+  // index of the next prefix of size |prefix_size| to merge.
+  static HashPrefix GetNextUnmergedPrefixForSize(
+      PrefixSize prefix_size,
+      const HashPrefixMap& hash_prefix_map,
+      const CounterMap& counter_map);
+
+  // Sets a value of 0 in |counter_map| for all keys in |hash_prefix_map|.
+  static void InitializeCounterMap(const HashPrefixMap& hash_prefix_map,
+                                   CounterMap* counter_map);
+
+  // Reserve the appropriate string size so that the string size of the merged
+  // list is exact.
+  static void ReserveSpaceInPrefixMap(const HashPrefixMap& other_prefixes_map,
+                                      HashPrefixMap* prefix_map_to_update);
+
+  // Updates the |additions_map| with the additions received in the partial
+  // update from the server.
+  static void UpdateHashPrefixMapFromAdditions(
+      const ::google::protobuf::RepeatedPtrField<ThreatEntrySet>& additions,
+      HashPrefixMap* additions_map);
+
+  // Merges the prefix map from the old store (|old_hash_prefix_map|) and the
+  // update (additions_map) to populate the prefix map for the current store.
+  // TODO(vakh): Process removals also.
+  void MergeUpdate(const HashPrefixMap& old_hash_prefix_map,
+                   const HashPrefixMap& additions_map);
 
   // Reads the state of the store from the file on disk and returns the reason
   // for the failure or reports success.
   StoreReadResult ReadFromDisk();
 
   // Writes the FULL_UPDATE |response| to disk as a V4StoreFileFormat proto.
   StoreWriteResult WriteToDisk(
       std::unique_ptr<ListUpdateResponse> response) const;
 
   // The state of the store as returned by the PVer4 server in the last applied
   // update response.
   std::string state_;
   const base::FilePath store_path_;
+  HashPrefixMap hash_prefix_map_;
   const scoped_refptr<base::SequencedTaskRunner> task_runner_;
 };
 
 std::ostream& operator<<(std::ostream& os, const V4Store& store);
 
 }  // namespace safe_browsing
 
 #endif  // COMPONENTS_SAFE_BROWSING_DB_V4_STORE_H_