binman: Support updating entries in an existing image

While it is useful and efficient to build images in a single pass from a
unified description, it is sometimes desirable to update the image later.

Add support for replace an existing file with one of the same size. This
avoids needing to repack the file. Support for more advanced updates will
come in future patches.

Signed-off-by: Simon Glass <sjg@chromium.org>
diff --git a/tools/binman/README b/tools/binman/README
index 756c6a0..3522354 100644
--- a/tools/binman/README
+++ b/tools/binman/README
@@ -557,6 +557,18 @@
     $ binman extract -i image.bin "*u-boot*" -O outdir
 
 
+Replacing files in an image
+---------------------------
+
+You can replace files in an existing firmware image created by binman, provided
+that there is an 'fdtmap' entry in the image. For example:
+
+    $ binman replace -i image.bin section/cbfs/u-boot
+
+which will write the contents of the file 'u-boot' from the current directory
+to the that entry. The file must be the same size as the entry being replaced.
+
+
 Logging
 -------
 
@@ -909,7 +921,6 @@
 - Allow easy building of images by specifying just the board name
 - Support building an image for a board (-b) more completely, with a
   configurable build directory
-- Support updating binaries in an image (with no size change / repacking)
 - Support updating binaries in an image (with repacking)
 - Support adding FITs to an image
 - Support for ARM Trusted Firmware (ATF)
diff --git a/tools/binman/control.py b/tools/binman/control.py
index 8700f48..ab94f9d 100644
--- a/tools/binman/control.py
+++ b/tools/binman/control.py
@@ -118,6 +118,32 @@
     return entry.ReadData(decomp)
 
 
+def WriteEntry(image_fname, entry_path, data, decomp=True):
+    """Replace an entry in an image
+
+    This replaces the data in a particular entry in an image. This size of the
+    new data must match the size of the old data
+
+    Args:
+        image_fname: Image filename to process
+        entry_path: Path to entry to extract
+        data: Data to replace with
+        decomp: True to compress the data if needed, False if data is
+            already compressed so should be used as is
+    """
+    tout.Info("WriteEntry '%s', file '%s'" % (entry_path, image_fname))
+    image = Image.FromFile(image_fname)
+    entry = image.FindEntryPath(entry_path)
+    state.PrepareFromLoadedData(image)
+    image.LoadData()
+    tout.Info('Writing data to %s' % entry.GetPath())
+    if not entry.WriteData(data, decomp):
+        entry.Raise('Entry data size does not match, but resize is disabled')
+    tout.Info('Processing image')
+    ProcessImage(image, update_fdt=True, write_map=False, get_contents=False)
+    tout.Info('WriteEntry done')
+
+
 def ExtractEntries(image_fname, output_fname, outdir, entry_paths,
                    decomp=True):
     """Extract the data from one or more entries and write it to files
@@ -238,7 +264,7 @@
     return images
 
 
-def ProcessImage(image, update_fdt, write_map):
+def ProcessImage(image, update_fdt, write_map, get_contents=True):
     """Perform all steps for this image, including checking and # writing it.
 
     This means that errors found with a later image will be reported after
@@ -249,8 +275,11 @@
         image: Image to process
         update_fdt: True to update the FDT wth entry offsets, etc.
         write_map: True to write a map file
+        get_contents: True to get the image contents from files, etc., False if
+            the contents is already present
     """
-    image.GetEntryContents()
+    if get_contents:
+        image.GetEntryContents()
     image.GetEntryOffsets()
 
     # We need to pack the entries to figure out where everything
diff --git a/tools/binman/entry.py b/tools/binman/entry.py
index ddf52d8..07d5838 100644
--- a/tools/binman/entry.py
+++ b/tools/binman/entry.py
@@ -698,6 +698,7 @@
 
     def LoadData(self, decomp=True):
         data = self.ReadData(decomp)
+        self.contents_size = len(data)
         self.ProcessContentsUpdate(data)
         self.Detail('Loaded data size %x' % len(data))
 
@@ -708,3 +709,25 @@
             Image object containing this entry
         """
         return self.section.GetImage()
+
+    def WriteData(self, data, decomp=True):
+        """Write the data to an entry in the image
+
+        This is used when the image has been read in and we want to replace the
+        data for a particular entry in that image.
+
+        The image must be re-packed and written out afterwards.
+
+        Args:
+            data: Data to replace it with
+            decomp: True to compress the data if needed, False if data is
+                already compressed so should be used as is
+
+        Returns:
+            True if the data did not result in a resize of this entry, False if
+                 the entry must be resized
+        """
+        self.contents_size = self.size
+        ok = self.ProcessContentsUpdate(data)
+        self.Detail('WriteData: size=%x, ok=%s' % (len(data), ok))
+        return ok
diff --git a/tools/binman/ftest.py b/tools/binman/ftest.py
index 4715328..e201b74 100644
--- a/tools/binman/ftest.py
+++ b/tools/binman/ftest.py
@@ -2334,7 +2334,7 @@
         image_fname = tools.GetOutputFilename('image.bin')
         image = Image.FromFile(image_fname)
         self.assertTrue(isinstance(image, Image))
-        self.assertEqual('image', image.image_name)
+        self.assertEqual('image', image.image_name[-5:])
 
     def testReadImageFail(self):
         """Test failing to read an image image's FDT map"""
@@ -2720,6 +2720,106 @@
         self.assertEqual(len(U_BOOT_DATA), entry.contents_size)
         self.assertEqual(len(U_BOOT_DATA), entry.size)
 
+    def _RunReplaceCmd(self, entry_name, data, decomp=True):
+        """Replace an entry in an image
+
+        This writes the entry data to update it, then opens the updated file and
+        returns the value that it now finds there.
+
+        Args:
+            entry_name: Entry name to replace
+            data: Data to replace it with
+            decomp: True to compress the data if needed, False if data is
+                already compressed so should be used as is
+
+        Returns:
+            Tuple:
+                data from entry
+                data from fdtmap (excluding header)
+        """
+        dtb_data = self._DoReadFileDtb('132_replace.dts', use_real_dtb=True,
+                                       update_dtb=True)[1]
+
+        self.assertIn('image', control.images)
+        image = control.images['image']
+        entries = image.GetEntries()
+        orig_dtb_data = entries['u-boot-dtb'].data
+        orig_fdtmap_data = entries['fdtmap'].data
+
+        image_fname = tools.GetOutputFilename('image.bin')
+        updated_fname = tools.GetOutputFilename('image-updated.bin')
+        tools.WriteFile(updated_fname, tools.ReadFile(image_fname))
+        control.WriteEntry(updated_fname, entry_name, data, decomp)
+        data = control.ReadEntry(updated_fname, entry_name, decomp)
+
+        # The DT data should not change
+        new_dtb_data = entries['u-boot-dtb'].data
+        self.assertEqual(new_dtb_data, orig_dtb_data)
+        new_fdtmap_data = entries['fdtmap'].data
+        self.assertEqual(new_fdtmap_data, orig_fdtmap_data)
+
+        return data, orig_fdtmap_data[fdtmap.FDTMAP_HDR_LEN:]
+
+    def testReplaceSimple(self):
+        """Test replacing a single file"""
+        expected = b'x' * len(U_BOOT_DATA)
+        data, expected_fdtmap = self._RunReplaceCmd('u-boot', expected)
+        self.assertEqual(expected, data)
+
+        # Test that the state looks right. There should be an FDT for the fdtmap
+        # that we jsut read back in, and it should match what we find in the
+        # 'control' tables. Checking for an FDT that does not exist should
+        # return None.
+        path, fdtmap = state.GetFdtContents('fdtmap')
+        self.assertIsNone(path)
+        self.assertEqual(expected_fdtmap, fdtmap)
+
+        dtb = state.GetFdtForEtype('fdtmap')
+        self.assertEqual(dtb.GetContents(), fdtmap)
+
+        missing_path, missing_fdtmap = state.GetFdtContents('missing')
+        self.assertIsNone(missing_path)
+        self.assertIsNone(missing_fdtmap)
+
+        missing_dtb = state.GetFdtForEtype('missing')
+        self.assertIsNone(missing_dtb)
+
+        self.assertEqual('/binman', state.fdt_path_prefix)
+
+    def testReplaceResizeFail(self):
+        """Test replacing a file by something larger"""
+        expected = U_BOOT_DATA + b'x'
+        with self.assertRaises(ValueError) as e:
+            self._RunReplaceCmd('u-boot', expected)
+        self.assertIn("Node '/u-boot': Entry data size does not match, but resize is disabled",
+                      str(e.exception))
+
+    def testReplaceMulti(self):
+        """Test replacing entry data where multiple images are generated"""
+        data = self._DoReadFileDtb('133_replace_multi.dts', use_real_dtb=True,
+                                   update_dtb=True)[0]
+        expected = b'x' * len(U_BOOT_DATA)
+        updated_fname = tools.GetOutputFilename('image-updated.bin')
+        tools.WriteFile(updated_fname, data)
+        entry_name = 'u-boot'
+        control.WriteEntry(updated_fname, entry_name, expected)
+        data = control.ReadEntry(updated_fname, entry_name)
+        self.assertEqual(expected, data)
+
+        # Check the state looks right.
+        self.assertEqual('/binman/image', state.fdt_path_prefix)
+
+        # Now check we can write the first image
+        image_fname = tools.GetOutputFilename('first-image.bin')
+        updated_fname = tools.GetOutputFilename('first-updated.bin')
+        tools.WriteFile(updated_fname, tools.ReadFile(image_fname))
+        entry_name = 'u-boot'
+        control.WriteEntry(updated_fname, entry_name, expected)
+        data = control.ReadEntry(updated_fname, entry_name)
+        self.assertEqual(expected, data)
+
+        # Check the state looks right.
+        self.assertEqual('/binman/first-image', state.fdt_path_prefix)
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tools/binman/image.py b/tools/binman/image.py
index c990818..5185b68 100644
--- a/tools/binman/image.py
+++ b/tools/binman/image.py
@@ -10,6 +10,7 @@
 from collections import OrderedDict
 import fnmatch
 from operator import attrgetter
+import os
 import re
 import sys
 
@@ -96,6 +97,8 @@
         image.fdtmap_dtb = dtb
         image.fdtmap_data = fdtmap_data
         image._data = data
+        image._filename = fname
+        image.image_name, _ = os.path.splitext(fname)
         return image
 
     def Raise(self, msg):
diff --git a/tools/binman/state.py b/tools/binman/state.py
index 34907d9..08e6279 100644
--- a/tools/binman/state.py
+++ b/tools/binman/state.py
@@ -8,8 +8,10 @@
 import hashlib
 import re
 
+import fdt
 import os
 import tools
+import tout
 
 # Records the device-tree files known to binman, keyed by entry type (e.g.
 # 'u-boot-spl-dtb'). These are the output FDT files, which can be updated by
@@ -22,6 +24,9 @@
 #       Entry object, or None if not known
 output_fdt_info = {}
 
+# Prefix to add to an fdtmap path to turn it into a path to the /binman node
+fdt_path_prefix = ''
+
 # Arguments passed to binman to provide arguments to entries
 entry_args = {}
 
@@ -55,24 +60,6 @@
         return None
     return value[0]
 
-def GetEntryForEtype(etype):
-    """Get the Entry for a particular device-tree filename
-
-    Binman keeps track of at least one device-tree file called u-boot.dtb but
-    can also have others (e.g. for SPL). This function looks up the given
-    filename and returns the associated Fdt object.
-
-    Args:
-        etype: Entry type of device tree (e.g. 'u-boot-dtb')
-
-    Returns:
-        Entry object associated with the entry type, if present in the image
-    """
-    value = output_fdt_info.get(etype);
-    if not value:
-        return None
-    return value[2]
-
 def GetFdtPath(etype):
     """Get the full pathname of a particular Fdt object
 
@@ -153,7 +140,7 @@
         images: List of images being used
         dtb: Main dtb
     """
-    global output_fdt_info, main_dtb
+    global output_fdt_info, main_dtb, fdt_path_prefix
     # Import these here in case libfdt.py is not available, in which case
     # the above help option still works.
     import fdt
@@ -165,6 +152,7 @@
     # was handled just above.
     main_dtb = dtb
     output_fdt_info.clear()
+    fdt_path_prefix = ''
     output_fdt_info['u-boot-dtb'] = [dtb, 'u-boot.dtb', None]
     output_fdt_info['u-boot-spl-dtb'] = [dtb, 'spl/u-boot-spl.dtb', None]
     output_fdt_info['u-boot-tpl-dtb'] = [dtb, 'tpl/u-boot-tpl.dtb', None]
@@ -182,13 +170,57 @@
             other_dtb = fdt.FdtScan(out_fname)
             output_fdt_info[etype] = [other_dtb, out_fname, entry]
 
+def PrepareFromLoadedData(image):
+    """Get device tree files ready for use with a loaded image
+
+    Loaded images are different from images that are being created by binman,
+    since there is generally already an fdtmap and we read the description from
+    that. This provides the position and size of every entry in the image with
+    no calculation required.
+
+    This function uses the same output_fdt_info[] as Prepare(). It finds the
+    device tree files, adds a reference to the fdtmap and sets the FDT path
+    prefix to translate from the fdtmap (where the root node is the image node)
+    to the normal device tree (where the image node is under a /binman node).
+
+    Args:
+        images: List of images being used
+    """
+    global output_fdt_info, main_dtb, fdt_path_prefix
+
+    tout.Info('Preparing device trees')
+    output_fdt_info.clear()
+    fdt_path_prefix = ''
+    output_fdt_info['fdtmap'] = [image.fdtmap_dtb, 'u-boot.dtb', None]
+    main_dtb = None
+    tout.Info("   Found device tree type 'fdtmap' '%s'" % image.fdtmap_dtb.name)
+    for etype, value in image.GetFdts().items():
+        entry, fname = value
+        out_fname = tools.GetOutputFilename('%s.dtb' % entry.etype)
+        tout.Info("   Found device tree type '%s' at '%s' path '%s'" %
+                  (etype, out_fname, entry.GetPath()))
+        entry._filename = entry.GetDefaultFilename()
+        data = entry.ReadData()
+
+        tools.WriteFile(out_fname, data)
+        dtb = fdt.Fdt(out_fname)
+        dtb.Scan()
+        image_node = dtb.GetNode('/binman')
+        if 'multiple-images' in image_node.props:
+            image_node = dtb.GetNode('/binman/%s' % image.image_node)
+        fdt_path_prefix = image_node.path
+        output_fdt_info[etype] = [dtb, None, entry]
+    tout.Info("   FDT path prefix '%s'" % fdt_path_prefix)
+
+
 def GetAllFdts():
     """Yield all device tree files being used by binman
 
     Yields:
         Device trees being used (U-Boot proper, SPL, TPL)
     """
-    yield main_dtb
+    if main_dtb:
+        yield main_dtb
     for etype in output_fdt_info:
         dtb = output_fdt_info[etype][0]
         if dtb != main_dtb:
@@ -211,7 +243,7 @@
     yield node
     for dtb, fname, _ in output_fdt_info.values():
         if dtb != node.GetFdt():
-            other_node = dtb.GetNode(node.path)
+            other_node = dtb.GetNode(fdt_path_prefix + node.path)
             if other_node:
                 yield other_node
 
diff --git a/tools/binman/test/132_replace.dts b/tools/binman/test/132_replace.dts
new file mode 100644
index 0000000..6ebdcda
--- /dev/null
+++ b/tools/binman/test/132_replace.dts
@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: GPL-2.0+
+
+/dts-v1/;
+
+/ {
+	#address-cells = <1>;
+	#size-cells = <1>;
+
+	binman {
+		size = <0xc00>;
+		u-boot {
+		};
+		fdtmap {
+		};
+		u-boot-dtb {
+		};
+		image-header {
+			location = "end";
+		};
+	};
+};
diff --git a/tools/binman/test/133_replace_multi.dts b/tools/binman/test/133_replace_multi.dts
new file mode 100644
index 0000000..38b2f39
--- /dev/null
+++ b/tools/binman/test/133_replace_multi.dts
@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: GPL-2.0+
+
+/dts-v1/;
+
+/ {
+	#address-cells = <1>;
+	#size-cells = <1>;
+
+	binman {
+		multiple-images;
+		first-image {
+			size = <0xc00>;
+			u-boot {
+			};
+			fdtmap {
+			};
+			u-boot-dtb {
+			};
+			image-header {
+				location = "end";
+			};
+		};
+
+		image {
+			fdtmap {
+			};
+			u-boot {
+			};
+			u-boot-dtb {
+			};
+		};
+	};
+};