2 files changed, 18 insertions, 22 deletions

diff --git a/twrp-functions.cpp b/twrp-functions.cpp
index c396a07fd..f773df6d0 100644
--- a/twrp-functions.cpp
+++ b/twrp-functions.cpp
@@ -44,14 +44,12 @@
 #include "variables.h"
 #include "bootloader.h"
 #include "cutils/properties.h"
-#ifdef ANDROID_RB_POWEROFF
-	#include "cutils/android_reboot.h"
-#endif
+#include "cutils/android_reboot.h"
+#include <sys/reboot.h>
 #endif // ndef BUILD_TWRPTAR_MAIN
 #ifndef TW_EXCLUDE_ENCRYPTED_BACKUPS
 	#include "openaes/inc/oaes_lib.h"
 #endif
-#include "cutils/android_reboot.h"
 
 extern "C" {
 	#include "libcrecovery/common.h"
@@ -520,48 +518,51 @@ int TWFunc::tw_reboot(RebootCommand command)
 {
 	// Always force a sync before we reboot
 	sync();
+	Update_Log_File();
 
 	switch (command) {
 		case rb_current:
 		case rb_system:
-			Update_Log_File();
 			Update_Intent_File("s");
 			sync();
 			check_and_run_script("/sbin/rebootsystem.sh", "reboot system");
+#ifdef ANDROID_RB_PROPERTY
+			return property_set(ANDROID_RB_PROPERTY, "reboot,");
+#elif defined(ANDROID_RB_RESTART)
+			return android_reboot(ANDROID_RB_RESTART, 0, 0);
+#else
 			return reboot(RB_AUTOBOOT);
+#endif
 		case rb_recovery:
 			check_and_run_script("/sbin/rebootrecovery.sh", "reboot recovery");
 #ifdef ANDROID_RB_PROPERTY
-			property_set(ANDROID_RB_PROPERTY, "reboot,recovery");
+			return property_set(ANDROID_RB_PROPERTY, "reboot,recovery");
 #else
 			return __reboot(LINUX_REBOOT_MAGIC1, LINUX_REBOOT_MAGIC2, LINUX_REBOOT_CMD_RESTART2, (void*) "recovery");
 #endif
-			sleep(5);
-			return 0;
 		case rb_bootloader:
 			check_and_run_script("/sbin/rebootbootloader.sh", "reboot bootloader");
 #ifdef ANDROID_RB_PROPERTY
-			property_set(ANDROID_RB_PROPERTY, "reboot,bootloader");
+			return property_set(ANDROID_RB_PROPERTY, "reboot,bootloader");
 #else
 			return __reboot(LINUX_REBOOT_MAGIC1, LINUX_REBOOT_MAGIC2, LINUX_REBOOT_CMD_RESTART2, (void*) "bootloader");
 #endif
-			sleep(5);
-			return 0;
 		case rb_poweroff:
 			check_and_run_script("/sbin/poweroff.sh", "power off");
-#ifdef ANDROID_RB_POWEROFF
-			android_reboot(ANDROID_RB_POWEROFF, 0, 0);
-#endif
+#ifdef ANDROID_RB_PROPERTY
+			return property_set(ANDROID_RB_PROPERTY, "shutdown,");
+#elif defined(ANDROID_RB_POWEROFF)
+			return android_reboot(ANDROID_RB_POWEROFF, 0, 0);
+#else
 			return reboot(RB_POWER_OFF);
+#endif
 		case rb_download:
 			check_and_run_script("/sbin/rebootdownload.sh", "reboot download");
 #ifdef ANDROID_RB_PROPERTY
-			property_set(ANDROID_RB_PROPERTY, "reboot,download");
+			return property_set(ANDROID_RB_PROPERTY, "reboot,download");
 #else
 			return __reboot(LINUX_REBOOT_MAGIC1, LINUX_REBOOT_MAGIC2, LINUX_REBOOT_CMD_RESTART2, (void*) "download");
 #endif
-			sleep(5);
-			return 0;
 		default:
 			return -1;
 	}
diff --git a/twrp.cpp b/twrp.cpp
index ef694bf24..ab57f3713 100644
--- a/twrp.cpp
+++ b/twrp.cpp
@@ -369,10 +369,5 @@ int main(int argc, char **argv) {
 	else
 		TWFunc::tw_reboot(rb_system);
 
-#ifdef ANDROID_RB_RESTART
-	android_reboot(ANDROID_RB_RESTART, 0, 0);
-#else
-	reboot(RB_AUTOBOOT);
-#endif
 	return 0;
 }