OLD | NEW |
1 // Copyright (c) 2011 The Chromium Authors. All rights reserved. | 1 // Copyright (c) 2011 The Chromium Authors. All rights reserved. |
2 // Use of this source code is governed by a BSD-style license that can be | 2 // Use of this source code is governed by a BSD-style license that can be |
3 // found in the LICENSE file. | 3 // found in the LICENSE file. |
4 | 4 |
5 #include "chrome/browser/safe_browsing/prefix_set.h" | 5 #include "chrome/browser/safe_browsing/prefix_set.h" |
6 | 6 |
7 #include <algorithm> | 7 #include <algorithm> |
8 | 8 |
| 9 #include "base/file_util.h" |
9 #include "base/logging.h" | 10 #include "base/logging.h" |
| 11 #include "base/md5.h" |
10 #include "base/rand_util.h" | 12 #include "base/rand_util.h" |
| 13 #include "base/scoped_ptr.h" |
| 14 #include "base/scoped_temp_dir.h" |
11 #include "testing/gtest/include/gtest/gtest.h" | 15 #include "testing/gtest/include/gtest/gtest.h" |
| 16 #include "testing/platform_test.h" |
12 | 17 |
13 namespace { | 18 namespace { |
14 | 19 |
15 SBPrefix GenPrefix() { | 20 class PrefixSetTest : public PlatformTest { |
16 return static_cast<SBPrefix>(base::RandUint64()); | 21 protected: |
| 22 // Constants for the v1 format. |
| 23 static const size_t kMagicOffset = 0 * sizeof(uint32); |
| 24 static const size_t kVersionOffset = 1 * sizeof(uint32); |
| 25 static const size_t kIndexSizeOffset = 2 * sizeof(uint32); |
| 26 static const size_t kDeltasSizeOffset = 3 * sizeof(uint32); |
| 27 static const size_t kPayloadOffset = 4 * sizeof(uint32); |
| 28 |
| 29 // Generate a set of random prefixes to share between tests. For |
| 30 // most tests this generation was a large fraction of the test time. |
| 31 static void SetUpTestCase() { |
| 32 for (size_t i = 0; i < 50000; ++i) { |
| 33 const SBPrefix prefix = static_cast<SBPrefix>(base::RandUint64()); |
| 34 shared_prefixes_.push_back(prefix); |
| 35 } |
| 36 |
| 37 // Sort for use with PrefixSet constructor. |
| 38 std::sort(shared_prefixes_.begin(), shared_prefixes_.end()); |
| 39 } |
| 40 |
| 41 // Check that all elements of |prefixes| are in |prefix_set|, and |
| 42 // that nearby elements are not (for lack of a more sensible set of |
| 43 // items to check for absence). |
| 44 static void CheckPrefixes(safe_browsing::PrefixSet* prefix_set, |
| 45 const std::vector<SBPrefix> &prefixes) { |
| 46 // The set can generate the prefixes it believes it has, so that's |
| 47 // a good starting point. |
| 48 std::set<SBPrefix> check(prefixes.begin(), prefixes.end()); |
| 49 std::vector<SBPrefix> prefixes_copy; |
| 50 prefix_set->GetPrefixes(&prefixes_copy); |
| 51 EXPECT_EQ(prefixes_copy.size(), check.size()); |
| 52 EXPECT_TRUE(std::equal(check.begin(), check.end(), prefixes_copy.begin())); |
| 53 |
| 54 for (size_t i = 0; i < prefixes.size(); ++i) { |
| 55 EXPECT_TRUE(prefix_set->Exists(prefixes[i])); |
| 56 |
| 57 const SBPrefix left_sibling = prefixes[i] - 1; |
| 58 if (check.count(left_sibling) == 0) |
| 59 EXPECT_FALSE(prefix_set->Exists(left_sibling)); |
| 60 |
| 61 const SBPrefix right_sibling = prefixes[i] + 1; |
| 62 if (check.count(right_sibling) == 0) |
| 63 EXPECT_FALSE(prefix_set->Exists(right_sibling)); |
| 64 } |
| 65 } |
| 66 |
| 67 // Generate a |PrefixSet| file from |shared_prefixes_|, store it in |
| 68 // a temporary file, and return the filename in |filenamep|. |
| 69 // Returns |true| on success. |
| 70 bool GetPrefixSetFile(FilePath* filenamep) { |
| 71 if (!temp_dir_.IsValid() && !temp_dir_.CreateUniqueTempDir()) |
| 72 return false; |
| 73 |
| 74 FilePath filename = temp_dir_.path().AppendASCII("PrefixSetTest"); |
| 75 |
| 76 safe_browsing::PrefixSet prefix_set(shared_prefixes_); |
| 77 if (!prefix_set.WriteFile(filename)) |
| 78 return false; |
| 79 |
| 80 *filenamep = filename; |
| 81 return true; |
| 82 } |
| 83 |
| 84 // Helper function to read the int32 value at |offset|, increment it |
| 85 // by |inc|, and write it back in place. |fp| should be opened in |
| 86 // r+ mode. |
| 87 static void IncrementIntAt(FILE* fp, long offset, int inc) { |
| 88 int32 value = 0; |
| 89 |
| 90 ASSERT_NE(-1, fseek(fp, offset, SEEK_SET)); |
| 91 ASSERT_EQ(1U, fread(&value, sizeof(value), 1, fp)); |
| 92 |
| 93 value += inc; |
| 94 |
| 95 ASSERT_NE(-1, fseek(fp, offset, SEEK_SET)); |
| 96 ASSERT_EQ(1U, fwrite(&value, sizeof(value), 1, fp)); |
| 97 } |
| 98 |
| 99 // Helper function to re-generated |fp|'s checksum to be correct for |
| 100 // the file's contents. |fp| should be opened in r+ mode. |
| 101 static void CleanChecksum(FILE* fp) { |
| 102 MD5Context context; |
| 103 MD5Init(&context); |
| 104 |
| 105 ASSERT_NE(-1, fseek(fp, 0, SEEK_END)); |
| 106 long file_size = ftell(fp); |
| 107 |
| 108 size_t payload_size = static_cast<size_t>(file_size) - sizeof(MD5Digest); |
| 109 size_t digested_size = 0; |
| 110 ASSERT_NE(-1, fseek(fp, 0, SEEK_SET)); |
| 111 while (digested_size < payload_size) { |
| 112 char buf[1024]; |
| 113 size_t nitems = std::min(payload_size - digested_size, sizeof(buf)); |
| 114 ASSERT_EQ(nitems, fread(buf, 1, nitems, fp)); |
| 115 MD5Update(&context, &buf, nitems); |
| 116 digested_size += nitems; |
| 117 } |
| 118 ASSERT_EQ(digested_size, payload_size); |
| 119 ASSERT_EQ(static_cast<long>(digested_size), ftell(fp)); |
| 120 |
| 121 MD5Digest new_digest; |
| 122 MD5Final(&new_digest, &context); |
| 123 ASSERT_NE(-1, fseek(fp, digested_size, SEEK_SET)); |
| 124 ASSERT_EQ(1U, fwrite(&new_digest, sizeof(new_digest), 1, fp)); |
| 125 ASSERT_EQ(file_size, ftell(fp)); |
| 126 } |
| 127 |
| 128 // Open |filename| and increment the int32 at |offset| by |inc|. |
| 129 // Then re-generate the checksum to account for the new contents. |
| 130 void ModifyAndCleanChecksum(const FilePath& filename, long offset, int inc) { |
| 131 int64 size_64; |
| 132 ASSERT_TRUE(file_util::GetFileSize(filename, &size_64)); |
| 133 |
| 134 file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b")); |
| 135 IncrementIntAt(file.get(), offset, inc); |
| 136 CleanChecksum(file.get()); |
| 137 file.reset(); |
| 138 |
| 139 int64 new_size_64; |
| 140 ASSERT_TRUE(file_util::GetFileSize(filename, &new_size_64)); |
| 141 ASSERT_EQ(new_size_64, size_64); |
| 142 } |
| 143 |
| 144 // Tests should not modify this shared resource. |
| 145 static std::vector<SBPrefix> shared_prefixes_; |
| 146 |
| 147 ScopedTempDir temp_dir_; |
| 148 }; |
| 149 |
| 150 std::vector<SBPrefix> PrefixSetTest::shared_prefixes_; |
| 151 |
| 152 // Test that a small sparse random input works. |
| 153 TEST_F(PrefixSetTest, Baseline) { |
| 154 safe_browsing::PrefixSet prefix_set(shared_prefixes_); |
| 155 CheckPrefixes(&prefix_set, shared_prefixes_); |
17 } | 156 } |
18 | 157 |
19 // Test that a small sparse random input works. | 158 // Test that the empty set doesn't appear to have anything in it. |
20 TEST(PrefixSetTest, Baseline) { | 159 TEST_F(PrefixSetTest, Empty) { |
21 std::vector<SBPrefix> prefixes; | 160 const std::vector<SBPrefix> empty; |
22 | 161 safe_browsing::PrefixSet prefix_set(empty); |
23 static const size_t kCount = 50000; | 162 for (size_t i = 0; i < shared_prefixes_.size(); ++i) { |
24 | 163 EXPECT_FALSE(prefix_set.Exists(shared_prefixes_[i])); |
25 for (size_t i = 0; i < kCount; ++i) { | |
26 prefixes.push_back(GenPrefix()); | |
27 } | |
28 | |
29 std::sort(prefixes.begin(), prefixes.end()); | |
30 safe_browsing::PrefixSet prefix_set(prefixes); | |
31 | |
32 // Check that |GetPrefixes()| returns exactly the same set of | |
33 // prefixes as was passed in. | |
34 std::set<SBPrefix> check(prefixes.begin(), prefixes.end()); | |
35 std::vector<SBPrefix> prefixes_copy; | |
36 prefix_set.GetPrefixes(&prefixes_copy); | |
37 EXPECT_EQ(prefixes_copy.size(), check.size()); | |
38 EXPECT_TRUE(std::equal(check.begin(), check.end(), prefixes_copy.begin())); | |
39 | |
40 // Check that the set flags all of the inputs, and also check items | |
41 // just above and below the inputs to make sure they aren't there. | |
42 for (size_t i = 0; i < prefixes.size(); ++i) { | |
43 EXPECT_TRUE(prefix_set.Exists(prefixes[i])); | |
44 | |
45 const SBPrefix left_sibling = prefixes[i] - 1; | |
46 if (check.count(left_sibling) == 0) | |
47 EXPECT_FALSE(prefix_set.Exists(left_sibling)); | |
48 | |
49 const SBPrefix right_sibling = prefixes[i] + 1; | |
50 if (check.count(right_sibling) == 0) | |
51 EXPECT_FALSE(prefix_set.Exists(right_sibling)); | |
52 } | 164 } |
53 } | 165 } |
54 | 166 |
55 // Test that the empty set doesn't appear to have anything in it. | |
56 TEST(PrefixSetTest, Empty) { | |
57 std::vector<SBPrefix> prefixes; | |
58 safe_browsing::PrefixSet prefix_set(prefixes); | |
59 for (size_t i = 0; i < 500; ++i) { | |
60 EXPECT_FALSE(prefix_set.Exists(GenPrefix())); | |
61 } | |
62 } | |
63 | |
64 // Use artificial inputs to test various edge cases in Exists(). | 167 // Use artificial inputs to test various edge cases in Exists(). |
65 // Items before the lowest item aren't present. Items after the | 168 // Items before the lowest item aren't present. Items after the |
66 // largest item aren't present. Create a sequence of items with | 169 // largest item aren't present. Create a sequence of items with |
67 // deltas above and below 2^16, and make sure they're all present. | 170 // deltas above and below 2^16, and make sure they're all present. |
68 // Create a very long sequence with deltas below 2^16 to test crossing | 171 // Create a very long sequence with deltas below 2^16 to test crossing |
69 // |kMaxRun|. | 172 // |kMaxRun|. |
70 TEST(PrefixSetTest, EdgeCases) { | 173 TEST_F(PrefixSetTest, EdgeCases) { |
71 std::vector<SBPrefix> prefixes; | 174 std::vector<SBPrefix> prefixes; |
72 | 175 |
73 const SBPrefix kVeryPositive = 1000 * 1000 * 1000; | 176 const SBPrefix kVeryPositive = 1000 * 1000 * 1000; |
74 const SBPrefix kVeryNegative = -kVeryPositive; | 177 const SBPrefix kVeryNegative = -kVeryPositive; |
75 | 178 |
76 // Put in a very negative prefix. | 179 // Put in a very negative prefix. |
77 SBPrefix prefix = kVeryNegative; | 180 SBPrefix prefix = kVeryNegative; |
78 prefixes.push_back(prefix); | 181 prefixes.push_back(prefix); |
79 | 182 |
80 // Add a sequence with very large deltas. | 183 // Add a sequence with very large deltas. |
(...skipping 44 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... |
125 // check items just above and below the inputs to make sure they | 228 // check items just above and below the inputs to make sure they |
126 // aren't present. | 229 // aren't present. |
127 for (size_t i = 0; i < prefixes.size(); ++i) { | 230 for (size_t i = 0; i < prefixes.size(); ++i) { |
128 EXPECT_TRUE(prefix_set.Exists(prefixes[i])); | 231 EXPECT_TRUE(prefix_set.Exists(prefixes[i])); |
129 | 232 |
130 EXPECT_FALSE(prefix_set.Exists(prefixes[i] - 1)); | 233 EXPECT_FALSE(prefix_set.Exists(prefixes[i] - 1)); |
131 EXPECT_FALSE(prefix_set.Exists(prefixes[i] + 1)); | 234 EXPECT_FALSE(prefix_set.Exists(prefixes[i] + 1)); |
132 } | 235 } |
133 } | 236 } |
134 | 237 |
| 238 // Similar to Baseline test, but write the set out to a file and read |
| 239 // it back in before testing. |
| 240 TEST_F(PrefixSetTest, ReadWrite) { |
| 241 FilePath filename; |
| 242 ASSERT_TRUE(GetPrefixSetFile(&filename)); |
| 243 |
| 244 scoped_ptr<safe_browsing::PrefixSet> |
| 245 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
| 246 ASSERT_TRUE(prefix_set.get()); |
| 247 |
| 248 CheckPrefixes(prefix_set.get(), shared_prefixes_); |
| 249 } |
| 250 |
| 251 // Check that |CleanChecksum()| makes an acceptable checksum. |
| 252 TEST_F(PrefixSetTest, CorruptionHelpers) { |
| 253 FilePath filename; |
| 254 ASSERT_TRUE(GetPrefixSetFile(&filename)); |
| 255 |
| 256 // This will modify data in |index_|, which will fail the digest check. |
| 257 file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b")); |
| 258 IncrementIntAt(file.get(), kPayloadOffset, 1); |
| 259 file.reset(); |
| 260 scoped_ptr<safe_browsing::PrefixSet> |
| 261 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
| 262 ASSERT_FALSE(prefix_set.get()); |
| 263 |
| 264 // Fix up the checksum and it will read successfully (though the |
| 265 // data will be wrong). |
| 266 file.reset(file_util::OpenFile(filename, "r+b")); |
| 267 CleanChecksum(file.get()); |
| 268 file.reset(); |
| 269 prefix_set.reset(safe_browsing::PrefixSet::LoadFile(filename)); |
| 270 ASSERT_TRUE(prefix_set.get()); |
| 271 } |
| 272 |
| 273 // Bad |index_| size is caught by the sanity check. |
| 274 TEST_F(PrefixSetTest, CorruptionMagic) { |
| 275 FilePath filename; |
| 276 ASSERT_TRUE(GetPrefixSetFile(&filename)); |
| 277 |
| 278 ASSERT_NO_FATAL_FAILURE( |
| 279 ModifyAndCleanChecksum(filename, kMagicOffset, 1)); |
| 280 scoped_ptr<safe_browsing::PrefixSet> |
| 281 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
| 282 ASSERT_FALSE(prefix_set.get()); |
| 283 } |
| 284 |
| 285 // Bad |index_| size is caught by the sanity check. |
| 286 TEST_F(PrefixSetTest, CorruptionVersion) { |
| 287 FilePath filename; |
| 288 ASSERT_TRUE(GetPrefixSetFile(&filename)); |
| 289 |
| 290 ASSERT_NO_FATAL_FAILURE( |
| 291 ModifyAndCleanChecksum(filename, kVersionOffset, 1)); |
| 292 scoped_ptr<safe_browsing::PrefixSet> |
| 293 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
| 294 ASSERT_FALSE(prefix_set.get()); |
| 295 } |
| 296 |
| 297 // Bad |index_| size is caught by the sanity check. |
| 298 TEST_F(PrefixSetTest, CorruptionIndexSize) { |
| 299 FilePath filename; |
| 300 ASSERT_TRUE(GetPrefixSetFile(&filename)); |
| 301 |
| 302 ASSERT_NO_FATAL_FAILURE( |
| 303 ModifyAndCleanChecksum(filename, kIndexSizeOffset, 1)); |
| 304 scoped_ptr<safe_browsing::PrefixSet> |
| 305 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
| 306 ASSERT_FALSE(prefix_set.get()); |
| 307 } |
| 308 |
| 309 // Bad |deltas_| size is caught by the sanity check. |
| 310 TEST_F(PrefixSetTest, CorruptionDeltasSize) { |
| 311 FilePath filename; |
| 312 ASSERT_TRUE(GetPrefixSetFile(&filename)); |
| 313 |
| 314 ASSERT_NO_FATAL_FAILURE( |
| 315 ModifyAndCleanChecksum(filename, kDeltasSizeOffset, 1)); |
| 316 scoped_ptr<safe_browsing::PrefixSet> |
| 317 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
| 318 ASSERT_FALSE(prefix_set.get()); |
| 319 } |
| 320 |
| 321 // Test that the digest catches corruption in the middle of the file |
| 322 // (in the payload between the header and the digest). |
| 323 TEST_F(PrefixSetTest, CorruptionPayload) { |
| 324 FilePath filename; |
| 325 ASSERT_TRUE(GetPrefixSetFile(&filename)); |
| 326 |
| 327 file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b")); |
| 328 ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), 666, 1)); |
| 329 file.reset(); |
| 330 scoped_ptr<safe_browsing::PrefixSet> |
| 331 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
| 332 ASSERT_FALSE(prefix_set.get()); |
| 333 } |
| 334 |
| 335 // Test corruption in the digest itself. |
| 336 TEST_F(PrefixSetTest, CorruptionDigest) { |
| 337 FilePath filename; |
| 338 ASSERT_TRUE(GetPrefixSetFile(&filename)); |
| 339 |
| 340 int64 size_64; |
| 341 ASSERT_TRUE(file_util::GetFileSize(filename, &size_64)); |
| 342 file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b")); |
| 343 long digest_offset = static_cast<long>(size_64 - sizeof(MD5Digest)); |
| 344 ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), digest_offset, 1)); |
| 345 file.reset(); |
| 346 scoped_ptr<safe_browsing::PrefixSet> |
| 347 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
| 348 ASSERT_FALSE(prefix_set.get()); |
| 349 } |
| 350 |
| 351 // Test excess data after the digest (fails the size test). |
| 352 TEST_F(PrefixSetTest, CorruptionExcess) { |
| 353 FilePath filename; |
| 354 ASSERT_TRUE(GetPrefixSetFile(&filename)); |
| 355 |
| 356 // Add some junk to the trunk. |
| 357 file_util::ScopedFILE file(file_util::OpenFile(filename, "ab")); |
| 358 const char buf[] = "im in ur base, killing ur d00dz."; |
| 359 ASSERT_EQ(strlen(buf), fwrite(buf, 1, strlen(buf), file.get())); |
| 360 file.reset(); |
| 361 scoped_ptr<safe_browsing::PrefixSet> |
| 362 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
| 363 ASSERT_FALSE(prefix_set.get()); |
| 364 } |
| 365 |
135 } // namespace | 366 } // namespace |
OLD | NEW |