Properly save custom header data

Ensure adding custom data upgrades to KDBX4
Implement review feedback

1 files changed, 30 insertions, 10 deletions

diff --git a/tests/TestKdbx4.cpp b/tests/TestKdbx4.cpp
index 8fb0f219b..7de3702a0 100644
--- a/tests/TestKdbx4.cpp
+++ b/tests/TestKdbx4.cpp
@@ -131,6 +131,7 @@ void TestKdbx4::testFormat400Upgrade()
 {
     QFETCH(Uuid, kdfUuid);
     QFETCH(Uuid, cipherUuid);
+    QFETCH(bool, addCustomData);
     QFETCH(quint32, expectedVersion);
 
     QScopedPointer<Database> sourceDb(new Database());
@@ -147,6 +148,12 @@ void TestKdbx4::testFormat400Upgrade()
     // upgrade to KDBX 4 by changing KDF and Cipher
     sourceDb->changeKdf(KeePass2::uuidToKdf(kdfUuid));
     sourceDb->setCipher(cipherUuid);
+
+    if (addCustomData) {
+        sourceDb->metadata()->customData()->set("CustomPublicData", "Hey look, I turned myself into a pickle!");
+        sourceDb->rootGroup()->customData()->set("CustomGroupData", "I just killed my family! I don't care who they were!");
+    }
+
     KeePass2Writer writer;
     writer.writeDatabase(&buffer, sourceDb.data());
     if (writer.hasError()) {
@@ -165,26 +172,39 @@ void TestKdbx4::testFormat400Upgrade()
     QCOMPARE(targetDb->metadata()->name(), sourceDb->metadata()->name());
 
     QCOMPARE(reader.version(), expectedVersion);
-    QCOMPARE(targetDb->kdf()->uuid(), sourceDb->kdf()->uuid());
     QCOMPARE(targetDb->cipher(), cipherUuid);
+    QCOMPARE(*targetDb->metadata()->customData(), *sourceDb->metadata()->customData());
+    QCOMPARE(*targetDb->rootGroup()->customData(), *sourceDb->rootGroup()->customData());
 }
 
 void TestKdbx4::testFormat400Upgrade_data()
 {
     QTest::addColumn<Uuid>("kdfUuid");
     QTest::addColumn<Uuid>("cipherUuid");
+    QTest::addColumn<bool>("addCustomData");
     QTest::addColumn<quint32>("expectedVersion");
 
     auto constexpr kdbx3 = KeePass2::FILE_VERSION_3_1 & KeePass2::FILE_VERSION_CRITICAL_MASK;
     auto constexpr kdbx4 = KeePass2::FILE_VERSION_4   & KeePass2::FILE_VERSION_CRITICAL_MASK;
 
-    QTest::newRow("Argon2           + AES")      << KeePass2::KDF_ARGON2    << KeePass2::CIPHER_AES      << kdbx4;
-    QTest::newRow("AES-KDF          + AES")      << KeePass2::KDF_AES_KDBX4 << KeePass2::CIPHER_AES      << kdbx4;
-    QTest::newRow("AES-KDF (legacy) + AES")      << KeePass2::KDF_AES_KDBX3 << KeePass2::CIPHER_AES      << kdbx3;
-    QTest::newRow("Argon2           + ChaCha20") << KeePass2::KDF_ARGON2    << KeePass2::CIPHER_CHACHA20 << kdbx4;
-    QTest::newRow("AES-KDF          + ChaCha20") << KeePass2::KDF_AES_KDBX4 << KeePass2::CIPHER_CHACHA20 << kdbx4;
-    QTest::newRow("AES-KDF (legacy) + ChaCha20") << KeePass2::KDF_AES_KDBX3 << KeePass2::CIPHER_CHACHA20 << kdbx3;
-    QTest::newRow("Argon2           + Twofish")  << KeePass2::KDF_ARGON2    << KeePass2::CIPHER_TWOFISH  << kdbx4;
-    QTest::newRow("AES-KDF          + Twofish")  << KeePass2::KDF_AES_KDBX4 << KeePass2::CIPHER_TWOFISH  << kdbx4;
-    QTest::newRow("AES-KDF (legacy) + Twofish")  << KeePass2::KDF_AES_KDBX3 << KeePass2::CIPHER_TWOFISH  << kdbx3;
+    QTest::newRow("Argon2           + AES")                   << KeePass2::KDF_ARGON2    << KeePass2::CIPHER_AES       << false << kdbx4;
+    QTest::newRow("AES-KDF          + AES")                   << KeePass2::KDF_AES_KDBX4 << KeePass2::CIPHER_AES       << false << kdbx4;
+    QTest::newRow("AES-KDF (legacy) + AES")                   << KeePass2::KDF_AES_KDBX3 << KeePass2::CIPHER_AES       << false << kdbx3;
+    QTest::newRow("Argon2           + AES     + CustomData")  << KeePass2::KDF_ARGON2    << KeePass2::CIPHER_AES       << true  << kdbx4;
+    QTest::newRow("AES-KDF          + AES     + CustomData")  << KeePass2::KDF_AES_KDBX4 << KeePass2::CIPHER_AES       << true  << kdbx4;
+    QTest::newRow("AES-KDF (legacy) + AES     + CustomData")  << KeePass2::KDF_AES_KDBX3 << KeePass2::CIPHER_AES       << true  << kdbx4;
+
+    QTest::newRow("Argon2           + ChaCha20")              << KeePass2::KDF_ARGON2    << KeePass2::CIPHER_CHACHA20  << false << kdbx4;
+    QTest::newRow("AES-KDF          + ChaCha20")              << KeePass2::KDF_AES_KDBX4 << KeePass2::CIPHER_CHACHA20  << false << kdbx4;
+    QTest::newRow("AES-KDF (legacy) + ChaCha20")              << KeePass2::KDF_AES_KDBX3 << KeePass2::CIPHER_CHACHA20  << false << kdbx3;
+    QTest::newRow("Argon2           + ChaCha20 + CustomData") << KeePass2::KDF_ARGON2    << KeePass2::CIPHER_CHACHA20  << true  << kdbx4;
+    QTest::newRow("AES-KDF          + ChaCha20 + CustomData") << KeePass2::KDF_AES_KDBX4 << KeePass2::CIPHER_CHACHA20  << true  << kdbx4;
+    QTest::newRow("AES-KDF (legacy) + ChaCha20 + CustomData") << KeePass2::KDF_AES_KDBX3 << KeePass2::CIPHER_CHACHA20  << true  << kdbx4;
+
+    QTest::newRow("Argon2           + Twofish")               << KeePass2::KDF_ARGON2    << KeePass2::CIPHER_TWOFISH   << false << kdbx4;
+    QTest::newRow("AES-KDF          + Twofish")               << KeePass2::KDF_AES_KDBX4 << KeePass2::CIPHER_TWOFISH   << false << kdbx4;
+    QTest::newRow("AES-KDF (legacy) + Twofish")               << KeePass2::KDF_AES_KDBX3 << KeePass2::CIPHER_TWOFISH   << false << kdbx3;
+    QTest::newRow("Argon2           + Twofish  + CustomData") << KeePass2::KDF_ARGON2    << KeePass2::CIPHER_TWOFISH   << true  << kdbx4;
+    QTest::newRow("AES-KDF          + Twofish  + CustomData") << KeePass2::KDF_AES_KDBX4 << KeePass2::CIPHER_TWOFISH   << true  << kdbx4;
+    QTest::newRow("AES-KDF (legacy) + Twofish  + CustomData") << KeePass2::KDF_AES_KDBX3 << KeePass2::CIPHER_TWOFISH   << true  << kdbx4;
 }