From 35c207dd8c31215fe11b56dbe2542674ccee6b6b Mon Sep 17 00:00:00 2001
From: vit9696 <vit9696@users.noreply.github.com>
Date: Sat, 23 Jun 2018 19:46:56 +0300
Subject: [PATCH] More proper solution to #137

---
 ffsengine.cpp | 26 ++++++++++++++++++++++++++
 treeitem.cpp  | 18 +-----------------
 2 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/ffsengine.cpp b/ffsengine.cpp
index ea7c172..b9e9c26 100644
--- a/ffsengine.cpp
+++ b/ffsengine.cpp
@@ -2453,6 +2453,32 @@ void FfsEngine::rebasePeiFiles(const QModelIndex & index)
             }
         }
     }
+
+    // Rebase VTF in subsequent volumes.
+    QModelIndex parent = index.parent();
+    while (parent.isValid() && model->type(parent) != Types::Volume)
+        parent = parent.parent();
+    if (parent.isValid()) {
+        QModelIndex volumeContainer = parent.parent();
+        // Iterate over volumes starting from the one after.
+        for (int i = parent.row() + 1; i < model->rowCount(volumeContainer); i++) {
+            QModelIndex currentVolumeIndex = volumeContainer.child(i, 0);
+            // Iterate over files within each volume after the current one.
+            for (int j = 0; j < model->rowCount(currentVolumeIndex); j++) {
+                QModelIndex currentFileIndex = currentVolumeIndex.child(j, 0);
+                if (model->header(currentFileIndex).left(sizeof(EFI_GUID)) == EFI_FFS_VOLUME_TOP_FILE_GUID) {
+                    for (int k = 0; k < model->rowCount(currentFileIndex); k++) {
+                        QModelIndex currentSectionIndex = currentFileIndex.child(k, 0);
+                        // If section stores PE32 or TE image
+                        if (model->subtype(currentSectionIndex) == EFI_SECTION_PE32 || model->subtype(currentSectionIndex) == EFI_SECTION_TE)
+                            // Set rebase action
+                            if (model->action(currentSectionIndex) != Actions::Remove)
+                                model->setAction(currentSectionIndex, Actions::Rebase);
+                    }
+                }
+            }
+        }
+    }
 }
 
 UINT8 FfsEngine::insert(const QModelIndex & index, const QByteArray & object, const UINT8 mode)
diff --git a/treeitem.cpp b/treeitem.cpp
index cdee8c2..2c3cc34 100644
--- a/treeitem.cpp
+++ b/treeitem.cpp
@@ -209,22 +209,6 @@ void TreeItem::setAction(const UINT8 action)
 
     // Set rebuild action for parent, if it has no action now
     if (parentItem && parentItem->type() != Types::Root
-        && parentItem->action() == Actions::NoAction) {
+        && parentItem->action() == Actions::NoAction)
         parentItem->setAction(Actions::Rebuild);
-
-        // Set rebuild action for subsequent items after parent.
-        // This is a little ugly, but fixes UEFIReplace image corruption,
-        // where one cannot manually choose the targets for rebuild.
-        // See: https://github.com/LongSoft/UEFITool/issues/137
-        TreeItem *grandParent = parentItem->parentItem;
-        if (grandParent) {
-            QList<TreeItem *> &parentCousins = grandParent->childItems;
-            int count = parentCousins.count();
-            for (int i = parentCousins.indexOf(parentItem) + 1; i < count; i++) {
-                TreeItem *parentCousin = parentCousins.value(i, NULL);
-                if (parentCousin && parentCousin->action() == Actions::NoAction)
-                    parentCousin->setAction(Actions::Rebuild);
-            }
-        }
-    }
 }