Fix IV reuse in encrypted single-key updates

Append encrypted updates so existing compaction handles IV rotation and sync, avoiding a synchronous flush on every update.

Add regression tests. Profiled on iPhone.

Core/MMKV_IO.cpp

@@ -1086,7 +1086,9 @@ KVHolderRet_t MMKV::doOverrideDataWithKey(const MMBuffer &data,
         return make_pair(false, KeyValueHolder());
     }
 
-    if (!checkSizeForOverride(size)) {
+    // Encrypted rewrites would reuse the IV. Append instead; compaction rotates
+    // the IV and synchronizes the rewritten data and metadata when needed.
+    if (m_crypter || !checkSizeForOverride(size)) {
         return doAppendDataWithKey(data, keyData, isDataHolder, originKeyLength);
     }
 

POSIX/demo/CMakeLists.txt

@@ -82,6 +82,11 @@ if(BUILD_TESTING)
         target_link_libraries(TestRemoveStorage core pthread ${CMAKE_DL_LIBS})
         set_target_properties(TestRemoveStorage PROPERTIES CXX_STANDARD 20)
         add_test(NAME TestRemoveStorage COMMAND TestRemoveStorage)
+        add_executable(TestSingleKeyIV TestSingleKeyIV.cpp)
+        target_include_directories(TestSingleKeyIV PRIVATE ../../Core)
+        target_link_libraries(TestSingleKeyIV core pthread)
+        set_target_properties(TestSingleKeyIV PROPERTIES CXX_STANDARD 20)
+        add_test(NAME TestSingleKeyIV COMMAND TestSingleKeyIV)
     endif()
 endif()
 

POSIX/demo/TestSingleKeyIV.cpp

@@ -0,0 +1,317 @@
+// 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;
+    }
+}