Ensure that apps launched by shim get correct BS->Exit() behavior

Right now applications run by shim get our wrapper for Exit(), but it
doesn't do as much cleanup as it should - shim itself also exits, but
currently is not doing all the cleanup it should be doing.

This changes it so all of shim's cleanup is also performed.

Based on a patch and lots of review from Gary Lin.

Signed-off-by: Peter Jones <pjones@redhat.com>

replacements.c

@@ -73,7 +73,6 @@ unhook_system_services(void)
 	if (!systab)
 		return;
 
-	systab->BootServices->Exit = system_exit;
 	systab->BootServices->LoadImage = system_load_image;
 	systab->BootServices->StartImage = system_start_image;
 	systab->BootServices->ExitBootServices = system_exit_boot_services;
@@ -163,18 +162,30 @@ exit_boot_services(EFI_HANDLE image_key, UINTN map_key)
 
 static EFI_STATUS EFIAPI
 do_exit(EFI_HANDLE ImageHandle, EFI_STATUS ExitStatus,
-     UINTN ExitDataSize, CHAR16 *ExitData)
+	UINTN ExitDataSize, CHAR16 *ExitData)
 {
 	EFI_STATUS status;
-	unhook_system_services();
 
-	status = systab->BootServices->Exit(ImageHandle, ExitStatus, ExitDataSize, ExitData);
+	shim_fini();
-	if (EFI_ERROR(status))
+
-		hook_system_services(systab);
+	status = systab->BootServices->Exit(ImageHandle, ExitStatus,
+					    ExitDataSize, ExitData);
+	if (EFI_ERROR(status)) {
+		EFI_STATUS status2 = shim_init();
+
+		if (EFI_ERROR(status2)) {
+			Print(L"Something has gone seriously wrong: %r\n",
+				status2);
+			Print(L"shim cannot continue, sorry.\n");
+			systab->BootServices->Stall(5000000);
+			systab->RuntimeServices->ResetSystem(
+				EfiResetShutdown,
+				EFI_SECURITY_VIOLATION, 0, NULL);
+		}
+	}
 	return status;
 }
 
-
 void
 hook_system_services(EFI_SYSTEM_TABLE *local_systab)
 {
@@ -201,6 +212,18 @@ hook_system_services(EFI_SYSTEM_TABLE *local_systab)
 	 * and b) we can unwrap when we're done. */
 	system_exit_boot_services = systab->BootServices->ExitBootServices;
 	systab->BootServices->ExitBootServices = exit_boot_services;
+}
+
+void
+unhook_exit(void)
+{
+	systab->BootServices->Exit = system_exit;
+}
+
+void
+hook_exit(EFI_SYSTEM_TABLE *local_systab)
+{
+	systab = local_systab;
 
 	/* we need to hook Exit() so that we can allow users to quit the
 	 * bootloader and still e.g. start a new one or run an internal

replacements.h

@@ -41,6 +41,9 @@ extern int loader_is_participating;
 extern void hook_system_services(EFI_SYSTEM_TABLE *local_systab);
 extern void unhook_system_services(void);
 
+extern void hook_exit(EFI_SYSTEM_TABLE *local_systab);
+extern void unhook_exit(void);
+
 extern EFI_STATUS install_shim_protocols(void);
 extern void uninstall_shim_protocols(void);
 

shim.c

@@ -54,6 +54,7 @@
 #define MOK_MANAGER L"\\MokManager.efi"
 
 static EFI_SYSTEM_TABLE *systab;
+static EFI_HANDLE image_handle;
 static EFI_STATUS (EFIAPI *entry_point) (EFI_HANDLE image_handle, EFI_SYSTEM_TABLE *system_table);
 
 static CHAR16 *second_stage;
@@ -1809,7 +1810,6 @@ EFI_STATUS check_mok_request(EFI_HANDLE image_handle)
 /*
  * Verify that MokSBState is valid, and if appropriate set insecure mode
  */
-
 static EFI_STATUS check_mok_sb (void)
 {
 	EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
@@ -2042,7 +2042,62 @@ uninstall_shim_protocols(void)
 			  &shim_lock_guid, &shim_lock_interface);
 }
 
-EFI_STATUS efi_main (EFI_HANDLE image_handle, EFI_SYSTEM_TABLE *passed_systab)
+EFI_STATUS
+shim_init(void)
+{
+	EFI_STATUS status = EFI_SUCCESS;
+	setup_console(1);
+	setup_verbosity();
+	dprinta(shim_version);
+
+	/* Set the second stage loader */
+	set_second_stage (image_handle);
+
+	if (secure_mode()) {
+		if (vendor_cert_size || vendor_dbx_size) {
+			/*
+			 * If shim includes its own certificates then ensure
+			 * that anything it boots has performed some
+			 * validation of the next image.
+			 */
+			hook_system_services(systab);
+			loader_is_participating = 0;
+		}
+
+		hook_exit(systab);
+
+		status = install_shim_protocols();
+	}
+	return status;
+}
+
+void
+shim_fini(void)
+{
+	if (secure_mode()) {
+		/*
+		 * Remove our protocols
+		 */
+		uninstall_shim_protocols();
+
+		/*
+		 * Remove our hooks from system services.
+		 */
+		unhook_system_services();
+		unhook_exit();
+	}
+
+	/*
+	 * Free the space allocated for the alternative 2nd stage loader
+	 */
+	if (load_options_size > 0 && second_stage)
+		FreePool(second_stage);
+
+	setup_console(0);
+}
+
+EFI_STATUS efi_main (EFI_HANDLE passed_image_handle,
+		     EFI_SYSTEM_TABLE *passed_systab)
 {
 	EFI_STATUS efi_status;
 
@@ -2062,84 +2117,59 @@ EFI_STATUS efi_main (EFI_HANDLE image_handle, EFI_SYSTEM_TABLE *passed_systab)
 	shim_lock_interface.Context = shim_read_header;
 
 	systab = passed_systab;
+	image_handle = passed_image_handle;
 
 	/*
 	 * Ensure that gnu-efi functions are available
 	 */
 	InitializeLib(image_handle, systab);
 
-	setup_console(1);
-	setup_verbosity();
-
-	dprinta(shim_version);
-
-	/* Set the second stage loader */
-	set_second_stage (image_handle);
-
 	/*
 	 * Check whether the user has configured the system to run in
 	 * insecure mode
 	 */
 	check_mok_sb();
 
+	efi_status = shim_init();
+	if (EFI_ERROR(efi_status)) {
+		Print(L"Something has gone seriously wrong: %r\n", efi_status);
+		Print(L"shim cannot continue, sorry.\n");
+		systab->BootServices->Stall(5000000);
+		systab->RuntimeServices->ResetSystem(EfiResetShutdown,
+						     EFI_SECURITY_VIOLATION,
+						     0, NULL);
+	}
+
 	/*
 	 * Tell the user that we're in insecure mode if necessary
 	 */
 	if (user_insecure_mode) {
 		Print(L"Booting in insecure mode\n");
 		uefi_call_wrapper(BS->Stall, 1, 2000000);
-	} else if (secure_mode()) {
-		if (vendor_cert_size || vendor_dbx_size) {
-			/*
-			 * If shim includes its own certificates then ensure
-			 * that anything it boots has performed some
-			 * validation of the next image.
-			 */
-			hook_system_services(systab);
-			loader_is_participating = 0;
-		}
 	}
 
-	efi_status = install_shim_protocols();
-	if (EFI_ERROR(efi_status))
-		return efi_status;
-
 	/*
 	 * Enter MokManager if necessary
 	 */
 	efi_status = check_mok_request(image_handle);
 
 	/*
-	 * Copy the MOK list to a runtime variable so the kernel can make
+	 * Copy the MOK list to a runtime variable so the kernel can
-	 * use of it
+	 * make use of it
 	 */
 	efi_status = mirror_mok_list();
 
 	/*
-	 * Create the runtime MokIgnoreDB variable so the kernel can make
+	 * Create the runtime MokIgnoreDB variable so the kernel can
-	 * use of it
+	 * make use of it
 	 */
 	efi_status = mok_ignore_db();
 
 	/*
 	 * Hand over control to the second stage bootloader
 	 */
-
 	efi_status = init_grub(image_handle);
 
-	uninstall_shim_protocols();
+	shim_fini();
-	/*
-	 * Remove our hooks from system services.
-	 */
-	unhook_system_services();
-
-	/*
-	 * Free the space allocated for the alternative 2nd stage loader
-	 */
-	if (load_options_size > 0)
-		FreePool(second_stage);
-
-	setup_console(0);
-
 	return efi_status;
 }

shim.h

@@ -34,3 +34,6 @@ typedef struct _SHIM_LOCK {
 	EFI_SHIM_LOCK_HASH Hash;
 	EFI_SHIM_LOCK_CONTEXT Context;
 } SHIM_LOCK;
+
+extern EFI_STATUS shim_init(void);
+extern void shim_fini(void);