From 6aa8000e7400c3861d7ede0ee4e206e19085116e Mon Sep 17 00:00:00 2001
From: Stefan Herbrechtsmeier <stefan.herbrechtsmeier@weidmueller.com>
Date: Fri, 19 Aug 2022 16:25:25 +0200
Subject: [PATCH] binman: Add length header attribute to dtb entry

Add an optional length header attribute to the device tree blob entry
class based on the compressed data header from the utilities to compress
and decompress data.

If needed the header could be enabled with the following
attribute beside the compress attribute:
  prepend = "length";

The header was introduced as part of commit eb0f4a4cb402 ("binman:
Support replacing data in a cbfs") to allow device tree entries to be
larger than the compressed contents. Regarding the commit "this is
necessary to cope with a compressed device tree being updated in such a
way that it shrinks after the entry size is already set (an obscure
case)". This case need to be fixed without influence any compressed data
by itself.

Signed-off-by: Stefan Herbrechtsmeier <stefan.herbrechtsmeier@weidmueller.com>
Reviewed-by: Simon Glass <sjg@chromium.org>
---
 tools/binman/entries.rst                      |  3 ++
 tools/binman/etype/blob_dtb.py                | 27 +++++++++++++
 tools/binman/ftest.py                         | 39 +++++++++++++++++++
 .../test/235_compress_dtb_prepend_invalid.dts | 17 ++++++++
 .../test/236_compress_dtb_prepend_length.dts  | 19 +++++++++
 5 files changed, 105 insertions(+)
 create mode 100644 tools/binman/test/235_compress_dtb_prepend_invalid.dts
 create mode 100644 tools/binman/test/236_compress_dtb_prepend_length.dts

diff --git a/tools/binman/entries.rst b/tools/binman/entries.rst
index 4dda34ca01f..b3613d7cbd0 100644
--- a/tools/binman/entries.rst
+++ b/tools/binman/entries.rst
@@ -216,6 +216,9 @@ 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.
 
+Additional Properties / Entry arguments:
+    - prepend: Header type to use:
+        length: 32-bit length header
 
 
 .. _etype_blob_ext:
diff --git a/tools/binman/etype/blob_dtb.py b/tools/binman/etype/blob_dtb.py
index 4159e3032a1..5a6a4547485 100644
--- a/tools/binman/etype/blob_dtb.py
+++ b/tools/binman/etype/blob_dtb.py
@@ -7,6 +7,8 @@
 
 from binman.entry import Entry
 from binman.etype.blob import Entry_blob
+from dtoc import fdt_util
+import struct
 
 # This is imported if needed
 state = None
@@ -17,6 +19,9 @@ class Entry_blob_dtb(Entry_blob):
     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.
+
+    Additional attributes:
+        prepend: Header used (e.g. 'length')
     """
     def __init__(self, section, etype, node):
         # Put this here to allow entry-docs and help to work without libfdt
@@ -24,6 +29,14 @@ class Entry_blob_dtb(Entry_blob):
         from binman import state
 
         super().__init__(section, etype, node)
+        self.prepend = None
+
+    def ReadNode(self):
+        super().ReadNode()
+        self.prepend = fdt_util.GetString(self._node, 'prepend')
+        if self.prepend and self.prepend not in ['length']:
+            self.Raise("Invalid prepend in '%s': '%s'" %
+                       (self._node.name, self.prepend))
 
     def ObtainContents(self):
         """Get the device-tree from the list held by the 'state' module"""
@@ -58,3 +71,17 @@ class Entry_blob_dtb(Entry_blob):
         # will still return the old contents
         state.UpdateFdtContents(self.GetFdtEtype(), data)
         return ok
+
+    def CompressData(self, indata):
+        data = super().CompressData(indata)
+        if self.prepend == 'length':
+            hdr = struct.pack('<I', len(data))
+            data = hdr + data
+        return data
+
+    def DecompressData(self, indata):
+        if self.prepend == 'length':
+            data_len = struct.unpack('<I', indata[:4])[0]
+            indata = indata[4:4 + data_len]
+        data = super().DecompressData(indata)
+        return data
diff --git a/tools/binman/ftest.py b/tools/binman/ftest.py
index 9bce102e77b..7a3e4f8ae0b 100644
--- a/tools/binman/ftest.py
+++ b/tools/binman/ftest.py
@@ -5807,6 +5807,45 @@ fdt         fdtmap                Extract the devicetree blob from the fdtmap
         expect = U_BOOT_SPL_DATA + U_BOOT_DATA
         self.assertEqual(expect, data[:len(expect)])
 
+    def testCompressDtbPrependInvalid(self):
+        """Test that invalid header is detected"""
+        with self.assertRaises(ValueError) as e:
+            self._DoReadFileDtb('235_compress_dtb_prepend_invalid.dts')
+        self.assertIn("Node '/binman/u-boot-dtb': Invalid prepend in "
+                      "'u-boot-dtb': 'invalid'", str(e.exception))
+
+    def testCompressDtbPrependLength(self):
+        """Test that compress with length header works as expected"""
+        data = self._DoReadFileRealDtb('236_compress_dtb_prepend_length.dts')
+        image = control.images['image']
+        entries = image.GetEntries()
+        self.assertIn('u-boot-dtb', entries)
+        u_boot_dtb = entries['u-boot-dtb']
+        self.assertIn('fdtmap', entries)
+        fdtmap = entries['fdtmap']
+
+        image_fname = tools.get_output_filename('image.bin')
+        orig = control.ReadEntry(image_fname, 'u-boot-dtb')
+        dtb = fdt.Fdt.FromData(orig)
+        dtb.Scan()
+        props = self._GetPropTree(dtb, ['size', 'uncomp-size'])
+        expected = {
+            'u-boot:size': len(U_BOOT_DATA),
+            'u-boot-dtb:uncomp-size': len(orig),
+            'u-boot-dtb:size': u_boot_dtb.size,
+            'fdtmap:size': fdtmap.size,
+            'size': len(data),
+            }
+        self.assertEqual(expected, props)
+
+        # Check implementation
+        self.assertEqual(U_BOOT_DATA, data[:len(U_BOOT_DATA)])
+        rest = data[len(U_BOOT_DATA):]
+        comp_data_len = struct.unpack('<I', rest[:4])[0]
+        comp_data = rest[4:4 + comp_data_len]
+        orig2 = self._decompress(comp_data)
+        self.assertEqual(orig, orig2)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tools/binman/test/235_compress_dtb_prepend_invalid.dts b/tools/binman/test/235_compress_dtb_prepend_invalid.dts
new file mode 100644
index 00000000000..ee32670a913
--- /dev/null
+++ b/tools/binman/test/235_compress_dtb_prepend_invalid.dts
@@ -0,0 +1,17 @@
+// SPDX-License-Identifier: GPL-2.0+
+
+/dts-v1/;
+
+/ {
+	#address-cells = <1>;
+	#size-cells = <1>;
+
+	binman {
+		u-boot {
+		};
+		u-boot-dtb {
+			compress = "lz4";
+			prepend = "invalid";
+		};
+	};
+};
diff --git a/tools/binman/test/236_compress_dtb_prepend_length.dts b/tools/binman/test/236_compress_dtb_prepend_length.dts
new file mode 100644
index 00000000000..1570233637a
--- /dev/null
+++ b/tools/binman/test/236_compress_dtb_prepend_length.dts
@@ -0,0 +1,19 @@
+// SPDX-License-Identifier: GPL-2.0+
+
+/dts-v1/;
+
+/ {
+	#address-cells = <1>;
+	#size-cells = <1>;
+
+	binman {
+		u-boot {
+		};
+		u-boot-dtb {
+			compress = "lz4";
+			prepend = "length";
+		};
+		fdtmap {
+		};
+	};
+};