Implement custom LZMA dictionary size support, fixes #154

LongSoft · Dec 30, 2018 · 4bee991 · 4bee991
1 parent e6f84f9
commit 4bee991
Show file tree

Hide file tree

Showing 9 changed files with 93 additions and 33 deletions.
diff --git a/LZMA/LzmaCompress.c b/LZMA/LzmaCompress.c
@@ -64,7 +64,8 @@ LzmaCompress (
     CONST UINT8  *Source,
     UINT32       SourceSize,
     UINT8        *Destination,
-    UINT32       *DestinationSize
+    UINT32       *DestinationSize,
+    UINT32       DictionarySize
 )
 {
     SRes              LzmaResult;
@@ -79,8 +80,7 @@ LzmaCompress (
     }
 
     LzmaEncProps_Init(&props);
-    // TODO: need to detect this instead of hardcoding
-    props.dictSize = LZMA_DICTIONARY_SIZE;
+    props.dictSize = DictionarySize;
     props.level = 9;
     props.fb = 273;
 

diff --git a/LZMA/LzmaCompress.h b/LZMA/LzmaCompress.h
@@ -21,7 +21,7 @@
 extern "C" {
 #endif
 
-#define LZMA_DICTIONARY_SIZE 0x800000
+#define DEFAULT_LZMA_DICTIONARY_SIZE 0x800000
 #define _LZMA_SIZE_OPT
 
     EFI_STATUS
@@ -30,7 +30,8 @@ extern "C" {
             const UINT8  *Source,
             UINT32       SourceSize,
             UINT8        *Destination,
-            UINT32       *DestinationSize
+            UINT32       *DestinationSize,
+            UINT32       DictionarySize
         );
 
 #ifdef __cplusplus

diff --git a/ffsengine.cpp b/ffsengine.cpp
@@ -1590,8 +1590,16 @@ UINT8 FfsEngine::parseSection(const QByteArray & section, QModelIndex & index, c
             .arg(compressionTypeToQString(algorithm))
             .hexarg(compressedSectionHeader->UncompressedLength).arg(compressedSectionHeader->UncompressedLength);
 
+        UINT32 dictionarySize = DEFAULT_LZMA_DICTIONARY_SIZE;
+        if (algorithm == COMPRESSION_ALGORITHM_LZMA) {
+            // Dictionary size is stored in bytes 1-4 of LZMA-compressed data
+            dictionarySize = *(UINT32*)(body.constData() + 1);
+            info += tr("\nLZMA dictionary size: %1h").hexarg(dictionarySize);
+        }
+
         // Add tree item
         index = model->addItem(Types::Section, sectionHeader->Type, algorithm, name, "", info, header, body, parent, mode);
+        model->setDictionarySize(index, dictionarySize);
 
         // Show message
         if (!parseCurrentSection)
@@ -1632,6 +1640,8 @@ UINT8 FfsEngine::parseSection(const QByteArray & section, QModelIndex & index, c
             .hexarg2(guidDefinedSectionHeader->Attributes, 4);
 
         UINT8 algorithm = COMPRESSION_ALGORITHM_NONE;
+        UINT32 dictionarySize = DEFAULT_LZMA_DICTIONARY_SIZE;
+
         // Check if section requires processing
         if (guidDefinedSectionHeader->Attributes & EFI_GUIDED_SECTION_PROCESSING_REQUIRED) {
             // Tiano compressed section
@@ -1664,6 +1674,10 @@ UINT8 FfsEngine::parseSection(const QByteArray & section, QModelIndex & index, c
                 if (algorithm == COMPRESSION_ALGORITHM_LZMA) {
                     info += tr("\nCompression type: LZMA");
                     info += tr("\nDecompressed size: %1h (%2)").hexarg(processed.length()).arg(processed.length());
+
+                    // Dictionary size is stored in bytes 1-4 of LZMA-compressed data
+                    dictionarySize = *(UINT32*)(body.constData() + 1);
+                    info += tr("\nLZMA dictionary size: %1h").hexarg(dictionarySize);
                 }
                 else
                     info += tr("\nCompression type: unknown");
@@ -1743,6 +1757,7 @@ UINT8 FfsEngine::parseSection(const QByteArray & section, QModelIndex & index, c
 
         // Add tree item
         index = model->addItem(Types::Section, sectionHeader->Type, algorithm, name, "", info, header, body, parent, mode);
+        model->setDictionarySize(index, dictionarySize);
 
         // Show messages
         if (msgUnknownGuid)
@@ -2164,6 +2179,7 @@ UINT8 FfsEngine::create(const QModelIndex & index, const UINT8 type, const QByte
     QByteArray created;
     UINT8 result;
     QModelIndex fileIndex;
+    UINT32 defaultDictionarySize = DEFAULT_LZMA_DICTIONARY_SIZE;
 
     if (!index.isValid() || !index.parent().isValid())
         return ERR_INVALID_PARAMETER;
@@ -2344,18 +2360,21 @@ UINT8 FfsEngine::create(const QModelIndex & index, const UINT8 type, const QByte
             sectionHeader->UncompressedLength = body.size();
 
             // Set compression type
-            if (algorithm == COMPRESSION_ALGORITHM_NONE)
+            if (algorithm == COMPRESSION_ALGORITHM_NONE) {
                 sectionHeader->CompressionType = EFI_NOT_COMPRESSED;
-            else if (algorithm == COMPRESSION_ALGORITHM_EFI11 || algorithm == COMPRESSION_ALGORITHM_TIANO)
+            }
+            else if (algorithm == COMPRESSION_ALGORITHM_EFI11 || algorithm == COMPRESSION_ALGORITHM_TIANO) {
                 sectionHeader->CompressionType = EFI_STANDARD_COMPRESSION;
-            else if (algorithm == COMPRESSION_ALGORITHM_LZMA || algorithm == COMPRESSION_ALGORITHM_IMLZMA)
+            }
+            else if (algorithm == COMPRESSION_ALGORITHM_LZMA || algorithm == COMPRESSION_ALGORITHM_IMLZMA) {
                 sectionHeader->CompressionType = EFI_CUSTOMIZED_COMPRESSION;
+            }
             else
                 return ERR_UNKNOWN_COMPRESSION_ALGORITHM;
 
             // Compress body
             QByteArray compressed;
-            result = compress(body, algorithm, compressed);
+            result = compress(body, algorithm, defaultDictionarySize, compressed);
             if (result)
                 return result;
 
@@ -2381,7 +2400,7 @@ UINT8 FfsEngine::create(const QModelIndex & index, const UINT8 type, const QByte
         case EFI_SECTION_GUID_DEFINED:{
             // Compress body
             QByteArray compressed;
-            result = compress(body, algorithm, compressed);
+            result = compress(body, algorithm, defaultDictionarySize, compressed);
             if (result)
                 return result;
 
@@ -2843,7 +2862,7 @@ UINT8 FfsEngine::decompress(const QByteArray & compressedData, const UINT8 compr
     }
 }
 
-UINT8 FfsEngine::compress(const QByteArray & data, const UINT8 algorithm, QByteArray & compressedData)
+UINT8 FfsEngine::compress(const QByteArray & data, const UINT8 algorithm, const UINT32 dictionarySize, QByteArray & compressedData)
 {
     UINT8* compressed;
 
@@ -2933,10 +2952,10 @@ UINT8 FfsEngine::compress(const QByteArray & data, const UINT8 algorithm, QByteA
     case COMPRESSION_ALGORITHM_LZMA:
     {
         UINT32 compressedSize = 0;
-        if (LzmaCompress((const UINT8*)data.constData(), data.size(), NULL, &compressedSize) != ERR_BUFFER_TOO_SMALL)
+        if (LzmaCompress((const UINT8*)data.constData(), data.size(), NULL, &compressedSize, dictionarySize) != ERR_BUFFER_TOO_SMALL)
             return ERR_CUSTOMIZED_COMPRESSION_FAILED;
         compressed = new UINT8[compressedSize];
-        if (LzmaCompress((const UINT8*)data.constData(), data.size(), compressed, &compressedSize) != ERR_SUCCESS) {
+        if (LzmaCompress((const UINT8*)data.constData(), data.size(), compressed, &compressedSize, dictionarySize) != ERR_SUCCESS) {
             delete[] compressed;
             return ERR_CUSTOMIZED_COMPRESSION_FAILED;
         }
@@ -2953,10 +2972,10 @@ UINT8 FfsEngine::compress(const QByteArray & data, const UINT8 algorithm, QByteA
         UINT32 headerSize = sizeOfSectionHeader(sectionHeader);
         header = data.left(headerSize);
         QByteArray newData = data.mid(headerSize);
-        if (LzmaCompress((const UINT8*)newData.constData(), newData.size(), NULL, &compressedSize) != ERR_BUFFER_TOO_SMALL)
+        if (LzmaCompress((const UINT8*)newData.constData(), newData.size(), NULL, &compressedSize, dictionarySize) != ERR_BUFFER_TOO_SMALL)
             return ERR_CUSTOMIZED_COMPRESSION_FAILED;
         compressed = new UINT8[compressedSize];
-        if (LzmaCompress((const UINT8*)newData.constData(), newData.size(), compressed, &compressedSize) != ERR_SUCCESS) {
+        if (LzmaCompress((const UINT8*)newData.constData(), newData.size(), compressed, &compressedSize, dictionarySize) != ERR_SUCCESS) {
             delete[] compressed;
             return ERR_CUSTOMIZED_COMPRESSION_FAILED;
         }
@@ -3919,29 +3938,37 @@ UINT8 FfsEngine::reconstructSection(const QModelIndex& index, const UINT32 base,
                 EFI_COMPRESSION_SECTION* compessionHeader = (EFI_COMPRESSION_SECTION*)header.data();
                 // Set new uncompressed size
                 compessionHeader->UncompressedLength = reconstructed.size();
-                // Compress new section body
-                QByteArray compressed;
-                result = compress(reconstructed, model->compression(index), compressed);
-                if (result)
-                    return result;
+
                 // Correct compression type
-                if (model->compression(index) == COMPRESSION_ALGORITHM_NONE)
+                if (model->compression(index) == COMPRESSION_ALGORITHM_NONE) {
                     compessionHeader->CompressionType = EFI_NOT_COMPRESSED;
-                else if (model->compression(index) == COMPRESSION_ALGORITHM_LZMA || model->compression(index) == COMPRESSION_ALGORITHM_IMLZMA)
-                    compessionHeader->CompressionType = EFI_CUSTOMIZED_COMPRESSION;
-                else if (model->compression(index) == COMPRESSION_ALGORITHM_EFI11 || model->compression(index) == COMPRESSION_ALGORITHM_TIANO)
+                }
+                else if (model->compression(index) == COMPRESSION_ALGORITHM_EFI11 || model->compression(index) == COMPRESSION_ALGORITHM_TIANO) {
                     compessionHeader->CompressionType = EFI_STANDARD_COMPRESSION;
+                }
+                else if (model->compression(index) == COMPRESSION_ALGORITHM_LZMA) {
+                    compessionHeader->CompressionType = EFI_CUSTOMIZED_COMPRESSION;
+                }
+                else if (model->compression(index) == COMPRESSION_ALGORITHM_IMLZMA) {
+                    compessionHeader->CompressionType = EFI_CUSTOMIZED_COMPRESSION;
+                }
                 else
                     return ERR_UNKNOWN_COMPRESSION_ALGORITHM;
 
+                // Compress new section body
+                QByteArray compressed;
+                result = compress(reconstructed, model->compression(index), model->dictionarySize(index), compressed);
+                if (result)
+                    return result;
+
                 // Replace new section body
                 reconstructed = compressed;
             }
             else if (model->subtype(index) == EFI_SECTION_GUID_DEFINED) {
                 EFI_GUID_DEFINED_SECTION* guidDefinedHeader = (EFI_GUID_DEFINED_SECTION*)header.data();
                 // Compress new section body
                 QByteArray compressed;
-                result = compress(reconstructed, model->compression(index), compressed);
+                result = compress(reconstructed, model->compression(index), model->dictionarySize(index), compressed);
                 if (result)
                     return result;
                 // Check for authentication status valid attribute

diff --git a/ffsengine.h b/ffsengine.h
@@ -76,7 +76,7 @@ class FfsEngine : public QObject
 
     // Compression routines
     UINT8 decompress(const QByteArray & compressed, const UINT8 compressionType, QByteArray & decompressedData, UINT8 * algorithm = NULL);
-    UINT8 compress(const QByteArray & data, const UINT8 algorithm, QByteArray & compressedData);
+    UINT8 compress(const QByteArray & data, const UINT8 algorithm, const UINT32 dictionarySize, QByteArray & compressedData);
 
     // Construction routines
     UINT8 reconstructImageFile(QByteArray &reconstructed);

diff --git a/treeitem.cpp b/treeitem.cpp
@@ -23,6 +23,7 @@ TreeItem::TreeItem(const UINT8 type, const UINT8 subtype, const UINT8 compressio
     itemType(type),
     itemSubtype(subtype),
     itemCompression(compression),
+    itemDictionarySize(0),
     itemName(name),
     itemText(text),
     itemInfo(info),
@@ -198,6 +199,16 @@ UINT8 TreeItem::action() const
     return itemAction;
 }
 
+UINT32 TreeItem::dictionarySize() const
+{
+    return itemDictionarySize;
+}
+
+void TreeItem::setDictionarySize(const UINT32 dictionarySize)
+{
+    itemDictionarySize = dictionarySize;
+}
+
 void TreeItem::setAction(const UINT8 action)
 {
     itemAction = action;

diff --git a/treeitem.h b/treeitem.h
@@ -72,12 +72,16 @@ class TreeItem
 
     UINT8 compression() const;
 
+    UINT32 dictionarySize() const;
+    void setDictionarySize(const UINT32 dictionarySize);
+
 private:
     QList<TreeItem*> childItems;
     UINT8      itemAction;
     UINT8      itemType;
     UINT8      itemSubtype;
     UINT8      itemCompression;
+    UINT32     itemDictionarySize;
     QString    itemName;
     QString    itemText;
     QString    itemInfo;

diff --git a/treemodel.cpp b/treemodel.cpp
@@ -218,6 +218,14 @@ UINT8 TreeModel::compression(const QModelIndex &index) const
     return item->compression();
 }
 
+UINT32 TreeModel::dictionarySize(const QModelIndex &index) const
+{
+    if (!index.isValid())
+        return 0;
+    TreeItem *item = static_cast<TreeItem*>(index.internalPointer());
+    return item->dictionarySize();
+}
+
 void TreeModel::setSubtype(const QModelIndex & index, const UINT8 subtype)
 {
     if (!index.isValid())
@@ -268,6 +276,16 @@ void TreeModel::setAction(const QModelIndex &index, const UINT8 action)
     emit dataChanged(this->index(0, 0), index);
 }
 
+void TreeModel::setDictionarySize(const QModelIndex &index, const UINT32 dictionarySize)
+{
+    if (!index.isValid())
+        return;
+
+    TreeItem *item = static_cast<TreeItem*>(index.internalPointer());
+    item->setDictionarySize(dictionarySize);
+    emit dataChanged(this->index(0, 0), index);
+}
+
 QModelIndex TreeModel::addItem(const UINT8 type, const UINT8 subtype, const UINT8 compression,
     const QString & name, const QString & text, const QString & info,
     const QByteArray & header, const QByteArray & body, const QModelIndex & parent, const UINT8 mode)

diff --git a/treemodel.h b/treemodel.h
@@ -47,7 +47,7 @@ class TreeModel : public QAbstractItemModel
     void setSubtype(const QModelIndex &index, const UINT8 subtype);
     void setName(const QModelIndex &index, const QString &name);
     void setText(const QModelIndex &index, const QString &text);
-    void setParsingData(const QModelIndex &index, const QByteArray &data);
+    void setDictionarySize(const QModelIndex &index, const UINT32 dictionarySize);
 
     QString name(const QModelIndex &index) const;
     QString text(const QModelIndex &index) const;
@@ -58,10 +58,9 @@ class TreeModel : public QAbstractItemModel
     bool hasEmptyHeader(const QModelIndex &index) const;
     QByteArray body(const QModelIndex &index) const;
     bool hasEmptyBody(const QModelIndex &index) const;
-    QByteArray parsingData(const QModelIndex &index) const;
-    bool hasEmptyParsingData(const QModelIndex &index) const;
     UINT8 action(const QModelIndex &index) const;
     UINT8 compression(const QModelIndex &index) const;
+    UINT32 dictionarySize(const QModelIndex &index) const;
 
     QModelIndex addItem(const UINT8 type, const UINT8 subtype = 0, const UINT8 compression = COMPRESSION_ALGORITHM_NONE,
         const QString & name = QString(), const QString & text = QString(), const QString & info = QString(),

diff --git a/version.h b/version.h
@@ -11,9 +11,9 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 
 */
 
-#ifndef __VERSION_H__
-#define __VERSION_H__
+#ifndef VERSION_H
+#define VERSION_H
 
-#define PROGRAM_VERSION "0.25.1"
+#define PROGRAM_VERSION "0.26.0"
 
-#endif
+#endif // VERSION_H