[safe browsing] Serialize full hashes with prefix set.

chrome/browser/safe_browsing/prefix_set_unittest.cc

 // Copyright (c) 2012 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 "chrome/browser/safe_browsing/prefix_set.h"
 
 #include <algorithm>
 #include <iterator>
 
 #include "base/file_util.h"
@@ skipping 19 matching lines @@
 
 namespace safe_browsing {
 
 class PrefixSetTest : public PlatformTest {
  protected:
   // Constants for the v1 format.
   static const size_t kMagicOffset = 0 * sizeof(uint32);
   static const size_t kVersionOffset = 1 * sizeof(uint32);
   static const size_t kIndexSizeOffset = 2 * sizeof(uint32);
   static const size_t kDeltasSizeOffset = 3 * sizeof(uint32);
-  static const size_t kPayloadOffset = 4 * sizeof(uint32);
+  static const size_t kFullHashesSizeOffset = 4 * sizeof(uint32);
+  static const size_t kPayloadOffset = 5 * sizeof(uint32);
 
   // Generate a set of random prefixes to share between tests.  For
   // most tests this generation was a large fraction of the test time.
   //
   // The set should contain sparse areas where adjacent items are more
   // than 2^16 apart, and dense areas where adjacent items are less
   // than 2^16 apart.
   static void SetUpTestCase() {
     // Distribute clusters of prefixes.
     for (size_t i = 0; i < 250; ++i) {
@@ skipping 340 matching lines @@
   // Test writing and reading an empty set.
   {
     std::vector<SBPrefix> prefixes;
     PrefixSetBuilder builder(prefixes);
     ASSERT_TRUE(builder.GetPrefixSetNoHashes()->WriteFile(filename));
 
     scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
     ASSERT_TRUE(prefix_set.get());
     CheckPrefixes(*prefix_set, prefixes);
   }
+
+  // Test that full hashes are persisted.
+  {
+    std::vector<SBFullHash> hashes;
+    hashes.push_back(SBFullHashForString("one"));
+    hashes.push_back(SBFullHashForString("two"));
+    hashes.push_back(SBFullHashForString("three"));
+
+    std::vector<SBPrefix> prefixes(shared_prefixes_);
+
+    // Remove any collisions from the prefixes.
+    for (size_t i = 0; i < hashes.size(); ++i) {
+      std::vector<SBPrefix>::iterator iter =
+          std::lower_bound(prefixes.begin(), prefixes.end(), hashes[i].prefix);
+      if (iter != prefixes.end() && *iter == hashes[i].prefix)
+        prefixes.erase(iter);
+    }
+
+    PrefixSetBuilder builder(prefixes);
+    ASSERT_TRUE(builder.GetPrefixSet(hashes)->WriteFile(filename));
+
+    scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
+    ASSERT_TRUE(prefix_set.get());
+    CheckPrefixes(*prefix_set, prefixes);
+
+    EXPECT_TRUE(prefix_set->Exists(hashes[0]));
+    EXPECT_TRUE(prefix_set->Exists(hashes[1]));
+    EXPECT_TRUE(prefix_set->Exists(hashes[2]));
+    EXPECT_FALSE(prefix_set->PrefixExists(hashes[0].prefix));
+    EXPECT_FALSE(prefix_set->PrefixExists(hashes[1].prefix));
+    EXPECT_FALSE(prefix_set->PrefixExists(hashes[2].prefix));
+  }
 }
 
 // Check that |CleanChecksum()| makes an acceptable checksum.
 TEST_F(PrefixSetTest, CorruptionHelpers) {
   base::FilePath filename;
   ASSERT_TRUE(GetPrefixSetFile(&filename));
 
   // This will modify data in |index_|, which will fail the digest check.
   base::ScopedFILE file(base::OpenFile(filename, "r+b"));
   IncrementIntAt(file.get(), kPayloadOffset, 1);
@@ skipping 47 matching lines @@
 TEST_F(PrefixSetTest, CorruptionDeltasSize) {
   base::FilePath filename;
   ASSERT_TRUE(GetPrefixSetFile(&filename));
 
   ASSERT_NO_FATAL_FAILURE(
       ModifyAndCleanChecksum(filename, kDeltasSizeOffset, 1));
   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
   ASSERT_FALSE(prefix_set.get());
 }
 
+// Bad |full_hashes_| size is caught by the sanity check.
+TEST_F(PrefixSetTest, CorruptionFullHashesSize) {
+  base::FilePath filename;
+  ASSERT_TRUE(GetPrefixSetFile(&filename));
+
+  ASSERT_NO_FATAL_FAILURE(
+      ModifyAndCleanChecksum(filename, kFullHashesSizeOffset, 1));
+  scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
+  ASSERT_FALSE(prefix_set.get());
+}
+
 // Test that the digest catches corruption in the middle of the file
 // (in the payload between the header and the digest).
 TEST_F(PrefixSetTest, CorruptionPayload) {
   base::FilePath filename;
   ASSERT_TRUE(GetPrefixSetFile(&filename));
 
   base::ScopedFILE file(base::OpenFile(filename, "r+b"));
   ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), 666, 1));
   file.reset();
   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
@@ skipping 182 matching lines @@
   ASSERT_TRUE(PathService::Get(chrome::DIR_TEST_DATA, &golden_path));
   golden_path = golden_path.AppendASCII("SafeBrowsing");
   golden_path = golden_path.AppendASCII(kBasename);
 
   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(golden_path);
   ASSERT_TRUE(prefix_set.get());
   CheckPrefixes(*prefix_set, ref_prefixes);
 }
 #endif
 
+// Test that a golden v3 file can be read by the current code.  All platforms
+// generating v3 files are little-endian, so there is no point to testing this
+// transition if/when a big-endian port is added.
+#if defined(ARCH_CPU_LITTLE_ENDIAN)
+TEST_F(PrefixSetTest, Version3) {
+  std::vector<SBPrefix> ref_prefixes;
+  ASSERT_TRUE(ReadReferencePrefixes(&ref_prefixes));
+
+  const char kBasename[] = "PrefixSetVersion3";
+  base::FilePath golden_path;
+  ASSERT_TRUE(PathService::Get(chrome::DIR_TEST_DATA, &golden_path));
+  golden_path = golden_path.AppendASCII("SafeBrowsing");
+  golden_path = golden_path.AppendASCII(kBasename);
+
+  scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(golden_path);
+  ASSERT_TRUE(prefix_set.get());
+  CheckPrefixes(*prefix_set, ref_prefixes);
+
+  const SBFullHash kHash1 = SBFullHashForString("www.evil.com/malware.html");
+  const SBFullHash kHash2 = SBFullHashForString("www.evil.com/phishing.html");
+
+  EXPECT_TRUE(prefix_set->Exists(kHash1));
+  EXPECT_TRUE(prefix_set->Exists(kHash2));
+  EXPECT_FALSE(prefix_set->PrefixExists(kHash1.prefix));
+  EXPECT_FALSE(prefix_set->PrefixExists(kHash2.prefix));
+}
+#endif
+
 }  // namespace safe_browsing