MB-65545: [cbcrypto] Add utility function maybeRewriteFiles

trondn · trondn · commit c0e3e530c8d8 · 2025-03-04T17:42:17.000Z
Add a function to allow iteration over files in a directory and potentially rewrite the files with or without encryption Change-Id: I0cbd23273ba86595f2ad7d792cda1d104d2778e9 Reviewed-on: https://review.couchbase.org/c/platform/+/224189 Tested-by: Build Bot <build@couchbase.com> Reviewed-by: Pavlos Georgiou <pavlos.georgiou@couchbase.com>
diff --git a/cbcrypto/file_utilities.cc b/cbcrypto/file_utilities.cc
@@ -9,12 +9,43 @@
  */
 
 #include <cbcrypto/encrypted_file_header.h>
+#include <cbcrypto/file_reader.h>
 #include <cbcrypto/file_utilities.h>
+#include <cbcrypto/file_writer.h>
+#include <fmt/format.h>
 #include <nlohmann/json.hpp>
+#include <platform/dirutils.h>
+
 #include <fstream>
 
 namespace cb::crypto {
 
+static std::string getEncryptionKey(const std::filesystem::path& path) {
+    std::array<char, sizeof(EncryptedFileHeader)> buffer;
+    auto size = file_size(path);
+    if (size < buffer.size()) {
+        // No header present
+        return {};
+    }
+
+    std::ifstream input;
+    input.exceptions(std::ifstream::failbit | std::ifstream::badbit);
+    input.open(path, std::ios::binary);
+    input.read(buffer.data(), buffer.size());
+    input.close();
+
+    auto* header = reinterpret_cast<EncryptedFileHeader*>(buffer.data());
+    if (!header->is_encrypted()) {
+        throw std::logic_error(
+                "File with .cef extension does not have correct magic");
+    }
+
+    if (!header->is_supported()) {
+        throw std::logic_error("File with .cef extension is not supported");
+    }
+    return std::string{header->get_id()};
+}
+
 std::unordered_set<std::string> findDeksInUse(
         const std::filesystem::path& directory,
         const std::function<bool(const std::filesystem::path&)>& filefilter,
@@ -28,39 +59,15 @@ std::unordered_set<std::string> findDeksInUse(
             continue;
         }
 
-        std::array<char, sizeof(cb::crypto::EncryptedFileHeader)> buffer;
         try {
-            auto size = file_size(path);
-            if (size < buffer.size()) {
-                // No header present
-                continue;
+            auto key = getEncryptionKey(path);
+            if (!key.empty()) {
+                deks.insert(key);
             }
-
-            std::ifstream input;
-            input.exceptions(std::ifstream::failbit | std::ifstream::badbit);
-            input.open(path, std::ios::binary);
-            input.read(buffer.data(), buffer.size());
-            input.close();
         } catch (const std::exception& e) {
             error("Failed to get deks from",
                   {{"path", path.string()}, {"error", e.what()}});
-            continue;
-        }
-
-        auto* header = reinterpret_cast<cb::crypto::EncryptedFileHeader*>(
-                buffer.data());
-        if (!header->is_encrypted()) {
-            error("Logfile with .cef extension does not have correct magic",
-                  {{"path", path.string()}});
-            continue;
-        }
-
-        if (!header->is_supported()) {
-            error("Logfile with .cef extension is not supported",
-                  {{"path", path.string()}});
-            continue;
         }
-        deks.insert(std::string(header->get_id()));
     }
 
     if (ec) {
@@ -71,4 +78,68 @@ std::unordered_set<std::string> findDeksInUse(
     return deks;
 }
 
+void maybeRewriteFiles(
+        const std::filesystem::path& directory,
+        const std::function<bool(const std::filesystem::path&,
+                                 std::string_view)>& filefilter,
+        SharedEncryptionKey encryption_key,
+        const std::function<SharedEncryptionKey(std::string_view)>&
+                key_lookup_function,
+        const std::function<void(std::string_view, const nlohmann::json&)>&
+                error,
+        std::string_view unencrypted_extension) {
+    std::error_code ec;
+    for (const auto& p : std::filesystem::directory_iterator(directory, ec)) {
+        auto path = p.path();
+        std::string key;
+        if (path.extension() == ".cef") {
+            try {
+                key = getEncryptionKey(path);
+            } catch (const std::exception& e) {
+                error("Failed to get deks from",
+                      {{"path", path.string()}, {"error", e.what()}});
+                continue;
+            }
+        }
+
+        if (!filefilter(path, key)) {
+            continue;
+        }
+
+        auto reader = FileReader::create(path, key_lookup_function);
+        std::filesystem::path tmpfile = cb::io::mktemp(path.string());
+        auto writer = FileWriter::create(encryption_key, tmpfile, 64 * 1024);
+
+        bool eof = false;
+        do {
+            try {
+                auto message = reader->nextChunk();
+                if (message.empty()) {
+                    eof = true;
+                } else {
+                    writer->write(message);
+                }
+            } catch (const std::underflow_error&) {
+                error("Partial chunk detected", {{"path", path.string()}});
+                eof = true;
+            }
+        } while (!eof);
+        writer->flush();
+        writer->close();
+        if (encryption_key && path.extension() != ".cef") {
+            auto next = path;
+            next.replace_extension(".cef");
+            rename(tmpfile, next);
+            remove(path);
+        } else if (!encryption_key && path.extension() == ".cef") {
+            auto next = path;
+            next.replace_extension(unencrypted_extension);
+            rename(tmpfile, next);
+            remove(path);
+        } else {
+            rename(tmpfile, path);
+        }
+    }
+}
+
 } // namespace cb::crypto
diff --git a/cbcrypto/file_utilities_test.cc b/cbcrypto/file_utilities_test.cc
@@ -9,6 +9,7 @@
  */
 
 #include "cbcrypto/file_utilities.h"
+#include "cbcrypto/key_store.h"
 
 #include <cbcrypto/common.h>
 #include <cbcrypto/file_reader.h>
@@ -25,26 +26,42 @@ using namespace std::string_literals;
 class FileUtilitiesTest : public ::testing::Test {
 protected:
     void SetUp() override {
+        keystore = {};
+        keystore.setActiveKey(DataEncryptionKey::generate());
         dir = cb::io::mkdtemp("FileUtilitiesTest");
     }
 
     void TearDown() override {
         remove_all(dir);
     }
 
-    void create_file(const std::string& name, const std::string& content) {
-        std::shared_ptr<DataEncryptionKey> key = DataEncryptionKey::generate();
+    void create_file(const std::string& name,
+                     std::string_view content,
+                     const bool encrypted = true,
+                     const std::size_t max_chunk_size =
+                             std::numeric_limits<std::size_t>::max()) {
+        std::shared_ptr<DataEncryptionKey> key;
+        if (encrypted) {
+            key = DataEncryptionKey::generate();
+            keystore.add(key);
+        }
         auto writer = FileWriter::create(key, dir / name);
-        writer->write(content);
+        while (!content.empty()) {
+            auto chunk =
+                    content.substr(0, std::min(content.size(), max_chunk_size));
+            content.remove_prefix(chunk.size());
+            writer->write(chunk);
+        }
         writer->flush();
         writer.reset();
 
-        files[name] = key->getId();
+        files[name] = encrypted ? key->getId() : "";
     }
 
     std::filesystem::path dir;
 
     std::unordered_map<std::string, std::string> files;
+    KeyStore keystore;
 };
 
 TEST_F(FileUtilitiesTest, findDeksInUse) {
@@ -84,7 +101,7 @@ TEST_F(FileUtilitiesTest, findDeksInUse) {
     bool error = false;
     keys = findDeksInUse(
             dir / "no-such-directory",
-            [](const std::filesystem::path& p) { return true; },
+            [](const std::filesystem::path&) { return true; },
             [&error](auto message, const auto& json) {
                 EXPECT_EQ("Error occurred while traversing directory", message);
                 EXPECT_TRUE(json.contains("error")) << json.dump();
@@ -93,3 +110,104 @@ TEST_F(FileUtilitiesTest, findDeksInUse) {
     EXPECT_TRUE(error) << "Error callback not called";
     EXPECT_TRUE(keys.empty());
 }
+
+TEST_F(FileUtilitiesTest, rewriteUnencryptedToEncrypted) {
+    create_file("file.txt", "This is the content", false);
+
+    maybeRewriteFiles(
+            dir,
+            [](const auto& path, auto) {
+                return path.filename().string() == "file.txt";
+            },
+            keystore.getActiveKey(),
+            [this](auto id) { return keystore.lookup(id); },
+            [](std::string_view, const nlohmann::json&) {});
+
+    auto file = dir / "file.txt";
+    EXPECT_FALSE(std::filesystem::exists(file));
+    file = dir / "file.cef";
+    EXPECT_TRUE(std::filesystem::exists(file));
+
+    const auto reader = FileReader::create(
+            file, [this](auto id) { return keystore.lookup(id); });
+    EXPECT_TRUE(reader->is_encrypted());
+    EXPECT_EQ("This is the content", reader->read());
+}
+
+TEST_F(FileUtilitiesTest, rewriteEncryptedToUnencrypted) {
+    create_file("file.cef", "This is the content");
+    maybeRewriteFiles(
+            dir,
+            [](const auto& path, auto) {
+                return path.filename().string() == "file.cef";
+            },
+            {},
+            [this](auto id) { return keystore.lookup(id); },
+            [](std::string_view, const nlohmann::json&) {});
+
+    auto file = dir / "file.cef";
+    EXPECT_FALSE(std::filesystem::exists(file));
+    file = dir / "file.txt";
+    EXPECT_TRUE(std::filesystem::exists(file));
+
+    const auto reader = FileReader::create(
+            file, [this](auto id) { return keystore.lookup(id); });
+    EXPECT_FALSE(reader->is_encrypted());
+    EXPECT_EQ("This is the content", reader->read());
+}
+
+TEST_F(FileUtilitiesTest, rewriteUsingCertainKey) {
+    create_file("file1.cef", "This is the content");
+    const auto key_id = files["file1.cef"];
+
+    maybeRewriteFiles(
+            dir,
+            [&key_id](const auto&, auto id) { return key_id == id; },
+            keystore.getActiveKey(),
+            [this](auto id) { return keystore.lookup(id); },
+            [](std::string_view, const nlohmann::json&) {});
+
+    std::string requested;
+    const auto reader =
+            FileReader::create(dir / "file1.cef", [this, &requested](auto id) {
+                requested = id;
+                return keystore.lookup(id);
+            });
+
+    EXPECT_TRUE(reader->is_encrypted());
+    EXPECT_EQ("This is the content", reader->read());
+    EXPECT_EQ(requested, keystore.getActiveKey()->getId());
+}
+
+TEST_F(FileUtilitiesTest, rewriteTruncatedFile) {
+    std::string data(1024 * 1024, 'a');
+    create_file("file1.cef", data, true, 1024);
+
+    const auto filename = dir / "file1.cef";
+    const auto size = file_size(filename);
+    // truncate the file with a partial final block
+    resize_file(filename, size - 512);
+
+    maybeRewriteFiles(
+            dir,
+            [](const auto& path, auto) {
+                return path.filename() == "file1.cef";
+            },
+            keystore.getActiveKey(),
+            [this](auto id) { return keystore.lookup(id); },
+            [](std::string_view, const nlohmann::json&) {});
+
+    std::string requested;
+    const auto reader =
+            FileReader::create(filename, [this, &requested](auto id) {
+                requested = id;
+                return keystore.lookup(id);
+            });
+
+    EXPECT_TRUE(reader->is_encrypted());
+    const auto content = reader->read();
+    EXPECT_EQ(content.size(), data.size() - 1024);
+    data.resize(data.size() - 1024);
+    EXPECT_EQ(data, content);
+    EXPECT_EQ(requested, keystore.getActiveKey()->getId());
+}
diff --git a/include/cbcrypto/file_utilities.h b/include/cbcrypto/file_utilities.h
@@ -16,6 +16,8 @@
 #include <unordered_set>
 
 namespace cb::crypto {
+struct DataEncryptionKey;
+using SharedEncryptionKey = std::shared_ptr<const DataEncryptionKey>;
 
 /**
  * Find all the DEKs in use in the specified directory
@@ -31,4 +33,29 @@ std::unordered_set<std::string> findDeksInUse(
         const std::function<void(std::string_view, const nlohmann::json&)>&
                 error);
 
+/**
+ * Iterate over all files in the specified directory and potentially
+ * rewrite all files.
+ *
+ * @param directory The directory to scan
+ * @param filefilter a function to filter out files we're not interested in
+ *                   (return true to inspect the file, false to skip it)
+ * @param encryption_key The key to use when rewriting the files (if empty the
+ * file will be written unencrypted)
+ * @param key_lookup_function A function used to look up encryption keys from
+ *            the id
+ * @param error a callback to add log messages when errors occurs
+ * @param unencrypted_extension The extension to use for unencrypted files
+ */
+void maybeRewriteFiles(
+        const std::filesystem::path& directory,
+        const std::function<bool(const std::filesystem::path&,
+                                 std::string_view)>& filefilter,
+        SharedEncryptionKey encryption_key,
+        const std::function<SharedEncryptionKey(std::string_view)>&
+                key_lookup_function,
+        const std::function<void(std::string_view, const nlohmann::json&)>&
+                error,
+        std::string_view unencrypted_extension = ".txt");
+
 } // namespace cb::crypto
