5 files changed, 27 insertions, 18 deletions

diff --git a/src/core/config.c b/src/core/config.c
index 420268276..220b7cdb1 100644
--- a/src/core/config.c
+++ b/src/core/config.c
@@ -205,3 +205,10 @@ REQUIRE_OBJECT ( gdbidt );
 REQUIRE_OBJECT ( gdbudp );
 REQUIRE_OBJECT ( gdbstub_cmd );
 #endif
+
+/*
+ * Drag in objects that are always required, but not dragged in via
+ * symbol dependencies.
+ *
+ */
+REQUIRE_OBJECT ( device );
diff --git a/src/core/device.c b/src/core/device.c
index b1b148e85..84915c2da 100644
--- a/src/core/device.c
+++ b/src/core/device.c
@@ -20,6 +20,7 @@
 #include <gpxe/list.h>
 #include <gpxe/tables.h>
 #include <gpxe/device.h>
+#include <gpxe/init.h>
 
 /**
  * @file
@@ -68,13 +69,11 @@ static void rootdev_remove ( struct root_device *rootdev ) {
 /**
  * Probe all devices
  *
- * @ret rc		Return status code
- *
  * This initiates probing for all devices in the system.  After this
  * call, the device hierarchy will be populated, and all hardware
  * should be ready to use.
  */
-int probe_devices ( void ) {
+static void probe_devices ( void ) {
 	struct root_device *rootdev;
 	int rc;
 
@@ -84,19 +83,28 @@ int probe_devices ( void ) {
 		if ( ( rc = rootdev_probe ( rootdev ) ) != 0 )
 			list_del ( &rootdev->dev.siblings );
 	}
-	return 0;
 }
 
 /**
  * Remove all devices
  *
  */
-void remove_devices ( void ) {
+static void remove_devices ( int flags ) {
 	struct root_device *rootdev;
 	struct root_device *tmp;
 
+	if ( flags & SHUTDOWN_KEEP_DEVICES ) {
+		DBG ( "Refusing to remove devices on shutdown\n" );
+		return;
+	}
+
 	list_for_each_entry_safe ( rootdev, tmp, &devices, dev.siblings ) {
 		rootdev_remove ( rootdev );
 		list_del ( &rootdev->dev.siblings );
 	}
 }
+
+struct startup_fn startup_devices __startup_fn ( STARTUP_NORMAL ) = {
+	.startup = probe_devices,
+	.shutdown = remove_devices,
+};
diff --git a/src/core/init.c b/src/core/init.c
index ed91bf367..50e199cec 100644
--- a/src/core/init.c
+++ b/src/core/init.c
@@ -79,17 +79,14 @@ void startup ( void ) {
 			startup_fn->startup();
 	}
 
-	/* Probe for all devices.  Treated separately because nothing
-	 * else will drag in device.o
-	 */
-	probe_devices();
-
 	started = 1;
 }
 
 /**
  * Shut down gPXE
  *
+ * @v flags		Shutdown behaviour flags
+ *
  * This function reverses the actions of startup(), and leaves gPXE in
  * a state ready to be removed from memory.  You may call startup()
  * again after calling shutdown().
@@ -97,20 +94,17 @@ void startup ( void ) {
  * Call this function only once, before either exiting main() or
  * starting up a non-returnable image.
  */
-void shutdown ( void ) {
+void shutdown ( int flags ) {
 	struct startup_fn *startup_fn;
 
 	if ( ! started )
 		return;
 
-	/* Remove all devices */
-	remove_devices();
-
 	/* Call registered shutdown functions (in reverse order) */
 	for ( startup_fn = startup_fns_end - 1 ; startup_fn >= startup_fns ;
 	      startup_fn-- ) {
 		if ( startup_fn->shutdown )
-			startup_fn->shutdown();
+			startup_fn->shutdown ( flags );
 	}
 
 	started = 0;
diff --git a/src/core/main.c b/src/core/main.c
index ca62db252..d58922617 100644
--- a/src/core/main.c
+++ b/src/core/main.c
@@ -62,7 +62,7 @@ __cdecl int main ( void ) {
 			shell();
 	}
 
-	shutdown();
+	shutdown ( SHUTDOWN_EXIT | shutdown_exit_flags );
 
 	return 0;
 }
diff --git a/src/core/serial.c b/src/core/serial.c
index 54c229549..97640f931 100644
--- a/src/core/serial.c
+++ b/src/core/serial.c
@@ -224,7 +224,7 @@ static void serial_init ( void ) {
  *	Cleanup our use of the serial port, in particular flush the
  *	output buffer so we don't accidentially lose characters.
  */
-static void serial_fini ( void ) {
+static void serial_fini ( int flags __unused ) {
 	int i, status;
 	/* Flush the output buffer to avoid dropping characters,
 	 * if we are reinitializing the serial port.
@@ -247,6 +247,6 @@ struct init_fn serial_init_fn __init_fn ( INIT_SERIAL ) = {
 };
 
 /** Serial driver startup function */
-struct startup_fn serial_startup_fn __startup_fn ( STARTUP_NORMAL ) = {
+struct startup_fn serial_startup_fn __startup_fn ( STARTUP_EARLY ) = {
 	.shutdown = serial_fini,
 };