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// A read-only set implementation for |SBPrefix| items.  Prefixes are
6// sorted and stored as 16-bit deltas from the previous prefix.  An
7// index structure provides quick random access, and also handles
8// cases where 16 bits cannot encode a delta.
9//
10// For example, the sequence {20, 25, 41, 65432, 150000, 160000} would
11// be stored as:
12//  A pair {20, 0} in |index_|.
13//  5, 16, 65391 in |deltas_|.
14//  A pair {150000, 3} in |index_|.
15//  10000 in |deltas_|.
16// |index_.size()| will be 2, |deltas_.size()| will be 4.
17//
18// This structure is intended for storage of sparse uniform sets of
19// prefixes of a certain size.  As of this writing, my safe-browsing
20// database contains:
21//   653132 add prefixes
22//   6446 are duplicates (from different chunks)
23//   24301 w/in 2^8 of the prior prefix
24//   622337 w/in 2^16 of the prior prefix
25//   47 further than 2^16 from the prior prefix
26// For this input, the memory usage is approximately 2 bytes per
27// prefix, a bit over 1.2M.  The bloom filter used 25 bits per prefix,
28// a bit over 1.9M on this data.
29//
30// Experimenting with random selections of the above data, storage
31// size drops almost linearly as prefix count drops, until the index
32// overhead starts to become a problem a bit under 200k prefixes.  The
33// memory footprint gets worse than storing the raw prefix data around
34// 75k prefixes.  Fortunately, the actual memory footprint also falls.
35// If the prefix count increases the memory footprint should increase
36// approximately linearly.  The worst-case would be 2^16 items all
37// 2^16 apart, which would need 512k (versus 256k to store the raw
38// data).
39//
40// The on-disk format looks like:
41//         4 byte magic number
42//         4 byte version number
43//         4 byte |index_.size()|
44//         4 byte |deltas_.size()|
45//     n * 8 byte |&index_[0]..&index_[n]|
46//     m * 2 byte |&deltas_[0]..&deltas_[m]|
47//        16 byte digest
48
49#ifndef CHROME_BROWSER_SAFE_BROWSING_PREFIX_SET_H_
50#define CHROME_BROWSER_SAFE_BROWSING_PREFIX_SET_H_
51
52#include <utility>
53#include <vector>
54
55#include "base/gtest_prod_util.h"
56#include "base/memory/scoped_ptr.h"
57#include "chrome/browser/safe_browsing/safe_browsing_util.h"
58
59namespace base {
60class FilePath;
61}
62
63namespace safe_browsing {
64
65class PrefixSet {
66 public:
67  ~PrefixSet();
68
69  // |true| if |hash| is in the hashes passed to the set's builder, or if
70  // |hash.prefix| is one of the prefixes passed to the set's builder.
71  bool Exists(const SBFullHash& hash) const;
72
73  // Persist the set on disk.
74  static scoped_ptr<PrefixSet> LoadFile(const base::FilePath& filter_name);
75  bool WriteFile(const base::FilePath& filter_name) const;
76
77 private:
78  friend class PrefixSetBuilder;
79
80  friend class PrefixSetTest;
81  FRIEND_TEST_ALL_PREFIXES(PrefixSetTest, AllBig);
82  FRIEND_TEST_ALL_PREFIXES(PrefixSetTest, EdgeCases);
83  FRIEND_TEST_ALL_PREFIXES(PrefixSetTest, Empty);
84  FRIEND_TEST_ALL_PREFIXES(PrefixSetTest, FullHashBuild);
85  FRIEND_TEST_ALL_PREFIXES(PrefixSetTest, IntMinMax);
86  FRIEND_TEST_ALL_PREFIXES(PrefixSetTest, OneElement);
87  FRIEND_TEST_ALL_PREFIXES(PrefixSetTest, ReadWrite);
88  FRIEND_TEST_ALL_PREFIXES(PrefixSetTest, ReadWriteSigned);
89  FRIEND_TEST_ALL_PREFIXES(PrefixSetTest, Version3);
90
91  FRIEND_TEST_ALL_PREFIXES(SafeBrowsingStoreFileTest, BasicStore);
92  FRIEND_TEST_ALL_PREFIXES(SafeBrowsingStoreFileTest, DeleteChunks);
93  FRIEND_TEST_ALL_PREFIXES(SafeBrowsingStoreFileTest, DetectsCorruption);
94  FRIEND_TEST_ALL_PREFIXES(SafeBrowsingStoreFileTest, Empty);
95  FRIEND_TEST_ALL_PREFIXES(SafeBrowsingStoreFileTest, PrefixMinMax);
96  FRIEND_TEST_ALL_PREFIXES(SafeBrowsingStoreFileTest, SubKnockout);
97  FRIEND_TEST_ALL_PREFIXES(SafeBrowsingStoreFileTest, Version7);
98  FRIEND_TEST_ALL_PREFIXES(SafeBrowsingStoreFileTest, Version8);
99
100  // Maximum number of consecutive deltas to encode before generating
101  // a new index entry.  This helps keep the worst-case performance
102  // for |Exists()| under control.
103  static const size_t kMaxRun = 100;
104
105  // Helpers to make |index_| easier to deal with.
106  typedef std::pair<SBPrefix, uint32> IndexPair;
107  typedef std::vector<IndexPair> IndexVector;
108  static bool PrefixLess(const IndexPair& a, const IndexPair& b);
109
110  // Helper to let |PrefixSetBuilder| add a run of data.  |index_prefix| is
111  // added to |index_|, with the other elements added into |deltas_|.
112  void AddRun(SBPrefix index_prefix,
113              const uint16* run_begin, const uint16* run_end);
114
115  // |true| if |prefix| is one of the prefixes passed to the set's builder.
116  // Provided for testing purposes.
117  bool PrefixExists(SBPrefix prefix) const;
118
119  // Regenerate the vector of prefixes passed to the constructor into
120  // |prefixes|.  Prefixes will be added in sorted order.  Useful for testing.
121  void GetPrefixes(std::vector<SBPrefix>* prefixes) const;
122
123  // Used by |PrefixSetBuilder|.
124  PrefixSet();
125
126  // Helper for |LoadFile()|.  Steals vector contents using |swap()|.
127  PrefixSet(IndexVector* index,
128            std::vector<uint16>* deltas,
129            std::vector<SBFullHash>* full_hashes);
130
131  // Top-level index of prefix to offset in |deltas_|.  Each pair
132  // indicates a base prefix and where the deltas from that prefix
133  // begin in |deltas_|.  The deltas for a pair end at the next pair's
134  // index into |deltas_|.
135  IndexVector index_;
136
137  // Deltas which are added to the prefix in |index_| to generate
138  // prefixes.  Deltas are only valid between consecutive items from
139  // |index_|, or the end of |deltas_| for the last |index_| pair.
140  std::vector<uint16> deltas_;
141
142  // Full hashes ordered by SBFullHashLess.
143  std::vector<SBFullHash> full_hashes_;
144
145  DISALLOW_COPY_AND_ASSIGN(PrefixSet);
146};
147
148// Helper to incrementally build a PrefixSet from a stream of sorted prefixes.
149class PrefixSetBuilder {
150 public:
151  PrefixSetBuilder();
152  ~PrefixSetBuilder();
153
154  // Helper for unit tests and format conversion.
155  explicit PrefixSetBuilder(const std::vector<SBPrefix>& prefixes);
156
157  // Add a prefix to the set.  Prefixes must arrive in ascending order.
158  // Duplicate prefixes are dropped.
159  void AddPrefix(SBPrefix prefix);
160
161  // Flush any buffered prefixes, and return the final PrefixSet instance.
162  // |hashes| are sorted and stored in |full_hashes_|.  Any call other than the
163  // destructor is illegal after this call.
164  scoped_ptr<PrefixSet> GetPrefixSet(const std::vector<SBFullHash>& hashes);
165
166  // Helper for clients which only track prefixes.  Calls GetPrefixSet() with
167  // empty hash vector.
168  scoped_ptr<PrefixSet> GetPrefixSetNoHashes();
169
170 private:
171  // Encode a run of deltas for |AddRun()|.  The run is broken by a too-large
172  // delta, or kMaxRun, whichever comes first.
173  void EmitRun();
174
175  // Buffers prefixes until enough are avaliable to emit a run.
176  std::vector<SBPrefix> buffer_;
177
178  // The PrefixSet being built.
179  scoped_ptr<PrefixSet> prefix_set_;
180};
181
182}  // namespace safe_browsing
183
184#endif  // CHROME_BROWSER_SAFE_BROWSING_PREFIX_SET_H_
185