2 files changed, 64 insertions, 11 deletions

diff --git a/src/core/image.c b/src/core/image.c
index 666ee3de8..1bf5447f0 100644
--- a/src/core/image.c
+++ b/src/core/image.c
@@ -327,6 +327,14 @@ int image_exec ( struct image *image ) {
 	if ( image->flags & IMAGE_AUTO_UNREGISTER )
 		unregister_image ( image );
 
+	/* Debug message for tail-recursion.  Placed here because the
+	 * image_put() may end up freeing the image.
+	 */
+	if ( replacement ) {
+		DBGC ( image, "IMAGE %s replacing self with IMAGE %s\n",
+		       image->name, replacement->name );
+	}
+
 	/* Drop temporary reference to the original image */
 	image_put ( image );
 
@@ -338,12 +346,8 @@ int image_exec ( struct image *image ) {
 	uri_put ( old_cwuri );
 
 	/* Tail-recurse into replacement image, if one exists */
-	if ( replacement ) {
-		DBGC ( image, "IMAGE <freed> replacing self with IMAGE %s\n",
-		       replacement->name );
-		if ( ( rc = image_exec ( replacement ) ) != 0 )
-			return rc;
-	}
+	if ( replacement )
+		return image_exec ( replacement );
 
 	return rc;
 }
diff --git a/src/hci/commands/image_cmd.c b/src/hci/commands/image_cmd.c
index 38d814cde..0a4232f27 100644
--- a/src/hci/commands/image_cmd.c
+++ b/src/hci/commands/image_cmd.c
@@ -26,6 +26,7 @@ FILE_LICENCE ( GPL2_OR_LATER );
 #include <ipxe/image.h>
 #include <ipxe/command.h>
 #include <ipxe/parseopt.h>
+#include <ipxe/shell.h>
 #include <usr/imgmgmt.h>
 
 /** @file
@@ -38,6 +39,8 @@ FILE_LICENCE ( GPL2_OR_LATER );
 struct imgsingle_options {
 	/** Image name */
 	const char *name;
+	/** Replace image */
+	int replace;
 	/** Free image after execution */
 	int autofree;
 };
@@ -46,6 +49,8 @@ struct imgsingle_options {
 static struct option_descriptor imgsingle_opts[] = {
 	OPTION_DESC ( "name", 'n', required_argument,
 		      struct imgsingle_options, name, parse_string ),
+	OPTION_DESC ( "replace", 'r', no_argument,
+		      struct imgsingle_options, replace, parse_flag ),
 	OPTION_DESC ( "autofree", 'a', no_argument,
 		      struct imgsingle_options, autofree, parse_flag ),
 };
@@ -66,7 +71,8 @@ struct imgsingle_descriptor {
 	/** Pre-action to take upon image, or NULL */
 	void ( * preaction ) ( struct image *image );
 	/** Action to take upon image, or NULL */
-	int ( * action ) ( struct image *image );
+	int ( * action ) ( struct image *image,
+			   struct imgsingle_options *opts );
 	/** Verb to describe action */
 	const char *verb;
 };
@@ -145,7 +151,7 @@ static int imgsingle_exec ( int argc, char **argv,
 
 	/* Carry out command action, if applicable */
 	if ( desc->action ) {
-		if ( ( rc = desc->action ( image ) ) != 0 ) {
+		if ( ( rc = desc->action ( image, &opts ) ) != 0 ) {
 			printf ( "Could not %s: %s\n",
 				 desc->verb, strerror ( rc ) );
 			goto err_action;
@@ -188,11 +194,23 @@ static int imgfetch_exec ( int argc, char **argv ) {
 	return imgsingle_exec ( argc, argv, &imgfetch_desc );
 }
 
+/**
+ * "imgselect" command action
+ *
+ * @v image		Image
+ * @v opts		Options
+ * @ret rc		Return status code
+ */
+static int imgselect ( struct image *image,
+		       struct imgsingle_options *opts __unused ) {
+	return image_select ( image );
+}
+
 /** "imgselect" family command descriptor */
 struct imgsingle_descriptor imgselect_desc = {
 	.cmd = &imgsingle_cmd,
 	.acquire = imgacquire,
-	.action = image_select,
+	.action = imgselect,
 	.verb = "select",
 };
 
@@ -211,13 +229,44 @@ static int imgselect_exec ( int argc, char **argv ) {
 static struct command_descriptor imgexec_cmd =
 	COMMAND_DESC ( struct imgsingle_options, imgsingle_opts,
 		       0, MAX_ARGUMENTS,
-		       "[--autofree] [<uri|image> [<arguments>...]]" );
+		       "[--autofree] [--replace] "
+		       "[<uri|image> [<arguments>...]]" );
+
+/**
+ * "imgexec" command action
+ *
+ * @v image		Image
+ * @v opts		Options
+ * @ret rc		Return status code
+ */
+static int imgexec ( struct image *image, struct imgsingle_options *opts ) {
+	int rc;
+
+	/* Perform replacement or execution as applicable */
+	if ( opts->replace ) {
+
+		/* Try to replace image */
+		if ( ( rc = image_replace ( image ) ) != 0 )
+			return rc;
+
+		/* Stop script and tail-recurse into replacement image */
+		shell_stop ( SHELL_STOP_COMMAND_SEQUENCE );
+
+	} else {
+
+		/* Try to execute image */
+		if ( ( rc = image_exec ( image ) ) != 0 )
+			return rc;
+	}
+
+	return 0;
+}
 
 /** "imgexec" family command descriptor */
 struct imgsingle_descriptor imgexec_desc = {
 	.cmd = &imgexec_cmd,
 	.acquire = imgacquire,
-	.action = image_exec,
+	.action = imgexec,
 	.verb = "boot",
 };