5 files changed, 23 insertions, 18 deletions

diff --git a/src/arch/i386/image/com32.c b/src/arch/i386/image/com32.c
index 27ccbc02..3e81f485 100644
--- a/src/arch/i386/image/com32.c
+++ b/src/arch/i386/image/com32.c
@@ -122,6 +122,7 @@ static int com32_exec ( struct image *image ) {
 		DBGC ( image, "COM32 %p: exited to run kernel %p\n",
 		       image, comboot_replacement_image );
 		image->replacement = comboot_replacement_image;
+		comboot_replacement_image = NULL;
 		image_autoload ( image->replacement );
 		break;
 
diff --git a/src/arch/i386/image/comboot.c b/src/arch/i386/image/comboot.c
index 3ca9d28f..e0452647 100644
--- a/src/arch/i386/image/comboot.c
+++ b/src/arch/i386/image/comboot.c
@@ -191,6 +191,7 @@ static int comboot_exec ( struct image *image ) {
 		DBGC ( image, "COMBOOT %p: exited to run kernel %p\n",
 		       image, comboot_replacement_image );
 		image->replacement = comboot_replacement_image;
+		comboot_replacement_image = NULL;
 		image_autoload ( image->replacement );
 		break;
 
diff --git a/src/arch/i386/interface/syslinux/comboot_call.c b/src/arch/i386/interface/syslinux/comboot_call.c
index 02b2250b..f37aa2ae 100644
--- a/src/arch/i386/interface/syslinux/comboot_call.c
+++ b/src/arch/i386/interface/syslinux/comboot_call.c
@@ -157,7 +157,7 @@ void comboot_force_text_mode ( void ) {
  * Fetch kernel and optional initrd
  */
 static int comboot_fetch_kernel ( char *kernel_file, char *cmdline ) {
-	struct image *kernel;
+	struct image *kernel = NULL;
 	struct image *initrd = NULL;
 	char *initrd_file;
 	int rc;
@@ -181,13 +181,13 @@ static int comboot_fetch_kernel ( char *kernel_file, char *cmdline ) {
 		if ( ! initrd ) {
 			DBG ( "COMBOOT: could not allocate initrd\n" );
 			rc = -ENOMEM;
-			goto err_alloc_initrd;
+			goto out;
 		}
 		if ( ( rc = imgfetch ( initrd, initrd_file,
 				       register_image ) ) != 0 ) {
 			DBG ( "COMBOOT: could not fetch initrd: %s\n",
 			      strerror ( rc ) );
-			goto err_fetch_initrd;
+			goto out;
 		}
 
 		/* Restore space after initrd name, if applicable */
@@ -202,32 +202,31 @@ static int comboot_fetch_kernel ( char *kernel_file, char *cmdline ) {
 	if ( ! kernel ) {
 		DBG ( "COMBOOT: could not allocate kernel\n" );
 		rc = -ENOMEM;
-		goto err_alloc_kernel;
+		goto out;
 	}
 	if ( ( rc = imgfetch ( kernel, kernel_file,
 			       register_image ) ) != 0 ) {
 		DBG ( "COMBOOT: could not fetch kernel: %s\n",
 		      strerror ( rc ) );
-		goto err_fetch_kernel;
+		goto out;
 	}
 	if ( ( rc = image_set_cmdline ( kernel, cmdline ) ) != 0 ) {
 		DBG ( "COMBOOT: could not set kernel command line: %s\n",
 		      strerror ( rc ) );
-		goto err_set_cmdline;
+		goto out;
 	}
 
 	/* Store kernel as replacement image */
-	comboot_replacement_image = kernel;
-
-	return 0;
-
- err_set_cmdline:
- err_fetch_kernel:
+	assert ( comboot_replacement_image == NULL );
+	comboot_replacement_image = image_get ( kernel );
+
+ out:
+	/* Drop image references unconditionally; either we want to
+	 * discard them, or they have been registered and we should
+	 * drop out local reference.
+	 */
 	image_put ( kernel );
- err_alloc_kernel:
- err_fetch_initrd:
 	image_put ( initrd );
- err_alloc_initrd:
 	return rc;
 }
 
diff --git a/src/core/image.c b/src/core/image.c
index 741b0547..277d09a9 100644
--- a/src/core/image.c
+++ b/src/core/image.c
@@ -275,8 +275,7 @@ int image_exec ( struct image *image ) {
 	/* Pick up replacement image before we drop the original
 	 * image's temporary reference.
 	 */
-	if ( ( replacement = image->replacement ) != NULL )
-		image_get ( replacement );
+	replacement = image->replacement;
 
 	/* Drop temporary reference to the original image */
 	image_put ( image );
diff --git a/src/include/gpxe/image.h b/src/include/gpxe/image.h
index 0163b082..f8b1482e 100644
--- a/src/include/gpxe/image.h
+++ b/src/include/gpxe/image.h
@@ -53,7 +53,12 @@ struct image {
 	 * style similar to a Unix exec() call) should return from its
 	 * exec() method with the replacement image set to point to
 	 * the new image.  The new image must already be in a suitable
-	 * state for execution.
+	 * state for execution (i.e. loaded).
+	 *
+	 * If an image unregisters itself as a result of being
+	 * executed, it must make sure that its replacement image (if
+	 * any) is registered, otherwise the replacement is likely to
+	 * be freed before it can be executed.
 	 */
 	struct image *replacement;
 };