[PATCH v2 5/5] binman: Update image positions of FIT subentries

[Alper Nebi Yasak <alpernebiyasak@gmail.com>]

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>
---

Changes in v2:
- Check missing_bintools list instead of catching Fdt exceptions
- Add tag: "Reviewed-by: Simon Glass <sjg@chromium.org>"

 tools/binman/etype/fit.py |  51 +++++++++++++++++
 tools/binman/ftest.py     | 112 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 163 insertions(+)

diff --git a/tools/binman/etype/fit.py b/tools/binman/etype/fit.py
index 0f8c88a9720e..5497f036a26d 100644
--- a/tools/binman/etype/fit.py
+++ b/tools/binman/etype/fit.py
@@ -291,6 +291,57 @@ def _BuildInput(self, fdt):
         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 f16798960b84..ab988565c335 100644
--- a/tools/binman/ftest.py
+++ b/tools/binman/ftest.py
@@ -3770,6 +3770,62 @@ def testSimpleFitExpandsSubentries(self):
 
         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 @@ def testFitExternal(self):
         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):
-- 
2.34.1