Properly save custom header data

Ensure adding custom data upgrades to KDBX4
Implement review feedback

src/format/Kdbx4Reader.cpp

@@ -209,9 +209,14 @@ bool Kdbx4Reader::readHeaderField(StoreDataStream& device)
         break;
     }
 
-    case KeePass2::HeaderFieldID::PublicCustomData:
-        m_db->setPublicCustomData(fieldData);
+    case KeePass2::HeaderFieldID::PublicCustomData: {
+        QBuffer variantBuffer;
+        variantBuffer.setBuffer(&fieldData);
+        variantBuffer.open(QBuffer::ReadOnly);
+        QVariantMap data = readVariantMap(&variantBuffer);
+        m_db->setPublicCustomData(data);
         break;
+    }
 
     case KeePass2::HeaderFieldID::ProtectedStreamKey:
     case KeePass2::HeaderFieldID::TransformRounds:
@@ -291,10 +296,10 @@ bool Kdbx4Reader::readInnerHeaderField(QIODevice* device)
 }
 
 /**
- * Helper method for reading KDF parameters into variant map.
+ * Helper method for reading a serialized variant map.
  *
  * @param device input device
- * @return filled variant map
+ * @return de-serialized variant map
  */
 QVariantMap Kdbx4Reader::readVariantMap(QIODevice* device)
 {

src/format/Kdbx4Writer.cpp

@@ -22,6 +22,8 @@
 
 #include "streams/HmacBlockStream.h"
 #include "core/Database.h"
+#include "core/Metadata.h"
+#include "core/CustomData.h"
 #include "crypto/CryptoHash.h"
 #include "crypto/Random.h"
 #include "format/KeePass2RandomStream.h"
@@ -86,10 +88,14 @@ bool Kdbx4Writer::writeDatabase(QIODevice* device, Database* db)
             raiseError(tr("Failed to serialize KDF parameters variant map"));
             return false;
         }
-        QByteArray publicCustomData = db->publicCustomData();
+
         CHECK_RETURN_FALSE(writeHeaderField<quint32>(&header, KeePass2::HeaderFieldID::KdfParameters, kdfParamBytes));
+        QVariantMap publicCustomData = db->publicCustomData();
         if (!publicCustomData.isEmpty()) {
-            CHECK_RETURN_FALSE(writeHeaderField<quint32>(&header, KeePass2::HeaderFieldID::PublicCustomData, publicCustomData));
+            QByteArray serialized;
+            serializeVariantMap(publicCustomData, serialized);
+            CHECK_RETURN_FALSE(writeHeaderField<quint32>(
+                &header, KeePass2::HeaderFieldID::PublicCustomData, serialized));
         }
 
         CHECK_RETURN_FALSE(writeHeaderField<quint32>(&header, KeePass2::HeaderFieldID::EndOfHeader, endOfHeader));
@@ -230,7 +236,7 @@ void Kdbx4Writer::writeAttachments(QIODevice* device, Database* db)
 }
 
 /**
- * Serialize KDF parameter variant map to byte array.
+ * Serialize variant map to byte array.
  *
  * @param map input variant map
  * @param outputBytes output byte array

src/format/KdbxXmlReader.cpp

@@ -397,7 +397,7 @@ void KdbxXmlReader::parseBinaries()
     }
 }
 
-void KdbxXmlReader::parseCustomData(CustomData *customData)
+void KdbxXmlReader::parseCustomData(CustomData* customData)
 {
     Q_ASSERT(m_xml.isStartElement() && m_xml.name() == "CustomData");
 
@@ -410,7 +410,7 @@ void KdbxXmlReader::parseCustomData(CustomData *customData)
     }
 }
 
-void KdbxXmlReader::parseCustomDataItem(CustomData *customData)
+void KdbxXmlReader::parseCustomDataItem(CustomData* customData)
 {
     Q_ASSERT(m_xml.isStartElement() && m_xml.name() == "Item");
 
@@ -748,7 +748,7 @@ Entry* KdbxXmlReader::parseEntry(bool history)
             }
             continue;
         }
-        if (m_xml.name() == "CustomData" ){
+        if (m_xml.name() == "CustomData") {
             parseCustomData(entry->customData());
             continue;
         }

src/format/KdbxXmlReader.h

@@ -66,8 +66,8 @@ protected:
     virtual void parseCustomIcons();
     virtual void parseIcon();
     virtual void parseBinaries();
-    virtual void parseCustomData(CustomData *customData);
-    virtual void parseCustomDataItem(CustomData *customData);
+    virtual void parseCustomData(CustomData* customData);
+    virtual void parseCustomDataItem(CustomData* customData);
     virtual bool parseRoot();
     virtual Group* parseGroup();
     virtual void parseDeletedObjects();

src/format/KdbxXmlWriter.cpp

@@ -129,7 +129,9 @@ void KdbxXmlWriter::writeMetadata()
     if (m_kdbxVersion < KeePass2::FILE_VERSION_4) {
         writeBinaries();
     }
-    writeCustomData(m_meta->customData());
+    if (m_kdbxVersion >= KeePass2::FILE_VERSION_4) {
+        writeCustomData(m_meta->customData());
+    }
 
     m_xml.writeEndElement();
 }
@@ -218,8 +220,11 @@ void KdbxXmlWriter::writeBinaries()
     m_xml.writeEndElement();
 }
 
-void KdbxXmlWriter::writeCustomData(const CustomData *customData)
+void KdbxXmlWriter::writeCustomData(const CustomData* customData)
 {
+    if (customData->isEmpty()) {
+        return;
+    }
     m_xml.writeStartElement("CustomData");
 
     const QList<QString> keyList = customData->keys();
@@ -276,7 +281,7 @@ void KdbxXmlWriter::writeGroup(const Group* group)
 
     writeUuid("LastTopVisibleEntry", group->lastTopVisibleEntry());
 
-    if (!group->customData()->isEmpty()){
+    if (m_kdbxVersion >= KeePass2::FILE_VERSION_4) {
         writeCustomData(group->customData());
     }
 
@@ -405,7 +410,7 @@ void KdbxXmlWriter::writeEntry(const Entry* entry)
 
     writeAutoType(entry);
 
-    if (!entry->customData()->isEmpty()){
+    if (m_kdbxVersion >= KeePass2::FILE_VERSION_4) {
         writeCustomData(entry->customData());
     }
 

src/format/KdbxXmlWriter.h

@@ -51,7 +51,7 @@ private:
     void writeCustomIcons();
     void writeIcon(const Uuid& uuid, const QImage& icon);
     void writeBinaries();
-    void writeCustomData(const CustomData *customData);
+    void writeCustomData(const CustomData* customData);
     void writeCustomDataItem(const QString& key, const QString& value);
     void writeRoot();
     void writeGroup(const Group* group);

src/format/KeePass2Writer.cpp

@@ -19,6 +19,8 @@
 #include <QFile>
 
 #include "core/Database.h"
+#include "core/Group.h"
+#include "core/Metadata.h"
 #include "crypto/kdf/AesKdf.h"
 #include "format/KeePass2Writer.h"
 #include "format/Kdbx3Writer.h"
@@ -48,12 +50,24 @@ bool KeePass2Writer::writeDatabase(const QString& filename, Database* db)
  * @param db source database
  * @return true on success
  */
+
 bool KeePass2Writer::writeDatabase(QIODevice* device, Database* db) {
     m_error = false;
     m_errorStr.clear();
 
     // determine KDBX3 vs KDBX4
-    if (db->kdf()->uuid() == KeePass2::KDF_AES_KDBX3 && db->publicCustomData().isEmpty()) {
+    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) {
         m_version = KeePass2::FILE_VERSION_3_1;
         m_writer.reset(new Kdbx3Writer());
     } else {