binman: Update image positions of FIT subentries

Binman keeps track of positions of each entry in the final image, but
currently this data is wrong for things included in FIT entries,
especially since a previous patch makes FIT a subclass of Section and
inherit its implementation.

There are three ways to put data into a FIT image. It can be directly
included as a "data" property, or it can be external to the FIT image
represented by an offset-size pair of properties. This external offset
is either "data-position" from the start of the FIT or "data-offset"
from the end of the FIT, and the size is "data-size" for both. However,
binman doesn't use the "data-offset" method while building FIT entries.

According to the Section docstring, its subclasses should calculate and
set the correct offsets and sizes in SetImagePos() method. Do this for
FIT subentries for the three ways mentioned above, and add tests for the
two ways binman can pack them in.

Signed-off-by: Alper Nebi Yasak <alpernebiyasak@gmail.com>
Reviewed-by: Simon Glass <sjg@chromium.org>

2 files changed, 163 insertions, 0 deletions

diff --git a/tools/binman/etype/fit.py b/tools/binman/etype/fit.py
index cd7ebc571e..1e957023f3 100644
--- a/tools/binman/etype/fit.py
+++ b/tools/binman/etype/fit.py
@@ -293,6 +293,57 @@ class Entry_fit(Entry_section):
         data = fdt.GetContents()
         return data
 
+    def SetImagePos(self, image_pos):
+        """Set the position in the image
+
+        This sets each subentry's offsets, sizes and positions-in-image
+        according to where they ended up in the packed FIT file.
+
+        Args:
+            image_pos: Position of this entry in the image
+        """
+        super().SetImagePos(image_pos)
+
+        # If mkimage is missing we'll have empty data,
+        # which will cause a FDT_ERR_BADMAGIC error
+        if self.mkimage in self.missing_bintools:
+            return
+
+        fdt = Fdt.FromData(self.GetData())
+        fdt.Scan()
+
+        for path, section in self._entries.items():
+            node = fdt.GetNode(path)
+
+            data_prop = node.props.get("data")
+            data_pos = fdt_util.GetInt(node, "data-position")
+            data_offset = fdt_util.GetInt(node, "data-offset")
+            data_size = fdt_util.GetInt(node, "data-size")
+
+            # Contents are inside the FIT
+            if data_prop is not None:
+                # GetOffset() returns offset of a fdt_property struct,
+                # which has 3 fdt32_t members before the actual data.
+                offset = data_prop.GetOffset() + 12
+                size = len(data_prop.bytes)
+
+            # External offset from the base of the FIT
+            elif data_pos is not None:
+                offset = data_pos
+                size = data_size
+
+            # External offset from the end of the FIT, not used in binman
+            elif data_offset is not None: # pragma: no cover
+                offset = fdt.GetFdtObj().totalsize() + data_offset
+                size = data_size
+
+            # This should never happen
+            else: # pragma: no cover
+                self.Raise("%s: missing data properties" % (path))
+
+            section.SetOffsetSize(offset, size)
+            section.SetImagePos(self.image_pos)
+
     def AddBintools(self, tools):
         super().AddBintools(tools)
         self.mkimage = self.AddBintool(tools, 'mkimage')
diff --git a/tools/binman/ftest.py b/tools/binman/ftest.py
index b6801b7275..f8e73c6f9b 100644
--- a/tools/binman/ftest.py
+++ b/tools/binman/ftest.py
@@ -3770,6 +3770,62 @@ class TestFunctional(unittest.TestCase):
 
         self._CheckSimpleFitData(fit_data, U_BOOT_EXP_DATA, U_BOOT_SPL_DTB_DATA)
 
+    def testSimpleFitImagePos(self):
+        """Test that we have correct image-pos for FIT subentries"""
+        data, _, _, out_dtb_fname = self._DoReadFileDtb('161_fit.dts',
+                                                        update_dtb=True)
+        dtb = fdt.Fdt(out_dtb_fname)
+        dtb.Scan()
+        props = self._GetPropTree(dtb, BASE_DTB_PROPS + REPACK_DTB_PROPS)
+
+        self.assertEqual({
+            'image-pos': 0,
+            'offset': 0,
+            'size': 1890,
+
+            'u-boot:image-pos': 0,
+            'u-boot:offset': 0,
+            'u-boot:size': 4,
+
+            'fit:image-pos': 4,
+            'fit:offset': 4,
+            'fit:size': 1840,
+
+            'fit/images/kernel:image-pos': 160,
+            'fit/images/kernel:offset': 156,
+            'fit/images/kernel:size': 4,
+
+            'fit/images/kernel/u-boot:image-pos': 160,
+            'fit/images/kernel/u-boot:offset': 0,
+            'fit/images/kernel/u-boot:size': 4,
+
+            'fit/images/fdt-1:image-pos': 456,
+            'fit/images/fdt-1:offset': 452,
+            'fit/images/fdt-1:size': 6,
+
+            'fit/images/fdt-1/u-boot-spl-dtb:image-pos': 456,
+            'fit/images/fdt-1/u-boot-spl-dtb:offset': 0,
+            'fit/images/fdt-1/u-boot-spl-dtb:size': 6,
+
+            'u-boot-nodtb:image-pos': 1844,
+            'u-boot-nodtb:offset': 1844,
+            'u-boot-nodtb:size': 46,
+        }, props)
+
+        # Actually check the data is where we think it is
+        for node, expected in [
+            ("u-boot", U_BOOT_DATA),
+            ("fit/images/kernel", U_BOOT_DATA),
+            ("fit/images/kernel/u-boot", U_BOOT_DATA),
+            ("fit/images/fdt-1", U_BOOT_SPL_DTB_DATA),
+            ("fit/images/fdt-1/u-boot-spl-dtb", U_BOOT_SPL_DTB_DATA),
+            ("u-boot-nodtb", U_BOOT_NODTB_DATA),
+        ]:
+            image_pos = props[f"{node}:image-pos"]
+            size = props[f"{node}:size"]
+            self.assertEqual(len(expected), size)
+            self.assertEqual(expected, data[image_pos:image_pos+size])
+
     def testFitExternal(self):
         """Test an image with an FIT with external images"""
         data = self._DoReadFile('162_fit_external.dts')
@@ -3798,6 +3854,62 @@ class TestFunctional(unittest.TestCase):
         self.assertEqual(U_BOOT_DATA + b'aa',
                          data[actual_pos:actual_pos + external_data_size])
 
+    def testFitExternalImagePos(self):
+        """Test that we have correct image-pos for external FIT subentries"""
+        data, _, _, out_dtb_fname = self._DoReadFileDtb('162_fit_external.dts',
+                                                        update_dtb=True)
+        dtb = fdt.Fdt(out_dtb_fname)
+        dtb.Scan()
+        props = self._GetPropTree(dtb, BASE_DTB_PROPS + REPACK_DTB_PROPS)
+
+        self.assertEqual({
+            'image-pos': 0,
+            'offset': 0,
+            'size': 1082,
+
+            'u-boot:image-pos': 0,
+            'u-boot:offset': 0,
+            'u-boot:size': 4,
+
+            'fit:size': 1032,
+            'fit:offset': 4,
+            'fit:image-pos': 4,
+
+            'fit/images/kernel:size': 4,
+            'fit/images/kernel:offset': 1024,
+            'fit/images/kernel:image-pos': 1028,
+
+            'fit/images/kernel/u-boot:size': 4,
+            'fit/images/kernel/u-boot:offset': 0,
+            'fit/images/kernel/u-boot:image-pos': 1028,
+
+            'fit/images/fdt-1:size': 2,
+            'fit/images/fdt-1:offset': 1028,
+            'fit/images/fdt-1:image-pos': 1032,
+
+            'fit/images/fdt-1/_testing:size': 2,
+            'fit/images/fdt-1/_testing:offset': 0,
+            'fit/images/fdt-1/_testing:image-pos': 1032,
+
+            'u-boot-nodtb:image-pos': 1036,
+            'u-boot-nodtb:offset': 1036,
+            'u-boot-nodtb:size': 46,
+         }, props)
+
+        # Actually check the data is where we think it is
+        for node, expected in [
+            ("u-boot", U_BOOT_DATA),
+            ("fit/images/kernel", U_BOOT_DATA),
+            ("fit/images/kernel/u-boot", U_BOOT_DATA),
+            ("fit/images/fdt-1", b'aa'),
+            ("fit/images/fdt-1/_testing", b'aa'),
+            ("u-boot-nodtb", U_BOOT_NODTB_DATA),
+        ]:
+            image_pos = props[f"{node}:image-pos"]
+            size = props[f"{node}:size"]
+            self.assertEqual(len(expected), size)
+            self.assertEqual(expected, data[image_pos:image_pos+size])
+
     def testFitMissing(self):
         """Test that binman still produces a FIT image if mkimage is missing"""
         with test_util.capture_sys_output() as (_, stderr):