chrome/browser/safe_browsing/prefix_set_unittest.cc

   1 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #include "chrome/browser/safe_browsing/prefix_set.h"
   6
   7 #include <algorithm>
   8 #include <iterator>
   9 #include <set>
  10 #include <string>
  11
  12 #include "base/file_util.h"
  13 #include "base/files/scoped_file.h"
  14 #include "base/files/scoped_temp_dir.h"
  15 #include "base/logging.h"
  16 #include "base/md5.h"
  17 #include "base/memory/scoped_ptr.h"
  18 #include "base/path_service.h"
  19 #include "base/rand_util.h"
  20 #include "base/strings/string_number_conversions.h"
  21 #include "base/strings/string_util.h"
  22 #include "chrome/common/chrome_paths.h"
  23 #include "testing/gtest/include/gtest/gtest.h"
  24 #include "testing/platform_test.h"
  25
  26 namespace {
  27
  28 const SBPrefix kHighBitClear = 1000u * 1000u * 1000u;
  29 const SBPrefix kHighBitSet = 3u * 1000u * 1000u * 1000u;
  30
  31 }  // namespace
  32
  33 namespace safe_browsing {
  34
  35 class PrefixSetTest : public PlatformTest {
  36  protected:
  37   // Constants for the v1 format.
  38   static const size_t kMagicOffset = 0 * sizeof(uint32);
  39   static const size_t kVersionOffset = 1 * sizeof(uint32);
  40   static const size_t kIndexSizeOffset = 2 * sizeof(uint32);
  41   static const size_t kDeltasSizeOffset = 3 * sizeof(uint32);
  42   static const size_t kFullHashesSizeOffset = 4 * sizeof(uint32);
  43   static const size_t kPayloadOffset = 5 * sizeof(uint32);
  44
  45   // Generate a set of random prefixes to share between tests.  For
  46   // most tests this generation was a large fraction of the test time.
  47   //
  48   // The set should contain sparse areas where adjacent items are more
  49   // than 2^16 apart, and dense areas where adjacent items are less
  50   // than 2^16 apart.
  51   static void SetUpTestCase() {
  52     // Distribute clusters of prefixes.
  53     for (size_t i = 0; i < 250; ++i) {
  54       // Unsigned for overflow characteristics.
  55       const uint32 base = static_cast<uint32>(base::RandUint64());
  56       for (size_t j = 0; j < 10; ++j) {
  57         const uint32 delta = static_cast<uint32>(base::RandUint64() & 0xFFFF);
  58         const SBPrefix prefix = static_cast<SBPrefix>(base + delta);
  59         shared_prefixes_.push_back(prefix);
  60       }
  61     }
  62
  63     // Lay down a sparsely-distributed layer.
  64     const size_t count = shared_prefixes_.size();
  65     for (size_t i = 0; i < count; ++i) {
  66       const SBPrefix prefix = static_cast<SBPrefix>(base::RandUint64());
  67       shared_prefixes_.push_back(prefix);
  68     }
  69
  70     // Sort for use with PrefixSet constructor.
  71     std::sort(shared_prefixes_.begin(), shared_prefixes_.end());
  72   }
  73
  74   // Check that all elements of |prefixes| are in |prefix_set|, and
  75   // that nearby elements are not (for lack of a more sensible set of
  76   // items to check for absence).
  77   static void CheckPrefixes(const PrefixSet& prefix_set,
  78                             const std::vector<SBPrefix> &prefixes) {
  79     // The set can generate the prefixes it believes it has, so that's
  80     // a good starting point.
  81     std::set<SBPrefix> check(prefixes.begin(), prefixes.end());
  82     std::vector<SBPrefix> prefixes_copy;
  83     prefix_set.GetPrefixes(&prefixes_copy);
  84     EXPECT_EQ(prefixes_copy.size(), check.size());
  85     EXPECT_TRUE(std::equal(check.begin(), check.end(), prefixes_copy.begin()));
  86
  87     for (size_t i = 0; i < prefixes.size(); ++i) {
  88       EXPECT_TRUE(prefix_set.PrefixExists(prefixes[i]));
  89
  90       const SBPrefix left_sibling = prefixes[i] - 1;
  91       if (check.count(left_sibling) == 0)
  92         EXPECT_FALSE(prefix_set.PrefixExists(left_sibling));
  93
  94       const SBPrefix right_sibling = prefixes[i] + 1;
  95       if (check.count(right_sibling) == 0)
  96         EXPECT_FALSE(prefix_set.PrefixExists(right_sibling));
  97     }
  98   }
  99
 100   // Generate a |PrefixSet| file from |shared_prefixes_|, store it in
 101   // a temporary file, and return the filename in |filenamep|.
 102   // Returns |true| on success.
 103   bool GetPrefixSetFile(base::FilePath* filenamep) {
 104     if (!temp_dir_.IsValid() && !temp_dir_.CreateUniqueTempDir())
 105       return false;
 106
 107     base::FilePath filename = temp_dir_.path().AppendASCII("PrefixSetTest");
 108
 109     PrefixSetBuilder builder(shared_prefixes_);
 110     if (!builder.GetPrefixSetNoHashes()->WriteFile(filename))
 111       return false;
 112
 113     *filenamep = filename;
 114     return true;
 115   }
 116
 117   // Helper function to read the uint32 value at |offset|, increment it
 118   // by |inc|, and write it back in place.  |fp| should be opened in
 119   // r+ mode.
 120   static void IncrementIntAt(FILE* fp, long offset, int inc) {
 121     uint32 value = 0;
 122
 123     ASSERT_NE(-1, fseek(fp, offset, SEEK_SET));
 124     ASSERT_EQ(1U, fread(&value, sizeof(value), 1, fp));
 125
 126     value += inc;
 127
 128     ASSERT_NE(-1, fseek(fp, offset, SEEK_SET));
 129     ASSERT_EQ(1U, fwrite(&value, sizeof(value), 1, fp));
 130   }
 131
 132   // Helper function to re-generated |fp|'s checksum to be correct for
 133   // the file's contents.  |fp| should be opened in r+ mode.
 134   static void CleanChecksum(FILE* fp) {
 135     base::MD5Context context;
 136     base::MD5Init(&context);
 137
 138     ASSERT_NE(-1, fseek(fp, 0, SEEK_END));
 139     long file_size = ftell(fp);
 140
 141     using base::MD5Digest;
 142     size_t payload_size = static_cast<size_t>(file_size) - sizeof(MD5Digest);
 143     size_t digested_size = 0;
 144     ASSERT_NE(-1, fseek(fp, 0, SEEK_SET));
 145     while (digested_size < payload_size) {
 146       char buf[1024];
 147       size_t nitems = std::min(payload_size - digested_size, sizeof(buf));
 148       ASSERT_EQ(nitems, fread(buf, 1, nitems, fp));
 149       base::MD5Update(&context, base::StringPiece(buf, nitems));
 150       digested_size += nitems;
 151     }
 152     ASSERT_EQ(digested_size, payload_size);
 153     ASSERT_EQ(static_cast<long>(digested_size), ftell(fp));
 154
 155     base::MD5Digest new_digest;
 156     base::MD5Final(&new_digest, &context);
 157     ASSERT_NE(-1, fseek(fp, digested_size, SEEK_SET));
 158     ASSERT_EQ(1U, fwrite(&new_digest, sizeof(new_digest), 1, fp));
 159     ASSERT_EQ(file_size, ftell(fp));
 160   }
 161
 162   // Open |filename| and increment the uint32 at |offset| by |inc|.
 163   // Then re-generate the checksum to account for the new contents.
 164   void ModifyAndCleanChecksum(const base::FilePath& filename, long offset,
 165                               int inc) {
 166     int64 size_64;
 167     ASSERT_TRUE(base::GetFileSize(filename, &size_64));
 168
 169     base::ScopedFILE file(base::OpenFile(filename, "r+b"));
 170     IncrementIntAt(file.get(), offset, inc);
 171     CleanChecksum(file.get());
 172     file.reset();
 173
 174     int64 new_size_64;
 175     ASSERT_TRUE(base::GetFileSize(filename, &new_size_64));
 176     ASSERT_EQ(new_size_64, size_64);
 177   }
 178
 179   // Fill |prefixes| with values read from a reference file.  The reference file
 180   // was generated from a specific |shared_prefixes_|.
 181   bool ReadReferencePrefixes(std::vector<SBPrefix>* prefixes) {
 182     const char kRefname[] = "PrefixSetRef";
 183     base::FilePath ref_path;
 184     if (!PathService::Get(chrome::DIR_TEST_DATA, &ref_path))
 185       return false;
 186     ref_path = ref_path.AppendASCII("SafeBrowsing");
 187     ref_path = ref_path.AppendASCII(kRefname);
 188
 189     base::ScopedFILE file(base::OpenFile(ref_path, "r"));
 190     if (!file.get())
 191       return false;
 192     char buf[1024];
 193     while (fgets(buf, sizeof(buf), file.get())) {
 194       std::string trimmed;
 195       if (base::TRIM_TRAILING !=
 196           base::TrimWhitespace(buf, base::TRIM_ALL, &trimmed))
 197         return false;
 198       unsigned prefix;
 199       if (!base::StringToUint(trimmed, &prefix))
 200         return false;
 201       prefixes->push_back(prefix);
 202     }
 203     return true;
 204   }
 205
 206   // Tests should not modify this shared resource.
 207   static std::vector<SBPrefix> shared_prefixes_;
 208
 209   base::ScopedTempDir temp_dir_;
 210 };
 211
 212 std::vector<SBPrefix> PrefixSetTest::shared_prefixes_;
 213
 214 // Test that a small sparse random input works.
 215 TEST_F(PrefixSetTest, Baseline) {
 216   PrefixSetBuilder builder(shared_prefixes_);
 217   CheckPrefixes(*builder.GetPrefixSetNoHashes(), shared_prefixes_);
 218 }
 219
 220 // Test that the empty set doesn't appear to have anything in it.
 221 TEST_F(PrefixSetTest, Empty) {
 222   const std::vector<SBPrefix> empty;
 223   PrefixSetBuilder builder(empty);
 224   scoped_ptr<PrefixSet> prefix_set = builder.GetPrefixSetNoHashes();
 225   for (size_t i = 0; i < shared_prefixes_.size(); ++i) {
 226     EXPECT_FALSE(prefix_set->PrefixExists(shared_prefixes_[i]));
 227   }
 228 }
 229
 230 // Single-element set should work fine.
 231 TEST_F(PrefixSetTest, OneElement) {
 232   const std::vector<SBPrefix> prefixes(100, 0u);
 233   PrefixSetBuilder builder(prefixes);
 234   scoped_ptr<PrefixSet> prefix_set = builder.GetPrefixSetNoHashes();
 235   EXPECT_FALSE(prefix_set->PrefixExists(static_cast<SBPrefix>(-1)));
 236   EXPECT_TRUE(prefix_set->PrefixExists(prefixes[0]));
 237   EXPECT_FALSE(prefix_set->PrefixExists(1u));
 238
 239   // Check that |GetPrefixes()| returns the same set of prefixes as
 240   // was passed in.
 241   std::vector<SBPrefix> prefixes_copy;
 242   prefix_set->GetPrefixes(&prefixes_copy);
 243   EXPECT_EQ(1U, prefixes_copy.size());
 244   EXPECT_EQ(prefixes[0], prefixes_copy[0]);
 245 }
 246
 247 // Edges of the 32-bit integer range.
 248 TEST_F(PrefixSetTest, IntMinMax) {
 249   std::vector<SBPrefix> prefixes;
 250
 251   // Using bit patterns rather than portable constants because this
 252   // really is testing how the entire 32-bit integer range is handled.
 253   prefixes.push_back(0x00000000);
 254   prefixes.push_back(0x0000FFFF);
 255   prefixes.push_back(0x7FFF0000);
 256   prefixes.push_back(0x7FFFFFFF);
 257   prefixes.push_back(0x80000000);
 258   prefixes.push_back(0x8000FFFF);
 259   prefixes.push_back(0xFFFF0000);
 260   prefixes.push_back(0xFFFFFFFF);
 261
 262   std::sort(prefixes.begin(), prefixes.end());
 263   PrefixSetBuilder builder(prefixes);
 264   scoped_ptr<PrefixSet> prefix_set = builder.GetPrefixSetNoHashes();
 265
 266   // Check that |GetPrefixes()| returns the same set of prefixes as
 267   // was passed in.
 268   std::vector<SBPrefix> prefixes_copy;
 269   prefix_set->GetPrefixes(&prefixes_copy);
 270   ASSERT_EQ(prefixes_copy.size(), prefixes.size());
 271   EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
 272                          prefixes_copy.begin()));
 273 }
 274
 275 // A range with only large deltas.
 276 TEST_F(PrefixSetTest, AllBig) {
 277   std::vector<SBPrefix> prefixes;
 278
 279   const unsigned kDelta = 10 * 1000 * 1000;
 280   for (SBPrefix prefix = kHighBitSet;
 281        prefix < kHighBitClear; prefix += kDelta) {
 282     prefixes.push_back(prefix);
 283   }
 284
 285   std::sort(prefixes.begin(), prefixes.end());
 286   PrefixSetBuilder builder(prefixes);
 287   scoped_ptr<PrefixSet> prefix_set = builder.GetPrefixSetNoHashes();
 288
 289   // Check that |GetPrefixes()| returns the same set of prefixes as
 290   // was passed in.
 291   std::vector<SBPrefix> prefixes_copy;
 292   prefix_set->GetPrefixes(&prefixes_copy);
 293   prefixes.erase(std::unique(prefixes.begin(), prefixes.end()), prefixes.end());
 294   EXPECT_EQ(prefixes_copy.size(), prefixes.size());
 295   EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
 296                          prefixes_copy.begin()));
 297 }
 298
 299 // Use artificial inputs to test various edge cases in PrefixExists().  Items
 300 // before the lowest item aren't present.  Items after the largest item aren't
 301 // present.  Create a sequence of items with deltas above and below 2^16, and
 302 // make sure they're all present.  Create a very long sequence with deltas below
 303 // 2^16 to test crossing |kMaxRun|.
 304 TEST_F(PrefixSetTest, EdgeCases) {
 305   std::vector<SBPrefix> prefixes;
 306
 307   // Put in a high-bit prefix.
 308   SBPrefix prefix = kHighBitSet;
 309   prefixes.push_back(prefix);
 310
 311   // Add a sequence with very large deltas.
 312   unsigned delta = 100 * 1000 * 1000;
 313   for (int i = 0; i < 10; ++i) {
 314     prefix += delta;
 315     prefixes.push_back(prefix);
 316   }
 317
 318   // Add a sequence with deltas that start out smaller than the
 319   // maximum delta, and end up larger.  Also include some duplicates.
 320   delta = 256 * 256 - 100;
 321   for (int i = 0; i < 200; ++i) {
 322     prefix += delta;
 323     prefixes.push_back(prefix);
 324     prefixes.push_back(prefix);
 325     delta++;
 326   }
 327
 328   // Add a long sequence with deltas smaller than the maximum delta,
 329   // so a new index item will be injected.
 330   delta = 256 * 256 - 1;
 331   prefix = kHighBitClear - delta * 1000;
 332   prefixes.push_back(prefix);
 333   for (int i = 0; i < 1000; ++i) {
 334     prefix += delta;
 335     prefixes.push_back(prefix);
 336     delta--;
 337   }
 338
 339   std::sort(prefixes.begin(), prefixes.end());
 340   PrefixSetBuilder builder(prefixes);
 341   scoped_ptr<PrefixSet> prefix_set = builder.GetPrefixSetNoHashes();
 342
 343   // Check that |GetPrefixes()| returns the same set of prefixes as
 344   // was passed in.
 345   std::vector<SBPrefix> prefixes_copy;
 346   prefix_set->GetPrefixes(&prefixes_copy);
 347   prefixes.erase(std::unique(prefixes.begin(), prefixes.end()), prefixes.end());
 348   EXPECT_EQ(prefixes_copy.size(), prefixes.size());
 349   EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
 350                          prefixes_copy.begin()));
 351
 352   // Items before and after the set are not present, and don't crash.
 353   EXPECT_FALSE(prefix_set->PrefixExists(kHighBitSet - 100));
 354   EXPECT_FALSE(prefix_set->PrefixExists(kHighBitClear + 100));
 355
 356   // Check that the set correctly flags all of the inputs, and also
 357   // check items just above and below the inputs to make sure they
 358   // aren't present.
 359   for (size_t i = 0; i < prefixes.size(); ++i) {
 360     EXPECT_TRUE(prefix_set->PrefixExists(prefixes[i]));
 361
 362     EXPECT_FALSE(prefix_set->PrefixExists(prefixes[i] - 1));
 363     EXPECT_FALSE(prefix_set->PrefixExists(prefixes[i] + 1));
 364   }
 365 }
 366
 367 // Test writing a prefix set to disk and reading it back in.
 368 TEST_F(PrefixSetTest, ReadWrite) {
 369   base::FilePath filename;
 370
 371   // Write the sample prefix set out, read it back in, and check all
 372   // the prefixes.  Leaves the path in |filename|.
 373   {
 374     ASSERT_TRUE(GetPrefixSetFile(&filename));
 375     scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
 376     ASSERT_TRUE(prefix_set.get());
 377     CheckPrefixes(*prefix_set, shared_prefixes_);
 378   }
 379
 380   // Test writing and reading a very sparse set containing no deltas.
 381   {
 382     std::vector<SBPrefix> prefixes;
 383     prefixes.push_back(kHighBitClear);
 384     prefixes.push_back(kHighBitSet);
 385
 386     PrefixSetBuilder builder(prefixes);
 387     ASSERT_TRUE(builder.GetPrefixSetNoHashes()->WriteFile(filename));
 388
 389     scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
 390     ASSERT_TRUE(prefix_set.get());
 391     CheckPrefixes(*prefix_set, prefixes);
 392   }
 393
 394   // Test writing and reading an empty set.
 395   {
 396     std::vector<SBPrefix> prefixes;
 397     PrefixSetBuilder builder(prefixes);
 398     ASSERT_TRUE(builder.GetPrefixSetNoHashes()->WriteFile(filename));
 399
 400     scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
 401     ASSERT_TRUE(prefix_set.get());
 402     CheckPrefixes(*prefix_set, prefixes);
 403   }
 404
 405   // Test that full hashes are persisted.
 406   {
 407     std::vector<SBFullHash> hashes;
 408     hashes.push_back(SBFullHashForString("one"));
 409     hashes.push_back(SBFullHashForString("two"));
 410     hashes.push_back(SBFullHashForString("three"));
 411
 412     std::vector<SBPrefix> prefixes(shared_prefixes_);
 413
 414     // Remove any collisions from the prefixes.
 415     for (size_t i = 0; i < hashes.size(); ++i) {
 416       std::vector<SBPrefix>::iterator iter =
 417           std::lower_bound(prefixes.begin(), prefixes.end(), hashes[i].prefix);
 418       if (iter != prefixes.end() && *iter == hashes[i].prefix)
 419         prefixes.erase(iter);
 420     }
 421
 422     PrefixSetBuilder builder(prefixes);
 423     ASSERT_TRUE(builder.GetPrefixSet(hashes)->WriteFile(filename));
 424
 425     scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
 426     ASSERT_TRUE(prefix_set.get());
 427     CheckPrefixes(*prefix_set, prefixes);
 428
 429     EXPECT_TRUE(prefix_set->Exists(hashes[0]));
 430     EXPECT_TRUE(prefix_set->Exists(hashes[1]));
 431     EXPECT_TRUE(prefix_set->Exists(hashes[2]));
 432     EXPECT_FALSE(prefix_set->PrefixExists(hashes[0].prefix));
 433     EXPECT_FALSE(prefix_set->PrefixExists(hashes[1].prefix));
 434     EXPECT_FALSE(prefix_set->PrefixExists(hashes[2].prefix));
 435   }
 436 }
 437
 438 // Check that |CleanChecksum()| makes an acceptable checksum.
 439 TEST_F(PrefixSetTest, CorruptionHelpers) {
 440   base::FilePath filename;
 441   ASSERT_TRUE(GetPrefixSetFile(&filename));
 442
 443   // This will modify data in |index_|, which will fail the digest check.
 444   base::ScopedFILE file(base::OpenFile(filename, "r+b"));
 445   IncrementIntAt(file.get(), kPayloadOffset, 1);
 446   file.reset();
 447   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
 448   ASSERT_FALSE(prefix_set.get());
 449
 450   // Fix up the checksum and it will read successfully (though the
 451   // data will be wrong).
 452   file.reset(base::OpenFile(filename, "r+b"));
 453   CleanChecksum(file.get());
 454   file.reset();
 455   prefix_set = PrefixSet::LoadFile(filename);
 456   ASSERT_TRUE(prefix_set.get());
 457 }
 458
 459 // Bad magic is caught by the sanity check.
 460 TEST_F(PrefixSetTest, CorruptionMagic) {
 461   base::FilePath filename;
 462   ASSERT_TRUE(GetPrefixSetFile(&filename));
 463
 464   ASSERT_NO_FATAL_FAILURE(
 465       ModifyAndCleanChecksum(filename, kMagicOffset, 1));
 466   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
 467   ASSERT_FALSE(prefix_set.get());
 468 }
 469
 470 // Bad version is caught by the sanity check.
 471 TEST_F(PrefixSetTest, CorruptionVersion) {
 472   base::FilePath filename;
 473   ASSERT_TRUE(GetPrefixSetFile(&filename));
 474
 475   ASSERT_NO_FATAL_FAILURE(
 476       ModifyAndCleanChecksum(filename, kVersionOffset, 10));
 477   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
 478   ASSERT_FALSE(prefix_set.get());
 479 }
 480
 481 // Bad |index_| size is caught by the sanity check.
 482 TEST_F(PrefixSetTest, CorruptionIndexSize) {
 483   base::FilePath filename;
 484   ASSERT_TRUE(GetPrefixSetFile(&filename));
 485
 486   ASSERT_NO_FATAL_FAILURE(
 487       ModifyAndCleanChecksum(filename, kIndexSizeOffset, 1));
 488   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
 489   ASSERT_FALSE(prefix_set.get());
 490 }
 491
 492 // Bad |deltas_| size is caught by the sanity check.
 493 TEST_F(PrefixSetTest, CorruptionDeltasSize) {
 494   base::FilePath filename;
 495   ASSERT_TRUE(GetPrefixSetFile(&filename));
 496
 497   ASSERT_NO_FATAL_FAILURE(
 498       ModifyAndCleanChecksum(filename, kDeltasSizeOffset, 1));
 499   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
 500   ASSERT_FALSE(prefix_set.get());
 501 }
 502
 503 // Bad |full_hashes_| size is caught by the sanity check.
 504 TEST_F(PrefixSetTest, CorruptionFullHashesSize) {
 505   base::FilePath filename;
 506   ASSERT_TRUE(GetPrefixSetFile(&filename));
 507
 508   ASSERT_NO_FATAL_FAILURE(
 509       ModifyAndCleanChecksum(filename, kFullHashesSizeOffset, 1));
 510   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
 511   ASSERT_FALSE(prefix_set.get());
 512 }
 513
 514 // Test that the digest catches corruption in the middle of the file
 515 // (in the payload between the header and the digest).
 516 TEST_F(PrefixSetTest, CorruptionPayload) {
 517   base::FilePath filename;
 518   ASSERT_TRUE(GetPrefixSetFile(&filename));
 519
 520   base::ScopedFILE file(base::OpenFile(filename, "r+b"));
 521   ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), 666, 1));
 522   file.reset();
 523   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
 524   ASSERT_FALSE(prefix_set.get());
 525 }
 526
 527 // Test corruption in the digest itself.
 528 TEST_F(PrefixSetTest, CorruptionDigest) {
 529   base::FilePath filename;
 530   ASSERT_TRUE(GetPrefixSetFile(&filename));
 531
 532   int64 size_64;
 533   ASSERT_TRUE(base::GetFileSize(filename, &size_64));
 534   base::ScopedFILE file(base::OpenFile(filename, "r+b"));
 535   long digest_offset = static_cast<long>(size_64 - sizeof(base::MD5Digest));
 536   ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), digest_offset, 1));
 537   file.reset();
 538   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
 539   ASSERT_FALSE(prefix_set.get());
 540 }
 541
 542 // Test excess data after the digest (fails the size test).
 543 TEST_F(PrefixSetTest, CorruptionExcess) {
 544   base::FilePath filename;
 545   ASSERT_TRUE(GetPrefixSetFile(&filename));
 546
 547   // Add some junk to the trunk.
 548   base::ScopedFILE file(base::OpenFile(filename, "ab"));
 549   const char buf[] = "im in ur base, killing ur d00dz.";
 550   ASSERT_EQ(strlen(buf), fwrite(buf, 1, strlen(buf), file.get()));
 551   file.reset();
 552   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
 553   ASSERT_FALSE(prefix_set.get());
 554 }
 555
 556 // Test that files which had 64-bit size_t are discarded.
 557 TEST_F(PrefixSetTest, SizeTRecovery) {
 558   base::FilePath filename;
 559   ASSERT_TRUE(GetPrefixSetFile(&filename));
 560
 561   // Open the file for rewrite.
 562   base::ScopedFILE file(base::OpenFile(filename, "r+b"));
 563
 564   // Leave existing magic and version.
 565   ASSERT_NE(-1, fseek(file.get(), sizeof(uint32) * 2, SEEK_SET));
 566
 567   // Indicate two index values and two deltas.
 568   uint32 val = 2;
 569   ASSERT_EQ(sizeof(val), fwrite(&val, 1, sizeof(val), file.get()));
 570   ASSERT_EQ(sizeof(val), fwrite(&val, 1, sizeof(val), file.get()));
 571
 572   // Write two index values with 64-bit "size_t".
 573   std::pair<SBPrefix, uint64> item;
 574   memset(&item, 0, sizeof(item));  // Includes any padding.
 575   item.first = 17;
 576   item.second = 0;
 577   ASSERT_EQ(sizeof(item), fwrite(&item, 1, sizeof(item), file.get()));
 578   item.first = 100042;
 579   item.second = 1;
 580   ASSERT_EQ(sizeof(item), fwrite(&item, 1, sizeof(item), file.get()));
 581
 582   // Write two delta values.
 583   uint16 delta = 23;
 584   ASSERT_EQ(sizeof(delta), fwrite(&delta, 1, sizeof(delta), file.get()));
 585   ASSERT_EQ(sizeof(delta), fwrite(&delta, 1, sizeof(delta), file.get()));
 586
 587   // Leave space for the digest at the end, and regenerate it.
 588   base::MD5Digest dummy = { { 0 } };
 589   ASSERT_EQ(sizeof(dummy), fwrite(&dummy, 1, sizeof(dummy), file.get()));
 590   ASSERT_TRUE(base::TruncateFile(file.get()));
 591   CleanChecksum(file.get());
 592   file.reset();  // Flush updates.
 593
 594   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
 595   ASSERT_FALSE(prefix_set.get());
 596 }
 597
 598 // Test Exists() against full hashes passed to builder.
 599 TEST_F(PrefixSetTest, FullHashBuild) {
 600   const SBFullHash kHash1 = SBFullHashForString("one");
 601   const SBFullHash kHash2 = SBFullHashForString("two");
 602   const SBFullHash kHash3 = SBFullHashForString("three");
 603   const SBFullHash kHash4 = SBFullHashForString("four");
 604   const SBFullHash kHash5 = SBFullHashForString("five");
 605   const SBFullHash kHash6 = SBFullHashForString("six");
 606
 607   std::vector<SBPrefix> prefixes;
 608   prefixes.push_back(kHash1.prefix);
 609   prefixes.push_back(kHash2.prefix);
 610   std::sort(prefixes.begin(), prefixes.end());
 611
 612   std::vector<SBFullHash> hashes;
 613   hashes.push_back(kHash4);
 614   hashes.push_back(kHash5);
 615
 616   PrefixSetBuilder builder(prefixes);
 617   scoped_ptr<PrefixSet> prefix_set = builder.GetPrefixSet(hashes);
 618
 619   EXPECT_TRUE(prefix_set->Exists(kHash1));
 620   EXPECT_TRUE(prefix_set->Exists(kHash2));
 621   EXPECT_FALSE(prefix_set->Exists(kHash3));
 622   EXPECT_TRUE(prefix_set->Exists(kHash4));
 623   EXPECT_TRUE(prefix_set->Exists(kHash5));
 624   EXPECT_FALSE(prefix_set->Exists(kHash6));
 625
 626   EXPECT_TRUE(prefix_set->PrefixExists(kHash1.prefix));
 627   EXPECT_TRUE(prefix_set->PrefixExists(kHash2.prefix));
 628   EXPECT_FALSE(prefix_set->PrefixExists(kHash3.prefix));
 629   EXPECT_FALSE(prefix_set->PrefixExists(kHash4.prefix));
 630   EXPECT_FALSE(prefix_set->PrefixExists(kHash5.prefix));
 631   EXPECT_FALSE(prefix_set->PrefixExists(kHash6.prefix));
 632 }
 633
 634 // Test that a version 1 file is discarded on read.
 635 TEST_F(PrefixSetTest, ReadSigned) {
 636   base::FilePath filename;
 637   ASSERT_TRUE(GetPrefixSetFile(&filename));
 638
 639   // Open the file for rewrite.
 640   base::ScopedFILE file(base::OpenFile(filename, "r+b"));
 641
 642   // Leave existing magic.
 643   ASSERT_NE(-1, fseek(file.get(), sizeof(uint32), SEEK_SET));
 644
 645   // Version 1.
 646   uint32 version = 1;
 647   ASSERT_EQ(sizeof(version), fwrite(&version, 1, sizeof(version), file.get()));
 648
 649   // Indicate two index values and two deltas.
 650   uint32 val = 2;
 651   ASSERT_EQ(sizeof(val), fwrite(&val, 1, sizeof(val), file.get()));
 652   ASSERT_EQ(sizeof(val), fwrite(&val, 1, sizeof(val), file.get()));
 653
 654   std::pair<int32, uint32> item;
 655   memset(&item, 0, sizeof(item));  // Includes any padding.
 656   item.first = -1000;
 657   item.second = 0;
 658   ASSERT_EQ(sizeof(item), fwrite(&item, 1, sizeof(item), file.get()));
 659   item.first = 1000;
 660   item.second = 1;
 661   ASSERT_EQ(sizeof(item), fwrite(&item, 1, sizeof(item), file.get()));
 662
 663   // Write two delta values.
 664   uint16 delta = 23;
 665   ASSERT_EQ(sizeof(delta), fwrite(&delta, 1, sizeof(delta), file.get()));
 666   ASSERT_EQ(sizeof(delta), fwrite(&delta, 1, sizeof(delta), file.get()));
 667
 668   // Leave space for the digest at the end, and regenerate it.
 669   base::MD5Digest dummy = { { 0 } };
 670   ASSERT_EQ(sizeof(dummy), fwrite(&dummy, 1, sizeof(dummy), file.get()));
 671   ASSERT_TRUE(base::TruncateFile(file.get()));
 672   CleanChecksum(file.get());
 673   file.reset();  // Flush updates.
 674
 675   scoped_ptr<safe_browsing::PrefixSet>
 676       prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
 677   ASSERT_FALSE(prefix_set.get());
 678 }
 679
 680 // Test that a golden v2 file is discarded on read.  All platforms generating v2
 681 // files are little-endian, so there is no point to testing this transition
 682 // if/when a big-endian port is added.
 683 #if defined(ARCH_CPU_LITTLE_ENDIAN)
 684 TEST_F(PrefixSetTest, Version2) {
 685   std::vector<SBPrefix> ref_prefixes;
 686   ASSERT_TRUE(ReadReferencePrefixes(&ref_prefixes));
 687
 688   const char kBasename[] = "PrefixSetVersion2";
 689   base::FilePath golden_path;
 690   ASSERT_TRUE(PathService::Get(chrome::DIR_TEST_DATA, &golden_path));
 691   golden_path = golden_path.AppendASCII("SafeBrowsing");
 692   golden_path = golden_path.AppendASCII(kBasename);
 693
 694   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(golden_path);
 695   ASSERT_FALSE(prefix_set.get());
 696 }
 697 #endif
 698
 699 // Test that a golden v3 file can be read by the current code.  All platforms
 700 // generating v3 files are little-endian, so there is no point to testing this
 701 // transition if/when a big-endian port is added.
 702 #if defined(ARCH_CPU_LITTLE_ENDIAN)
 703 TEST_F(PrefixSetTest, Version3) {
 704   std::vector<SBPrefix> ref_prefixes;
 705   ASSERT_TRUE(ReadReferencePrefixes(&ref_prefixes));
 706
 707   const char kBasename[] = "PrefixSetVersion3";
 708   base::FilePath golden_path;
 709   ASSERT_TRUE(PathService::Get(chrome::DIR_TEST_DATA, &golden_path));
 710   golden_path = golden_path.AppendASCII("SafeBrowsing");
 711   golden_path = golden_path.AppendASCII(kBasename);
 712
 713   scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(golden_path);
 714   ASSERT_TRUE(prefix_set.get());
 715   CheckPrefixes(*prefix_set, ref_prefixes);
 716
 717   const SBFullHash kHash1 = SBFullHashForString("www.evil.com/malware.html");
 718   const SBFullHash kHash2 = SBFullHashForString("www.evil.com/phishing.html");
 719
 720   EXPECT_TRUE(prefix_set->Exists(kHash1));
 721   EXPECT_TRUE(prefix_set->Exists(kHash2));
 722   EXPECT_FALSE(prefix_set->PrefixExists(kHash1.prefix));
 723   EXPECT_FALSE(prefix_set->PrefixExists(kHash2.prefix));
 724 }
 725 #endif
 726
 727 }  // namespace safe_browsing