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 pf |prefixes| are in |prefix_set|, and | |
lzheng
2011/03/08 01:15:51
pf -> of?
Scott Hess - ex-Googler
2011/03/08 01:26:41
'pf' is a well-known alternative spelling, but I c
| |
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 } | |
lzheng
2011/03/08 01:15:51
empty line needed here.
Scott Hess - ex-Googler
2011/03/08 01:26:41
Done.
| |
83 // Helper function to read the int32 value at |offset|, increment it | |
84 // by |inc|, and write it back in place. |fp| should be opened in | |
85 // r+ mode. | |
86 static void IncrementIntAt(FILE* fp, long offset, int inc) { | |
87 int32 value = 0; | |
88 | |
89 ASSERT_NE(-1, fseek(fp, offset, SEEK_SET)); | |
90 ASSERT_EQ(1U, fread(&value, sizeof(value), 1, fp)); | |
91 | |
92 value += inc; | |
93 | |
94 ASSERT_NE(-1, fseek(fp, offset, SEEK_SET)); | |
95 ASSERT_EQ(1U, fwrite(&value, sizeof(value), 1, fp)); | |
96 } | |
97 | |
98 // Helper function to re-generated |fp|'s checksum to be correct for | |
99 // the file's contents. |fp| should be opened in r+ mode. | |
100 static void CleanChecksum(FILE* fp) { | |
101 MD5Context context; | |
102 MD5Init(&context); | |
103 | |
104 ASSERT_NE(-1, fseek(fp, 0, SEEK_END)); | |
105 long file_size = ftell(fp); | |
106 | |
107 size_t payload_size = static_cast<size_t>(file_size) - sizeof(MD5Digest); | |
108 size_t digested_size = 0; | |
109 ASSERT_NE(-1, fseek(fp, 0, SEEK_SET)); | |
110 while (digested_size < payload_size) { | |
111 char buf[1024]; | |
112 size_t nitems = std::min(payload_size - digested_size, sizeof(buf)); | |
113 ASSERT_EQ(nitems, fread(buf, 1, nitems, fp)); | |
114 MD5Update(&context, &buf, nitems); | |
115 digested_size += nitems; | |
116 } | |
117 ASSERT_EQ(digested_size, payload_size); | |
118 ASSERT_EQ(static_cast<long>(digested_size), ftell(fp)); | |
119 | |
120 MD5Digest new_digest; | |
121 MD5Final(&new_digest, &context); | |
122 ASSERT_NE(-1, fseek(fp, digested_size, SEEK_SET)); | |
123 ASSERT_EQ(1U, fwrite(&new_digest, sizeof(new_digest), 1, fp)); | |
124 ASSERT_EQ(file_size, ftell(fp)); | |
125 } | |
126 | |
127 // Open |filename| and increment the int32 at |offset| by |inc|. | |
128 // Then re-generate the checksum to account for the new contents. | |
129 void ModifyAndCleanChecksum(const FilePath& filename, long offset, int inc) { | |
130 int64 size_64; | |
131 ASSERT_TRUE(file_util::GetFileSize(filename, &size_64)); | |
132 | |
133 file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b")); | |
134 IncrementIntAt(file.get(), offset, inc); | |
135 CleanChecksum(file.get()); | |
136 file.reset(); | |
137 | |
138 int64 new_size_64; | |
139 ASSERT_TRUE(file_util::GetFileSize(filename, &new_size_64)); | |
140 ASSERT_EQ(new_size_64, size_64); | |
141 } | |
142 | |
143 // Tests should not modify this shared resource. | |
144 static std::vector<SBPrefix> shared_prefixes_; | |
145 | |
146 ScopedTempDir temp_dir_; | |
147 }; | |
148 | |
149 std::vector<SBPrefix> PrefixSetTest::shared_prefixes_; | |
150 | |
151 // Test that a small sparse random input works. | |
152 TEST_F(PrefixSetTest, Baseline) { | |
153 safe_browsing::PrefixSet prefix_set(shared_prefixes_); | |
154 CheckPrefixes(&prefix_set, shared_prefixes_); | |
17 } | 155 } |
18 | 156 |
19 // Test that a small sparse random input works. | 157 // Test that the empty set doesn't appear to have anything in it. |
20 TEST(PrefixSetTest, Baseline) { | 158 TEST_F(PrefixSetTest, Empty) { |
21 std::vector<SBPrefix> prefixes; | 159 const std::vector<SBPrefix> empty; |
22 | 160 safe_browsing::PrefixSet prefix_set(empty); |
23 static const size_t kCount = 50000; | 161 for (size_t i = 0; i < shared_prefixes_.size(); ++i) { |
24 | 162 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 } | 163 } |
53 } | 164 } |
54 | 165 |
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(). | 166 // Use artificial inputs to test various edge cases in Exists(). |
65 // Items before the lowest item aren't present. Items after the | 167 // Items before the lowest item aren't present. Items after the |
66 // largest item aren't present. Create a sequence of items with | 168 // 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. | 169 // 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 | 170 // Create a very long sequence with deltas below 2^16 to test crossing |
69 // |kMaxRun|. | 171 // |kMaxRun|. |
70 TEST(PrefixSetTest, EdgeCases) { | 172 TEST_F(PrefixSetTest, EdgeCases) { |
71 std::vector<SBPrefix> prefixes; | 173 std::vector<SBPrefix> prefixes; |
72 | 174 |
73 const SBPrefix kVeryPositive = 1000 * 1000 * 1000; | 175 const SBPrefix kVeryPositive = 1000 * 1000 * 1000; |
74 const SBPrefix kVeryNegative = -kVeryPositive; | 176 const SBPrefix kVeryNegative = -kVeryPositive; |
75 | 177 |
76 // Put in a very negative prefix. | 178 // Put in a very negative prefix. |
77 SBPrefix prefix = kVeryNegative; | 179 SBPrefix prefix = kVeryNegative; |
78 prefixes.push_back(prefix); | 180 prefixes.push_back(prefix); |
79 | 181 |
80 // Add a sequence with very large deltas. | 182 // 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 | 227 // check items just above and below the inputs to make sure they |
126 // aren't present. | 228 // aren't present. |
127 for (size_t i = 0; i < prefixes.size(); ++i) { | 229 for (size_t i = 0; i < prefixes.size(); ++i) { |
128 EXPECT_TRUE(prefix_set.Exists(prefixes[i])); | 230 EXPECT_TRUE(prefix_set.Exists(prefixes[i])); |
129 | 231 |
130 EXPECT_FALSE(prefix_set.Exists(prefixes[i] - 1)); | 232 EXPECT_FALSE(prefix_set.Exists(prefixes[i] - 1)); |
131 EXPECT_FALSE(prefix_set.Exists(prefixes[i] + 1)); | 233 EXPECT_FALSE(prefix_set.Exists(prefixes[i] + 1)); |
132 } | 234 } |
133 } | 235 } |
134 | 236 |
237 // Similar to Baseline test, but write the set out to a file and read | |
238 // it back in before testing. | |
239 TEST_F(PrefixSetTest, ReadWrite) { | |
240 FilePath filename; | |
241 ASSERT_TRUE(GetPrefixSetFile(&filename)); | |
242 | |
243 scoped_ptr<safe_browsing::PrefixSet> | |
244 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); | |
245 ASSERT_TRUE(prefix_set.get()); | |
246 | |
247 CheckPrefixes(prefix_set.get(), shared_prefixes_); | |
248 } | |
249 | |
250 // Check that |CleanChecksum()| makes an acceptable checksum. | |
251 TEST_F(PrefixSetTest, CorruptionHelpers) { | |
252 FilePath filename; | |
253 ASSERT_TRUE(GetPrefixSetFile(&filename)); | |
254 | |
255 // This will modify data in |index_|, which will fail the digest check. | |
256 file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b")); | |
257 IncrementIntAt(file.get(), kPayloadOffset, 1); | |
258 file.reset(); | |
259 scoped_ptr<safe_browsing::PrefixSet> | |
260 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); | |
261 ASSERT_FALSE(prefix_set.get()); | |
262 | |
263 // Fix up the checksum and it will read successfully (though the | |
264 // data will be wrong). | |
265 file.reset(file_util::OpenFile(filename, "r+b")); | |
266 CleanChecksum(file.get()); | |
267 file.reset(); | |
268 prefix_set.reset(safe_browsing::PrefixSet::LoadFile(filename)); | |
269 ASSERT_TRUE(prefix_set.get()); | |
270 } | |
271 | |
272 // Bad |index_| size is caught by the sanity check. | |
273 TEST_F(PrefixSetTest, CorruptionMagic) { | |
274 FilePath filename; | |
275 ASSERT_TRUE(GetPrefixSetFile(&filename)); | |
276 | |
277 ASSERT_NO_FATAL_FAILURE( | |
278 ModifyAndCleanChecksum(filename, kMagicOffset, 1)); | |
279 scoped_ptr<safe_browsing::PrefixSet> | |
280 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); | |
281 ASSERT_FALSE(prefix_set.get()); | |
282 } | |
283 | |
284 // Bad |index_| size is caught by the sanity check. | |
285 TEST_F(PrefixSetTest, CorruptionVersion) { | |
286 FilePath filename; | |
287 ASSERT_TRUE(GetPrefixSetFile(&filename)); | |
288 | |
289 ASSERT_NO_FATAL_FAILURE( | |
290 ModifyAndCleanChecksum(filename, kVersionOffset, 1)); | |
291 scoped_ptr<safe_browsing::PrefixSet> | |
292 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); | |
293 ASSERT_FALSE(prefix_set.get()); | |
294 } | |
295 | |
296 // Bad |index_| size is caught by the sanity check. | |
297 TEST_F(PrefixSetTest, CorruptionIndexSize) { | |
298 FilePath filename; | |
299 ASSERT_TRUE(GetPrefixSetFile(&filename)); | |
300 | |
301 ASSERT_NO_FATAL_FAILURE( | |
302 ModifyAndCleanChecksum(filename, kIndexSizeOffset, 1)); | |
303 scoped_ptr<safe_browsing::PrefixSet> | |
304 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); | |
305 ASSERT_FALSE(prefix_set.get()); | |
306 } | |
307 | |
308 // Bad |deltas_| size is caught by the sanity check. | |
309 TEST_F(PrefixSetTest, CorruptionDeltasSize) { | |
310 FilePath filename; | |
311 ASSERT_TRUE(GetPrefixSetFile(&filename)); | |
312 | |
313 ASSERT_NO_FATAL_FAILURE( | |
314 ModifyAndCleanChecksum(filename, kDeltasSizeOffset, 1)); | |
315 scoped_ptr<safe_browsing::PrefixSet> | |
316 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); | |
317 ASSERT_FALSE(prefix_set.get()); | |
318 } | |
319 | |
320 // Test that the digest catches corruption in the middle of the file | |
321 // (in the payload between the header and the digest). | |
322 TEST_F(PrefixSetTest, CorruptionPayload) { | |
323 FilePath filename; | |
324 ASSERT_TRUE(GetPrefixSetFile(&filename)); | |
325 | |
326 file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b")); | |
327 ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), 666, 1)); | |
328 file.reset(); | |
329 scoped_ptr<safe_browsing::PrefixSet> | |
330 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); | |
331 ASSERT_FALSE(prefix_set.get()); | |
332 } | |
333 | |
334 // Test corruption in the digest itself. | |
335 TEST_F(PrefixSetTest, CorruptionDigest) { | |
336 FilePath filename; | |
337 ASSERT_TRUE(GetPrefixSetFile(&filename)); | |
338 | |
339 int64 size_64; | |
340 ASSERT_TRUE(file_util::GetFileSize(filename, &size_64)); | |
341 file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b")); | |
342 long digest_offset = static_cast<long>(size_64 - sizeof(MD5Digest)); | |
343 ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), digest_offset, 1)); | |
344 file.reset(); | |
345 scoped_ptr<safe_browsing::PrefixSet> | |
346 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); | |
347 ASSERT_FALSE(prefix_set.get()); | |
348 } | |
349 | |
350 // Test excess data after the digest (fails the size test). | |
351 TEST_F(PrefixSetTest, CorruptionExcess) { | |
352 FilePath filename; | |
353 ASSERT_TRUE(GetPrefixSetFile(&filename)); | |
354 | |
355 // Add some junk to the trunk. | |
356 file_util::ScopedFILE file(file_util::OpenFile(filename, "ab")); | |
357 const char buf[] = "im in ur base, killing ur d00dz."; | |
358 ASSERT_EQ(strlen(buf), fwrite(buf, 1, strlen(buf), file.get())); | |
359 file.reset(); | |
360 scoped_ptr<safe_browsing::PrefixSet> | |
361 prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); | |
362 ASSERT_FALSE(prefix_set.get()); | |
363 } | |
364 | |
135 } // namespace | 365 } // namespace |
OLD | NEW |