diff options
6 files changed, 104 insertions, 12 deletions
diff --git a/chrome/browser/safe_browsing/safe_browsing_store.h b/chrome/browser/safe_browsing/safe_browsing_store.h index 86eb9a1..b099d9f 100644 --- a/chrome/browser/safe_browsing/safe_browsing_store.h +++ b/chrome/browser/safe_browsing/safe_browsing_store.h @@ -174,6 +174,9 @@ class SafeBrowsingStore { // successful. virtual bool Delete() = 0; + // Get all Add prefixes out from the store. + virtual bool GetAddPrefixes(std::vector<SBAddPrefix>* add_prefixes) = 0; + // Start an update. None of the following methods should be called // unless this returns true. If this returns true, the update // should be terminated by FinishUpdate() or CancelUpdate(). diff --git a/chrome/browser/safe_browsing/safe_browsing_store_file.cc b/chrome/browser/safe_browsing/safe_browsing_store_file.cc index f03080b..1791c32 100644 --- a/chrome/browser/safe_browsing/safe_browsing_store_file.cc +++ b/chrome/browser/safe_browsing/safe_browsing_store_file.cc @@ -40,6 +40,18 @@ bool FileRewind(FILE* fp) { return rv == 0; } +// Move file read pointer forward by |bytes| relative to current position. +bool FileSkip(size_t bytes, FILE* fp) { + // Although fseek takes negative values, for this case, we only want + // to skip forward. + DCHECK(static_cast<long>(bytes) > 0); + if (static_cast<long>(bytes) < 0) + return false; + int rv = fseek(fp, static_cast<long>(bytes), SEEK_CUR); + DCHECK_EQ(rv, 0); + return rv == 0; +} + // Read an array of |nmemb| items from |fp| into |ptr|, and fold the // input data into the checksum in |context|, if non-NULL. Return // true on success. @@ -169,6 +181,21 @@ bool FileHeaderSanityCheck(const FilePath& filename, return true; } +// This a helper function that reads header to |header|. Returns true if the +// magic number is correct and santiy check passes. +bool ReadAndVerifyHeader(const FilePath& filename, + FILE* fp, + FileHeader* header, + MD5Context* context) { + if (!ReadArray(header, 1, fp, context)) + return false; + if (header->magic != kFileMagic || header->version != kFileVersion) + return false; + if (!FileHeaderSanityCheck(filename, *header)) + return false; + return true; +} + } // namespace // static @@ -263,6 +290,28 @@ bool SafeBrowsingStoreFile::WriteAddPrefix(int32 chunk_id, SBPrefix prefix) { return true; } +bool SafeBrowsingStoreFile::GetAddPrefixes( + std::vector<SBAddPrefix>* add_prefixes) { + add_prefixes->clear(); + + file_util::ScopedFILE file(file_util::OpenFile(filename_, "rb")); + if (file.get() == NULL) return false; + + FileHeader header; + if (!ReadAndVerifyHeader(filename_, file.get(), &header, NULL)) + return OnCorruptDatabase(); + + size_t add_prefix_offset = header.add_chunk_count * sizeof(int32) + + header.sub_chunk_count * sizeof(int32); + if (!FileSkip(add_prefix_offset, file.get())) + return false; + + if (!ReadToVector(add_prefixes, header.add_prefix_count, file.get(), NULL)) + return false; + + return true; +} + bool SafeBrowsingStoreFile::WriteAddHash(int32 chunk_id, base::Time receive_time, const SBFullHash& full_hash) { @@ -442,6 +491,8 @@ bool SafeBrowsingStoreFile::DoUpdate( std::vector<SBAddFullHash>* add_full_hashes_result) { DCHECK(old_store_.get() || file_.get() || empty_); DCHECK(new_file_.get()); + CHECK(add_prefixes_result); + CHECK(add_full_hashes_result); std::vector<SBAddPrefix> add_prefixes; std::vector<SBSubPrefix> sub_prefixes; @@ -471,7 +522,6 @@ bool SafeBrowsingStoreFile::DoUpdate( if (!old_store_->CancelUpdate()) return OnCorruptDatabase(); } else if (!empty_) { - // Read |file_| into the vectors. DCHECK(file_.get()); if (!FileRewind(file_.get())) @@ -482,13 +532,7 @@ bool SafeBrowsingStoreFile::DoUpdate( // Read the file header and make sure it looks right. FileHeader header; - if (!ReadArray(&header, 1, file_.get(), &context)) - return OnCorruptDatabase(); - - if (header.magic != kFileMagic || header.version != kFileVersion) - return OnCorruptDatabase(); - - if (!FileHeaderSanityCheck(filename_, header)) + if (!ReadAndVerifyHeader(filename_, file_.get(), &header, &context)) return OnCorruptDatabase(); // Re-read the chunks-seen data to get to the later data in the @@ -518,6 +562,7 @@ bool SafeBrowsingStoreFile::DoUpdate( MD5Digest file_digest; if (!ReadArray(&file_digest, 1, file_.get(), NULL)) return OnCorruptDatabase(); + if (0 != memcmp(&file_digest, &calculated_digest, sizeof(file_digest))) return OnCorruptDatabase(); @@ -666,6 +711,9 @@ bool SafeBrowsingStoreFile::FinishUpdate( const std::set<SBPrefix>& prefix_misses, std::vector<SBAddPrefix>* add_prefixes_result, std::vector<SBAddFullHash>* add_full_hashes_result) { + DCHECK(add_prefixes_result); + DCHECK(add_full_hashes_result); + bool ret = DoUpdate(pending_adds, prefix_misses, add_prefixes_result, add_full_hashes_result); diff --git a/chrome/browser/safe_browsing/safe_browsing_store_file.h b/chrome/browser/safe_browsing/safe_browsing_store_file.h index 95dc030..02a2920 100644 --- a/chrome/browser/safe_browsing/safe_browsing_store_file.h +++ b/chrome/browser/safe_browsing/safe_browsing_store_file.h @@ -124,7 +124,11 @@ class SafeBrowsingStoreFile : public SafeBrowsingStore { // Delete any on-disk files, including the permanent storage. virtual bool Delete(); + // Get all Add prefixes from the store. + virtual bool GetAddPrefixes(std::vector<SBAddPrefix>* add_prefixes); + virtual bool BeginChunk(); + virtual bool WriteAddPrefix(int32 chunk_id, SBPrefix prefix); virtual bool WriteAddHash(int32 chunk_id, base::Time receive_time, @@ -136,10 +140,8 @@ class SafeBrowsingStoreFile : public SafeBrowsingStore { virtual bool FinishChunk(); virtual bool BeginUpdate(); - virtual bool DoUpdate(const std::vector<SBAddFullHash>& pending_adds, - const std::set<SBPrefix>& prefix_misses, - std::vector<SBAddPrefix>* add_prefixes_result, - std::vector<SBAddFullHash>* add_full_hashes_result); + // Store updates with pending add full hashes in file store and + // return |add_prefixes_result| and |add_full_hashes_result|. virtual bool FinishUpdate(const std::vector<SBAddFullHash>& pending_adds, const std::set<SBPrefix>& prefix_misses, std::vector<SBAddPrefix>* add_prefixes_result, @@ -163,6 +165,11 @@ class SafeBrowsingStoreFile : public SafeBrowsingStore { } private: + // Update store file with pending full hashes. + virtual bool DoUpdate(const std::vector<SBAddFullHash>& pending_adds, + const std::set<SBPrefix>& prefix_misses, + std::vector<SBAddPrefix>* add_prefixes_result, + std::vector<SBAddFullHash>* add_full_hashes_result); // Enumerate different format-change events for histogramming // purposes. DO NOT CHANGE THE ORDERING OF THESE VALUES. // TODO(shess): Remove this once the format change is complete. diff --git a/chrome/browser/safe_browsing/safe_browsing_store_file_unittest.cc b/chrome/browser/safe_browsing/safe_browsing_store_file_unittest.cc index 452a52e..01977a7 100644 --- a/chrome/browser/safe_browsing/safe_browsing_store_file_unittest.cc +++ b/chrome/browser/safe_browsing/safe_browsing_store_file_unittest.cc @@ -185,6 +185,18 @@ void LoadStore(SafeBrowsingStore* store) { EXPECT_TRUE(store->FinishUpdate(pending_adds, prefix_misses, &add_prefixes, &add_hashes)); EXPECT_EQ(3U, add_prefixes.size()); + EXPECT_EQ(1U, add_hashes.size()); + + // Make sure add prefixes are correct. + std::vector<SBAddPrefix> in_store_add_prefixes; + EXPECT_TRUE(store->GetAddPrefixes(&in_store_add_prefixes)); + ASSERT_EQ(3U, in_store_add_prefixes.size()); + EXPECT_EQ(kPrefix1, in_store_add_prefixes[0].prefix); + EXPECT_EQ(kAddChunk1, in_store_add_prefixes[0].chunk_id); + EXPECT_EQ(kPrefix2, in_store_add_prefixes[1].prefix); + EXPECT_EQ(kAddChunk1, in_store_add_prefixes[1].chunk_id); + EXPECT_EQ(kPrefix3, in_store_add_prefixes[2].prefix); + EXPECT_EQ(kAddChunk2, in_store_add_prefixes[2].chunk_id); } // Verify that the store looks like what results from LoadStore(), and @@ -230,6 +242,7 @@ void UpdateStore(SafeBrowsingStore* store) { EXPECT_TRUE(store->FinishUpdate(pending_adds, prefix_misses, &add_prefixes, &add_hashes)); EXPECT_EQ(2U, add_prefixes.size()); + EXPECT_EQ(0U, add_hashes.size()); } // Verify that the expected UpdateStore() data is present. @@ -250,6 +263,15 @@ void CheckStore(SafeBrowsingStore* store) { EXPECT_EQ(kSubChunk2, sub_chunks[1]); EXPECT_TRUE(store->CancelUpdate()); + + // Make sure add prefixes are correct. + std::vector<SBAddPrefix> add_prefixes; + EXPECT_TRUE(store->GetAddPrefixes(&add_prefixes)); + ASSERT_EQ(2U, add_prefixes.size()); + EXPECT_EQ(kPrefix2, add_prefixes[0].prefix); + EXPECT_EQ(kAddChunk1, add_prefixes[0].chunk_id); + EXPECT_EQ(kPrefix5, add_prefixes[1].prefix); + EXPECT_EQ(kAddChunk3, add_prefixes[1].chunk_id); } // Verify that the migration sequence works as expected in the diff --git a/chrome/browser/safe_browsing/safe_browsing_store_sqlite.cc b/chrome/browser/safe_browsing/safe_browsing_store_sqlite.cc index d687570..cf3369a 100644 --- a/chrome/browser/safe_browsing/safe_browsing_store_sqlite.cc +++ b/chrome/browser/safe_browsing/safe_browsing_store_sqlite.cc @@ -84,6 +84,15 @@ void SafeBrowsingStoreSqlite::Init(const FilePath& filename, corruption_callback_.reset(corruption_callback); } +bool SafeBrowsingStoreSqlite::GetAddPrefixes( + std::vector<SBAddPrefix>* add_prefixes) { + add_prefixes->clear(); + if (!Open()) return false; + bool ret = ReadAddPrefixes(add_prefixes); + Close(); + return ret; +} + bool SafeBrowsingStoreSqlite::OnCorruptDatabase() { if (corruption_callback_.get()) corruption_callback_->Run(); diff --git a/chrome/browser/safe_browsing/safe_browsing_store_sqlite.h b/chrome/browser/safe_browsing/safe_browsing_store_sqlite.h index 9093c4a..b577b08 100644 --- a/chrome/browser/safe_browsing/safe_browsing_store_sqlite.h +++ b/chrome/browser/safe_browsing/safe_browsing_store_sqlite.h @@ -32,6 +32,9 @@ class SafeBrowsingStoreSqlite : public SafeBrowsingStore { virtual bool BeginChunk() { return true; } + + // Get all Add prefixes out from the store. + virtual bool GetAddPrefixes(std::vector<SBAddPrefix>* add_prefixes); virtual bool WriteAddPrefix(int32 chunk_id, SBPrefix prefix) { const std::vector<SBAddPrefix> prefixes(1, SBAddPrefix(chunk_id, prefix)); return WriteAddPrefixes(prefixes); |