blob: 0b84790dd655aa643c9cba50604bab0d0628a145 [file] [log] [blame]
// Copyright 2016 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#include <iostream>
#include <memory>
#include <unordered_map>
#include "base/at_exit.h"
#include "base/command_line.h"
#include "base/files/file_path.h"
#include "base/format_macros.h"
#include "base/logging.h"
#include "base/macros.h"
#include "base/md5.h"
#include "base/message_loop/message_loop.h"
#include "base/run_loop.h"
#include "base/strings/string_number_conversions.h"
#include "base/strings/string_piece.h"
#include "base/strings/stringprintf.h"
#include "base/task_scheduler/task_scheduler.h"
#include "net/base/io_buffer.h"
#include "net/base/test_completion_callback.h"
#include "net/disk_cache/disk_cache.h"
#include "net/http/http_cache.h"
#include "net/http/http_response_headers.h"
#include "net/http/http_util.h"
using disk_cache::Backend;
using disk_cache::Entry;
namespace {
struct EntryData {
std::string url;
std::string mime_type;
int size;
};
constexpr int kResponseInfoIndex = 0;
constexpr int kResponseContentIndex = 1;
const char* const kCommandNames[] = {
"stop", "get_size", "list_keys", "get_stream_for_key",
"delete_stream", "delete_key", "update_raw_headers", "list_dups",
};
// Prints the command line help.
void PrintHelp() {
std::cout << "cachetool <cache_path> <cache_backend_type> <subcommand> "
<< std::endl
<< std::endl;
std::cout << "Available cache backend types: simple, blockfile" << std::endl;
std::cout << "Available subcommands:" << std::endl;
std::cout << " batch: Starts cachetool to process serialized commands "
<< "passed down by the standard input and return commands output "
<< "in the stdout until the stop command is received." << std::endl;
std::cout << " delete_key <key>: Delete key from cache." << std::endl;
std::cout << " delete_stream <key> <index>: Delete a particular stream of a"
<< " given key." << std::endl;
std::cout << " get_size: Calculate the total size of the cache in bytes."
<< std::endl;
std::cout << " get_stream <key> <index>: Print a particular stream for a"
<< " given key." << std::endl;
std::cout << " list_keys: List all keys in the cache." << std::endl;
std::cout << " list_dups: List all resources with duplicate bodies in the "
<< "cache." << std::endl;
std::cout << " update_raw_headers <key>: Update stdin as the key's raw "
<< "response headers." << std::endl;
std::cout << " stop: Verify that the cache can be opened and return, "
<< "confirming the cache exists and is of the right type."
<< std::endl;
std::cout << "Expected values of <index> are:" << std::endl;
std::cout << " 0 (HTTP response headers)" << std::endl;
std::cout << " 1 (transport encoded content)" << std::endl;
std::cout << " 2 (compiled content)" << std::endl;
}
// Generic command input/output.
class CommandMarshal {
public:
explicit CommandMarshal(Backend* cache_backend)
: command_failed_(false), cache_backend_(cache_backend) {}
virtual ~CommandMarshal() {}
// Reads the next command's name to execute.
virtual std::string ReadCommandName() = 0;
// Reads the next parameter as an integer.
virtual int ReadInt() = 0;
// Reads the next parameter as stream index.
int ReadStreamIndex() {
if (has_failed())
return -1;
int index = ReadInt();
if (index < 0 || index > 2) {
ReturnFailure("Invalid stream index.");
return -1;
}
return index;
}
// Reads the next parameter as an string.
virtual std::string ReadString() = 0;
// Reads the next parameter from stdin as string.
virtual std::string ReadBufferedString() = 0;
// Communicates back an integer.
virtual void ReturnInt(int integer) = 0;
// Communicates back a string.
virtual void ReturnString(const std::string& string) = 0;
// Communicates back a buffer.
virtual void ReturnBuffer(net::GrowableIOBuffer* buffer) = 0;
// Communicates back command failure.
virtual void ReturnFailure(const std::string& error_msg) = 0;
// Communicates back command success.
virtual void ReturnSuccess() { DCHECK(!command_failed_); }
// Returns whether the command has failed.
inline bool has_failed() { return command_failed_; }
// Returns the opened cache backend.
Backend* cache_backend() { return cache_backend_; }
protected:
bool command_failed_;
Backend* const cache_backend_;
};
// Command line input/output that is user readable.
class ProgramArgumentCommandMarshal final : public CommandMarshal {
public:
ProgramArgumentCommandMarshal(Backend* cache_backend,
base::CommandLine::StringVector args)
: CommandMarshal(cache_backend), command_line_args_(args), args_id_(0) {}
// Implements CommandMarshal.
std::string ReadCommandName() override {
if (args_id_ == 0)
return ReadString();
else if (args_id_ == command_line_args_.size())
return "stop";
else if (!has_failed())
ReturnFailure("Command line arguments to long.");
return "";
}
// Implements CommandMarshal.
int ReadInt() override {
std::string interger_str = ReadString();
int interger = -1;
if (!base::StringToInt(interger_str, &interger)) {
ReturnFailure("Couldn't parse integer.");
return 0;
}
return interger;
}
// Implements CommandMarshal.
std::string ReadString() override {
if (args_id_ < command_line_args_.size())
return command_line_args_[args_id_++];
if (!has_failed())
ReturnFailure("Command line arguments to short.");
return "";
}
// Implements CommandMarshal.
std::string ReadBufferedString() override {
std::ostringstream raw_headers_stream;
for (std::string line; std::getline(std::cin, line);)
raw_headers_stream << line << std::endl;
return raw_headers_stream.str();
}
// Implements CommandMarshal.
void ReturnInt(int integer) override {
DCHECK(!has_failed());
std::cout << integer << std::endl;
}
// Implements CommandMarshal.
void ReturnString(const std::string& string) override {
DCHECK(!has_failed());
std::cout << string << std::endl;
}
// Implements CommandMarshal.
void ReturnBuffer(net::GrowableIOBuffer* buffer) override {
DCHECK(!has_failed());
std::cout.write(buffer->data(), buffer->offset());
}
// Implements CommandMarshal.
void ReturnFailure(const std::string& error_msg) override {
DCHECK(!has_failed());
std::cerr << error_msg << std::endl;
command_failed_ = true;
}
private:
const base::CommandLine::StringVector command_line_args_;
size_t args_id_;
};
// Online command input/output that receives pickled commands from stdin and
// returns their results back in stdout. Send the stop command to properly exit
// cachetool's main loop.
class StreamCommandMarshal final : public CommandMarshal {
public:
explicit StreamCommandMarshal(Backend* cache_backend)
: CommandMarshal(cache_backend) {}
// Implements CommandMarshal.
std::string ReadCommandName() override {
if (has_failed())
return "";
std::cout.flush();
size_t command_id = static_cast<size_t>(std::cin.get());
if (command_id >= arraysize(kCommandNames)) {
ReturnFailure("Unknown command.");
return "";
}
return kCommandNames[command_id];
}
// Implements CommandMarshal.
int ReadInt() override {
if (has_failed())
return -1;
int integer = -1;
std::cin.read(reinterpret_cast<char*>(&integer), sizeof(integer));
return integer;
}
// Implements CommandMarshal.
std::string ReadString() override {
if (has_failed())
return "";
int string_size = ReadInt();
if (string_size <= 0) {
if (string_size < 0)
ReturnFailure("Size of string is negative.");
return "";
}
std::vector<char> tmp_buffer(string_size + 1);
std::cin.read(&tmp_buffer[0], string_size);
tmp_buffer[string_size] = 0;
return std::string(&tmp_buffer[0], string_size);
}
// Implements CommandMarshal.
std::string ReadBufferedString() override { return ReadString(); }
// Implements CommandMarshal.
void ReturnInt(int integer) override {
DCHECK(!command_failed_);
std::cout.write(reinterpret_cast<char*>(&integer), sizeof(integer));
}
// Implements CommandMarshal.
void ReturnString(const std::string& string) override {
ReturnInt(string.size());
std::cout.write(string.c_str(), string.size());
}
// Implements CommandMarshal.
void ReturnBuffer(net::GrowableIOBuffer* buffer) override {
ReturnInt(buffer->offset());
std::cout.write(buffer->StartOfBuffer(), buffer->offset());
}
// Implements CommandMarshal.
void ReturnFailure(const std::string& error_msg) override {
ReturnString(error_msg);
command_failed_ = true;
}
// Implements CommandMarshal.
void ReturnSuccess() override { ReturnInt(0); }
};
// Gets the cache's size.
void GetSize(CommandMarshal* command_marshal) {
net::TestCompletionCallback cb;
int rv = command_marshal->cache_backend()->CalculateSizeOfAllEntries(
cb.callback());
rv = cb.GetResult(rv);
if (rv < 0)
return command_marshal->ReturnFailure("Couldn't get cache size.");
command_marshal->ReturnSuccess();
command_marshal->ReturnInt(rv);
}
// Prints all of a cache's keys to stdout.
bool ListKeys(CommandMarshal* command_marshal) {
std::unique_ptr<Backend::Iterator> entry_iterator =
command_marshal->cache_backend()->CreateIterator();
Entry* entry = nullptr;
net::TestCompletionCallback cb;
int rv = entry_iterator->OpenNextEntry(&entry, cb.callback());
command_marshal->ReturnSuccess();
while (cb.GetResult(rv) == net::OK) {
std::string url = entry->GetKey();
command_marshal->ReturnString(url);
entry->Close();
entry = nullptr;
rv = entry_iterator->OpenNextEntry(&entry, cb.callback());
}
command_marshal->ReturnString("");
return true;
}
bool GetResponseInfoForEntry(disk_cache::Entry* entry,
net::HttpResponseInfo* response_info) {
int size = entry->GetDataSize(kResponseInfoIndex);
if (size == 0)
return false;
scoped_refptr<net::IOBuffer> buffer = new net::IOBufferWithSize(size);
net::TestCompletionCallback cb;
int bytes_read = 0;
while (true) {
int rv = entry->ReadData(kResponseInfoIndex, bytes_read, buffer.get(), size,
cb.callback());
rv = cb.GetResult(rv);
if (rv < 0) {
entry->Close();
return false;
}
if (rv == 0) {
bool truncated_response_info = false;
net::HttpCache::ParseResponseInfo(buffer->data(), size, response_info,
&truncated_response_info);
return !truncated_response_info;
}
bytes_read += rv;
}
NOTREACHED();
return false;
}
std::string GetMD5ForResponseBody(disk_cache::Entry* entry) {
if (entry->GetDataSize(kResponseContentIndex) == 0)
return "";
const int kInitBufferSize = 80 * 1024;
scoped_refptr<net::IOBuffer> buffer =
new net::IOBufferWithSize(kInitBufferSize);
net::TestCompletionCallback cb;
base::MD5Context ctx;
base::MD5Init(&ctx);
int bytes_read = 0;
while (true) {
int rv = entry->ReadData(kResponseContentIndex, bytes_read, buffer.get(),
kInitBufferSize, cb.callback());
rv = cb.GetResult(rv);
if (rv < 0) {
entry->Close();
return "";
}
if (rv == 0) {
base::MD5Digest digest;
base::MD5Final(&digest, &ctx);
return base::MD5DigestToBase16(digest);
}
bytes_read += rv;
MD5Update(&ctx, base::StringPiece(buffer->data(), rv));
}
NOTREACHED();
return "";
}
void ListDups(CommandMarshal* command_marshal) {
std::unique_ptr<Backend::Iterator> entry_iterator =
command_marshal->cache_backend()->CreateIterator();
Entry* entry = nullptr;
net::TestCompletionCallback cb;
int rv = entry_iterator->OpenNextEntry(&entry, cb.callback());
command_marshal->ReturnSuccess();
std::unordered_map<std::string, std::vector<EntryData>> md5_entries;
int total_entries = 0;
while (cb.GetResult(rv) == net::OK) {
total_entries += 1;
net::HttpResponseInfo response_info;
if (!GetResponseInfoForEntry(entry, &response_info)) {
entry->Close();
entry = nullptr;
rv = entry_iterator->OpenNextEntry(&entry, cb.callback());
continue;
}
std::string hash = GetMD5ForResponseBody(entry);
if (hash.empty()) {
// Sparse entries and empty bodies are skipped.
entry->Close();
entry = nullptr;
rv = entry_iterator->OpenNextEntry(&entry, cb.callback());
continue;
}
EntryData entry_data;
entry_data.url = entry->GetKey();
entry_data.size = entry->GetDataSize(kResponseContentIndex);
if (response_info.headers)
response_info.headers->GetMimeType(&entry_data.mime_type);
auto iter = md5_entries.find(hash);
if (iter == md5_entries.end())
md5_entries.insert(
std::make_pair(hash, std::vector<EntryData>{entry_data}));
else
iter->second.push_back(entry_data);
entry->Close();
entry = nullptr;
rv = entry_iterator->OpenNextEntry(&entry, cb.callback());
}
// Print the duplicates and collect stats.
int total_duped_entries = 0;
int64_t total_duped_bytes = 0u;
for (const auto& hash_and_entries : md5_entries) {
if (hash_and_entries.second.size() == 1)
continue;
int dups = hash_and_entries.second.size() - 1;
total_duped_entries += dups;
total_duped_bytes += hash_and_entries.second[0].size * dups;
for (const auto& entry : hash_and_entries.second) {
std::string out = base::StringPrintf(
"%d, %s, %s", entry.size, entry.url.c_str(), entry.mime_type.c_str());
command_marshal->ReturnString(out);
}
}
// Print the stats.
rv = command_marshal->cache_backend()->CalculateSizeOfAllEntries(
cb.callback());
rv = cb.GetResult(rv);
LOG(ERROR) << "Wasted bytes = " << total_duped_bytes;
LOG(ERROR) << "Wasted entries = " << total_duped_entries;
LOG(ERROR) << "Total entries = " << total_entries;
LOG(ERROR) << "Cache size = " << rv;
LOG(ERROR) << "Percentage of cache wasted = " << total_duped_bytes * 100 / rv;
}
// Gets a key's stream to a buffer.
scoped_refptr<net::GrowableIOBuffer> GetStreamForKeyBuffer(
CommandMarshal* command_marshal,
const std::string& key,
int index) {
DCHECK(!command_marshal->has_failed());
Entry* cache_entry;
net::TestCompletionCallback cb;
int rv = command_marshal->cache_backend()->OpenEntry(key, &cache_entry,
cb.callback());
if (cb.GetResult(rv) != net::OK) {
command_marshal->ReturnFailure("Couldn't find key's entry.");
return nullptr;
}
const int kInitBufferSize = 8192;
scoped_refptr<net::GrowableIOBuffer> buffer(new net::GrowableIOBuffer());
buffer->SetCapacity(kInitBufferSize);
while (true) {
rv = cache_entry->ReadData(index, buffer->offset(), buffer.get(),
buffer->capacity() - buffer->offset(),
cb.callback());
rv = cb.GetResult(rv);
if (rv < 0) {
cache_entry->Close();
command_marshal->ReturnFailure("Stream read error.");
return nullptr;
}
buffer->set_offset(buffer->offset() + rv);
if (rv == 0)
break;
buffer->SetCapacity(buffer->offset() * 2);
}
cache_entry->Close();
return buffer;
}
// Prints a key's stream to stdout.
void GetStreamForKey(CommandMarshal* command_marshal) {
std::string key = command_marshal->ReadString();
int index = command_marshal->ReadInt();
if (command_marshal->has_failed())
return;
scoped_refptr<net::GrowableIOBuffer> buffer(
GetStreamForKeyBuffer(command_marshal, key, index));
if (command_marshal->has_failed())
return;
if (index == kResponseInfoIndex) {
net::HttpResponseInfo response_info;
bool truncated_response_info = false;
net::HttpCache::ParseResponseInfo(buffer->StartOfBuffer(), buffer->offset(),
&response_info, &truncated_response_info);
if (truncated_response_info)
return command_marshal->ReturnFailure("Truncated HTTP response.");
command_marshal->ReturnSuccess();
command_marshal->ReturnString(
net::HttpUtil::ConvertHeadersBackToHTTPResponse(
response_info.headers->raw_headers()));
} else {
command_marshal->ReturnSuccess();
command_marshal->ReturnBuffer(buffer.get());
}
}
// Sets stdin as the key's raw response headers.
void UpdateRawResponseHeaders(CommandMarshal* command_marshal) {
std::string key = command_marshal->ReadString();
std::string raw_headers = command_marshal->ReadBufferedString();
if (command_marshal->has_failed())
return;
scoped_refptr<net::GrowableIOBuffer> buffer(
GetStreamForKeyBuffer(command_marshal, key, kResponseInfoIndex));
if (command_marshal->has_failed())
return;
net::HttpResponseInfo response_info;
bool truncated_response_info = false;
net::HttpCache::ParseResponseInfo(buffer->StartOfBuffer(), buffer->offset(),
&response_info, &truncated_response_info);
if (truncated_response_info)
return command_marshal->ReturnFailure("Truncated HTTP response.");
response_info.headers = new net::HttpResponseHeaders(raw_headers);
scoped_refptr<net::PickledIOBuffer> data(new net::PickledIOBuffer());
response_info.Persist(data->pickle(), false, false);
data->Done();
Entry* cache_entry;
net::TestCompletionCallback cb;
int rv = command_marshal->cache_backend()->OpenEntry(key, &cache_entry,
cb.callback());
CHECK(cb.GetResult(rv) == net::OK);
int data_len = data->pickle()->size();
rv = cache_entry->WriteData(kResponseInfoIndex, 0, data.get(), data_len,
cb.callback(), true);
if (cb.GetResult(rv) != data_len)
return command_marshal->ReturnFailure("Couldn't write headers.");
command_marshal->ReturnSuccess();
cache_entry->Close();
}
// Deletes a specified key stream from the cache.
void DeleteStreamForKey(CommandMarshal* command_marshal) {
std::string key = command_marshal->ReadString();
int index = command_marshal->ReadInt();
if (command_marshal->has_failed())
return;
Entry* cache_entry;
net::TestCompletionCallback cb;
int rv = command_marshal->cache_backend()->OpenEntry(key, &cache_entry,
cb.callback());
if (cb.GetResult(rv) != net::OK)
return command_marshal->ReturnFailure("Couldn't find key's entry.");
scoped_refptr<net::StringIOBuffer> buffer(new net::StringIOBuffer(""));
rv = cache_entry->WriteData(index, 0, buffer.get(), 0, cb.callback(), true);
if (cb.GetResult(rv) != net::OK)
return command_marshal->ReturnFailure("Couldn't delete key stream.");
command_marshal->ReturnSuccess();
cache_entry->Close();
}
// Deletes a specified key from the cache.
void DeleteKey(CommandMarshal* command_marshal) {
std::string key = command_marshal->ReadString();
if (command_marshal->has_failed())
return;
net::TestCompletionCallback cb;
int rv = command_marshal->cache_backend()->DoomEntry(key, cb.callback());
if (cb.GetResult(rv) != net::OK)
command_marshal->ReturnFailure("Couldn't delete key.");
else
command_marshal->ReturnSuccess();
}
// Executes all command from the |command_marshal|.
bool ExecuteCommands(CommandMarshal* command_marshal) {
while (!command_marshal->has_failed()) {
std::string subcommand(command_marshal->ReadCommandName());
if (command_marshal->has_failed())
break;
if (subcommand == "stop") {
command_marshal->ReturnSuccess();
return true;
} else if (subcommand == "batch") {
StreamCommandMarshal stream_command_marshal(
command_marshal->cache_backend());
return ExecuteCommands(&stream_command_marshal);
} else if (subcommand == "delete_key") {
DeleteKey(command_marshal);
} else if (subcommand == "delete_stream") {
DeleteStreamForKey(command_marshal);
} else if (subcommand == "get_size") {
GetSize(command_marshal);
} else if (subcommand == "get_stream") {
GetStreamForKey(command_marshal);
} else if (subcommand == "list_keys") {
ListKeys(command_marshal);
} else if (subcommand == "update_raw_headers") {
UpdateRawResponseHeaders(command_marshal);
} else if (subcommand == "list_dups") {
ListDups(command_marshal);
} else {
// The wrong subcommand is originated from the command line.
command_marshal->ReturnFailure("Unknown command.");
PrintHelp();
}
}
return false;
}
} // namespace
int main(int argc, char* argv[]) {
base::AtExitManager at_exit_manager;
base::MessageLoopForIO message_loop;
base::CommandLine::Init(argc, argv);
const base::CommandLine& command_line =
*base::CommandLine::ForCurrentProcess();
base::CommandLine::StringVector args = command_line.GetArgs();
if (args.size() < 3U) {
PrintHelp();
return 1;
}
base::TaskScheduler::CreateAndStartWithDefaultParams("cachetool");
base::FilePath cache_path(args[0]);
std::string cache_backend_type(args[1]);
net::BackendType backend_type;
if (cache_backend_type == "simple") {
backend_type = net::CACHE_BACKEND_SIMPLE;
} else if (cache_backend_type == "blockfile") {
backend_type = net::CACHE_BACKEND_BLOCKFILE;
} else {
std::cerr << "Unknown cache type." << std::endl;
PrintHelp();
return 1;
}
std::unique_ptr<Backend> cache_backend;
net::TestCompletionCallback cb;
int rv = disk_cache::CreateCacheBackend(net::DISK_CACHE, backend_type,
cache_path, INT_MAX, false, nullptr,
&cache_backend, cb.callback());
if (cb.GetResult(rv) != net::OK) {
std::cerr << "Invalid cache." << std::endl;
return 1;
}
ProgramArgumentCommandMarshal program_argument_marshal(
cache_backend.get(),
base::CommandLine::StringVector(args.begin() + 2, args.end()));
bool successful_commands = ExecuteCommands(&program_argument_marshal);
base::RunLoop().RunUntilIdle();
cache_backend = nullptr;
disk_cache::FlushCacheThreadForTesting();
base::RunLoop().RunUntilIdle();
return !successful_commands;
}