chrome/browser/safe_browsing/prefix_set.cc

   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
   9 #include "base/file_util.h"
  10 #include "base/files/scoped_file.h"
  11 #include "base/logging.h"
  12 #include "base/md5.h"
  13 #include "base/metrics/histogram.h"
  14 #include "base/metrics/sparse_histogram.h"
  15
  16 namespace {
  17
  18 // |kMagic| should be reasonably unique, and not match itself across
  19 // endianness changes.  I generated this value with:
  20 // md5 -qs chrome/browser/safe_browsing/prefix_set.cc | colrm 9
  21 static uint32 kMagic = 0x864088dd;
  22
  23 // Version history:
  24 // Version 1: b6cb7cfe/r74487 by shess@chromium.org on 2011-02-10
  25 // Version 2: 2b59b0a6/r253924 by shess@chromium.org on 2014-02-27
  26 // Version 3: ????????/r?????? by shess@chromium.org on 2014-04-??
  27
  28 // Version 2 layout is identical to version 1.  The sort order of |index_|
  29 // changed from |int32| to |uint32| to match the change of |SBPrefix|.
  30 // Version 3 adds storage for full hashes.
  31 static uint32 kVersion = 3;
  32 static uint32 kDeprecatedVersion = 1;  // And lower.
  33
  34 typedef struct {
  35   uint32 magic;
  36   uint32 version;
  37   uint32 index_size;
  38   uint32 deltas_size;
  39 } FileHeader_v2;
  40
  41 typedef struct {
  42   uint32 magic;
  43   uint32 version;
  44   uint32 index_size;
  45   uint32 deltas_size;
  46   uint32 full_hashes_size;
  47 } FileHeader;
  48
  49 // Common std::vector<> implementations add capacity by multiplying from the
  50 // current size (usually either by 2 or 1.5) to satisfy push_back() running in
  51 // amortized constant time.  This is not necessary for insert() at end(), but
  52 // AFAICT it seems true for some implementations.  SBPrefix values should
  53 // uniformly cover the 32-bit space, so the final size can be estimated given a
  54 // subset of the input.
  55 //
  56 // |kEstimateThreshold| is when estimates start converging.  Results are strong
  57 // starting around 1<<27.  1<<30 is chosen to prevent the degenerate case of
  58 // resizing capacity from >50% to 100%.
  59 //
  60 // TODO(shess): I'm sure there is math in the world to describe good settings
  61 // for estimating the size of a uniformly-distributed set of integers from a
  62 // sorted subset.  I do not have such math in me, so I assumed that my current
  63 // organic database of prefixes was scale-free, and wrote a script to see how
  64 // often given slop values would always suffice for given strides.  At 1<<30,
  65 // .5% slop was sufficient to cover all cases (though the code below uses 1%).
  66 //
  67 // TODO(shess): A smaller threshold uses less transient space in reallocation.
  68 // 1<<30 uses between 125% and 150%, 1<<29 between 112% and 125%, etc.  The cost
  69 // is that a smaller threshold needs more slop (locked down for the long term).
  70 // 1<<29 worked well with 1%, 1<<27 worked well with 2%.
  71 const SBPrefix kEstimateThreshold = 1 << 30;
  72 size_t EstimateFinalCount(SBPrefix current_prefix, size_t current_count) {
  73   // estimated_count / current_count == estimated_max / current_prefix
  74   // For large input sets, estimated_max of 2^32 is close enough.
  75   const size_t estimated_prefix_count = static_cast<size_t>(
  76       (static_cast<uint64>(current_count) << 32) / current_prefix);
  77
  78   // The estimate has an error bar, if the final total is below the estimate, no
  79   // harm, but if it is above a capacity resize will happen at nearly 100%.  Add
  80   // some slop to make sure all cases are covered.
  81   return estimated_prefix_count + estimated_prefix_count / 100;
  82 }
  83
  84 }  // namespace
  85
  86 namespace safe_browsing {
  87
  88 // For |std::upper_bound()| to find a prefix w/in a vector of pairs.
  89 // static
  90 bool PrefixSet::PrefixLess(const IndexPair& a, const IndexPair& b) {
  91   return a.first < b.first;
  92 }
  93
  94 PrefixSet::PrefixSet() {
  95 }
  96
  97 PrefixSet::PrefixSet(IndexVector* index,
  98                      std::vector<uint16>* deltas,
  99                      std::vector<SBFullHash>* full_hashes) {
 100   DCHECK(index && deltas && full_hashes);
 101   index_.swap(*index);
 102   deltas_.swap(*deltas);
 103   full_hashes_.swap(*full_hashes);
 104 }
 105
 106 PrefixSet::~PrefixSet() {}
 107
 108 bool PrefixSet::PrefixExists(SBPrefix prefix) const {
 109   if (index_.empty())
 110     return false;
 111
 112   // Find the first position after |prefix| in |index_|.
 113   IndexVector::const_iterator iter =
 114       std::upper_bound(index_.begin(), index_.end(),
 115                        IndexPair(prefix, 0), PrefixLess);
 116
 117   // |prefix| comes before anything that's in the set.
 118   if (iter == index_.begin())
 119     return false;
 120
 121   // Capture the upper bound of our target entry's deltas.
 122   const size_t bound = (iter == index_.end() ? deltas_.size() : iter->second);
 123
 124   // Back up to the entry our target is in.
 125   --iter;
 126
 127   // All prefixes in |index_| are in the set.
 128   SBPrefix current = iter->first;
 129   if (current == prefix)
 130     return true;
 131
 132   // Scan forward accumulating deltas while a match is possible.
 133   for (size_t di = iter->second; di < bound && current < prefix; ++di) {
 134     current += deltas_[di];
 135   }
 136
 137   return current == prefix;
 138 }
 139
 140 bool PrefixSet::Exists(const SBFullHash& hash) const {
 141   if (std::binary_search(full_hashes_.begin(), full_hashes_.end(),
 142                          hash, SBFullHashLess)) {
 143     return true;
 144   }
 145   return PrefixExists(hash.prefix);
 146 }
 147
 148 void PrefixSet::GetPrefixes(std::vector<SBPrefix>* prefixes) const {
 149   prefixes->reserve(index_.size() + deltas_.size());
 150
 151   for (size_t ii = 0; ii < index_.size(); ++ii) {
 152     // The deltas for this |index_| entry run to the next index entry,
 153     // or the end of the deltas.
 154     const size_t deltas_end =
 155         (ii + 1 < index_.size()) ? index_[ii + 1].second : deltas_.size();
 156
 157     SBPrefix current = index_[ii].first;
 158     prefixes->push_back(current);
 159     for (size_t di = index_[ii].second; di < deltas_end; ++di) {
 160       current += deltas_[di];
 161       prefixes->push_back(current);
 162     }
 163   }
 164 }
 165
 166 // static
 167 scoped_ptr<PrefixSet> PrefixSet::LoadFile(const base::FilePath& filter_name) {
 168   int64 size_64;
 169   if (!base::GetFileSize(filter_name, &size_64))
 170     return scoped_ptr<PrefixSet>();
 171   using base::MD5Digest;
 172   // TODO(shess): Revert to sizeof(FileHeader) for sanity check once v2 is
 173   // deprecated.
 174   if (size_64 < static_cast<int64>(sizeof(FileHeader_v2) + sizeof(MD5Digest)))
 175     return scoped_ptr<PrefixSet>();
 176
 177   base::ScopedFILE file(base::OpenFile(filter_name, "rb"));
 178   if (!file.get())
 179     return scoped_ptr<PrefixSet>();
 180
 181   FileHeader header;
 182   size_t read = fread(&header, sizeof(header), 1, file.get());
 183   if (read != 1)
 184     return scoped_ptr<PrefixSet>();
 185
 186   // The file looks valid, start building the digest.
 187   base::MD5Context context;
 188   base::MD5Init(&context);
 189   base::MD5Update(&context, base::StringPiece(reinterpret_cast<char*>(&header),
 190                                               sizeof(header)));
 191
 192   if (header.magic != kMagic)
 193     return scoped_ptr<PrefixSet>();
 194
 195   // Track version read to inform removal of support for older versions.
 196   UMA_HISTOGRAM_SPARSE_SLOWLY("SB2.PrefixSetVersionRead", header.version);
 197
 198   // TODO(shess): <http://crbug.com/368044> for removing v2 support.
 199   size_t header_size = sizeof(header);
 200   if (header.version <= kDeprecatedVersion) {
 201     return scoped_ptr<PrefixSet>();
 202   } else if (header.version == 2) {
 203     // Rewind the file and restart building the digest with the old header
 204     // structure.
 205     FileHeader_v2 v2_header;
 206     if (0 != fseek(file.get(), 0, SEEK_SET))
 207       return scoped_ptr<PrefixSet>();
 208
 209     size_t read = fread(&v2_header, sizeof(v2_header), 1, file.get());
 210     if (read != 1)
 211       return scoped_ptr<PrefixSet>();
 212
 213     base::MD5Init(&context);
 214     base::MD5Update(&context,
 215                     base::StringPiece(reinterpret_cast<char*>(&v2_header),
 216                                       sizeof(v2_header)));
 217
 218     // The current header is a superset of the old header, fill it in with the
 219     // information read.
 220     header.magic = v2_header.magic;
 221     header.version = v2_header.version;
 222     header.index_size = v2_header.index_size;
 223     header.deltas_size = v2_header.deltas_size;
 224     header.full_hashes_size = 0;
 225     header_size = sizeof(v2_header);
 226   } else if (header.version != kVersion) {
 227     return scoped_ptr<PrefixSet>();
 228   }
 229
 230   IndexVector index;
 231   const size_t index_bytes = sizeof(index[0]) * header.index_size;
 232
 233   std::vector<uint16> deltas;
 234   const size_t deltas_bytes = sizeof(deltas[0]) * header.deltas_size;
 235
 236   std::vector<SBFullHash> full_hashes;
 237   const size_t full_hashes_bytes =
 238       sizeof(full_hashes[0]) * header.full_hashes_size;
 239
 240   // Check for bogus sizes before allocating any space.
 241   const size_t expected_bytes = header_size +
 242       index_bytes + deltas_bytes + full_hashes_bytes + sizeof(MD5Digest);
 243   if (static_cast<int64>(expected_bytes) != size_64)
 244     return scoped_ptr<PrefixSet>();
 245
 246   // Read the index vector.  Herb Sutter indicates that vectors are
 247   // guaranteed to be contiuguous, so reading to where element 0 lives
 248   // is valid.
 249   if (header.index_size) {
 250     index.resize(header.index_size);
 251     read = fread(&(index[0]), sizeof(index[0]), index.size(), file.get());
 252     if (read != index.size())
 253       return scoped_ptr<PrefixSet>();
 254     base::MD5Update(&context,
 255                     base::StringPiece(reinterpret_cast<char*>(&(index[0])),
 256                                       index_bytes));
 257   }
 258
 259   // Read vector of deltas.
 260   if (header.deltas_size) {
 261     deltas.resize(header.deltas_size);
 262     read = fread(&(deltas[0]), sizeof(deltas[0]), deltas.size(), file.get());
 263     if (read != deltas.size())
 264       return scoped_ptr<PrefixSet>();
 265     base::MD5Update(&context,
 266                     base::StringPiece(reinterpret_cast<char*>(&(deltas[0])),
 267                                       deltas_bytes));
 268   }
 269
 270   // Read vector of full hashes.
 271   if (header.full_hashes_size) {
 272     full_hashes.resize(header.full_hashes_size);
 273     read = fread(&(full_hashes[0]), sizeof(full_hashes[0]), full_hashes.size(),
 274                  file.get());
 275     if (read != full_hashes.size())
 276       return scoped_ptr<PrefixSet>();
 277     base::MD5Update(&context,
 278                     base::StringPiece(
 279                         reinterpret_cast<char*>(&(full_hashes[0])),
 280                         full_hashes_bytes));
 281   }
 282
 283   base::MD5Digest calculated_digest;
 284   base::MD5Final(&calculated_digest, &context);
 285
 286   base::MD5Digest file_digest;
 287   read = fread(&file_digest, sizeof(file_digest), 1, file.get());
 288   if (read != 1)
 289     return scoped_ptr<PrefixSet>();
 290
 291   if (0 != memcmp(&file_digest, &calculated_digest, sizeof(file_digest)))
 292     return scoped_ptr<PrefixSet>();
 293
 294   // Steals vector contents using swap().
 295   return scoped_ptr<PrefixSet>(new PrefixSet(&index, &deltas, &full_hashes));
 296 }
 297
 298 bool PrefixSet::WriteFile(const base::FilePath& filter_name) const {
 299   FileHeader header;
 300   header.magic = kMagic;
 301   header.version = kVersion;
 302   header.index_size = static_cast<uint32>(index_.size());
 303   header.deltas_size = static_cast<uint32>(deltas_.size());
 304   header.full_hashes_size = static_cast<uint32>(full_hashes_.size());
 305
 306   // Sanity check that the 32-bit values never mess things up.
 307   if (static_cast<size_t>(header.index_size) != index_.size() ||
 308       static_cast<size_t>(header.deltas_size) != deltas_.size() ||
 309       static_cast<size_t>(header.full_hashes_size) != full_hashes_.size()) {
 310     NOTREACHED();
 311     return false;
 312   }
 313
 314   base::ScopedFILE file(base::OpenFile(filter_name, "wb"));
 315   if (!file.get())
 316     return false;
 317
 318   base::MD5Context context;
 319   base::MD5Init(&context);
 320
 321   // TODO(shess): The I/O code in safe_browsing_store_file.cc would
 322   // sure be useful about now.
 323   size_t written = fwrite(&header, sizeof(header), 1, file.get());
 324   if (written != 1)
 325     return false;
 326   base::MD5Update(&context, base::StringPiece(reinterpret_cast<char*>(&header),
 327                                               sizeof(header)));
 328
 329   // As for reads, the standard guarantees the ability to access the
 330   // contents of the vector by a pointer to an element.
 331   if (index_.size()) {
 332     const size_t index_bytes = sizeof(index_[0]) * index_.size();
 333     written = fwrite(&(index_[0]), sizeof(index_[0]), index_.size(),
 334                      file.get());
 335     if (written != index_.size())
 336       return false;
 337     base::MD5Update(&context,
 338                     base::StringPiece(
 339                         reinterpret_cast<const char*>(&(index_[0])),
 340                         index_bytes));
 341   }
 342
 343   if (deltas_.size()) {
 344     const size_t deltas_bytes = sizeof(deltas_[0]) * deltas_.size();
 345     written = fwrite(&(deltas_[0]), sizeof(deltas_[0]), deltas_.size(),
 346                      file.get());
 347     if (written != deltas_.size())
 348       return false;
 349     base::MD5Update(&context,
 350                     base::StringPiece(
 351                         reinterpret_cast<const char*>(&(deltas_[0])),
 352                         deltas_bytes));
 353   }
 354
 355   if (full_hashes_.size()) {
 356     const size_t elt_size = sizeof(full_hashes_[0]);
 357     const size_t elts = full_hashes_.size();
 358     const size_t full_hashes_bytes = elt_size * elts;
 359     written = fwrite(&(full_hashes_[0]), elt_size, elts, file.get());
 360     if (written != elts)
 361       return false;
 362     base::MD5Update(&context,
 363                     base::StringPiece(
 364                         reinterpret_cast<const char*>(&(full_hashes_[0])),
 365                         full_hashes_bytes));
 366   }
 367
 368   base::MD5Digest digest;
 369   base::MD5Final(&digest, &context);
 370   written = fwrite(&digest, sizeof(digest), 1, file.get());
 371   if (written != 1)
 372     return false;
 373
 374   // TODO(shess): Can this code check that the close was successful?
 375   file.reset();
 376
 377   return true;
 378 }
 379
 380 void PrefixSet::AddRun(SBPrefix index_prefix,
 381                        const uint16* run_begin, const uint16* run_end) {
 382   // Preempt organic capacity decisions for |delta_| once strong estimates can
 383   // be made.
 384   if (index_prefix > kEstimateThreshold &&
 385       deltas_.capacity() < deltas_.size() + (run_end - run_begin)) {
 386     deltas_.reserve(EstimateFinalCount(index_prefix, deltas_.size()));
 387   }
 388
 389   index_.push_back(std::make_pair(index_prefix, deltas_.size()));
 390   deltas_.insert(deltas_.end(), run_begin, run_end);
 391 }
 392
 393 PrefixSetBuilder::PrefixSetBuilder()
 394     : prefix_set_(new PrefixSet()) {
 395 }
 396
 397 PrefixSetBuilder::PrefixSetBuilder(const std::vector<SBPrefix>& prefixes)
 398     : prefix_set_(new PrefixSet()) {
 399   for (size_t i = 0; i < prefixes.size(); ++i) {
 400     AddPrefix(prefixes[i]);
 401   }
 402 }
 403
 404 PrefixSetBuilder::~PrefixSetBuilder() {
 405 }
 406
 407 scoped_ptr<PrefixSet> PrefixSetBuilder::GetPrefixSet(
 408     const std::vector<SBFullHash>& hashes) {
 409   DCHECK(prefix_set_.get());
 410
 411   // Flush runs until buffered data is gone.
 412   while (!buffer_.empty()) {
 413     EmitRun();
 414   }
 415
 416   // Precisely size |index_| for read-only.  It's 50k-60k, so minor savings, but
 417   // they're almost free.
 418   PrefixSet::IndexVector(prefix_set_->index_).swap(prefix_set_->index_);
 419
 420   prefix_set_->full_hashes_ = hashes;
 421   std::sort(prefix_set_->full_hashes_.begin(), prefix_set_->full_hashes_.end(),
 422             SBFullHashLess);
 423
 424   return prefix_set_.Pass();
 425 }
 426
 427 scoped_ptr<PrefixSet> PrefixSetBuilder::GetPrefixSetNoHashes() {
 428   return GetPrefixSet(std::vector<SBFullHash>()).Pass();
 429 }
 430
 431 void PrefixSetBuilder::EmitRun() {
 432   DCHECK(prefix_set_.get());
 433
 434   SBPrefix prev_prefix = buffer_[0];
 435   uint16 run[PrefixSet::kMaxRun];
 436   size_t run_pos = 0;
 437
 438   size_t i;
 439   for (i = 1; i < buffer_.size() && run_pos < PrefixSet::kMaxRun; ++i) {
 440     // Calculate the delta.  |unsigned| is mandatory, because the
 441     // sorted_prefixes could be more than INT_MAX apart.
 442     DCHECK_GT(buffer_[i], prev_prefix);
 443     const unsigned delta = buffer_[i] - prev_prefix;
 444     const uint16 delta16 = static_cast<uint16>(delta);
 445
 446     // Break the run if the delta doesn't fit.
 447     if (delta != static_cast<unsigned>(delta16))
 448       break;
 449
 450     // Continue the run of deltas.
 451     run[run_pos++] = delta16;
 452     DCHECK_EQ(static_cast<unsigned>(run[run_pos - 1]), delta);
 453
 454     prev_prefix = buffer_[i];
 455   }
 456   prefix_set_->AddRun(buffer_[0], run, run + run_pos);
 457   buffer_.erase(buffer_.begin(), buffer_.begin() + i);
 458 }
 459
 460 void PrefixSetBuilder::AddPrefix(SBPrefix prefix) {
 461   DCHECK(prefix_set_.get());
 462
 463   if (buffer_.empty()) {
 464     DCHECK(prefix_set_->index_.empty());
 465     DCHECK(prefix_set_->deltas_.empty());
 466   } else {
 467     // Drop duplicates.
 468     if (buffer_.back() == prefix)
 469       return;
 470
 471     DCHECK_LT(buffer_.back(), prefix);
 472   }
 473   buffer_.push_back(prefix);
 474
 475   // Flush buffer when a run can be constructed.  +1 for the index item, and +1
 476   // to leave at least one item in the buffer for dropping duplicates.
 477   if (buffer_.size() > PrefixSet::kMaxRun + 2)
 478     EmitRun();
 479 }
 480
 481 }  // namespace safe_browsing