From aa86d1abfa7144fee538f3ffae638ad1004182ee Mon Sep 17 00:00:00 2001
From: Michael Scire <SciresM@gmail.com>
Date: Sat, 10 Nov 2018 02:42:07 -0800
Subject: [PATCH] fatal: Implement PowerControlTask

---
 stratosphere/fatal/source/fatal_main.cpp      |  6 ++
 stratosphere/fatal/source/fatal_task.cpp      |  2 +-
 .../fatal/source/fatal_task_power.cpp         | 75 +++++++++++++++++++
 .../fatal/source/fatal_task_power.hpp         |  3 +
 4 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/stratosphere/fatal/source/fatal_main.cpp b/stratosphere/fatal/source/fatal_main.cpp
index e9fb47c44..e01cbd7e8 100644
--- a/stratosphere/fatal/source/fatal_main.cpp
+++ b/stratosphere/fatal/source/fatal_main.cpp
@@ -77,6 +77,11 @@ void __appInit(void) {
         std::abort();
     }
     
+    rc = psmInitialize();
+    if (R_FAILED(rc)) {
+        std::abort();
+    }
+    
     rc = spsmInitialize();
     if (R_FAILED(rc)) {
         std::abort();
@@ -88,6 +93,7 @@ void __appInit(void) {
 void __appExit(void) {
     /* Cleanup services. */
     spsmExit();
+    psmExit();
     bpcExit();
     pminfoExit();
     setsysExit();
diff --git a/stratosphere/fatal/source/fatal_task.cpp b/stratosphere/fatal/source/fatal_task.cpp
index 407b78e8c..43108e88b 100644
--- a/stratosphere/fatal/source/fatal_task.cpp
+++ b/stratosphere/fatal/source/fatal_task.cpp
@@ -52,7 +52,7 @@ static void RunTask(IFatalTask *task) {
 
 void RunFatalTasks(FatalContext *ctx, u64 title_id, bool error_report, Event *erpt_event, Event *battery_event) {
     RunTask(new ErrorReportTask(ctx, title_id, error_report, erpt_event));
-    /* TODO: RunTask(new PowerControlTask(ctx, title_id, erpt_event, battery_event)); */
+    RunTask(new PowerControlTask(ctx, title_id, erpt_event, battery_event));
     /* TODO: RunTask(new ShowFatalTask(ctx, title_id, battery_event)); */
     /* TODO: RunTask(new StopSoundTask(ctx, title_id)); */
     /* TODO: RunTask(new BacklightControlTask(ctx, title_id)); */
diff --git a/stratosphere/fatal/source/fatal_task_power.cpp b/stratosphere/fatal/source/fatal_task_power.cpp
index 410fc158a..2ebd2489e 100644
--- a/stratosphere/fatal/source/fatal_task_power.cpp
+++ b/stratosphere/fatal/source/fatal_task_power.cpp
@@ -17,6 +17,76 @@
 #include <switch.h>
 #include "fatal_task_power.hpp"
 
+bool PowerControlTask::TryShutdown() {
+    /* Set a timeout of 30 seconds. */
+    TimeoutHelper timeout_helper(30000000000UL); 
+    bool cancel_shutdown = false;
+    PsmBatteryVoltageState bv_state = PsmBatteryVoltageState_Normal;
+    
+    while (true) {
+        if (timeout_helper.TimedOut()) {
+            break;
+        }
+        
+        if (R_FAILED(psmGetBatteryVoltageState(&bv_state)) || bv_state == PsmBatteryVoltageState_NeedsShutdown) {
+            break;
+        }
+        
+        if (bv_state == PsmBatteryVoltageState_Normal) {
+            cancel_shutdown = true;
+            break;
+        }
+        
+        /* Query voltage state every 5 seconds, for 30 seconds. */
+        svcSleepThread(5000000000UL);
+    }
+    
+    if (!cancel_shutdown) {
+        bpcShutdownSystem();
+        return true;
+    } else {
+        return false;
+    }
+}
+
+void PowerControlTask::MonitorBatteryState() {
+    PsmBatteryVoltageState bv_state = PsmBatteryVoltageState_Normal;
+    
+    /* Check the battery state, and shutdown on low voltage. */
+    if (R_FAILED(psmGetBatteryVoltageState(&bv_state)) || bv_state == PsmBatteryVoltageState_NeedsShutdown) {
+        /* Wait a second for the error report task to finish. */
+        eventWait(this->erpt_event, TimeoutHelper::NsToTick(1000000000UL));
+        this->TryShutdown();
+        return;
+    }
+    
+    /* Signal we've checked the battery at least once. */
+    eventFire(this->battery_event);
+    
+    while (true) {
+        if (R_FAILED(psmGetBatteryVoltageState(&bv_state))) {
+            bv_state = PsmBatteryVoltageState_NeedsShutdown;
+        }
+        
+        switch (bv_state) {
+            case PsmBatteryVoltageState_NeedsShutdown:
+            case PsmBatteryVoltageState_NeedsSleep:
+                {
+                    bool shutdown = this->TryShutdown();
+                    if (shutdown) {
+                        return;
+                    }
+                }
+                break;
+            default:
+                break;
+        }
+        
+        /* Query voltage state every 5 seconds. */
+        svcSleepThread(5000000000UL);
+    }
+}
+
 void PowerButtonObserveTask::WaitForPowerButton() {
     /* Wait up to a second for error report generation to finish. */
     eventWait(this->erpt_event, TimeoutHelper::NsToTick(1000000000UL));
@@ -36,6 +106,11 @@ void PowerButtonObserveTask::WaitForPowerButton() {
     }
 }
 
+Result PowerControlTask::Run() {
+    MonitorBatteryState();
+    return 0;
+}
+
 Result PowerButtonObserveTask::Run() {
     WaitForPowerButton();
     return 0;
diff --git a/stratosphere/fatal/source/fatal_task_power.hpp b/stratosphere/fatal/source/fatal_task_power.hpp
index bdd5b7155..d31e78752 100644
--- a/stratosphere/fatal/source/fatal_task_power.hpp
+++ b/stratosphere/fatal/source/fatal_task_power.hpp
@@ -23,6 +23,9 @@ class PowerControlTask : public IFatalTask {
     private:
         Event *erpt_event;
         Event *battery_event;
+    private:
+        bool TryShutdown();
+        void MonitorBatteryState();
     public:
         PowerControlTask(FatalContext *ctx, u64 title_id, Event *er_evt, Event *bt_evt) : IFatalTask(ctx, title_id), erpt_event(er_evt), battery_event(bt_evt) { }
         virtual Result Run() override;