3 files changed, 20 insertions, 7 deletions

diff --git a/sys/boot/common/boot.c b/sys/boot/common/boot.c
index f69d4302d16c..4fded2c7c50a 100644
--- a/sys/boot/common/boot.c
+++ b/sys/boot/common/boot.c
@@ -53,6 +53,7 @@ command_boot(int argc, char *argv[])
     struct loaded_module	*km;
     char			*cp;
     int				try;
+    int				i;
     
     /*
      * See if the user has specified an explicit kernel to boot.
@@ -109,6 +110,11 @@ command_boot(int argc, char *argv[])
     if (archsw.arch_autoload() != 0)
 	return(CMD_ERROR);
 
+    /* Call cleanup routines */
+    for (i = 0; devsw[i] != NULL; ++i)
+	if (devsw[i]->dv_cleanup != NULL)
+	    (devsw[i]->dv_cleanup)();
+
     /* Call the exec handler from the loader matching the kernel */
     module_formats[km->m_loader]->l_exec(km);
     return(CMD_ERROR);
diff --git a/sys/boot/i386/libi386/pxe.c b/sys/boot/i386/libi386/pxe.c
index a2d56ff09390..9d5ab08a1a49 100644
--- a/sys/boot/i386/libi386/pxe.c
+++ b/sys/boot/i386/libi386/pxe.c
@@ -128,7 +128,8 @@ struct devsw pxedisk = {
 	pxe_open, 
 	pxe_close, 
 	noioctl,
-	pxe_print
+	pxe_print,
+	pxe_cleanup
 };
 
 /*
@@ -325,14 +326,15 @@ pxe_cleanup(void)
 	t_PXENV_UNDI_SHUTDOWN *undi_shutdown_p =
 	    (t_PXENV_UNDI_SHUTDOWN *)scratch_buffer;
 
-	pxe_call(PXENV_UNLOAD_STACK);
-	if (unload_stack_p->Status != 0)
-		panic("pxe_cleanup: UNLOAD_STACK failed");
-	
 	pxe_call(PXENV_UNDI_SHUTDOWN);
 	if (undi_shutdown_p->Status != 0)
-		panic("pxe_cleanup: UNDI_SHUTDOWN failed");
-	printf("All cleaned up!\n");
+		panic("pxe_cleanup: UNDI_SHUTDOWN failed %x",
+		    undi_shutdown_p->Status);
+
+	pxe_call(PXENV_UNLOAD_STACK);
+	if (unload_stack_p->Status != 0)
+		panic("pxe_cleanup: UNLOAD_STACK failed %x",
+		    unload_stack_p->Status);
 }
 
 void
diff --git a/sys/boot/i386/loader/main.c b/sys/boot/i386/loader/main.c
index db35e597d08e..721a664ff2c2 100644
--- a/sys/boot/i386/loader/main.c
+++ b/sys/boot/i386/loader/main.c
@@ -225,6 +225,11 @@ COMMAND_SET(reboot, "reboot", "reboot the system", command_reboot);
 static int
 command_reboot(int argc, char *argv[])
 {
+    int i;
+
+    for (i = 0; devsw[i] != NULL; ++i)
+	if (devsw[i]->dv_cleanup != NULL)
+	    (devsw[i]->dv_cleanup)();
 
     printf("Rebooting...\n");
     delay(1000000);