This source file includes following definitions.
- SetUpTestCase
- CheckPrefixes
- GetPrefixSetFile
- IncrementIntAt
- CleanChecksum
- ModifyAndCleanChecksum
- TEST_F
- TEST_F
- TEST_F
- TEST_F
- TEST_F
- TEST_F
- TEST_F
- TEST_F
- TEST_F
- TEST_F
- TEST_F
- TEST_F
- TEST_F
- TEST_F
- TEST_F
- TEST_F
- TEST_F
#include "chrome/browser/safe_browsing/prefix_set.h"
#include <algorithm>
#include <iterator>
#include "base/file_util.h"
#include "base/files/scoped_file.h"
#include "base/files/scoped_temp_dir.h"
#include "base/logging.h"
#include "base/md5.h"
#include "base/memory/scoped_ptr.h"
#include "base/rand_util.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "testing/platform_test.h"
namespace {
const SBPrefix kHighBitClear = 1000u * 1000u * 1000u;
const SBPrefix kHighBitSet = 3u * 1000u * 1000u * 1000u;
class PrefixSetTest : public PlatformTest {
protected:
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);
static void SetUpTestCase() {
for (size_t i = 0; i < 250; ++i) {
const uint32 base = static_cast<uint32>(base::RandUint64());
for (size_t j = 0; j < 10; ++j) {
const uint32 delta = static_cast<uint32>(base::RandUint64() & 0xFFFF);
const SBPrefix prefix = static_cast<SBPrefix>(base + delta);
shared_prefixes_.push_back(prefix);
}
}
const size_t count = shared_prefixes_.size();
for (size_t i = 0; i < count; ++i) {
const SBPrefix prefix = static_cast<SBPrefix>(base::RandUint64());
shared_prefixes_.push_back(prefix);
}
std::sort(shared_prefixes_.begin(), shared_prefixes_.end());
}
static void CheckPrefixes(const safe_browsing::PrefixSet& prefix_set,
const std::vector<SBPrefix> &prefixes) {
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()));
for (size_t i = 0; i < prefixes.size(); ++i) {
EXPECT_TRUE(prefix_set.Exists(prefixes[i]));
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));
}
}
bool GetPrefixSetFile(base::FilePath* filenamep) {
if (!temp_dir_.IsValid() && !temp_dir_.CreateUniqueTempDir())
return false;
base::FilePath filename = temp_dir_.path().AppendASCII("PrefixSetTest");
safe_browsing::PrefixSetBuilder builder(shared_prefixes_);
if (!builder.GetPrefixSet()->WriteFile(filename))
return false;
*filenamep = filename;
return true;
}
static void IncrementIntAt(FILE* fp, long offset, int inc) {
uint32 value = 0;
ASSERT_NE(-1, fseek(fp, offset, SEEK_SET));
ASSERT_EQ(1U, fread(&value, sizeof(value), 1, fp));
value += inc;
ASSERT_NE(-1, fseek(fp, offset, SEEK_SET));
ASSERT_EQ(1U, fwrite(&value, sizeof(value), 1, fp));
}
static void CleanChecksum(FILE* fp) {
base::MD5Context context;
base::MD5Init(&context);
ASSERT_NE(-1, fseek(fp, 0, SEEK_END));
long file_size = ftell(fp);
using base::MD5Digest;
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));
base::MD5Update(&context, base::StringPiece(buf, nitems));
digested_size += nitems;
}
ASSERT_EQ(digested_size, payload_size);
ASSERT_EQ(static_cast<long>(digested_size), ftell(fp));
base::MD5Digest new_digest;
base::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));
}
void ModifyAndCleanChecksum(const base::FilePath& filename, long offset,
int inc) {
int64 size_64;
ASSERT_TRUE(base::GetFileSize(filename, &size_64));
base::ScopedFILE file(base::OpenFile(filename, "r+b"));
IncrementIntAt(file.get(), offset, inc);
CleanChecksum(file.get());
file.reset();
int64 new_size_64;
ASSERT_TRUE(base::GetFileSize(filename, &new_size_64));
ASSERT_EQ(new_size_64, size_64);
}
static std::vector<SBPrefix> shared_prefixes_;
base::ScopedTempDir temp_dir_;
};
std::vector<SBPrefix> PrefixSetTest::shared_prefixes_;
TEST_F(PrefixSetTest, Baseline) {
safe_browsing::PrefixSetBuilder builder(shared_prefixes_);
CheckPrefixes(*builder.GetPrefixSet(), shared_prefixes_);
}
TEST_F(PrefixSetTest, Empty) {
const std::vector<SBPrefix> empty;
safe_browsing::PrefixSetBuilder builder(empty);
scoped_ptr<safe_browsing::PrefixSet> prefix_set = builder.GetPrefixSet();
for (size_t i = 0; i < shared_prefixes_.size(); ++i) {
EXPECT_FALSE(prefix_set->Exists(shared_prefixes_[i]));
}
}
TEST_F(PrefixSetTest, OneElement) {
const std::vector<SBPrefix> prefixes(100, 0u);
safe_browsing::PrefixSetBuilder builder(prefixes);
scoped_ptr<safe_browsing::PrefixSet> prefix_set = builder.GetPrefixSet();
EXPECT_FALSE(prefix_set->Exists(static_cast<SBPrefix>(-1)));
EXPECT_TRUE(prefix_set->Exists(prefixes[0]));
EXPECT_FALSE(prefix_set->Exists(1u));
std::vector<SBPrefix> prefixes_copy;
prefix_set->GetPrefixes(&prefixes_copy);
EXPECT_EQ(1U, prefixes_copy.size());
EXPECT_EQ(prefixes[0], prefixes_copy[0]);
}
TEST_F(PrefixSetTest, IntMinMax) {
std::vector<SBPrefix> prefixes;
prefixes.push_back(0x00000000);
prefixes.push_back(0x0000FFFF);
prefixes.push_back(0x7FFF0000);
prefixes.push_back(0x7FFFFFFF);
prefixes.push_back(0x80000000);
prefixes.push_back(0x8000FFFF);
prefixes.push_back(0xFFFF0000);
prefixes.push_back(0xFFFFFFFF);
std::sort(prefixes.begin(), prefixes.end());
safe_browsing::PrefixSetBuilder builder(prefixes);
scoped_ptr<safe_browsing::PrefixSet> prefix_set = builder.GetPrefixSet();
std::vector<SBPrefix> prefixes_copy;
prefix_set->GetPrefixes(&prefixes_copy);
ASSERT_EQ(prefixes_copy.size(), prefixes.size());
EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
prefixes_copy.begin()));
}
TEST_F(PrefixSetTest, AllBig) {
std::vector<SBPrefix> prefixes;
const unsigned kDelta = 10 * 1000 * 1000;
for (SBPrefix prefix = kHighBitSet;
prefix < kHighBitClear; prefix += kDelta) {
prefixes.push_back(prefix);
}
std::sort(prefixes.begin(), prefixes.end());
safe_browsing::PrefixSetBuilder builder(prefixes);
scoped_ptr<safe_browsing::PrefixSet> prefix_set = builder.GetPrefixSet();
std::vector<SBPrefix> prefixes_copy;
prefix_set->GetPrefixes(&prefixes_copy);
prefixes.erase(std::unique(prefixes.begin(), prefixes.end()), prefixes.end());
EXPECT_EQ(prefixes_copy.size(), prefixes.size());
EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
prefixes_copy.begin()));
}
TEST_F(PrefixSetTest, EdgeCases) {
std::vector<SBPrefix> prefixes;
SBPrefix prefix = kHighBitSet;
prefixes.push_back(prefix);
unsigned delta = 100 * 1000 * 1000;
for (int i = 0; i < 10; ++i) {
prefix += delta;
prefixes.push_back(prefix);
}
delta = 256 * 256 - 100;
for (int i = 0; i < 200; ++i) {
prefix += delta;
prefixes.push_back(prefix);
prefixes.push_back(prefix);
delta++;
}
delta = 256 * 256 - 1;
prefix = kHighBitClear - delta * 1000;
prefixes.push_back(prefix);
for (int i = 0; i < 1000; ++i) {
prefix += delta;
prefixes.push_back(prefix);
delta--;
}
std::sort(prefixes.begin(), prefixes.end());
safe_browsing::PrefixSetBuilder builder(prefixes);
scoped_ptr<safe_browsing::PrefixSet> prefix_set = builder.GetPrefixSet();
std::vector<SBPrefix> prefixes_copy;
prefix_set->GetPrefixes(&prefixes_copy);
prefixes.erase(std::unique(prefixes.begin(), prefixes.end()), prefixes.end());
EXPECT_EQ(prefixes_copy.size(), prefixes.size());
EXPECT_TRUE(std::equal(prefixes.begin(), prefixes.end(),
prefixes_copy.begin()));
EXPECT_FALSE(prefix_set->Exists(kHighBitSet - 100));
EXPECT_FALSE(prefix_set->Exists(kHighBitClear + 100));
for (size_t i = 0; i < prefixes.size(); ++i) {
EXPECT_TRUE(prefix_set->Exists(prefixes[i]));
EXPECT_FALSE(prefix_set->Exists(prefixes[i] - 1));
EXPECT_FALSE(prefix_set->Exists(prefixes[i] + 1));
}
}
TEST_F(PrefixSetTest, ReadWrite) {
base::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, shared_prefixes_);
}
{
std::vector<SBPrefix> prefixes;
prefixes.push_back(kHighBitClear);
prefixes.push_back(kHighBitSet);
safe_browsing::PrefixSetBuilder builder(prefixes);
ASSERT_TRUE(builder.GetPrefixSet()->WriteFile(filename));
scoped_ptr<safe_browsing::PrefixSet> prefix_set =
safe_browsing::PrefixSet::LoadFile(filename);
ASSERT_TRUE(prefix_set.get());
CheckPrefixes(*prefix_set, prefixes);
}
{
std::vector<SBPrefix> prefixes;
safe_browsing::PrefixSetBuilder builder(prefixes);
ASSERT_TRUE(builder.GetPrefixSet()->WriteFile(filename));
scoped_ptr<safe_browsing::PrefixSet> prefix_set =
safe_browsing::PrefixSet::LoadFile(filename);
ASSERT_TRUE(prefix_set.get());
CheckPrefixes(*prefix_set, prefixes);
}
}
TEST_F(PrefixSetTest, CorruptionHelpers) {
base::FilePath filename;
ASSERT_TRUE(GetPrefixSetFile(&filename));
base::ScopedFILE file(base::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());
file.reset(base::OpenFile(filename, "r+b"));
CleanChecksum(file.get());
file.reset();
prefix_set = safe_browsing::PrefixSet::LoadFile(filename);
ASSERT_TRUE(prefix_set.get());
}
TEST_F(PrefixSetTest, CorruptionMagic) {
base::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());
}
TEST_F(PrefixSetTest, CorruptionVersion) {
base::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());
}
TEST_F(PrefixSetTest, CorruptionIndexSize) {
base::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());
}
TEST_F(PrefixSetTest, CorruptionDeltasSize) {
base::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_F(PrefixSetTest, CorruptionPayload) {
base::FilePath filename;
ASSERT_TRUE(GetPrefixSetFile(&filename));
base::ScopedFILE file(base::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_F(PrefixSetTest, CorruptionDigest) {
base::FilePath filename;
ASSERT_TRUE(GetPrefixSetFile(&filename));
int64 size_64;
ASSERT_TRUE(base::GetFileSize(filename, &size_64));
base::ScopedFILE file(base::OpenFile(filename, "r+b"));
long digest_offset = static_cast<long>(size_64 - sizeof(base::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_F(PrefixSetTest, CorruptionExcess) {
base::FilePath filename;
ASSERT_TRUE(GetPrefixSetFile(&filename));
base::ScopedFILE file(base::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());
}
TEST_F(PrefixSetTest, SizeTRecovery) {
base::FilePath filename;
ASSERT_TRUE(GetPrefixSetFile(&filename));
base::ScopedFILE file(base::OpenFile(filename, "r+b"));
ASSERT_NE(-1, fseek(file.get(), sizeof(uint32) * 2, SEEK_SET));
uint32 val = 2;
ASSERT_EQ(sizeof(val), fwrite(&val, 1, sizeof(val), file.get()));
ASSERT_EQ(sizeof(val), fwrite(&val, 1, sizeof(val), file.get()));
std::pair<SBPrefix, uint64> item;
memset(&item, 0, sizeof(item));
item.first = 17;
item.second = 0;
ASSERT_EQ(sizeof(item), fwrite(&item, 1, sizeof(item), file.get()));
item.first = 100042;
item.second = 1;
ASSERT_EQ(sizeof(item), fwrite(&item, 1, sizeof(item), file.get()));
uint16 delta = 23;
ASSERT_EQ(sizeof(delta), fwrite(&delta, 1, sizeof(delta), file.get()));
ASSERT_EQ(sizeof(delta), fwrite(&delta, 1, sizeof(delta), file.get()));
base::MD5Digest dummy = { { 0 } };
ASSERT_EQ(sizeof(dummy), fwrite(&dummy, 1, sizeof(dummy), file.get()));
ASSERT_TRUE(base::TruncateFile(file.get()));
CleanChecksum(file.get());
file.reset();
scoped_ptr<safe_browsing::PrefixSet> prefix_set =
safe_browsing::PrefixSet::LoadFile(filename);
ASSERT_FALSE(prefix_set.get());
}
TEST_F(PrefixSetTest, ReadWriteSigned) {
base::FilePath filename;
ASSERT_TRUE(GetPrefixSetFile(&filename));
base::ScopedFILE file(base::OpenFile(filename, "r+b"));
ASSERT_NE(-1, fseek(file.get(), sizeof(uint32), SEEK_SET));
uint32 version = 1;
ASSERT_EQ(sizeof(version), fwrite(&version, 1, sizeof(version), file.get()));
uint32 val = 2;
ASSERT_EQ(sizeof(val), fwrite(&val, 1, sizeof(val), file.get()));
ASSERT_EQ(sizeof(val), fwrite(&val, 1, sizeof(val), file.get()));
std::pair<int32, uint32> item;
memset(&item, 0, sizeof(item));
item.first = -1000;
item.second = 0;
ASSERT_EQ(sizeof(item), fwrite(&item, 1, sizeof(item), file.get()));
item.first = 1000;
item.second = 1;
ASSERT_EQ(sizeof(item), fwrite(&item, 1, sizeof(item), file.get()));
uint16 delta = 23;
ASSERT_EQ(sizeof(delta), fwrite(&delta, 1, sizeof(delta), file.get()));
ASSERT_EQ(sizeof(delta), fwrite(&delta, 1, sizeof(delta), file.get()));
base::MD5Digest dummy = { { 0 } };
ASSERT_EQ(sizeof(dummy), fwrite(&dummy, 1, sizeof(dummy), file.get()));
ASSERT_TRUE(base::TruncateFile(file.get()));
CleanChecksum(file.get());
file.reset();
scoped_ptr<safe_browsing::PrefixSet> prefix_set =
safe_browsing::PrefixSet::LoadFile(filename);
ASSERT_TRUE(prefix_set.get());
EXPECT_TRUE(prefix_set->Exists(1000u));
EXPECT_TRUE(prefix_set->Exists(1023u));
EXPECT_TRUE(prefix_set->Exists(static_cast<uint32>(-1000)));
EXPECT_TRUE(prefix_set->Exists(static_cast<uint32>(-1000 + 23)));
std::vector<SBPrefix> prefixes_copy;
prefix_set->GetPrefixes(&prefixes_copy);
EXPECT_EQ(prefixes_copy.size(), 4u);
EXPECT_EQ(prefixes_copy[0], 1000u);
EXPECT_EQ(prefixes_copy[1], 1023u);
EXPECT_EQ(prefixes_copy[2], static_cast<uint32>(-1000));
EXPECT_EQ(prefixes_copy[3], static_cast<uint32>(-1000 + 23));
}
}