From 17c744c3eaf69e195eee291b06d0c2577b334741 Mon Sep 17 00:00:00 2001
From: Rui Miguel Silva <rui.silva@linaro.org>
Date: Mon, 12 Jun 2023 09:09:19 +0100
Subject: [PATCH] corstone1000: set kernel_addr based on boot_idx

We need to distinguish between boot banks and from which
partition to load the kernel+initramfs to memory.

For that, fetch the boot index, fetch the correspondent
partition, calculate the correct kernel address and
then set the env variable kernel_addr with that value.

Signed-off-by: Rui Miguel Silva <rui.silva@linaro.org>
---
 board/armltd/corstone1000/corstone1000.c | 56 +++++++++++++++++++++++-
 configs/corstone1000_defconfig           |  1 +
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/board/armltd/corstone1000/corstone1000.c b/board/armltd/corstone1000/corstone1000.c
index 1bead7a0a8..a4567449f1 100644
--- a/board/armltd/corstone1000/corstone1000.c
+++ b/board/armltd/corstone1000/corstone1000.c
@@ -5,16 +5,24 @@
  * Rui Miguel Silva <rui.silva@linaro.org>
  */
 
+#include <blk.h>
 #include <common.h>
 #include <cpu_func.h>
 #include <dm.h>
 #include <env.h>
-#include <env.h>
 #include <netdev.h>
+#include <nvmxip.h>
+#include <part.h>
 #include <dm/platform_data/serial_pl01x.h>
 #include <asm/armv8/mmu.h>
 #include <asm/global_data.h>
 
+#define CORSTONE1000_KERNEL_PARTS 2
+#define CORSTONE1000_KERNEL_PRIMARY "kernel_primary"
+#define CORSTONE1000_KERNEL_SECONDARY "kernel_secondary"
+
+static int corstone1000_boot_idx;
+
 static struct mm_region corstone1000_mem_map[] = {
 	{
 		/* CVM */
@@ -103,6 +111,52 @@ void fwu_plat_get_bootidx(uint *boot_idx)
 		*boot_idx = CONFIG_FWU_NUM_BANKS;
 		log_err("corstone1000: failed to read active index\n");
 	}
+}
+
+int board_late_init(void)
+{
+	struct disk_partition part_info;
+	struct udevice *dev, *bdev;
+	struct nvmxip_plat *plat;
+	struct blk_desc *desc;
+	int ret;
+
+	ret = uclass_first_device_err(UCLASS_NVMXIP, &dev);
+	if (ret < 0) {
+		log_err("Cannot find kernel device\n");
+		return ret;
+	}
+
+	plat = dev_get_plat(dev);
+	device_find_first_child(dev, &bdev);
+	desc = dev_get_uclass_plat(bdev);
+	ret = fwu_get_active_index(&corstone1000_boot_idx);
+	if (ret < 0) {
+		log_err("corstone1000: failed to read boot index\n");
+		return ret;
+	}
+
+	if (!corstone1000_boot_idx)
+		ret = part_get_info_by_name(desc, CORSTONE1000_KERNEL_PRIMARY,
+					    &part_info);
+	else
+		ret = part_get_info_by_name(desc, CORSTONE1000_KERNEL_SECONDARY,
+					    &part_info);
+
+	if (ret < 0) {
+		log_err("failed to fetch kernel partition index: %d\n",
+			corstone1000_boot_idx);
+		return ret;
+	}
+
+	ret = 0;
+
+	ret |= env_set_hex("kernel_addr", plat->phys_base +
+			   (part_info.start * part_info.blksz));
+	ret |= env_set_hex("kernel_size", part_info.size * part_info.blksz);
+
+	if (ret < 0)
+		log_err("failed to setup kernel addr and size\n");
 
 	return ret;
 }
diff --git a/configs/corstone1000_defconfig b/configs/corstone1000_defconfig
index 2d391048cd..5be5335bdf 100644
--- a/configs/corstone1000_defconfig
+++ b/configs/corstone1000_defconfig
@@ -20,6 +20,7 @@ CONFIG_CONSOLE_RECORD=y
 CONFIG_LOGLEVEL=7
 # CONFIG_DISPLAY_CPUINFO is not set
 # CONFIG_DISPLAY_BOARDINFO is not set
+CONFIG_BOARD_LATE_INIT=y
 CONFIG_SYS_MAXARGS=64
 CONFIG_SYS_CBSIZE=512
 # CONFIG_CMD_CONSOLE is not set
-- 
2.39.5