[CHERRY-PICK][Rebase & FF] Pick up FMMT related fixes from edk2

BaseTools/Source/Python/FMMT/core/FMMTOperation.py

@@ -4,6 +4,8 @@
 # Copyright (c) 2021-, Intel Corporation. All rights reserved.<BR>
 # SPDX-License-Identifier: BSD-2-Clause-Patent
 ##
+import copy
+
 from core.FMMTParser import *
 from core.FvHandler import *
 from utils.FvLayoutPrint import *
@@ -162,7 +164,8 @@ def ReplaceFfs(inputfile: str, Ffs_name: str, newffsfile: str, outputfile: str,
                 FmmtParser.WholeFvTree.Findlist.remove(FmmtParser.WholeFvTree.Findlist[index])
     if FmmtParser.WholeFvTree.Findlist != []:
         for TargetFfs in FmmtParser.WholeFvTree.Findlist:
-            FfsMod = FvHandler(newFmmtParser.WholeFvTree.Child[0], TargetFfs)
+            NewFfs = copy.deepcopy(newFmmtParser.WholeFvTree.Child[0])
+            FfsMod = FvHandler(NewFfs, TargetFfs)
             Status = FfsMod.ReplaceFfs()
     else:
         logger.error('Target Ffs not found!!!')

BaseTools/Source/Python/FMMT/core/FvHandler.py

@@ -145,6 +145,25 @@ def ModifyFvSystemGuid(TargetFv) -> None:
         else:
             TargetFv.Data.Data += struct2stream(item.Data.Header)+ item.Data.Data + item.Data.PadData
 
+def SetFvFreeSpace(TargetFv, NewFreeSpace: int) -> None:
+    FreeSpaceTree = (
+        TargetFv.Child[-1]
+        if TargetFv.Child[-1].type == FFS_FREE_SPACE
+        else None
+    )
+    if NewFreeSpace:
+        if FreeSpaceTree is None:
+            FreeSpaceTree = BIOSTREE('FREE_SPACE')
+            FreeSpaceTree.type = FFS_FREE_SPACE
+            FreeSpaceTree.Data = FreeSpaceNode(b'\xff' * NewFreeSpace)
+            TargetFv.insertChild(FreeSpaceTree)
+        else:
+            FreeSpaceTree.Data.Data = b'\xff' * NewFreeSpace
+            FreeSpaceTree.Data.Size = NewFreeSpace
+    elif FreeSpaceTree is not None:
+        TargetFv.Child.remove(FreeSpaceTree)
+    TargetFv.Data.Free_Space = NewFreeSpace
+
 class FvHandler:
     def __init__(self, NewFfs, TargetFfs=None) -> None:
         self.NewFfs = NewFfs
@@ -283,12 +302,10 @@ def ModifyTest(self, ParTree, Needed_Space: int) -> None:
                         BlockSize = ParTree.Data.Header.BlockMap[0].Length
                         New_Add_Len = BlockSize - Needed_Space%BlockSize
                         if New_Add_Len % BlockSize:
-                            ParTree.Child[-1].Data.Data = b'\xff' * New_Add_Len
-                            ParTree.Data.Free_Space = New_Add_Len
+                            SetFvFreeSpace(ParTree, New_Add_Len)
                             Needed_Space += New_Add_Len
                         else:
-                            ParTree.Child.remove(ParTree.Child[-1])
-                            ParTree.Data.Free_Space = 0
+                            SetFvFreeSpace(ParTree, 0)
                         ParTree.Data.Size += Needed_Space
                         ParTree.Data.Header.FvLength = ParTree.Data.Size
                 ModifyFvSystemGuid(ParTree)
@@ -451,15 +468,14 @@ def ReplaceFfs(self) -> bool:
                     New_Add_Len = BlockSize - Needed_Space%BlockSize
                     Target_index = TargetFv.Child.index(self.TargetFfs)
                     if New_Add_Len % BlockSize:
-                        TargetFv.Child[-1].Data.Data = b'\xff' * New_Add_Len
-                        TargetFv.Data.Free_Space = New_Add_Len
+                        SetFvFreeSpace(TargetFv, New_Add_Len)
                         Needed_Space += New_Add_Len
                         TargetFv.insertChild(self.NewFfs, Target_index)
                         TargetFv.Child.remove(self.TargetFfs)
                     else:
+                        SetFvFreeSpace(TargetFv, 0)
                         TargetFv.Child.remove(self.TargetFfs)
-                        TargetFv.Data.Free_Space = 0
-                        TargetFv.insertChild(self.NewFfs)
+                        TargetFv.insertChild(self.NewFfs, Target_index)
                     # Encapsulate the Fv Data for update.
                     TargetFv.Data.Data = b''
                     for item in TargetFv.Child:
@@ -488,11 +504,7 @@ def ReplaceFfs(self) -> bool:
                 TargetFv.insertChild(self.NewFfs, Target_index)
             # If TargetFv do not have free space, create free space for Fv.
             else:
-                New_Free_Space_Tree = BIOSTREE('FREE_SPACE')
-                New_Free_Space_Tree.type = FFS_FREE_SPACE
-                New_Free_Space_Tree.Data = FfsNode(b'\xff' * New_Free_Space)
-                TargetFv.Data.Free_Space = New_Free_Space
-                TargetFv.insertChild(New_Free_Space)
+                SetFvFreeSpace(TargetFv, New_Free_Space)
                 Target_index = TargetFv.Child.index(self.TargetFfs)
                 TargetFv.Child.remove(self.TargetFfs)
                 TargetFv.insertChild(self.NewFfs, Target_index)

BaseTools/Tests/FMMTTests.py

@@ -0,0 +1,289 @@
+## @file
+# Unit tests for the Firmware Module Management Tool.
+#
+# Copyright (c) Microsoft Corporation.
+#
+# SPDX-License-Identifier: BSD-2-Clause-Patent
+##
+
+import logging
+import os
+import sys
+import tempfile
+import unittest
+import uuid
+from pathlib import Path
+from types import SimpleNamespace
+from unittest import mock
+
+import TestTools
+
+
+_TESTS_DIR = Path(__file__).resolve().parent
+_PYTHON_SOURCE = _TESTS_DIR.parent / "Source" / "Python"
+_FMMT_SOURCE = _PYTHON_SOURCE / "FMMT"
+sys.path.insert(0, str(_FMMT_SOURCE))
+sys.path.insert(0, str(_PYTHON_SOURCE))
+os.environ.setdefault("FmmtConfPath", "")
+
+# FMMT opens a build log when imported. Unit tests do not need a persistent log.
+with mock.patch("logging.FileHandler", return_value=logging.NullHandler()):
+    import core.FMMTOperation as FMMTOperation
+    from core.BiosTree import BIOSTREE
+    from core.BiosTree import FFS_FREE_SPACE
+    from core.BiosTree import FFS_TREE
+    from core.BiosTree import SEC_FV_TREE
+    from core.BiosTreeNode import FreeSpaceNode
+    from core.FvHandler import FvHandler
+    from FirmwareStorageFormat.Common import ModifyGuidFormat
+    from FirmwareStorageFormat.FfsFileHeader import EFI_FFS_FILE_HEADER
+
+
+class _FirmwareVolumeData:
+    def __init__(self, size, free_space=0, block_size=0x40):
+        self.Data = b""
+        self.Free_Space = free_space
+        self.Header = SimpleNamespace(
+            Attributes=0,
+            BlockMap=[SimpleNamespace(Length=block_size)],
+            ExtHeaderOffset=0,
+            FileSystemGuid=ModifyGuidFormat(
+                "00000000-0000-0000-0000-000000000000"
+            ),
+            FvLength=size,
+        )
+        self.Size = size
+
+    def ModCheckSum(self):
+        pass
+
+    def ModExtHeaderData(self):
+        pass
+
+    def ModFvExt(self):
+        pass
+
+    def ModFvSize(self):
+        pass
+
+
+class _TestFvHandler(FvHandler):
+    def CompressData(self, target_tree):
+        self.Status = True
+
+    def ModifyTest(self, parent_tree, needed_space):
+        self.Status = True
+
+
+class _TestRecursiveFvHandler(FvHandler):
+    def CompressData(self, target_tree):
+        self.Status = True
+
+
+def _create_ffs(name, size, data=b""):
+    node = BIOSTREE(name)
+    node.type = FFS_TREE
+    header = EFI_FFS_FILE_HEADER()
+    header.Name = ModifyGuidFormat(str(name))
+    header.Size[0] = size & 0xFF
+    header.Size[1] = (size >> 8) & 0xFF
+    header.Size[2] = (size >> 16) & 0xFF
+    node.Data = SimpleNamespace(
+        Data=data,
+        Header=header,
+        Name=name,
+        PadData=b"",
+        Size=size,
+    )
+    return node
+
+
+class FmmtReplaceTests(unittest.TestCase):
+    def test_replace_multiple_matches_uses_independent_ffs_trees(self):
+        target_guid = uuid.UUID("11111111-2222-3333-4444-555555555555")
+        targets = [SimpleNamespace(), SimpleNamespace()]
+        replacement = SimpleNamespace(
+            Data=SimpleNamespace(PadData=b"", Size=0x20),
+            Parent=None,
+        )
+
+        class _Root:
+            def __init__(self):
+                self.Findlist = []
+
+            def FindNode(self, name, matches):
+                matches.extend(targets)
+
+        class _Parser:
+            def __init__(self, root):
+                self.FinalData = b"output"
+                self.WholeFvTree = root
+
+            def Encapsulation(self, root, compress_status):
+                pass
+
+            def ParserFromRoot(self, root, data):
+                pass
+
+        input_parser = _Parser(_Root())
+        replacement_parser = _Parser(SimpleNamespace(Child=[replacement]))
+        replacements = []
+
+        class _Handler:
+            def __init__(self, new_ffs, target_ffs):
+                replacements.append(new_ffs)
+                new_ffs.Parent = target_ffs
+
+            def ReplaceFfs(self):
+                return True
+
+        with tempfile.TemporaryDirectory() as directory:
+            input_path = Path(directory) / "input.fd"
+            new_ffs_path = Path(directory) / "new.ffs"
+            output_path = Path(directory) / "output.fd"
+            input_path.write_bytes(b"input")
+            new_ffs_path.write_bytes(b"replacement")
+
+            with mock.patch.object(
+                FMMTOperation,
+                "FMMTParser",
+                side_effect=[input_parser, replacement_parser],
+            ), mock.patch.object(FMMTOperation, "FvHandler", _Handler):
+                FMMTOperation.ReplaceFfs(
+                    str(input_path),
+                    target_guid,
+                    str(new_ffs_path),
+                    str(output_path),
+                )
+
+        self.assertEqual(2, len(replacements))
+        self.assertIsNot(replacements[0], replacements[1])
+        self.assertIs(targets[0], replacements[0].Parent)
+        self.assertIs(targets[1], replacements[1].Parent)
+
+    def test_smaller_replacement_creates_free_space_node(self):
+        prefix = _create_ffs(
+            uuid.UUID("10000000-0000-0000-0000-000000000001"), 0x18
+        )
+        target = _create_ffs(
+            uuid.UUID("20000000-0000-0000-0000-000000000002"), 0x40
+        )
+        replacement = _create_ffs(target.Data.Name, 0x18)
+        volume = BIOSTREE("FV")
+        volume.type = SEC_FV_TREE
+        volume.Data = _FirmwareVolumeData(0x80)
+        volume.insertChild(prefix)
+        volume.insertChild(target)
+
+        self.assertTrue(_TestFvHandler(replacement, target).ReplaceFfs())
+
+        self.assertIs(replacement, volume.Child[1])
+        free_space = volume.Child[-1]
+        self.assertEqual(FFS_FREE_SPACE, free_space.type)
+        self.assertIsInstance(free_space.Data, FreeSpaceNode)
+        self.assertEqual(b"\xFF" * 0x28, free_space.Data.Data)
+
+    def test_growing_nested_fv_preserves_final_ffs(self):
+        prefix = _create_ffs(
+            uuid.UUID("10000000-0000-0000-0000-000000000001"), 0x18, b"prefix"
+        )
+        target = _create_ffs(
+            uuid.UUID("20000000-0000-0000-0000-000000000002"), 0x20, b"target"
+        )
+        trailing = _create_ffs(
+            uuid.UUID("30000000-0000-0000-0000-000000000003"),
+            0x18,
+            b"trailing",
+        )
+        replacement = _create_ffs(target.Data.Name, 0x30, b"replacement")
+        volume = BIOSTREE("FV")
+        volume.type = SEC_FV_TREE
+        volume.Data = _FirmwareVolumeData(0x80)
+        volume.insertChild(prefix)
+        volume.insertChild(target)
+        volume.insertChild(trailing)
+
+        self.assertTrue(_TestFvHandler(replacement, target).ReplaceFfs())
+
+        self.assertEqual(b"trailing", trailing.Data.Data)
+        self.assertIn(trailing, volume.Child)
+        self.assertIs(replacement, volume.Child[1])
+        free_space = volume.Child[-1]
+        self.assertEqual(FFS_FREE_SPACE, free_space.type)
+        self.assertIsInstance(free_space.Data, FreeSpaceNode)
+        self.assertEqual(b"\xFF" * 0x30, free_space.Data.Data)
+
+    def test_block_aligned_growth_of_only_ffs_does_not_crash(self):
+        target = _create_ffs(
+            uuid.UUID("20000000-0000-0000-0000-000000000002"), 0x20, b"target"
+        )
+        replacement = _create_ffs(target.Data.Name, 0x60, b"replacement")
+        volume = BIOSTREE("FV")
+        volume.type = SEC_FV_TREE
+        volume.Data = _FirmwareVolumeData(0x40)
+        volume.insertChild(target)
+
+        self.assertTrue(_TestFvHandler(replacement, target).ReplaceFfs())
+
+        self.assertEqual([replacement], volume.Child)
+        self.assertEqual(0, volume.Data.Free_Space)
+
+    def test_block_aligned_growth_preserves_ffs_order(self):
+        prefix = _create_ffs(
+            uuid.UUID("10000000-0000-0000-0000-000000000001"), 0x18, b"prefix"
+        )
+        target = _create_ffs(
+            uuid.UUID("20000000-0000-0000-0000-000000000002"), 0x20, b"target"
+        )
+        trailing = _create_ffs(
+            uuid.UUID("30000000-0000-0000-0000-000000000003"),
+            0x18,
+            b"trailing",
+        )
+        replacement = _create_ffs(target.Data.Name, 0x60, b"replacement")
+        volume = BIOSTREE("FV")
+        volume.type = SEC_FV_TREE
+        volume.Data = _FirmwareVolumeData(0x80)
+        volume.insertChild(prefix)
+        volume.insertChild(target)
+        volume.insertChild(trailing)
+
+        self.assertTrue(_TestFvHandler(replacement, target).ReplaceFfs())
+
+        self.assertEqual([prefix, replacement, trailing], volume.Child)
+        self.assertEqual(0, volume.Data.Free_Space)
+
+    def test_growing_ancestor_fv_preserves_final_ffs(self):
+        prefix = _create_ffs(
+            uuid.UUID("10000000-0000-0000-0000-000000000001"), 0x18, b"prefix"
+        )
+        trailing = _create_ffs(
+            uuid.UUID("30000000-0000-0000-0000-000000000003"),
+            0x18,
+            b"trailing",
+        )
+        volume = BIOSTREE("FV")
+        volume.type = SEC_FV_TREE
+        volume.Data = _FirmwareVolumeData(0x80)
+        volume.insertChild(prefix)
+        volume.insertChild(trailing)
+
+        parent = BIOSTREE("ROOT")
+        parent.type = FMMTOperation.ROOT_TREE
+        parent.insertChild(volume)
+
+        _TestRecursiveFvHandler(None).ModifyTest(volume, 0x10)
+
+        self.assertEqual(b"trailing", trailing.Data.Data)
+        self.assertIn(trailing, volume.Child)
+        free_space = volume.Child[-1]
+        self.assertEqual(FFS_FREE_SPACE, free_space.type)
+        self.assertIsInstance(free_space.Data, FreeSpaceNode)
+        self.assertEqual(b"\xFF" * 0x30, free_space.Data.Data)
+
+
+TheTestSuite = TestTools.MakeTheTestSuite(locals())
+
+
+if __name__ == "__main__":
+    unittest.TextTestRunner(verbosity=2).run(TheTestSuite())

BaseTools/Tests/PythonToolsTests.py

@@ -20,6 +20,8 @@ def TheTestSuite():
     suites.append(CheckPythonSyntax.TheTestSuite())
     import CheckUnicodeSourceFiles
     suites.append(CheckUnicodeSourceFiles.TheTestSuite())
+    import FMMTTests
+    suites.append(FMMTTests.TheTestSuite())
     return unittest.TestSuite(suites)
 
 if __name__ == '__main__':