binman: Support updating the dtb in an ELF file

WIth EFI we must embed the devicetree in an ELF image so that it is loaded
as part of the executable file. We want it to include the binman
definition in there also, which in some cases cannot be created until the
ELF (u-boot) is built. Add an option to binman to support writing the
updated dtb to the ELF file u-boot.out

This is useful with the EFI app, which is always packaged as an ELF file.

Signed-off-by: Simon Glass <sjg@chromium.org>
diff --git a/tools/binman/binman.rst b/tools/binman/binman.rst
index 614df54..35de93b 100644
--- a/tools/binman/binman.rst
+++ b/tools/binman/binman.rst
@@ -818,6 +818,42 @@
    # u_boot_dtsi_options_debug = $(u_boot_dtsi_options_raw)
 
 
+Updating an ELF file
+====================
+
+For the EFI app, where U-Boot is loaded from UEFI and runs as an app, there is
+no way to update the devicetree after U-Boot is built. Normally this works by
+creating a new u-boot.dtb.out with he updated devicetree, which is automatically
+built into the output image. With ELF this is not possible since the ELF is
+not part of an image, just a stand-along file. We must create an updated ELF
+file with the new devicetree.
+
+This is handled by the --update-fdt-in-elf option. It takes four arguments,
+separated by comma:
+
+   infile     - filename of input ELF file, e.g. 'u-boot's
+   outfile    - filename of output ELF file, e.g. 'u-boot.out'
+   begin_sym - symbol at the start of the embedded devicetree, e.g.
+   '__dtb_dt_begin'
+   end_sym   - symbol at the start of the embedded devicetree, e.g.
+   '__dtb_dt_end'
+
+When this flag is used, U-Boot does all the normal packaging, but as an
+additional step, it creates a new ELF file with the new devicetree embedded in
+it.
+
+If logging is enabled you will see a message like this::
+
+   Updating file 'u-boot' with data length 0x400a (16394) between symbols
+   '__dtb_dt_begin' and '__dtb_dt_end'
+
+There must be enough space for the updated devicetree. If not, an error like
+the following is produced::
+
+   ValueError: Not enough space in 'u-boot' for data length 0x400a (16394);
+   size is 0x1744 (5956)
+
+
 Entry Documentation
 ===================
 
diff --git a/tools/binman/cmdline.py b/tools/binman/cmdline.py
index d6156df..23729f1 100644
--- a/tools/binman/cmdline.py
+++ b/tools/binman/cmdline.py
@@ -71,6 +71,8 @@
              'given')
     build_parser.add_argument('-u', '--update-fdt', action='store_true',
         default=False, help='Update the binman node with offset/size info')
+    build_parser.add_argument('--update-fdt-in-elf', type=str,
+        help='Update an ELF file with the output dtb: infile,outfile,begin_sym,end_sym')
 
     entry_parser = subparsers.add_parser('entry-docs',
         help='Write out entry documentation (see entries.rst)')
diff --git a/tools/binman/control.py b/tools/binman/control.py
index 0dbcbc2..a56e65a 100644
--- a/tools/binman/control.py
+++ b/tools/binman/control.py
@@ -595,6 +595,13 @@
             tools.FinaliseOutputDir()
         return 0
 
+    elf_params = None
+    if args.update_fdt_in_elf:
+        elf_params = args.update_fdt_in_elf.split(',')
+        if len(elf_params) != 4:
+            raise ValueError('Invalid args %s to --update-fdt-in-elf: expected infile,outfile,begin_sym,end_sym' %
+                             elf_params)
+
     # Try to figure out which device tree contains our image description
     if args.dt:
         dtb_fname = args.dt
@@ -641,6 +648,10 @@
             for dtb_item in state.GetAllFdts():
                 tools.WriteFile(dtb_item._fname, dtb_item.GetContents())
 
+            if elf_params:
+                data = state.GetFdtForEtype('u-boot-dtb').GetContents()
+                elf.UpdateFile(*elf_params, data)
+
             if missing:
                 tout.Warning("\nSome images are invalid")
 
diff --git a/tools/binman/elf.py b/tools/binman/elf.py
index 4aca4f8..de2bb46 100644
--- a/tools/binman/elf.py
+++ b/tools/binman/elf.py
@@ -348,3 +348,24 @@
                 segment.data()[offset:])
     return ElfInfo(output, data_start, elf.header['e_entry'] + virt_to_phys,
                    mem_end - data_start)
+
+def UpdateFile(infile, outfile, start_sym, end_sym, insert):
+    tout.Notice("Creating file '%s' with data length %#x (%d) between symbols '%s' and '%s'" %
+                (outfile, len(insert), len(insert), start_sym, end_sym))
+    syms = GetSymbolFileOffset(infile, [start_sym, end_sym])
+    if len(syms) != 2:
+        raise ValueError("Expected two symbols '%s' and '%s': got %d: %s" %
+                         (start_sym, end_sym, len(syms),
+                          ','.join(syms.keys())))
+
+    size = syms[end_sym].offset - syms[start_sym].offset
+    if len(insert) > size:
+        raise ValueError("Not enough space in '%s' for data length %#x (%d); size is %#x (%d)" %
+                         (infile, len(insert), len(insert), size, size))
+
+    data = tools.ReadFile(infile)
+    newdata = data[:syms[start_sym].offset]
+    newdata += insert + tools.GetBytes(0, size - len(insert))
+    newdata += data[syms[end_sym].offset:]
+    tools.WriteFile(outfile, newdata)
+    tout.Info('Written to offset %#x' % syms[start_sym].offset)
diff --git a/tools/binman/ftest.py b/tools/binman/ftest.py
index 39a4b94..6be0037 100644
--- a/tools/binman/ftest.py
+++ b/tools/binman/ftest.py
@@ -309,7 +309,7 @@
                     entry_args=None, images=None, use_real_dtb=False,
                     use_expanded=False, verbosity=None, allow_missing=False,
                     extra_indirs=None, threads=None,
-                    test_section_timeout=False):
+                    test_section_timeout=False, update_fdt_in_elf=None):
         """Run binman with a given test file
 
         Args:
@@ -336,6 +336,7 @@
                 single-threaded)
             test_section_timeout: True to force the first time to timeout, as
                 used in testThreadTimeout()
+            update_fdt_in_elf: Value to pass with --update-fdt-in-elf=xxx
 
         Returns:
             int return code, 0 on success
@@ -368,6 +369,8 @@
                 args.append('-a%s=%s' % (arg, value))
         if allow_missing:
             args.append('-M')
+        if update_fdt_in_elf:
+            args += ['--update-fdt-in-elf', update_fdt_in_elf]
         if images:
             for image in images:
                 args += ['-i', image]
@@ -4580,6 +4583,84 @@
         self.assertIn('read:', stdout.getvalue())
         self.assertIn('compress:', stdout.getvalue())
 
+    def testUpdateFdtInElf(self):
+        """Test that we can update the devicetree in an ELF file"""
+        infile = elf_fname = self.ElfTestFile('u_boot_binman_embed')
+        outfile = os.path.join(self._indir, 'u-boot.out')
+        begin_sym = 'dtb_embed_begin'
+        end_sym = 'dtb_embed_end'
+        retcode = self._DoTestFile(
+            '060_fdt_update.dts', update_dtb=True,
+            update_fdt_in_elf=','.join([infile,outfile,begin_sym,end_sym]))
+        self.assertEqual(0, retcode)
+
+        # Check that the output file does in fact contact a dtb with the binman
+        # definition in the correct place
+        syms = elf.GetSymbolFileOffset(infile,
+                                       ['dtb_embed_begin', 'dtb_embed_end'])
+        data = tools.ReadFile(outfile)
+        dtb_data = data[syms['dtb_embed_begin'].offset:
+                        syms['dtb_embed_end'].offset]
+
+        dtb = fdt.Fdt.FromData(dtb_data)
+        dtb.Scan()
+        props = self._GetPropTree(dtb, BASE_DTB_PROPS + REPACK_DTB_PROPS)
+        self.assertEqual({
+            'image-pos': 0,
+            'offset': 0,
+            '_testing:offset': 32,
+            '_testing:size': 2,
+            '_testing:image-pos': 32,
+            'section@0/u-boot:offset': 0,
+            'section@0/u-boot:size': len(U_BOOT_DATA),
+            'section@0/u-boot:image-pos': 0,
+            'section@0:offset': 0,
+            'section@0:size': 16,
+            'section@0:image-pos': 0,
+
+            'section@1/u-boot:offset': 0,
+            'section@1/u-boot:size': len(U_BOOT_DATA),
+            'section@1/u-boot:image-pos': 16,
+            'section@1:offset': 16,
+            'section@1:size': 16,
+            'section@1:image-pos': 16,
+            'size': 40
+        }, props)
+
+    def testUpdateFdtInElfInvalid(self):
+        """Test that invalid args are detected with --update-fdt-in-elf"""
+        with self.assertRaises(ValueError) as e:
+            self._DoTestFile('060_fdt_update.dts', update_fdt_in_elf='fred')
+        self.assertIn("Invalid args ['fred'] to --update-fdt-in-elf",
+                      str(e.exception))
+
+    def testUpdateFdtInElfNoSyms(self):
+        """Test that missing symbols are detected with --update-fdt-in-elf"""
+        infile = elf_fname = self.ElfTestFile('u_boot_binman_embed')
+        outfile = ''
+        begin_sym = 'wrong_begin'
+        end_sym = 'wrong_end'
+        with self.assertRaises(ValueError) as e:
+            self._DoTestFile(
+                '060_fdt_update.dts',
+                update_fdt_in_elf=','.join([infile,outfile,begin_sym,end_sym]))
+        self.assertIn("Expected two symbols 'wrong_begin' and 'wrong_end': got 0:",
+                      str(e.exception))
+
+    def testUpdateFdtInElfTooSmall(self):
+        """Test that an over-large dtb is detected with --update-fdt-in-elf"""
+        infile = elf_fname = self.ElfTestFile('u_boot_binman_embed_sm')
+        outfile = os.path.join(self._indir, 'u-boot.out')
+        begin_sym = 'dtb_embed_begin'
+        end_sym = 'dtb_embed_end'
+        with self.assertRaises(ValueError) as e:
+            self._DoTestFile(
+                '060_fdt_update.dts', update_dtb=True,
+                update_fdt_in_elf=','.join([infile,outfile,begin_sym,end_sym]))
+        self.assertRegex(
+            str(e.exception),
+            "Not enough space in '.*u_boot_binman_embed_sm' for data length.*")
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tools/binman/test/Makefile b/tools/binman/test/Makefile
index 6e6cc6d..387ba16 100644
--- a/tools/binman/test/Makefile
+++ b/tools/binman/test/Makefile
@@ -28,10 +28,12 @@
 LDS_BINMAN := -T $(SRC)u_boot_binman_syms.lds
 LDS_BINMAN_BAD := -T $(SRC)u_boot_binman_syms_bad.lds
 LDS_BINMAN_X86 := -T $(SRC)u_boot_binman_syms_x86.lds
+LDS_BINMAN_EMBED := -T $(SRC)u_boot_binman_embed.lds
 
 TARGETS = u_boot_ucode_ptr u_boot_no_ucode_ptr bss_data \
 	u_boot_binman_syms u_boot_binman_syms.bin u_boot_binman_syms_bad \
-	u_boot_binman_syms_size u_boot_binman_syms_x86 embed_data
+	u_boot_binman_syms_size u_boot_binman_syms_x86 embed_data \
+	u_boot_binman_embed u_boot_binman_embed_sm
 
 all: $(TARGETS)
 
@@ -62,6 +64,12 @@
 u_boot_binman_syms_size: CFLAGS += $(LDS_BINMAN)
 u_boot_binman_syms_size: u_boot_binman_syms_size.c
 
+u_boot_binman_embed: CFLAGS += $(LDS_BINMAN_EMBED)
+u_boot_binman_embed: u_boot_binman_embed.c
+
+u_boot_binman_embed_sm: CFLAGS += $(LDS_BINMAN_EMBED)
+u_boot_binman_embed_sm: u_boot_binman_embed_sm.c
+
 clean:
 	rm -f $(TARGETS)
 
diff --git a/tools/binman/test/u_boot_binman_embed.c b/tools/binman/test/u_boot_binman_embed.c
new file mode 100644
index 0000000..75874bb
--- /dev/null
+++ b/tools/binman/test/u_boot_binman_embed.c
@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: GPL-2.0+
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Simple program to embed a devicetree. This is used by binman tests.
+ */
+
+int __attribute__((section(".mydtb"))) dtb_data[4096];
+
+int main(void)
+{
+	return 0;
+}
diff --git a/tools/binman/test/u_boot_binman_embed.lds b/tools/binman/test/u_boot_binman_embed.lds
new file mode 100644
index 0000000..e213fa8
--- /dev/null
+++ b/tools/binman/test/u_boot_binman_embed.lds
@@ -0,0 +1,29 @@
+/* SPDX-License-Identifier: GPL-2.0+ */
+/*
+ * Copyright (c) 2016 Google, Inc
+ */
+
+OUTPUT_FORMAT("elf32-i386", "elf32-i386", "elf32-i386")
+OUTPUT_ARCH(i386)
+ENTRY(_start)
+
+SECTIONS
+{
+	. = 0x00000000;
+	_start = .;
+
+	. = ALIGN(4);
+	.text :
+	{
+		*(.text*)
+	}
+
+	. = ALIGN(4);
+	.data : {
+		dtb_embed_begin = .;
+		KEEP(*(.mydtb));
+		dtb_embed_end = .;
+	}
+	.interp : { *(.interp*) }
+
+}
diff --git a/tools/binman/test/u_boot_binman_embed_sm.c b/tools/binman/test/u_boot_binman_embed_sm.c
new file mode 100644
index 0000000..ae245d7
--- /dev/null
+++ b/tools/binman/test/u_boot_binman_embed_sm.c
@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: GPL-2.0+
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Simple program to embed a devicetree. This is used by binman tests.
+ */
+
+int __attribute__((section(".mydtb"))) dtb_data[16];
+
+int main(void)
+{
+	return 0;
+}