src/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
  10 #include "base/file_util.h"
  11 #include "base/files/scoped_temp_dir.h"
  12 #include "base/logging.h"
  13 #include "base/md5.h"
  14 #include "base/memory/scoped_ptr.h"
  15 #include "base/rand_util.h"
  16 #include "testing/gtest/include/gtest/gtest.h"
  17 #include "testing/platform_test.h"
  18
  19 namespace {
  20
  21 class PrefixSetTest : public PlatformTest {
  22  protected:
  23   // Constants for the v1 format.
  24   static const size_t kMagicOffset = 0 * sizeof(uint32);
  25   static const size_t kVersionOffset = 1 * sizeof(uint32);
  26   static const size_t kIndexSizeOffset = 2 * sizeof(uint32);
  27   static const size_t kDeltasSizeOffset = 3 * sizeof(uint32);
  28   static const size_t kPayloadOffset = 4 * sizeof(uint32);
  29
  30   // Generate a set of random prefixes to share between tests.  For
  31   // most tests this generation was a large fraction of the test time.
  32   static void SetUpTestCase() {
  33     for (size_t i = 0; i < 50000; ++i) {
  34       const SBPrefix prefix = static_cast<SBPrefix>(base::RandUint64());
  35       shared_prefixes_.push_back(prefix);
  36     }
  37
  38     // Sort for use with PrefixSet constructor.
  39     std::sort(shared_prefixes_.begin(), shared_prefixes_.end());
  40   }
  41
  42   // Check that all elements of |prefixes| are in |prefix_set|, and
  43   // that nearby elements are not (for lack of a more sensible set of
  44   // items to check for absence).
  45   static void CheckPrefixes(const safe_browsing::PrefixSet& prefix_set,
  46                             const std::vector<SBPrefix> &prefixes) {
  47     // The set can generate the prefixes it believes it has, so that's
  48     // a good starting point.
  49     std::set<SBPrefix> check(prefixes.begin(), prefixes.end());
  50     std::vector<SBPrefix> prefixes_copy;
  51     prefix_set.GetPrefixes(&prefixes_copy);
  52     EXPECT_EQ(prefixes_copy.size(), check.size());
  53     EXPECT_TRUE(std::equal(check.begin(), check.end(), prefixes_copy.begin()));
  54
  55     for (size_t i = 0; i < prefixes.size(); ++i) {
  56       EXPECT_TRUE(prefix_set.Exists(prefixes[i]));
  57
  58       const SBPrefix left_sibling = prefixes[i] - 1;
  59       if (check.count(left_sibling) == 0)
  60         EXPECT_FALSE(prefix_set.Exists(left_sibling));
  61
  62       const SBPrefix right_sibling = prefixes[i] + 1;
  63       if (check.count(right_sibling) == 0)
  64         EXPECT_FALSE(prefix_set.Exists(right_sibling));
  65     }
  66   }
  67
  68   // Generate a |PrefixSet| file from |shared_prefixes_|, store it in
  69   // a temporary file, and return the filename in |filenamep|.
  70   // Returns |true| on success.
  71   bool GetPrefixSetFile(base::FilePath* filenamep) {
  72     if (!temp_dir_.IsValid() && !temp_dir_.CreateUniqueTempDir())
  73       return false;
  74
  75     base::FilePath filename = temp_dir_.path().AppendASCII("PrefixSetTest");
  76
  77     safe_browsing::PrefixSet prefix_set(shared_prefixes_);
  78     if (!prefix_set.WriteFile(filename))
  79       return false;
  80
  81     *filenamep = filename;
  82     return true;
  83   }
  84
  85   // Helper function to read the int32 value at |offset|, increment it
  86   // by |inc|, and write it back in place.  |fp| should be opened in
  87   // r+ mode.
  88   static void IncrementIntAt(FILE* fp, long offset, int inc) {
  89     int32 value = 0;
  90
  91     ASSERT_NE(-1, fseek(fp, offset, SEEK_SET));
  92     ASSERT_EQ(1U, fread(&value, sizeof(value), 1, fp));
  93
  94     value += inc;
  95
  96     ASSERT_NE(-1, fseek(fp, offset, SEEK_SET));
  97     ASSERT_EQ(1U, fwrite(&value, sizeof(value), 1, fp));
  98   }
  99
 100   // Helper function to re-generated |fp|'s checksum to be correct for
 101   // the file's contents.  |fp| should be opened in r+ mode.
 102   static void CleanChecksum(FILE* fp) {
 103     base::MD5Context context;
 104     base::MD5Init(&context);
 105
 106     ASSERT_NE(-1, fseek(fp, 0, SEEK_END));
 107     long file_size = ftell(fp);
 108
 109     using base::MD5Digest;
 110     size_t payload_size = static_cast<size_t>(file_size) - sizeof(MD5Digest);
 111     size_t digested_size = 0;
 112     ASSERT_NE(-1, fseek(fp, 0, SEEK_SET));
 113     while (digested_size < payload_size) {
 114       char buf[1024];
 115       size_t nitems = std::min(payload_size - digested_size, sizeof(buf));
 116       ASSERT_EQ(nitems, fread(buf, 1, nitems, fp));
 117       base::MD5Update(&context, base::StringPiece(buf, nitems));
 118       digested_size += nitems;
 119     }
 120     ASSERT_EQ(digested_size, payload_size);
 121     ASSERT_EQ(static_cast<long>(digested_size), ftell(fp));
 122
 123     base::MD5Digest new_digest;
 124     base::MD5Final(&new_digest, &context);
 125     ASSERT_NE(-1, fseek(fp, digested_size, SEEK_SET));
 126     ASSERT_EQ(1U, fwrite(&new_digest, sizeof(new_digest), 1, fp));
 127     ASSERT_EQ(file_size, ftell(fp));
 128   }
 129
 130   // Open |filename| and increment the int32 at |offset| by |inc|.
 131   // Then re-generate the checksum to account for the new contents.
 132   void ModifyAndCleanChecksum(const base::FilePath& filename, long offset,
 133                               int inc) {
 134     int64 size_64;
 135     ASSERT_TRUE(file_util::GetFileSize(filename, &size_64));
 136
 137     file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b"));
 138     IncrementIntAt(file.get(), offset, inc);
 139     CleanChecksum(file.get());
 140     file.reset();
 141
 142     int64 new_size_64;
 143     ASSERT_TRUE(file_util::GetFileSize(filename, &new_size_64));
 144     ASSERT_EQ(new_size_64, size_64);
 145   }
 146
 147   // Tests should not modify this shared resource.
 148   static std::vector<SBPrefix> shared_prefixes_;
 149
 150   base::ScopedTempDir temp_dir_;
 151 };
 152
 153 std::vector<SBPrefix> PrefixSetTest::shared_prefixes_;
 154
 155 // Test that a small sparse random input works.
 156 TEST_F(PrefixSetTest, Baseline) {
 157   safe_browsing::PrefixSet prefix_set(shared_prefixes_);
 158   CheckPrefixes(prefix_set, shared_prefixes_);
 159 }
 160
 161 // Test that the empty set doesn't appear to have anything in it.
 162 TEST_F(PrefixSetTest, Empty) {
 163   const std::vector<SBPrefix> empty;
 164   safe_browsing::PrefixSet prefix_set(empty);
 165   for (size_t i = 0; i < shared_prefixes_.size(); ++i) {
 166     EXPECT_FALSE(prefix_set.Exists(shared_prefixes_[i]));
 167   }
 168 }
 169
 170 // Single-element set should work fine.
 171 TEST_F(PrefixSetTest, OneElement) {
 172   const std::vector<SBPrefix> prefixes(100, 0);
 173   safe_browsing::PrefixSet prefix_set(prefixes);
 174   EXPECT_FALSE(prefix_set.Exists(-1));
 175   EXPECT_TRUE(prefix_set.Exists(prefixes[0]));
 176   EXPECT_FALSE(prefix_set.Exists(1));
 177
 178   // Check that |GetPrefixes()| returns the same set of prefixes as
 179   // was passed in.
 180   std::vector<SBPrefix> prefixes_copy;
 181   prefix_set.GetPrefixes(&prefixes_copy);
 182   EXPECT_EQ(1U, prefixes_copy.size());
 183   EXPECT_EQ(prefixes[0], prefixes_copy[0]);
 184 }
 185
 186 // Edges of the 32-bit integer range.
 187 TEST_F(PrefixSetTest, IntMinMax) {
 188   std::vector<SBPrefix> prefixes;
 189
 190   // Using bit patterns rather than portable constants because this
 191   // really is testing how the entire 32-bit integer range is handled.
 192   prefixes.push_back(0x00000000);
 193   prefixes.push_back(0x0000FFFF);
 194   prefixes.push_back(0x7FFF0000);
 195   prefixes.push_back(0x7FFFFFFF);
 196   prefixes.push_back(0x80000000);
 197   prefixes.push_back(0x8000FFFF);
 198   prefixes.push_back(0xFFFF0000);
 199   prefixes.push_back(0xFFFFFFFF);
 200
 201   std::sort(prefixes.begin(), prefixes.end());
 202   safe_browsing::PrefixSet prefix_set(prefixes);
 203
 204   // Check that |GetPrefixes()| returns the same set of prefixes as
 205   // was passed in.
 206   std::vector<SBPrefix> prefixes_copy;
 207   prefix_set.GetPrefixes(&prefixes_copy);
 208   ASSERT_EQ(prefixes_copy.size(), prefixes.size());
 209   EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
 210                          prefixes_copy.begin()));
 211 }
 212
 213 // A range with only large deltas.
 214 TEST_F(PrefixSetTest, AllBig) {
 215   std::vector<SBPrefix> prefixes;
 216
 217   const SBPrefix kVeryPositive = 1000 * 1000 * 1000;
 218   const SBPrefix kVeryNegative = -kVeryPositive;
 219   const unsigned kDelta = 10 * 1000 * 1000;
 220
 221   for (SBPrefix prefix = kVeryNegative;
 222        prefix < kVeryPositive; prefix += kDelta) {
 223     prefixes.push_back(prefix);
 224   }
 225
 226   std::sort(prefixes.begin(), prefixes.end());
 227   safe_browsing::PrefixSet prefix_set(prefixes);
 228
 229   // Check that |GetPrefixes()| returns the same set of prefixes as
 230   // was passed in.
 231   std::vector<SBPrefix> prefixes_copy;
 232   prefix_set.GetPrefixes(&prefixes_copy);
 233   prefixes.erase(std::unique(prefixes.begin(), prefixes.end()), prefixes.end());
 234   EXPECT_EQ(prefixes_copy.size(), prefixes.size());
 235   EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
 236                          prefixes_copy.begin()));
 237 }
 238
 239 // Use artificial inputs to test various edge cases in Exists().
 240 // Items before the lowest item aren't present.  Items after the
 241 // largest item aren't present.  Create a sequence of items with
 242 // deltas above and below 2^16, and make sure they're all present.
 243 // Create a very long sequence with deltas below 2^16 to test crossing
 244 // |kMaxRun|.
 245 TEST_F(PrefixSetTest, EdgeCases) {
 246   std::vector<SBPrefix> prefixes;
 247
 248   const SBPrefix kVeryPositive = 1000 * 1000 * 1000;
 249   const SBPrefix kVeryNegative = -kVeryPositive;
 250
 251   // Put in a very negative prefix.
 252   SBPrefix prefix = kVeryNegative;
 253   prefixes.push_back(prefix);
 254
 255   // Add a sequence with very large deltas.
 256   unsigned delta = 100 * 1000 * 1000;
 257   for (int i = 0; i < 10; ++i) {
 258     prefix += delta;
 259     prefixes.push_back(prefix);
 260   }
 261
 262   // Add a sequence with deltas that start out smaller than the
 263   // maximum delta, and end up larger.  Also include some duplicates.
 264   delta = 256 * 256 - 100;
 265   for (int i = 0; i < 200; ++i) {
 266     prefix += delta;
 267     prefixes.push_back(prefix);
 268     prefixes.push_back(prefix);
 269     delta++;
 270   }
 271
 272   // Add a long sequence with deltas smaller than the maximum delta,
 273   // so a new index item will be injected.
 274   delta = 256 * 256 - 1;
 275   prefix = kVeryPositive - delta * 1000;
 276   prefixes.push_back(prefix);
 277   for (int i = 0; i < 1000; ++i) {
 278     prefix += delta;
 279     prefixes.push_back(prefix);
 280     delta--;
 281   }
 282
 283   std::sort(prefixes.begin(), prefixes.end());
 284   safe_browsing::PrefixSet prefix_set(prefixes);
 285
 286   // Check that |GetPrefixes()| returns the same set of prefixes as
 287   // was passed in.
 288   std::vector<SBPrefix> prefixes_copy;
 289   prefix_set.GetPrefixes(&prefixes_copy);
 290   prefixes.erase(std::unique(prefixes.begin(), prefixes.end()), prefixes.end());
 291   EXPECT_EQ(prefixes_copy.size(), prefixes.size());
 292   EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
 293                          prefixes_copy.begin()));
 294
 295   // Items before and after the set are not present, and don't crash.
 296   EXPECT_FALSE(prefix_set.Exists(kVeryNegative - 100));
 297   EXPECT_FALSE(prefix_set.Exists(kVeryPositive + 100));
 298
 299   // Check that the set correctly flags all of the inputs, and also
 300   // check items just above and below the inputs to make sure they
 301   // aren't present.
 302   for (size_t i = 0; i < prefixes.size(); ++i) {
 303     EXPECT_TRUE(prefix_set.Exists(prefixes[i]));
 304
 305     EXPECT_FALSE(prefix_set.Exists(prefixes[i] - 1));
 306     EXPECT_FALSE(prefix_set.Exists(prefixes[i] + 1));
 307   }
 308 }
 309
 310 // Test writing a prefix set to disk and reading it back in.
 311 TEST_F(PrefixSetTest, ReadWrite) {
 312   base::FilePath filename;
 313
 314   // Write the sample prefix set out, read it back in, and check all
 315   // the prefixes.  Leaves the path in |filename|.
 316   {
 317     ASSERT_TRUE(GetPrefixSetFile(&filename));
 318     scoped_ptr<safe_browsing::PrefixSet>
 319         prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
 320     ASSERT_TRUE(prefix_set.get());
 321     CheckPrefixes(*prefix_set, shared_prefixes_);
 322   }
 323
 324   // Test writing and reading a very sparse set containing no deltas.
 325   {
 326     const SBPrefix kVeryPositive = 1000 * 1000 * 1000;
 327     const SBPrefix kVeryNegative = -kVeryPositive;
 328
 329     std::vector<SBPrefix> prefixes;
 330     prefixes.push_back(kVeryNegative);
 331     prefixes.push_back(kVeryPositive);
 332
 333     safe_browsing::PrefixSet prefix_set_to_write(prefixes);
 334     ASSERT_TRUE(prefix_set_to_write.WriteFile(filename));
 335     scoped_ptr<safe_browsing::PrefixSet>
 336         prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
 337     ASSERT_TRUE(prefix_set.get());
 338     CheckPrefixes(*prefix_set, prefixes);
 339   }
 340
 341   // Test writing and reading an empty set.
 342   {
 343     std::vector<SBPrefix> prefixes;
 344     safe_browsing::PrefixSet prefix_set_to_write(prefixes);
 345     ASSERT_TRUE(prefix_set_to_write.WriteFile(filename));
 346     scoped_ptr<safe_browsing::PrefixSet>
 347         prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
 348     ASSERT_TRUE(prefix_set.get());
 349     CheckPrefixes(*prefix_set, prefixes);
 350   }
 351 }
 352
 353 // Check that |CleanChecksum()| makes an acceptable checksum.
 354 TEST_F(PrefixSetTest, CorruptionHelpers) {
 355   base::FilePath filename;
 356   ASSERT_TRUE(GetPrefixSetFile(&filename));
 357
 358   // This will modify data in |index_|, which will fail the digest check.
 359   file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b"));
 360   IncrementIntAt(file.get(), kPayloadOffset, 1);
 361   file.reset();
 362   scoped_ptr<safe_browsing::PrefixSet>
 363       prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
 364   ASSERT_FALSE(prefix_set.get());
 365
 366   // Fix up the checksum and it will read successfully (though the
 367   // data will be wrong).
 368   file.reset(file_util::OpenFile(filename, "r+b"));
 369   CleanChecksum(file.get());
 370   file.reset();
 371   prefix_set.reset(safe_browsing::PrefixSet::LoadFile(filename));
 372   ASSERT_TRUE(prefix_set.get());
 373 }
 374
 375 // Bad |index_| size is caught by the sanity check.
 376 TEST_F(PrefixSetTest, CorruptionMagic) {
 377   base::FilePath filename;
 378   ASSERT_TRUE(GetPrefixSetFile(&filename));
 379
 380   ASSERT_NO_FATAL_FAILURE(
 381       ModifyAndCleanChecksum(filename, kMagicOffset, 1));
 382   scoped_ptr<safe_browsing::PrefixSet>
 383       prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
 384   ASSERT_FALSE(prefix_set.get());
 385 }
 386
 387 // Bad |index_| size is caught by the sanity check.
 388 TEST_F(PrefixSetTest, CorruptionVersion) {
 389   base::FilePath filename;
 390   ASSERT_TRUE(GetPrefixSetFile(&filename));
 391
 392   ASSERT_NO_FATAL_FAILURE(
 393       ModifyAndCleanChecksum(filename, kVersionOffset, 1));
 394   scoped_ptr<safe_browsing::PrefixSet>
 395       prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
 396   ASSERT_FALSE(prefix_set.get());
 397 }
 398
 399 // Bad |index_| size is caught by the sanity check.
 400 TEST_F(PrefixSetTest, CorruptionIndexSize) {
 401   base::FilePath filename;
 402   ASSERT_TRUE(GetPrefixSetFile(&filename));
 403
 404   ASSERT_NO_FATAL_FAILURE(
 405       ModifyAndCleanChecksum(filename, kIndexSizeOffset, 1));
 406   scoped_ptr<safe_browsing::PrefixSet>
 407       prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
 408   ASSERT_FALSE(prefix_set.get());
 409 }
 410
 411 // Bad |deltas_| size is caught by the sanity check.
 412 TEST_F(PrefixSetTest, CorruptionDeltasSize) {
 413   base::FilePath filename;
 414   ASSERT_TRUE(GetPrefixSetFile(&filename));
 415
 416   ASSERT_NO_FATAL_FAILURE(
 417       ModifyAndCleanChecksum(filename, kDeltasSizeOffset, 1));
 418   scoped_ptr<safe_browsing::PrefixSet>
 419       prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
 420   ASSERT_FALSE(prefix_set.get());
 421 }
 422
 423 // Test that the digest catches corruption in the middle of the file
 424 // (in the payload between the header and the digest).
 425 TEST_F(PrefixSetTest, CorruptionPayload) {
 426   base::FilePath filename;
 427   ASSERT_TRUE(GetPrefixSetFile(&filename));
 428
 429   file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b"));
 430   ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), 666, 1));
 431   file.reset();
 432   scoped_ptr<safe_browsing::PrefixSet>
 433       prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
 434   ASSERT_FALSE(prefix_set.get());
 435 }
 436
 437 // Test corruption in the digest itself.
 438 TEST_F(PrefixSetTest, CorruptionDigest) {
 439   base::FilePath filename;
 440   ASSERT_TRUE(GetPrefixSetFile(&filename));
 441
 442   int64 size_64;
 443   ASSERT_TRUE(file_util::GetFileSize(filename, &size_64));
 444   file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b"));
 445   long digest_offset = static_cast<long>(size_64 - sizeof(base::MD5Digest));
 446   ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), digest_offset, 1));
 447   file.reset();
 448   scoped_ptr<safe_browsing::PrefixSet>
 449       prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
 450   ASSERT_FALSE(prefix_set.get());
 451 }
 452
 453 // Test excess data after the digest (fails the size test).
 454 TEST_F(PrefixSetTest, CorruptionExcess) {
 455   base::FilePath filename;
 456   ASSERT_TRUE(GetPrefixSetFile(&filename));
 457
 458   // Add some junk to the trunk.
 459   file_util::ScopedFILE file(file_util::OpenFile(filename, "ab"));
 460   const char buf[] = "im in ur base, killing ur d00dz.";
 461   ASSERT_EQ(strlen(buf), fwrite(buf, 1, strlen(buf), file.get()));
 462   file.reset();
 463   scoped_ptr<safe_browsing::PrefixSet>
 464       prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
 465   ASSERT_FALSE(prefix_set.get());
 466 }
 467
 468 }  // namespace