From 0d4aff55bc3f778e673e3f8185a69366ee11d407 Mon Sep 17 00:00:00 2001
From: Janek Bevendorff <janek@jbev.net>
Date: Wed, 28 Feb 2018 23:02:45 +0100
Subject: [PATCH] Don't upgrade to KDBX 4 when CustomData are present only in
 meta data section, resolves #1565

---
 src/format/KdbxXmlWriter.cpp  |  4 +---
 src/format/KeePass2Writer.cpp | 44 ++++++++++++++++++++++++-----------
 src/format/KeePass2Writer.h   |  1 +
 tests/TestKdbx4.cpp           |  4 +++-
 4 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/src/format/KdbxXmlWriter.cpp b/src/format/KdbxXmlWriter.cpp
index 8f405aca..a546f317 100644
--- a/src/format/KdbxXmlWriter.cpp
+++ b/src/format/KdbxXmlWriter.cpp
@@ -129,9 +129,7 @@ void KdbxXmlWriter::writeMetadata()
     if (m_kdbxVersion < KeePass2::FILE_VERSION_4) {
         writeBinaries();
     }
-    if (m_kdbxVersion >= KeePass2::FILE_VERSION_4) {
-        writeCustomData(m_meta->customData());
-    }
+    writeCustomData(m_meta->customData());
 
     m_xml.writeEndElement();
 }
diff --git a/src/format/KeePass2Writer.cpp b/src/format/KeePass2Writer.cpp
index 68cf3af9..7d5f470d 100644
--- a/src/format/KeePass2Writer.cpp
+++ b/src/format/KeePass2Writer.cpp
@@ -43,6 +43,36 @@ bool KeePass2Writer::writeDatabase(const QString& filename, Database* db)
     return writeDatabase(&file, db);
 }
 
+/**
+ * @return true if the database should upgrade to KDBX4.
+ */
+bool KeePass2Writer::implicitUpgradeNeeded(Database const* db) const
+{
+    if (!db->publicCustomData().isEmpty()) {
+        return true;
+    }
+
+    for (const auto& group: db->rootGroup()->groupsRecursive(true)) {
+        if (group->customData() && !group->customData()->isEmpty()) {
+            return true;
+        }
+
+        for (const auto& entry: group->entries()) {
+            if (entry->customData() && !entry->customData()->isEmpty()) {
+                return true;
+            }
+
+            for (const auto& historyItem: entry->historyItems()) {
+                if (historyItem->customData() && !historyItem->customData()->isEmpty()) {
+                    return true;
+                }
+            }
+        }
+    }
+
+    return false;
+}
+
 /**
  * Write a database to a device in KDBX format.
  *
@@ -55,19 +85,7 @@ bool KeePass2Writer::writeDatabase(QIODevice* device, Database* db) {
     m_error = false;
     m_errorStr.clear();
 
-    // determine KDBX3 vs KDBX4
-    bool hasCustomData = !db->publicCustomData().isEmpty() || (db->metadata()->customData() && !db->metadata()->customData()->isEmpty());
-    if (!hasCustomData) {
-        for (const auto& entry: db->rootGroup()->entriesRecursive(true)) {
-            if ((entry->customData() && !entry->customData()->isEmpty()) ||
-                (entry->group() && entry->group()->customData() && !entry->group()->customData()->isEmpty())) {
-                hasCustomData = true;
-                break;
-            }
-        }
-    }
-
-    if (db->kdf()->uuid() == KeePass2::KDF_AES_KDBX3 && !hasCustomData) {
+    if (db->kdf()->uuid() == KeePass2::KDF_AES_KDBX3 && !implicitUpgradeNeeded(db)) {
         m_version = KeePass2::FILE_VERSION_3_1;
         m_writer.reset(new Kdbx3Writer());
     } else {
diff --git a/src/format/KeePass2Writer.h b/src/format/KeePass2Writer.h
index 98daed5e..f024d4a8 100644
--- a/src/format/KeePass2Writer.h
+++ b/src/format/KeePass2Writer.h
@@ -42,6 +42,7 @@ public:
 
 private:
     void raiseError(const QString& errorMessage);
+    bool implicitUpgradeNeeded(Database const* db) const;
 
     bool m_error = false;
     QString m_errorStr = "";
diff --git a/tests/TestKdbx4.cpp b/tests/TestKdbx4.cpp
index 24a07aa6..08b24c47 100644
--- a/tests/TestKdbx4.cpp
+++ b/tests/TestKdbx4.cpp
@@ -149,8 +149,10 @@ void TestKdbx4::testFormat400Upgrade()
     sourceDb->changeKdf(KeePass2::uuidToKdf(kdfUuid));
     sourceDb->setCipher(cipherUuid);
 
+    // CustomData in meta should not cause any version change
+    sourceDb->metadata()->customData()->set("CustomPublicData", "Hey look, I turned myself into a pickle!");
     if (addCustomData) {
-        sourceDb->metadata()->customData()->set("CustomPublicData", "Hey look, I turned myself into a pickle!");
+        // this, however, should
         sourceDb->rootGroup()->customData()->set("CustomGroupData", "I just killed my family! I don't care who they were!");
     }