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
3// found in the LICENSE file.
4
5#include "chrome/browser/safe_browsing/prefix_set.h"
6
7#include <algorithm>
8
9#include "base/file_util.h"
10#include "base/logging.h"
11#include "base/md5.h"
12#include "base/memory/scoped_ptr.h"
13#include "base/memory/scoped_temp_dir.h"
14#include "base/rand_util.h"
15#include "testing/gtest/include/gtest/gtest.h"
16#include "testing/platform_test.h"
17
18namespace {
19
20class PrefixSetTest : public PlatformTest {
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
150std::vector<SBPrefix> PrefixSetTest::shared_prefixes_;
151
152// Test that a small sparse random input works.
153TEST_F(PrefixSetTest, Baseline) {
154  safe_browsing::PrefixSet prefix_set(shared_prefixes_);
155  CheckPrefixes(&prefix_set, shared_prefixes_);
156}
157
158// Test that the empty set doesn't appear to have anything in it.
159TEST_F(PrefixSetTest, Empty) {
160  const std::vector<SBPrefix> empty;
161  safe_browsing::PrefixSet prefix_set(empty);
162  for (size_t i = 0; i < shared_prefixes_.size(); ++i) {
163    EXPECT_FALSE(prefix_set.Exists(shared_prefixes_[i]));
164  }
165}
166
167// Single-element set should work fine.
168TEST_F(PrefixSetTest, OneElement) {
169  const std::vector<SBPrefix> prefixes(100, 0);
170  safe_browsing::PrefixSet prefix_set(prefixes);
171  EXPECT_FALSE(prefix_set.Exists(-1));
172  EXPECT_TRUE(prefix_set.Exists(prefixes[0]));
173  EXPECT_FALSE(prefix_set.Exists(1));
174
175  // Check that |GetPrefixes()| returns the same set of prefixes as
176  // was passed in.
177  std::vector<SBPrefix> prefixes_copy;
178  prefix_set.GetPrefixes(&prefixes_copy);
179  EXPECT_EQ(1U, prefixes_copy.size());
180  EXPECT_EQ(prefixes[0], prefixes_copy[0]);
181}
182
183// Edges of the 32-bit integer range.
184TEST_F(PrefixSetTest, IntMinMax) {
185  std::vector<SBPrefix> prefixes;
186
187  // Using bit patterns rather than portable constants because this
188  // really is testing how the entire 32-bit integer range is handled.
189  prefixes.push_back(0x00000000);
190  prefixes.push_back(0x0000FFFF);
191  prefixes.push_back(0x7FFF0000);
192  prefixes.push_back(0x7FFFFFFF);
193  prefixes.push_back(0x80000000);
194  prefixes.push_back(0x8000FFFF);
195  prefixes.push_back(0xFFFF0000);
196  prefixes.push_back(0xFFFFFFFF);
197
198  std::sort(prefixes.begin(), prefixes.end());
199  safe_browsing::PrefixSet prefix_set(prefixes);
200
201  // Check that |GetPrefixes()| returns the same set of prefixes as
202  // was passed in.
203  std::vector<SBPrefix> prefixes_copy;
204  prefix_set.GetPrefixes(&prefixes_copy);
205  ASSERT_EQ(prefixes_copy.size(), prefixes.size());
206  EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
207                         prefixes_copy.begin()));
208}
209
210// A range with only large deltas.
211TEST_F(PrefixSetTest, AllBig) {
212  std::vector<SBPrefix> prefixes;
213
214  const SBPrefix kVeryPositive = 1000 * 1000 * 1000;
215  const SBPrefix kVeryNegative = -kVeryPositive;
216  const unsigned kDelta = 10 * 1000 * 1000;
217
218  for (SBPrefix prefix = kVeryNegative;
219       prefix < kVeryPositive; prefix += kDelta) {
220    prefixes.push_back(prefix);
221  }
222
223  std::sort(prefixes.begin(), prefixes.end());
224  safe_browsing::PrefixSet prefix_set(prefixes);
225
226  // Check that |GetPrefixes()| returns the same set of prefixes as
227  // was passed in.
228  std::vector<SBPrefix> prefixes_copy;
229  prefix_set.GetPrefixes(&prefixes_copy);
230  prefixes.erase(std::unique(prefixes.begin(), prefixes.end()), prefixes.end());
231  EXPECT_EQ(prefixes_copy.size(), prefixes.size());
232  EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
233                         prefixes_copy.begin()));
234}
235
236// Use artificial inputs to test various edge cases in Exists().
237// Items before the lowest item aren't present.  Items after the
238// largest item aren't present.  Create a sequence of items with
239// deltas above and below 2^16, and make sure they're all present.
240// Create a very long sequence with deltas below 2^16 to test crossing
241// |kMaxRun|.
242TEST_F(PrefixSetTest, EdgeCases) {
243  std::vector<SBPrefix> prefixes;
244
245  const SBPrefix kVeryPositive = 1000 * 1000 * 1000;
246  const SBPrefix kVeryNegative = -kVeryPositive;
247
248  // Put in a very negative prefix.
249  SBPrefix prefix = kVeryNegative;
250  prefixes.push_back(prefix);
251
252  // Add a sequence with very large deltas.
253  unsigned delta = 100 * 1000 * 1000;
254  for (int i = 0; i < 10; ++i) {
255    prefix += delta;
256    prefixes.push_back(prefix);
257  }
258
259  // Add a sequence with deltas that start out smaller than the
260  // maximum delta, and end up larger.  Also include some duplicates.
261  delta = 256 * 256 - 100;
262  for (int i = 0; i < 200; ++i) {
263    prefix += delta;
264    prefixes.push_back(prefix);
265    prefixes.push_back(prefix);
266    delta++;
267  }
268
269  // Add a long sequence with deltas smaller than the maximum delta,
270  // so a new index item will be injected.
271  delta = 256 * 256 - 1;
272  prefix = kVeryPositive - delta * 1000;
273  prefixes.push_back(prefix);
274  for (int i = 0; i < 1000; ++i) {
275    prefix += delta;
276    prefixes.push_back(prefix);
277    delta--;
278  }
279
280  std::sort(prefixes.begin(), prefixes.end());
281  safe_browsing::PrefixSet prefix_set(prefixes);
282
283  // Check that |GetPrefixes()| returns the same set of prefixes as
284  // was passed in.
285  std::vector<SBPrefix> prefixes_copy;
286  prefix_set.GetPrefixes(&prefixes_copy);
287  prefixes.erase(std::unique(prefixes.begin(), prefixes.end()), prefixes.end());
288  EXPECT_EQ(prefixes_copy.size(), prefixes.size());
289  EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
290                         prefixes_copy.begin()));
291
292  // Items before and after the set are not present, and don't crash.
293  EXPECT_FALSE(prefix_set.Exists(kVeryNegative - 100));
294  EXPECT_FALSE(prefix_set.Exists(kVeryPositive + 100));
295
296  // Check that the set correctly flags all of the inputs, and also
297  // check items just above and below the inputs to make sure they
298  // aren't present.
299  for (size_t i = 0; i < prefixes.size(); ++i) {
300    EXPECT_TRUE(prefix_set.Exists(prefixes[i]));
301
302    EXPECT_FALSE(prefix_set.Exists(prefixes[i] - 1));
303    EXPECT_FALSE(prefix_set.Exists(prefixes[i] + 1));
304  }
305}
306
307// Similar to Baseline test, but write the set out to a file and read
308// it back in before testing.
309TEST_F(PrefixSetTest, ReadWrite) {
310  FilePath filename;
311  ASSERT_TRUE(GetPrefixSetFile(&filename));
312
313  scoped_ptr<safe_browsing::PrefixSet>
314      prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
315  ASSERT_TRUE(prefix_set.get());
316
317  CheckPrefixes(prefix_set.get(), shared_prefixes_);
318}
319
320// Check that |CleanChecksum()| makes an acceptable checksum.
321TEST_F(PrefixSetTest, CorruptionHelpers) {
322  FilePath filename;
323  ASSERT_TRUE(GetPrefixSetFile(&filename));
324
325  // This will modify data in |index_|, which will fail the digest check.
326  file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b"));
327  IncrementIntAt(file.get(), kPayloadOffset, 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  // Fix up the checksum and it will read successfully (though the
334  // data will be wrong).
335  file.reset(file_util::OpenFile(filename, "r+b"));
336  CleanChecksum(file.get());
337  file.reset();
338  prefix_set.reset(safe_browsing::PrefixSet::LoadFile(filename));
339  ASSERT_TRUE(prefix_set.get());
340}
341
342// Bad |index_| size is caught by the sanity check.
343TEST_F(PrefixSetTest, CorruptionMagic) {
344  FilePath filename;
345  ASSERT_TRUE(GetPrefixSetFile(&filename));
346
347  ASSERT_NO_FATAL_FAILURE(
348      ModifyAndCleanChecksum(filename, kMagicOffset, 1));
349  scoped_ptr<safe_browsing::PrefixSet>
350      prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
351  ASSERT_FALSE(prefix_set.get());
352}
353
354// Bad |index_| size is caught by the sanity check.
355TEST_F(PrefixSetTest, CorruptionVersion) {
356  FilePath filename;
357  ASSERT_TRUE(GetPrefixSetFile(&filename));
358
359  ASSERT_NO_FATAL_FAILURE(
360      ModifyAndCleanChecksum(filename, kVersionOffset, 1));
361  scoped_ptr<safe_browsing::PrefixSet>
362      prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
363  ASSERT_FALSE(prefix_set.get());
364}
365
366// Bad |index_| size is caught by the sanity check.
367TEST_F(PrefixSetTest, CorruptionIndexSize) {
368  FilePath filename;
369  ASSERT_TRUE(GetPrefixSetFile(&filename));
370
371  ASSERT_NO_FATAL_FAILURE(
372      ModifyAndCleanChecksum(filename, kIndexSizeOffset, 1));
373  scoped_ptr<safe_browsing::PrefixSet>
374      prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
375  ASSERT_FALSE(prefix_set.get());
376}
377
378// Bad |deltas_| size is caught by the sanity check.
379TEST_F(PrefixSetTest, CorruptionDeltasSize) {
380  FilePath filename;
381  ASSERT_TRUE(GetPrefixSetFile(&filename));
382
383  ASSERT_NO_FATAL_FAILURE(
384      ModifyAndCleanChecksum(filename, kDeltasSizeOffset, 1));
385  scoped_ptr<safe_browsing::PrefixSet>
386      prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
387  ASSERT_FALSE(prefix_set.get());
388}
389
390// Test that the digest catches corruption in the middle of the file
391// (in the payload between the header and the digest).
392TEST_F(PrefixSetTest, CorruptionPayload) {
393  FilePath filename;
394  ASSERT_TRUE(GetPrefixSetFile(&filename));
395
396  file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b"));
397  ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), 666, 1));
398  file.reset();
399  scoped_ptr<safe_browsing::PrefixSet>
400      prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
401  ASSERT_FALSE(prefix_set.get());
402}
403
404// Test corruption in the digest itself.
405TEST_F(PrefixSetTest, CorruptionDigest) {
406  FilePath filename;
407  ASSERT_TRUE(GetPrefixSetFile(&filename));
408
409  int64 size_64;
410  ASSERT_TRUE(file_util::GetFileSize(filename, &size_64));
411  file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b"));
412  long digest_offset = static_cast<long>(size_64 - sizeof(MD5Digest));
413  ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), digest_offset, 1));
414  file.reset();
415  scoped_ptr<safe_browsing::PrefixSet>
416      prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
417  ASSERT_FALSE(prefix_set.get());
418}
419
420// Test excess data after the digest (fails the size test).
421TEST_F(PrefixSetTest, CorruptionExcess) {
422  FilePath filename;
423  ASSERT_TRUE(GetPrefixSetFile(&filename));
424
425  // Add some junk to the trunk.
426  file_util::ScopedFILE file(file_util::OpenFile(filename, "ab"));
427  const char buf[] = "im in ur base, killing ur d00dz.";
428  ASSERT_EQ(strlen(buf), fwrite(buf, 1, strlen(buf), file.get()));
429  file.reset();
430  scoped_ptr<safe_browsing::PrefixSet>
431      prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
432  ASSERT_FALSE(prefix_set.get());
433}
434
435// TODO(shess): Remove once the problem is debugged.  But, until then,
436// make sure the accessors work!
437TEST_F(PrefixSetTest, DebuggingAccessors) {
438  std::vector<SBPrefix> prefixes;
439  std::unique_copy(shared_prefixes_.begin(), shared_prefixes_.end(),
440                   std::back_inserter(prefixes));
441  safe_browsing::PrefixSet prefix_set(prefixes);
442
443  EXPECT_EQ(prefixes.size(), prefix_set.GetSize());
444  EXPECT_FALSE(prefix_set.IsDeltaAt(0));
445  for (size_t i = 1; i < prefixes.size(); ++i) {
446    const int delta = prefixes[i] - prefixes[i - 1];
447    if (delta > 0xFFFF) {
448      EXPECT_FALSE(prefix_set.IsDeltaAt(i));
449    } else {
450      ASSERT_TRUE(prefix_set.IsDeltaAt(i));
451      EXPECT_EQ(delta, prefix_set.DeltaAt(i));
452    }
453  }
454}
455
456}  // namespace
457