[PATCH RFC 2/2] spl: opensbi: fix OS image address retrieval from FIT

[Nikita Shubin]

The current implementation of spl_invoke_opensbi() attempts to locate
the next-stage OS (U-Boot proper or Linux) by parsing the FDT blob
pointed to by spl_image->fdt_addr. This FDT, however, is the one that
will be passed to the OS (e.g. a kernel device tree), not the FIT image
itself. As a result, when a FIT image containing multiple components
(firmware, loadables, fdt) is loaded, U-Boot SPL cannot find the OS node
because it searches in the wrong FDT.

Fix this by storing the load addresses of U-Boot and Linux in
spl_image_info during the loadables phase (spl_load_fit_image()).
In spl_invoke_opensbi(), use these stored addresses if they are
non-zero, and fall back to parsing the FDT only if no address is stored,
maintaining backward compatibility.

This ensures that the correct entry point is passed to OpenSBI,
allowing proper boot flow with FIT images.

Signed-off-by: Nikita Shubin <nikita.shubin at maquefel.me>
---
 common/spl/spl_fit.c     |  8 ++++++++
 common/spl/spl_opensbi.c | 33 +++++++++++++++++++--------------
 include/spl.h            |  8 ++++++++
 3 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/common/spl/spl_fit.c b/common/spl/spl_fit.c
index 46ebcabe56a..3da1533c791 100644
--- a/common/spl/spl_fit.c
+++ b/common/spl/spl_fit.c
@@ -1031,6 +1031,14 @@ int spl_load_fit_image(struct spl_image_info *spl_image,
 				     &img_data, &img_len);
 		if (ret < 0)
 			return ret;
+
+#if CONFIG_IS_ENABLED(OPENSBI)
+		if (fit_image_check_os((const void *)header, ret, IH_OS_LINUX))
+			spl_image->kernel_addr = img_data;
+
+		if (fit_image_check_os((const void *)header, ret, IH_OS_U_BOOT))
+			spl_image->uboot_addr = img_data;
+#endif
 	}
 	spl_image->flags |= SPL_FIT_FOUND;
 
diff --git a/common/spl/spl_opensbi.c b/common/spl/spl_opensbi.c
index 0ed6afeacc6..f7bef67bd89 100644
--- a/common/spl/spl_opensbi.c
+++ b/common/spl/spl_opensbi.c
@@ -46,8 +46,8 @@ static int spl_opensbi_find_os_node(void *blob, int *uboot_node, int os_type)
 
 void __noreturn spl_invoke_opensbi(struct spl_image_info *spl_image)
 {
-	int ret, os_node;
-	ulong os_entry;
+	int ret = 0, os_node;
+	ulong os_entry = 0;
 	int os_type;
 	typedef void __noreturn (*opensbi_entry_t)(ulong hartid, ulong dtb, ulong info);
 	opensbi_entry_t opensbi_entry;
@@ -81,22 +81,27 @@ void __noreturn spl_invoke_opensbi(struct spl_image_info *spl_image)
 	 * The next os image default is u-boot proper, once enable
 	 * OpenSBI OS boot mode, the OS image should be linux.
 	 */
-	if (CONFIG_IS_ENABLED(LOAD_FIT_OPENSBI_OS_BOOT))
+	if (CONFIG_IS_ENABLED(LOAD_FIT_OPENSBI_OS_BOOT)) {
 		os_type = IH_OS_LINUX;
-	else
+		os_entry = spl_image->kernel_addr;
+	} else {
 		os_type = IH_OS_U_BOOT;
-
-	ret = spl_opensbi_find_os_node(spl_image->fdt_addr, &os_node, os_type);
-	if (ret) {
-		pr_err("Can't find %s node for opensbi, %d\n",
-		       genimg_get_os_name(os_type), ret);
-		hang();
+		os_entry = spl_image->uboot_addr;
 	}
 
-	/* Get U-Boot entry point */
-	ret = fit_image_get_entry(spl_image->fdt_addr, os_node, &os_entry);
-	if (ret)
-		ret = fit_image_get_load(spl_image->fdt_addr, os_node, &os_entry);
+	if (!os_entry) {
+		ret = spl_opensbi_find_os_node(spl_image->fdt_addr, &os_node, os_type);
+		if (ret) {
+			pr_err("Can't find %s node for opensbi, %d\n",
+				genimg_get_os_name(os_type), ret);
+			hang();
+		}
+
+		/* Get U-Boot entry point */
+		ret = fit_image_get_entry(spl_image->fdt_addr, os_node, &os_entry);
+		if (ret)
+			ret = fit_image_get_load(spl_image->fdt_addr, os_node, &os_entry);
+	}
 
 	/* Prepare opensbi_info object */
 	opensbi_info.magic = FW_DYNAMIC_INFO_MAGIC_VALUE;
diff --git a/include/spl.h b/include/spl.h
index 5078d7525ab..3f4bfc35941 100644
--- a/include/spl.h
+++ b/include/spl.h
@@ -275,6 +275,10 @@ enum spl_sandbox_flags {
  *	if @fdt_size is non-zero
  * @fdt_start: Pointer to the FDT to be copied (must be set up before calling
  *	spl_reloc_jump()
+ * @uboot_addr: Store U-Boot proper addr during "loadables" phase
+ *  (RISC-V only)
+ * @kernel_addr: Store Linux addr during "loadables" phase
+ *  (RISC-V only)
  * @rcode_buf: Buffer to hold the relocating-jump code
  * @stack_prot: Pointer to the stack-protection value, used to ensure the stack
  *	does not overflow
@@ -288,6 +292,10 @@ struct spl_image_info {
 	ulong entry_point;
 #if CONFIG_IS_ENABLED(LOAD_FIT) || CONFIG_IS_ENABLED(LOAD_FIT_FULL)
 	void *fdt_addr;
+#if CONFIG_IS_ENABLED(OPENSBI)
+	ulong uboot_addr;
+	ulong kernel_addr;
+#endif
 #endif
 #if defined(CONFIG_BOOTM_OPTEE) && defined(CONFIG_ARM) && !defined(CONFIG_ARM64)
 	ulong optee_addr;

-- 
2.52.0