3 files changed, 36 insertions, 14 deletions
diff --git a/src/include/ipxe/device.h b/src/include/ipxe/device.h
index 068268ba4..635ce59cc 100644
--- a/src/include/ipxe/device.h
+++ b/src/include/ipxe/device.h
@@ -112,6 +112,24 @@ struct root_driver {
 /** Declare a root device */
 #define __root_device __table_entry ( ROOT_DEVICES, 01 )
 
+extern int device_keep_count;
+
+/**
+ * Prevent devices from being removed on shutdown
+ *
+ */
+static inline void devices_get ( void ) {
+	device_keep_count++;
+}
+
+/**
+ * Allow devices to be removed on shutdown
+ *
+ */
+static inline void devices_put ( void ) {
+	device_keep_count--;
+}
+
 extern struct device * identify_device ( struct interface *intf );
 #define identify_device_TYPE( object_type ) \
 	typeof ( struct device * ( object_type ) )
diff --git a/src/include/ipxe/init.h b/src/include/ipxe/init.h
index 10ff8695d..954cda453 100644
--- a/src/include/ipxe/init.h
+++ b/src/include/ipxe/init.h
@@ -32,16 +32,6 @@ struct init_fn {
 
 /** @} */
 
-/** Shutdown flags */
-enum shutdown_flags {
-	/** Shutdown is in order to exit (return to iPXE's caller) */
-	SHUTDOWN_EXIT = 0x0001,
-	/** Shutdown is in order to boot an OS */
-	SHUTDOWN_BOOT = 0x0002,
-	/** Do not remove devices */
-	SHUTDOWN_KEEP_DEVICES = 0x0004,
-};
-
 /**
  * A startup/shutdown function
  *
@@ -50,7 +40,7 @@ enum shutdown_flags {
  */
 struct startup_fn {
 	void ( * startup ) ( void );
-	void ( * shutdown ) ( int flags );
+	void ( * shutdown ) ( int booting );
 };
 
 /** Startup/shutdown function table */
@@ -76,6 +66,22 @@ struct startup_fn {
 
 extern void initialise ( void );
 extern void startup ( void );
-extern void shutdown ( int flags );
+extern void shutdown ( int booting );
+
+/**
+ * Shut down system for OS boot
+ *
+ */
+static inline void shutdown_boot ( void ) {
+	shutdown ( 1 );
+}
+
+/**
+ * Shut down system for exit back to firmware
+ *
+ */
+static inline void shutdown_exit ( void ) {
+	shutdown ( 0 );
+}
 
 #endif /* _IPXE_INIT_H */
diff --git a/src/include/usr/autoboot.h b/src/include/usr/autoboot.h
index e9abf833f..32763beee 100644
--- a/src/include/usr/autoboot.h
+++ b/src/include/usr/autoboot.h
@@ -12,8 +12,6 @@ FILE_LICENCE ( GPL2_OR_LATER );
 #include <ipxe/in.h>
 struct net_device;
 
-extern int shutdown_exit_flags;
-
 extern int netboot ( struct net_device *netdev );
 extern int autoboot ( void );
 extern int boot_next_server_and_filename ( struct in_addr next_server,