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 469
