Index: chrome/browser/safe_browsing/prefix_set_unittest.cc |
diff --git a/chrome/browser/safe_browsing/prefix_set_unittest.cc b/chrome/browser/safe_browsing/prefix_set_unittest.cc |
index f7be5dad77b8307ac91797802484fe6ed159c59b..154711720350ba24a7f4bbb8413922ce4f0ad91f 100644 |
--- a/chrome/browser/safe_browsing/prefix_set_unittest.cc |
+++ b/chrome/browser/safe_browsing/prefix_set_unittest.cc |
@@ -6,58 +6,161 @@ |
#include <algorithm> |
+#include "base/file_util.h" |
#include "base/logging.h" |
+#include "base/md5.h" |
#include "base/rand_util.h" |
+#include "base/scoped_ptr.h" |
+#include "base/scoped_temp_dir.h" |
#include "testing/gtest/include/gtest/gtest.h" |
+#include "testing/platform_test.h" |
namespace { |
-SBPrefix GenPrefix() { |
- return static_cast<SBPrefix>(base::RandUint64()); |
-} |
+class PrefixSetTest : public PlatformTest { |
+ protected: |
+ // Constants for the v1 format. |
+ static const size_t kMagicOffset = 0 * sizeof(uint32); |
+ static const size_t kVersionOffset = 1 * sizeof(uint32); |
+ static const size_t kIndexSizeOffset = 2 * sizeof(uint32); |
+ static const size_t kDeltasSizeOffset = 3 * sizeof(uint32); |
+ static const size_t kPayloadOffset = 4 * sizeof(uint32); |
-// Test that a small sparse random input works. |
-TEST(PrefixSetTest, Baseline) { |
- std::vector<SBPrefix> prefixes; |
+ // Generate a set of random prefixes to share between tests. For |
+ // most tests this generation was a large fraction of the test time. |
+ static void SetUpTestCase() { |
+ for (size_t i = 0; i < 50000; ++i) { |
+ const SBPrefix prefix = static_cast<SBPrefix>(base::RandUint64()); |
+ shared_prefixes_.push_back(prefix); |
+ } |
+ |
+ // Sort for use with PrefixSet constructor. |
+ std::sort(shared_prefixes_.begin(), shared_prefixes_.end()); |
+ } |
+ |
+ // Check that all elements of |prefixes| are in |prefix_set|, and |
+ // that nearby elements are not (for lack of a more sensible set of |
+ // items to check for absence). |
+ static void CheckPrefixes(safe_browsing::PrefixSet* prefix_set, |
+ const std::vector<SBPrefix> &prefixes) { |
+ // The set can generate the prefixes it believes it has, so that's |
+ // a good starting point. |
+ std::set<SBPrefix> check(prefixes.begin(), prefixes.end()); |
+ std::vector<SBPrefix> prefixes_copy; |
+ prefix_set->GetPrefixes(&prefixes_copy); |
+ EXPECT_EQ(prefixes_copy.size(), check.size()); |
+ EXPECT_TRUE(std::equal(check.begin(), check.end(), prefixes_copy.begin())); |
- static const size_t kCount = 50000; |
+ for (size_t i = 0; i < prefixes.size(); ++i) { |
+ EXPECT_TRUE(prefix_set->Exists(prefixes[i])); |
- for (size_t i = 0; i < kCount; ++i) { |
- prefixes.push_back(GenPrefix()); |
+ const SBPrefix left_sibling = prefixes[i] - 1; |
+ if (check.count(left_sibling) == 0) |
+ EXPECT_FALSE(prefix_set->Exists(left_sibling)); |
+ |
+ const SBPrefix right_sibling = prefixes[i] + 1; |
+ if (check.count(right_sibling) == 0) |
+ EXPECT_FALSE(prefix_set->Exists(right_sibling)); |
+ } |
} |
- std::sort(prefixes.begin(), prefixes.end()); |
- safe_browsing::PrefixSet prefix_set(prefixes); |
+ // Generate a |PrefixSet| file from |shared_prefixes_|, store it in |
+ // a temporary file, and return the filename in |filenamep|. |
+ // Returns |true| on success. |
+ bool GetPrefixSetFile(FilePath* filenamep) { |
+ if (!temp_dir_.IsValid() && !temp_dir_.CreateUniqueTempDir()) |
+ return false; |
- // Check that |GetPrefixes()| returns exactly the same set of |
- // prefixes as was passed in. |
- std::set<SBPrefix> check(prefixes.begin(), prefixes.end()); |
- std::vector<SBPrefix> prefixes_copy; |
- prefix_set.GetPrefixes(&prefixes_copy); |
- EXPECT_EQ(prefixes_copy.size(), check.size()); |
- EXPECT_TRUE(std::equal(check.begin(), check.end(), prefixes_copy.begin())); |
+ FilePath filename = temp_dir_.path().AppendASCII("PrefixSetTest"); |
- // Check that the set flags all of the inputs, and also check items |
- // just above and below the inputs to make sure they aren't there. |
- for (size_t i = 0; i < prefixes.size(); ++i) { |
- EXPECT_TRUE(prefix_set.Exists(prefixes[i])); |
+ safe_browsing::PrefixSet prefix_set(shared_prefixes_); |
+ if (!prefix_set.WriteFile(filename)) |
+ return false; |
+ |
+ *filenamep = filename; |
+ return true; |
+ } |
+ |
+ // Helper function to read the int32 value at |offset|, increment it |
+ // by |inc|, and write it back in place. |fp| should be opened in |
+ // r+ mode. |
+ static void IncrementIntAt(FILE* fp, long offset, int inc) { |
+ int32 value = 0; |
- const SBPrefix left_sibling = prefixes[i] - 1; |
- if (check.count(left_sibling) == 0) |
- EXPECT_FALSE(prefix_set.Exists(left_sibling)); |
+ ASSERT_NE(-1, fseek(fp, offset, SEEK_SET)); |
+ ASSERT_EQ(1U, fread(&value, sizeof(value), 1, fp)); |
- const SBPrefix right_sibling = prefixes[i] + 1; |
- if (check.count(right_sibling) == 0) |
- EXPECT_FALSE(prefix_set.Exists(right_sibling)); |
+ value += inc; |
+ |
+ ASSERT_NE(-1, fseek(fp, offset, SEEK_SET)); |
+ ASSERT_EQ(1U, fwrite(&value, sizeof(value), 1, fp)); |
+ } |
+ |
+ // Helper function to re-generated |fp|'s checksum to be correct for |
+ // the file's contents. |fp| should be opened in r+ mode. |
+ static void CleanChecksum(FILE* fp) { |
+ MD5Context context; |
+ MD5Init(&context); |
+ |
+ ASSERT_NE(-1, fseek(fp, 0, SEEK_END)); |
+ long file_size = ftell(fp); |
+ |
+ size_t payload_size = static_cast<size_t>(file_size) - sizeof(MD5Digest); |
+ size_t digested_size = 0; |
+ ASSERT_NE(-1, fseek(fp, 0, SEEK_SET)); |
+ while (digested_size < payload_size) { |
+ char buf[1024]; |
+ size_t nitems = std::min(payload_size - digested_size, sizeof(buf)); |
+ ASSERT_EQ(nitems, fread(buf, 1, nitems, fp)); |
+ MD5Update(&context, &buf, nitems); |
+ digested_size += nitems; |
+ } |
+ ASSERT_EQ(digested_size, payload_size); |
+ ASSERT_EQ(static_cast<long>(digested_size), ftell(fp)); |
+ |
+ MD5Digest new_digest; |
+ MD5Final(&new_digest, &context); |
+ ASSERT_NE(-1, fseek(fp, digested_size, SEEK_SET)); |
+ ASSERT_EQ(1U, fwrite(&new_digest, sizeof(new_digest), 1, fp)); |
+ ASSERT_EQ(file_size, ftell(fp)); |
} |
+ |
+ // Open |filename| and increment the int32 at |offset| by |inc|. |
+ // Then re-generate the checksum to account for the new contents. |
+ void ModifyAndCleanChecksum(const FilePath& filename, long offset, int inc) { |
+ int64 size_64; |
+ ASSERT_TRUE(file_util::GetFileSize(filename, &size_64)); |
+ |
+ file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b")); |
+ IncrementIntAt(file.get(), offset, inc); |
+ CleanChecksum(file.get()); |
+ file.reset(); |
+ |
+ int64 new_size_64; |
+ ASSERT_TRUE(file_util::GetFileSize(filename, &new_size_64)); |
+ ASSERT_EQ(new_size_64, size_64); |
+ } |
+ |
+ // Tests should not modify this shared resource. |
+ static std::vector<SBPrefix> shared_prefixes_; |
+ |
+ ScopedTempDir temp_dir_; |
+}; |
+ |
+std::vector<SBPrefix> PrefixSetTest::shared_prefixes_; |
+ |
+// Test that a small sparse random input works. |
+TEST_F(PrefixSetTest, Baseline) { |
+ safe_browsing::PrefixSet prefix_set(shared_prefixes_); |
+ CheckPrefixes(&prefix_set, shared_prefixes_); |
} |
// Test that the empty set doesn't appear to have anything in it. |
-TEST(PrefixSetTest, Empty) { |
- std::vector<SBPrefix> prefixes; |
- safe_browsing::PrefixSet prefix_set(prefixes); |
- for (size_t i = 0; i < 500; ++i) { |
- EXPECT_FALSE(prefix_set.Exists(GenPrefix())); |
+TEST_F(PrefixSetTest, Empty) { |
+ const std::vector<SBPrefix> empty; |
+ safe_browsing::PrefixSet prefix_set(empty); |
+ for (size_t i = 0; i < shared_prefixes_.size(); ++i) { |
+ EXPECT_FALSE(prefix_set.Exists(shared_prefixes_[i])); |
} |
} |
@@ -67,7 +170,7 @@ TEST(PrefixSetTest, Empty) { |
// deltas above and below 2^16, and make sure they're all present. |
// Create a very long sequence with deltas below 2^16 to test crossing |
// |kMaxRun|. |
-TEST(PrefixSetTest, EdgeCases) { |
+TEST_F(PrefixSetTest, EdgeCases) { |
std::vector<SBPrefix> prefixes; |
const SBPrefix kVeryPositive = 1000 * 1000 * 1000; |
@@ -132,4 +235,132 @@ TEST(PrefixSetTest, EdgeCases) { |
} |
} |
+// Similar to Baseline test, but write the set out to a file and read |
+// it back in before testing. |
+TEST_F(PrefixSetTest, ReadWrite) { |
+ FilePath filename; |
+ ASSERT_TRUE(GetPrefixSetFile(&filename)); |
+ |
+ scoped_ptr<safe_browsing::PrefixSet> |
+ prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
+ ASSERT_TRUE(prefix_set.get()); |
+ |
+ CheckPrefixes(prefix_set.get(), shared_prefixes_); |
+} |
+ |
+// Check that |CleanChecksum()| makes an acceptable checksum. |
+TEST_F(PrefixSetTest, CorruptionHelpers) { |
+ FilePath filename; |
+ ASSERT_TRUE(GetPrefixSetFile(&filename)); |
+ |
+ // This will modify data in |index_|, which will fail the digest check. |
+ file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b")); |
+ IncrementIntAt(file.get(), kPayloadOffset, 1); |
+ file.reset(); |
+ scoped_ptr<safe_browsing::PrefixSet> |
+ prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
+ ASSERT_FALSE(prefix_set.get()); |
+ |
+ // Fix up the checksum and it will read successfully (though the |
+ // data will be wrong). |
+ file.reset(file_util::OpenFile(filename, "r+b")); |
+ CleanChecksum(file.get()); |
+ file.reset(); |
+ prefix_set.reset(safe_browsing::PrefixSet::LoadFile(filename)); |
+ ASSERT_TRUE(prefix_set.get()); |
+} |
+ |
+// Bad |index_| size is caught by the sanity check. |
+TEST_F(PrefixSetTest, CorruptionMagic) { |
+ FilePath filename; |
+ ASSERT_TRUE(GetPrefixSetFile(&filename)); |
+ |
+ ASSERT_NO_FATAL_FAILURE( |
+ ModifyAndCleanChecksum(filename, kMagicOffset, 1)); |
+ scoped_ptr<safe_browsing::PrefixSet> |
+ prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
+ ASSERT_FALSE(prefix_set.get()); |
+} |
+ |
+// Bad |index_| size is caught by the sanity check. |
+TEST_F(PrefixSetTest, CorruptionVersion) { |
+ FilePath filename; |
+ ASSERT_TRUE(GetPrefixSetFile(&filename)); |
+ |
+ ASSERT_NO_FATAL_FAILURE( |
+ ModifyAndCleanChecksum(filename, kVersionOffset, 1)); |
+ scoped_ptr<safe_browsing::PrefixSet> |
+ prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
+ ASSERT_FALSE(prefix_set.get()); |
+} |
+ |
+// Bad |index_| size is caught by the sanity check. |
+TEST_F(PrefixSetTest, CorruptionIndexSize) { |
+ FilePath filename; |
+ ASSERT_TRUE(GetPrefixSetFile(&filename)); |
+ |
+ ASSERT_NO_FATAL_FAILURE( |
+ ModifyAndCleanChecksum(filename, kIndexSizeOffset, 1)); |
+ scoped_ptr<safe_browsing::PrefixSet> |
+ prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
+ ASSERT_FALSE(prefix_set.get()); |
+} |
+ |
+// Bad |deltas_| size is caught by the sanity check. |
+TEST_F(PrefixSetTest, CorruptionDeltasSize) { |
+ FilePath filename; |
+ ASSERT_TRUE(GetPrefixSetFile(&filename)); |
+ |
+ ASSERT_NO_FATAL_FAILURE( |
+ ModifyAndCleanChecksum(filename, kDeltasSizeOffset, 1)); |
+ scoped_ptr<safe_browsing::PrefixSet> |
+ prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
+ ASSERT_FALSE(prefix_set.get()); |
+} |
+ |
+// Test that the digest catches corruption in the middle of the file |
+// (in the payload between the header and the digest). |
+TEST_F(PrefixSetTest, CorruptionPayload) { |
+ FilePath filename; |
+ ASSERT_TRUE(GetPrefixSetFile(&filename)); |
+ |
+ file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b")); |
+ ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), 666, 1)); |
+ file.reset(); |
+ scoped_ptr<safe_browsing::PrefixSet> |
+ prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
+ ASSERT_FALSE(prefix_set.get()); |
+} |
+ |
+// Test corruption in the digest itself. |
+TEST_F(PrefixSetTest, CorruptionDigest) { |
+ FilePath filename; |
+ ASSERT_TRUE(GetPrefixSetFile(&filename)); |
+ |
+ int64 size_64; |
+ ASSERT_TRUE(file_util::GetFileSize(filename, &size_64)); |
+ file_util::ScopedFILE file(file_util::OpenFile(filename, "r+b")); |
+ long digest_offset = static_cast<long>(size_64 - sizeof(MD5Digest)); |
+ ASSERT_NO_FATAL_FAILURE(IncrementIntAt(file.get(), digest_offset, 1)); |
+ file.reset(); |
+ scoped_ptr<safe_browsing::PrefixSet> |
+ prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
+ ASSERT_FALSE(prefix_set.get()); |
+} |
+ |
+// Test excess data after the digest (fails the size test). |
+TEST_F(PrefixSetTest, CorruptionExcess) { |
+ FilePath filename; |
+ ASSERT_TRUE(GetPrefixSetFile(&filename)); |
+ |
+ // Add some junk to the trunk. |
+ file_util::ScopedFILE file(file_util::OpenFile(filename, "ab")); |
+ const char buf[] = "im in ur base, killing ur d00dz."; |
+ ASSERT_EQ(strlen(buf), fwrite(buf, 1, strlen(buf), file.get())); |
+ file.reset(); |
+ scoped_ptr<safe_browsing::PrefixSet> |
+ prefix_set(safe_browsing::PrefixSet::LoadFile(filename)); |
+ ASSERT_FALSE(prefix_set.get()); |
+} |
+ |
} // namespace |