diff --git a/BaseTools/Source/Python/FMMT/core/FMMTOperation.py b/BaseTools/Source/Python/FMMT/core/FMMTOperation.py index 4ed976d5fb9..76cf0c44fb4 100644 --- a/BaseTools/Source/Python/FMMT/core/FMMTOperation.py +++ b/BaseTools/Source/Python/FMMT/core/FMMTOperation.py @@ -4,6 +4,8 @@ # Copyright (c) 2021-, Intel Corporation. All rights reserved.
# 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!!!') diff --git a/BaseTools/Source/Python/FMMT/core/FvHandler.py b/BaseTools/Source/Python/FMMT/core/FvHandler.py index a1bd4685c67..c82fa4f0d12 100644 --- a/BaseTools/Source/Python/FMMT/core/FvHandler.py +++ b/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) diff --git a/BaseTools/Tests/FMMTTests.py b/BaseTools/Tests/FMMTTests.py new file mode 100644 index 00000000000..8905b4361e8 --- /dev/null +++ b/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()) diff --git a/BaseTools/Tests/PythonToolsTests.py b/BaseTools/Tests/PythonToolsTests.py index 05b27ab0333..4c26d1809dc 100644 --- a/BaseTools/Tests/PythonToolsTests.py +++ b/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__':