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