feat(tlc): add --align argument

Extended the command line interface to receive an alignment
argument.

TLC tool will align the data of the TEs accordingly.

Signed-off-by: J-Alves <joao.alves@arm.com>
Change-Id: I281b0b4c1851d58377bf6b31fcee03ee2f53367b

tools/tlc/tests/test_cli.py

@@ -17,6 +17,7 @@ from unittest import mock
 import pytest
 import yaml
 from click.testing import CliRunner
+from conftest import generate_random_bytes
 
 from tlc.cli import cli
 from tlc.te import TransferEntry
@@ -32,6 +33,22 @@ def test_create_empty_tl(tmpdir):
     assert TransferList.fromfile(test_file) is not None
 
 
+@pytest.mark.parametrize("align", [4, 6, 12, 13])
+def test_create_with_align(align, tlcrunner, tmpdir):
+    tl_file = tmpdir.join("tl.bin").strpath
+    tlcrunner.invoke(cli, ["create", "-s", "10000", "-a", align, tl_file])
+
+    blob = tmpdir.join("blob.bin")
+
+    blob.write_binary(generate_random_bytes(0x200))
+    tlcrunner.invoke(cli, ["add", "--entry", 1, blob.strpath, tl_file])
+
+    tl = TransferList.fromfile(tl_file)
+    te = tl.entries[-1]
+    assert tl.alignment == align
+    assert (te.offset + te.hdr_size) % (1 << align) == 0
+
+
 def test_create_with_fdt(tmpdir):
     runner = CliRunner()
     fdt = tmpdir.join("fdt.dtb")
@@ -69,6 +86,20 @@ def test_add_multiple_entries(tlcrunner, tlc_entries, tmptlstr):
     assert len(tl.entries) == len(tlc_entries)
 
 
+@pytest.mark.parametrize("align", [4, 6, 12, 13])
+def test_cli_add_entry_with_align(align, tlcrunner, tmpdir, tmptlstr):
+    blob = tmpdir.join("blob.bin")
+    blob.write_binary(bytes(0x100))
+
+    tlcrunner.invoke(cli, ["add", "--align", align, "--entry", 1, blob, tmptlstr])
+    tl = TransferList.fromfile(tmptlstr)
+    te = tl.entries[-1]
+
+    print(tl, *(te for te in tl.entries), sep="\n---------------\n")
+    assert (te.offset + te.hdr_size) % (1 << align) == 0
+    assert tl.alignment == align
+
+
 def test_info(tlcrunner, tmptlstr, tmpfdt):
     tlcrunner.invoke(cli, ["add", "--entry", "0", "/dev/null", tmptlstr])
     tlcrunner.invoke(cli, ["add", "--fdt", tmpfdt.strpath, tmptlstr])

tools/tlc/tests/test_transfer_list.py

@@ -157,7 +157,7 @@ def test_single_te_transfer_list(tag_id, data, tmpdir):
         assert f.read(te.data_size) == te.data
 
 
-def test_write_multiple_tes_to_file(tmpdir, random_entries):
+def test_write_multiple_tes_to_file(tmpdir, random_entries, random_entry):
     """Check that we can create a TL with multiple TE's."""
     test_file = tmpdir.join("test_tl_blob.bin")
     tl = TransferList(0x4000)
@@ -166,6 +166,10 @@ def test_write_multiple_tes_to_file(tmpdir, random_entries):
     for tag_id, data in _test_entries:
         tl.add_transfer_entry(tag_id, data)
 
+    # Add a few entries with special alignment requirements
+    blob_id, blob = random_entry(0x200)
+    tl.add_transfer_entry(blob_id, blob, data_align=12)
+
     tl.write_to_file(test_file)
 
     with open(test_file, "rb") as f:
@@ -180,6 +184,13 @@ def test_write_multiple_tes_to_file(tmpdir, random_entries):
             data_size = int.from_bytes(f.read(4), "little")
             assert f.read(data_size) == data
 
+        f.seek(int(math.ceil(f.tell() / (1 << 12)) * (1 << 12)) - 8)
+        assert int.from_bytes(f.read(3), "little") == blob_id
+        assert int.from_bytes(f.read(1), "little") == TransferEntry.hdr_size
+        # Make sure the data in the TE matches the data in the original case
+        data_size = int.from_bytes(f.read(4), "little")
+        assert f.read(data_size) == blob
+
         # padding is added to align TE's, make sure padding is added to the size of
         # the TL by checking we don't overflow.
         assert f.tell() <= tl.size

tools/tlc/tlc/cli.py

@@ -26,6 +26,14 @@ def cli():
 
 @cli.command()
 @click.argument("filename", type=click.Path(dir_okay=False))
+@click.option(
+    "-a",
+    "--align",
+    type=int,
+    default=3,
+    show_default=True,
+    help="Set alignment in powers of 2 (e.g., -a 3 for 8 byte alignment).",
+)
 @click.option(
     "-s", "--size", default=0x1000, type=int, help="Maximum size of the Transfer List"
 )
@@ -51,7 +59,7 @@ def cli():
     type=click.Path(exists=True),
     help="Create the transfer list from a YAML config file.",
 )
-def create(filename, size, fdt, entry, flags, from_yaml):
+def create(filename, align, size, fdt, entry, flags, from_yaml):
     """Create a new Transfer List."""
     try:
         if from_yaml:
@@ -60,12 +68,12 @@ def create(filename, size, fdt, entry, flags, from_yaml):
 
             tl = TransferList.from_dict(config)
         else:
-            tl = TransferList(size)
+            tl = TransferList(size, alignment=align)
 
             entry = (*entry, (1, fdt)) if fdt else entry
 
             for id, path in entry:
-                tl.add_transfer_entry_from_file(id, path)
+                tl.add_transfer_entry_from_file(id, path, data_align=align)
     except MemoryError as mem_excp:
         raise MemoryError(
             "TL max size exceeded, consider increasing with the option -s"
@@ -133,19 +141,24 @@ def remove(filename, tags):
 
 
 @cli.command()
-@click.argument("filename", type=click.Path(exists=True, dir_okay=False))
+@click.option(
+    "-a",
+    "--align",
+    type=int,
+    help="Set alignment in powers of 2 (e.g., -a 3 for 8 byte alignment).",
+)
 @click.option(
     "--entry",
     type=(int, click.Path(exists=True)),
     multiple=True,
     help="A tag ID and the corresponding path to a binary blob in the form <id> <path-to-blob>.",
 )
-def add(filename, entry):
+@click.argument("filename", type=click.Path(exists=True, dir_okay=False))
+def add(align, entry, filename):
     """Update an existing Transfer List with given images."""
     tl = TransferList.fromfile(filename)
-
     for id, path in entry:
-        tl.add_transfer_entry_from_file(id, path)
+        tl.add_transfer_entry_from_file(id, path, data_align=align)
 
     tl.write_to_file(filename)
 

tools/tlc/tlc/tl.py

@@ -86,11 +86,14 @@ class TransferList:
     granule = 8
 
     def __init__(
-        self, max_size: int = hdr_size, flags: int = TRANSFER_LIST_ENABLE_CHECKSUM
+        self,
+        max_size: int = hdr_size,
+        flags: int = TRANSFER_LIST_ENABLE_CHECKSUM,
+        alignment: int = 3,
     ) -> None:
         assert max_size >= self.hdr_size
         self.checksum: int = 0
-        self.alignment: int = 3
+        self.alignment: int = alignment
         self.size = self.hdr_size
         self.total_size = max_size
         self.flags = flags
@@ -163,10 +166,14 @@ class TransferList:
         # get settings from config and set defaults
         max_size = config.get("max_size", 0x1000)
         has_checksum = config.get("has_checksum", True)
+        align = config.get("alignment", None)
 
         flags = TRANSFER_LIST_ENABLE_CHECKSUM if has_checksum else 0
 
-        tl = cls(max_size, flags)
+        if align:
+            tl = cls(max_size, flags, alignment=align)
+        else:
+            tl = cls(max_size, flags)
 
         for entry in config["entries"]:
             tl.add_transfer_entry_from_dict(entry)
@@ -327,8 +334,12 @@ class TransferList:
         te_format = transfer_entry_formats[tag_id]
         tag_name = te_format["tag_name"]
 
+        align = entry.get("alignment", None)
+
         if "blob_file_path" in entry:
-            return self.add_transfer_entry_from_file(tag_id, entry["blob_file_path"])
+            return self.add_transfer_entry_from_file(
+                tag_id, entry["blob_file_path"], data_align=align
+            )
         elif tag_name == "tpm_event_log_table":
             with open(entry["event_log"], "rb") as f:
                 event_log_data = f.read()
@@ -336,7 +347,7 @@ class TransferList:
             flags_bytes = entry["flags"].to_bytes(4, "little")
             data = flags_bytes + event_log_data
 
-            return self.add_transfer_entry(tag_id, data)
+            return self.add_transfer_entry(tag_id, data, data_align=align)
         elif tag_name == "exec_ep_info":
             return self.add_entry_point_info_transfer_entry(entry)
         elif "format" in te_format and "fields" in te_format:
@@ -347,9 +358,11 @@ class TransferList:
         else:
             raise ValueError(f"Invalid transfer entry {entry}.")
 
-    def add_transfer_entry_from_file(self, tag_id: int, path: Path) -> TransferEntry:
+    def add_transfer_entry_from_file(
+        self, tag_id: int, path: Path, data_align: int = 0
+    ) -> TransferEntry:
         with open(path, "rb") as f:
-            return self.add_transfer_entry(tag_id, f.read())
+            return self.add_transfer_entry(tag_id, f.read(), data_align=data_align)
 
     def write_to_file(self, file: Path) -> None:
         """Write the contents of the TL to a file."""