binman: Allow entries to expand after packing

Add support for detecting entries that change size after they have already
been packed, and re-running packing when it happens.

This removes the limitation that entry size cannot change after
PackEntries() is called.

Signed-off-by: Simon Glass <sjg@chromium.org>
diff --git a/tools/binman/README b/tools/binman/README
index abbf809..9860633 100644
--- a/tools/binman/README
+++ b/tools/binman/README
@@ -566,7 +566,8 @@
 The default implementatoin does nothing. This can be overriden to adjust the
 contents of an entry in some way. For example, it would be possible to create
 an entry containing a hash of the contents of some other entries. At this
-stage the offset and size of entries should not be adjusted.
+stage the offset and size of entries should not be adjusted unless absolutely
+necessary, since it requires a repack (going back to PackEntries()).
 
 10. WriteSymbols() - write the value of symbols into the U-Boot SPL binary.
 See 'Access to binman entry offsets at run time' below for a description of
diff --git a/tools/binman/bsection.py b/tools/binman/bsection.py
index f49a6e9..9047e55 100644
--- a/tools/binman/bsection.py
+++ b/tools/binman/bsection.py
@@ -45,6 +45,8 @@
         _name_prefix: Prefix to add to the name of all entries within this
             section
         _entries: OrderedDict() of entries
+        _orig_offset: Original offset value read from node
+        _orig_size: Original size value read from node
     """
     def __init__(self, name, parent_section, node, image, test=False):
         global entry
@@ -76,6 +78,8 @@
         """Read properties from the section node"""
         self._offset = fdt_util.GetInt(self._node, 'offset')
         self._size = fdt_util.GetInt(self._node, 'size')
+        self._orig_offset = self._offset
+        self._orig_size = self._size
         self._align_size = fdt_util.GetInt(self._node, 'align-size')
         if tools.NotPowerOfTwo(self._align_size):
             self._Raise("Alignment size %s must be a power of two" %
@@ -257,6 +261,13 @@
             for name, info in offset_dict.items():
                 self._SetEntryOffsetSize(name, *info)
 
+    def ResetForPack(self):
+        """Reset offset/size fields so that packing can be done again"""
+        self._offset = self._orig_offset
+        self._size = self._orig_size
+        for entry in self._entries.values():
+            entry.ResetForPack()
+
     def PackEntries(self):
         """Pack all entries into the section"""
         offset = self._skip_at_start
@@ -325,6 +336,7 @@
         for entry in self._entries.values():
             if not entry.ProcessContents():
                 sizes_ok = False
+                print("Entry '%s' size change" % self._node.path)
         return sizes_ok
 
     def WriteSymbols(self):
diff --git a/tools/binman/control.py b/tools/binman/control.py
index 9022cf7..35faf11 100644
--- a/tools/binman/control.py
+++ b/tools/binman/control.py
@@ -170,21 +170,42 @@
                 # completed and written, but that does not seem important.
                 image.GetEntryContents()
                 image.GetEntryOffsets()
-                try:
-                    image.PackEntries()
-                    image.CheckSize()
-                    image.CheckEntries()
-                except Exception as e:
-                    if args.map:
-                        fname = image.WriteMap()
-                        print("Wrote map file '%s' to show errors"  % fname)
-                    raise
-                image.SetImagePos()
-                if args.update_fdt:
-                    image.SetCalculatedProperties()
-                    for dtb_item in state.GetFdts():
-                        dtb_item.Sync()
-                image.ProcessEntryContents()
+
+                # We need to pack the entries to figure out where everything
+                # should be placed. This sets the offset/size of each entry.
+                # However, after packing we call ProcessEntryContents() which
+                # may result in an entry changing size. In that case we need to
+                # do another pass. Since the device tree often contains the
+                # final offset/size information we try to make space for this in
+                # AddMissingProperties() above. However, if the device is
+                # compressed we cannot know this compressed size in advance,
+                # since changing an offset from 0x100 to 0x104 (for example) can
+                # alter the compressed size of the device tree. So we need a
+                # third pass for this.
+                passes = 3
+                for pack_pass in range(passes):
+                    try:
+                        image.PackEntries()
+                        image.CheckSize()
+                        image.CheckEntries()
+                    except Exception as e:
+                        if args.map:
+                            fname = image.WriteMap()
+                            print("Wrote map file '%s' to show errors"  % fname)
+                        raise
+                    image.SetImagePos()
+                    if args.update_fdt:
+                        image.SetCalculatedProperties()
+                        for dtb_item in state.GetFdts():
+                            dtb_item.Sync()
+                    sizes_ok = image.ProcessEntryContents()
+                    if sizes_ok:
+                        break
+                    image.ResetForPack()
+                if not sizes_ok:
+                    image.Raise('Entries expanded after packing (tried %s passes)' %
+                                passes)
+
                 image.WriteSymbols()
                 image.BuildImage()
                 if args.map:
diff --git a/tools/binman/entry.py b/tools/binman/entry.py
index 7db1402..e38cb71 100644
--- a/tools/binman/entry.py
+++ b/tools/binman/entry.py
@@ -61,6 +61,8 @@
         pad_after: Number of pad bytes after the contents, 0 if none
         data: Contents of entry (string of bytes)
         compress: Compression algoithm used (e.g. 'lz4'), 'none' if none
+        orig_offset: Original offset value read from node
+        orig_size: Original size value read from node
     """
     def __init__(self, section, etype, node, read_node=True, name_prefix=''):
         self.section = section
@@ -153,6 +155,9 @@
             self.Raise("Please use 'offset' instead of 'pos'")
         self.offset = fdt_util.GetInt(self._node, 'offset')
         self.size = fdt_util.GetInt(self._node, 'size')
+        self.orig_offset = self.offset
+        self.orig_size = self.size
+
         self.align = fdt_util.GetInt(self._node, 'align')
         if tools.NotPowerOfTwo(self.align):
             raise ValueError("Node '%s': Alignment %s must be a power of two" %
@@ -255,9 +260,16 @@
             ValueError if the new data size is not the same as the old
         """
         size_ok = True
-        if len(data) != self.contents_size:
+        new_size = len(data)
+        if state.AllowEntryExpansion():
+            if new_size > self.contents_size:
+                print("Entry '%s' size change from %#x to %#x" % (
+                    self._node.path, self.contents_size, new_size))
+                # self.data will indicate the new size needed
+                size_ok = False
+        elif new_size != self.contents_size:
             self.Raise('Cannot update entry size from %d to %d' %
-                       (self.contents_size, len(data)))
+                       (self.contents_size, new_size))
         self.SetContents(data)
         return size_ok
 
@@ -271,6 +283,11 @@
         # No contents by default: subclasses can implement this
         return True
 
+    def ResetForPack(self):
+        """Reset offset/size fields so that packing can be done again"""
+        self.offset = self.orig_offset
+        self.size = self.orig_size
+
     def Pack(self, offset):
         """Figure out how to pack the entry into the section
 
diff --git a/tools/binman/etype/_testing.py b/tools/binman/etype/_testing.py
index 2204362..ae24fe8 100644
--- a/tools/binman/etype/_testing.py
+++ b/tools/binman/etype/_testing.py
@@ -50,6 +50,8 @@
                                                     'bad-update-contents')
         self.return_contents_once = fdt_util.GetBool(self._node,
                                                      'return-contents-once')
+        self.bad_update_contents_twice = fdt_util.GetBool(self._node,
+                                                    'bad-update-contents-twice')
 
         # Set to True when the entry is ready to process the FDT.
         self.process_fdt_ready = False
@@ -71,11 +73,12 @@
         if self.force_bad_datatype:
             self.GetEntryArgsOrProps([EntryArg('test-bad-datatype-arg', bool)])
         self.return_contents = True
+        self.contents = b'a'
 
     def ObtainContents(self):
         if self.return_unknown_contents or not self.return_contents:
             return False
-        self.data = b'a'
+        self.data = self.contents
         self.contents_size = len(self.data)
         if self.return_contents_once:
             self.return_contents = False
@@ -90,7 +93,11 @@
         if self.bad_update_contents:
             # Request to update the contents with something larger, to cause a
             # failure.
-            return self.ProcessContentsUpdate('aa')
+            if self.bad_update_contents_twice:
+                self.contents += b'a'
+            else:
+                self.contents = b'aa'
+            return self.ProcessContentsUpdate(self.contents)
         return True
 
     def ProcessFdt(self, fdt):
diff --git a/tools/binman/etype/section.py b/tools/binman/etype/section.py
index 51eddcd..23bf221 100644
--- a/tools/binman/etype/section.py
+++ b/tools/binman/etype/section.py
@@ -64,6 +64,11 @@
         self._section.GetEntryOffsets()
         return {}
 
+    def ResetForPack(self):
+        """Reset offset/size fields so that packing can be done again"""
+        self._section.ResetForPack()
+        Entry.ResetForPack(self)
+
     def Pack(self, offset):
         """Pack all entries into the section"""
         self._section.PackEntries()
diff --git a/tools/binman/etype/u_boot_with_ucode_ptr.py b/tools/binman/etype/u_boot_with_ucode_ptr.py
index 4104bf8..cb7dbc6 100644
--- a/tools/binman/etype/u_boot_with_ucode_ptr.py
+++ b/tools/binman/etype/u_boot_with_ucode_ptr.py
@@ -49,7 +49,7 @@
     def ProcessContents(self):
         # If the image does not need microcode, there is nothing to do
         if not self.target_offset:
-            return
+            return True
 
         # Get the offset of the microcode
         ucode_entry = self.section.FindEntryType('u-boot-ucode')
diff --git a/tools/binman/ftest.py b/tools/binman/ftest.py
index 1c91734..aae8dbc 100644
--- a/tools/binman/ftest.py
+++ b/tools/binman/ftest.py
@@ -1220,10 +1220,14 @@
 
     def testBadChangeSize(self):
         """Test that trying to change the size of an entry fails"""
-        with self.assertRaises(ValueError) as e:
-            self._DoReadFile('059_change_size.dts', True)
-        self.assertIn("Node '/binman/_testing': Cannot update entry size from "
-                      '1 to 2', str(e.exception))
+        try:
+            state.SetAllowEntryExpansion(False)
+            with self.assertRaises(ValueError) as e:
+                self._DoReadFile('059_change_size.dts', True)
+            self.assertIn("Node '/binman/_testing': Cannot update entry size from 1 to 2",
+                          str(e.exception))
+        finally:
+            state.SetAllowEntryExpansion(True)
 
     def testUpdateFdt(self):
         """Test that we can update the device tree with offset/size info"""
@@ -2117,6 +2121,27 @@
         self.assertIn("Invalid location 'None', expected 'start' or 'end'",
                       str(e.exception))
 
+    def testEntryExpand(self):
+        """Test expanding an entry after it is packed"""
+        data = self._DoReadFile('121_entry_expand.dts')
+        self.assertEqual(b'aa', data[:2])
+        self.assertEqual(U_BOOT_DATA, data[2:2 + len(U_BOOT_DATA)])
+        self.assertEqual(b'aa', data[-2:])
+
+    def testEntryExpandBad(self):
+        """Test expanding an entry after it is packed, twice"""
+        with self.assertRaises(ValueError) as e:
+            self._DoReadFile('122_entry_expand_twice.dts')
+        self.assertIn("Image '/binman': Entries expanded after packing",
+                      str(e.exception))
+
+    def testEntryExpandSection(self):
+        """Test expanding an entry within a section after it is packed"""
+        data = self._DoReadFile('123_entry_expand_section.dts')
+        self.assertEqual(b'aa', data[:2])
+        self.assertEqual(U_BOOT_DATA, data[2:2 + len(U_BOOT_DATA)])
+        self.assertEqual(b'aa', data[-2:])
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tools/binman/image.py b/tools/binman/image.py
index c8bce39..6339d02 100644
--- a/tools/binman/image.py
+++ b/tools/binman/image.py
@@ -55,6 +55,10 @@
             self._filename = filename
         self._section = bsection.Section('main-section', None, self._node, self)
 
+    def Raise(self, msg):
+        """Convenience function to raise an error referencing an image"""
+        raise ValueError("Image '%s': %s" % (self._node.path, msg))
+
     def GetFdtSet(self):
         """Get the set of device tree files used by this image"""
         return self._section.GetFdtSet()
@@ -100,6 +104,10 @@
         """
         self._section.GetEntryOffsets()
 
+    def ResetForPack(self):
+        """Reset offset/size fields so that packing can be done again"""
+        self._section.ResetForPack()
+
     def PackEntries(self):
         """Pack all entries into the image"""
         self._section.PackEntries()
diff --git a/tools/binman/test/121_entry_expand.dts b/tools/binman/test/121_entry_expand.dts
new file mode 100644
index 0000000..ebb7816
--- /dev/null
+++ b/tools/binman/test/121_entry_expand.dts
@@ -0,0 +1,20 @@
+/dts-v1/;
+
+/ {
+	#address-cells = <1>;
+	#size-cells = <1>;
+
+	binman {
+		_testing {
+			bad-update-contents;
+		};
+
+		u-boot {
+		};
+
+		_testing2 {
+			type = "_testing";
+			bad-update-contents;
+		};
+	};
+};
diff --git a/tools/binman/test/122_entry_expand_twice.dts b/tools/binman/test/122_entry_expand_twice.dts
new file mode 100644
index 0000000..258cf85
--- /dev/null
+++ b/tools/binman/test/122_entry_expand_twice.dts
@@ -0,0 +1,21 @@
+/dts-v1/;
+
+/ {
+	#address-cells = <1>;
+	#size-cells = <1>;
+
+	binman {
+		_testing {
+			bad-update-contents;
+			bad-update-contents-twice;
+		};
+
+		u-boot {
+		};
+
+		_testing2 {
+			type = "_testing";
+			bad-update-contents;
+		};
+	};
+};
diff --git a/tools/binman/test/123_entry_expand_section.dts b/tools/binman/test/123_entry_expand_section.dts
new file mode 100644
index 0000000..046f723
--- /dev/null
+++ b/tools/binman/test/123_entry_expand_section.dts
@@ -0,0 +1,22 @@
+/dts-v1/;
+
+/ {
+	#address-cells = <1>;
+	#size-cells = <1>;
+
+	binman {
+		_testing {
+			bad-update-contents;
+		};
+
+		u-boot {
+		};
+
+		section {
+			_testing2 {
+				type = "_testing";
+				bad-update-contents;
+			};
+		};
+	};
+};