1 files changed, 26 insertions, 18 deletions

diff --git a/gpxe/src/arch/i386/image/com32.c b/gpxe/src/arch/i386/image/com32.c
index da604625..d1b9a59f 100644
--- a/gpxe/src/arch/i386/image/com32.c
+++ b/gpxe/src/arch/i386/image/com32.c
@@ -52,7 +52,7 @@ static int com32_exec ( struct image *image ) {
 	int state;
 	uint32_t avail_mem_top;
 
-	state = setjmp ( comboot_return );
+	state = rmsetjmp ( comboot_return );
 
 	switch ( state ) {
 	case 0: /* First time through; invoke COM32 program */
@@ -70,20 +70,20 @@ static int com32_exec ( struct image *image ) {
 		}
 
 		DBGC ( image, "COM32 %p: available memory top = 0x%x\n",
-			 image, (int)avail_mem_top );
+		       image, avail_mem_top );
 
 		assert ( avail_mem_top != 0 );
 
 		com32_external_esp = phys_to_virt ( avail_mem_top );
 
 		/* Hook COMBOOT API interrupts */
-		hook_comboot_interrupts( );
+		hook_comboot_interrupts();
 
-		/* Temporarily de-register image, so that a "boot" command
-		 * doesn't throw us into an execution loop.  Hold a reference
-		 * to avoid the image's being freed.
+		/* Unregister image, so that a "boot" command doesn't
+		 * throw us into an execution loop.  We never
+		 * reregister ourselves; COMBOOT images expect to be
+		 * removed on exit.
 		 */
-		image_get ( image );
 		unregister_image ( image );
 
 		__asm__ __volatile__ (
@@ -111,25 +111,33 @@ static int com32_exec ( struct image *image ) {
 			/* %6 */ "r" ( COM32_START_PHYS )
 		:
 			"memory" );
+		DBGC ( image, "COM32 %p: returned\n", image );
 		break;
 
-	case COMBOOT_RETURN_RUN_KERNEL:
-		DBGC ( image, "COM32 %p: returned to run kernel...\n", image );
-		comboot_run_kernel ( );
+	case COMBOOT_EXIT:
+		DBGC ( image, "COM32 %p: exited\n", image );
 		break;
 
-	case COMBOOT_RETURN_EXIT:
+	case COMBOOT_EXIT_RUN_KERNEL:
+		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;
 
-	}
-
-	comboot_force_text_mode ( );
+	case COMBOOT_EXIT_COMMAND:
+		DBGC ( image, "COM32 %p: exited after executing command\n",
+		       image );
+		break;
 
-	DBGC ( image, "COM32 %p returned\n", image );
+	default:
+		assert ( 0 );
+		break;
+	}
 
-	/* Re-register image and return */
-	register_image ( image );
-	image_put ( image );
+	unhook_comboot_interrupts();
+	comboot_force_text_mode();
 
 	return 0;
 }