binman: Support updating all device tree files

Binman currently supports updating the main device tree with things like
the position of each entry. Extend this support to SPL and TPL as well,
since they may need (a subset of) this information.

Also adjust DTB output files to have a .out extension since this seems
clearer than having a .dtb extension with 'out' in the name somwhere.

Also add a few missing comments and update the DT setup code to use
ReadFile and WriteFile().

Signed-off-by: Simon Glass <sjg@chromium.org>
diff --git a/tools/binman/README.entries b/tools/binman/README.entries
index 5cb52a9..091fb5c 100644
--- a/tools/binman/README.entries
+++ b/tools/binman/README.entries
@@ -26,6 +26,15 @@
 
 
 
+Entry: blob-dtb: A blob that holds a device tree
+------------------------------------------------
+
+This is a blob containing a device tree. The contents of the blob are
+obtained from the list of available device-tree files, managed by the
+'state' module.
+
+
+
 Entry: blob-named-by-arg: A blob entry which gets its filename property from its subclass
 -----------------------------------------------------------------------------------------
 
@@ -309,6 +318,9 @@
 U-Boot. U-Boot needs this to know what devices are present and which drivers
 to activate.
 
+Note: This is mostly an internal entry type, used by others. This allows
+binman to know which entries contain a device tree.
+
 
 
 Entry: u-boot-dtb-with-ucode: A U-Boot device tree file, with the microcode removed
diff --git a/tools/binman/control.py b/tools/binman/control.py
index 34ec74b..e326456 100644
--- a/tools/binman/control.py
+++ b/tools/binman/control.py
@@ -116,10 +116,8 @@
             # output into a file in our output directly. Then scan it for use
             # in binman.
             dtb_fname = fdt_util.EnsureCompiled(dtb_fname)
-            fname = tools.GetOutputFilename('u-boot-out.dtb')
-            with open(dtb_fname) as infd:
-                with open(fname, 'wb') as outfd:
-                    outfd.write(infd.read())
+            fname = tools.GetOutputFilename('u-boot.dtb.out')
+            tools.WriteFile(fname, tools.ReadFile(dtb_fname))
             dtb = fdt.FdtScan(fname)
 
             node = _FindBinmanNode(dtb)
diff --git a/tools/binman/entry.py b/tools/binman/entry.py
index e5f55774..f922107 100644
--- a/tools/binman/entry.py
+++ b/tools/binman/entry.py
@@ -192,6 +192,17 @@
         state.SetInt(self._node, 'image-pos', self.image_pos)
 
     def ProcessFdt(self, fdt):
+        """Allow entries to adjust the device tree
+
+        Some entries need to adjust the device tree for their purposes. This
+        may involve adding or deleting properties.
+
+        Returns:
+            True if processing is complete
+            False if processing could not be completed due to a dependency.
+                This will cause the entry to be retried after others have been
+                called
+        """
         return True
 
     def SetPrefix(self, prefix):
diff --git a/tools/binman/etype/blob_dtb.py b/tools/binman/etype/blob_dtb.py
new file mode 100644
index 0000000..cc5b4a3
--- /dev/null
+++ b/tools/binman/etype/blob_dtb.py
@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: GPL-2.0+
+# Copyright (c) 2018 Google, Inc
+# Written by Simon Glass <sjg@chromium.org>
+#
+# Entry-type module for U-Boot device tree files
+#
+
+import state
+
+from entry import Entry
+from blob import Entry_blob
+
+class Entry_blob_dtb(Entry_blob):
+    """A blob that holds a device tree
+
+    This is a blob containing a device tree. The contents of the blob are
+    obtained from the list of available device-tree files, managed by the
+    'state' module.
+    """
+    def __init__(self, section, etype, node):
+        Entry_blob.__init__(self, section, etype, node)
+
+    def ObtainContents(self):
+        """Get the device-tree from the list held by the 'state' module"""
+        self._filename = self.GetDefaultFilename()
+        self._pathname, data = state.GetFdtContents(self._filename)
+        self.SetContents(data)
+        return True
+
+    def ProcessContents(self):
+        """Re-read the DTB contents so that we get any calculated properties"""
+        _, data = state.GetFdtContents(self._filename)
+        self.SetContents(data)
diff --git a/tools/binman/etype/section.py b/tools/binman/etype/section.py
index f5b2ed6..a30cc91 100644
--- a/tools/binman/etype/section.py
+++ b/tools/binman/etype/section.py
@@ -34,6 +34,9 @@
         Entry.__init__(self, section, etype, node)
         self._section = bsection.Section(node.name, node)
 
+    def GetFdtSet(self):
+        return self._section.GetFdtSet()
+
     def ProcessFdt(self, fdt):
         return self._section.ProcessFdt(fdt)
 
diff --git a/tools/binman/etype/u_boot_dtb.py b/tools/binman/etype/u_boot_dtb.py
index fb3dd1c..6263c4e 100644
--- a/tools/binman/etype/u_boot_dtb.py
+++ b/tools/binman/etype/u_boot_dtb.py
@@ -6,9 +6,9 @@
 #
 
 from entry import Entry
-from blob import Entry_blob
+from blob_dtb import Entry_blob_dtb
 
-class Entry_u_boot_dtb(Entry_blob):
+class Entry_u_boot_dtb(Entry_blob_dtb):
     """U-Boot device tree
 
     Properties / Entry arguments:
@@ -17,9 +17,12 @@
     This is the U-Boot device tree, containing configuration information for
     U-Boot. U-Boot needs this to know what devices are present and which drivers
     to activate.
+
+    Note: This is mostly an internal entry type, used by others. This allows
+    binman to know which entries contain a device tree.
     """
     def __init__(self, section, etype, node):
-        Entry_blob.__init__(self, section, etype, node)
+        Entry_blob_dtb.__init__(self, section, etype, node)
 
     def GetDefaultFilename(self):
         return 'u-boot.dtb'
diff --git a/tools/binman/etype/u_boot_dtb_with_ucode.py b/tools/binman/etype/u_boot_dtb_with_ucode.py
index 3fab766..73f5fbf 100644
--- a/tools/binman/etype/u_boot_dtb_with_ucode.py
+++ b/tools/binman/etype/u_boot_dtb_with_ucode.py
@@ -6,11 +6,11 @@
 #
 
 from entry import Entry
-from blob import Entry_blob
+from blob_dtb import Entry_blob_dtb
 import state
 import tools
 
-class Entry_u_boot_dtb_with_ucode(Entry_blob):
+class Entry_u_boot_dtb_with_ucode(Entry_blob_dtb):
     """A U-Boot device tree file, with the microcode removed
 
     Properties / Entry arguments:
@@ -25,7 +25,7 @@
     it available to u_boot_ucode.
     """
     def __init__(self, section, etype, node):
-        Entry_blob.__init__(self, section, etype, node)
+        Entry_blob_dtb.__init__(self, section, etype, node)
         self.ucode_data = ''
         self.collate = False
         self.ucode_offset = None
@@ -69,7 +69,7 @@
 
     def ObtainContents(self):
         # Call the base class just in case it does something important.
-        Entry_blob.ObtainContents(self)
+        Entry_blob_dtb.ObtainContents(self)
         self._pathname = state.GetFdtPath(self._filename)
         self.ReadBlobContents()
         if self.ucode:
diff --git a/tools/binman/etype/u_boot_spl_dtb.py b/tools/binman/etype/u_boot_spl_dtb.py
index cb29ba3..e735464 100644
--- a/tools/binman/etype/u_boot_spl_dtb.py
+++ b/tools/binman/etype/u_boot_spl_dtb.py
@@ -6,9 +6,9 @@
 #
 
 from entry import Entry
-from blob import Entry_blob
+from blob_dtb import Entry_blob_dtb
 
-class Entry_u_boot_spl_dtb(Entry_blob):
+class Entry_u_boot_spl_dtb(Entry_blob_dtb):
     """U-Boot SPL device tree
 
     Properties / Entry arguments:
@@ -19,7 +19,7 @@
     to activate.
     """
     def __init__(self, section, etype, node):
-        Entry_blob.__init__(self, section, etype, node)
+        Entry_blob_dtb.__init__(self, section, etype, node)
 
     def GetDefaultFilename(self):
         return 'spl/u-boot-spl.dtb'
diff --git a/tools/binman/etype/u_boot_tpl_dtb.py b/tools/binman/etype/u_boot_tpl_dtb.py
index 9c4e668..bdeb0f7 100644
--- a/tools/binman/etype/u_boot_tpl_dtb.py
+++ b/tools/binman/etype/u_boot_tpl_dtb.py
@@ -6,9 +6,9 @@
 #
 
 from entry import Entry
-from blob import Entry_blob
+from blob_dtb import Entry_blob_dtb
 
-class Entry_u_boot_tpl_dtb(Entry_blob):
+class Entry_u_boot_tpl_dtb(Entry_blob_dtb):
     """U-Boot TPL device tree
 
     Properties / Entry arguments:
@@ -19,7 +19,7 @@
     to activate.
     """
     def __init__(self, section, etype, node):
-        Entry_blob.__init__(self, section, etype, node)
+        Entry_blob_dtb.__init__(self, section, etype, node)
 
     def GetDefaultFilename(self):
         return 'tpl/u-boot-tpl.dtb'
diff --git a/tools/binman/ftest.py b/tools/binman/ftest.py
index 75e9a21..6bfef7b 100644
--- a/tools/binman/ftest.py
+++ b/tools/binman/ftest.py
@@ -225,8 +225,26 @@
             TestFunctional._MakeInputFile(outfile, data)
             return data
 
+    def _GetDtbContentsForSplTpl(self, dtb_data, name):
+        """Create a version of the main DTB for SPL or SPL
+
+        For testing we don't actually have different versions of the DTB. With
+        U-Boot we normally run fdtgrep to remove unwanted nodes, but for tests
+        we don't normally have any unwanted nodes.
+
+        We still want the DTBs for SPL and TPL to be different though, since
+        otherwise it is confusing to know which one we are looking at. So add
+        an 'spl' or 'tpl' property to the top-level node.
+        """
+        dtb = fdt.Fdt.FromData(dtb_data)
+        dtb.Scan()
+        dtb.GetNode('/binman').AddZeroProp(name)
+        dtb.Sync(auto_resize=True)
+        dtb.Pack()
+        return dtb.GetContents()
+
     def _DoReadFileDtb(self, fname, use_real_dtb=False, map=False,
-                       update_dtb=False, entry_args=None):
+                       update_dtb=False, entry_args=None, reset_dtbs=True):
         """Run binman and return the resulting image
 
         This runs binman with a given test file and then reads the resulting
@@ -256,12 +274,21 @@
         # Use the compiled test file as the u-boot-dtb input
         if use_real_dtb:
             dtb_data = self._SetupDtb(fname)
+            infile = os.path.join(self._indir, 'u-boot.dtb')
+
+            # For testing purposes, make a copy of the DT for SPL and TPL. Add
+            # a node indicating which it is, so aid verification.
+            for name in ['spl', 'tpl']:
+                dtb_fname = '%s/u-boot-%s.dtb' % (name, name)
+                outfile = os.path.join(self._indir, dtb_fname)
+                TestFunctional._MakeInputFile(dtb_fname,
+                        self._GetDtbContentsForSplTpl(dtb_data, name))
 
         try:
             retcode = self._DoTestFile(fname, map=map, update_dtb=update_dtb,
-                                       entry_args=entry_args)
+                    entry_args=entry_args, use_real_dtb=use_real_dtb)
             self.assertEqual(0, retcode)
-            out_dtb_fname = state.GetFdtPath('u-boot.dtb')
+            out_dtb_fname = tools.GetOutputFilename('u-boot.dtb.out')
 
             # Find the (only) image, read it and return its contents
             image = control.images['image']
@@ -277,7 +304,7 @@
                 return fd.read(), dtb_data, map_data, out_dtb_fname
         finally:
             # Put the test file back
-            if use_real_dtb:
+            if reset_dtbs and use_real_dtb:
                 self._ResetDtbs()
 
     def _DoReadFile(self, fname, use_real_dtb=False):
@@ -1404,6 +1431,80 @@
         self.assertFalse(os.path.exists(tools.GetOutputFilename('image1.bin')))
         self.assertTrue(os.path.exists(tools.GetOutputFilename('image2.bin')))
 
+    def testUpdateFdtAll(self):
+        """Test that all device trees are updated with offset/size info"""
+        data, _, _, _ = self._DoReadFileDtb('82_fdt_update_all.dts',
+                                            use_real_dtb=True, update_dtb=True)
+
+        base_expected = {
+            'section:image-pos': 0,
+            'u-boot-tpl-dtb:size': 513,
+            'u-boot-spl-dtb:size': 513,
+            'u-boot-spl-dtb:offset': 493,
+            'image-pos': 0,
+            'section/u-boot-dtb:image-pos': 0,
+            'u-boot-spl-dtb:image-pos': 493,
+            'section/u-boot-dtb:size': 493,
+            'u-boot-tpl-dtb:image-pos': 1006,
+            'section/u-boot-dtb:offset': 0,
+            'section:size': 493,
+            'offset': 0,
+            'section:offset': 0,
+            'u-boot-tpl-dtb:offset': 1006,
+            'size': 1519
+        }
+
+        # We expect three device-tree files in the output, one after the other.
+        # Read them in sequence. We look for an 'spl' property in the SPL tree,
+        # and 'tpl' in the TPL tree, to make sure they are distinct from the
+        # main U-Boot tree. All three should have the same postions and offset.
+        start = 0
+        for item in ['', 'spl', 'tpl']:
+            dtb = fdt.Fdt.FromData(data[start:])
+            dtb.Scan()
+            props = self._GetPropTree(dtb, ['offset', 'size', 'image-pos',
+                                            'spl', 'tpl'])
+            expected = dict(base_expected)
+            if item:
+                expected[item] = 0
+            self.assertEqual(expected, props)
+            start += dtb._fdt_obj.totalsize()
+
+    def testUpdateFdtOutput(self):
+        """Test that output DTB files are updated"""
+        try:
+            data, dtb_data, _, _ = self._DoReadFileDtb('82_fdt_update_all.dts',
+                    use_real_dtb=True, update_dtb=True, reset_dtbs=False)
+
+            # Unfortunately, compiling a source file always results in a file
+            # called source.dtb (see fdt_util.EnsureCompiled()). The test
+            # source file (e.g. test/75_fdt_update_all.dts) thus does not enter
+            # binman as a file called u-boot.dtb. To fix this, copy the file
+            # over to the expected place.
+            #tools.WriteFile(os.path.join(self._indir, 'u-boot.dtb'),
+                    #tools.ReadFile(tools.GetOutputFilename('source.dtb')))
+            start = 0
+            for fname in ['u-boot.dtb.out', 'spl/u-boot-spl.dtb.out',
+                          'tpl/u-boot-tpl.dtb.out']:
+                dtb = fdt.Fdt.FromData(data[start:])
+                size = dtb._fdt_obj.totalsize()
+                pathname = tools.GetOutputFilename(os.path.split(fname)[1])
+                outdata = tools.ReadFile(pathname)
+                name = os.path.split(fname)[0]
+
+                if name:
+                    orig_indata = self._GetDtbContentsForSplTpl(dtb_data, name)
+                else:
+                    orig_indata = dtb_data
+                self.assertNotEqual(outdata, orig_indata,
+                        "Expected output file '%s' be updated" % pathname)
+                self.assertEqual(outdata, data[start:start + size],
+                        "Expected output file '%s' to match output image" %
+                        pathname)
+                start += size
+        finally:
+            self._ResetDtbs()
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tools/binman/image.py b/tools/binman/image.py
index 1fb5eb6..bfb2229 100644
--- a/tools/binman/image.py
+++ b/tools/binman/image.py
@@ -70,6 +70,11 @@
         self._section.AddMissingProperties()
 
     def ProcessFdt(self, fdt):
+        """Allow entries to adjust the device tree
+
+        Some entries need to adjust the device tree for their purposes. This
+        may involve adding or deleting properties.
+        """
         return self._section.ProcessFdt(fdt)
 
     def GetEntryContents(self):
diff --git a/tools/binman/state.py b/tools/binman/state.py
index b27eb07..09ead44 100644
--- a/tools/binman/state.py
+++ b/tools/binman/state.py
@@ -59,6 +59,29 @@
     """
     return fdt_files[fname]._fname
 
+def GetFdtContents(fname):
+    """Looks up the FDT pathname and contents
+
+    This is used to obtain the Fdt pathname and contents when needed by an
+    entry. It supports a 'fake' dtb, allowing tests to substitute test data for
+    the real dtb.
+
+    Args:
+        fname: Filename to look up (e.g. 'u-boot.dtb').
+
+    Returns:
+        tuple:
+            pathname to Fdt
+            Fdt data (as bytes)
+    """
+    if fname in fdt_files and not use_fake_dtb:
+        pathname = GetFdtPath(fname)
+        data = GetFdt(fname).GetContents()
+    else:
+        pathname = tools.GetInputFilename(fname)
+        data = tools.ReadFile(pathname)
+    return pathname, data
+
 def SetEntryArgs(args):
     """Set the value of the entry args
 
@@ -133,6 +156,8 @@
         Device trees being used (U-Boot proper, SPL, TPL)
     """
     yield main_dtb
+    for other_fname in fdt_subset:
+        yield fdt_files[other_fname]
 
 def GetUpdateNodes(node):
     """Yield all the nodes that need to be updated in all device trees
@@ -149,6 +174,11 @@
             is node, SPL and TPL)
     """
     yield node
+    for dtb in fdt_files.values():
+        if dtb != node.GetFdt():
+            other_node = dtb.GetNode(node.path)
+            if other_node:
+                yield other_node
 
 def AddZeroProp(node, prop):
     """Add a new property to affected device trees with an integer value of 0.
diff --git a/tools/binman/test/82_fdt_update_all.dts b/tools/binman/test/82_fdt_update_all.dts
new file mode 100644
index 0000000..284975c
--- /dev/null
+++ b/tools/binman/test/82_fdt_update_all.dts
@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: GPL-2.0+
+/dts-v1/;
+
+/ {
+	#address-cells = <1>;
+	#size-cells = <1>;
+
+	binman {
+		section {
+			u-boot-dtb {
+			};
+		};
+		u-boot-spl-dtb {
+		};
+		u-boot-tpl-dtb {
+		};
+	};
+};