1 files changed, 22 insertions, 10 deletions
diff --git a/src/format/KeePass2Repair.cpp b/src/format/KeePass2Repair.cpp
index 8d18457d4..47ed59dc9 100644
--- a/src/format/KeePass2Repair.cpp
+++ b/src/format/KeePass2Repair.cpp
@@ -19,12 +19,13 @@
 #include "KeePass2Repair.h"
 
 #include <QBuffer>
-#include <QScopedPointer>
-#include <QRegExp>
 
+#include "core/Group.h"
+#include "format/KeePass2.h"
 #include "format/KeePass2RandomStream.h"
 #include "format/KeePass2Reader.h"
-#include "format/KeePass2XmlReader.h"
+#include "format/Kdbx4Reader.h"
+#include "format/KdbxXmlReader.h"
 
 KeePass2Repair::RepairOutcome KeePass2Repair::repairDatabase(QIODevice* device, const CompositeKey& key)
 {
@@ -38,7 +39,7 @@ KeePass2Repair::RepairOutcome KeePass2Repair::repairDatabase(QIODevice* device,
         return qMakePair(NothingTodo, nullptr);
     }
 
-    QByteArray xmlData = reader.xmlData();
+    QByteArray xmlData = reader.reader()->xmlData();
     if (!db || xmlData.isEmpty()) {
         m_errorStr = reader.errorString();
         return qMakePair(UnableToOpen, nullptr);
@@ -59,7 +60,7 @@ KeePass2Repair::RepairOutcome KeePass2Repair::repairDatabase(QIODevice* device,
 
     // try to fix broken databases because of bug #392
     for (int i = (xmlData.size() - 1); i >= 0; i--) {
-        quint8 ch = static_cast<quint8>(xmlData.at(i));
+        auto ch = static_cast<quint8>(xmlData.at(i));
         if (ch < 0x20 && ch != 0x09 && ch != 0x0A && ch != 0x0D) {
             xmlData.remove(i, 1);
             repairAction = true;
@@ -71,14 +72,25 @@ KeePass2Repair::RepairOutcome KeePass2Repair::repairDatabase(QIODevice* device,
         return qMakePair(RepairFailed, nullptr);
     }
 
-    KeePass2RandomStream randomStream;
-    randomStream.init(reader.streamKey());
-    KeePass2XmlReader xmlReader;
+    KeePass2RandomStream randomStream(reader.reader()->protectedStreamAlgo());
+    randomStream.init(reader.reader()->streamKey());
+    bool hasError;
+
     QBuffer buffer(&xmlData);
     buffer.open(QIODevice::ReadOnly);
-    xmlReader.readDatabase(&buffer, db.data(), &randomStream);
+    if ((reader.version() & KeePass2::FILE_VERSION_CRITICAL_MASK) < KeePass2::FILE_VERSION_4) {
+        KdbxXmlReader xmlReader(KeePass2::FILE_VERSION_3_1);
+        xmlReader.readDatabase(&buffer, db.data(), &randomStream);
+        hasError = xmlReader.hasError();
+    } else {
+        auto reader4 = reader.reader().staticCast<Kdbx4Reader>();
+        QHash<QString, QByteArray> pool = reader4->binaryPool();
+        KdbxXmlReader xmlReader(KeePass2::FILE_VERSION_4, pool);
+        xmlReader.readDatabase(&buffer, db.data(), &randomStream);
+        hasError = xmlReader.hasError();
+    }
 
-    if (xmlReader.hasError()) {
+    if (hasError) {
         return qMakePair(RepairFailed, nullptr);
     }
     else {