1 files changed, 39 insertions, 4 deletions

diff --git a/cmd/extension_board.c b/cmd/extension_board.c
index f94abd612d5..2b672d888c6 100644
--- a/cmd/extension_board.c
+++ b/cmd/extension_board.c
@@ -5,7 +5,9 @@
  */
 
 #include <common.h>
+#include <bootdev.h>
 #include <command.h>
+#include <dm.h>
 #include <malloc.h>
 #include <extension_board.h>
 #include <mapmem.h>
@@ -80,8 +82,7 @@ static int do_extension_list(struct cmd_tbl *cmdtp, int flag,
 	return CMD_RET_SUCCESS;
 }
 
-static int do_extension_scan(struct cmd_tbl *cmdtp, int flag,
-			     int argc, char *const argv[])
+static int extension_scan(bool show)
 {
 	struct extension *extension, *next;
 	int extension_num;
@@ -91,12 +92,23 @@ static int do_extension_scan(struct cmd_tbl *cmdtp, int flag,
 		free(extension);
 	}
 	extension_num = extension_board_scan(&extension_list);
+	if (show && extension_num >= 0)
+		printf("Found %d extension board(s).\n", extension_num);
+
+	/* either the number of extensions, or -ve for error */
+	return extension_num;
+}
+
+
+static int do_extension_scan(struct cmd_tbl *cmdtp, int flag,
+			     int argc, char *const argv[])
+{
+	int extension_num;
 
+	extension_num = extension_scan(true);
 	if (extension_num < 0)
 		return CMD_RET_FAILURE;
 
-	printf("Found %d extension board(s).\n", extension_num);
-
 	return CMD_RET_SUCCESS;
 }
 
@@ -166,3 +178,26 @@ U_BOOT_CMD(extension, 3, 1, do_extensionops,
 	"extension list - lists available extension(s) board(s)\n"
 	"extension apply <extension number|all> - applies DT overlays corresponding to extension boards\n"
 );
+
+static int extension_bootdev_hunt(struct bootdev_hunter *info, bool show)
+{
+	int ret;
+
+	ret = env_set_hex("extension_overlay_addr",
+			  env_get_hex("fdtoverlay_addr_r", 0));
+	if (ret)
+		return log_msg_ret("env", ret);
+
+	ret = extension_scan(show);
+	if (ret < 0)
+		return log_msg_ret("ext", ret);
+
+	return 0;
+}
+
+/* extensions should have a uclass - for now we use UCLASS_SIMPLE_BUS uclass */
+BOOTDEV_HUNTER(extension_bootdev_hunter) = {
+	.prio		= BOOTDEVP_1_PRE_SCAN,
+	.uclass		= UCLASS_SIMPLE_BUS,
+	.hunt		= extension_bootdev_hunt,
+};