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#include <set>
10#include <string>
11
12#include "base/files/file_util.h"
13#include "base/files/scoped_file.h"
14#include "base/files/scoped_temp_dir.h"
15#include "base/logging.h"
16#include "base/md5.h"
17#include "base/memory/scoped_ptr.h"
18#include "base/path_service.h"
19#include "base/rand_util.h"
20#include "base/strings/string_number_conversions.h"
21#include "base/strings/string_util.h"
22#include "chrome/common/chrome_paths.h"
23#include "testing/gtest/include/gtest/gtest.h"
24#include "testing/platform_test.h"
25
26namespace {
27
28const SBPrefix kHighBitClear = 1000u * 1000u * 1000u;
29const SBPrefix kHighBitSet = 3u * 1000u * 1000u * 1000u;
30
31}  // namespace
32
33namespace safe_browsing {
34
35class PrefixSetTest : public PlatformTest {
36 protected:
37  // Constants for the v1 format.
38  static const size_t kMagicOffset = 0 * sizeof(uint32);
39  static const size_t kVersionOffset = 1 * sizeof(uint32);
40  static const size_t kIndexSizeOffset = 2 * sizeof(uint32);
41  static const size_t kDeltasSizeOffset = 3 * sizeof(uint32);
42  static const size_t kFullHashesSizeOffset = 4 * sizeof(uint32);
43  static const size_t kPayloadOffset = 5 * sizeof(uint32);
44
45  // Generate a set of random prefixes to share between tests.  For
46  // most tests this generation was a large fraction of the test time.
47  //
48  // The set should contain sparse areas where adjacent items are more
49  // than 2^16 apart, and dense areas where adjacent items are less
50  // than 2^16 apart.
51  static void SetUpTestCase() {
52    // Distribute clusters of prefixes.
53    for (size_t i = 0; i < 250; ++i) {
54      // Unsigned for overflow characteristics.
55      const uint32 base = static_cast<uint32>(base::RandUint64());
56      for (size_t j = 0; j < 10; ++j) {
57        const uint32 delta = static_cast<uint32>(base::RandUint64() & 0xFFFF);
58        const SBPrefix prefix = static_cast<SBPrefix>(base + delta);
59        shared_prefixes_.push_back(prefix);
60      }
61    }
62
63    // Lay down a sparsely-distributed layer.
64    const size_t count = shared_prefixes_.size();
65    for (size_t i = 0; i < count; ++i) {
66      const SBPrefix prefix = static_cast<SBPrefix>(base::RandUint64());
67      shared_prefixes_.push_back(prefix);
68    }
69
70    // Sort for use with PrefixSet constructor.
71    std::sort(shared_prefixes_.begin(), shared_prefixes_.end());
72  }
73
74  // Check that all elements of |prefixes| are in |prefix_set|, and
75  // that nearby elements are not (for lack of a more sensible set of
76  // items to check for absence).
77  static void CheckPrefixes(const PrefixSet& prefix_set,
78                            const std::vector<SBPrefix> &prefixes) {
79    // The set can generate the prefixes it believes it has, so that's
80    // a good starting point.
81    std::set<SBPrefix> check(prefixes.begin(), prefixes.end());
82    std::vector<SBPrefix> prefixes_copy;
83    prefix_set.GetPrefixes(&prefixes_copy);
84    EXPECT_EQ(prefixes_copy.size(), check.size());
85    EXPECT_TRUE(std::equal(check.begin(), check.end(), prefixes_copy.begin()));
86
87    for (size_t i = 0; i < prefixes.size(); ++i) {
88      EXPECT_TRUE(prefix_set.PrefixExists(prefixes[i]));
89
90      const SBPrefix left_sibling = prefixes[i] - 1;
91      if (check.count(left_sibling) == 0)
92        EXPECT_FALSE(prefix_set.PrefixExists(left_sibling));
93
94      const SBPrefix right_sibling = prefixes[i] + 1;
95      if (check.count(right_sibling) == 0)
96        EXPECT_FALSE(prefix_set.PrefixExists(right_sibling));
97    }
98  }
99
100  // Generate a |PrefixSet| file from |shared_prefixes_|, store it in
101  // a temporary file, and return the filename in |filenamep|.
102  // Returns |true| on success.
103  bool GetPrefixSetFile(base::FilePath* filenamep) {
104    if (!temp_dir_.IsValid() && !temp_dir_.CreateUniqueTempDir())
105      return false;
106
107    base::FilePath filename = temp_dir_.path().AppendASCII("PrefixSetTest");
108
109    PrefixSetBuilder builder(shared_prefixes_);
110    if (!builder.GetPrefixSetNoHashes()->WriteFile(filename))
111      return false;
112
113    *filenamep = filename;
114    return true;
115  }
116
117  // Helper function to read the uint32 value at |offset|, increment it
118  // by |inc|, and write it back in place.  |fp| should be opened in
119  // r+ mode.
120  static void IncrementIntAt(FILE* fp, long offset, int inc) {
121    uint32 value = 0;
122
123    ASSERT_NE(-1, fseek(fp, offset, SEEK_SET));
124    ASSERT_EQ(1U, fread(&value, sizeof(value), 1, fp));
125
126    value += inc;
127
128    ASSERT_NE(-1, fseek(fp, offset, SEEK_SET));
129    ASSERT_EQ(1U, fwrite(&value, sizeof(value), 1, fp));
130  }
131
132  // Helper function to re-generated |fp|'s checksum to be correct for
133  // the file's contents.  |fp| should be opened in r+ mode.
134  static void CleanChecksum(FILE* fp) {
135    base::MD5Context context;
136    base::MD5Init(&context);
137
138    ASSERT_NE(-1, fseek(fp, 0, SEEK_END));
139    long file_size = ftell(fp);
140
141    using base::MD5Digest;
142    size_t payload_size = static_cast<size_t>(file_size) - sizeof(MD5Digest);
143    size_t digested_size = 0;
144    ASSERT_NE(-1, fseek(fp, 0, SEEK_SET));
145    while (digested_size < payload_size) {
146      char buf[1024];
147      size_t nitems = std::min(payload_size - digested_size, sizeof(buf));
148      ASSERT_EQ(nitems, fread(buf, 1, nitems, fp));
149      base::MD5Update(&context, base::StringPiece(buf, nitems));
150      digested_size += nitems;
151    }
152    ASSERT_EQ(digested_size, payload_size);
153    ASSERT_EQ(static_cast<long>(digested_size), ftell(fp));
154
155    base::MD5Digest new_digest;
156    base::MD5Final(&new_digest, &context);
157    ASSERT_NE(-1, fseek(fp, digested_size, SEEK_SET));
158    ASSERT_EQ(1U, fwrite(&new_digest, sizeof(new_digest), 1, fp));
159    ASSERT_EQ(file_size, ftell(fp));
160  }
161
162  // Open |filename| and increment the uint32 at |offset| by |inc|.
163  // Then re-generate the checksum to account for the new contents.
164  void ModifyAndCleanChecksum(const base::FilePath& filename, long offset,
165                              int inc) {
166    int64 size_64;
167    ASSERT_TRUE(base::GetFileSize(filename, &size_64));
168
169    base::ScopedFILE file(base::OpenFile(filename, "r+b"));
170    IncrementIntAt(file.get(), offset, inc);
171    CleanChecksum(file.get());
172    file.reset();
173
174    int64 new_size_64;
175    ASSERT_TRUE(base::GetFileSize(filename, &new_size_64));
176    ASSERT_EQ(new_size_64, size_64);
177  }
178
179  // Fill |prefixes| with values read from a reference file.  The reference file
180  // was generated from a specific |shared_prefixes_|.
181  bool ReadReferencePrefixes(std::vector<SBPrefix>* prefixes) {
182    const char kRefname[] = "PrefixSetRef";
183    base::FilePath ref_path;
184    if (!PathService::Get(chrome::DIR_TEST_DATA, &ref_path))
185      return false;
186    ref_path = ref_path.AppendASCII("SafeBrowsing");
187    ref_path = ref_path.AppendASCII(kRefname);
188
189    base::ScopedFILE file(base::OpenFile(ref_path, "r"));
190    if (!file.get())
191      return false;
192    char buf[1024];
193    while (fgets(buf, sizeof(buf), file.get())) {
194      std::string trimmed;
195      if (base::TRIM_TRAILING !=
196          base::TrimWhitespace(buf, base::TRIM_ALL, &trimmed))
197        return false;
198      unsigned prefix;
199      if (!base::StringToUint(trimmed, &prefix))
200        return false;
201      prefixes->push_back(prefix);
202    }
203    return true;
204  }
205
206  // Tests should not modify this shared resource.
207  static std::vector<SBPrefix> shared_prefixes_;
208
209  base::ScopedTempDir temp_dir_;
210};
211
212std::vector<SBPrefix> PrefixSetTest::shared_prefixes_;
213
214// Test that a small sparse random input works.
215TEST_F(PrefixSetTest, Baseline) {
216  PrefixSetBuilder builder(shared_prefixes_);
217  CheckPrefixes(*builder.GetPrefixSetNoHashes(), shared_prefixes_);
218}
219
220// Test that the empty set doesn't appear to have anything in it.
221TEST_F(PrefixSetTest, Empty) {
222  const std::vector<SBPrefix> empty;
223  PrefixSetBuilder builder(empty);
224  scoped_ptr<PrefixSet> prefix_set = builder.GetPrefixSetNoHashes();
225  for (size_t i = 0; i < shared_prefixes_.size(); ++i) {
226    EXPECT_FALSE(prefix_set->PrefixExists(shared_prefixes_[i]));
227  }
228}
229
230// Single-element set should work fine.
231TEST_F(PrefixSetTest, OneElement) {
232  const std::vector<SBPrefix> prefixes(100, 0u);
233  PrefixSetBuilder builder(prefixes);
234  scoped_ptr<PrefixSet> prefix_set = builder.GetPrefixSetNoHashes();
235  EXPECT_FALSE(prefix_set->PrefixExists(static_cast<SBPrefix>(-1)));
236  EXPECT_TRUE(prefix_set->PrefixExists(prefixes[0]));
237  EXPECT_FALSE(prefix_set->PrefixExists(1u));
238
239  // Check that |GetPrefixes()| returns the same set of prefixes as
240  // was passed in.
241  std::vector<SBPrefix> prefixes_copy;
242  prefix_set->GetPrefixes(&prefixes_copy);
243  EXPECT_EQ(1U, prefixes_copy.size());
244  EXPECT_EQ(prefixes[0], prefixes_copy[0]);
245}
246
247// Edges of the 32-bit integer range.
248TEST_F(PrefixSetTest, IntMinMax) {
249  std::vector<SBPrefix> prefixes;
250
251  // Using bit patterns rather than portable constants because this
252  // really is testing how the entire 32-bit integer range is handled.
253  prefixes.push_back(0x00000000);
254  prefixes.push_back(0x0000FFFF);
255  prefixes.push_back(0x7FFF0000);
256  prefixes.push_back(0x7FFFFFFF);
257  prefixes.push_back(0x80000000);
258  prefixes.push_back(0x8000FFFF);
259  prefixes.push_back(0xFFFF0000);
260  prefixes.push_back(0xFFFFFFFF);
261
262  std::sort(prefixes.begin(), prefixes.end());
263  PrefixSetBuilder builder(prefixes);
264  scoped_ptr<PrefixSet> prefix_set = builder.GetPrefixSetNoHashes();
265
266  // Check that |GetPrefixes()| returns the same set of prefixes as
267  // was passed in.
268  std::vector<SBPrefix> prefixes_copy;
269  prefix_set->GetPrefixes(&prefixes_copy);
270  ASSERT_EQ(prefixes_copy.size(), prefixes.size());
271  EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
272                         prefixes_copy.begin()));
273}
274
275// A range with only large deltas.
276TEST_F(PrefixSetTest, AllBig) {
277  std::vector<SBPrefix> prefixes;
278
279  const unsigned kDelta = 10 * 1000 * 1000;
280  for (SBPrefix prefix = kHighBitSet;
281       prefix < kHighBitClear; prefix += kDelta) {
282    prefixes.push_back(prefix);
283  }
284
285  std::sort(prefixes.begin(), prefixes.end());
286  PrefixSetBuilder builder(prefixes);
287  scoped_ptr<PrefixSet> prefix_set = builder.GetPrefixSetNoHashes();
288
289  // Check that |GetPrefixes()| returns the same set of prefixes as
290  // was passed in.
291  std::vector<SBPrefix> prefixes_copy;
292  prefix_set->GetPrefixes(&prefixes_copy);
293  prefixes.erase(std::unique(prefixes.begin(), prefixes.end()), prefixes.end());
294  EXPECT_EQ(prefixes_copy.size(), prefixes.size());
295  EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
296                         prefixes_copy.begin()));
297}
298
299// Use artificial inputs to test various edge cases in PrefixExists().  Items
300// before the lowest item aren't present.  Items after the largest item aren't
301// present.  Create a sequence of items with deltas above and below 2^16, and
302// make sure they're all present.  Create a very long sequence with deltas below
303// 2^16 to test crossing |kMaxRun|.
304TEST_F(PrefixSetTest, EdgeCases) {
305  std::vector<SBPrefix> prefixes;
306
307  // Put in a high-bit prefix.
308  SBPrefix prefix = kHighBitSet;
309  prefixes.push_back(prefix);
310
311  // Add a sequence with very large deltas.
312  unsigned delta = 100 * 1000 * 1000;
313  for (int i = 0; i < 10; ++i) {
314    prefix += delta;
315    prefixes.push_back(prefix);
316  }
317
318  // Add a sequence with deltas that start out smaller than the
319  // maximum delta, and end up larger.  Also include some duplicates.
320  delta = 256 * 256 - 100;
321  for (int i = 0; i < 200; ++i) {
322    prefix += delta;
323    prefixes.push_back(prefix);
324    prefixes.push_back(prefix);
325    delta++;
326  }
327
328  // Add a long sequence with deltas smaller than the maximum delta,
329  // so a new index item will be injected.
330  delta = 256 * 256 - 1;
331  prefix = kHighBitClear - delta * 1000;
332  prefixes.push_back(prefix);
333  for (int i = 0; i < 1000; ++i) {
334    prefix += delta;
335    prefixes.push_back(prefix);
336    delta--;
337  }
338
339  std::sort(prefixes.begin(), prefixes.end());
340  PrefixSetBuilder builder(prefixes);
341  scoped_ptr<PrefixSet> prefix_set = builder.GetPrefixSetNoHashes();
342
343  // Check that |GetPrefixes()| returns the same set of prefixes as
344  // was passed in.
345  std::vector<SBPrefix> prefixes_copy;
346  prefix_set->GetPrefixes(&prefixes_copy);
347  prefixes.erase(std::unique(prefixes.begin(), prefixes.end()), prefixes.end());
348  EXPECT_EQ(prefixes_copy.size(), prefixes.size());
349  EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
350                         prefixes_copy.begin()));
351
352  // Items before and after the set are not present, and don't crash.
353  EXPECT_FALSE(prefix_set->PrefixExists(kHighBitSet - 100));
354  EXPECT_FALSE(prefix_set->PrefixExists(kHighBitClear + 100));
355
356  // Check that the set correctly flags all of the inputs, and also
357  // check items just above and below the inputs to make sure they
358  // aren't present.
359  for (size_t i = 0; i < prefixes.size(); ++i) {
360    EXPECT_TRUE(prefix_set->PrefixExists(prefixes[i]));
361
362    EXPECT_FALSE(prefix_set->PrefixExists(prefixes[i] - 1));
363    EXPECT_FALSE(prefix_set->PrefixExists(prefixes[i] + 1));
364  }
365}
366
367// Test writing a prefix set to disk and reading it back in.
368TEST_F(PrefixSetTest, ReadWrite) {
369  base::FilePath filename;
370
371  // Write the sample prefix set out, read it back in, and check all
372  // the prefixes.  Leaves the path in |filename|.
373  {
374    ASSERT_TRUE(GetPrefixSetFile(&filename));
375    scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
376    ASSERT_TRUE(prefix_set.get());
377    CheckPrefixes(*prefix_set, shared_prefixes_);
378  }
379
380  // Test writing and reading a very sparse set containing no deltas.
381  {
382    std::vector<SBPrefix> prefixes;
383    prefixes.push_back(kHighBitClear);
384    prefixes.push_back(kHighBitSet);
385
386    PrefixSetBuilder builder(prefixes);
387    ASSERT_TRUE(builder.GetPrefixSetNoHashes()->WriteFile(filename));
388
389    scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
390    ASSERT_TRUE(prefix_set.get());
391    CheckPrefixes(*prefix_set, prefixes);
392  }
393
394  // Test writing and reading an empty set.
395  {
396    std::vector<SBPrefix> prefixes;
397    PrefixSetBuilder builder(prefixes);
398    ASSERT_TRUE(builder.GetPrefixSetNoHashes()->WriteFile(filename));
399
400    scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
401    ASSERT_TRUE(prefix_set.get());
402    CheckPrefixes(*prefix_set, prefixes);
403  }
404
405  // Test that full hashes are persisted.
406  {
407    std::vector<SBFullHash> hashes;
408    hashes.push_back(SBFullHashForString("one"));
409    hashes.push_back(SBFullHashForString("two"));
410    hashes.push_back(SBFullHashForString("three"));
411
412    std::vector<SBPrefix> prefixes(shared_prefixes_);
413
414    // Remove any collisions from the prefixes.
415    for (size_t i = 0; i < hashes.size(); ++i) {
416      std::vector<SBPrefix>::iterator iter =
417          std::lower_bound(prefixes.begin(), prefixes.end(), hashes[i].prefix);
418      if (iter != prefixes.end() && *iter == hashes[i].prefix)
419        prefixes.erase(iter);
420    }
421
422    PrefixSetBuilder builder(prefixes);
423    ASSERT_TRUE(builder.GetPrefixSet(hashes)->WriteFile(filename));
424
425    scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
426    ASSERT_TRUE(prefix_set.get());
427    CheckPrefixes(*prefix_set, prefixes);
428
429    EXPECT_TRUE(prefix_set->Exists(hashes[0]));
430    EXPECT_TRUE(prefix_set->Exists(hashes[1]));
431    EXPECT_TRUE(prefix_set->Exists(hashes[2]));
432    EXPECT_FALSE(prefix_set->PrefixExists(hashes[0].prefix));
433    EXPECT_FALSE(prefix_set->PrefixExists(hashes[1].prefix));
434    EXPECT_FALSE(prefix_set->PrefixExists(hashes[2].prefix));
435  }
436}
437
438// Check that |CleanChecksum()| makes an acceptable checksum.
439TEST_F(PrefixSetTest, CorruptionHelpers) {
440  base::FilePath filename;
441  ASSERT_TRUE(GetPrefixSetFile(&filename));
442
443  // This will modify data in |index_|, which will fail the digest check.
444  base::ScopedFILE file(base::OpenFile(filename, "r+b"));
445  IncrementIntAt(file.get(), kPayloadOffset, 1);
446  file.reset();
447  scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
448  ASSERT_FALSE(prefix_set.get());
449
450  // Fix up the checksum and it will read successfully (though the
451  // data will be wrong).
452  file.reset(base::OpenFile(filename, "r+b"));
453  CleanChecksum(file.get());
454  file.reset();
455  prefix_set = PrefixSet::LoadFile(filename);
456  ASSERT_TRUE(prefix_set.get());
457}
458
459// Bad magic is caught by the sanity check.
460TEST_F(PrefixSetTest, CorruptionMagic) {
461  base::FilePath filename;
462  ASSERT_TRUE(GetPrefixSetFile(&filename));
463
464  ASSERT_NO_FATAL_FAILURE(
465      ModifyAndCleanChecksum(filename, kMagicOffset, 1));
466  scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
467  ASSERT_FALSE(prefix_set.get());
468}
469
470// Bad version is caught by the sanity check.
471TEST_F(PrefixSetTest, CorruptionVersion) {
472  base::FilePath filename;
473  ASSERT_TRUE(GetPrefixSetFile(&filename));
474
475  ASSERT_NO_FATAL_FAILURE(
476      ModifyAndCleanChecksum(filename, kVersionOffset, 10));
477  scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
478  ASSERT_FALSE(prefix_set.get());
479}
480
481// Bad |index_| size is caught by the sanity check.
482TEST_F(PrefixSetTest, CorruptionIndexSize) {
483  base::FilePath filename;
484  ASSERT_TRUE(GetPrefixSetFile(&filename));
485
486  ASSERT_NO_FATAL_FAILURE(
487      ModifyAndCleanChecksum(filename, kIndexSizeOffset, 1));
488  scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
489  ASSERT_FALSE(prefix_set.get());
490}
491
492// Bad |deltas_| size is caught by the sanity check.
493TEST_F(PrefixSetTest, CorruptionDeltasSize) {
494  base::FilePath filename;
495  ASSERT_TRUE(GetPrefixSetFile(&filename));
496
497  ASSERT_NO_FATAL_FAILURE(
498      ModifyAndCleanChecksum(filename, kDeltasSizeOffset, 1));
499  scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
500  ASSERT_FALSE(prefix_set.get());
501}
502
503// Bad |full_hashes_| size is caught by the sanity check.
504TEST_F(PrefixSetTest, CorruptionFullHashesSize) {
505  base::FilePath filename;
506  ASSERT_TRUE(GetPrefixSetFile(&filename));
507
508  ASSERT_NO_FATAL_FAILURE(
509      ModifyAndCleanChecksum(filename, kFullHashesSizeOffset, 1));
510  scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
511  ASSERT_FALSE(prefix_set.get());
512}
513
514// Test that the digest catches corruption in the middle of the file
515// (in the payload between the header and the digest).
516TEST_F(PrefixSetTest, CorruptionPayload) {
517  base::FilePath filename;
518  ASSERT_TRUE(GetPrefixSetFile(&filename));
519
520  base::ScopedFILE file(base::OpenFile(filename, "r+b"));
521  ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), 666, 1));
522  file.reset();
523  scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
524  ASSERT_FALSE(prefix_set.get());
525}
526
527// Test corruption in the digest itself.
528TEST_F(PrefixSetTest, CorruptionDigest) {
529  base::FilePath filename;
530  ASSERT_TRUE(GetPrefixSetFile(&filename));
531
532  int64 size_64;
533  ASSERT_TRUE(base::GetFileSize(filename, &size_64));
534  base::ScopedFILE file(base::OpenFile(filename, "r+b"));
535  long digest_offset = static_cast<long>(size_64 - sizeof(base::MD5Digest));
536  ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), digest_offset, 1));
537  file.reset();
538  scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
539  ASSERT_FALSE(prefix_set.get());
540}
541
542// Test excess data after the digest (fails the size test).
543TEST_F(PrefixSetTest, CorruptionExcess) {
544  base::FilePath filename;
545  ASSERT_TRUE(GetPrefixSetFile(&filename));
546
547  // Add some junk to the trunk.
548  base::ScopedFILE file(base::OpenFile(filename, "ab"));
549  const char buf[] = "im in ur base, killing ur d00dz.";
550  ASSERT_EQ(strlen(buf), fwrite(buf, 1, strlen(buf), file.get()));
551  file.reset();
552  scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
553  ASSERT_FALSE(prefix_set.get());
554}
555
556// Test that files which had 64-bit size_t are discarded.
557TEST_F(PrefixSetTest, SizeTRecovery) {
558  base::FilePath filename;
559  ASSERT_TRUE(GetPrefixSetFile(&filename));
560
561  // Open the file for rewrite.
562  base::ScopedFILE file(base::OpenFile(filename, "r+b"));
563
564  // Leave existing magic and version.
565  ASSERT_NE(-1, fseek(file.get(), sizeof(uint32) * 2, SEEK_SET));
566
567  // Indicate two index values and two deltas.
568  uint32 val = 2;
569  ASSERT_EQ(sizeof(val), fwrite(&val, 1, sizeof(val), file.get()));
570  ASSERT_EQ(sizeof(val), fwrite(&val, 1, sizeof(val), file.get()));
571
572  // Write two index values with 64-bit "size_t".
573  std::pair<SBPrefix, uint64> item;
574  memset(&item, 0, sizeof(item));  // Includes any padding.
575  item.first = 17;
576  item.second = 0;
577  ASSERT_EQ(sizeof(item), fwrite(&item, 1, sizeof(item), file.get()));
578  item.first = 100042;
579  item.second = 1;
580  ASSERT_EQ(sizeof(item), fwrite(&item, 1, sizeof(item), file.get()));
581
582  // Write two delta values.
583  uint16 delta = 23;
584  ASSERT_EQ(sizeof(delta), fwrite(&delta, 1, sizeof(delta), file.get()));
585  ASSERT_EQ(sizeof(delta), fwrite(&delta, 1, sizeof(delta), file.get()));
586
587  // Leave space for the digest at the end, and regenerate it.
588  base::MD5Digest dummy = { { 0 } };
589  ASSERT_EQ(sizeof(dummy), fwrite(&dummy, 1, sizeof(dummy), file.get()));
590  ASSERT_TRUE(base::TruncateFile(file.get()));
591  CleanChecksum(file.get());
592  file.reset();  // Flush updates.
593
594  scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(filename);
595  ASSERT_FALSE(prefix_set.get());
596}
597
598// Test Exists() against full hashes passed to builder.
599TEST_F(PrefixSetTest, FullHashBuild) {
600  const SBFullHash kHash1 = SBFullHashForString("one");
601  const SBFullHash kHash2 = SBFullHashForString("two");
602  const SBFullHash kHash3 = SBFullHashForString("three");
603  const SBFullHash kHash4 = SBFullHashForString("four");
604  const SBFullHash kHash5 = SBFullHashForString("five");
605  const SBFullHash kHash6 = SBFullHashForString("six");
606
607  std::vector<SBPrefix> prefixes;
608  prefixes.push_back(kHash1.prefix);
609  prefixes.push_back(kHash2.prefix);
610  std::sort(prefixes.begin(), prefixes.end());
611
612  std::vector<SBFullHash> hashes;
613  hashes.push_back(kHash4);
614  hashes.push_back(kHash5);
615
616  PrefixSetBuilder builder(prefixes);
617  scoped_ptr<PrefixSet> prefix_set = builder.GetPrefixSet(hashes);
618
619  EXPECT_TRUE(prefix_set->Exists(kHash1));
620  EXPECT_TRUE(prefix_set->Exists(kHash2));
621  EXPECT_FALSE(prefix_set->Exists(kHash3));
622  EXPECT_TRUE(prefix_set->Exists(kHash4));
623  EXPECT_TRUE(prefix_set->Exists(kHash5));
624  EXPECT_FALSE(prefix_set->Exists(kHash6));
625
626  EXPECT_TRUE(prefix_set->PrefixExists(kHash1.prefix));
627  EXPECT_TRUE(prefix_set->PrefixExists(kHash2.prefix));
628  EXPECT_FALSE(prefix_set->PrefixExists(kHash3.prefix));
629  EXPECT_FALSE(prefix_set->PrefixExists(kHash4.prefix));
630  EXPECT_FALSE(prefix_set->PrefixExists(kHash5.prefix));
631  EXPECT_FALSE(prefix_set->PrefixExists(kHash6.prefix));
632}
633
634// Test that a version 1 file is discarded on read.
635TEST_F(PrefixSetTest, ReadSigned) {
636  base::FilePath filename;
637  ASSERT_TRUE(GetPrefixSetFile(&filename));
638
639  // Open the file for rewrite.
640  base::ScopedFILE file(base::OpenFile(filename, "r+b"));
641
642  // Leave existing magic.
643  ASSERT_NE(-1, fseek(file.get(), sizeof(uint32), SEEK_SET));
644
645  // Version 1.
646  uint32 version = 1;
647  ASSERT_EQ(sizeof(version), fwrite(&version, 1, sizeof(version), file.get()));
648
649  // Indicate two index values and two deltas.
650  uint32 val = 2;
651  ASSERT_EQ(sizeof(val), fwrite(&val, 1, sizeof(val), file.get()));
652  ASSERT_EQ(sizeof(val), fwrite(&val, 1, sizeof(val), file.get()));
653
654  std::pair<int32, uint32> item;
655  memset(&item, 0, sizeof(item));  // Includes any padding.
656  item.first = -1000;
657  item.second = 0;
658  ASSERT_EQ(sizeof(item), fwrite(&item, 1, sizeof(item), file.get()));
659  item.first = 1000;
660  item.second = 1;
661  ASSERT_EQ(sizeof(item), fwrite(&item, 1, sizeof(item), file.get()));
662
663  // Write two delta values.
664  uint16 delta = 23;
665  ASSERT_EQ(sizeof(delta), fwrite(&delta, 1, sizeof(delta), file.get()));
666  ASSERT_EQ(sizeof(delta), fwrite(&delta, 1, sizeof(delta), file.get()));
667
668  // Leave space for the digest at the end, and regenerate it.
669  base::MD5Digest dummy = { { 0 } };
670  ASSERT_EQ(sizeof(dummy), fwrite(&dummy, 1, sizeof(dummy), file.get()));
671  ASSERT_TRUE(base::TruncateFile(file.get()));
672  CleanChecksum(file.get());
673  file.reset();  // Flush updates.
674
675  scoped_ptr<safe_browsing::PrefixSet>
676      prefix_set(safe_browsing::PrefixSet::LoadFile(filename));
677  ASSERT_FALSE(prefix_set.get());
678}
679
680// Test that a golden v2 file is discarded on read.  All platforms generating v2
681// files are little-endian, so there is no point to testing this transition
682// if/when a big-endian port is added.
683#if defined(ARCH_CPU_LITTLE_ENDIAN)
684TEST_F(PrefixSetTest, Version2) {
685  std::vector<SBPrefix> ref_prefixes;
686  ASSERT_TRUE(ReadReferencePrefixes(&ref_prefixes));
687
688  const char kBasename[] = "PrefixSetVersion2";
689  base::FilePath golden_path;
690  ASSERT_TRUE(PathService::Get(chrome::DIR_TEST_DATA, &golden_path));
691  golden_path = golden_path.AppendASCII("SafeBrowsing");
692  golden_path = golden_path.AppendASCII(kBasename);
693
694  scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(golden_path);
695  ASSERT_FALSE(prefix_set.get());
696}
697#endif
698
699// Test that a golden v3 file can be read by the current code.  All platforms
700// generating v3 files are little-endian, so there is no point to testing this
701// transition if/when a big-endian port is added.
702#if defined(ARCH_CPU_LITTLE_ENDIAN)
703TEST_F(PrefixSetTest, Version3) {
704  std::vector<SBPrefix> ref_prefixes;
705  ASSERT_TRUE(ReadReferencePrefixes(&ref_prefixes));
706
707  const char kBasename[] = "PrefixSetVersion3";
708  base::FilePath golden_path;
709  ASSERT_TRUE(PathService::Get(chrome::DIR_TEST_DATA, &golden_path));
710  golden_path = golden_path.AppendASCII("SafeBrowsing");
711  golden_path = golden_path.AppendASCII(kBasename);
712
713  scoped_ptr<PrefixSet> prefix_set = PrefixSet::LoadFile(golden_path);
714  ASSERT_TRUE(prefix_set.get());
715  CheckPrefixes(*prefix_set, ref_prefixes);
716
717  const SBFullHash kHash1 = SBFullHashForString("www.evil.com/malware.html");
718  const SBFullHash kHash2 = SBFullHashForString("www.evil.com/phishing.html");
719
720  EXPECT_TRUE(prefix_set->Exists(kHash1));
721  EXPECT_TRUE(prefix_set->Exists(kHash2));
722  EXPECT_FALSE(prefix_set->PrefixExists(kHash1.prefix));
723  EXPECT_FALSE(prefix_set->PrefixExists(kHash2.prefix));
724}
725#endif
726
727}  // namespace safe_browsing
728