mirror of
https://github.com/Tencent/MMKV.git
synced 2026-10-17 04:30:56 +08:00
guoling
acadaa53a1
Append encrypted updates so existing compaction handles IV rotation and sync, avoiding a synchronous flush on every update. Add regression tests. Profiled on iPhone.
318 lines
14 KiB
C++
318 lines
14 KiB
C++
// Exercise single-key overwrites through the public API and inspect the
|
|
// persisted IV/ciphertext, rather than testing an isolated encryption helper.
|
|
#include "MMKV.h"
|
|
#include "MMBuffer.h"
|
|
#include "MMKVMetaInfo.hpp"
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <fstream>
|
|
#include <iterator>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <vector>
|
|
#include <unistd.h>
|
|
|
|
using namespace mmkv;
|
|
|
|
static void require(bool condition, const char *message) {
|
|
if (!condition) {
|
|
throw std::runtime_error(message);
|
|
}
|
|
}
|
|
|
|
static bool passed = true;
|
|
static void check(bool condition, const char *message) {
|
|
if (!condition) {
|
|
std::fprintf(stderr, "test single-key IV: %s\n", message);
|
|
passed = false;
|
|
}
|
|
}
|
|
|
|
static constexpr const char *storageIDs[] = {
|
|
"aes128", "aes256", "growth", "expiry", "plain", "borrowed", "transitions", "compaction"};
|
|
|
|
struct OwnedDirectory {
|
|
std::string path;
|
|
explicit OwnedDirectory(const char *base) {
|
|
path = std::string(base && *base ? base : "/tmp") + "/mmkv-single-key-iv.XXXXXX";
|
|
require(mkdtemp(path.data()) != nullptr, "cannot create owned scratch directory");
|
|
}
|
|
~OwnedDirectory() {
|
|
MMKV::onExit();
|
|
// Only these files inside the directory created by this invocation
|
|
// belong to the test. Never clean pre-existing temp directories.
|
|
for (const char *id : storageIDs) {
|
|
unlink((path + "/" + id).c_str());
|
|
unlink((path + "/" + id + ".crc").c_str());
|
|
}
|
|
if (rmdir(path.c_str()) != 0) {
|
|
std::fprintf(stderr, "test single-key IV: cannot remove owned directory %s\n", path.c_str());
|
|
passed = false;
|
|
}
|
|
}
|
|
};
|
|
|
|
static std::vector<uint8_t> readFile(const std::string &path) {
|
|
std::ifstream input(path, std::ios::binary);
|
|
require(input.is_open(), "cannot open snapshot file");
|
|
std::vector<uint8_t> result((std::istreambuf_iterator<char>(input)), std::istreambuf_iterator<char>());
|
|
require(!input.bad(), "cannot read snapshot file");
|
|
return result;
|
|
}
|
|
|
|
struct Snapshot {
|
|
std::array<uint8_t, AES_IV_LEN> iv;
|
|
std::vector<uint8_t> ciphertext;
|
|
};
|
|
|
|
static Snapshot snapshot(MMKV *kv, const std::string &root) {
|
|
kv->sync();
|
|
auto metaBytes = readFile(root + "/" + kv->mmapID() + ".crc");
|
|
require(metaBytes.size() >= sizeof(MMKVMetaInfo), "truncated metadata snapshot");
|
|
MMKVMetaInfo meta;
|
|
std::memcpy(&meta, metaBytes.data(), sizeof(meta));
|
|
auto data = readFile(root + "/" + kv->mmapID());
|
|
require(meta.m_actualSize == kv->actualSize(), "persisted actual size does not match instance");
|
|
require(data.size() >= sizeof(uint32_t) + meta.m_actualSize, "truncated data snapshot");
|
|
data.resize(sizeof(uint32_t) + meta.m_actualSize);
|
|
Snapshot result;
|
|
std::memcpy(result.iv.data(), meta.m_vector, result.iv.size());
|
|
result.ciphertext = std::move(data);
|
|
return result;
|
|
}
|
|
|
|
static MMKV *openKV(const char *id, const MMKVConfig &config) {
|
|
auto kv = MMKV::mmkvWithID(id, config);
|
|
require(kv != nullptr, "cannot open storage");
|
|
return kv;
|
|
}
|
|
|
|
static void verifyString(MMKV *kv, const std::string &expected) {
|
|
std::string actual;
|
|
require(kv->getString("k", actual) && actual == expected, "decoded value differs after overwrite");
|
|
require(kv->count() == 1, "overwrite did not preserve a single key");
|
|
}
|
|
|
|
static MMKV *reloadAndReopen(MMKV *kv, const char *id, const MMKVConfig &config, const std::string &expected) {
|
|
verifyString(kv, expected);
|
|
kv->clearMemoryCache();
|
|
verifyString(kv, expected);
|
|
require(MMKV::isFileValid(id), "persisted overwrite fails CRC validation");
|
|
kv->close();
|
|
kv = openKV(id, config);
|
|
verifyString(kv, expected);
|
|
return kv;
|
|
}
|
|
|
|
static bool preservesCiphertextAndAppends(const Snapshot &before, const Snapshot &after) {
|
|
// The raw four-byte size header changes as entries are appended. Every
|
|
// previously active encrypted byte must retain its position and content.
|
|
return after.ciphertext.size() > before.ciphertext.size() &&
|
|
std::equal(before.ciphertext.begin() + sizeof(uint32_t), before.ciphertext.end(),
|
|
after.ciphertext.begin() + sizeof(uint32_t));
|
|
}
|
|
|
|
static void checkEncryptedReplacement(const Snapshot &before, const Snapshot &after) {
|
|
check(after.iv != before.iv || preservesCiphertextAndAppends(before, after),
|
|
"encrypted replacement rewrote existing ciphertext without changing the IV");
|
|
}
|
|
|
|
static void securityRegression(const Snapshot &first, const Snapshot &second) {
|
|
if (first.iv == second.iv) {
|
|
// Historical entries remain when replacement appends. Their identical
|
|
// prefix is safe only if the new entry starts beyond the previous end.
|
|
check(preservesCiphertextAndAppends(first, second),
|
|
"same-IV replacement did not preserve ciphertext and append beyond its active end");
|
|
return;
|
|
}
|
|
// Disk offsets 0..3 hold the unencrypted size; 4..7 are the random
|
|
// MiniPB size holder. With key "k" and a ten-byte string, 8..11 are
|
|
// fixed framing and 12..19 are eight value bytes in the first AES block.
|
|
// CFB's first-block keystream depends only on the initial IV and key.
|
|
require(first.ciphertext.size() >= 20 && second.ciphertext.size() >= 20,
|
|
"snapshot too short for known-plaintext regression");
|
|
bool sameFraming = true;
|
|
bool recoveredValue = true;
|
|
for (size_t i = 8; i < 12; ++i) {
|
|
sameFraming &= first.ciphertext[i] == second.ciphertext[i];
|
|
}
|
|
for (size_t i = 12; i < 20; ++i) {
|
|
recoveredValue &= uint8_t('A' ^ first.ciphertext[i] ^ second.ciphertext[i]) == uint8_t('B');
|
|
}
|
|
check(!sameFraming, "fixed framing ciphertext is identical across overwrites");
|
|
check(!recoveredValue, "known plaintext recovers all eight bytes of the replacement value");
|
|
}
|
|
|
|
static void overwriteSizes(const std::string &root, const char *id, const MMKVConfig &config,
|
|
bool encrypted, bool inspectLeak = false) {
|
|
auto kv = openKV(id, config);
|
|
// Preallocated storage lets replacements cross direct, heap-backed, and
|
|
// offset-backed holder sizes without requiring capacity expansion.
|
|
require(kv->set(std::string(10, 'Z'), "k"), "initial write failed");
|
|
auto previous = snapshot(kv, root);
|
|
Snapshot knownPlaintext;
|
|
const std::vector<std::string> values = {
|
|
std::string(10, 'A'), std::string(10, 'B'), std::string(100, 'C'),
|
|
std::string(300, 'D'), std::string(32768, 'E'), std::string(3, 'F'), std::string(300, 'G')};
|
|
for (size_t i = 0; i < values.size(); ++i) {
|
|
require(kv->set(values[i], "k"), "overwrite failed");
|
|
auto current = snapshot(kv, root);
|
|
if (encrypted) {
|
|
checkEncryptedReplacement(previous, current);
|
|
}
|
|
if (inspectLeak && i == 0) {
|
|
knownPlaintext = current;
|
|
} else if (inspectLeak && i == 1) {
|
|
securityRegression(knownPlaintext, current);
|
|
}
|
|
kv = reloadAndReopen(kv, id, config, values[i]);
|
|
previous = std::move(current);
|
|
}
|
|
kv->close();
|
|
}
|
|
|
|
static void capacityGrowth(const std::string &root, const MMKVConfig &config) {
|
|
auto kv = openKV("growth", config);
|
|
require(kv->set(std::string(10, 'A'), "k"), "initial growth write failed");
|
|
auto initialCapacity = kv->totalSize();
|
|
const std::string large(initialCapacity * 4, 'B');
|
|
// The first oversized replacement expands storage; the next replacement
|
|
// exercises its existing append/compaction lifecycle at the larger size.
|
|
require(kv->set(large, "k"), "capacity expansion failed");
|
|
check(kv->totalSize() > initialCapacity, "growth scenario did not expand the file");
|
|
auto before = snapshot(kv, root);
|
|
kv = reloadAndReopen(kv, "growth", config, large);
|
|
const std::string replacement(large.size(), 'C');
|
|
require(kv->set(replacement, "k"), "overwrite after expansion failed");
|
|
auto after = snapshot(kv, root);
|
|
checkEncryptedReplacement(before, after);
|
|
kv = reloadAndReopen(kv, "growth", config, replacement);
|
|
kv->close();
|
|
}
|
|
|
|
static void borrowedInput(const std::string &root, const MMKVConfig &config) {
|
|
auto kv = openKV("borrowed", config);
|
|
require(kv->set(std::string(100, 'A'), "k"), "initial borrowed-input write failed");
|
|
auto before = snapshot(kv, root);
|
|
std::string source(100, 'B');
|
|
MMBuffer borrowed(source.data(), source.size(), MMBufferNoCopy);
|
|
require(kv->set(borrowed, "k"), "borrowed-input overwrite failed");
|
|
auto after = snapshot(kv, root);
|
|
checkEncryptedReplacement(before, after);
|
|
// Replacing the caller's bytes must not change the cached value.
|
|
source.assign(source.size(), 'X');
|
|
kv = reloadAndReopen(kv, "borrowed", config, std::string(100, 'B'));
|
|
kv->close();
|
|
}
|
|
|
|
static void dictionaryTransitions(const std::string &root, const MMKVConfig &config) {
|
|
auto kv = openKV("transitions", config);
|
|
require(kv->set(std::string(300, 'A'), "k"), "initial transition write failed");
|
|
auto initial = snapshot(kv, root);
|
|
require(kv->set(std::string(300, 'B'), "k"), "transition overwrite failed");
|
|
auto overwritten = snapshot(kv, root);
|
|
checkEncryptedReplacement(initial, overwritten);
|
|
// Appending after a replacement must continue from the current crypter state,
|
|
// while reads of the offset-backed first value retain its saved status.
|
|
require(kv->set(std::string(300, 'C'), "other"), "append after overwrite failed");
|
|
auto verifyTwoKeys = [](MMKV *instance) {
|
|
std::string first, second;
|
|
require(instance->getString("k", first) && first == std::string(300, 'B'),
|
|
"first value corrupted by append after overwrite");
|
|
require(instance->getString("other", second) && second == std::string(300, 'C'),
|
|
"appended value does not decode with new IV");
|
|
require(instance->count() == 2, "append did not retain both keys");
|
|
};
|
|
auto appended = snapshot(kv, root);
|
|
check(appended.iv == overwritten.iv, "ordinary append unexpectedly changed the IV");
|
|
verifyTwoKeys(kv);
|
|
kv->clearMemoryCache();
|
|
verifyTwoKeys(kv);
|
|
require(MMKV::isFileValid("transitions"), "append after overwrite fails CRC validation");
|
|
kv->close();
|
|
kv = openKV("transitions", config);
|
|
verifyTwoKeys(kv);
|
|
|
|
require(kv->removeValueForKey("other") && kv->removeValueForKey("k"), "cannot remove last keys");
|
|
require(kv->count() == 0 && kv->actualSize() > 0, "empty dictionary lacks retained log for override");
|
|
auto removed = snapshot(kv, root);
|
|
require(kv->set(std::string(300, 'D'), "k"), "insert after last-key removal failed");
|
|
auto inserted = snapshot(kv, root);
|
|
check(inserted.iv != removed.iv, "insert into emptied dictionary reused the persisted IV");
|
|
kv = reloadAndReopen(kv, "transitions", config, std::string(300, 'D'));
|
|
kv->close();
|
|
}
|
|
|
|
static void forcedCompaction(const std::string &root, MMKVConfig config) {
|
|
config.expectedCapacity = 0;
|
|
auto kv = openKV("compaction", config);
|
|
require(kv->set(std::string(300, 'A'), "k"), "initial compaction write failed");
|
|
const size_t valueSize = kv->totalSize() / 10;
|
|
auto previous = snapshot(kv, root);
|
|
bool rotatedIV = false;
|
|
std::string value;
|
|
// Roughly ten entries fill the initial mapping. Repeated replacements
|
|
// must compact old entries and persist a fresh IV at least once.
|
|
for (size_t i = 0; i < 32; ++i) {
|
|
value.assign(valueSize, char('B' + i % 24));
|
|
require(kv->set(value, "k"), "compaction replacement failed");
|
|
auto current = snapshot(kv, root);
|
|
checkEncryptedReplacement(previous, current);
|
|
verifyString(kv, value);
|
|
if (current.iv != previous.iv) {
|
|
rotatedIV = true;
|
|
kv = reloadAndReopen(kv, "compaction", config, value);
|
|
}
|
|
previous = std::move(current);
|
|
}
|
|
check(rotatedIV, "repeated replacements did not exercise compaction with a fresh IV");
|
|
kv = reloadAndReopen(kv, "compaction", config, value);
|
|
kv->close();
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
try {
|
|
require(argc <= 2, "usage: TestSingleKeyIV [scratch-directory]");
|
|
{
|
|
OwnedDirectory directory(argc == 2 ? argv[1] : std::getenv("TMPDIR"));
|
|
MMKV::initializeMMKV(directory.path, MMKVLogNone);
|
|
MMKVConfig config;
|
|
config.expectedCapacity = 128 * 1024;
|
|
#ifndef MMKV_DISABLE_CRYPT
|
|
const std::string key128 = "0123456789abcdef";
|
|
config.cryptKey = &key128;
|
|
overwriteSizes(directory.path, "aes128", config, true, true);
|
|
const std::string key256 = "0123456789abcdef0123456789abcdef";
|
|
config.cryptKey = &key256;
|
|
config.aes256 = true;
|
|
overwriteSizes(directory.path, "aes256", config, true, true);
|
|
config.cryptKey = &key128;
|
|
config.aes256 = false;
|
|
borrowedInput(directory.path, config);
|
|
dictionaryTransitions(directory.path, config);
|
|
forcedCompaction(directory.path, config);
|
|
config.enableKeyExpire = true;
|
|
config.expiredInSeconds = 3600;
|
|
overwriteSizes(directory.path, "expiry", config, true);
|
|
config.enableKeyExpire = std::nullopt;
|
|
config.expiredInSeconds = 0;
|
|
config.expectedCapacity = 0;
|
|
capacityGrowth(directory.path, config);
|
|
config.cryptKey = nullptr;
|
|
#endif
|
|
config.expectedCapacity = 128 * 1024;
|
|
overwriteSizes(directory.path, "plain", config, false);
|
|
}
|
|
if (passed) {
|
|
std::puts("test single-key IV: safe replacements, compaction IVs, reload/reopen, sizes and plain control passed");
|
|
}
|
|
return passed ? EXIT_SUCCESS : EXIT_FAILURE;
|
|
} catch (const std::exception &error) {
|
|
std::fprintf(stderr, "test single-key IV: %s\n", error.what());
|
|
return EXIT_FAILURE;
|
|
}
|
|
}
|